fix: stitch all frames into one snapshot (#49)

src/tools/snapshot.ts

@@ -155,5 +155,13 @@ export const screenshot: Tool = {
 };
 
 function refLocator(page: playwright.Page, ref: string): playwright.Locator {
-  return page.locator(`aria-ref=${ref}`);
+  let frame = page.frames()[0];
+  const match = ref.match(/^f(\d+)(.*)/);
+  if (match) {
+    const frameIndex = parseInt(match[1], 10);
+    frame = page.frames()[frameIndex];
+    ref = match[2];
+  }
+
+  return frame.locator(`aria-ref=${ref}`);
 }

src/tools/utils.ts

@@ -79,15 +79,24 @@ export async function runAndWait(context: Context, status: string, callback: (pa
   };
 }
 
+export async function captureAllFrameSnapshot(page: playwright.Page): Promise<string> {
+  const snapshots = await Promise.all(page.frames().map(frame => frame.locator('html').ariaSnapshot({ ref: true })));
+  const scopedSnapshots = snapshots.map((snapshot, frameIndex) => {
+    if (frameIndex === 0)
+      return snapshot;
+    return snapshot.replaceAll('[ref=', `[ref=f${frameIndex}`);
+  });
+  return scopedSnapshots.join('\n');
+}
+
 export async function captureAriaSnapshot(page: playwright.Page, status: string = ''): Promise<ToolResult> {
-  const snapshot = await page.locator('html').ariaSnapshot({ ref: true });
   return {
     content: [{ type: 'text', text: `${status ? `${status}\n` : ''}
 - Page URL: ${page.url()}
 - Page Title: ${await page.title()}
 - Page Snapshot
 \`\`\`yaml
-${snapshot}
+${await captureAllFrameSnapshot(page)}
 \`\`\`
 `
     }],

tests/basic.spec.ts

@@ -416,3 +416,37 @@ test('browser://console', async ({ server }) => {
     }),
   }));
 });
+
+test('stitched aria frames', async ({ server }) => {
+  const response = await server.send({
+    jsonrpc: '2.0',
+    id: 2,
+    method: 'tools/call',
+    params: {
+      name: 'browser_navigate',
+      arguments: {
+        url: 'data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe>',
+      },
+    },
+  });
+
+  expect(response).toEqual(expect.objectContaining({
+    id: 2,
+    result: {
+      content: [{
+        type: 'text',
+        text: `
+- Page URL: data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe>
+- Page Title: 
+- Page Snapshot
+\`\`\`yaml
+- document [ref=s1e2]:
+  - heading \"Hello\" [level=1] [ref=s1e4]
+- document [ref=f1s1e2]:
+  - heading \"World\" [level=1] [ref=f1s1e4]
+\`\`\`
+`,
+      }],
+    },
+  }));
+});