From cdfe2a7b944f24b9817ad83df363de18e64eed26 Mon Sep 17 00:00:00 2001 From: Kylejeong2 Date: Thu, 25 Sep 2025 11:37:04 -0700 Subject: [PATCH 1/3] ensuring screenshots are deleted by mapping sessionids to the screenshot then deleting them after --- src/mcp/resources.ts | 63 +++++++++++++++++++++++++++++++++++++++++ src/sessionManager.ts | 17 +++++++++++ src/stagehandStore.ts | 13 +++++++++ src/tools/screenshot.ts | 6 ++-- 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/src/mcp/resources.ts b/src/mcp/resources.ts index 0aaa1a3..312f242 100644 --- a/src/mcp/resources.ts +++ b/src/mcp/resources.ts @@ -13,6 +13,69 @@ 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(); +} + +/** + * Retry wrapper for clearing screenshots for a session. Uses exponential backoff. + * This protects against transient failures in any consumer of the map. + */ +export async function retryClearScreenshotsForSession( + sessionId: string, + options: { retries?: number; initialDelayMs?: number } = {}, +): Promise { + const retries = options.retries ?? 3; + const initialDelayMs = options.initialDelayMs ?? 50; + + let attempt = 0; + let delay = initialDelayMs; + // Attempt at least once + while (true) { + try { + clearScreenshotsForSession(sessionId); + return; + } catch (err) { + if (attempt >= retries) { + // Give up + throw err instanceof Error ? err : new Error(String(err)); + } + await new Promise((r) => setTimeout(r, delay)); + attempt += 1; + delay = Math.min(delay * 2, 1000); + } + } +} + /** * Handle listing resources request * @returns A list of available resources including screenshots diff --git a/src/sessionManager.ts b/src/sessionManager.ts index 10df670..7c98996 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 { retryClearScreenshotsForSession } from "./mcp/resources.js"; import type { BrowserSession } from "./types/types.js"; // Global state for managing browser sessions @@ -131,6 +132,13 @@ export async function createNewBrowserSession( ); setActiveSessionId(defaultSessionId); } + + // Purge any screenshots associated with both internal and Browserbase IDs + retryClearScreenshotsForSession(newSessionId).catch(() => {}); + const bbId = browserbaseSessionId; + if (bbId) { + retryClearScreenshotsForSession(bbId).catch(() => {}); + } }); // Add cookies to the context if they are provided in the config @@ -192,6 +200,12 @@ 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 + retryClearScreenshotsForSession(sessionIdToLog).catch(() => {}); + const bbId = session?.stagehand?.browserbaseSessionID; + if (bbId) { + retryClearScreenshotsForSession(bbId).catch(() => {}); + } } catch (closeError) { process.stderr.write( `[SessionManager] WARN - Error closing Stagehand for session ${sessionIdToLog}: ${ @@ -335,6 +349,9 @@ export async function cleanupSession(sessionId: string): Promise { // Remove from browsers map browsers.delete(sessionId); + // Always purge screenshots for this (internal) session id + await retryClearScreenshotsForSession(sessionId).catch(() => {}); + // 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..ee2230a 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 { retryClearScreenshotsForSession } from "./mcp/resources.js"; // Store for all active sessions const store = new Map(); @@ -110,6 +111,12 @@ 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 + retryClearScreenshotsForSession(id).catch(() => {}); + const bbId = session.metadata?.bbSessionId; + if (bbId) { + retryClearScreenshotsForSession(bbId).catch(() => {}); + } }; browser.on("disconnected", disconnectHandler); @@ -158,6 +165,12 @@ 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 + await retryClearScreenshotsForSession(id).catch(() => {}); + const bbId = session.metadata?.bbSessionId; + if (bbId) { + await retryClearScreenshotsForSession(bbId).catch(() => {}); + } } 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(); From 9d8e720cdf97cec37f1a53165ae626080d7380c2 Mon Sep 17 00:00:00 2001 From: Kylejeong2 Date: Thu, 25 Sep 2025 11:38:08 -0700 Subject: [PATCH 2/3] changesets --- CHANGELOG.md | 6 ++++++ package.json | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) 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", From 5f410656415e208eb51716a8deafb1d6e052defe Mon Sep 17 00:00:00 2001 From: Kylejeong2 Date: Thu, 25 Sep 2025 12:07:56 -0700 Subject: [PATCH 3/3] addressing greptile comments --- src/mcp/resources.ts | 30 ----------------------------- src/sessionManager.ts | 44 +++++++++++++++++++++++++++++++++---------- src/stagehandStore.ts | 30 ++++++++++++++++++++--------- 3 files changed, 55 insertions(+), 49 deletions(-) diff --git a/src/mcp/resources.ts b/src/mcp/resources.ts index 312f242..5b9a327 100644 --- a/src/mcp/resources.ts +++ b/src/mcp/resources.ts @@ -46,36 +46,6 @@ export function clearAllScreenshots() { sessionIdToScreenshotNames.clear(); } -/** - * Retry wrapper for clearing screenshots for a session. Uses exponential backoff. - * This protects against transient failures in any consumer of the map. - */ -export async function retryClearScreenshotsForSession( - sessionId: string, - options: { retries?: number; initialDelayMs?: number } = {}, -): Promise { - const retries = options.retries ?? 3; - const initialDelayMs = options.initialDelayMs ?? 50; - - let attempt = 0; - let delay = initialDelayMs; - // Attempt at least once - while (true) { - try { - clearScreenshotsForSession(sessionId); - return; - } catch (err) { - if (attempt >= retries) { - // Give up - throw err instanceof Error ? err : new Error(String(err)); - } - await new Promise((r) => setTimeout(r, delay)); - attempt += 1; - delay = Math.min(delay * 2, 1000); - } - } -} - /** * Handle listing resources request * @returns A list of available resources including screenshots diff --git a/src/sessionManager.ts b/src/sessionManager.ts index 7c98996..5093a7e 100644 --- a/src/sessionManager.ts +++ b/src/sessionManager.ts @@ -2,7 +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 { retryClearScreenshotsForSession } from "./mcp/resources.js"; +import { clearScreenshotsForSession } from "./mcp/resources.js"; import type { BrowserSession } from "./types/types.js"; // Global state for managing browser sessions @@ -134,10 +134,18 @@ export async function createNewBrowserSession( } // Purge any screenshots associated with both internal and Browserbase IDs - retryClearScreenshotsForSession(newSessionId).catch(() => {}); - const bbId = browserbaseSessionId; - if (bbId) { - retryClearScreenshotsForSession(bbId).catch(() => {}); + 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`, + ); } }); @@ -201,10 +209,18 @@ async function closeBrowserGracefully( `[SessionManager] Successfully closed Stagehand and browser for session: ${sessionIdToLog}\n`, ); // After close, purge any screenshots associated with both internal and Browserbase IDs - retryClearScreenshotsForSession(sessionIdToLog).catch(() => {}); - const bbId = session?.stagehand?.browserbaseSessionID; - if (bbId) { - retryClearScreenshotsForSession(bbId).catch(() => {}); + 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( @@ -350,7 +366,15 @@ export async function cleanupSession(sessionId: string): Promise { browsers.delete(sessionId); // Always purge screenshots for this (internal) session id - await retryClearScreenshotsForSession(sessionId).catch(() => {}); + 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) { diff --git a/src/stagehandStore.ts b/src/stagehandStore.ts index ee2230a..fd644d1 100644 --- a/src/stagehandStore.ts +++ b/src/stagehandStore.ts @@ -2,7 +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 { retryClearScreenshotsForSession } from "./mcp/resources.js"; +import { clearScreenshotsForSession } from "./mcp/resources.js"; // Store for all active sessions const store = new Map(); @@ -112,10 +112,16 @@ export const create = async ( process.stderr.write(`[StagehandStore] Session disconnected: ${id}\n`); store.delete(id); // Purge by internal store ID and Browserbase session ID - retryClearScreenshotsForSession(id).catch(() => {}); - const bbId = session.metadata?.bbSessionId; - if (bbId) { - retryClearScreenshotsForSession(bbId).catch(() => {}); + try { + clearScreenshotsForSession(id); + const bbId = session.metadata?.bbSessionId; + if (bbId) { + clearScreenshotsForSession(bbId); + } + } catch { + process.stderr.write( + `[StagehandStore] Error clearing screenshots for session ${id}\n`, + ); } }; @@ -166,10 +172,16 @@ 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 - await retryClearScreenshotsForSession(id).catch(() => {}); - const bbId = session.metadata?.bbSessionId; - if (bbId) { - await retryClearScreenshotsForSession(bbId).catch(() => {}); + 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(