Skip to content

Commit 777333e

Browse files
aberemia24OrKoN
authored andcommitted
feat: add filePath parameter to take_snapshot and evaluate_script
Add optional filePath parameter to save tool output directly to files instead of including in response. Features: - take_snapshot can save accessibility tree to file - evaluate_script can save JSON result to file - Reduces token usage for large outputs - Enables offline analysis and processing - Backwards compatible - tools work as before when filePath is omitted Implementation: - Added filePath optional parameter to both tool schemas - Write formatted output to file using writeFile from node:fs/promises - Updated Context interface to expose createTextSnapshot and getTextSnapshot - Includes comprehensive tests for file saving functionality Resolves #153
1 parent 139ce60 commit 777333e

File tree

5 files changed

+117
-7
lines changed

5 files changed

+117
-7
lines changed

src/tools/ToolDefinition.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import type {Dialog, ElementHandle, Page} from 'puppeteer-core';
88
import z from 'zod';
99

10+
import type {TextSnapshot} from '../McpContext.js';
1011
import type {TraceResult} from '../trace-processing/parse.js';
1112

1213
import type {ToolCategories} from './categories.js';
@@ -79,6 +80,8 @@ export type Context = Readonly<{
7980
filename: string,
8081
): Promise<{filename: string}>;
8182
waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
83+
createTextSnapshot(): Promise<void>;
84+
getTextSnapshot(): TextSnapshot | null;
8285
}>;
8386

8487
export function defineTool<Schema extends z.ZodRawShape>(

src/tools/script.ts

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
* Copyright 2025 Google LLC
44
* SPDX-License-Identifier: Apache-2.0
55
*/
6+
import {writeFile} from 'node:fs/promises';
7+
68
import type {JSHandle} from 'puppeteer-core';
79
import z from 'zod';
810

@@ -42,6 +44,12 @@ Example with arguments: \`(el) => {
4244
)
4345
.optional()
4446
.describe(`An optional list of arguments to pass to the function.`),
47+
filePath: z
48+
.string()
49+
.optional()
50+
.describe(
51+
'The absolute path, or a path relative to the current working directory, to save the result to instead of including it in the response.',
52+
),
4553
},
4654
handler: async (request, response, context) => {
4755
const page = context.getSelectedPage();
@@ -59,10 +67,18 @@ Example with arguments: \`(el) => {
5967
},
6068
...args,
6169
);
62-
response.appendResponseLine('Script ran on page and returned:');
63-
response.appendResponseLine('```json');
64-
response.appendResponseLine(`${result}`);
65-
response.appendResponseLine('```');
70+
71+
if (request.params.filePath) {
72+
await writeFile(request.params.filePath, result);
73+
response.appendResponseLine(
74+
`Saved script result to ${request.params.filePath}.`,
75+
);
76+
} else {
77+
response.appendResponseLine('Script ran on page and returned:');
78+
response.appendResponseLine('```json');
79+
response.appendResponseLine(`${result}`);
80+
response.appendResponseLine('```');
81+
}
6682
});
6783
} finally {
6884
Promise.allSettled(args.map(arg => arg.dispose())).catch(() => {

src/tools/snapshot.ts

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,12 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7+
import {writeFile} from 'node:fs/promises';
8+
79
import {Locator} from 'puppeteer-core';
810
import z from 'zod';
911

12+
import {formatA11ySnapshot} from '../formatters/snapshotFormatter.js';
1013
import {ToolCategories} from './categories.js';
1114
import {defineTool, timeoutSchema} from './ToolDefinition.js';
1215

@@ -18,9 +21,33 @@ identifier (uid). Always use the latest snapshot. Prefer taking a snapshot over
1821
category: ToolCategories.DEBUGGING,
1922
readOnlyHint: true,
2023
},
21-
schema: {},
22-
handler: async (_request, response) => {
23-
response.setIncludeSnapshot(true);
24+
schema: {
25+
filePath: z
26+
.string()
27+
.optional()
28+
.describe(
29+
'The absolute path, or a path relative to the current working directory, to save the snapshot to instead of including it in the response.',
30+
),
31+
},
32+
handler: async (request, response, context) => {
33+
await context.createTextSnapshot();
34+
const snapshot = context.getTextSnapshot();
35+
36+
if (!snapshot) {
37+
response.appendResponseLine('No snapshot data available.');
38+
return;
39+
}
40+
41+
const formattedSnapshot = formatA11ySnapshot(snapshot.root);
42+
43+
if (request.params.filePath) {
44+
await writeFile(request.params.filePath, formattedSnapshot);
45+
response.appendResponseLine(
46+
`Saved snapshot to ${request.params.filePath}.`,
47+
);
48+
} else {
49+
response.setIncludeSnapshot(true);
50+
}
2451
},
2552
});
2653

tests/tools/script.test.ts

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66
import assert from 'node:assert';
7+
import fs from 'node:fs/promises';
8+
import path from 'node:path';
79
import {describe, it} from 'node:test';
810

911
import {evaluateScript} from '../../src/tools/script.js';
@@ -152,5 +154,38 @@ describe('script', () => {
152154
assert.strictEqual(JSON.parse(lineEvaluation), true);
153155
});
154156
});
157+
158+
it('saves result to file when filePath is provided', async () => {
159+
const testFilePath = path.join(process.cwd(), 'test-script-result.json');
160+
161+
await withBrowser(async (response, context) => {
162+
const page = context.getSelectedPage();
163+
await page.setContent(html`<div id="data">{"key":"value"}</div>`);
164+
165+
await evaluateScript.handler(
166+
{
167+
params: {
168+
function: String(() => ({result: 'test', number: 42})),
169+
filePath: testFilePath,
170+
},
171+
},
172+
response,
173+
context,
174+
);
175+
176+
assert.strictEqual(
177+
response.responseLines[0],
178+
`Saved script result to ${testFilePath}.`,
179+
);
180+
assert.strictEqual(response.responseLines.length, 1);
181+
182+
// Verify file was created and contains expected JSON
183+
const fileContent = await fs.readFile(testFilePath, 'utf-8');
184+
const parsed = JSON.parse(fileContent);
185+
assert.deepEqual(parsed, {result: 'test', number: 42});
186+
});
187+
188+
await fs.unlink(testFilePath);
189+
});
155190
});
156191
});

tests/tools/snapshot.test.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66
import assert from 'node:assert';
7+
import fs from 'node:fs/promises';
8+
import path from 'node:path';
79
import {describe, it} from 'node:test';
810

911
import {takeSnapshot, waitFor} from '../../src/tools/snapshot.js';
@@ -17,6 +19,33 @@ describe('snapshot', () => {
1719
assert.ok(response.includeSnapshot);
1820
});
1921
});
22+
23+
it('saves snapshot to file when filePath is provided', async () => {
24+
const testFilePath = path.join(process.cwd(), 'test-snapshot.txt');
25+
26+
await withBrowser(async (response, context) => {
27+
const page = context.getSelectedPage();
28+
await page.setContent(html`<button>Click me</button>`);
29+
30+
await takeSnapshot.handler(
31+
{params: {filePath: testFilePath}},
32+
response,
33+
context,
34+
);
35+
36+
assert.strictEqual(response.includeSnapshot, false);
37+
assert.strictEqual(
38+
response.responseLines[0],
39+
`Saved snapshot to ${testFilePath}.`,
40+
);
41+
42+
// Verify file was created and contains button text
43+
const fileContent = await fs.readFile(testFilePath, 'utf-8');
44+
assert.ok(fileContent.includes('Click me'));
45+
});
46+
47+
await fs.unlink(testFilePath);
48+
});
2049
});
2150
describe('browser_wait_for', () => {
2251
it('should work', async () => {

0 commit comments

Comments
 (0)