feat: support saving snapshots to file

Author: OrKoN

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,8 @@ 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;
   }
 
   setIncludeNetworkRequests(
@@ -158,12 +160,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 +171,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 screenshot to ${this.#snapshotParams.filePath}.`;
+        } else {
+          formattedSnapshot = formatA11ySnapshot(snapshot.root);
+        }
+      }
     }
 
     const bodies: {
@@ -281,6 +293,7 @@ export class McpResponse implements Response {
       bodies,
       consoleData,
       consoleListData,
+      formattedSnapshot,
     });
   }
 
@@ -294,6 +307,7 @@ export class McpResponse implements Response {
       };
       consoleData: ConsoleMessageData | undefined;
       consoleListData: ConsoleMessageData[] | undefined;
+      formattedSnapshot: string | undefined;
     },
   ): Array<TextContent | ImageContent> {
     const response = [`# ${toolName} response`];
@@ -339,13 +353,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(true);
     } 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(true);
     } finally {
       void handle.dispose();
     }
@@ -159,7 +159,7 @@ export const fill = defineTool({
       );
     });
     response.appendResponseLine(`Successfully filled out the element`);
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot(true);
   },
 });
 
@@ -184,7 +184,7 @@ export const drag = defineTool({
         await toHandle.drop(fromHandle);
       });
       response.appendResponseLine(`Successfully dragged an element`);
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot(true);
     } 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(true);
   },
 });
 
@@ -264,7 +264,7 @@ export const uploadFile = defineTool({
           );
         }
       }
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot(true);
       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(true);
   },
 });

src/tools/snapshot.ts

@@ -24,9 +24,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 +57,6 @@ export const waitFor = defineTool({
       `Element with text "${request.params.text}" found.`,
     );
 
-    response.setIncludeSnapshot(true);
+    response.includeSnapshot(true);
   },
 });

tests/McpResponse.test.ts

@@ -54,7 +54,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(true);
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(
@@ -80,7 +80,7 @@ uid=1_0 RootWebArea
         /></label>`,
       );
       await page.focus('input');
-      response.setIncludeSnapshot(true);
+      response.includeSnapshot(true);
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(
@@ -99,7 +99,7 @@ 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(true, true);
       const result = await response.handle('test', context);
       assert.equal(result[0].type, 'text');
       assert.strictEqual(