chore: only include visible iframes, keep frame locators in own array (#60)

As discussed:

- hides invisible frames from snapshot
- keep our own frame locator array, so we don't rely on `page.frames()`
ordering to be stable

src/context.ts

@@ -23,6 +23,7 @@ export class Context {
   private _page: playwright.Page | undefined;
   private _console: playwright.ConsoleMessage[] = [];
   private _createPagePromise: Promise<playwright.Page> | undefined;
+  private _lastSnapshotFrames: playwright.FrameLocator[] = [];
 
   constructor(userDataDir: string, launchOptions?: playwright.LaunchOptions) {
     this._userDataDir = userDataDir;
@@ -90,4 +91,42 @@ export class Context {
     const [page] = context.pages();
     return { page };
   }
+
+  async allFramesSnapshot() {
+    const page = this.existingPage();
+    const visibleFrames = await page.locator('iframe').filter({ visible: true }).all();
+    this._lastSnapshotFrames = visibleFrames.map(frame => frame.contentFrame());
+
+    const snapshots = await Promise.all([
+      page.locator('html').ariaSnapshot({ ref: true }),
+      ...this._lastSnapshotFrames.map(async (frame, index) => {
+        const snapshot = await frame.locator('html').ariaSnapshot({ ref: true });
+        const args = [];
+        const src = await frame.owner().getAttribute('src');
+        if (src)
+          args.push(`src=${src}`);
+        const name = await frame.owner().getAttribute('name');
+        if (name)
+          args.push(`name=${name}`);
+        return `\n# iframe ${args.join(' ')}\n` + snapshot.replaceAll('[ref=', `[ref=f${index}`);
+      })
+    ]);
+
+    return snapshots.join('\n');
+  }
+
+  refLocator(ref: string): playwright.Locator {
+    const page = this.existingPage();
+    let frame: playwright.Frame | playwright.FrameLocator = page.mainFrame();
+    const match = ref.match(/^f(\d+)(.*)/);
+    if (match) {
+      const frameIndex = parseInt(match[1], 10);
+      if (!this._lastSnapshotFrames[frameIndex])
+        throw new Error(`Frame does not exist. Provide ref from the most current snapshot.`);
+      frame = this._lastSnapshotFrames[frameIndex];
+      ref = match[2];
+    }
+
+    return frame.locator(`aria-ref=${ref}`);
+  }
 }

src/tools/common.ts

@@ -41,7 +41,7 @@ export const navigate: ToolFactory = snapshot => ({
     // Cap load event to 5 seconds, the page is operational at this point.
     await page.waitForLoadState('load', { timeout: 5000 }).catch(() => {});
     if (snapshot)
-      return captureAriaSnapshot(page);
+      return captureAriaSnapshot(context);
     return {
       content: [{
         type: 'text',

src/tools/snapshot.ts

@@ -30,7 +30,7 @@ export const snapshot: Tool = {
   },
 
   handle: async context => {
-    return await captureAriaSnapshot(context.existingPage());
+    return await captureAriaSnapshot(context);
   },
 };
 
@@ -48,7 +48,7 @@ export const click: Tool = {
 
   handle: async (context, params) => {
     const validatedParams = elementSchema.parse(params);
-    return runAndWait(context, `"${validatedParams.element}" clicked`, page => refLocator(page, validatedParams.ref).click(), true);
+    return runAndWait(context, `"${validatedParams.element}" clicked`, () => context.refLocator(validatedParams.ref).click(), true);
   },
 };
 
@@ -68,9 +68,9 @@ export const drag: Tool = {
 
   handle: async (context, params) => {
     const validatedParams = dragSchema.parse(params);
-    return runAndWait(context, `Dragged "${validatedParams.startElement}" to "${validatedParams.endElement}"`, async page => {
-      const startLocator = refLocator(page, validatedParams.startRef);
-      const endLocator = refLocator(page, validatedParams.endRef);
+    return runAndWait(context, `Dragged "${validatedParams.startElement}" to "${validatedParams.endElement}"`, async () => {
+      const startLocator = context.refLocator(validatedParams.startRef);
+      const endLocator = context.refLocator(validatedParams.endRef);
       await startLocator.dragTo(endLocator);
     }, true);
   },
@@ -85,7 +85,7 @@ export const hover: Tool = {
 
   handle: async (context, params) => {
     const validatedParams = elementSchema.parse(params);
-    return runAndWait(context, `Hovered over "${validatedParams.element}"`, page => refLocator(page, validatedParams.ref).hover(), true);
+    return runAndWait(context, `Hovered over "${validatedParams.element}"`, () => context.refLocator(validatedParams.ref).hover(), true);
   },
 };
 
@@ -103,8 +103,8 @@ export const type: Tool = {
 
   handle: async (context, params) => {
     const validatedParams = typeSchema.parse(params);
-    return await runAndWait(context, `Typed "${validatedParams.text}" into "${validatedParams.element}"`, async page => {
-      const locator = refLocator(page, validatedParams.ref);
+    return await runAndWait(context, `Typed "${validatedParams.text}" into "${validatedParams.element}"`, async () => {
+      const locator = context.refLocator(validatedParams.ref);
       await locator.fill(validatedParams.text);
       if (validatedParams.submit)
         await locator.press('Enter');
@@ -125,8 +125,8 @@ export const selectOption: Tool = {
 
   handle: async (context, params) => {
     const validatedParams = selectOptionSchema.parse(params);
-    return await runAndWait(context, `Selected option in "${validatedParams.element}"`, async page => {
-      const locator = refLocator(page, validatedParams.ref);
+    return await runAndWait(context, `Selected option in "${validatedParams.element}"`, async () => {
+      const locator = context.refLocator(validatedParams.ref);
       await locator.selectOption(validatedParams.values);
     }, true);
   },
@@ -153,15 +153,3 @@ export const screenshot: Tool = {
     };
   },
 };
-
-function refLocator(page: playwright.Page, ref: string): playwright.Locator {
-  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

@@ -74,29 +74,20 @@ async function waitForCompletion<R>(page: playwright.Page, callback: () => Promi
 export async function runAndWait(context: Context, status: string, callback: (page: playwright.Page) => Promise<any>, snapshot: boolean = false): Promise<ToolResult> {
   const page = context.existingPage();
   await waitForCompletion(page, () => callback(page));
-  return snapshot ? captureAriaSnapshot(page, status) : {
+  return snapshot ? captureAriaSnapshot(context, status) : {
     content: [{ type: 'text', text: status }],
   };
 }
 
-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> {
+export async function captureAriaSnapshot(context: Context, status: string = ''): Promise<ToolResult> {
+  const page = context.existingPage();
   return {
     content: [{ type: 'text', text: `${status ? `${status}\n` : ''}
 - Page URL: ${page.url()}
 - Page Title: ${await page.title()}
 - Page Snapshot
 \`\`\`yaml
-${await captureAllFrameSnapshot(page)}
+${await context.allFramesSnapshot()}
 \`\`\`
 `
     }],

tests/basic.spec.ts

@@ -427,7 +427,7 @@ test('stitched aria frames', async ({ server }) => {
     params: {
       name: 'browser_navigate',
       arguments: {
-        url: 'data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe>',
+        url: 'data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe><iframe src="data:text/html,<h1>Should be invisible</h1>" style="display: none;"></iframe>',
       },
     },
   });
@@ -438,14 +438,16 @@ test('stitched aria frames', async ({ server }) => {
       content: [{
         type: 'text',
         text: `
-- Page URL: data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe>
+- Page URL: data:text/html,<h1>Hello</h1><iframe src="data:text/html,<h1>World</h1>"></iframe><iframe src="data:text/html,<h1>Should be invisible</h1>" style="display: none;"></iframe>
 - Page Title: 
 - Page Snapshot
 \`\`\`yaml
 - document [ref=s1e2]:
   - heading \"Hello\" [level=1] [ref=s1e4]
-- document [ref=f1s1e2]:
-  - heading \"World\" [level=1] [ref=f1s1e4]
+
+# iframe src=data:text/html,<h1>World</h1>
+- document [ref=f0s1e2]:
+  - heading \"World\" [level=1] [ref=f0s1e4]
 \`\`\`
 `,
       }],