Skip to content

Commit 398813c

Browse files
OrKoNOrKoN
committed
chore: structured snapshot
1 parent 79ab800 commit 398813c

3 files changed

Lines changed: 193 additions & 80 deletions

File tree

‎src/McpResponse.ts‎

Lines changed: 4 additions & 5 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,15 @@ 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(
199-
new TextEncoder().encode(
200-
formatSnapshotNode(snapshot.root, snapshot),
201-
),
200+
new TextEncoder().encode(formatter.toString()),
202201
this.#snapshotParams.filePath,
203202
);
204203
formattedSnapshot = `Saved snapshot to ${this.#snapshotParams.filePath}.`;
205204
} else {
206-
formattedSnapshot = formatSnapshotNode(snapshot.root, snapshot);
205+
formattedSnapshot = formatter.toString();
207206
}
208207
}
209208
}

‎src/formatters/snapshotFormatter.ts‎

Lines changed: 136 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,94 +1,164 @@
11
/**
22
* @license
3-
* Copyright 2025 Google LLC
3+
* Copyright 2026 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+
]);

‎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)