fix(browser): reconnect after tool timeout

Author: MagMueller

packages/bcode-browser/src/browser-execute.ts

@@ -49,6 +49,10 @@
 import fs from "fs/promises"
 import path from "path"
 import { Effect, Schema } from "effect"
+import {
+  getSessionRecoveryState,
+  type SessionRecoveryState,
+} from "./cdp/session"
 import { SessionStore } from "./session-store"
 import { Skills } from "./skills"
 
@@ -58,6 +62,7 @@ const MAX_TIMEOUT_OUTPUT_BYTES = 8 * 1024
 const TIMEOUT_OUTPUT_TRUNCATED = "[partial console output truncated; showing final bytes]\n"
 const v4Connections = new Map<string, Promise<void>>()
 const v4Bootstrapped = new Set<string>()
+const v4TimeoutRecovery = new Map<string, SessionRecoveryState>()
 
 // Tail-cap the captured output for the timeout error: last 8 KiB, snapped
 // forward to a UTF-8 sequence start so multibyte characters survive the cut.
@@ -274,6 +279,9 @@ export const make = Effect.fn("BrowserExecute.make")(function* (dataDir: string)
             // never evicted. A concurrent same-sessionID call would share the
             // retired object — acceptable for v1, opencode serializes tool
             // calls within an assistant message.
+            const recoveryState = getSessionRecoveryState(session)
+            if (recoveryState) v4TimeoutRecovery.set(ctx.sessionID, recoveryState)
+            else v4TimeoutRecovery.delete(ctx.sessionID)
             SessionStore.invalidate(ctx.sessionID, session, error)
             return Effect.fail(error)
           }),
@@ -286,27 +294,33 @@ export const make = Effect.fn("BrowserExecute.make")(function* (dataDir: string)
 
 async function ensureCloudConnected(sessionID: string, session: ReturnType<typeof SessionStore.get>) {
   if (!process.env.V4_RUN_ID || (!process.env.BU_CDP_WS && !process.env.BU_CDP_URL)) return
-  if (v4Bootstrapped.has(sessionID)) return
+  const recoveryState = v4TimeoutRecovery.get(sessionID)
+  if (v4Bootstrapped.has(sessionID) && !recoveryState) return
 
   const existing = v4Connections.get(sessionID)
   if (existing) return existing
 
   const connecting = (async () => {
-    if (!session.isConnected()) await session.connect()
+    if (!session.isConnected()) {
+      await session.connect(recoveryState ? { wsUrl: recoveryState.wsUrl } : {})
+    }
     if (session.getActiveSession()) return
-    const page = (await session.domains.Target.getTargets({})).targetInfos.find(
-      (target) => target.type === "page" && !target.url.startsWith("chrome://"),
-    )
+    const targets = (await session.domains.Target.getTargets({})).targetInfos
+    const page = targets.find((target) => target.targetId === recoveryState?.targetId)
+      ?? targets.find((target) => target.type === "page" && !target.url.startsWith("chrome://"))
+    if (recoveryState && !page) {
+      throw new Error("No page target available after browser_execute timeout")
+    }
     if (page) await session.use(page.targetId)
+    if (recoveryState) v4TimeoutRecovery.delete(sessionID)
   })()
   v4Connections.set(sessionID, connecting)
   try {
     await connecting
   } finally {
-    // One automatic attempt per logical BrowserCode session. A later disconnect
-    // (including after timeout replacement) must be surfaced:
-    // BU_CDP_WS is the browser selected at run start, not necessarily a newer
-    // browser the agent explicitly switched to during this run.
+    // The initial bootstrap happens once. Only a browser_execute timeout may
+    // trigger another automatic connection, using the exact endpoint that the
+    // timed-out Session was driving. Other disconnects remain explicit.
     v4Bootstrapped.add(sessionID)
     if (v4Connections.get(sessionID) === connecting) v4Connections.delete(sessionID)
   }

packages/bcode-browser/src/cdp/session.ts

@@ -25,6 +25,20 @@ export type ConnectOptions = {
   timeoutMs?: number;
 };
 
+export type SessionRecoveryState = {
+  wsUrl: string;
+  targetId?: string;
+};
+
+const recoveryStates = new WeakMap<Session, SessionRecoveryState>();
+
+export const getSessionRecoveryState = (
+  session: Session,
+): SessionRecoveryState | undefined => {
+  const state = recoveryStates.get(session);
+  return state ? { ...state } : undefined;
+};
+
 /** A Chromium-based browser detected as running on this machine. */
 export type DetectedBrowser = {
   /** Short label, e.g. 'Google Chrome', 'Brave', 'Comet'. */
@@ -126,6 +140,7 @@ export class Session implements Transport {
       const previousWs = this.ws;
       this.ws = ws;
       this.activeSessionId = undefined;
+      recoveryStates.delete(this);
       if (previousWs) {
         for (const [, p] of this.pending) p.reject(new Error('CDP connection replaced'));
         this.pending.clear();
@@ -145,6 +160,7 @@ export class Session implements Transport {
           finish(new Error('CDP connection superseded'));
           return;
         }
+        if (!this.invalidatedError) recoveryStates.set(this, { wsUrl });
         finish(this.invalidatedError);
       });
       ws.addEventListener('error', (e) => finish(new Error(`WS error: ${(e as any)?.message ?? 'connect failed (likely 403, permission not granted, or port closed)'}`)));
@@ -170,6 +186,7 @@ export class Session implements Transport {
   }
 
   close(): void {
+    recoveryStates.delete(this);
     this.ws?.close();
   }
 
@@ -180,7 +197,7 @@ export class Session implements Transport {
    * orphan keeps running and would otherwise share this object (and its
    * socket) with the next tool call, interleaving two authors on one
    * transport. Invalidation rejects all future `connect`/`_call` attempts
-   * and closes the socket (the close handler rejects in-flight calls);
+   * and closes the socket after rejecting in-flight calls;
    * `SessionStore.invalidate` removes the entry so the next call gets a
    * fresh Session.
    */
@@ -190,6 +207,8 @@ export class Session implements Transport {
     const ws = this.ws;
     this.ws = undefined;
     this.activeSessionId = undefined;
+    for (const [, p] of this.pending) p.reject(error);
+    this.pending.clear();
     try { ws?.close(); } catch { /* ignore */ }
   }
 
@@ -200,12 +219,16 @@ export class Session implements Transport {
   async use(targetId: string): Promise<string> {
     const r = await this._call('Target.attachToTarget', { targetId, flatten: true }) as { sessionId: string };
     this.activeSessionId = r.sessionId;
+    const state = recoveryStates.get(this);
+    if (state) recoveryStates.set(this, { ...state, targetId });
     return r.sessionId;
   }
 
   /** Set the active sessionId directly (e.g. one you already attached). */
   setActiveSession(sessionId: string | undefined): void {
     this.activeSessionId = sessionId;
+    const state = recoveryStates.get(this);
+    if (state) recoveryStates.set(this, { wsUrl: state.wsUrl });
   }
 
   getActiveSession(): string | undefined {

packages/bcode-browser/test/browser-auto-connect.test.ts

@@ -7,20 +7,22 @@ import { BrowserExecute } from "../src/browser-execute";
 import { SessionStore } from "../src/session-store";
 
 let connections = 0;
+let alternateConnections = 0;
 let closedConnections = 0;
 let attachedCalls = 0;
 let pageCallsWithSession = 0;
 let latestSocket: { close(): void } | undefined;
-const server = Bun.serve({
+const server = Bun.serve<{ path: string }>({
   port: 0,
   fetch(req, srv) {
-    return srv.upgrade(req)
+    return srv.upgrade(req, { data: { path: new URL(req.url).pathname } })
       ? undefined
       : new Response("upgrade required", { status: 426 });
   },
   websocket: {
     open(ws) {
       connections++;
+      if (ws.data.path === "/alternate") alternateConnections++;
       latestSocket = ws;
     },
     message(ws, message) {
@@ -77,6 +79,7 @@ afterAll(() => {
 });
 
 const wsUrl = `ws://127.0.0.1:${server.port}/`;
+const alternateWsUrl = `ws://127.0.0.1:${server.port}/alternate`;
 
 const withEnv = async <T>(
   vars: Record<string, string | undefined>,
@@ -276,8 +279,10 @@ test("an explicit browser switch retires the old socket and target attachment",
   );
 });
 
-test("a timeout replacement does not auto-attach the run-start browser again", async () => {
+test("a timeout replacement reconnects the same browser and target", async () => {
   connections = 0;
+  attachedCalls = 0;
+  pageCallsWithSession = 0;
   await withEnv(
     { V4_RUN_ID: "run-timeout", BU_CDP_WS: wsUrl, BU_CDP_URL: undefined },
     () =>
@@ -295,22 +300,77 @@ test("a timeout replacement does not auto-attach the run-start browser again", a
           ),
         ).rejects.toThrow("browser_execute timed out");
 
-        await expect(
-          Effect.runPromise(
-            impl.execute(
-              {
-                description: "Do not silently return to the run-start browser",
-                code: "return await session.Page.navigate({ url: 'https://sap.com' })",
-              },
-              { sessionID, workspaceDir },
-            ),
+        const recovered = await Effect.runPromise(
+          impl.execute(
+            {
+              description: "Continue on the same browser",
+              code: "return await session.Page.navigate({ url: 'https://sap.com' })",
+            },
+            { sessionID, workspaceDir },
           ),
-        ).rejects.toThrow("Not connected. Call session.connect(...) first.");
-        expect(connections).toBe(1);
+        );
+        expect(JSON.parse(recovered.result)).toEqual({});
+        expect(connections).toBe(2);
+        expect(attachedCalls).toBe(2);
+        expect(pageCallsWithSession).toBe(1);
       }),
   );
 });
 
+test("timeout recovery preserves an explicit browser switch", async () => {
+  connections = 0;
+  attachedCalls = 0;
+  alternateConnections = 0;
+  pageCallsWithSession = 0;
+  await withEnv(
+    {
+      V4_RUN_ID: "run-switched-timeout",
+      BU_CDP_WS: wsUrl,
+      BU_CDP_URL: undefined,
+    },
+    () =>
+      withBrowserExecute(
+        "switched-timeout",
+        async (impl, sessionID, workspaceDir) => {
+          const run = (code: string, timeout?: number) =>
+            Effect.runPromise(
+              impl.execute(
+                { description: "Drive switched browser", code, timeout },
+                { sessionID, workspaceDir },
+              ),
+            );
+
+          await run(
+            "return await session.Page.navigate({ url: 'https://sap.com' })",
+          );
+          await run(`
+          await session.connect({ wsUrl: ${JSON.stringify(alternateWsUrl)} })
+          const page = (await session.Target.getTargets({})).targetInfos[0]
+          await session.use(page.targetId)
+          return await session.Page.navigate({ url: "https://example.com" })
+        `);
+          await expect(
+            run("await new Promise(resolve => setTimeout(resolve, 100))", 10),
+          ).rejects.toThrow("browser_execute timed out");
+
+          expect(
+            JSON.parse(
+              (
+                await run(
+                  "return await session.Page.navigate({ url: 'https://example.com/continued' })",
+                )
+              ).result,
+            ),
+          ).toEqual({});
+          expect(connections).toBe(3);
+          expect(alternateConnections).toBe(2);
+          expect(attachedCalls).toBe(3);
+          expect(pageCallsWithSession).toBe(3);
+        },
+      ),
+  );
+});
+
 test("sessions without a provisioned endpoint still require connect", async () => {
   await withEnv(
     { V4_RUN_ID: undefined, BU_CDP_WS: undefined, BU_CDP_URL: undefined },