fix: clear session screenshots in the mcp (#120)

* ensuring screenshots are deleted by mapping sessionids to the screenshot then deleting them after

* changesets

* addressing greptile comments

Author: Kylejeong2

CHANGELOG.md

@@ -1,5 +1,11 @@
 # @browserbasehq/mcp-server-browserbase
 
+## 2.1.2
+
+### Patch Changes
+
+- fixing screenshot map behavior
+
 ## 2.1.1
 
 ### Patch Changes

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",

src/mcp/resources.ts

@@ -13,6 +13,39 @@ export const RESOURCE_TEMPLATES = [];
 // Store screenshots in a map
 export const screenshots = new Map<string, string>();
 
+// 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<string, Set<string>>();
+
+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

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<void> {
   // 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;

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<string, StagehandSession>();
@@ -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<void> => {
 
     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}: ${

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();