feat: implement takeViewportScreenshot tool (#26)

shadowusr · web-flow · commit 20ac3b1129b3 · 2025-08-28T15:04:00.000+03:00
* feat: implement takeViewportScreenshot tool
diff --git a/README.md b/README.md
@@ -327,4 +327,12 @@ Capture a DOM snapshot of the current page with configurable filtering options.
 
 **Note:** By default, only useful tags and attributes are included in snapshots. The response will indicate what was omitted. Use the filtering options only if you need specific content that's not included by default.
 
+### `takeViewportScreenshot`
+Capture a PNG screenshot of the current browser viewport.
+
+- **Parameters:**
+  - `filePath` (string, optional): Path to save the screenshot (defaults to tmp directory with timestamp in file name)
+
+**Note:** Screenshots are saved as PNG files. If no filePath is provided, the screenshot will be saved to the system's temporary directory with a filePath like `viewport-{timestamp}.png`.
+
 </details>
diff --git a/src/browser-context.ts b/src/browser-context.ts
@@ -71,6 +71,5 @@ export class BrowserContext {
             /* empty */
         }
         return false;
-        // return this._browser !== null;
     }
 }
diff --git a/src/tools/index.ts b/src/tools/index.ts
@@ -5,6 +5,7 @@ import { clickOnElement } from "./click-on-element.js";
 import { hoverElement } from "./hover-element.js";
 import { typeIntoElement } from "./type-into-element.js";
 import { takePageSnapshot } from "./take-page-snapshot.js";
+import { takeViewportScreenshot } from "./take-viewport-screenshot.js";
 import { listTabs } from "./list-tabs.js";
 import { switchToTab } from "./switch-to-tab.js";
 import { openNewTab } from "./open-new-tab.js";
@@ -19,6 +20,7 @@ export const tools = [
     hoverElement,
     typeIntoElement,
     takePageSnapshot,
+    takeViewportScreenshot,
     listTabs,
     switchToTab,
     openNewTab,
diff --git a/src/tools/take-viewport-screenshot.ts b/src/tools/take-viewport-screenshot.ts
@@ -0,0 +1,53 @@
+import { z } from "zod";
+import { ToolDefinition } from "../types.js";
+import { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { contextProvider } from "../context-provider.js";
+import { createBrowserStateResponse, createErrorResponse } from "../responses/index.js";
+import path from "path";
+import fs from "fs/promises";
+import os from "os";
+
+export const takeViewportScreenshotSchema = {
+    filePath: z.string().optional().describe("Path to save the screenshot (defaults to tmp directory)"),
+};
+
+const takeViewportScreenshotCb: ToolCallback<typeof takeViewportScreenshotSchema> = async args => {
+    try {
+        const context = contextProvider.getContext();
+        const browser = await context.browser.get();
+
+        const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
+        const defaultFilePath = path.join(os.tmpdir(), `viewport-${timestamp}.png`);
+        const filePath = args.filePath || defaultFilePath;
+
+        const screenshotDir = path.dirname(filePath);
+        await fs.mkdir(screenshotDir, { recursive: true });
+
+        await browser.saveScreenshot(filePath);
+
+        const fileStats = await fs.stat(filePath);
+        const fileSizeKB = Math.round(fileStats.size / 1024);
+
+        const additionalInfo = `Screenshot saved: ${filePath} (${fileSizeKB} KB)`;
+
+        return await createBrowserStateResponse(browser, {
+            action: "Viewport screenshot captured successfully",
+            testplaneCode: `await browser.saveScreenshot("${filePath}");`,
+            additionalInfo,
+            isSnapshotNeeded: false,
+        });
+    } catch (error) {
+        console.error("Error taking viewport screenshot:", error);
+        return createErrorResponse("Error taking viewport screenshot", error instanceof Error ? error : undefined);
+    }
+};
+
+export const takeViewportScreenshot: ToolDefinition<typeof takeViewportScreenshotSchema> = {
+    name: "takeViewportScreenshot",
+    description:
+        "Capture a PNG screenshot of the current browser viewport. " +
+        "Strongly prefer capturing text-based snapshots using takePageSnapshot tool. " +
+        "Only use to test for visual changes when text-based snapshots are not useful.",
+    schema: takeViewportScreenshotSchema,
+    cb: takeViewportScreenshotCb,
+};
diff --git a/test/tools/take-viewport-screenshot.test.ts b/test/tools/take-viewport-screenshot.test.ts
@@ -0,0 +1,238 @@
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { describe, it, expect, beforeAll, beforeEach, afterEach, afterAll } from "vitest";
+import { startClient } from "../utils";
+import { INTEGRATION_TEST_TIMEOUT } from "../constants";
+import { PlaygroundServer } from "../test-server";
+import fs from "fs/promises";
+import path from "path";
+import os from "os";
+
+describe(
+    "tools/takeViewportScreenshot",
+    () => {
+        let client: Client;
+        let playgroundUrl: string;
+        let testServer: PlaygroundServer;
+
+        beforeAll(async () => {
+            testServer = new PlaygroundServer();
+            playgroundUrl = await testServer.start();
+        }, 20000);
+
+        afterAll(async () => {
+            if (testServer) {
+                await testServer.stop();
+            }
+        });
+
+        beforeEach(async () => {
+            client = await startClient();
+        });
+
+        afterEach(async () => {
+            if (client) {
+                await client.close();
+            }
+        });
+
+        describe("takeViewportScreenshot tool availability", () => {
+            it("should list takeViewportScreenshot tool in available tools", async () => {
+                const tools = await client.listTools();
+
+                const screenshotTool = tools.tools.find(tool => tool.name === "takeViewportScreenshot");
+
+                expect(screenshotTool).toBeDefined();
+                expect(screenshotTool?.description).toContain(
+                    "Capture a PNG screenshot of the current browser viewport",
+                );
+                expect(screenshotTool?.inputSchema.properties).toHaveProperty("filePath");
+            });
+        });
+
+        describe("takeViewportScreenshot tool execution", () => {
+            it("should capture screenshot with default filePath in tmp directory", async () => {
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: {},
+                });
+
+                expect(result.isError).toBe(false);
+                expect(result.content).toBeDefined();
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                expect(content).toHaveLength(1);
+                expect(content[0].type).toBe("text");
+
+                const responseText = content[0].text;
+
+                expect(responseText).toContain("✅ Viewport screenshot captured successfully");
+                expect(responseText).toContain(os.tmpdir());
+                expect(responseText).toContain("viewport-");
+                expect(responseText).toContain(".png");
+                expect(responseText).toContain("KB)");
+
+                const filePathMatch = responseText.match(/Screenshot saved: ([^\s]+)/);
+                expect(filePathMatch).toBeTruthy();
+
+                if (filePathMatch) {
+                    const filePath = filePathMatch[1];
+                    const fileExists = await fs
+                        .access(filePath)
+                        .then(() => true)
+                        .catch(() => false);
+                    expect(fileExists).toBe(true);
+
+                    if (fileExists) {
+                        const stats = await fs.stat(filePath);
+                        expect(stats.size).toBeGreaterThan(0);
+                        await fs.unlink(filePath).catch(() => {});
+                    }
+                }
+            });
+
+            it("should capture screenshot with custom filePath", async () => {
+                const customFilePath = path.join(os.tmpdir(), "test-screenshot.png");
+
+                await fs.unlink(customFilePath).catch(() => {});
+
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: { filePath: customFilePath },
+                });
+
+                expect(result.isError).toBe(false);
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                const responseText = content[0].text;
+
+                expect(responseText).toContain("✅ Viewport screenshot captured successfully");
+                expect(responseText).toContain(`Screenshot saved: ${customFilePath}`);
+                expect(responseText).toContain(`await browser.saveScreenshot("${customFilePath}")`);
+
+                const fileExists = await fs
+                    .access(customFilePath)
+                    .then(() => true)
+                    .catch(() => false);
+                expect(fileExists).toBe(true);
+
+                if (fileExists) {
+                    const stats = await fs.stat(customFilePath);
+                    expect(stats.size).toBeGreaterThan(0);
+                    await fs.unlink(customFilePath).catch(() => {});
+                }
+            });
+
+            it("should create directory if it doesn't exist", async () => {
+                const testDir = path.join(os.tmpdir(), "test-screenshots-" + Date.now());
+                const customFilePath = path.join(testDir, "nested", "screenshot.png");
+
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: { filePath: customFilePath },
+                });
+
+                expect(result.isError).toBe(false);
+
+                const fileExists = await fs
+                    .access(customFilePath)
+                    .then(() => true)
+                    .catch(() => false);
+                expect(fileExists).toBe(true);
+
+                if (fileExists) {
+                    await fs.rm(testDir, { recursive: true, force: true }).catch(() => {});
+                }
+            });
+
+            it("should include browser tabs information in response", async () => {
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: {},
+                });
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                const responseText = content[0].text;
+
+                expect(responseText).toContain("## Browser Tabs");
+                expect(responseText).toContain("Element Click Test Playground");
+            });
+
+            it("should not include page snapshot in response", async () => {
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: {},
+                });
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                const responseText = content[0].text;
+
+                expect(responseText).not.toContain("## Current Tab Snapshot");
+            });
+        });
+
+        describe("takeViewportScreenshot error handling", () => {
+            it("should handle screenshot capture when no page is loaded", async () => {
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: {},
+                });
+
+                expect(result.isError).toBe(false);
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                expect(content).toHaveLength(1);
+
+                const responseText = content[0].text;
+                expect(responseText).toContain("Screenshot saved:");
+            });
+
+            it("should handle invalid file path gracefully", async () => {
+                const invalidPath = "/invalid\0path/screenshot.png";
+
+                await client.callTool({
+                    name: "navigate",
+                    arguments: { url: playgroundUrl },
+                });
+
+                const result = await client.callTool({
+                    name: "takeViewportScreenshot",
+                    arguments: { filePath: invalidPath },
+                });
+
+                expect(result.isError).toBe(true);
+
+                const content = result.content as Array<{ type: string; text: string }>;
+                const responseText = content[0].text;
+
+                expect(responseText).toContain("❌");
+                expect(responseText).toContain("Error taking viewport screenshot");
+            });
+        });
+    },
+    INTEGRATION_TEST_TIMEOUT,
+);

Original file line number	Diff line number	Diff line change
`@@ -71,6 +71,5 @@ export class BrowserContext {`
`71`	`71`	`/* empty */`
`72`	`72`	`}`
`73`	`73`	`return false;`
`74`		`- // return this._browser !== null;`
`75`	`74`	`}`
`76`	`75`	`}`