From c120651ce609163e13839913e1148fe474a91d11 Mon Sep 17 00:00:00 2001 From: teamleaderleo Date: Wed, 30 Sep 2026 05:52:36 -0400 Subject: [PATCH 1/3] test: reproduce session reconnect after chat disposal --- agent-chat/src/connection.ts | 25 +++++++ agent-chat/src/session.ts | 23 +++--- agent-chat/test/connection.test.ts | 109 +++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+), 13 deletions(-) create mode 100644 agent-chat/src/connection.ts create mode 100644 agent-chat/test/connection.test.ts diff --git a/agent-chat/src/connection.ts b/agent-chat/src/connection.ts new file mode 100644 index 000000000000..4d0617c9a2ab --- /dev/null +++ b/agent-chat/src/connection.ts @@ -0,0 +1,25 @@ +interface SessionConnection { + createSocket(): WebSocket; + onSocket(socket: WebSocket | null): void; + onOpen(): void; + onMessage(event: MessageEvent): void; +} + +/** 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; + const connect = () => { + const ws = callbacks.createSocket(); + socket = ws; + callbacks.onSocket(ws); + ws.onopen = callbacks.onOpen; + ws.onmessage = callbacks.onMessage; + ws.onclose = () => { if (!closed) setTimeout(connect, 800); }; + }; + connect(); + return () => { + closed = true; + socket?.close(); + }; +} diff --git a/agent-chat/src/session.ts b/agent-chat/src/session.ts index 6a28dce5ed30..e9972679d13d 100644 --- a/agent-chat/src/session.ts +++ b/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]); diff --git a/agent-chat/test/connection.test.ts b/agent-chat/test/connection.test.ts new file mode 100644 index 000000000000..ef1075d2dea2 --- /dev/null +++ b/agent-chat/test/connection.test.ts @@ -0,0 +1,109 @@ +import { openSessionConnection } from "../src/connection"; + +class FakeSocket { + onopen: (() => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onclose: (() => void) | null = null; + closed = false; + close() { this.closed = true; this.onclose?.(); } + asWebSocket() { return this as unknown as WebSocket; } +} + +const originalSetTimeout = globalThis.setTimeout; +const originalClearTimeout = globalThis.clearTimeout; +const timers = new Map void>(); +let timerId = 0; +globalThis.setTimeout = ((callback: () => void, delay: number) => { + if (delay !== 800) throw new Error(`unexpected reconnect delay: ${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() { + const sockets: FakeSocket[] = []; + let current: WebSocket | null = null; + let opens = 0; + const messages: string[] = []; + const disconnect = openSessionConnection({ + 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(); + const leakedTimers = timers.size; + // Drain even a leaked retry to demonstrate that it must not open a socket. + flushTimers(); + if (disposed.sockets.length !== 1) { + throw new Error("disposed chat reopened a WebSocket from a pending reconnect"); + } + if (leakedTimers !== 0) throw new Error("disposed chat retained its reconnect timer"); + if (disposed.current() !== null) throw new Error("disposed chat retained a sendable socket reference"); + + 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 (timers.size !== 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 (timers.size !== 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 (timers.size !== 0 || recovering.current() !== null) throw new Error("connection cleanup left owned state behind"); + + const mounted = client(); + mounted.disconnect(); + mounted.disconnect(); + if (!mounted.sockets[0].closed || timers.size !== 0 || mounted.current() !== null) { + throw new Error("normal connection cleanup failed or scheduled a retry"); + } + console.log("session connection lifecycle assertions passed"); +} finally { + globalThis.setTimeout = originalSetTimeout; + globalThis.clearTimeout = originalClearTimeout; +} From 0a3f36bca4a28a37fe207c67eae7ca4bf9045097 Mon Sep 17 00:00:00 2001 From: teamleaderleo Date: Wed, 30 Sep 2026 05:53:41 -0400 Subject: [PATCH 2/3] fix: prevent stale Cloud chat reconnects --- agent-chat/src/connection.ts | 22 ++++++++++++++++++---- agent-chat/test/connection.test.ts | 11 ++++++----- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/agent-chat/src/connection.ts b/agent-chat/src/connection.ts index 4d0617c9a2ab..fe34d66a777f 100644 --- a/agent-chat/src/connection.ts +++ b/agent-chat/src/connection.ts @@ -9,17 +9,31 @@ interface SessionConnection { export function openSessionConnection(callbacks: SessionConnection): () => void { let closed = false; let socket: WebSocket | null = null; + let retry: ReturnType | null = null; const connect = () => { + if (closed) return; + retry = null; const ws = callbacks.createSocket(); socket = ws; callbacks.onSocket(ws); - ws.onopen = callbacks.onOpen; - ws.onmessage = callbacks.onMessage; - ws.onclose = () => { if (!closed) setTimeout(connect, 800); }; + 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, 800); + }; }; connect(); return () => { closed = true; - socket?.close(); + if (retry !== null) clearTimeout(retry); + retry = null; + const ws = socket; + socket = null; + callbacks.onSocket(null); + ws?.close(); }; } diff --git a/agent-chat/test/connection.test.ts b/agent-chat/test/connection.test.ts index ef1075d2dea2..80f30f077953 100644 --- a/agent-chat/test/connection.test.ts +++ b/agent-chat/test/connection.test.ts @@ -13,6 +13,7 @@ const originalSetTimeout = globalThis.setTimeout; const originalClearTimeout = globalThis.clearTimeout; const timers = new Map void>(); let timerId = 0; +const timerCount = () => timers.size; globalThis.setTimeout = ((callback: () => void, delay: number) => { if (delay !== 800) throw new Error(`unexpected reconnect delay: ${delay}`); const id = ++timerId; @@ -49,7 +50,7 @@ try { const disposed = client(); disposed.sockets[0].close(); disposed.disconnect(); - const leakedTimers = timers.size; + const leakedTimers = timerCount(); // Drain even a leaked retry to demonstrate that it must not open a socket. flushTimers(); if (disposed.sockets.length !== 1) { @@ -68,7 +69,7 @@ try { first.close(); staleClose(); if (recovering.current() !== null) throw new Error("closed connection still owns the send socket"); - if (timers.size !== 1) throw new Error("one closed socket scheduled multiple reconnects"); + 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]; @@ -81,7 +82,7 @@ try { if (recovering.messages.join(",") !== "first history,replacement history") { throw new Error("stale socket replaced the recovered chat's history"); } - if (timers.size !== 0) throw new Error("stale socket scheduled another connection"); + 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. @@ -94,12 +95,12 @@ try { 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 (timers.size !== 0 || recovering.current() !== null) throw new Error("connection cleanup left owned state behind"); + if (timerCount() !== 0 || recovering.current() !== null) throw new Error("connection cleanup left owned state behind"); const mounted = client(); mounted.disconnect(); mounted.disconnect(); - if (!mounted.sockets[0].closed || timers.size !== 0 || mounted.current() !== null) { + if (!mounted.sockets[0].closed || timerCount() !== 0 || mounted.current() !== null) { throw new Error("normal connection cleanup failed or scheduled a retry"); } console.log("session connection lifecycle assertions passed"); From fe108531097ba41e83df3173d821e56fa574b337 Mon Sep 17 00:00:00 2001 From: Leo Li Date: Wed, 30 Sep 2026 03:33:00 -0700 Subject: [PATCH 3/3] fix(agent-chat): survive a throwing socket constructor, detach on cleanup A review of this branch found two ways the connection helper can strand or retain a view, both of which this now covers. A throw from `createSocket()` was unhandled. The retry timer is the only thing that calls `connect()`, and `connect()` clears `retry` before constructing, so a throw inside the timer left no socket and no pending timer: nothing remained to reconnect. `new WebSocket()` throws on a bad URL and on an opaque origin, and agent-chat renders inside a workspace tab, so the view would sit at `ready` with every send silently returning false and no way back. The throw is caught and re-arms the retry, which is the behaviour the PR already claimed. Cleanup closed the socket without detaching its handlers. They capture the whole `useSession` closure graph, every setState and every ref, and the browser keeps the socket alive until the close handshake finishes, which a server that never answers stretches to a TCP timeout. The identity gates would ignore the callbacks anyway, so nothing was holding that state except the disposed socket. They are nulled before `close()` now. Test changes, all confirmed to fail against the previous source: - The disposed socket's handlers are null after cleanup, and it is closed exactly once across two `disconnect()` calls. - A constructor that throws once publishes no socket, arms a retry, and the retry connects. - The fake timer records delays instead of throwing on an unexpected one, and the 800 ms constant is asserted once at the end. Whoever adds backoff should get a named assertion, not a thrown string from a timer factory. - Dropped two assertions from the first block. `disposed.sockets.length` could not fail there, because `disconnect()` had already cleared the timer, so there was nothing for `flushTimers()` to drain and the comment claiming otherwise was wrong. `disposed.current()` could not fail either, since the close had already pushed null through `onSocket`. Both properties are genuinely covered by the queued-retry and mounted blocks. Co-Authored-By: Claude Opus 5.5 (1M context) --- agent-chat/src/connection.ts | 28 ++++++++++++++-- agent-chat/test/connection.test.ts | 53 ++++++++++++++++++++++-------- 2 files changed, 64 insertions(+), 17 deletions(-) diff --git a/agent-chat/src/connection.ts b/agent-chat/src/connection.ts index fe34d66a777f..5cb44716ef58 100644 --- a/agent-chat/src/connection.ts +++ b/agent-chat/src/connection.ts @@ -5,6 +5,8 @@ interface SessionConnection { 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; @@ -13,7 +15,17 @@ export function openSessionConnection(callbacks: SessionConnection): () => void const connect = () => { if (closed) return; retry = null; - const ws = callbacks.createSocket(); + 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; @@ -23,7 +35,7 @@ export function openSessionConnection(callbacks: SessionConnection): () => void if (!isCurrent()) return; socket = null; callbacks.onSocket(null); - retry = setTimeout(connect, 800); + retry = setTimeout(connect, RETRY_DELAY_MS); }; }; connect(); @@ -34,6 +46,16 @@ export function openSessionConnection(callbacks: SessionConnection): () => void const ws = socket; socket = null; callbacks.onSocket(null); - ws?.close(); + 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(); + } }; } diff --git a/agent-chat/test/connection.test.ts b/agent-chat/test/connection.test.ts index 80f30f077953..ef9799f33894 100644 --- a/agent-chat/test/connection.test.ts +++ b/agent-chat/test/connection.test.ts @@ -5,17 +5,19 @@ class FakeSocket { onmessage: ((event: MessageEvent) => void) | null = null; onclose: (() => void) | null = null; closed = false; - close() { this.closed = true; this.onclose?.(); } + 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 void>(); +const delays: number[] = []; let timerId = 0; const timerCount = () => timers.size; globalThis.setTimeout = ((callback: () => void, delay: number) => { - if (delay !== 800) throw new Error(`unexpected reconnect delay: ${delay}`); + delays.push(delay); const id = ++timerId; timers.set(id, callback); return id; @@ -28,17 +30,17 @@ function flushTimers() { for (const callback of pending) callback(); } -function client() { +function client(createSocket?: () => WebSocket) { const sockets: FakeSocket[] = []; let current: WebSocket | null = null; let opens = 0; const messages: string[] = []; const disconnect = openSessionConnection({ - createSocket: () => { + createSocket: createSocket ?? (() => { const socket = new FakeSocket(); sockets.push(socket); return socket.asWebSocket(); - }, + }), onSocket: (socket) => { current = socket; }, onOpen: () => { opens++; }, onMessage: (event) => { messages.push(event.data); }, @@ -50,14 +52,7 @@ try { const disposed = client(); disposed.sockets[0].close(); disposed.disconnect(); - const leakedTimers = timerCount(); - // Drain even a leaked retry to demonstrate that it must not open a socket. - flushTimers(); - if (disposed.sockets.length !== 1) { - throw new Error("disposed chat reopened a WebSocket from a pending reconnect"); - } - if (leakedTimers !== 0) throw new Error("disposed chat retained its reconnect timer"); - if (disposed.current() !== null) throw new Error("disposed chat retained a sendable socket reference"); + if (timerCount() !== 0) throw new Error("disposed chat retained its reconnect timer"); const recovering = client(); const first = recovering.sockets[0]; @@ -98,11 +93,41 @@ try { 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 (!mounted.sockets[0].closed || timerCount() !== 0 || mounted.current() !== null) { + 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;