Prevent stale Cloud agent-chat reconnects

agent-chat/src/connection.ts

@@ -0,0 +1,61 @@
+interface SessionConnection {
+  createSocket(): WebSocket;
+  onSocket(socket: WebSocket | null): void;
+  onOpen(): void;
+  onMessage(event: MessageEvent): void;
+}
+
+const RETRY_DELAY_MS = 800;
+
+/** Owns the session socket and its reconnect lifecycle for one mounted view. */
+export function openSessionConnection(callbacks: SessionConnection): () => void {
+  let closed = false;
+  let socket: WebSocket | null = null;
+  let retry: ReturnType<typeof setTimeout> | null = null;
+  const connect = () => {
+    if (closed) return;
+    retry = null;
+    let ws: WebSocket;
+    try {
+      ws = callbacks.createSocket();
+    } catch {
+      // The constructor throws on a bad URL or an opaque origin, and this call
+      // is the only thing that ever re-arms a retry. Letting it escape the
+      // timer leaves no socket and no pending timer, so the view would sit at
+      // ready with every send returning false and nothing to recover it.
+      retry = setTimeout(connect, RETRY_DELAY_MS);
+      return;
+    }
+    socket = ws;
+    callbacks.onSocket(ws);
+    const isCurrent = () => !closed && socket === ws;
+    ws.onopen = () => { if (isCurrent()) callbacks.onOpen(); };
+    ws.onmessage = (event) => { if (isCurrent()) callbacks.onMessage(event); };
+    ws.onclose = () => {
+      if (!isCurrent()) return;
+      socket = null;
+      callbacks.onSocket(null);
+      retry = setTimeout(connect, RETRY_DELAY_MS);
+    };
+  };
+  connect();
+  return () => {
+    closed = true;
+    if (retry !== null) clearTimeout(retry);
+    retry = null;
+    const ws = socket;
+    socket = null;
+    callbacks.onSocket(null);
+    if (ws) {
+      // Detach before closing. These handlers capture the whole session
+      // closure graph, and the browser keeps the socket alive until the close
+      // handshake finishes, which a server that never answers stretches to a
+      // TCP timeout. The gates would ignore the callbacks anyway, so the only
+      // thing holding them is the unmounted view's state.
+      ws.onopen = null;
+      ws.onmessage = null;
+      ws.onclose = null;
+      ws.close();
+    }
+  };
+}

agent-chat/src/session.ts

@@ -1,6 +1,7 @@
 // Client-side session state: one WebSocket, one session per page.
 import { useCallback, useEffect, useRef, useState } from "react";
 import { applyThemeVars } from "./theme";
+import { openSessionConnection } from "./connection";
 import type { HarnessRecommendation, HarnessCatalogs } from "../harness-contract";
 import { latestRouteStatus, normalizeRouteStatus, type RouteHealth, type RoutePhase, type RouteStatus } from "../route-status";
 
@@ -436,19 +437,18 @@ export function useSession(): SessionState {
   }, []);
 
   useEffect(() => {
-    let closed = false;
-    const connect = () => {
-      const ws = new WebSocket((location.protocol === "https:" ? "wss://" : "ws://") + location.host + appPath("/ws"));
-      wsRef.current = ws;
-      ws.onopen = () => {
+    const disconnect = openSessionConnection({
+      createSocket: () => new WebSocket((location.protocol === "https:" ? "wss://" : "ws://") + location.host + appPath("/ws")),
+      onSocket: (ws) => { wsRef.current = ws; },
+      onOpen: () => {
         const pending = pendingStartRef.current;
         if (sessionIdRef.current) sendRaw({ op: "subscribe", sessionId: sessionIdRef.current });
         else if (pending && !pending.failed) {
           sendRaw({ op: "start", requestId: pending.requestId, conversationId: pending.conversationId, provider: pending.provider, cwd: pending.cwd, prompt: pending.prompt, options: pending.options });
           armPendingStartTimeout();
         }
-      };
-      ws.onmessage = (e) => {
+      },
+      onMessage: (e) => {
         const msg = JSON.parse(e.data);
         switch (msg.kind) {
           case "hello": {
@@ -660,15 +660,12 @@ export function useSession(): SessionState {
             }
             break;
         }
-      };
-      ws.onclose = () => { if (!closed) setTimeout(connect, 800); };
-    };
-    connect();
+      },
+    });
     return () => {
-      closed = true;
+      disconnect();
       clearPendingStartTimeout();
       closeHandoffWindow();
-      wsRef.current?.close();
     };
   }, [armPendingStartTimeout, clearPendingStartTimeout, closeHandoffWindow, failPendingStart, sendRaw]);
 

agent-chat/test/connection.test.ts

@@ -0,0 +1,135 @@
+import { openSessionConnection } from "../src/connection";
+
+class FakeSocket {
+  onopen: (() => void) | null = null;
+  onmessage: ((event: MessageEvent) => void) | null = null;
+  onclose: (() => void) | null = null;
+  closed = false;
+  closes = 0;
+  close() { this.closed = true; this.closes++; this.onclose?.(); }
+  asWebSocket() { return this as unknown as WebSocket; }
+}
+
+const originalSetTimeout = globalThis.setTimeout;
+const originalClearTimeout = globalThis.clearTimeout;
+const timers = new Map<number, () => void>();
+const delays: number[] = [];
+let timerId = 0;
+const timerCount = () => timers.size;
+globalThis.setTimeout = ((callback: () => void, delay: number) => {
+  delays.push(delay);
+  const id = ++timerId;
+  timers.set(id, callback);
+  return id;
+}) as unknown as typeof setTimeout;
+globalThis.clearTimeout = ((id: number) => { timers.delete(id); }) as typeof clearTimeout;
+
+function flushTimers() {
+  const pending = [...timers.values()];
+  timers.clear();
+  for (const callback of pending) callback();
+}
+
+function client(createSocket?: () => WebSocket) {
+  const sockets: FakeSocket[] = [];
+  let current: WebSocket | null = null;
+  let opens = 0;
+  const messages: string[] = [];
+  const disconnect = openSessionConnection({
+    createSocket: createSocket ?? (() => {
+      const socket = new FakeSocket();
+      sockets.push(socket);
+      return socket.asWebSocket();
+    }),
+    onSocket: (socket) => { current = socket; },
+    onOpen: () => { opens++; },
+    onMessage: (event) => { messages.push(event.data); },
+  });
+  return { sockets, messages, disconnect, current: () => current, opens: () => opens };
+}
+
+try {
+  const disposed = client();
+  disposed.sockets[0].close();
+  disposed.disconnect();
+  if (timerCount() !== 0) throw new Error("disposed chat retained its reconnect timer");
+
+  const recovering = client();
+  const first = recovering.sockets[0];
+  const staleOpen = first.onopen!;
+  const staleMessage = first.onmessage!;
+  const staleClose = first.onclose!;
+  first.onopen?.();
+  first.onmessage?.({ data: "first history" } as MessageEvent);
+  first.close();
+  staleClose();
+  if (recovering.current() !== null) throw new Error("closed connection still owns the send socket");
+  if (timerCount() !== 1) throw new Error("one closed socket scheduled multiple reconnects");
+  flushTimers();
+  if (recovering.sockets.length !== 2) throw new Error("connection failed to open exactly one replacement socket");
+  const second = recovering.sockets[1];
+  second.onopen?.();
+  second.onmessage?.({ data: "replacement history" } as MessageEvent);
+  staleOpen();
+  staleMessage({ data: "stale history" } as MessageEvent);
+  staleClose();
+  if (recovering.opens() !== 2) throw new Error("stale socket resubscribed the current chat");
+  if (recovering.messages.join(",") !== "first history,replacement history") {
+    throw new Error("stale socket replaced the recovered chat's history");
+  }
+  if (timerCount() !== 0) throw new Error("stale socket scheduled another connection");
+  if (recovering.current() !== second.asWebSocket()) throw new Error("stale socket cleared the replacement connection");
+
+  // The retry may already be queued for execution when the view is disposed.
+  second.close();
+  const queuedRetry = [...timers.values()][0];
+  recovering.disconnect();
+  queuedRetry();
+  staleOpen();
+  staleMessage({ data: "late history" } as MessageEvent);
+  if (recovering.sockets.length !== 2 || recovering.opens() !== 2 || recovering.messages.length !== 2) {
+    throw new Error("disposed connection accepted a queued retry or stale callback");
+  }
+  if (timerCount() !== 0 || recovering.current() !== null) throw new Error("connection cleanup left owned state behind");
+
+  const mounted = client();
+  const only = mounted.sockets[0];
+  mounted.disconnect();
+  // Detached before closing: these handlers capture the whole session closure
+  // graph, and the socket outlives cleanup until the close handshake finishes.
+  if (only.onopen !== null || only.onmessage !== null || only.onclose !== null) {
+    throw new Error("cleanup left handlers attached to the disposed socket");
+  }
+  mounted.disconnect();
+  if (!only.closed || only.closes !== 1) throw new Error("cleanup did not close the socket exactly once");
+  if (timerCount() !== 0 || mounted.current() !== null) {
+    throw new Error("normal connection cleanup failed or scheduled a retry");
+  }
+
+  // A throwing constructor is the one failure that can strand the view: the
+  // retry timer is the only thing that calls connect(), so if the throw
+  // escapes it there is no socket and no timer left to recover from.
+  let attempts = 0;
+  const flakySockets: FakeSocket[] = [];
+  const flaky = client(() => {
+    attempts++;
+    if (attempts === 1) throw new Error("SecurityError");
+    const socket = new FakeSocket();
+    flakySockets.push(socket);
+    return socket.asWebSocket();
+  });
+  if (flaky.current() !== null) throw new Error("a failed connect published a socket");
+  if (timerCount() !== 1) throw new Error("a throwing constructor left no retry armed");
+  flushTimers();
+  if (attempts !== 2 || flakySockets.length !== 1) throw new Error("the retry after a throw did not reconnect");
+  if (flaky.current() !== flakySockets[0].asWebSocket()) throw new Error("the recovered socket was not published");
+  flaky.disconnect();
+
+  if (delays.length === 0 || delays.some((delay) => delay !== 800)) {
+    throw new Error(`reconnect delay changed: ${[...new Set(delays)].join(",")}`);
+  }
+  console.log("session connection lifecycle assertions passed");
+} finally {
+  globalThis.setTimeout = originalSetTimeout;
+  globalThis.clearTimeout = originalClearTimeout;
+}