Skip to content

Commit 4bc19b8

Browse files
OrKoNOrKoN
committed
chore: structured snapshot
1 parent 79ab800 commit 4bc19b8

3 files changed

Lines changed: 185 additions & 81 deletions

File tree

‎src/McpResponse.ts‎

Lines changed: 4 additions & 3 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 {
@@ -194,16 +194,17 @@ export class McpResponse implements Response {
194194
);
195195
const snapshot = context.getTextSnapshot();
196196
if (snapshot) {
197+
const formatter = new SnapshotFormatter(snapshot);
197198
if (this.#snapshotParams.filePath) {
198199
await context.saveFile(
199200
new TextEncoder().encode(
200-
formatSnapshotNode(snapshot.root, snapshot),
201+
formatter.toString(),
201202
),
202203
this.#snapshotParams.filePath,
203204
);
204205
formattedSnapshot = `Saved snapshot to ${this.#snapshotParams.filePath}.`;
205206
} else {
206-
formattedSnapshot = formatSnapshotNode(snapshot.root, snapshot);
207+
formattedSnapshot = formatter.toString();
207208
}
208209
}
209210
}

‎src/formatters/snapshotFormatter.ts‎

Lines changed: 128 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,94 +1,153 @@
1-
/**
2-
* @license
3-
* Copyright 2025 Google LLC
4-
* SPDX-License-Identifier: Apache-2.0
5-
*/
6-
71
import type {TextSnapshot, TextSnapshotNode} from '../McpContext.js';
82

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

16-
if (depth === 0) {
1713
// Top-level content of the snapshot.
1814
if (
19-
snapshot?.verbose &&
20-
snapshot?.hasSelectedElement &&
21-
!snapshot.selectedElementUid
15+
this.snapshot.verbose &&
16+
this.snapshot.hasSelectedElement &&
17+
!this.snapshot.selectedElementUid
2218
) {
2319
chunks.push(`Note: there is a selected element in the DevTools Elements panel but it is not included into the current a11y tree snapshot.
2420
Get a verbose snapshot to include all elements if you are interested in the selected element.\n\n`);
2521
}
26-
}
2722

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);
23+
chunks.push(this.#formatNode(root, 0));
24+
return chunks.join('');
25+
}
3726

38-
for (const child of root.children) {
39-
chunks.push(formatSnapshotNode(child, snapshot, depth + 1));
27+
toJSON(): object {
28+
if (!this.snapshot) {
29+
return {};
30+
}
31+
return this.#nodeToJSON(this.snapshot.root);
4032
}
4133

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

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-
);
46+
for (const child of node.children) {
47+
chunks.push(this.#formatNode(child, depth + 1));
48+
}
49+
return chunks.join('');
5450
}
55-
if (serializedAXNodeRoot.name) {
56-
attributes.push(`"${serializedAXNodeRoot.name}"`);
51+
52+
#nodeToJSON(node: TextSnapshotNode): object {
53+
const rawAttrs = this.#getAttributesMap(node);
54+
const children = node.children.map(child => this.#nodeToJSON(child));
55+
const result: Record<string, unknown> = {
56+
...rawAttrs,
57+
};
58+
if (children.length > 0) {
59+
result.children = children;
60+
}
61+
return result;
5762
}
5863

59-
const excluded = new Set([
60-
'id',
61-
'role',
62-
'name',
63-
'elementHandle',
64-
'children',
65-
'backendNodeId',
66-
]);
64+
#getAttributes(serializedAXNodeRoot: TextSnapshotNode): string[] {
65+
const attributes = [`uid=${serializedAXNodeRoot.id}`];
66+
67+
if (serializedAXNodeRoot.role) {
68+
attributes.push(
69+
serializedAXNodeRoot.role === 'none'
70+
? 'ignored'
71+
: serializedAXNodeRoot.role,
72+
);
73+
}
74+
if (serializedAXNodeRoot.name) {
75+
attributes.push(`"${serializedAXNodeRoot.name}"`);
76+
}
6777

68-
const booleanPropertyMap: Record<string, string> = {
69-
disabled: 'disableable',
70-
expanded: 'expandable',
71-
focused: 'focusable',
72-
selected: 'selectable',
73-
};
78+
const simpleAttrs = this.#getAttributesMap(serializedAXNodeRoot, /* excludeSpecial */ true);
79+
80+
for (const attr of Object.keys(serializedAXNodeRoot).sort()) {
81+
if (excludedAttributes.has(attr)) {
82+
continue;
83+
}
84+
85+
const mapped = booleanPropertyMap[attr];
86+
if (mapped && simpleAttrs[mapped]) {
87+
attributes.push(mapped);
88+
}
89+
90+
const val = simpleAttrs[attr];
91+
if (val === true) {
92+
attributes.push(attr);
93+
} else if (typeof val === 'string' || typeof val === 'number') {
94+
attributes.push(`${attr}="${val}"`);
95+
}
96+
}
97+
98+
return attributes;
99+
}
74100

75-
for (const attr of Object.keys(serializedAXNodeRoot).sort()) {
76-
if (excluded.has(attr)) {
77-
continue;
101+
#getAttributesMap(node: TextSnapshotNode, excludeSpecial = false): Record<string, unknown> {
102+
const result: Record<string, unknown> = {};
103+
if (!excludeSpecial) {
104+
result.id = node.id;
105+
if (node.role) result.role = node.role;
106+
if (node.name) result.name = node.name;
78107
}
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]);
108+
109+
// Re-implementing the exact logic from original function for #getAttributes to be safe:
110+
return {
111+
...result,
112+
...this.#extractedAttributes(node),
113+
};
114+
}
115+
116+
#extractedAttributes(node: TextSnapshotNode): Record<string, unknown> {
117+
const result: Record<string, unknown> = {};
118+
119+
for (const attr of Object.keys(node).sort()) {
120+
if (excludedAttributes.has(attr)) {
121+
continue;
85122
}
86-
if (value) {
87-
attributes.push(attr);
123+
const value = (node as unknown as Record<string, unknown>)[attr];
124+
if (typeof value === 'boolean') {
125+
if (booleanPropertyMap[attr]) {
126+
result[booleanPropertyMap[attr]] = true;
127+
}
128+
if (value) {
129+
result[attr] = true;
130+
}
131+
} else if (typeof value === 'string' || typeof value === 'number') {
132+
result[attr] = value;
88133
}
89-
} else if (typeof value === 'string' || typeof value === 'number') {
90-
attributes.push(`${attr}="${value}"`);
91134
}
135+
return result;
92136
}
93-
return attributes;
94137
}
138+
139+
const booleanPropertyMap: Record<string, string> = {
140+
disabled: 'disableable',
141+
expanded: 'expandable',
142+
focused: 'focusable',
143+
selected: 'selectable',
144+
};
145+
146+
const excludedAttributes = new Set([
147+
'id',
148+
'role',
149+
'name',
150+
'elementHandle',
151+
'children',
152+
'backendNodeId',
153+
]);

‎tests/formatters/snapshotFormatter.test.ts‎

Lines changed: 53 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ import {describe, it} from 'node:test';
99

1010
import type {ElementHandle} from 'puppeteer-core';
1111

12-
import {formatSnapshotNode} from '../../src/formatters/snapshotFormatter.js';
13-
import type {TextSnapshotNode} from '../../src/McpContext.js';
12+
import {SnapshotFormatter} from '../../src/formatters/snapshotFormatter.js';
13+
import type {TextSnapshot, TextSnapshotNode} from '../../src/McpContext.js';
1414

1515
describe('snapshotFormatter', () => {
1616
it('formats a snapshot with value properties', () => {
@@ -35,7 +35,8 @@ describe('snapshotFormatter', () => {
3535
},
3636
};
3737

38-
const formatted = formatSnapshotNode(node);
38+
const formatter = new SnapshotFormatter({ root: node } as TextSnapshot);
39+
const formatted = formatter.toString();
3940
assert.strictEqual(
4041
formatted,
4142
`uid=1_1 textbox "textbox" value="value"
@@ -66,7 +67,8 @@ describe('snapshotFormatter', () => {
6667
},
6768
};
6869

69-
const formatted = formatSnapshotNode(node);
70+
const formatter = new SnapshotFormatter({ root: node } as TextSnapshot);
71+
const formatted = formatter.toString();
7072
assert.strictEqual(
7173
formatted,
7274
`uid=1_1 button "button" disableable disabled
@@ -97,7 +99,8 @@ describe('snapshotFormatter', () => {
9799
},
98100
};
99101

100-
const formatted = formatSnapshotNode(node);
102+
const formatter = new SnapshotFormatter({ root: node } as TextSnapshot);
103+
const formatted = formatter.toString();
101104
assert.strictEqual(
102105
formatted,
103106
`uid=1_1 checkbox "checkbox" checked
@@ -139,7 +142,8 @@ describe('snapshotFormatter', () => {
139142
},
140143
};
141144

142-
const formatted = formatSnapshotNode(node);
145+
const formatter = new SnapshotFormatter({ root: node } as TextSnapshot);
146+
const formatted = formatter.toString();
143147
assert.strictEqual(
144148
formatted,
145149
`uid=1_1 root "root"
@@ -171,13 +175,14 @@ describe('snapshotFormatter', () => {
171175
},
172176
};
173177

174-
const formatted = formatSnapshotNode(node, {
178+
const formatter = new SnapshotFormatter({
175179
snapshotId: '1',
176180
root: node,
177181
idToNode: new Map(),
178182
hasSelectedElement: true,
179183
verbose: false,
180184
});
185+
const formatted = formatter.toString();
181186

182187
t.assert.snapshot?.(formatted);
183188
});
@@ -204,13 +209,14 @@ describe('snapshotFormatter', () => {
204209
},
205210
};
206211

207-
const formatted = formatSnapshotNode(node, {
212+
const formatter = new SnapshotFormatter({
208213
snapshotId: '1',
209214
root: node,
210215
idToNode: new Map(),
211216
hasSelectedElement: true,
212217
verbose: true,
213218
});
219+
const formatted = formatter.toString();
214220

215221
t.assert.snapshot?.(formatted);
216222
});
@@ -237,15 +243,53 @@ describe('snapshotFormatter', () => {
237243
},
238244
};
239245

240-
const formatted = formatSnapshotNode(node, {
246+
const formatter = new SnapshotFormatter({
241247
snapshotId: '1',
242248
root: node,
243249
idToNode: new Map(),
244250
hasSelectedElement: true,
245251
selectedElementUid: '1_1',
246252
verbose: false,
247253
});
254+
const formatted = formatter.toString();
248255

249256
t.assert.snapshot?.(formatted);
250257
});
258+
259+
it('toJSON returns expected structure', () => {
260+
const node: TextSnapshotNode = {
261+
id: '1_1',
262+
role: 'root',
263+
name: 'root',
264+
children: [
265+
{
266+
id: '1_2',
267+
role: 'button',
268+
name: 'button',
269+
disabled: true,
270+
children: [],
271+
elementHandle: async () => null,
272+
},
273+
],
274+
elementHandle: async () => null,
275+
};
276+
277+
const formatter = new SnapshotFormatter({ root: node } as TextSnapshot);
278+
const json = formatter.toJSON();
279+
280+
assert.deepStrictEqual(json, {
281+
id: '1_1',
282+
role: 'root',
283+
name: 'root',
284+
children: [
285+
{
286+
id: '1_2',
287+
role: 'button',
288+
name: 'button',
289+
disableable: true,
290+
disabled: true,
291+
},
292+
],
293+
});
294+
});
251295
});

0 commit comments

Comments
 (0)