pi-browser-harness 0.11.0 → 0.13.0

package/src/daemon/bridge.ts

@@ -1,10 +1,28 @@
-import WebSocket from "ws";
-import { discoverWsUrl } from "../cdp/discovery";
 import { isCdpRawMessage } from "../cdp/types";
 import type { CdpRawMessage } from "../cdp/types";
 import type { WireRequest, WireResponse, WireEvent } from "./protocol";
-import { CDP_CONNECT_TIMEOUT_MS, CDP_COMMAND_TIMEOUT_MS } from "./protocol";
+import { CDP_COMMAND_TIMEOUT_MS, DAEMON_RECONNECT_METHOD } from "./protocol";
 import { asString, isRecord } from "../util/guards";
+import { type ChromeConnectionDependencies, type ConnectionStatus, createChromeConnection } from "./connection";
+
+// The text a tool sees when the daemon has no Chrome connection to offer; it tells the model
+// (and through it the user) whether to wait, to look at Chrome, or to run /browser-setup.
+export const describeDisconnected = (status: ConnectionStatus): string => {
+  switch (status.state) {
+    case "open":
+      return "Chrome not connected";
+    case "connecting":
+      return `Chrome not connected: a connection attempt has been waiting ${Math.round(status.pendingForMs / 1000)}s. `
+        + "If Chrome is showing an \"Allow remote debugging\" prompt, accept it; no further prompts will be opened while this attempt is pending.";
+    case "cooldown":
+      return `Chrome not connected: the last attempt failed (${status.lastFailure}). `
+        + `The next automatic attempt is in ${Math.ceil(status.retryInMs / 1000)}s; run /browser-setup to retry now.`;
+    case "disconnected":
+      return status.lastFailure ? `Chrome not connected: ${status.lastFailure}` : "Chrome not connected";
+    case "stopped":
+      return "Chrome not connected: the daemon is shutting down";
+  }
+};
 
 export type SendToClient = (clientId: string, msg: WireResponse | WireEvent) => void;
 
@@ -13,14 +31,16 @@ export type EventHandler = (event: WireEvent, targetClientIds: string[]) => void
 export type CloseHandler = () => void;
 
 export type CdpBridge = {
-  start(): Promise<void>;
-  stop(): Promise<void>;
-  handleRequest(req: WireRequest, clientId: string, send: SendToClient): Promise<void>;
-  isAlive(): boolean;
-  getSessionOwner(sessionId: string): string | undefined;
-  removeClient(clientId: string): void;
-  onEvent(handler: EventHandler): void;
-  onClose(handler: CloseHandler): void;
+  readonly start: () => Promise<void>;
+  readonly stop: () => Promise<void>;
+  readonly handleRequest: (req: WireRequest, clientId: string, send: SendToClient) => Promise<void>;
+  readonly isAlive: () => boolean;
+  // Not used by the daemon itself; test/manual/bridge-test.ts asserts session ownership through it.
+  readonly getSessionOwner: (sessionId: string) => string | undefined;
+  readonly removeClient: (clientId: string) => void;
+  readonly onEvent: (handler: EventHandler) => void;
+  // Fires when an established Chrome connection drops; never for an attempt that failed to open.
+  readonly onClose: (handler: CloseHandler) => void;
 };
 
 export type InFlight = {
@@ -29,14 +49,16 @@ export type InFlight = {
   readonly send: SendToClient;
   readonly timer: ReturnType<typeof setTimeout>;
   readonly isAttach: boolean;
+  // The Chrome connection the request went out on, so that only its loss fails the request.
+  readonly generation: number;
 };
 
 // One map, keyed by the daemon-side id: the client that asked, the id it used, and everything needed to answer it.
 export type IdMultiplexer = {
-  allocate(entry: Omit<InFlight, "timer">, arm: (daemonId: number) => InFlight["timer"]): number;
-  take(daemonId: number): InFlight | undefined;
-  takeAll(): ReadonlyArray<InFlight>;
-  clearClient(clientId: string): void;
+  readonly allocate: (entry: Omit<InFlight, "timer">, arm: (daemonId: number) => InFlight["timer"]) => number;
+  readonly take: (daemonId: number) => InFlight | undefined;
+  // Removes every matching entry and cancels its timeout, so nothing taken can be answered later.
+  readonly takeWhere: (pred: (entry: InFlight) => boolean) => ReadonlyArray<InFlight>;
 };
 
 export const createIdMultiplexer = (): IdMultiplexer => {
@@ -56,29 +78,25 @@ export const createIdMultiplexer = (): IdMultiplexer => {
       clearTimeout(e.timer);
       return e;
     },
-    takeAll() {
-      const all = [...inFlight.values()];
-      inFlight.clear();
-      for (const e of all) clearTimeout(e.timer);
-      return all;
-    },
-    // A disconnected client must not be answered later: dropping its entries also cancels their timeouts.
-    clearClient(clientId) {
+    takeWhere(pred) {
+      const taken: InFlight[] = [];
       for (const [daemonId, e] of inFlight) {
-        if (e.clientId !== clientId) continue;
+        if (!pred(e)) continue;
         clearTimeout(e.timer);
         inFlight.delete(daemonId);
+        taken.push(e);
       }
+      return taken;
     },
   };
 };
 
 export type EventRouter = {
-  record(clientId: string, sessionId: string): void;
-  release(sessionId: string): void;
-  getOwner(sessionId: string): string | undefined;
-  removeClient(clientId: string): void;
-  route(event: WireEvent): string[];
+  readonly record: (clientId: string, sessionId: string) => void;
+  readonly release: (sessionId: string) => void;
+  readonly getOwner: (sessionId: string) => string | undefined;
+  readonly removeClient: (clientId: string) => void;
+  readonly route: (event: WireEvent) => string[];
 };
 
 export const createEventRouter = (): EventRouter => {
@@ -117,34 +135,22 @@ export const createEventRouter = (): EventRouter => {
   };
 };
 
-const CONNECT_WAIT_MS = 15_000;
+const errorResponse = (id: number, message: string): WireResponse => ({
+  type: "response",
+  id,
+  error: { code: -32000, message },
+});
 
-export const createCdpBridge = (): CdpBridge => {
-  let ws: WebSocket | null = null;
-  let wsUrl: string | null = null;
+const failAll = (entries: ReadonlyArray<InFlight>, message: string): void => {
+  for (const e of entries) e.send(e.clientId, errorResponse(e.localId, message));
+};
+
+export const createCdpBridge = (options: ChromeConnectionDependencies = {}): CdpBridge => {
   const mux = createIdMultiplexer();
   const router = createEventRouter();
   let eventHandler: EventHandler | null = null;
   let closeHandler: CloseHandler | null = null;
 
-  const connectedWaiters: Array<() => void> = [];
-  const signalConnected = (): void => {
-    for (const w of connectedWaiters.splice(0)) w();
-  };
-
-  const waitForConnection = (timeoutMs: number): Promise<void> =>
-    new Promise((resolve) => {
-      let settled = false;
-      const finish = (): void => {
-        if (settled) return;
-        settled = true;
-        clearTimeout(timer);
-        resolve();
-      };
-      const timer = setTimeout(finish, timeoutMs);
-      connectedWaiters.push(finish);
-    });
-
   const onChromeMessage = (raw: string): void => {
     let parsed: unknown;
     try { parsed = JSON.parse(raw); } catch { return; }
@@ -154,23 +160,20 @@ export const createCdpBridge = (): CdpBridge => {
     if (msg.id !== undefined) {
       const cb = mux.take(msg.id);
       if (!cb) return;
-      const localId = cb.localId;
 
       if (msg.error) {
         cb.send(cb.clientId, {
           type: "response",
-          id: localId,
+          id: cb.localId,
           error: { code: msg.error.code ?? -1, message: msg.error.message },
         });
-      } else {
-        if (cb.isAttach && msg.result) {
-          const sessionId = isRecord(msg.result) ? asString(msg.result["sessionId"]) : undefined;
-          if (sessionId) {
-            router.record(cb.clientId, sessionId);
-          }
-        }
-        cb.send(cb.clientId, { type: "response", id: localId, result: msg.result });
+        return;
       }
+      if (cb.isAttach && isRecord(msg.result)) {
+        const sessionId = asString(msg.result["sessionId"]);
+        if (sessionId !== undefined) router.record(cb.clientId, sessionId);
+      }
+      cb.send(cb.clientId, { type: "response", id: cb.localId, result: msg.result });
       return;
     }
 
@@ -195,112 +198,27 @@ export const createCdpBridge = (): CdpBridge => {
     }
   };
 
-  let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
-  let reconnectAttempt = 0;
-  let stopped = false;
-
-  const tryConnect = (): void => {
-    if (stopped) return;
-    if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
-
-    const attempt = async (): Promise<void> => {
-      if (ws && ws.readyState === WebSocket.OPEN) return;
-
-      const cachedUrl = wsUrl;
-      let url: string;
-      if (cachedUrl === null || cachedUrl === "" || reconnectAttempt > 0) {
-        const d = await discoverWsUrl();
-        if (!d.success) { scheduleRetry(); return; }
-        url = d.data;
-        wsUrl = url;
-      } else {
-        url = cachedUrl;
-      }
-
-      const settledPromise = new Promise<void>((settle) => {
-        let settled = false;
-        const settleOnce = (): void => {
-          if (settled) return;
-          settled = true;
-          settle();
-        };
-
-        let sock: WebSocket;
-        try {
-          sock = new WebSocket(url, { perMessageDeflate: false });
-        } catch {
-          scheduleRetry();
-          return;
-        }
-
-        const timer = setTimeout(() => {
-          sock.close();
-          settleOnce();
-        }, CDP_CONNECT_TIMEOUT_MS);
-
-        sock.on("open", () => {
-          clearTimeout(timer);
-          ws = sock;
-          reconnectAttempt = 0;
-          ws.send(JSON.stringify({ id: 0, method: "Target.setDiscoverTargets", params: { discover: true } }));
-          console.log("[pi-browser-daemon] Connected to Chrome ✓");
-          signalConnected();
-          settleOnce();
-        });
-
-        sock.on("message", (data: WebSocket.Data) => {
-          onChromeMessage(typeof data === "string" ? data : data.toString());
-        });
-
-        sock.on("error", () => {
-          clearTimeout(timer);
-          settleOnce();
-        });
-
-        sock.on("close", () => {
-          clearTimeout(timer);
-          ws = null;
-          for (const cb of mux.takeAll()) {
-            cb.send(cb.clientId, {
-              type: "response",
-              id: cb.localId,
-              error: { code: -32000, message: "Chrome disconnected" },
-            });
-          }
-          closeHandler?.();
-          scheduleRetry();
-        });
-      });
-
-      await settledPromise;
-      if (!ws || ws.readyState !== WebSocket.OPEN) {
-        scheduleRetry();
-      }
-    };
-
-    attempt().catch(() => scheduleRetry());
-  };
-
-  const scheduleRetry = (): void => {
-    if (stopped) return;
-    const delay = Math.min(1000 * Math.pow(2, Math.min(reconnectAttempt, 6)), 60_000);
-    reconnectAttempt++;
-    wsUrl = null;
-    reconnectTimer = setTimeout(tryConnect, delay);
-  };
+  const connection = createChromeConnection(options, {
+    onMessage: onChromeMessage,
+    onEnd: (generation, reason) => {
+      // Requests that went out on this socket can never be answered now; fail them at once
+      // rather than leaving them to the command timeout.
+      failAll(mux.takeWhere((e) => e.generation === generation), "Chrome disconnected");
+      if (reason === "dropped") closeHandler?.();
+    },
+  });
 
-  const start = async (): Promise<void> => {
-    stopped = false;
-    tryConnect();
+  // Kicks off the first attempt without waiting for it: the daemon serves its IPC socket
+  // meanwhile, and handleRequest() joins the attempt in flight.
+  const start = (): Promise<void> => {
+    void connection.start();
+    return Promise.resolve();
   };
 
-  const stop = async (): Promise<void> => {
-    stopped = true;
-    if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
-    mux.takeAll();
-    signalConnected();
-    if (ws) { try { ws.close(1000, "Shutdown"); } catch {} ws = null; }
-    wsUrl = null;
+  const stop = (): Promise<void> => {
+    connection.stop();
+    failAll(mux.takeWhere(() => true), "Chrome disconnected");
+    return Promise.resolve();
   };
 
   const handleRequest = async (
@@ -308,31 +226,38 @@ export const createCdpBridge = (): CdpBridge => {
     clientId: string,
     send: SendToClient,
   ): Promise<void> => {
-    // Every queued request awaits the same connect signal rather than each spinning its own poll loop.
-    if (!ws || ws.readyState !== WebSocket.OPEN) await waitForConnection(CONNECT_WAIT_MS);
-
-    if (!ws || ws.readyState !== WebSocket.OPEN) {
-      send(clientId, {
-        type: "response",
-        id: req.id,
-        error: { code: -32000, message: "Chrome not connected" },
-      });
+    // The one request the daemon answers itself: a human asked for a fresh dial.
+    if (req.method === DAEMON_RECONNECT_METHOD) {
+      const open = await connection.reconnectNow();
+      send(clientId, open
+        ? { type: "response", id: req.id, result: { connected: true } }
+        : errorResponse(req.id, describeDisconnected(connection.status())));
+      return;
+    }
+    // Connection attempts happen only in response to explicit browser work.
+    if (!connection.isOpen() && !(await connection.connect())) {
+      send(clientId, errorResponse(req.id, describeDisconnected(connection.status())));
       return;
     }
 
-    const isDetach = req.method === "Target.detachFromTarget";
-    if (isDetach && req.sessionId) router.release(req.sessionId);
+    // The session being detached travels in params (src/client.ts sends the envelope sessionId as null).
+    if (req.method === "Target.detachFromTarget") {
+      const sessionId = asString(req.params?.["sessionId"]);
+      if (sessionId !== undefined) router.release(sessionId);
+    }
 
     const daemonId = mux.allocate(
-      { clientId, localId: req.id, send, isAttach: req.method === "Target.attachToTarget" },
+      {
+        clientId,
+        localId: req.id,
+        send,
+        isAttach: req.method === "Target.attachToTarget",
+        generation: connection.generation(),
+      },
       (id) =>
         setTimeout(() => {
           mux.take(id);
-          send(clientId, {
-            type: "response",
-            id: req.id,
-            error: { code: -32000, message: `Timeout after ${CDP_COMMAND_TIMEOUT_MS}ms: ${req.method}` },
-          });
+          send(clientId, errorResponse(req.id, `Timeout after ${CDP_COMMAND_TIMEOUT_MS}ms: ${req.method}`));
         }, CDP_COMMAND_TIMEOUT_MS),
     );
 
@@ -343,15 +268,10 @@ export const createCdpBridge = (): CdpBridge => {
     };
     if (req.sessionId) payload["sessionId"] = req.sessionId;
 
-    try {
-      ws.send(JSON.stringify(payload));
-    } catch (e) {
+    const sent = connection.send(JSON.stringify(payload));
+    if (!sent.success) {
       mux.take(daemonId);
-      send(clientId, {
-        type: "response",
-        id: req.id,
-        error: { code: -32000, message: e instanceof Error ? e.message : String(e) },
-      });
+      send(clientId, errorResponse(req.id, sent.error));
     }
   };
 
@@ -359,9 +279,13 @@ export const createCdpBridge = (): CdpBridge => {
     start,
     stop,
     handleRequest,
-    isAlive: () => ws !== null && ws.readyState === WebSocket.OPEN,
+    isAlive: () => connection.isOpen(),
     getSessionOwner: (sid) => router.getOwner(sid),
-    removeClient: (cid) => { mux.clearClient(cid); router.removeClient(cid); },
+    removeClient: (cid) => {
+      // A gone client must never be answered: drop its requests silently and release its sessions.
+      mux.takeWhere((e) => e.clientId === cid);
+      router.removeClient(cid);
+    },
     onEvent: (h) => { eventHandler = h; },
     onClose: (h) => { closeHandler = h; },
   };