diff --git a/src/sessionManager.ts b/src/sessionManager.ts index 3b1f538..a55abe0 100644 --- a/src/sessionManager.ts +++ b/src/sessionManager.ts @@ -375,6 +375,34 @@ export class SessionManager { const sessionObj = this.browsers.get(sessionId); if (!sessionObj) { + // The session isn't tracked in this process. This is expected when the + // MCP transport session was not preserved between tool calls (e.g. the + // request was routed to a different server replica, or the client did + // not reuse the Mcp-Session-Id). If the caller supplied an explicit + // Browserbase session ID, attempt to reconnect to that live session + // instead of failing, so state can be recovered across calls. + if (createIfMissing) { + process.stderr.write( + `[SessionManager] Session ${sessionId} not tracked locally, attempting to reconnect to Browserbase session.\n`, + ); + try { + return await this.createNewBrowserSession( + sessionId, + config, + sessionId, + ); + } catch (reconnectError) { + process.stderr.write( + `[SessionManager] WARN - Failed to reconnect to Browserbase session ${sessionId}: ${ + reconnectError instanceof Error + ? reconnectError.message + : String(reconnectError) + }\n`, + ); + return null; + } + } + process.stderr.write( `[SessionManager] WARN - Session not found in map: ${sessionId}\n`, ); diff --git a/src/tools/__tests__/tools.test.ts b/src/tools/__tests__/tools.test.ts index 9af1f4e..5be1637 100644 --- a/src/tools/__tests__/tools.test.ts +++ b/src/tools/__tests__/tools.test.ts @@ -98,6 +98,39 @@ describe("tool input schemas", () => { .success, ).toBe(true); }); + + it("start accepts an optional sessionId to resume a session", () => { + expect( + tool("start").schema.inputSchema.safeParse({ sessionId: "bb-session-123" }) + .success, + ).toBe(true); + }); + + it("navigate/act/observe/extract accept an optional sessionId", () => { + expect( + tool("navigate").schema.inputSchema.safeParse({ + url: "https://example.com", + sessionId: "bb-session-123", + }).success, + ).toBe(true); + expect( + tool("act").schema.inputSchema.safeParse({ + action: "click button", + sessionId: "bb-session-123", + }).success, + ).toBe(true); + expect( + tool("observe").schema.inputSchema.safeParse({ + instruction: "find links", + sessionId: "bb-session-123", + }).success, + ).toBe(true); + expect( + tool("extract").schema.inputSchema.safeParse({ + sessionId: "bb-session-123", + }).success, + ).toBe(true); + }); }); describe("tool capabilities", () => { diff --git a/src/tools/act.ts b/src/tools/act.ts index d616c7f..66757e4 100644 --- a/src/tools/act.ts +++ b/src/tools/act.ts @@ -5,6 +5,12 @@ import type { ToolActionResult } from "../types/types.js"; const ActInputSchema = z.object({ action: z.string().min(1), + sessionId: z + .string() + .optional() + .describe( + "The Browserbase session ID to act on. Defaults to the active session. Pass the sessionId returned by start to target a specific session.", + ), }); type ActInput = z.infer; @@ -21,7 +27,7 @@ async function handleAct( ): Promise { const action = async (): Promise => { try { - const stagehand = await context.getStagehand(); + const stagehand = await context.getStagehand(params.sessionId); const result = await stagehand.act(params.action); diff --git a/src/tools/extract.ts b/src/tools/extract.ts index e8b314e..9add4cc 100644 --- a/src/tools/extract.ts +++ b/src/tools/extract.ts @@ -5,6 +5,12 @@ import type { ToolActionResult } from "../types/types.js"; const ExtractInputSchema = z.object({ instruction: z.string().optional(), + sessionId: z + .string() + .optional() + .describe( + "The Browserbase session ID to act on. Defaults to the active session. Pass the sessionId returned by start to target a specific session.", + ), }); type ExtractInput = z.infer; @@ -21,7 +27,7 @@ async function handleExtract( ): Promise { const action = async (): Promise => { try { - const stagehand = await context.getStagehand(); + const stagehand = await context.getStagehand(params.sessionId); const extraction = params.instruction ? await stagehand.extract(params.instruction) diff --git a/src/tools/navigate.ts b/src/tools/navigate.ts index 08abe38..e49a302 100644 --- a/src/tools/navigate.ts +++ b/src/tools/navigate.ts @@ -5,6 +5,12 @@ import type { ToolActionResult } from "../types/types.js"; const NavigateInputSchema = z.object({ url: z.string().min(1), + sessionId: z + .string() + .optional() + .describe( + "The Browserbase session ID to act on. Defaults to the active session. Pass the sessionId returned by start to target a specific session.", + ), }); type NavigateInput = z.infer; @@ -21,7 +27,7 @@ async function handleNavigate( ): Promise { const action = async (): Promise => { try { - const stagehand = await context.getStagehand(); + const stagehand = await context.getStagehand(params.sessionId); const pages = stagehand.context.pages(); const page = pages[0]; diff --git a/src/tools/observe.ts b/src/tools/observe.ts index 2550ec3..b7d86a0 100644 --- a/src/tools/observe.ts +++ b/src/tools/observe.ts @@ -5,6 +5,12 @@ import type { ToolActionResult } from "../types/types.js"; const ObserveInputSchema = z.object({ instruction: z.string().min(1), + sessionId: z + .string() + .optional() + .describe( + "The Browserbase session ID to act on. Defaults to the active session. Pass the sessionId returned by start to target a specific session.", + ), }); type ObserveInput = z.infer; @@ -21,7 +27,7 @@ async function handleObserve( ): Promise { const action = async (): Promise => { try { - const stagehand = await context.getStagehand(); + const stagehand = await context.getStagehand(params.sessionId); const observations = await stagehand.observe(params.instruction); diff --git a/src/tools/session.ts b/src/tools/session.ts index 0f2f2e8..16db856 100644 --- a/src/tools/session.ts +++ b/src/tools/session.ts @@ -5,7 +5,16 @@ import type { ToolActionResult } from "../types/types.js"; import type { BrowserSession } from "../types/types.js"; // --- Tool: start --- -const StartInputSchema = z.object({}); +const StartInputSchema = z.object({ + sessionId: z + .string() + .optional() + .describe( + "An existing Browserbase session ID to resume. If omitted, a new session is created. Pass the sessionId returned by a previous start call to reconnect to the same browser.", + ), +}); + +type StartInput = z.infer; const startSchema: ToolSchema = { name: "start", @@ -13,15 +22,20 @@ const startSchema: ToolSchema = { inputSchema: StartInputSchema, }; -async function handleStart(context: Context): Promise { +async function handleStart( + context: Context, + params: StartInput, +): Promise { const action = async (): Promise => { try { const sessionManager = context.getSessionManager(); const config = context.config; - const targetSessionId = sessionManager.getDefaultSessionId(); + const targetSessionId = + params.sessionId ?? sessionManager.getDefaultSessionId(); - const session: BrowserSession = - await sessionManager.ensureDefaultSessionInternal(config); + const session: BrowserSession | null = params.sessionId + ? await sessionManager.getSession(params.sessionId, config, true) + : await sessionManager.ensureDefaultSessionInternal(config); if ( !session ||