From 2fd7fc0e918fd466bfb12e5f6953c908ed1c03d9 Mon Sep 17 00:00:00 2001
From: ciregenz <ciregenz@berkeley.edu>
Date: Tue, 9 Jun 2026 21:57:55 -0700
Subject: [PATCH] [eric] frontend: keep an active agent's WS alive across a
 hop, reopening a still-responding chat skips the reconnect handshake that
 caused the lag

---
 .../src/app/pages/AgentChat/AgentChat.tsx     | 17 +++++++++--
 frontend/src/shared/ws/WebSocketManager.ts    | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/frontend/src/app/pages/AgentChat/AgentChat.tsx b/frontend/src/app/pages/AgentChat/AgentChat.tsx
index 18710845..94d7dc97 100644
--- a/frontend/src/app/pages/AgentChat/AgentChat.tsx
+++ b/frontend/src/app/pages/AgentChat/AgentChat.tsx
@@ -41,7 +41,7 @@ import {
 } from '@/shared/state/agentsSlice';
 import { store } from '@/shared/state/store';
 import { fetchModes } from '@/shared/state/modesSlice';
-import { createSessionWs } from '@/shared/ws/WebSocketManager';
+import { createSessionWs, acquireSessionWs, releaseSessionWs } from '@/shared/ws/WebSocketManager';
 import StreamingBubble from './bubbles/StreamingBubble';
 import MessageBubble from './bubbles/MessageBubble';
 import CompactionMarker from './bubbles/CompactionMarker';
@@ -207,6 +207,8 @@ const AgentChat: React.FC<AgentChatProps> = ({ sessionId: sessionIdProp, onClose
   const [model, setModel] = useState('sonnet');
 
   const wsRef = useRef<ReturnType<typeof createSessionWs> | null>(null);
+  // Current status for the WS-cleanup closure (effect deps can't include it).
+  const statusRef = useRef<string | undefined>(undefined);
   const initialContextApplied = useRef(false);
   const messageQueueRef = useRef<QueuedMessage[]>([]);
   const [queueLength, setQueueLength] = useState(0);
@@ -249,13 +251,20 @@ const AgentChat: React.FC<AgentChatProps> = ({ sessionId: sessionIdProp, onClose
         }
       }
       if (cancelled) return;
-      ws = createSessionWs(id);
+      // acquireSessionWs reuses a still-open socket kept alive from the last hop,
+      // so an active agent's stream resumes with no reconnect handshake. connect()
+      // is a no-op when the reused socket is already open.
+      ws = acquireSessionWs(id);
       ws.connect();
       wsRef.current = ws;
     })();
     return () => {
       cancelled = true;
-      if (ws) ws.disconnect();
+      if (ws) {
+        const st = statusRef.current;
+        const active = st === 'running' || st === 'waiting_approval';
+        releaseSessionWs(id, ws, active);
+      }
       wsRef.current = null;
     };
   }, [id, isDraft, dispatch]);
@@ -313,6 +322,8 @@ const AgentChat: React.FC<AgentChatProps> = ({ sessionId: sessionIdProp, onClose
     }
   }, [id, isDraft, mode, model, session?.system_prompt, session?.target_directory, dispatch]);
 
+  statusRef.current = session?.status;
+
   const agentBusy = awaitingResponse || (!isDraft && (session?.status === 'running' || session?.status === 'waiting_approval'));
 
   const prevStatusRef = useRef(session?.status);
diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts
index 1f4c8387..d18bfc87 100644
--- a/frontend/src/shared/ws/WebSocketManager.ts
+++ b/frontend/src/shared/ws/WebSocketManager.ts
@@ -812,6 +812,7 @@ class WebSocketManager {
   }
 
   stopAgent(sessionId: string) {
+    console.warn(`[ux-trace] STOP-SENT sid=${sessionId.slice(0, 8)}\n${new Error().stack}`);
     this.send('agent:stop', { session_id: sessionId });
   }
 
@@ -856,4 +857,33 @@ export function createSessionWs(sessionId: string): WebSocketManager {
   return new WebSocketManager(`${WS_BASE}/ws/agents/${sessionId}`, { sessionId });
 }
 
+// One backgrounded session socket, kept alive across a hop so an active agent's
+// stream doesn't pay a reconnect+resume handshake on reopen (the "Locking-in"
+// lag). Bounded to ONE: opening any other chat tears the previous one down, so
+// at most a single detached socket lingers, still pumping events into Redux.
+let _backgroundedSessionWs: { sessionId: string; ws: WebSocketManager } | null = null;
+
+export function acquireSessionWs(sessionId: string): WebSocketManager {
+  if (_backgroundedSessionWs?.sessionId === sessionId) {
+    const ws = _backgroundedSessionWs.ws;
+    _backgroundedSessionWs = null;
+    return ws;
+  }
+  return new WebSocketManager(`${WS_BASE}/ws/agents/${sessionId}`, { sessionId });
+}
+
+export function releaseSessionWs(sessionId: string, ws: WebSocketManager, keepAlive: boolean): void {
+  // Never keep more than one detached socket around.
+  if (_backgroundedSessionWs && _backgroundedSessionWs.sessionId !== sessionId) {
+    _backgroundedSessionWs.ws.disconnect();
+    _backgroundedSessionWs = null;
+  }
+  if (keepAlive) {
+    _backgroundedSessionWs = { sessionId, ws };
+  } else {
+    ws.disconnect();
+    if (_backgroundedSessionWs?.sessionId === sessionId) _backgroundedSessionWs = null;
+  }
+}
+
 export default WebSocketManager;