fix: re-use node ids across snapshots (#814)

previously, Chrome DevTools MCP was strictly requiring the client to use
the latest snapshot. This change updates the behavior to be more
lenient. Previously issued uids will be recognized if the underlying
backend Id (loaderId + backendNodeId) still exists.

Related #726

Author: OrKoN

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');