chore(snapshot for ai): add timeout to limit waiting for iframes (#36712)

Co-authored-by: ogadra <61941819+ogadra@users.noreply.github.com>

Author: Skn0tt

packages/playwright-core/src/client/page.ts

@@ -832,8 +832,8 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     return result.pdf;
   }
 
-  async _snapshotForAI(): Promise<string> {
-    const result = await this._channel.snapshotForAI();
+  async _snapshotForAI(options: TimeoutOptions = {}): Promise<string> {
+    const result = await this._channel.snapshotForAI({ timeout: this._timeoutSettings.timeout(options) });
     return result.snapshot;
   }
 }

packages/playwright-core/src/protocol/validator.ts

@@ -1454,7 +1454,9 @@ scheme.PagePdfParams = tObject({
 scheme.PagePdfResult = tObject({
   pdf: tBinary,
 });
-scheme.PageSnapshotForAIParams = tOptional(tObject({}));
+scheme.PageSnapshotForAIParams = tObject({
+  timeout: tNumber,
+});
 scheme.PageSnapshotForAIResult = tObject({
   snapshot: tString,
 });

packages/protocol/src/channels.d.ts

@@ -2111,7 +2111,7 @@ export interface PageChannel extends PageEventTarget, EventTargetChannel {
   touchscreenTap(params: PageTouchscreenTapParams, progress?: Progress): Promise<PageTouchscreenTapResult>;
   accessibilitySnapshot(params: PageAccessibilitySnapshotParams, progress?: Progress): Promise<PageAccessibilitySnapshotResult>;
   pdf(params: PagePdfParams, progress?: Progress): Promise<PagePdfResult>;
-  snapshotForAI(params?: PageSnapshotForAIParams, progress?: Progress): Promise<PageSnapshotForAIResult>;
+  snapshotForAI(params: PageSnapshotForAIParams, progress?: Progress): Promise<PageSnapshotForAIResult>;
   startJSCoverage(params: PageStartJSCoverageParams, progress?: Progress): Promise<PageStartJSCoverageResult>;
   stopJSCoverage(params?: PageStopJSCoverageParams, progress?: Progress): Promise<PageStopJSCoverageResult>;
   startCSSCoverage(params: PageStartCSSCoverageParams, progress?: Progress): Promise<PageStartCSSCoverageResult>;
@@ -2540,8 +2540,12 @@ export type PagePdfOptions = {
 export type PagePdfResult = {
   pdf: Binary,
 };
-export type PageSnapshotForAIParams = {};
-export type PageSnapshotForAIOptions = {};
+export type PageSnapshotForAIParams = {
+  timeout: number,
+};
+export type PageSnapshotForAIOptions = {
+
+};
 export type PageSnapshotForAIResult = {
   snapshot: string,
 };

packages/protocol/src/protocol.yml

@@ -1918,6 +1918,8 @@ Page:
 
     snapshotForAI:
       internal: true
+      parameters:
+        timeout: number
       returns:
         snapshot: string
       flags:

tests/page/page-aria-snapshot-ai.spec.ts

@@ -19,8 +19,8 @@ import { asLocator } from 'playwright-core/lib/utils';
 
 import { test as it, expect, unshift } from './pageTest';
 
-function snapshotForAI(page: any): Promise<string> {
-  return page._snapshotForAI();
+function snapshotForAI(page: any, options?: { timeout?: number }): Promise<string> {
+  return page._snapshotForAI(options);
 }
 
 it('should generate refs', async ({ page }) => {
@@ -354,3 +354,23 @@ it('should mark iframe as active when it contains focused element', async ({ pag
         - textbox "Input in iframe" [active] [ref=f1e2]
   `);
 });
+
+it('return empty snapshot when iframe is not loaded', { annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/pull/36710' } }, async ({ page, server }) => {
+  await page.setContent(`
+    <div style="height: 5000px;">Test</div>
+    <iframe loading="lazy" src="${server.PREFIX}/frame.html"></iframe>
+  `);
+
+  // Wait for the iframe to load
+  await page.waitForSelector('iframe');
+
+  // Get the snapshot of the page
+  const snapshot = await snapshotForAI(page, { timeout: 100 });
+
+  // The iframe should be present but empty
+  expect(snapshot).toContainYaml(`
+    - generic [active] [ref=e1]:
+      - generic [ref=e2]: Test
+      - iframe [ref=e3]
+  `);
+});