fix: clear session screenshots in the mcp

[Kylejeong2]

what

On the remote host shttp URL on smithery, we had a global in-memory cache—with all of the screenshots. This fix is to make sure we map the screenshots to the MCP session ID and then delete and clear them afterwards when the session's done/closed.

[greptile-apps]

Greptile Overview

Summary

This PR implements session-based screenshot cleanup to prevent memory leaks by mapping screenshots to MCP session IDs and clearing them when sessions end.

• Added registerScreenshot() and clearScreenshotsForSession() functions to track screenshots by session ID
• Integrated cleanup calls across multiple session lifecycle points (disconnect, close, cleanup)
• Handles both internal session IDs and Browserbase session IDs for comprehensive cleanup
• Uses try-catch blocks with warning logs to prevent cleanup failures from affecting session management

The implementation addresses the global in-memory cache issue mentioned in the description, ensuring screenshots are properly cleaned up when sessions terminate.

Confidence Score: 3/5

• This PR is moderately safe to merge but has timing and error handling issues that should be addressed
• Score reflects solid core functionality but concerns about timing of cleanup calls during session creation and inconsistent error handling across modules
• src/sessionManager.ts needs attention for timing issues with cleanup during session creation

Important Files Changed

File Analysis

Filename	Score	Overview
src/mcp/resources.ts	4/5	Added session-based screenshot tracking with register/clear functions - clean implementation
src/sessionManager.ts	2/5	Added screenshot cleanup in multiple session lifecycle hooks - potential timing and error handling issues
src/stagehandStore.ts	3/5	Added screenshot cleanup on session disconnect/remove - good pattern but swallows errors

Sequence Diagram

sequenceDiagram
    participant Client
    participant SessionManager
    participant StagehandStore
    participant Resources
    participant ScreenshotTool

    Note over Client,Resources: Screenshot Creation Flow
    Client->>SessionManager: Create session
    SessionManager->>StagehandStore: create(sessionId)
    StagehandStore->>Resources: (session established)
    
    Client->>ScreenshotTool: Take screenshot
    ScreenshotTool->>Resources: registerScreenshot(sessionId, name, base64)
    Resources->>Resources: Map sessionId -> screenshot names
    
    Note over Client,Resources: Screenshot Cleanup Flow
    Client->>SessionManager: Disconnect/Close session
    SessionManager->>Resources: clearScreenshotsForSession(sessionId)
    SessionManager->>Resources: clearScreenshotsForSession(bbSessionId)
    
    alt StagehandStore disconnect
        StagehandStore->>Resources: clearScreenshotsForSession(sessionId)
        StagehandStore->>Resources: clearScreenshotsForSession(bbSessionId)
    end
    
    alt SessionManager cleanup
        SessionManager->>Resources: clearScreenshotsForSession(sessionId)
    end
    
    Resources->>Resources: Delete screenshots from memory

Loading

_{3 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: Screenshot cleanup is called in a disconnect handler that runs during session creation, but the session might not be fully established yet. This could lead to premature cleanup of screenshots that are still being used.

Suggested change

	// 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`,
	);
	// Only clear screenshots if the session was actually established
	if (session?.stagehand) {
	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`,
	);
	}
	}

[greptile-apps]

logic: Silent error handling might hide important failures during screenshot cleanup.

Suggested change

	}
	} catch {
	process.stderr.write(
	`[StagehandStore] Error clearing screenshots for session ${id}\n`,
	);
	}
	} catch (err) {
	process.stderr.write(
	`[StagehandStore] Error clearing screenshots for session ${id}: ${
	err instanceof Error ? err.message : String(err)
	}\n`,
	);
	}

[greptile-apps]

logic: Silent error handling might hide important failures during screenshot cleanup.

Suggested change

	process.stderr.write(
	`[StagehandStore] Error clearing screenshots for session ${id}\n`,
	);
	}
	} catch (error) {
	process.stderr.write(
	} catch (err) {
	process.stderr.write(
	`[StagehandStore] Error clearing screenshots for session ${id}: ${
	err instanceof Error ? err.message : String(err)
	}\n`,
	);
	}