hotfix for screenshots

Author: Kylejeong2

src/context.ts

@@ -2,7 +2,6 @@ import type { Stagehand } from "@browserbasehq/stagehand";
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { Config } from "../config.d.ts";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { listResources, readResource } from "./mcp/resources.js";
 import { SessionManager } from "./sessionManager.js";
 import type { MCPTool } from "./types/types.js";
 
@@ -17,6 +16,7 @@ export class Context {
   public readonly config: Config;
   private server: Server;
   private sessionManager: SessionManager;
+  private screenshots: Map<string, string> = new Map();
 
   // currentSessionId is a getter that delegates to SessionManager to ensure synchronization
   // This prevents desync between Context and SessionManager session tracking
@@ -104,19 +104,57 @@ export class Context {
     }
   }
 
+  /**
+   * Register a screenshot in this context's storage
+   */
+  registerScreenshot(name: string, data: string): void {
+    this.screenshots.set(name, data);
+  }
+
+  /**
+   * Clear all screenshots in this context
+   */
+  clearScreenshots(): void {
+    this.screenshots.clear();
+  }
+
   /**
    * List resources
    * Documentation: https://modelcontextprotocol.io/docs/concepts/resources
    */
   listResources() {
-    return listResources();
+    return {
+      resources: [
+        ...Array.from(this.screenshots.keys()).map((name) => ({
+          uri: `screenshot://${name}`,
+          mimeType: "image/png",
+          name: `Screenshot: ${name}`,
+        })),
+      ],
+    };
   }
 
   /**
    * Read a resource by URI
    * Documentation: https://modelcontextprotocol.io/docs/concepts/resources
    */
   readResource(uri: string) {
-    return readResource(uri);
+    if (uri.startsWith("screenshot://")) {
+      const name = uri.split("://")[1];
+      const screenshot = this.screenshots.get(name);
+      if (screenshot) {
+        return {
+          contents: [
+            {
+              uri,
+              mimeType: "image/png",
+              blob: screenshot,
+            },
+          ],
+        };
+      }
+    }
+
+    throw new Error(`Resource not found: ${uri}`);
   }
 }

src/mcp/resources.ts

@@ -2,95 +2,13 @@
  * Resources module for the Browserbase MCP server
  * Contains resources definitions and handlers for resource-related requests
  * Docs: https://modelcontextprotocol.io/docs/concepts/resources
+ *
+ * Note: Screenshot storage has been moved to Context class for per-session isolation.
+ * Each MCP session has its own Context instance with isolated screenshot storage.
  */
 
 // Define the resources
 export const RESOURCES = [];
 
 // Define the resource templates
 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
- */
-export function listResources() {
-  return {
-    resources: [
-      ...Array.from(screenshots.keys()).map((name) => ({
-        uri: `screenshot://${name}`,
-        mimeType: "image/png",
-        name: `Screenshot: ${name}`,
-      })),
-    ],
-  };
-}
-
-/**
- * Handle listing resource templates request
- * @returns An empty resource templates list response
- */
-export function listResourceTemplates() {
-  return { resourceTemplates: [] };
-}
-
-/**
- * Read a resource by its URI
- * @param uri The URI of the resource to read
- * @returns The resource content or throws if not found
- */
-export function readResource(uri: string) {
-  if (uri.startsWith("screenshot://")) {
-    const name = uri.split("://")[1];
-    const screenshot = screenshots.get(name);
-    if (screenshot) {
-      return {
-        contents: [
-          {
-            uri,
-            mimeType: "image/png",
-            blob: screenshot,
-          },
-        ],
-      };
-    }
-  }
-
-  throw new Error(`Resource not found: ${uri}`);
-}

src/sessionManager.ts

@@ -1,6 +1,5 @@
 import { Stagehand } from "@browserbasehq/stagehand";
 import type { Config } from "../config.d.ts";
-import { clearScreenshotsForSession } from "./mcp/resources.js";
 import type { BrowserSession, CreateSessionParams } from "./types/types.js";
 import { randomUUID } from "crypto";
 
@@ -243,16 +242,6 @@ export class SessionManager {
           process.stderr.write(
             `[SessionManager] Successfully closed Stagehand and browser for session: ${sessionIdToLog}\n`,
           );
-          // After close, purge any screenshots associated with this session
-          try {
-            clearScreenshotsForSession(sessionIdToLog);
-          } 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}: ${

src/tools/screenshot.ts

@@ -2,7 +2,6 @@ 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 { registerScreenshot } from "../mcp/resources.js";
 
 /**
  * Screenshot
@@ -49,9 +48,8 @@ async function handleScreenshot(
         : `screenshot-${new Date().toISOString().replace(/:/g, "-")}` +
           context.config.browserbaseProjectId;
 
-      // Associate with current mcp session id and store in memory /src/mcp/resources.ts
-      const sessionId = context.currentSessionId;
-      registerScreenshot(sessionId, name, screenshotBase64);
+      // Store screenshot in this context (scoped to this MCP session)
+      context.registerScreenshot(name, screenshotBase64);
 
       // Notify the client that the resources changed
       const serverInstance = context.getServer();