feat: support saving snapshots to file (#463)

Refs #153

---------

Co-authored-by: Nikolay Vitkov <[email protected]>

Author: OrKoN

docs/tool-reference.md

@@ -342,6 +342,7 @@ identifier (uid). Always use the latest snapshot. Prefer taking a snapshot over
 
 **Parameters:**
 
+- **filePath** (string) _(optional)_: The absolute path, or a path relative to the current working directory, to save the snapshot to instead of attaching it to the response.
 - **verbose** (boolean) _(optional)_: Whether to include all possible information available in the full a11y tree. Default is false.
 
 ---

src/McpResponse.ts

@@ -24,14 +24,17 @@ import type {
   TextContent,
 } from './third_party/index.js';
 import {handleDialog} from './tools/pages.js';
-import type {ImageContentData, Response} from './tools/ToolDefinition.js';
+import type {
+  ImageContentData,
+  Response,
+  SnapshotParams,
+} from './tools/ToolDefinition.js';
 import {paginate} from './utils/pagination.js';
 import type {PaginationOptions} from './utils/types.js';
 
 export class McpResponse implements Response {
   #includePages = false;
-  #includeSnapshot = false;
-  #includeVerboseSnapshot = false;
+  #snapshotParams?: SnapshotParams;
   #attachedNetworkRequestId?: number;
   #attachedConsoleMessageId?: number;
   #textResponseLines: string[] = [];
@@ -53,9 +56,10 @@ export class McpResponse implements Response {
     this.#includePages = value;
   }
 
-  setIncludeSnapshot(value: boolean, verbose = false): void {
-    this.#includeSnapshot = value;
-    this.#includeVerboseSnapshot = verbose;
+  includeSnapshot(params?: SnapshotParams): void {
+    this.#snapshotParams = params ?? {
+      verbose: false,
+    };
   }
 
   setIncludeNetworkRequests(
@@ -158,12 +162,8 @@ export class McpResponse implements Response {
     return this.#images;
   }
 
-  get includeSnapshot(): boolean {
-    return this.#includeSnapshot;
-  }
-
-  get includeVersboseSnapshot(): boolean {
-    return this.#includeVerboseSnapshot;
+  get snapshotParams(): SnapshotParams | undefined {
+    return this.#snapshotParams;
   }
 
   async handle(
@@ -173,8 +173,22 @@ export class McpResponse implements Response {
     if (this.#includePages) {
       await context.createPagesSnapshot();
     }
-    if (this.#includeSnapshot) {
-      await context.createTextSnapshot(this.#includeVerboseSnapshot);
+
+    let formattedSnapshot: string | undefined;
+    if (this.#snapshotParams) {
+      await context.createTextSnapshot(this.#snapshotParams.verbose);
+      const snapshot = context.getTextSnapshot();
+      if (snapshot) {
+        if (this.#snapshotParams.filePath) {
+          await context.saveFile(
+            new TextEncoder().encode(formatA11ySnapshot(snapshot.root)),
+            this.#snapshotParams.filePath,
+          );
+          formattedSnapshot = `Saved snapshot to ${this.#snapshotParams.filePath}.`;
+        } else {
+          formattedSnapshot = formatA11ySnapshot(snapshot.root);
+        }
+      }
     }
 
     const bodies: {
@@ -281,6 +295,7 @@ export class McpResponse implements Response {
       bodies,
       consoleData,
       consoleListData,
+      formattedSnapshot,
     });
   }
 
@@ -294,6 +309,7 @@ export class McpResponse implements Response {
       };
       consoleData: ConsoleMessageData | undefined;
       consoleListData: ConsoleMessageData[] | undefined;
+      formattedSnapshot: string | undefined;
     },
   ): Array<TextContent | ImageContent> {
     const response = [`# ${toolName} response`];
@@ -339,13 +355,9 @@ Call ${handleDialog.name} to handle it before continuing.`);
       response.push(...parts);
     }
 
-    if (this.#includeSnapshot) {
-      const snapshot = context.getTextSnapshot();
-      if (snapshot) {
-        const formattedSnapshot = formatA11ySnapshot(snapshot.root);
-        response.push('## Page content');
-        response.push(formattedSnapshot);
-      }
+    if (data.formattedSnapshot) {
+      response.push('## Page content');
+      response.push(data.formattedSnapshot);
     }
 
     response.push(...this.#formatNetworkRequestData(context, data.bodies));

src/tools/ToolDefinition.ts

@@ -42,6 +42,11 @@ export interface ImageContentData {
   mimeType: string;
 }
 
+export interface SnapshotParams {
+  verbose?: boolean;
+  filePath?: string;
+}
+
 export interface Response {
   appendResponseLine(value: string): void;
   setIncludePages(value: boolean): void;
@@ -59,8 +64,7 @@ export interface Response {
       includePreservedMessages?: boolean;
     },
   ): void;
-  setIncludeSnapshot(value: boolean): void;
-  setIncludeSnapshot(value: boolean, verbose?: boolean): void;
+  includeSnapshot(params?: SnapshotParams): void;
   attachImage(value: ImageContentData): void;
   attachNetworkRequest(reqid: number): void;
   attachConsoleMessage(msgid: number): void;

src/tools/input.ts

@@ -44,7 +44,7 @@ export const click = defineTool({
           ? `Successfully double clicked on the element`
           : `Successfully clicked on the element`,
       );
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
     } finally {
       void handle.dispose();
     }
@@ -73,7 +73,7 @@ export const hover = defineTool({
         await handle.asLocator().hover();
       });
       response.appendResponseLine(`Successfully hovered over the element`);
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
     } finally {
       void handle.dispose();
     }
@@ -159,7 +159,7 @@ export const fill = defineTool({
       );
     });
     response.appendResponseLine(`Successfully filled out the element`);
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot();
   },
 });
 
@@ -184,7 +184,7 @@ export const drag = defineTool({
         await toHandle.drop(fromHandle);
       });
       response.appendResponseLine(`Successfully dragged an element`);
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
     } finally {
       void fromHandle.dispose();
       void toHandle.dispose();
@@ -220,7 +220,7 @@ export const fillForm = defineTool({
       });
     }
     response.appendResponseLine(`Successfully filled out the form`);
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot();
   },
 });
 
@@ -264,7 +264,7 @@ export const uploadFile = defineTool({
           );
         }
       }
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
       response.appendResponseLine(`File uploaded from ${filePath}.`);
     } finally {
       void handle.dispose();
@@ -304,6 +304,6 @@ export const pressKey = defineTool({
     response.appendResponseLine(
       `Successfully pressed key: ${request.params.key}`,
     );
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot();
   },
 });

src/tools/screenshot.ts

@@ -15,7 +15,8 @@ export const screenshot = defineTool({
   description: `Take a screenshot of the page or element.`,
   annotations: {
     category: ToolCategory.DEBUGGING,
-    readOnlyHint: true,
+    // Not read-only due to filePath param.
+    readOnlyHint: false,
   },
   schema: {
     format: zod

src/tools/snapshot.ts

@@ -15,7 +15,8 @@ export const takeSnapshot = defineTool({
 identifier (uid). Always use the latest snapshot. Prefer taking a snapshot over taking a screenshot.`,
   annotations: {
     category: ToolCategory.DEBUGGING,
-    readOnlyHint: true,
+    // Not read-only due to filePath param.
+    readOnlyHint: false,
   },
   schema: {
     verbose: zod
@@ -24,9 +25,18 @@ identifier (uid). Always use the latest snapshot. Prefer taking a snapshot over
       .describe(
         'Whether to include all possible information available in the full a11y tree. Default is false.',
       ),
+    filePath: zod
+      .string()
+      .optional()
+      .describe(
+        'The absolute path, or a path relative to the current working directory, to save the snapshot to instead of attaching it to the response.',
+      ),
   },
   handler: async (request, response) => {
-    response.setIncludeSnapshot(true, request.params.verbose ?? false);
+    response.includeSnapshot({
+      verbose: request.params.verbose ?? false,
+      filePath: request.params.filePath,
+    });
   },
 });
 
@@ -48,6 +58,6 @@ export const waitFor = defineTool({
       `Element with text "${request.params.text}" found.`,
     );
 
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot();
   },
 });

tests/McpResponse.test.ts

@@ -4,6 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import assert from 'node:assert';
+import {readFile, rm} from 'node:fs/promises';
+import {tmpdir} from 'node:os';
+import {join} from 'node:path';
 import {describe, it} from 'node:test';
 
 import {getMockRequest, getMockResponse, html, withBrowser} from './utils.js';
@@ -54,7 +57,7 @@ Testing 2`,
       await page.setContent(`<!DOCTYPE html>
 <button>Click me</button><input type="text" value="Input">`);
       await page.focus('button');
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(
@@ -80,7 +83,7 @@ uid=1_0 RootWebArea
         /></label>`,
       );
       await page.focus('input');
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot();
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(
@@ -99,7 +102,9 @@ uid=1_0 RootWebArea "My test page"
     await withBrowser(async (response, context) => {
       const page = context.getSelectedPage();
       await page.setContent(html`<aside>test</aside>`);
-      response.setIncludeSnapshot(true, true);
+      response.includeSnapshot({
+        verbose: true,
+      });
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(
@@ -117,6 +122,42 @@ uid=1_0 RootWebArea "My test page"
     });
   });
 
+  it('saves snapshot to file', async () => {
+    const filePath = join(tmpdir(), 'test-screenshot.png');
+    try {
+      await withBrowser(async (response, context) => {
+        const page = context.getSelectedPage();
+        await page.setContent(html`<aside>test</aside>`);
+        response.includeSnapshot({
+          verbose: true,
+          filePath,
+        });
+        const result = await response.handle('test', context);
+        assert.equal(result[0].type, 'text');
+        console.log(result[0].text);
+        assert.strictEqual(
+          result[0].text,
+          `# test response
+## Page content
+Saved snapshot to ${filePath}.`,
+        );
+      });
+      const content = await readFile(filePath, 'utf-8');
+      assert.strictEqual(
+        content,
+        `uid=1_0 RootWebArea "My test page"
+  uid=1_1 ignored
+    uid=1_2 ignored
+      uid=1_3 complementary
+        uid=1_4 StaticText "test"
+          uid=1_5 InlineTextBox "test"
+`,
+      );
+    } finally {
+      await rm(filePath, {force: true});
+    }
+  });
+
   it('adds throttling setting when it is not null', async () => {
     await withBrowser(async (response, context) => {
       context.setNetworkConditions('Slow 3G');

tests/tools/input.test.ts

@@ -427,7 +427,7 @@ describe('input', () => {
         );
 
         assert.strictEqual(response.responseLines.length, 0);
-        assert.strictEqual(response.includeSnapshot, false);
+        assert.strictEqual(response.snapshotParams, undefined);
 
         await fs.unlink(testFilePath);
       });