diff --git a/src/McpContext.ts b/src/McpContext.ts index 8187e0eb..557f012f 100644 --- a/src/McpContext.ts +++ b/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(); + 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(); + const seenUniqueIds = new Set(); 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 { diff --git a/tests/McpContext.test.ts b/tests/McpContext.test.ts index 8d1be513..e3683152 100644 --- a/tests/McpContext.test.ts +++ b/tests/McpContext.test.ts @@ -20,8 +20,8 @@ describe('McpContext', () => { await withMcpContext(async (_response, context) => { const page = context.getSelectedPage(); await page.setContent( - html`Click me + `, @@ -29,15 +29,7 @@ describe('McpContext', () => { 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'); }); }); diff --git a/tests/McpResponse.test.ts b/tests/McpResponse.test.ts index 80ad376f..574b1309 100644 --- a/tests/McpResponse.test.ts +++ b/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`Click me + `, @@ -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` +
+ + Span 1 +
+ `); + 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` +
+ +
+ `, + ); + 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');