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

2025-07-25 16:02:26 +08:00 · 2025-04-01 23:47:53 +02:00 · 2025-04-01 23:47:53 +02:00 · 0a5518b252
commit 0a5518b252
parent 4f16786432
4 changed files with 107 additions and 73 deletions
--- a/package-lock.json
+++ b/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",
-      "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.52.0-alpha-1743011787000.tgz",
-      "integrity": "sha512-ikJR8JXof5IBvErrmIsR3ixov4nKlQe/6PSYK/R6eTEe6eoT+eEXlaNY4z6mn9dF02Z1zYGxzAbb8TvSvuwh4Q==",
+      "version": "1.52.0-alpha-1743163434000",
+      "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.52.0-alpha-1743163434000.tgz",
+      "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",
-      "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.52.0-alpha-1743011787000.tgz",
-      "integrity": "sha512-wg9Tu4ZDKJWo7hBKpeuD/XLtLOQ7fCCuBfekgUrPLStA12O3224E1fbp/xGFnmi47SF71Y8F6C2Beyd3gYFWlQ==",
+      "version": "1.52.0-alpha-1743163434000",
+      "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.52.0-alpha-1743163434000.tgz",
+      "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",
-      "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.52.0-alpha-1743011787000.tgz",
-      "integrity": "sha512-yOpMfKxTBRqdm50b52cojvTCNttWN+Xk6LXF+KU4ufcGwcRjUud1xdHmHHvQNFFanXM1MBYnDKsMkRvjPsuYOw==",
+      "version": "1.52.0-alpha-1743163434000",
+      "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.52.0-alpha-1743163434000.tgz",
+      "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",
--- a/package.json
+++ b/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",
--- a/src/context.ts
+++ b/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();
-    const visibleFrames = await page.locator('iframe').filter({ visible: true }).all();
-    this._lastSnapshotFrames = visibleFrames.map(frame => frame.contentFrame());
+    this._lastSnapshotFrames = [];
+    const yaml = await this._allFramesSnapshot(this.existingPage());
+    return yaml.toString().trim();
+  }

-    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}`);
-      })
-    ]);
+  private async _allFramesSnapshot(frame: playwright.Page | playwright.FrameLocator): Promise<yaml.Document> {
+    const frameIndex = this._lastSnapshotFrames.push(frame) - 1;
+    const snapshotString = await frame.locator('body').ariaSnapshot({ ref: true });
+    const snapshot = yaml.parseDocument(snapshotString);

-    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: playwright.Frame | playwright.FrameLocator = page.mainFrame();
+    let frame = this._lastSnapshotFrames[0];
    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}`);
  }
 }
--- a/tests/basic.spec.ts
+++ b/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=s2e4]
+- button "Submit" [ref=s2e3]
 \`\`\`
 `);
 });
@ -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=s2e4]:
-    - option "Foo" [ref=s2e5]
-    - option "Bar" [selected] [ref=s2e6]
+- combobox [ref=s2e3]:
+    - option "Foo" [ref=s2e4]
+    - option "Bar" [selected] [ref=s2e5]
 \`\`\`
 `);
 });
@ -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=s2e4]:
-    - option "Foo" [ref=s2e5]
-    - option "Bar" [selected] [ref=s2e6]
-    - option "Baz" [selected] [ref=s2e7]
+- listbox [ref=s2e3]:
+    - option "Foo" [ref=s2e4]
+    - option "Bar" [selected] [ref=s2e5]
+    - option "Baz" [selected] [ref=s2e6]
 \`\`\`
 `);
 });
@ -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(`
- 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
+  })).toContainTextContent(`
 \`\`\`yaml
- document [ref=s1e2]:
-  - heading "Hello" [level=1] [ref=s1e4]
-
-# iframe src=data:text/html,<h1>World</h1>
- document [ref=f0s1e2]:
-  - heading "World" [level=1] [ref=f0s1e4]
+- heading "Hello" [level=1] [ref=s1e3]
+- iframe [ref=s1e4]:
+    - button "World" [ref=f1s1e3]
+    - main [ref=f1s1e4]:
+        - iframe [ref=f1s1e5]:
+            - 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!
 \`\`\`
 `
  );