Skip to content

Commit 76720a6

Browse files
OrKoNOrKoN
committed
chore: structured snapshot
1 parent 14ff400 commit 76720a6

9 files changed

Lines changed: 371 additions & 184 deletions

File tree

‎GEMINI.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# Instructions
2+
3+
- use `npm run build` to run tsc and test build
4+
- use `npm run test` to run tests, run all tests to verify correctness

‎src/McpResponse.ts‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import {
1717
getShortDescriptionForRequest,
1818
getStatusFromRequest,
1919
} from './formatters/networkFormatter.js';
20-
import {formatSnapshotNode} from './formatters/snapshotFormatter.js';
20+
import {SnapshotFormatter} from './formatters/SnapshotFormatter.js';
2121
import type {McpContext} from './McpContext.js';
2222
import {DevTools} from './third_party/index.js';
2323
import type {
@@ -181,29 +181,28 @@ export class McpResponse implements Response {
181181
async handle(
182182
toolName: string,
183183
context: McpContext,
184-
): Promise<Array<TextContent | ImageContent>> {
184+
): Promise<{content: Array<TextContent | ImageContent>; structuredContent: object}> {
185185
if (this.#includePages) {
186186
await context.createPagesSnapshot();
187187
}
188188

189-
let formattedSnapshot: string | undefined;
189+
let snapshot: SnapshotFormatter | string | undefined;
190190
if (this.#snapshotParams) {
191191
await context.createTextSnapshot(
192192
this.#snapshotParams.verbose,
193193
this.#devToolsData,
194194
);
195-
const snapshot = context.getTextSnapshot();
196-
if (snapshot) {
195+
const textSnapshot = context.getTextSnapshot();
196+
if (textSnapshot) {
197+
const formatter = new SnapshotFormatter(textSnapshot);
197198
if (this.#snapshotParams.filePath) {
198199
await context.saveFile(
199-
new TextEncoder().encode(
200-
formatSnapshotNode(snapshot.root, snapshot),
201-
),
200+
new TextEncoder().encode(formatter.toString()),
202201
this.#snapshotParams.filePath,
203202
);
204-
formattedSnapshot = `Saved snapshot to ${this.#snapshotParams.filePath}.`;
203+
snapshot = this.#snapshotParams.filePath;
205204
} else {
206-
formattedSnapshot = formatSnapshotNode(snapshot.root, snapshot);
205+
snapshot = formatter;
207206
}
208207
}
209208
}
@@ -335,7 +334,7 @@ export class McpResponse implements Response {
335334
bodies,
336335
consoleData,
337336
consoleListData,
338-
formattedSnapshot,
337+
snapshot,
339338
});
340339
}
341340

@@ -349,9 +348,9 @@ export class McpResponse implements Response {
349348
};
350349
consoleData: ConsoleMessageData | undefined;
351350
consoleListData: ConsoleMessageData[] | undefined;
352-
formattedSnapshot: string | undefined;
351+
snapshot: SnapshotFormatter | string | undefined;
353352
},
354-
): Array<TextContent | ImageContent> {
353+
): {content: Array<TextContent | ImageContent>; structuredContent: object} {
355354
const response = [`# ${toolName} response`];
356355
for (const line of this.#textResponseLines) {
357356
response.push(line);
@@ -393,9 +392,20 @@ Call ${handleDialog.name} to handle it before continuing.`);
393392
response.push(...parts);
394393
}
395394

396-
if (data.formattedSnapshot) {
397-
response.push('## Latest page snapshot');
398-
response.push(data.formattedSnapshot);
395+
const structuredContent: {
396+
snapshot?: object;
397+
snapshotFilePath?: string;
398+
} = {};
399+
400+
if (data.snapshot) {
401+
if (typeof data.snapshot === 'string') {
402+
response.push(`Saved snapshot to ${data.snapshot}.`);
403+
structuredContent.snapshotFilePath = data.snapshot;
404+
} else {
405+
response.push('## Latest page snapshot');
406+
response.push(data.snapshot.toString());
407+
structuredContent.snapshot = data.snapshot.toJSON();
408+
}
399409
}
400410

401411
response.push(...this.#formatNetworkRequestData(context, data.bodies));
@@ -468,7 +478,10 @@ Call ${handleDialog.name} to handle it before continuing.`);
468478
} as const;
469479
});
470480

471-
return [text, ...images];
481+
return {
482+
content: [text, ...images],
483+
structuredContent,
484+
};
472485
}
473486

474487
#dataWithPagination<T>(data: T[], pagination?: PaginationOptions) {

‎src/formatters/snapshotFormatter.ts‎

Lines changed: 135 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -3,92 +3,162 @@
33
* Copyright 2025 Google LLC
44
* SPDX-License-Identifier: Apache-2.0
55
*/
6-
76
import type {TextSnapshot, TextSnapshotNode} from '../McpContext.js';
87

9-
export function formatSnapshotNode(
10-
root: TextSnapshotNode,
11-
snapshot?: TextSnapshot,
12-
depth = 0,
13-
): string {
14-
const chunks: string[] = [];
8+
export class SnapshotFormatter {
9+
constructor(public snapshot?: TextSnapshot) {}
10+
11+
toString(): string {
12+
if (!this.snapshot) {
13+
return '';
14+
}
15+
const chunks: string[] = [];
16+
const root = this.snapshot.root;
1517

16-
if (depth === 0) {
1718
// Top-level content of the snapshot.
1819
if (
19-
snapshot?.verbose &&
20-
snapshot?.hasSelectedElement &&
21-
!snapshot.selectedElementUid
20+
this.snapshot.verbose &&
21+
this.snapshot.hasSelectedElement &&
22+
!this.snapshot.selectedElementUid
2223
) {
2324
chunks.push(`Note: there is a selected element in the DevTools Elements panel but it is not included into the current a11y tree snapshot.
2425
Get a verbose snapshot to include all elements if you are interested in the selected element.\n\n`);
2526
}
27+
28+
chunks.push(this.#formatNode(root, 0));
29+
return chunks.join('');
2630
}
2731

28-
const attributes = getAttributes(root);
29-
const line =
30-
' '.repeat(depth * 2) +
31-
attributes.join(' ') +
32-
(root.id === snapshot?.selectedElementUid
33-
? ' [selected in the DevTools Elements panel]'
34-
: '') +
35-
'\n';
36-
chunks.push(line);
37-
38-
for (const child of root.children) {
39-
chunks.push(formatSnapshotNode(child, snapshot, depth + 1));
32+
toJSON(): object {
33+
if (!this.snapshot) {
34+
return {};
35+
}
36+
return this.#nodeToJSON(this.snapshot.root);
4037
}
4138

42-
return chunks.join('');
43-
}
39+
#formatNode(node: TextSnapshotNode, depth: number): string {
40+
const chunks: string[] = [];
41+
const attributes = this.#getAttributes(node);
42+
const line =
43+
' '.repeat(depth * 2) +
44+
attributes.join(' ') +
45+
(node.id === this.snapshot?.selectedElementUid
46+
? ' [selected in the DevTools Elements panel]'
47+
: '') +
48+
'\n';
49+
chunks.push(line);
4450

45-
function getAttributes(serializedAXNodeRoot: TextSnapshotNode): string[] {
46-
const attributes = [`uid=${serializedAXNodeRoot.id}`];
47-
if (serializedAXNodeRoot.role) {
48-
// To match representation in DevTools.
49-
attributes.push(
50-
serializedAXNodeRoot.role === 'none'
51-
? 'ignored'
52-
: serializedAXNodeRoot.role,
53-
);
51+
for (const child of node.children) {
52+
chunks.push(this.#formatNode(child, depth + 1));
53+
}
54+
return chunks.join('');
5455
}
55-
if (serializedAXNodeRoot.name) {
56-
attributes.push(`"${serializedAXNodeRoot.name}"`);
56+
57+
#nodeToJSON(node: TextSnapshotNode): object {
58+
const rawAttrs = this.#getAttributesMap(node);
59+
const children = node.children.map(child => this.#nodeToJSON(child));
60+
const result: Record<string, unknown> = {
61+
...rawAttrs,
62+
};
63+
if (children.length > 0) {
64+
result.children = children;
65+
}
66+
return result;
5767
}
5868

59-
const excluded = new Set([
60-
'id',
61-
'role',
62-
'name',
63-
'elementHandle',
64-
'children',
65-
'backendNodeId',
66-
]);
67-
68-
const booleanPropertyMap: Record<string, string> = {
69-
disabled: 'disableable',
70-
expanded: 'expandable',
71-
focused: 'focusable',
72-
selected: 'selectable',
73-
};
74-
75-
for (const attr of Object.keys(serializedAXNodeRoot).sort()) {
76-
if (excluded.has(attr)) {
77-
continue;
69+
#getAttributes(serializedAXNodeRoot: TextSnapshotNode): string[] {
70+
const attributes = [`uid=${serializedAXNodeRoot.id}`];
71+
72+
if (serializedAXNodeRoot.role) {
73+
attributes.push(
74+
serializedAXNodeRoot.role === 'none'
75+
? 'ignored'
76+
: serializedAXNodeRoot.role,
77+
);
78+
}
79+
if (serializedAXNodeRoot.name) {
80+
attributes.push(`"${serializedAXNodeRoot.name}"`);
7881
}
79-
const value = (serializedAXNodeRoot as unknown as Record<string, unknown>)[
80-
attr
81-
];
82-
if (typeof value === 'boolean') {
83-
if (booleanPropertyMap[attr]) {
84-
attributes.push(booleanPropertyMap[attr]);
82+
83+
const simpleAttrs = this.#getAttributesMap(
84+
serializedAXNodeRoot,
85+
/* excludeSpecial */ true,
86+
);
87+
88+
for (const attr of Object.keys(serializedAXNodeRoot).sort()) {
89+
if (excludedAttributes.has(attr)) {
90+
continue;
8591
}
86-
if (value) {
92+
93+
const mapped = booleanPropertyMap[attr];
94+
if (mapped && simpleAttrs[mapped]) {
95+
attributes.push(mapped);
96+
}
97+
98+
const val = simpleAttrs[attr];
99+
if (val === true) {
87100
attributes.push(attr);
101+
} else if (typeof val === 'string' || typeof val === 'number') {
102+
attributes.push(`${attr}="${val}"`);
88103
}
89-
} else if (typeof value === 'string' || typeof value === 'number') {
90-
attributes.push(`${attr}="${value}"`);
91104
}
105+
106+
return attributes;
107+
}
108+
109+
#getAttributesMap(
110+
node: TextSnapshotNode,
111+
excludeSpecial = false,
112+
): Record<string, unknown> {
113+
const result: Record<string, unknown> = {};
114+
if (!excludeSpecial) {
115+
result.id = node.id;
116+
if (node.role) result.role = node.role;
117+
if (node.name) result.name = node.name;
118+
}
119+
120+
// Re-implementing the exact logic from original function for #getAttributes to be safe:
121+
return {
122+
...result,
123+
...this.#extractedAttributes(node),
124+
};
125+
}
126+
127+
#extractedAttributes(node: TextSnapshotNode): Record<string, unknown> {
128+
const result: Record<string, unknown> = {};
129+
130+
for (const attr of Object.keys(node).sort()) {
131+
if (excludedAttributes.has(attr)) {
132+
continue;
133+
}
134+
const value = (node as unknown as Record<string, unknown>)[attr];
135+
if (typeof value === 'boolean') {
136+
if (booleanPropertyMap[attr]) {
137+
result[booleanPropertyMap[attr]] = true;
138+
}
139+
if (value) {
140+
result[attr] = true;
141+
}
142+
} else if (typeof value === 'string' || typeof value === 'number') {
143+
result[attr] = value;
144+
}
145+
}
146+
return result;
92147
}
93-
return attributes;
94148
}
149+
150+
const booleanPropertyMap: Record<string, string> = {
151+
disabled: 'disableable',
152+
expanded: 'expandable',
153+
focused: 'focusable',
154+
selected: 'selectable',
155+
};
156+
157+
const excludedAttributes = new Set([
158+
'id',
159+
'role',
160+
'name',
161+
'elementHandle',
162+
'children',
163+
'backendNodeId',
164+
]);

‎src/main.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ function registerTool(tool: ToolDefinition): void {
153153
response,
154154
context,
155155
);
156-
const content = await response.handle(tool.name, context);
156+
const {content} = await response.handle(tool.name, context);
157157
return {
158158
content,
159159
};

0 commit comments

Comments
 (0)