fix(browser): recover stale CDP sessions

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

@@ -9,6 +9,7 @@
 import { bindDomains, type Domains, type Transport } from './generated.ts';
 
 type Pending = {
+  ws: WebSocket;
   resolve: (v: unknown) => void;
   reject: (e: unknown) => void;
 };
@@ -46,6 +47,9 @@ export class Session implements Transport {
   private nextId = 1;
   private pending = new Map<number, Pending>();
   private activeSessionId: string | undefined;
+  private activeTargetId: string | undefined;
+  private reattachPromise?: Promise<void>;
+  private enabledDomains = new Map<string, Map<string, unknown>>();
   private eventListeners: Array<(method: string, params: unknown, sessionId?: string) => void> = [];
   private callResultListeners: Array<(method: string, params: unknown, result: unknown) => void> = [];
 
@@ -79,6 +83,11 @@ export class Session implements Transport {
    * and we connect directly to the supplied endpoint.
    */
   async connect(opts: ConnectOptions = {}): Promise<void> {
+    // No-argument connect is an ensure-connected operation. Reopening the
+    // same configured endpoint would discard the active target session and
+    // make the next page command run against the browser-level socket.
+    if (!opts.wsUrl && !opts.profileDir && this.isConnected()) return;
+
     const timeoutMs = opts.timeoutMs ?? 5_000;
     if (opts.wsUrl || opts.profileDir) {
       const wsUrl = await resolveWsUrl(opts, timeoutMs);
@@ -124,15 +133,33 @@ export class Session implements Transport {
         else res();
       };
       const timer = setTimeout(() => finish(new Error(`timed out after ${timeoutMs}ms`)), timeoutMs);
-      ws.addEventListener('open', () => finish());
+      ws.addEventListener('open', () => {
+        if (done) {
+          try { ws.close(); } catch { /* ignore */ }
+          return;
+        }
+        const previous = this.ws;
+        this.ws = ws;
+        this.activeSessionId = undefined;
+        this.activeTargetId = undefined;
+        this.enabledDomains.clear();
+        finish();
+        if (previous && previous !== ws) {
+          try { previous.close(); } catch { /* ignore */ }
+        }
+      });
       ws.addEventListener('error', (e) => finish(new Error(`WS error: ${(e as any)?.message ?? 'connect failed (likely 403, permission not granted, or port closed)'}`)));
-      ws.addEventListener('message', (e) => this.onMessage(String(e.data)));
+      ws.addEventListener('message', (e) => this.onMessage(String(e.data), ws));
       ws.addEventListener('close', () => {
-        for (const [, p] of this.pending) p.reject(new Error('CDP socket closed'));
-        this.pending.clear();
+        this.rejectPending(ws, new Error('CDP socket closed'));
+        if (this.ws === ws) {
+          this.ws = undefined;
+          this.activeSessionId = undefined;
+          this.activeTargetId = undefined;
+          this.enabledDomains.clear();
+        }
         finish(new Error('WS closed before open (likely 403 or port closed)'));
       });
-      this.ws = ws;
     });
   }
 
@@ -151,12 +178,14 @@ 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;
+    this.activeTargetId = targetId;
     return r.sessionId;
   }
 
   /** Set the active sessionId directly (e.g. one you already attached). */
   setActiveSession(sessionId: string | undefined): void {
     this.activeSessionId = sessionId;
+    this.activeTargetId = undefined;
   }
 
   getActiveSession(): string | undefined {
@@ -206,35 +235,133 @@ export class Session implements Transport {
   }
 
   // Transport implementation. Called by the generated domain bindings.
-  _call(method: string, params: unknown = {}): Promise<unknown> {
-    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+  async _call(method: string, params: unknown = {}): Promise<unknown> {
+    const browserLevel = isBrowserLevel(method);
+    const sentSessionId = browserLevel ? undefined : this.activeSessionId;
+    const sentTargetId = browserLevel ? undefined : this.activeTargetId;
+    try {
+      return await this.send(method, params, sentSessionId);
+    } catch (error) {
+      if (!sentSessionId || !isMissingSessionError(error)) throw error;
+
+      // Chrome explicitly rejected the command before executing it, so this is
+      // safe to retry once. Socket drops are deliberately not retried: Chrome
+      // may have applied a click or submission before the response was lost.
+      if (this.activeSessionId === sentSessionId) {
+        await this.reattachPage(sentSessionId, sentTargetId);
+      }
+      else if (this.reattachPromise) await this.reattachPromise;
+      if (!this.activeSessionId || this.activeSessionId === sentSessionId) throw error;
+      return this.send(method, params, this.activeSessionId);
+    }
+  }
+
+  private send(method: string, params: unknown, sessionId?: string): Promise<unknown> {
+    const ws = this.ws;
+    if (!ws || ws.readyState !== WebSocket.OPEN) {
       return Promise.reject(new Error('Not connected. Call session.connect(...) first.'));
     }
+
     const id = this.nextId++;
     const msg: Record<string, unknown> = { id, method, params: params ?? {} };
-    if (this.activeSessionId && !isBrowserLevel(method)) {
-      msg.sessionId = this.activeSessionId;
-    }
+    if (sessionId) msg.sessionId = sessionId;
     return new Promise((resolve, reject) => {
       this.pending.set(id, {
+        ws,
         resolve: (v) => {
+          this.recordDomainState(method, params, sessionId);
           for (const fn of this.callResultListeners) {
             try { fn(method, params, v); } catch { /* ignore */ }
           }
           resolve(v);
         },
         reject,
       });
-      this.ws!.send(JSON.stringify(msg));
+      try {
+        ws.send(JSON.stringify(msg));
+      } catch (error) {
+        this.pending.delete(id);
+        reject(error);
+      }
     });
   }
 
-  private onMessage(raw: string): void {
+  private async reattachPage(staleSessionId: string, staleTargetId?: string): Promise<void> {
+    if (this.reattachPromise) return this.reattachPromise;
+
+    const attempt = this.attachPage(staleSessionId, staleTargetId);
+    this.reattachPromise = attempt;
+    try {
+      await attempt;
+    } finally {
+      if (this.reattachPromise === attempt) this.reattachPromise = undefined;
+    }
+  }
+
+  private async attachPage(staleSessionId: string, staleTargetId?: string): Promise<void> {
+    if (!staleTargetId) {
+      if (this.activeSessionId === staleSessionId) {
+        this.activeSessionId = undefined;
+        this.activeTargetId = undefined;
+        this.enabledDomains.delete(staleSessionId);
+      }
+      throw new Error(
+        'CDP target session was lost and its target is unknown; command was not retried on another page.',
+      );
+    }
+    const domainsToRestore = [...(this.enabledDomains.get(staleSessionId)?.entries() ?? [])];
+    const { targetInfos } = await this.domains.Target.getTargets({});
+    const pages = targetInfos as PageTarget[];
+    const exactTarget = pages.find(
+      target => target.type === 'page' && target.targetId === staleTargetId,
+    );
+    if (!exactTarget) {
+      if (this.activeSessionId === staleSessionId) {
+        this.activeSessionId = undefined;
+        this.activeTargetId = undefined;
+        this.enabledDomains.delete(staleSessionId);
+      }
+      throw new Error(
+        `CDP target ${staleTargetId} was closed; command was not retried on another page.`,
+      );
+    }
+    const sessionId = await this.use(exactTarget.targetId);
+    await Promise.all(
+      domainsToRestore.map(
+        ([method, params]) => this.send(method, params, sessionId),
+      ),
+    );
+    this.enabledDomains.delete(staleSessionId);
+  }
+
+  private recordDomainState(method: string, params: unknown, sessionId?: string): void {
+    if (!sessionId) return;
+    const match = /^([^.]+)\.(enable|disable)$/.exec(method);
+    if (!match) return;
+    const [, domain, command] = match;
+    if (!domain || !command) return;
+    const enabled = this.enabledDomains.get(sessionId) ?? new Map<string, unknown>();
+    if (command === 'enable') enabled.set(`${domain}.enable`, params);
+    else enabled.delete(`${domain}.enable`);
+    if (enabled.size > 0) this.enabledDomains.set(sessionId, enabled);
+    else this.enabledDomains.delete(sessionId);
+  }
+
+  private rejectPending(ws: WebSocket, error: Error): void {
+    for (const [id, pending] of this.pending) {
+      if (pending.ws !== ws) continue;
+      this.pending.delete(id);
+      pending.reject(error);
+    }
+  }
+
+  private onMessage(raw: string, ws: WebSocket): void {
+    if (ws !== this.ws) return;
     let m: any;
     try { m = JSON.parse(raw); } catch { return; }
     if (typeof m.id === 'number') {
       const p = this.pending.get(m.id);
-      if (!p) return;
+      if (!p || p.ws !== ws) return;
       this.pending.delete(m.id);
       if (m.error) p.reject(new CdpError(m.error.code, m.error.message, m.error.data));
       else p.resolve(m.result);
@@ -258,6 +385,12 @@ function isBrowserLevel(method: string): boolean {
   return method.startsWith('Browser.') || method.startsWith('Target.');
 }
 
+function isMissingSessionError(error: unknown): boolean {
+  return error instanceof CdpError
+    && error.code === -32001
+    && error.message.includes('Session with given id not found');
+}
+
 /**
  * Resolve a WebSocket URL for one of the explicit connect forms:
  *   { wsUrl }      — passthrough.
@@ -423,4 +556,3 @@ async function tryReadDevToolsActivePort(
     return undefined;
   }
 }
-