fix: re-use node ids across snapshots

src/McpContext.ts

@@ -40,6 +40,7 @@ import {WaitForHelper} from './WaitForHelper.js';
 export interface TextSnapshotNode extends SerializedAXNode {
   id: string;
   backendNodeId?: number;
+  loaderId?: string;
   children: TextSnapshotNode[];
 }
 
@@ -129,6 +130,8 @@ export class McpContext implements Context {
   #locatorClass: typeof Locator;
   #options: McpContextOptions;
 
+  #uniqueBackendNodeIdToMcpId = new Map<string, string>();
+
   private constructor(
     browser: Browser,
     logger: Debugger,
@@ -440,14 +443,6 @@ export class McpContext implements Context {
         `No snapshot found. Use ${takeSnapshot.name} to capture one.`,
       );
     }
-    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');
@@ -589,10 +584,24 @@ export class McpContext implements Context {
     // will be used for the tree serialization and mapping ids back to nodes.
     let idCounter = 0;
     const idToNode = new Map<string, TextSnapshotNode>();
+    const seenUniqueIds = new Set<string>();
     const assignIds = (node: SerializedAXNode): TextSnapshotNode => {
+      let id = '';
+      // @ts-expect-error untyped loaderId & backendNodeId.
+      const uniqueBackendId = `${node.loaderId}_${node.backendNodeId}`;
+      if (this.#uniqueBackendNodeIdToMcpId.has(uniqueBackendId)) {
+        // Re-use MCP exposed ID if the uniqueId is the same.
+        id = this.#uniqueBackendNodeIdToMcpId.get(uniqueBackendId)!;
+      } else {
+        // Only generate a new ID if we have not seen the node before.
+        id = `${snapshotId}_${idCounter++}`;
+        this.#uniqueBackendNodeIdToMcpId.set(uniqueBackendId, id);
+      }
+      seenUniqueIds.add(uniqueBackendId);
+
       const nodeWithId: TextSnapshotNode = {
         ...node,
-        id: `${snapshotId}_${idCounter++}`,
+        id,
         children: node.children
           ? node.children.map(child => assignIds(child))
           : [],
@@ -626,6 +635,13 @@ export class McpContext implements Context {
         data?.cdpBackendNodeId,
       );
     }
+
+    // Clean up unique IDs that we did not see anymore.
+    for (const key of this.#uniqueBackendNodeIdToMcpId.keys()) {
+      if (!seenUniqueIds.has(key)) {
+        this.#uniqueBackendNodeIdToMcpId.delete(key);
+      }
+    }
   }
 
   getTextSnapshot(): TextSnapshot | null {

tests/McpContext.test.ts

@@ -20,24 +20,16 @@ describe('McpContext', () => {
     await withMcpContext(async (_response, context) => {
       const page = context.getSelectedPage();
       await page.setContent(
-        html`<button>Click me</button
-          ><input
+        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',
-        );
-      }
+      await context.getElementByUid('1_1');
     });
   });
 

tests/McpResponse.test.ts

@@ -16,6 +16,7 @@ import {
   traceResultIsSuccess,
 } from '../src/trace-processing/parse.js';
 
+import {serverHooks} from './server.js';
 import {loadTraceAsBuffer} from './trace-processing/fixtures/load.js';
 import {
   getImageContent,
@@ -63,8 +64,8 @@ describe('McpResponse', () => {
     await withMcpContext(async (response, context) => {
       const page = context.getSelectedPage();
       await page.setContent(
-        html`<button>Click me</button
-          ><input
+        html`<button>Click me</button>
+          <input
             type="text"
             value="Input"
           />`,
@@ -145,6 +146,111 @@ describe('McpResponse', () => {
     }
   });
 
+  it('preserves mapping ids across multiple snapshots', async () => {
+    await withMcpContext(async (response, context) => {
+      const page = context.getSelectedPage();
+      await page.setContent(html`
+        <div>
+          <button id="btn1">Button 1</button>
+          <span id="span1">Span 1</span>
+        </div>
+      `);
+      response.includeSnapshot();
+      // First snapshot
+      const res1 = await response.handle('test', context);
+      const text1 = getTextContent(res1.content[0]);
+      const btn1IdMatch = text1.match(/uid=(\S+) .*Button 1/);
+      const span1IdMatch = text1.match(/uid=(\S+) .*Span 1/);
+
+      assert.ok(btn1IdMatch, 'Button 1 ID not found in first snapshot');
+      assert.ok(span1IdMatch, 'Span 1 ID not found in first snapshot');
+
+      const btn1Id = btn1IdMatch[1];
+      const span1Id = span1IdMatch[1];
+
+      // Modify page: add a new element before the others to potentially shift indices if not stable
+      await page.evaluate(() => {
+        const newBtn = document.createElement('button');
+        newBtn.textContent = 'Button 2';
+        document.body.prepend(newBtn);
+      });
+
+      // Second snapshot
+      const res2 = await response.handle('test', context);
+      const text2 = getTextContent(res2.content[0]);
+
+      const btn1IdMatch2 = text2.match(/uid=(\S+) .*Button 1/);
+      const span1IdMatch2 = text2.match(/uid=(\S+) .*Span 1/);
+      const btn2IdMatch = text2.match(/uid=(\S+) .*Button 2/);
+
+      assert.ok(btn1IdMatch2, 'Button 1 ID not found in second snapshot');
+      assert.ok(span1IdMatch2, 'Span 1 ID not found in second snapshot');
+      assert.ok(btn2IdMatch, 'Button 2 ID not found in second snapshot');
+
+      assert.strictEqual(
+        btn1IdMatch2[1],
+        btn1Id,
+        'Button 1 ID changed between snapshots',
+      );
+      assert.strictEqual(
+        span1IdMatch2[1],
+        span1Id,
+        'Span 1 ID changed between snapshots',
+      );
+      assert.notStrictEqual(
+        btn2IdMatch[1],
+        btn1Id,
+        'Button 2 ID collides with Button 1',
+      );
+      assert.notStrictEqual(
+        btn2IdMatch[1],
+        btn1Id,
+        'Button 2 ID collides with Button 1',
+      );
+    });
+  });
+
+  describe('navigation', () => {
+    const server = serverHooks();
+
+    it('resets ids after navigation', async () => {
+      await withMcpContext(async (response, context) => {
+        server.addHtmlRoute(
+          '/page.html',
+          html`
+            <div>
+              <button id="btn1">Button 1</button>
+            </div>
+          `,
+        );
+        const page = context.getSelectedPage();
+        await page.goto(server.getRoute('/page.html'));
+
+        response.includeSnapshot();
+        const res1 = await response.handle('test', context);
+        const text1 = getTextContent(res1.content[0]);
+        const btn1IdMatch = text1.match(/uid=(\S+) .*Button 1/);
+        assert.ok(btn1IdMatch, 'Button 1 ID not found in first snapshot');
+        const btn1Id = btn1IdMatch[1];
+
+        // Navigate to the same page again (or meaningful navigation)
+        await page.goto(server.getRoute('/page.html'));
+
+        const res2 = await response.handle('test', context);
+        const text2 = getTextContent(res2.content[0]);
+        const btn1IdMatch2 = text2.match(/uid=(\S+) .*Button 1/);
+        assert.ok(btn1IdMatch2, 'Button 1 ID not found in second snapshot');
+        const btn1Id2 = btn1IdMatch2[1];
+
+        assert.notStrictEqual(
+          btn1Id2,
+          btn1Id,
+          'ID should reset after navigation',
+        );
+      });
+    });
+  });
+
   it('adds throttling setting when it is not null', async t => {
     await withMcpContext(async (response, context) => {
       context.setNetworkConditions('Slow 3G');