feat: validate uids

Author: OrKoN

src/McpContext.ts

@@ -21,10 +21,16 @@ import path from 'node:path';
 import {listPages} from './tools/pages.js';
 
 export interface TextSnapshotNode extends SerializedAXNode {
-  id: number;
+  id: string;
   children: TextSnapshotNode[];
 }
 
+export interface TextSnapshot {
+  root: TextSnapshotNode;
+  idToNode: Map<string, TextSnapshotNode>;
+  snapshotId: string;
+}
+
 export class McpContext implements Context {
   browser: Browser;
   logger: Debugger;
@@ -33,8 +39,7 @@ export class McpContext implements Context {
   #pages: Page[] = [];
   #selectedPageIdx = 0;
   // The most recent snapshot.
-  #textSnapshot: TextSnapshotNode | null = null;
-  #idToNodeMap = new Map<number, TextSnapshotNode>();
+  #textSnapshot: TextSnapshot | null = null;
   #networkCollector: NetworkCollector;
   #consoleCollector: PageCollector<ConsoleMessage | Error>;
 
@@ -43,6 +48,8 @@ export class McpContext implements Context {
   #cpuThrottlingRate = 1;
   #dialog?: Dialog;
 
+  #nextSnapshotId = 1;
+
   private constructor(browser: Browser, logger: Debugger) {
     this.browser = browser;
     this.logger = logger;
@@ -192,11 +199,19 @@ export class McpContext implements Context {
     newPage.setDefaultNavigationTimeout(10_000);
   }
 
-  async getElementByUid(uid: number): Promise<ElementHandle<Element>> {
-    if (!this.#idToNodeMap.size) {
+  async getElementByUid(uid: string): Promise<ElementHandle<Element>> {
+    if (!this.#textSnapshot?.idToNode.size) {
       throw new Error('No snapshot found. Use browser_snapshot to capture one');
     }
-    const node = this.#idToNodeMap.get(uid);
+    const [snapshotId] = uid.split('_');
+
+    if (this.#textSnapshot.snapshotId !== snapshotId) {
+      throw new Error(
+        'This uid is coming from a stale snapshot. Call take_snapshot to get a fresh snapshot.',
+      );
+    }
+
+    const node = this.#textSnapshot?.idToNode.get(uid);
     if (!node) {
       throw new Error('No such element found in the snapshot');
     }
@@ -222,35 +237,38 @@ export class McpContext implements Context {
   /**
    * Creates a text snapshot of a page.
    */
-  async createTextSnapshot(): Promise<TextSnapshotNode | null> {
+  async createTextSnapshot(): Promise<void> {
     const page = this.getSelectedPage();
     const rootNode = await page.accessibility.snapshot();
     if (!rootNode) {
-      return null;
+      return;
     }
 
     // Iterate through the whole accessibility node tree and assign node ids that
     // will be used for the tree serialization and mapping ids back to nodes.
     let idCounter = 0;
-    this.#idToNodeMap.clear();
+    const idToNode = new Map<string, TextSnapshotNode>();
     const assignIds = (node: SerializedAXNode): TextSnapshotNode => {
       const nodeWithId: TextSnapshotNode = {
         ...node,
-        id: idCounter++,
+        id: `${this.#nextSnapshotId}_${idCounter++}`,
         children: node.children
           ? node.children.map(child => assignIds(child))
           : [],
       };
-      this.#idToNodeMap.set(nodeWithId.id, nodeWithId);
+      idToNode.set(nodeWithId.id, nodeWithId);
       return nodeWithId;
     };
 
     const rootNodeWithId = assignIds(rootNode);
-    this.#textSnapshot = rootNodeWithId;
-    return rootNodeWithId;
+    this.#textSnapshot = {
+      root: rootNodeWithId,
+      snapshotId: String(this.#nextSnapshotId++),
+      idToNode,
+    };
   }
 
-  getTextSnapshot(): TextSnapshotNode | null {
+  getTextSnapshot(): TextSnapshot | null {
     return this.#textSnapshot;
   }
 

src/McpResponse.ts

@@ -145,9 +145,9 @@ Call browser_handle_dialog to handle it before continuing.`);
     }
 
     if (this.#includeSnapshot) {
-      const rootNode = context.getTextSnapshot();
-      if (rootNode) {
-        const formattedSnapshot = formatA11ySnapshot(rootNode);
+      const snapshot = context.getTextSnapshot();
+      if (snapshot) {
+        const formattedSnapshot = formatA11ySnapshot(snapshot.root);
         response.push('## Page content');
         response.push(formattedSnapshot);
       }

src/tools/ToolDefinition.ts

@@ -60,7 +60,7 @@ export type Context = Readonly<{
   getPageByIdx(idx: number): Page;
   newPage(): Promise<Page>;
   setSelectedPageIdx(idx: number): void;
-  getElementByUid(uid: number): Promise<ElementHandle<Element>>;
+  getElementByUid(uid: string): Promise<ElementHandle<Element>>;
   setNetworkConditions(conditions: string | null): void;
   setCpuThrottlingRate(rate: number): void;
   saveTemporaryFile(

src/tools/input.ts

@@ -19,7 +19,7 @@ export const click = defineTool({
   },
   schema: {
     uid: z
-      .number()
+      .string()
       .describe(
         'The uid of an element on the page from the page content snapshot',
       ),
@@ -58,7 +58,7 @@ export const hover = defineTool({
   },
   schema: {
     uid: z
-      .number()
+      .string()
       .describe(
         'The uid of an element on the page from the page content snapshot',
       ),
@@ -87,7 +87,7 @@ export const fill = defineTool({
   },
   schema: {
     uid: z
-      .number()
+      .string()
       .describe(
         'The uid of an element on the page from the page content snapshot',
       ),
@@ -115,8 +115,8 @@ export const drag = defineTool({
     readOnlyHint: false,
   },
   schema: {
-    from_uid: z.number().describe('The uid of the element to drag'),
-    to_uid: z.number().describe('The uid of the element to drop into'),
+    from_uid: z.string().describe('The uid of the element to drag'),
+    to_uid: z.string().describe('The uid of the element to drop into'),
   },
   handler: async (request, response, context) => {
     const fromHandle = await context.getElementByUid(request.params.from_uid);
@@ -147,7 +147,7 @@ export const fillForm = defineTool({
     elements: z
       .array(
         z.object({
-          uid: z.number().describe('The uid of the element to fill out'),
+          uid: z.string().describe('The uid of the element to fill out'),
           value: z.string().describe('Value for the element'),
         }),
       )
@@ -178,7 +178,7 @@ export const uploadFile = defineTool({
   },
   schema: {
     uid: z
-      .number()
+      .string()
       .describe(
         'The uid of the file input element or an element that will open file chooser on the page from the page content snapshot',
       ),

src/tools/screenshot.ts

@@ -22,7 +22,7 @@ export const screenshot = defineTool({
       .default('png')
       .describe('Type of format to save the screenshot as. Default is "png"'),
     uid: z
-      .number()
+      .string()
       .optional()
       .describe(
         'The uid of an element on the page from the page content snapshot. If omitted takes a pages screenshot.',

tests/McpContext.test.ts

@@ -0,0 +1,31 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import {describe, it} from 'node:test';
+import assert from 'assert';
+
+import {withBrowser} from './utils.js';
+
+describe('McpResponse', () => {
+  it('list pages', async () => {
+    await withBrowser(async (response, context) => {
+      const page = context.getSelectedPage();
+      await page.setContent(`<!DOCTYPE html>
+<button>Click me</button><input type="text" value="Input">`);
+      await context.createTextSnapshot();
+      assert.ok(await context.getElementByUid('1_1'));
+      await context.createTextSnapshot();
+      try {
+        await context.getElementByUid('1_1');
+        assert.fail('not reached');
+      } catch (err) {
+        assert.strict(
+          err.message,
+          'This uid is coming from a stale snapshot. Call take_snapshot to get a fresh snapshot',
+        );
+      }
+    });
+  });
+});

tests/McpResponse.test.ts

@@ -61,9 +61,9 @@ Testing 2`,
         result[0].text,
         `# test response
 ## Page content
-uid=0 RootWebArea ""
-  uid=1 button "Click me" focusable focused
-  uid=2 textbox "" value="Input"
+uid=1_0 RootWebArea ""
+  uid=1_1 button "Click me" focusable focused
+  uid=1_2 textbox "" value="Input"
 `,
       );
     });
@@ -87,9 +87,9 @@ uid=0 RootWebArea ""
         result[0].text,
         `# test response
 ## Page content
-uid=0 RootWebArea "My test page"
-  uid=1 StaticText "username"
-  uid=2 textbox "username" value="mcp" focusable focused
+uid=1_0 RootWebArea "My test page"
+  uid=1_1 StaticText "username"
+  uid=1_2 textbox "username" value="mcp" focusable focused
 `,
       );
     });

tests/formatters/snapshotFormatter.test.ts

@@ -13,13 +13,13 @@ import type {ElementHandle} from 'puppeteer-core';
 describe('snapshotFormatter', () => {
   it('formats a snapshot with value properties', () => {
     const snapshot: TextSnapshotNode = {
-      id: 1,
+      id: '1_1',
       role: 'textbox',
       name: 'textbox',
       value: 'value',
       children: [
         {
-          id: 2,
+          id: '1_2',
           role: 'statictext',
           name: 'text',
           children: [],
@@ -36,21 +36,21 @@ describe('snapshotFormatter', () => {
     const formatted = formatA11ySnapshot(snapshot);
     assert.strictEqual(
       formatted,
-      `uid=1 textbox "textbox" value="value"
-  uid=2 statictext "text"
+      `uid=1_1 textbox "textbox" value="value"
+  uid=1_2 statictext "text"
 `,
     );
   });
 
   it('formats a snapshot with boolean properties', () => {
     const snapshot: TextSnapshotNode = {
-      id: 1,
+      id: '1_1',
       role: 'button',
       name: 'button',
       disabled: true,
       children: [
         {
-          id: 2,
+          id: '1_2',
           role: 'statictext',
           name: 'text',
           children: [],
@@ -67,21 +67,21 @@ describe('snapshotFormatter', () => {
     const formatted = formatA11ySnapshot(snapshot);
     assert.strictEqual(
       formatted,
-      `uid=1 button "button" disableable disabled
-  uid=2 statictext "text"
+      `uid=1_1 button "button" disableable disabled
+  uid=1_2 statictext "text"
 `,
     );
   });
 
   it('formats a snapshot with checked properties', () => {
     const snapshot: TextSnapshotNode = {
-      id: 1,
+      id: '1_1',
       role: 'checkbox',
       name: 'checkbox',
       checked: true,
       children: [
         {
-          id: 2,
+          id: '1_2',
           role: 'statictext',
           name: 'text',
           children: [],
@@ -98,20 +98,20 @@ describe('snapshotFormatter', () => {
     const formatted = formatA11ySnapshot(snapshot);
     assert.strictEqual(
       formatted,
-      `uid=1 checkbox "checkbox" checked checked="true"
-  uid=2 statictext "text"
+      `uid=1_1 checkbox "checkbox" checked checked="true"
+  uid=1_2 statictext "text"
 `,
     );
   });
 
   it('formats a snapshot with multiple different type attributes', () => {
     const snapshot: TextSnapshotNode = {
-      id: 1,
+      id: '1_1',
       role: 'root',
       name: 'root',
       children: [
         {
-          id: 2,
+          id: '1_2',
           role: 'button',
           name: 'button',
           focused: true,
@@ -122,7 +122,7 @@ describe('snapshotFormatter', () => {
           },
         },
         {
-          id: 3,
+          id: '1_3',
           role: 'textbox',
           name: 'textbox',
           value: 'value',
@@ -140,9 +140,9 @@ describe('snapshotFormatter', () => {
     const formatted = formatA11ySnapshot(snapshot);
     assert.strictEqual(
       formatted,
-      `uid=1 root "root"
-  uid=2 button "button" disableable disabled focusable focused
-  uid=3 textbox "textbox" value="value"
+      `uid=1_1 root "root"
+  uid=1_2 button "button" disableable disabled focusable focused
+  uid=1_3 textbox "textbox" value="value"
 `,
     );
   });