From f2986bf86365dd45ad01c82e4e88b4b6454017d5 Mon Sep 17 00:00:00 2001 From: Gaurav Bhola Date: Thu, 2 Apr 2026 19:07:00 -0700 Subject: [PATCH] =?UTF-8?q?fix(web):=20address=20bugbot=20review=20?= =?UTF-8?q?=E2=80=94=20heartbeat,=20shutdown,=20runtime=20config,=20cleanu?= =?UTF-8?q?p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Use native WS ping frames for heartbeat so idle browser clients are not incorrectly disconnected (browser auto-responds to native ping with pong) - Pass runtime config to buildMuxWsUrl() so dynamic port/proxy path set via TERMINAL_WS_PATH env var is actually used (was fetched but never consumed) - Delay initial WS connect until runtime config fetch resolves, preventing race condition on first load - shutdown() now terminates all mux clients and closes the WSS, preventing orphaned PTY processes on restart - ws 'error' handler now unsubscribes terminal callbacks alongside session subscription to prevent leaks in edge cases where 'close' does not follow Co-Authored-By: Claude Opus 4.6 --- packages/web/server/direct-terminal-ws.ts | 8 +++ packages/web/server/mux-websocket.ts | 20 +++++-- .../hooks/__tests__/useSessionEvents.test.ts | 4 +- packages/web/src/providers/MuxProvider.tsx | 58 +++++++++---------- 4 files changed, 53 insertions(+), 37 deletions(-) diff --git a/packages/web/server/direct-terminal-ws.ts b/packages/web/server/direct-terminal-ws.ts index 75b2f5c2f..c9276c055 100644 --- a/packages/web/server/direct-terminal-ws.ts +++ b/packages/web/server/direct-terminal-ws.ts @@ -52,6 +52,14 @@ export function createDirectTerminalServer(tmuxPath?: string): DirectTerminalSer }); function shutdown() { + // Terminate all connected mux clients — this triggers their 'close' events + // which unsubscribe terminal callbacks and kill PTY processes. + if (muxWss) { + for (const client of muxWss.clients) { + client.terminate(); + } + muxWss.close(); + } server.close(); } diff --git a/packages/web/server/mux-websocket.ts b/packages/web/server/mux-websocket.ts index dc73efece..09ea8c633 100644 --- a/packages/web/server/mux-websocket.ts +++ b/packages/web/server/mux-websocket.ts @@ -424,26 +424,30 @@ export function createMuxWebSocket(tmuxPath?: string): WebSocketServer | null { let missedPongs = 0; const MAX_MISSED_PONGS = 3; - // Heartbeat: send pong every 15s + // Heartbeat: send native WebSocket ping every 15s. + // Browsers automatically respond to native pings with pong frames — + // no application-level code is needed on the client side. const heartbeatInterval = setInterval(() => { if (ws.readyState === WebSocket.OPEN) { - const msg: ServerMessage = { ch: "system", type: "pong" }; - ws.send(JSON.stringify(msg)); missedPongs += 1; - if (missedPongs >= MAX_MISSED_PONGS) { console.log("[MuxServer] Too many missed pongs, closing connection"); ws.close(1000, "Heartbeat timeout"); + return; } + ws.ping(); } }, 15_000); + // Native pong resets the missed counter + ws.on("pong", () => { + missedPongs = 0; + }); + /** * Handle incoming messages */ ws.on("message", (data) => { - // Reset missed pongs on any incoming message - missedPongs = 0; try { const msg = JSON.parse(data.toString("utf8")) as ClientMessage; @@ -555,6 +559,10 @@ export function createMuxWebSocket(tmuxPath?: string): WebSocketServer | null { clearInterval(heartbeatInterval); sessionUnsubscribe?.(); sessionUnsubscribe = null; + for (const unsub of subscriptions.values()) { + unsub(); + } + subscriptions.clear(); }); }); diff --git a/packages/web/src/hooks/__tests__/useSessionEvents.test.ts b/packages/web/src/hooks/__tests__/useSessionEvents.test.ts index 44228bd3c..3ca5f028a 100644 --- a/packages/web/src/hooks/__tests__/useSessionEvents.test.ts +++ b/packages/web/src/hooks/__tests__/useSessionEvents.test.ts @@ -674,7 +674,7 @@ describe("useSessionEvents", () => { it("seeds attention levels from initialAttentionLevels parameter", () => { const sessions = makeSessions(2); const initialLevels = { "session-0": "working" as const, "session-1": "respond" as const }; - const { result } = renderHook(() => useSessionEvents(sessions, null, undefined, initialLevels)); + const { result } = renderHook(() => useSessionEvents(sessions, null, undefined, undefined, initialLevels)); expect(result.current.sseAttentionLevels).toEqual({ "session-0": "working", "session-1": "respond", @@ -773,7 +773,7 @@ describe("useSessionEvents", () => { let currentLevels: Record | undefined = { "session-0": "respond" as const }; const { result, rerender } = renderHook(() => - useSessionEvents(currentSessions, null, undefined, currentLevels), + useSessionEvents(currentSessions, null, undefined, undefined, currentLevels), ); expect(result.current.sseAttentionLevels).toEqual({ "session-0": "respond" }); diff --git a/packages/web/src/providers/MuxProvider.tsx b/packages/web/src/providers/MuxProvider.tsx index 0ebfff48b..962ddedbd 100644 --- a/packages/web/src/providers/MuxProvider.tsx +++ b/packages/web/src/providers/MuxProvider.tsx @@ -48,25 +48,24 @@ function normalizePathValue(value: unknown): string | undefined { return trimmed; } -function buildMuxWsUrl(): string { +function buildMuxWsUrl(runtimeConfig: { directTerminalPort?: string; proxyWsPath?: string }): string { const loc = window.location; const protocol = loc.protocol === "https:" ? "wss:" : "ws:"; - // Check for proxy path env var - const proxyWsPath = process.env.NEXT_PUBLIC_TERMINAL_WS_PATH; + // Runtime proxy path takes priority (set by `ao start` via TERMINAL_WS_PATH env var) + const proxyWsPath = runtimeConfig.proxyWsPath ?? process.env.NEXT_PUBLIC_TERMINAL_WS_PATH; if (proxyWsPath) { - // Replace trailing /ws with /mux const basePath = proxyWsPath.replace(/\/ws\/?$/, ""); return `${protocol}//${loc.host}${basePath}/mux`; } - // Port-less or standard ports: use path-based + // Port-less or standard ports: use path-based routing (reverse proxy expected) if (loc.port === "" || loc.port === "443" || loc.port === "80") { return `${protocol}//${loc.hostname}/ao-terminal-mux`; } - // Direct port connection - const port = process.env.NEXT_PUBLIC_DIRECT_TERMINAL_PORT ?? "14801"; + // Direct port connection — prefer runtime-configured port, fall back to env/default + const port = runtimeConfig.directTerminalPort ?? process.env.NEXT_PUBLIC_DIRECT_TERMINAL_PORT ?? "14801"; return `${protocol}//${loc.hostname}:${port}/mux`; } @@ -85,25 +84,6 @@ export function MuxProvider({ children }: { children: ReactNode }) { const runtimeConfigRef = useRef<{ directTerminalPort?: string; proxyWsPath?: string }>({}); const sessionSubscribedRef = useRef(false); - // Fetch runtime config on mount - useEffect(() => { - const fetchRuntimeConfig = async () => { - try { - const res = await fetch("/api/runtime/terminal"); - if (res.ok) { - const data = (await res.json()) as RuntimeTerminalConfig; - runtimeConfigRef.current = { - directTerminalPort: normalizePortValue(data.directTerminalPort), - proxyWsPath: normalizePathValue(data.proxyWsPath), - }; - } - } catch { - // Ignore errors, fall back to defaults - } - }; - void fetchRuntimeConfig(); - }, []); - const connect = useCallback(() => { if (wsRef.current) { return; @@ -112,7 +92,7 @@ export function MuxProvider({ children }: { children: ReactNode }) { setStatus("connecting"); try { - const url = buildMuxWsUrl(); + const url = buildMuxWsUrl(runtimeConfigRef.current); console.log("[MuxProvider] Connecting to", url); const ws = new WebSocket(url); @@ -209,11 +189,31 @@ export function MuxProvider({ children }: { children: ReactNode }) { } }, []); - // Initial connection + // Fetch runtime config then connect. This ensures buildMuxWsUrl() has the + // server-configured port/path before the WebSocket is opened. useEffect(() => { - connect(); + let cancelled = false; + + const init = async () => { + try { + const res = await fetch("/api/runtime/terminal"); + if (res.ok) { + const data = (await res.json()) as RuntimeTerminalConfig; + runtimeConfigRef.current = { + directTerminalPort: normalizePortValue(data.directTerminalPort), + proxyWsPath: normalizePathValue(data.proxyWsPath), + }; + } + } catch { + // Ignore — fall back to env/default values + } + if (!cancelled) connect(); + }; + + void init(); return () => { + cancelled = true; if (reconnectTimer.current) { clearTimeout(reconnectTimer.current); }