From f36b98db30e91e2443002253814d4c81c40f345e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Padez?= <andre.padez@gmail.com>
Date: Fri, 24 Jul 2026 14:43:42 +0000
Subject: [PATCH] chat: heartbeat the chat WebSocket so long/idle turns don't
 drop the connection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Bun closes a WS idle for idleTimeout (60s) and only resets its timer on frames
received from the client — but during a chat turn the client only receives, so a
long tool call or a gap between turns tripped the timeout (clean close 1000,
mid-stream cuts, reconnects). Now each chat connection is pinged every 25s in the
WS open handler (cleared in close); the client auto-pongs at the protocol level,
resetting Bun's timer. Covers every client (web, mobile) with one server change.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 src/servers/api/pi/websocket.ts | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/servers/api/pi/websocket.ts b/src/servers/api/pi/websocket.ts
index ecb322cf..9183c639 100644
--- a/src/servers/api/pi/websocket.ts
+++ b/src/servers/api/pi/websocket.ts
@@ -54,6 +54,13 @@ export const resolveBaseCwd = (email: string, role: string, cwd?: string) => {
 
 const wsToSessionMap = new WeakMap<any, string>();
 
+// Per-connection heartbeat. Bun closes a WS idle for `idleTimeout` (60s), and its timer only resets
+// on frames *received* from the client — but during a chat turn the client only receives. So we ping
+// each connection every 25s; the client auto-pongs at the protocol level, which resets Bun's timer
+// (and keeps reverse proxies happy). Genuinely dead sockets still time out (no pong).
+const pingTimers = new WeakMap<ServerWebSocket<WSData>, ReturnType<typeof setInterval>>();
+const PING_INTERVAL_MS = 25_000;
+
 function sendToClient(ws: ServerWebSocket<WSData> | null, msg: ServerMessage): void {
   if (ws?.readyState === 1) {
     ws.send(JSON.stringify(msg));
@@ -61,7 +68,14 @@ function sendToClient(ws: ServerWebSocket<WSData> | null, msg: ServerMessage): v
 }
 
 export async function open(ws: ServerWebSocket<WSData>): Promise<void> {
-  // logger.info('WebSocket connection opened', { email: ws.data.email });
+  const timer = setInterval(() => {
+    try {
+      ws.ping(); // client auto-pongs → resets Bun's idleTimeout
+    } catch {
+      /* socket already gone */
+    }
+  }, PING_INTERVAL_MS);
+  pingTimers.set(ws, timer);
 }
 
 export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void {
@@ -86,7 +100,11 @@ export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void
 }
 
 export function close(ws: ServerWebSocket<WSData>): void {
-  // logger.info('WebSocket connection closed', { email: ws.data.email });
+  const timer = pingTimers.get(ws);
+  if (timer) {
+    clearInterval(timer);
+    pingTimers.delete(ws);
+  }
 
   const sessionId = wsToSessionMap.get(ws);
   if (sessionId) {