cherry-pick(#36203): fix: capture snapshot for ai during navigation (#36203)

Author: pavelfeldman

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

@@ -812,7 +812,7 @@ export class Page extends SdkObject {
 
   async snapshotForAI(metadata: CallMetadata): Promise<string> {
     this.lastSnapshotFrameIds = [];
-    const snapshot = await snapshotFrameForAI(this.mainFrame(), 0, this.lastSnapshotFrameIds);
+    const snapshot = await snapshotFrameForAI(metadata, this.mainFrame(), 0, this.lastSnapshotFrameIds);
     return snapshot.join('\n');
   }
 }
@@ -989,12 +989,24 @@ class FrameThrottler {
   }
 }
 
-async function snapshotFrameForAI(frame: frames.Frame, frameOrdinal: number, frameIds: string[]): Promise<string[]> {
-  const context = await frame._utilityContext();
-  const injectedScript = await context.injectedScript();
-  const snapshot = await injectedScript.evaluate((injected, refPrefix) => {
-    return injected.ariaSnapshot(injected.document.body, { forAI: true, refPrefix });
-  }, frameOrdinal ? 'f' + frameOrdinal : '');
+async function snapshotFrameForAI(metadata: CallMetadata, frame: frames.Frame, frameOrdinal: number, frameIds: string[]): Promise<string[]> {
+  // Only await the topmost navigations, inner frames will be empty when racing.
+  const controller = new ProgressController(metadata, frame);
+  const snapshot = await controller.run(progress => {
+    return frame.retryWithProgressAndTimeouts(progress, [1000, 2000, 4000, 8000], async continuePolling => {
+      try {
+        const context = await frame._utilityContext();
+        const injectedScript = await context.injectedScript();
+        return await injectedScript.evaluate((injected, refPrefix) => {
+          return injected.ariaSnapshot(injected.document.body, { forAI: true, refPrefix });
+        }, frameOrdinal ? 'f' + frameOrdinal : '');
+      } catch (e) {
+        if (js.isJavaScriptErrorInEvaluate(e))
+          throw e;
+        return continuePolling;
+      }
+    });
+  });
 
   const lines = snapshot.split('\n');
   const result = [];
@@ -1017,7 +1029,7 @@ async function snapshotFrameForAI(frame: frames.Frame, frameOrdinal: number, fra
     const frameOrdinal = frameIds.length + 1;
     frameIds.push(child.frame._id);
     try {
-      const childSnapshot = await snapshotFrameForAI(child.frame, frameOrdinal, frameIds);
+      const childSnapshot = await snapshotFrameForAI(metadata, child.frame, frameOrdinal, frameIds);
       result.push(line + ':', ...childSnapshot.map(l => leadingSpace + '  ' + l));
     } catch {
       result.push(line);

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

@@ -229,3 +229,14 @@ it('should gracefully fallback when child frame cant be captured', async ({ page
       - iframe [ref=e3]
   `);
 });
+
+it('should auto-wait for navigation', async ({ page, server }) => {
+  await page.goto(server.PREFIX + '/frames/frame.html');
+  const [, snapshot] = await Promise.all([
+    page.evaluate(() => window.location.reload()),
+    snapshotForAI(page)
+  ]);
+  expect(snapshot).toContainYaml(`
+    - generic [ref=e2]: Hi, I'm frame
+  `);
+});