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

Author: Kylejeong2

src/mcp/resources.ts

@@ -13,6 +13,69 @@ 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();
+}
+
+/**
+ * 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<void> {
+  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

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

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

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