feat(screenshot): adds ability to output screenshot to a specific path. Closes #152. Closes #153

Author: james2doyle

src/tools/screenshot.ts

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import {writeFile} from 'node:fs/promises';
 import type {ElementHandle, Page} from 'puppeteer-core';
 import z from 'zod';
 
@@ -34,6 +35,12 @@ export const screenshot = defineTool({
       .describe(
         'If set to true takes a screenshot of the full page instead of the currently visible viewport. Incompatible with uid.',
       ),
+    filePath: z
+      .string()
+      .optional()
+      .describe(
+        'The path to save the screenshot to. If provided, the screenshot will be saved to this path instead of being attached to the response.',
+      ),
   },
   handler: async (request, response, context) => {
     if (request.params.uid && request.params.fullPage) {
@@ -66,7 +73,12 @@ export const screenshot = defineTool({
       );
     }
 
-    if (screenshot.length >= 2_000_000) {
+    if (request.params.filePath) {
+      await writeFile(request.params.filePath, screenshot);
+      response.appendResponseLine(
+        `Saved screenshot to ${request.params.filePath}.`,
+      );
+    } else if (screenshot.length >= 2_000_000) {
       const {filename} = await context.saveTemporaryFile(
         screenshot,
         `image/${request.params.format}`,

tests/tools/screenshot.test.ts

@@ -4,6 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import assert from 'node:assert';
+import {rm, stat, mkdir, chmod} from 'node:fs/promises';
+import {tmpdir} from 'node:os';
+import {join} from 'node:path';
 import {describe, it} from 'node:test';
 
 import {screenshot} from '../../src/tools/screenshot.js';
@@ -108,5 +111,102 @@ describe('screenshot', () => {
         );
       });
     });
+
+    it('with filePath', async () => {
+      await withBrowser(async (response, context) => {
+        const filePath = join(tmpdir(), 'test-screenshot.png');
+        try {
+          const fixture = screenshots.basic;
+          const page = context.getSelectedPage();
+          await page.setContent(fixture.html);
+          await screenshot.handler(
+            {params: {format: 'png', filePath}},
+            response,
+            context,
+          );
+
+          assert.equal(response.images.length, 0);
+          assert.equal(
+            response.responseLines.at(0),
+            "Took a screenshot of the current page's viewport.",
+          );
+          assert.equal(
+            response.responseLines.at(1),
+            `Saved screenshot to ${filePath}.`,
+          );
+
+          const stats = await stat(filePath);
+          assert.ok(stats.isFile());
+          assert.ok(stats.size > 0);
+        } finally {
+          await rm(filePath, {force: true});
+        }
+      });
+    });
+
+    it('with unwritable filePath', async () => {
+      const dir = join(tmpdir(), 'readonly-dir-for-screenshot-test');
+      await mkdir(dir, {recursive: true});
+      await chmod(dir, 0o500);
+      const filePath = join(dir, 'test-screenshot.png');
+
+      try {
+        await withBrowser(async (response, context) => {
+          const fixture = screenshots.basic;
+          const page = context.getSelectedPage();
+          await page.setContent(fixture.html);
+          await screenshot.handler(
+            {params: {format: 'png', filePath}},
+            response,
+            context,
+          );
+
+          assert.equal(response.images.length, 0);
+          assert.equal(
+            response.responseLines.at(0),
+            "Took a screenshot of the current page's viewport.",
+          );
+          assert.ok(
+            response.responseLines
+              .at(1)
+              ?.startsWith(`Could not write screenshot to ${filePath}.`),
+            `Expected error message for unwritable path, but got: ${response.responseLines.at(
+              1,
+            )}`,
+          );
+        });
+      } finally {
+        await chmod(dir, 0o700);
+        await rm(dir, {recursive: true, force: true});
+      }
+    });
+
+    it('with malformed filePath', async () => {
+      await withBrowser(async (response, context) => {
+        const filePath = 'malformed\0path.png';
+        const fixture = screenshots.basic;
+        const page = context.getSelectedPage();
+        await page.setContent(fixture.html);
+        await screenshot.handler(
+          {params: {format: 'png', filePath}},
+          response,
+          context,
+        );
+
+        assert.equal(response.images.length, 0);
+        assert.equal(
+          response.responseLines.at(0),
+          "Took a screenshot of the current page's viewport.",
+        );
+        assert.ok(
+          response.responseLines
+            .at(1)
+            ?.startsWith(`Could not write screenshot to ${filePath}.`),
+          `Expected error message for malformed path, but got: ${response.responseLines.at(
+            1,
+          )}`,
+        );
+      });
+    });
   });
 });