chore: stitch together iframes into one tree (#71)

package-lock.json

@@ -11,7 +11,8 @@
       "dependencies": {
         "@modelcontextprotocol/sdk": "^1.6.1",
         "commander": "^13.1.0",
-        "playwright": "1.52.0-alpha-1743011787000",
+        "playwright": "^1.52.0-alpha-1743163434000",
+        "yaml": "^2.7.1",
         "zod-to-json-schema": "^3.24.4"
       },
       "bin": {
@@ -20,7 +21,7 @@
       "devDependencies": {
         "@eslint/eslintrc": "^3.2.0",
         "@eslint/js": "^9.19.0",
-        "@playwright/test": "1.52.0-alpha-1743011787000",
+        "@playwright/test": "^1.52.0-alpha-1743163434000",
         "@stylistic/eslint-plugin": "^3.0.1",
         "@types/node": "^22.13.10",
         "@typescript-eslint/eslint-plugin": "^8.26.1",
@@ -285,13 +286,13 @@
       }
     },
     "node_modules/@playwright/test": {
-      "version": "1.52.0-alpha-1743011787000",
+      "version": "1.52.0-alpha-1743163434000",
-      "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.52.0-alpha-1743011787000.tgz",
+      "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.52.0-alpha-1743163434000.tgz",
-      "integrity": "sha512-ikJR8JXof5IBvErrmIsR3ixov4nKlQe/6PSYK/R6eTEe6eoT+eEXlaNY4z6mn9dF02Z1zYGxzAbb8TvSvuwh4Q==",
+      "integrity": "sha512-4uBgNlJ6hgPtB8DrwQsgoKuVoe7j+nPqudna7CLXWCmmT3LYPMD5aOjGoBkszr+R9NejtKashq/bOi/ny9hsIA==",
       "dev": true,
       "license": "Apache-2.0",
       "dependencies": {
-        "playwright": "1.52.0-alpha-1743011787000"
+        "playwright": "1.52.0-alpha-1743163434000"
       },
       "bin": {
         "playwright": "cli.js"
@@ -3296,12 +3297,12 @@
       }
     },
     "node_modules/playwright": {
-      "version": "1.52.0-alpha-1743011787000",
+      "version": "1.52.0-alpha-1743163434000",
-      "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.52.0-alpha-1743011787000.tgz",
+      "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.52.0-alpha-1743163434000.tgz",
-      "integrity": "sha512-wg9Tu4ZDKJWo7hBKpeuD/XLtLOQ7fCCuBfekgUrPLStA12O3224E1fbp/xGFnmi47SF71Y8F6C2Beyd3gYFWlQ==",
+      "integrity": "sha512-4uYv49ekPjolydfFfTfFQ2z4URF9UZMVUXLy7aXam/tPxEQ5O7+jQC+yzrDMGmhcj5QkMnxjlyk7N2V9a0QLdQ==",
       "license": "Apache-2.0",
       "dependencies": {
-        "playwright-core": "1.52.0-alpha-1743011787000"
+        "playwright-core": "1.52.0-alpha-1743163434000"
       },
       "bin": {
         "playwright": "cli.js"
@@ -3314,9 +3315,9 @@
       }
     },
     "node_modules/playwright-core": {
-      "version": "1.52.0-alpha-1743011787000",
+      "version": "1.52.0-alpha-1743163434000",
-      "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.52.0-alpha-1743011787000.tgz",
+      "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.52.0-alpha-1743163434000.tgz",
-      "integrity": "sha512-yOpMfKxTBRqdm50b52cojvTCNttWN+Xk6LXF+KU4ufcGwcRjUud1xdHmHHvQNFFanXM1MBYnDKsMkRvjPsuYOw==",
+      "integrity": "sha512-Tn4u3Ywwjkh847/bYWlXIrNxv5DRJRDgtb+VYMXHvNCKkrxL6yfZ1ApIAYD7IAkkKH/KLTXszGWl3a/Z/KDfQA==",
       "license": "Apache-2.0",
       "bin": {
         "playwright-core": "cli.js"
@@ -4348,6 +4349,18 @@
       "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==",
       "license": "ISC"
     },
+    "node_modules/yaml": {
+      "version": "2.7.1",
+      "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.7.1.tgz",
+      "integrity": "sha512-10ULxpnOCQXxJvBgxsn9ptjq6uviG/htZKk9veJGhlqn3w/DxQ631zFF+nlQXLwmImeS5amR2dl2U8sg6U9jsQ==",
+      "license": "ISC",
+      "bin": {
+        "yaml": "bin.mjs"
+      },
+      "engines": {
+        "node": ">= 14"
+      }
+    },
     "node_modules/yocto-queue": {
       "version": "0.1.0",
       "resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-0.1.0.tgz",

package.json

@@ -32,18 +32,19 @@
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.6.1",
     "commander": "^13.1.0",
-    "playwright": "1.52.0-alpha-1743011787000",
+    "playwright": "^1.52.0-alpha-1743163434000",
+    "yaml": "^2.7.1",
     "zod-to-json-schema": "^3.24.4"
   },
   "devDependencies": {
     "@eslint/eslintrc": "^3.2.0",
     "@eslint/js": "^9.19.0",
-    "@playwright/test": "1.52.0-alpha-1743011787000",
+    "@playwright/test": "^1.52.0-alpha-1743163434000",
     "@stylistic/eslint-plugin": "^3.0.1",
+    "@types/node": "^22.13.10",
     "@typescript-eslint/eslint-plugin": "^8.26.1",
     "@typescript-eslint/parser": "^8.26.1",
     "@typescript-eslint/utils": "^8.26.1",
-    "@types/node": "^22.13.10",
     "eslint": "^9.19.0",
     "eslint-plugin-import": "^2.31.0",
     "eslint-plugin-notice": "^1.0.0",

src/context.ts

@@ -18,6 +18,7 @@ import { fork } from 'child_process';
 import path from 'path';
 
 import * as playwright from 'playwright';
+import yaml from 'yaml';
 
 export type ContextOptions = {
   browserName?: 'chromium' | 'firefox' | 'webkit';
@@ -34,7 +35,7 @@ export class Context {
   private _console: playwright.ConsoleMessage[] = [];
   private _createPagePromise: Promise<playwright.Page> | undefined;
   private _fileChooser: playwright.FileChooser | undefined;
-  private _lastSnapshotFrames: playwright.FrameLocator[] = [];
+  private _lastSnapshotFrames: (playwright.Page | playwright.FrameLocator)[] = [];
 
   constructor(options: ContextOptions) {
     this._options = options;
@@ -161,40 +162,57 @@ export class Context {
   }
 
   async allFramesSnapshot() {
-    const page = this.existingPage();
+    this._lastSnapshotFrames = [];
-    const visibleFrames = await page.locator('iframe').filter({ visible: true }).all();
+    const yaml = await this._allFramesSnapshot(this.existingPage());
-    this._lastSnapshotFrames = visibleFrames.map(frame => frame.contentFrame());
+    return yaml.toString().trim();
+  }
 
-    const snapshots = await Promise.all([
+  private async _allFramesSnapshot(frame: playwright.Page | playwright.FrameLocator): Promise<yaml.Document> {
-      page.locator('html').ariaSnapshot({ ref: true }),
+    const frameIndex = this._lastSnapshotFrames.push(frame) - 1;
-      ...this._lastSnapshotFrames.map(async (frame, index) => {
+    const snapshotString = await frame.locator('body').ariaSnapshot({ ref: true });
-        const snapshot = await frame.locator('html').ariaSnapshot({ ref: true });
+    const snapshot = yaml.parseDocument(snapshotString);
-        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');
+    const visit = async (node: any): Promise<unknown> => {
+      if (yaml.isPair(node)) {
+        await Promise.all([
+          visit(node.key).then(k => node.key = k),
+          visit(node.value).then(v => node.value = v)
+        ]);
+      } else if (yaml.isSeq(node) || yaml.isMap(node)) {
+        node.items = await Promise.all(node.items.map(visit));
+      } else if (yaml.isScalar(node)) {
+        if (typeof node.value === 'string') {
+          const value = node.value;
+          if (frameIndex > 0)
+            node.value = value.replace('[ref=', `[ref=f${frameIndex}`);
+          if (value.startsWith('iframe ')) {
+            const ref = value.match(/\[ref=(.*)\]/)?.[1];
+            if (ref) {
+              const childSnapshot = await this._allFramesSnapshot(frame.frameLocator(`aria-ref=${ref}`));
+              return snapshot.createPair(node.value, childSnapshot);
+            }
+          }
+        }
+      }
+
+      return node;
+    };
+    await visit(snapshot.contents);
+    return snapshot;
   }
 
   refLocator(ref: string): playwright.Locator {
-    const page = this.existingPage();
+    let frame = this._lastSnapshotFrames[0];
-    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];
     }
 
+    if (!frame)
+      throw new Error(`Frame does not exist. Provide ref from the most current snapshot.`);
+
     return frame.locator(`aria-ref=${ref}`);
   }
 }

tests/basic.spec.ts

@@ -79,7 +79,7 @@ test('test browser_navigate', async ({ client }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s1e2]: Hello, world!
+- text: Hello, world!
 \`\`\`
 `
   );
@@ -97,7 +97,7 @@ test('test browser_click', async ({ client }) => {
     name: 'browser_click',
     arguments: {
       element: 'Submit button',
-      ref: 's1e4',
+      ref: 's1e3',
     },
   })).toHaveTextContent(`"Submit button" clicked
 
@@ -105,8 +105,7 @@ test('test browser_click', async ({ client }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s2e2]:
+- button "Submit" [ref=s2e3]
-  - button "Submit" [ref=s2e4]
 \`\`\`
 `);
 });
@@ -133,7 +132,7 @@ test('test reopen browser', async ({ client }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s1e2]: Hello, world!
+- text: Hello, world!
 \`\`\`
 `);
 });
@@ -150,7 +149,7 @@ test('single option', async ({ client }) => {
     name: 'browser_select_option',
     arguments: {
       element: 'Select',
-      ref: 's1e4',
+      ref: 's1e3',
       values: ['bar'],
     },
   })).toHaveTextContent(`Selected option in "Select"
@@ -159,10 +158,9 @@ test('single option', async ({ client }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s2e2]:
+- combobox [ref=s2e3]:
-  - combobox [ref=s2e4]:
+    - option "Foo" [ref=s2e4]
-    - option "Foo" [ref=s2e5]
+    - option "Bar" [selected] [ref=s2e5]
-    - option "Bar" [selected] [ref=s2e6]
 \`\`\`
 `);
 });
@@ -179,7 +177,7 @@ test('multiple option', async ({ client }) => {
     name: 'browser_select_option',
     arguments: {
       element: 'Select',
-      ref: 's1e4',
+      ref: 's1e3',
       values: ['bar', 'baz'],
     },
   })).toHaveTextContent(`Selected option in "Select"
@@ -188,11 +186,10 @@ test('multiple option', async ({ client }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s2e2]:
+- listbox [ref=s2e3]:
-  - listbox [ref=s2e4]:
+    - option "Foo" [ref=s2e4]
-    - option "Foo" [ref=s2e5]
+    - option "Bar" [selected] [ref=s2e5]
-    - option "Bar" [selected] [ref=s2e6]
+    - option "Baz" [selected] [ref=s2e6]
-    - option "Baz" [selected] [ref=s2e7]
 \`\`\`
 `);
 });
@@ -219,21 +216,26 @@ test('stitched aria frames', async ({ client }) => {
   expect(await client.callTool({
     name: 'browser_navigate',
     arguments: {
-      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>',
+      url: `data:text/html,<h1>Hello</h1><iframe src="data:text/html,<button>World</button><main><iframe src='data:text/html,<p>Nested</p>'></iframe></main>"></iframe><iframe src="data:text/html,<h1>Should be invisible</h1>" style="display: none;"></iframe>`,
     },
-  })).toHaveTextContent(`
+  })).toContainTextContent(`
-- 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=s1e3]
-  - heading "Hello" [level=1] [ref=s1e4]
+- iframe [ref=s1e4]:
-
+    - button "World" [ref=f1s1e3]
-# iframe src=data:text/html,<h1>World</h1>
+    - main [ref=f1s1e4]:
-- document [ref=f0s1e2]:
+        - iframe [ref=f1s1e5]:
-  - heading "World" [level=1] [ref=f0s1e4]
+            - paragraph [ref=f2s1e3]: Nested
 \`\`\`
 `);
+
+  expect(await client.callTool({
+    name: 'browser_click',
+    arguments: {
+      element: 'World',
+      ref: 'f1s1e3',
+    },
+  })).toContainTextContent('"World" clicked');
 });
 
 test('browser_choose_file', async ({ client }) => {
@@ -242,13 +244,13 @@ test('browser_choose_file', async ({ client }) => {
     arguments: {
       url: 'data:text/html,<html><title>Title</title><input type="file" /><button>Button</button></html>',
     },
-  })).toContainTextContent('- textbox [ref=s1e4]');
+  })).toContainTextContent('- textbox [ref=s1e3]');
 
   expect(await client.callTool({
     name: 'browser_click',
     arguments: {
       element: 'Textbox',
-      ref: 's1e4',
+      ref: 's1e3',
     },
   })).toContainTextContent('There is a file chooser visible that requires browser_choose_file to be called');
 
@@ -264,7 +266,7 @@ test('browser_choose_file', async ({ client }) => {
     });
 
     expect(response).not.toContainTextContent('There is a file chooser visible that requires browser_choose_file to be called');
-    expect(response).toContainTextContent('textbox [ref=s3e4]: C:\\fakepath\\test.txt');
+    expect(response).toContainTextContent('textbox [ref=s3e3]: C:\\fakepath\\test.txt');
   }
 
   {
@@ -272,12 +274,12 @@ test('browser_choose_file', async ({ client }) => {
       name: 'browser_click',
       arguments: {
         element: 'Textbox',
-        ref: 's3e4',
+        ref: 's3e3',
       },
     });
 
     expect(response).toContainTextContent('There is a file chooser visible that requires browser_choose_file to be called');
-    expect(response).toContainTextContent('button "Button" [ref=s4e5]');
+    expect(response).toContainTextContent('button "Button" [ref=s4e4]');
   }
 
   {
@@ -285,7 +287,7 @@ test('browser_choose_file', async ({ client }) => {
       name: 'browser_click',
       arguments: {
         element: 'Button',
-        ref: 's4e5',
+        ref: 's4e4',
       },
     });
 
@@ -328,7 +330,7 @@ test('cdp server', async ({ cdpEndpoint, startClient }) => {
 - Page Title: Title
 - Page Snapshot
 \`\`\`yaml
-- document [ref=s1e2]: Hello, world!
+- text: Hello, world!
 \`\`\`
 `
   );