fix session manager bugs

Author: Kylejeong2

src/context.ts

@@ -3,7 +3,7 @@ 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 { getSession, defaultSessionId } from "./sessionManager.js";
+import { SessionManager } from "./sessionManager.js";
 import type { MCPTool, BrowserSession } from "./types/types.js";
 
 /**
@@ -16,24 +16,38 @@ import type { MCPTool, BrowserSession } from "./types/types.js";
 export class Context {
   public readonly config: Config;
   private server: Server;
-  public currentSessionId: string = defaultSessionId;
+  private sessionManager: SessionManager;
 
-  constructor(server: Server, config: Config) {
+  // currentSessionId is a getter that delegates to SessionManager to ensure synchronization
+  // This prevents desync between Context and SessionManager session tracking
+  public get currentSessionId(): string {
+    return this.sessionManager.getActiveSessionId();
+  }
+
+  constructor(server: Server, config: Config, contextId?: string) {
     this.server = server;
     this.config = config;
+    this.sessionManager = new SessionManager(contextId);
   }
 
   public getServer(): Server {
     return this.server;
   }
 
+  public getSessionManager(): SessionManager {
+    return this.sessionManager;
+  }
+
   /**
    * Gets the Stagehand instance for the current session from SessionManager
    */
   public async getStagehand(
     sessionId: string = this.currentSessionId,
   ): Promise<Stagehand> {
-    const session = await getSession(sessionId, this.config);
+    const session = await this.sessionManager.getSession(
+      sessionId,
+      this.config,
+    );
     if (!session) {
       throw new Error(`No session found for ID: ${sessionId}`);
     }
@@ -42,7 +56,10 @@ export class Context {
 
   public async getActivePage(): Promise<BrowserSession["page"] | null> {
     // Get page from session manager
-    const session = await getSession(this.currentSessionId, this.config);
+    const session = await this.sessionManager.getSession(
+      this.currentSessionId,
+      this.config,
+    );
     if (session && session.page && !session.page.isClosed()) {
       return session.page;
     }
@@ -53,7 +70,7 @@ export class Context {
   public async getActiveBrowser(
     createIfMissing: boolean = true,
   ): Promise<BrowserSession["browser"] | null> {
-    const session = await getSession(
+    const session = await this.sessionManager.getSession(
       this.currentSessionId,
       this.config,
       createIfMissing,

src/index.ts

@@ -1,6 +1,8 @@
 import * as dotenv from "dotenv";
 dotenv.config();
 
+import { randomUUID } from "crypto";
+
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z } from "zod";
 import type { MCPToolsArray } from "./types/types.js";
@@ -151,7 +153,8 @@ export default function ({ config }: { config: z.infer<typeof configSchema> }) {
   const internalConfig: Config = config as Config;
 
   // Create the context, passing server instance and config
-  const context = new Context(server.server, internalConfig);
+  const contextId = randomUUID();
+  const context = new Context(server.server, internalConfig, contextId);
 
   server.server.registerCapabilities({
     resources: {

src/program.ts

@@ -6,7 +6,6 @@ import { fileURLToPath } from "url";
 import createServerFunction from "./index.js";
 import { ServerList } from "./server.js";
 import { startHttpTransport, startStdioTransport } from "./transport.js";
-import * as stagehandStore from "./stagehandStore.js";
 
 import { resolveConfig } from "./config.js";
 
@@ -85,7 +84,8 @@ function setupExitWatchdog(serverList: ServerList) {
   const handleExit = async () => {
     setTimeout(() => process.exit(0), 15000);
     try {
-      await Promise.all([stagehandStore.removeAll(), serverList.closeAll()]);
+      // SessionManager within each server handles session cleanup
+      await serverList.closeAll();
     } catch (error) {
       console.error("Error during cleanup:", error);
     }