diff --git a/CHANGELOG.md b/CHANGELOG.md index 2804db7..2430253 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # @browserbasehq/mcp-server-browserbase +## 2.1.2 + +### Patch Changes + +- fixing screenshot map behavior + ## 2.1.1 ### Patch Changes diff --git a/package.json b/package.json index 978df8f..c25598d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@browserbasehq/mcp-server-browserbase", - "version": "2.1.1", + "version": "2.1.2", "description": "MCP server for AI web browser automation using Browserbase and Stagehand", "mcpName": "io.github.browserbase/mcp-server-browserbase", "license": "Apache-2.0", diff --git a/src/mcp/resources.ts b/src/mcp/resources.ts index 0aaa1a3..5b9a327 100644 --- a/src/mcp/resources.ts +++ b/src/mcp/resources.ts @@ -13,6 +13,39 @@ export const RESOURCE_TEMPLATES = []; // Store screenshots in a map export const screenshots = new Map(); +// Track screenshots by session so we can purge them on session end +// key: sessionId (internal/current session id), value: set of screenshot names +const sessionIdToScreenshotNames = new Map>(); + +export function registerScreenshot( + sessionId: string, + name: string, + base64: string, +) { + screenshots.set(name, base64); + let set = sessionIdToScreenshotNames.get(sessionId); + if (!set) { + set = new Set(); + sessionIdToScreenshotNames.set(sessionId, set); + } + set.add(name); +} + +export function clearScreenshotsForSession(sessionId: string) { + const set = sessionIdToScreenshotNames.get(sessionId); + if (set) { + for (const name of set) { + screenshots.delete(name); + } + sessionIdToScreenshotNames.delete(sessionId); + } +} + +export function clearAllScreenshots() { + screenshots.clear(); + sessionIdToScreenshotNames.clear(); +} + /** * Handle listing resources request * @returns A list of available resources including screenshots diff --git a/src/sessionManager.ts b/src/sessionManager.ts index 10df670..5093a7e 100644 --- a/src/sessionManager.ts +++ b/src/sessionManager.ts @@ -2,6 +2,7 @@ import { Page, BrowserContext } from "@browserbasehq/stagehand"; import type { Config } from "../config.d.ts"; import type { Cookie } from "playwright-core"; import { createStagehandInstance } from "./stagehandStore.js"; +import { clearScreenshotsForSession } from "./mcp/resources.js"; import type { BrowserSession } from "./types/types.js"; // Global state for managing browser sessions @@ -131,6 +132,21 @@ export async function createNewBrowserSession( ); setActiveSessionId(defaultSessionId); } + + // Purge any screenshots associated with both internal and Browserbase IDs + try { + clearScreenshotsForSession(newSessionId); + const bbId = browserbaseSessionId; + if (bbId) { + clearScreenshotsForSession(bbId); + } + } catch (err) { + process.stderr.write( + `[SessionManager] WARN - Failed to clear screenshots on disconnect for ${newSessionId}: ${ + err instanceof Error ? err.message : String(err) + }\n`, + ); + } }); // Add cookies to the context if they are provided in the config @@ -192,6 +208,20 @@ async function closeBrowserGracefully( process.stderr.write( `[SessionManager] Successfully closed Stagehand and browser for session: ${sessionIdToLog}\n`, ); + // After close, purge any screenshots associated with both internal and Browserbase IDs + try { + clearScreenshotsForSession(sessionIdToLog); + const bbId = session?.stagehand?.browserbaseSessionID; + if (bbId) { + clearScreenshotsForSession(bbId); + } + } catch (err) { + process.stderr.write( + `[SessionManager] WARN - Failed to clear screenshots after close for ${sessionIdToLog}: ${ + err instanceof Error ? err.message : String(err) + }\n`, + ); + } } catch (closeError) { process.stderr.write( `[SessionManager] WARN - Error closing Stagehand for session ${sessionIdToLog}: ${ @@ -335,6 +365,17 @@ export async function cleanupSession(sessionId: string): Promise { // Remove from browsers map browsers.delete(sessionId); + // Always purge screenshots for this (internal) session id + try { + clearScreenshotsForSession(sessionId); + } catch (err) { + process.stderr.write( + `[SessionManager] WARN - Failed to clear screenshots during cleanup for ${sessionId}: ${ + err instanceof Error ? err.message : String(err) + }\n`, + ); + } + // Clear default session reference if this was the default if (sessionId === defaultSessionId && defaultBrowserSession) { defaultBrowserSession = null; diff --git a/src/stagehandStore.ts b/src/stagehandStore.ts index 9f239a3..fd644d1 100644 --- a/src/stagehandStore.ts +++ b/src/stagehandStore.ts @@ -2,6 +2,7 @@ import { randomUUID } from "crypto"; import { Stagehand, Page } from "@browserbasehq/stagehand"; import { StagehandSession, CreateSessionParams } from "./types/types.js"; import type { Config } from "../config.d.ts"; +import { clearScreenshotsForSession } from "./mcp/resources.js"; // Store for all active sessions const store = new Map(); @@ -110,6 +111,18 @@ export const create = async ( const disconnectHandler = () => { process.stderr.write(`[StagehandStore] Session disconnected: ${id}\n`); store.delete(id); + // Purge by internal store ID and Browserbase session ID + try { + clearScreenshotsForSession(id); + const bbId = session.metadata?.bbSessionId; + if (bbId) { + clearScreenshotsForSession(bbId); + } + } catch { + process.stderr.write( + `[StagehandStore] Error clearing screenshots for session ${id}\n`, + ); + } }; browser.on("disconnected", disconnectHandler); @@ -158,6 +171,18 @@ export const remove = async (id: string): Promise => { await session.stagehand.close(); process.stderr.write(`[StagehandStore] Session closed: ${id}\n`); + // Purge by internal store ID and Browserbase session ID + try { + clearScreenshotsForSession(id); + const bbId = session.metadata?.bbSessionId; + if (bbId) { + clearScreenshotsForSession(bbId); + } + } catch { + process.stderr.write( + `[StagehandStore] Error clearing screenshots for session ${id}\n`, + ); + } } catch (error) { process.stderr.write( `[StagehandStore] Error closing session ${id}: ${ diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts index a680799..48a568c 100644 --- a/src/tools/screenshot.ts +++ b/src/tools/screenshot.ts @@ -2,7 +2,7 @@ import { z } from "zod"; import type { Tool, ToolSchema, ToolResult } from "./tool.js"; import type { Context } from "../context.js"; import type { ToolActionResult } from "../types/types.js"; -import { screenshots } from "../mcp/resources.js"; +import { registerScreenshot } from "../mcp/resources.js"; const ScreenshotInputSchema = z.object({ name: z.string().optional().describe("The name of the screenshot"), @@ -40,7 +40,9 @@ async function handleScreenshot( .replace(/:/g, "-")}` : `screenshot-${new Date().toISOString().replace(/:/g, "-")}` + context.config.browserbaseProjectId; - screenshots.set(name, screenshotBase64); + // Associate with current session id and store in memory + const sessionId = context.currentSessionId; + registerScreenshot(sessionId, name, screenshotBase64); // Notify the client that the resources changed const serverInstance = context.getServer();