fix: make Browserbase sessions re-attachable by ID (STG-2446)

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`,
       );

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", () => {

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<typeof ActInputSchema>;
@@ -21,7 +27,7 @@ async function handleAct(
 ): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     try {
-      const stagehand = await context.getStagehand();
+      const stagehand = await context.getStagehand(params.sessionId);
 
       const result = await stagehand.act(params.action);
 

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<typeof ExtractInputSchema>;
@@ -21,7 +27,7 @@ async function handleExtract(
 ): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     try {
-      const stagehand = await context.getStagehand();
+      const stagehand = await context.getStagehand(params.sessionId);
 
       const extraction = params.instruction
         ? await stagehand.extract(params.instruction)

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<typeof NavigateInputSchema>;
@@ -21,7 +27,7 @@ async function handleNavigate(
 ): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     try {
-      const stagehand = await context.getStagehand();
+      const stagehand = await context.getStagehand(params.sessionId);
 
       const pages = stagehand.context.pages();
       const page = pages[0];

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<typeof ObserveInputSchema>;
@@ -21,7 +27,7 @@ async function handleObserve(
 ): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     try {
-      const stagehand = await context.getStagehand();
+      const stagehand = await context.getStagehand(params.sessionId);
 
       const observations = await stagehand.observe(params.instruction);
 

src/tools/session.ts

@@ -5,23 +5,37 @@ 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<typeof StartInputSchema>;
 
 const startSchema: ToolSchema<typeof StartInputSchema> = {
   name: "start",
   description: "Create or reuse a Browserbase session",
   inputSchema: StartInputSchema,
 };
 
-async function handleStart(context: Context): Promise<ToolResult> {
+async function handleStart(
+  context: Context,
+  params: StartInput,
+): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     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 ||