stagehand tools using browser context similar to browserbase + stagehand init issues fixed

Author: Kylejeong2

stagehand/package.json

@@ -7,11 +7,16 @@
   "homepage": "https://modelcontextprotocol.io",
   "bugs": "https://github.com/modelcontextprotocol/servers/issues",
   "type": "module",
+  "main": "./cli.js",
   "bin": {
-    "mcp-server-stagehand": "dist/index.js"
+    "mcp-server-stagehand": "cli.js"
   },
   "files": [
-    "dist"
+    "dist",
+    "cli.js",
+    "index.d.ts",
+    "index.js",
+    "config.d.ts"
   ],
   "scripts": {
     "build": "tsc && shx chmod +x dist/*.js",

stagehand/src/config.ts

@@ -3,29 +3,10 @@ import fs from 'fs';
 import path from 'path';
 import { sanitizeForFilePath } from './tools/utils.js'; 
 import type { Cookie } from "playwright-core";
+import type { Config } from '../config.js';
 
 export type ToolCapability = 'core' | string; 
 
-export interface Config {
-  browserbaseApiKey?: string; 
-  browserbaseProjectId?: string; 
-  server?: {
-    port?: number;
-    host?: string;
-  };
-  proxies?: boolean;
-  advancedStealth?: boolean;
-  context?: {
-    contextId?: string;
-    persist?: boolean;
-  };
-  viewPort?: {
-    browserWidth?: number;
-    browserHeight?: number;
-  };
-  cookies?: Cookie[]; 
-}
-
 // Define Command Line Options Structure
 export type CLIOptions = {
   browserbaseApiKey?: string;

stagehand/src/context.ts

@@ -1,7 +1,7 @@
 import { Stagehand } from "@browserbasehq/stagehand";
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { Config } from "../config.js";
-import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { CallToolResult, TextContent, ImageContent } from "@modelcontextprotocol/sdk/types.js";
 import { handleToolCall } from "./tools/tools.js";
 import { listResources, readResource } from "./resources.js";
 import { 
@@ -12,11 +12,23 @@ import {
   setServerInstance,
   log 
 } from "./logging.js";
+import { 
+  getSession, 
+  getSessionReadOnly, 
+  defaultSessionId, 
+  type BrowserSession 
+} from "./sessionManager.js";
+
+export type ToolActionResult =
+  | { content?: (ImageContent | TextContent)[] }
+  | undefined
+  | void;
 
 export class Context {
-  private stagehand: Stagehand;
+  private stagehands = new Map<string, Stagehand>();
+  public readonly config: Config;
   private server: Server;
-  private config: Config;
+  public currentSessionId: string = defaultSessionId;
 
   constructor(server: Server, config: Config) {
     this.server = server;
@@ -28,26 +40,171 @@ export class Context {
     setupLogRotation();
     registerExitHandlers();
     scheduleLogRotation();
+  }
+
+  /**
+   * Gets the Stagehand instance for the current session, creating one if needed
+   */
+  public async getStagehand(sessionId: string = this.currentSessionId): Promise<Stagehand> {
+    let stagehand = this.stagehands.get(sessionId);
+    
+    if (!stagehand) {
+      // Create a new Stagehand instance for this session
+      stagehand = new Stagehand({
+        env: "BROWSERBASE",
+        logger: (logLine) => {
+          log(`Stagehand[${sessionId}]: ${logLine.message}`, 'info');
+        },
+      });
+      this.stagehands.set(sessionId, stagehand);
+    }
+
+    await stagehand.init();
+    
+    return stagehand;
+  }
+
+  /**
+   * Sets the Stagehand instance for a specific session
+   */
+  public setStagehand(sessionId: string, stagehand: Stagehand): void {
+    this.stagehands.set(sessionId, stagehand);
+  }
+
+  /**
+   * Removes the Stagehand instance for a specific session
+   */
+  public async removeStagehand(sessionId: string): Promise<void> {
+    const stagehand = this.stagehands.get(sessionId);
+    if (stagehand) {
+      try {
+        await stagehand.close();
+      } catch (error) {
+        log(`Error closing Stagehand for session ${sessionId}: ${error}`, 'error');
+      }
+      this.stagehands.delete(sessionId);
+    }
+  }
 
-    // Initialize Stagehand
-    this.stagehand = new Stagehand({
-      env: "BROWSERBASE",
-      logger: (logLine) => {
-        log(`Stagehand: ${logLine.message}`, 'info');
-      },
-    });
+  public async getActivePage(): Promise<BrowserSession["page"] | null> {
+    // Try to get page from Stagehand first (if available for this session)
+    const stagehand = this.stagehands.get(this.currentSessionId);
+    if (stagehand && stagehand.page && !stagehand.page.isClosed()) {
+      return stagehand.page;
+    }
+    
+    // Fallback to session manager
+    const session = await getSession(this.currentSessionId, this.config);
+    if (!session || !session.page || session.page.isClosed()) {
+      try {
+        const currentSession = await getSession(
+          this.currentSessionId,
+          this.config
+        );
+        if (
+          !currentSession ||
+          !currentSession.page ||
+          currentSession.page.isClosed()
+        ) {
+          return null;
+        }
+        return currentSession.page;
+      } catch (refreshError) {
+        return null;
+      }
+    }
+    return session.page;
+  }
+
+  // Will create a new default session if one doesn't exist
+  public async getActiveBrowser(): Promise<BrowserSession["browser"] | null> {
+    const session = await getSession(this.currentSessionId, this.config);
+    if (!session || !session.browser || !session.browser.isConnected()) {
+      try {
+        const currentSession = await getSession(
+          this.currentSessionId,
+          this.config
+        );
+        if (
+          !currentSession ||
+          !currentSession.browser ||
+          !currentSession.browser.isConnected()
+        ) {
+          return null;
+        }
+        return currentSession.browser;
+      } catch (refreshError) {
+        return null;
+      }
+    }
+    return session.browser;
+  }
+
+  /**
+   * Get the active browser without triggering session creation.
+   * This is a read-only operation used when we need to check for an existing browser
+   * without side effects (e.g., during close operations).
+   * @returns The browser if it exists and is connected, null otherwise
+   */
+  public getActiveBrowserReadOnly(): BrowserSession["browser"] | null {
+    const session = getSessionReadOnly(this.currentSessionId);
+    if (!session || !session.browser || !session.browser.isConnected()) {
+      return null;
+    }
+    return session.browser;
+  }
+
+  /**
+   * Get the active page without triggering session creation.
+   * This is a read-only operation used when we need to check for an existing page
+   * without side effects.
+   * @returns The page if it exists and is not closed, null otherwise
+   */
+  public getActivePageReadOnly(): BrowserSession["page"] | null {
+    const session = getSessionReadOnly(this.currentSessionId);
+    if (!session || !session.page || session.page.isClosed()) {
+      return null;
+    }
+    return session.page;
   }
 
   async run(tool: any, args: any): Promise<CallToolResult> {
     try {
-      log(`Executing tool: ${tool.name} with args: ${JSON.stringify(args)}`, 'info');
-      const result = await handleToolCall(tool.name, args, this.stagehand);
-      log(`Tool ${tool.name} completed successfully`, 'info');
-      return result;
+      log(`Executing tool: ${tool.schema.name} with args: ${JSON.stringify(args)}`, 'info');
+      
+      // Check if this tool has a handle method (new session tools)
+      // Only use handle method for session create and close tools
+      if ("handle" in tool && typeof tool.handle === "function" && 
+          (tool.schema.name === "browserbase_session_create" || tool.schema.name === "browserbase_session_close")) {
+        const toolResult = await tool.handle(this as any, args);
+        
+        if (toolResult?.action) {
+          const actionResult = await toolResult.action();
+          const content = actionResult?.content || [];
+          
+          return {
+            content: Array.isArray(content) ? content : [{ type: "text", text: "Action completed successfully." }],
+            isError: false,
+          };
+        } else {
+          return {
+            content: [{ type: "text", text: `${tool.schema.name} completed successfully.` }],
+            isError: false,
+          };
+        }
+      } else {
+        const stagehand = await this.getStagehand();
+        const result = await handleToolCall(tool.schema.name, args, stagehand);
+        log(`Tool ${tool.schema.name} completed successfully`, 'info');
+        return result;
+      }
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : String(error);
-      log(`Tool ${tool.name} failed: ${errorMessage}`, 'error');
-      throw error;
+      log(`Tool ${tool.schema?.name || 'unknown'} failed: ${errorMessage}`, 'error');
+      return {
+        content: [{ type: "text", text: `Error: ${errorMessage}` }],
+        isError: true,
+      };
     }
   }
 
@@ -61,10 +218,19 @@ export class Context {
 
   async close() {
     try {
-      await this.stagehand.close();
-      log('Stagehand context closed successfully', 'info');
+      // Close all Stagehand instances
+      for (const [sessionId, stagehand] of this.stagehands.entries()) {
+        try {
+          await stagehand.close();
+          log(`Closed Stagehand for session ${sessionId}`, 'info');
+        } catch (error) {
+          log(`Error closing Stagehand for session ${sessionId}: ${error}`, 'error');
+        }
+      }
+      this.stagehands.clear();
+      log('All Stagehand contexts closed successfully', 'info');
     } catch (error) {
-      log(`Error closing Stagehand context: ${error}`, 'error');
+      log(`Error closing Stagehand contexts: ${error}`, 'error');
     }
   }
 } 

stagehand/src/index.ts

@@ -10,17 +10,6 @@ import { zodToJsonSchema } from "zod-to-json-schema";
 import { Context } from "./context.js";
 import { TOOLS } from "./tools/tools.js";
 
-// Environment variables configuration
-const requiredEnvVars = {
-  BROWSERBASE_API_KEY: process.env.BROWSERBASE_API_KEY,
-  BROWSERBASE_PROJECT_ID: process.env.BROWSERBASE_PROJECT_ID,
-};
-
-// Validate required environment variables
-Object.entries(requiredEnvVars).forEach(([name, value]) => {
-  if (!value) throw new Error(`${name} environment variable is required`);
-});
-
 export async function createServer(config: Config): Promise<Server> {
   // Create the server
   const server = new Server(
@@ -31,6 +20,7 @@ export async function createServer(config: Config): Promise<Server> {
         tools: { list: true, call: true },
         prompts: { list: true, get: true },
         notifications: { resources: { list_changed: true } },
+        logging: {},
       },
     }
   ); 

stagehand/src/server.ts

@@ -15,10 +15,10 @@ export class ServerList {
   }
 
   async close(server: Server) {
+    await server.close();
     const index = this._servers.indexOf(server);
     if (index !== -1)
       this._servers.splice(index, 1);
-    await server.close();
   }
 
   async closeAll() {