fix(chat): survive socket reconnects — thread-key session/cancel + thread-room stream (#2493)

Co-authored-by: sanil-23 <sanil@alphahuman.xyz>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

app/src/services/__tests__/socketService.events.test.ts

@@ -10,7 +10,15 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 type EventHandlerMap = Record<string, (...args: unknown[]) => void>;
 
 // All mocks must be hoisted to module scope.
-const storeMock = { dispatch: vi.fn() };
+type ThreadStateShape = {
+  thread: { selectedThreadId: string | null; activeThreadId: string | null };
+};
+const storeMock = {
+  dispatch: vi.fn(),
+  getState: vi.fn(
+    (): ThreadStateShape => ({ thread: { selectedThreadId: null, activeThreadId: null } })
+  ),
+};
 vi.mock('../../store', () => ({ store: storeMock }));
 
 const setBackendMock = vi.fn((x: unknown) => ({ type: 'connectivity/setBackend', payload: x }));
@@ -85,6 +93,9 @@ describe('socketService — socket event handler dispatches (lines 212, 230, 237
   beforeEach(() => {
     vi.resetModules();
     storeMock.dispatch.mockClear();
+    storeMock.getState.mockReturnValue({
+      thread: { selectedThreadId: null, activeThreadId: null },
+    });
     setBackendMock.mockClear();
     getCoreRpcUrlMock.mockReset();
   });
@@ -116,6 +127,47 @@ describe('socketService — socket event handler dispatches (lines 212, 230, 237
     expect(connectedCall).toBeDefined();
   });
 
+  it('re-subscribes to the active thread room on connect (thread:subscribe)', async () => {
+    const { handlers, mockSocket } = buildMockSocket();
+
+    vi.doMock('socket.io-client', () => ({ io: vi.fn(() => mockSocket) }));
+    getCoreRpcUrlMock.mockResolvedValue('http://127.0.0.1:7788/rpc');
+    storeMock.getState.mockReturnValue({
+      thread: { selectedThreadId: 'thread-xyz', activeThreadId: null },
+    });
+
+    const { socketService } = await import('../socketService');
+    socketService.connect('jwt-test-thread-sub');
+
+    await pollUntil(() => expect(handlers['connect']).toBeDefined());
+
+    handlers['connect']!();
+
+    expect((mockSocket as { emit: ReturnType<typeof vi.fn> }).emit).toHaveBeenCalledWith(
+      'thread:subscribe',
+      { thread_id: 'thread-xyz' }
+    );
+  });
+
+  it('does not emit thread:subscribe on connect when no active thread', async () => {
+    const { handlers, mockSocket } = buildMockSocket();
+
+    vi.doMock('socket.io-client', () => ({ io: vi.fn(() => mockSocket) }));
+    getCoreRpcUrlMock.mockResolvedValue('http://127.0.0.1:7788/rpc');
+    // beforeEach already sets thread ids to null.
+
+    const { socketService } = await import('../socketService');
+    socketService.connect('jwt-test-no-thread');
+
+    await pollUntil(() => expect(handlers['connect']).toBeDefined());
+
+    handlers['connect']!();
+
+    const emitMock = (mockSocket as { emit: ReturnType<typeof vi.fn> }).emit;
+    const threadSub = emitMock.mock.calls.find(([ev]) => ev === 'thread:subscribe');
+    expect(threadSub).toBeUndefined();
+  });
+
   it('dispatches setBackend(disconnected) with reason when socket emits "disconnect" (line 230)', async () => {
     const { handlers, mockSocket } = buildMockSocket();
 

app/src/services/socketService.ts

@@ -240,6 +240,17 @@ class SocketService {
       store.dispatch(setStatusForUser({ userId: uid, status: 'connected' }));
       store.dispatch(setSocketIdForUser({ userId: uid, socketId }));
       store.dispatch(setBackend({ value: 'connected' }));
+
+      // Re-join the active thread's room so an in-flight turn's stream survives
+      // this (re)connection. Chat events are delivered to both the client_id
+      // room and a per-thread room (see socketio.rs `emit_web_channel_event`);
+      // because a reconnect produces a NEW client_id, the new socket must
+      // re-subscribe to the thread room to keep receiving the stream.
+      const threadState = store.getState().thread;
+      const activeThreadId = threadState?.selectedThreadId ?? threadState?.activeThreadId;
+      if (activeThreadId) {
+        this.socket?.emit('thread:subscribe', { thread_id: activeThreadId });
+      }
     });
 
     this.socket.on('ready', () => {

src/core/socketio.rs

@@ -251,6 +251,11 @@ struct ChatCancelPayload {
     thread_id: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct ThreadSubscribePayload {
+    thread_id: String,
+}
+
 /// Attaches the Socket.IO layer to the Axum router and sets up event handlers.
 ///
 /// It configures:
@@ -437,6 +442,31 @@ pub fn attach_socketio() -> (socketioxide::layer::SocketIoLayer, SocketIo) {
                     .await;
                 },
             );
+
+            // Handler for subscribing this socket to a thread's room.
+            //
+            // Chat-stream events are delivered to BOTH the initiating client's
+            // own room AND a per-thread room (`thread:<id>`). After a socket
+            // reconnects it has a NEW client_id, so it would miss an in-flight
+            // turn's remaining stream (delivered to the OLD client_id room). The
+            // frontend emits this on connect/reconnect for the active thread, so
+            // the new socket re-joins the thread room and keeps receiving the
+            // stream. Membership is dropped automatically on disconnect.
+            socket.on(
+                "thread:subscribe",
+                |socket: SocketRef, Data(payload): Data<ThreadSubscribePayload>| async move {
+                    if !socket_is_authed(&socket) {
+                        drop_unauthed(&socket, "thread:subscribe from unauthenticated socket");
+                        return;
+                    }
+                    let thread_id = payload.thread_id.trim();
+                    if thread_id.is_empty() {
+                        return;
+                    }
+                    let room = format!("thread:{thread_id}");
+                    join_room_logged(&socket, &room, &socket.id.to_string());
+                },
+            );
         },
     );
 
@@ -704,13 +734,23 @@ fn join_room_logged(socket: &SocketRef, room: &str, client_id: &str) {
 }
 
 fn emit_web_channel_event(io: &SocketIo, event: WebChannelEvent) {
-    let room = event.client_id.clone();
     let name = event.event.clone();
+    // Deliver to the initiating client's own room AND the per-thread room. The
+    // thread room lets a socket that reconnected with a new client_id (after
+    // re-subscribing via `thread:subscribe`) keep receiving an in-flight turn's
+    // stream. socket.io de-duplicates a socket present in multiple target rooms,
+    // so a socket in both receives each event exactly once (no double-render).
+    // "system" broadcasts and events without a thread_id keep the legacy
+    // single-room behavior.
+    let mut rooms: Vec<String> = vec![event.client_id.clone()];
+    if event.client_id != "system" && !event.thread_id.is_empty() {
+        rooms.push(format!("thread:{}", event.thread_id));
+    }
     if let Ok(payload) = serde_json::to_value(event) {
         log::debug!(
-            "[socketio] send event={} room={} thread_id={} request_id={}",
+            "[socketio] send event={} rooms={:?} thread_id={} request_id={}",
             name,
-            room,
+            rooms,
             payload
                 .get("thread_id")
                 .and_then(|v| v.as_str())
@@ -720,7 +760,7 @@ fn emit_web_channel_event(io: &SocketIo, event: WebChannelEvent) {
                 .and_then(|v| v.as_str())
                 .unwrap_or_default()
         );
-        emit_room_with_aliases(io, &room, &name, &payload);
+        emit_rooms_with_aliases(io, &rooms, &name, &payload);
     }
 }
 
@@ -741,10 +781,17 @@ fn emit_with_aliases(socket: &SocketRef, name: &str, payload: &serde_json::Value
     }
 }
 
-fn emit_room_with_aliases(io: &SocketIo, room: &str, name: &str, payload: &serde_json::Value) {
-    let _ = io.to(room.to_string()).emit(name, payload);
+fn emit_rooms_with_aliases(
+    io: &SocketIo,
+    rooms: &[String],
+    name: &str,
+    payload: &serde_json::Value,
+) {
+    // Emitting to multiple rooms in a single call delivers each event once per
+    // socket, even if a socket belongs to more than one of the target rooms.
+    let _ = io.to(rooms.to_vec()).emit(name, payload);
     if let Some(alias) = event_alias(name) {
-        let _ = io.to(room.to_string()).emit(alias, payload);
+        let _ = io.to(rooms.to_vec()).emit(alias, payload);
     }
 }
 

src/openhuman/channels/providers/web.rs

@@ -131,8 +131,17 @@ static BUDGET_ERROR_PATTERNS: Lazy<Vec<Regex>> = Lazy::new(|| {
     ]
 });
 
-fn key_for(client_id: &str, thread_id: &str) -> String {
-    format!("{client_id}::{thread_id}")
+/// Key for the per-thread runtime maps (`THREAD_SESSIONS`, `IN_FLIGHT`).
+///
+/// Keyed by `thread_id` ALONE — the stable, persistent identity of a
+/// conversation — NOT by the Socket.IO `client_id`, which is regenerated on
+/// every reconnect. Keying these maps by `client_id` previously orphaned a
+/// thread's cached session (conversation amnesia) and its in-flight task handle
+/// (Cancel became a no-op) whenever the socket reconnected with a new id. Event
+/// delivery still routes by `client_id` (the live socket); only the
+/// thread-owned runtime state keys off `thread_id`.
+fn key_for(thread_id: &str) -> String {
+    thread_id.to_string()
 }
 
 fn event_session_id_for(client_id: &str, thread_id: &str) -> String {
@@ -507,7 +516,7 @@ pub async fn start_chat(
         return Err(prompt_guard_user_message(prompt_decision.action).to_string());
     }
 
-    let map_key = key_for(&client_id, &thread_id);
+    let map_key = key_for(&thread_id);
 
     {
         let mut in_flight = IN_FLIGHT.lock().await;
@@ -724,7 +733,7 @@ pub async fn cancel_chat(client_id: &str, thread_id: &str) -> Result<Option<Stri
         return Err("thread_id is required".to_string());
     }
 
-    let map_key = key_for(client_id, thread_id);
+    let map_key = key_for(thread_id);
     let mut removed_request_id: Option<String> = None;
 
     {
@@ -792,7 +801,7 @@ async fn run_chat_task(
     let config = config_rpc::load_config_with_timeout().await?;
     let (_profiles_state, profile) =
         AgentProfileStore::new(config.workspace_dir.clone()).resolve(profile_id.as_deref())?;
-    let map_key = key_for(client_id, thread_id);
+    let map_key = key_for(thread_id);
     let model_override = normalize_model_override(profile.model_override.clone())
         .or_else(|| normalize_model_override(model_override));
     let temperature = profile.temperature.or(temperature);

src/openhuman/channels/providers/web_tests.rs

@@ -425,9 +425,16 @@ fn unknown_schema_returns_unknown_fallback() {
 // ── Helpers ───────────────────────────────────────────────────
 
 #[test]
-fn key_for_combines_client_id_and_thread_id() {
-    assert_eq!(key_for("c1", "t1"), "c1::t1");
-    assert_eq!(key_for("", ""), "::");
+fn key_for_is_thread_scoped_not_client_scoped() {
+    // Runtime maps (THREAD_SESSIONS, IN_FLIGHT) key by thread_id ALONE, so the
+    // key is stable across socket reconnects (which regenerate client_id).
+    // Regression guard for the conversation-amnesia / dead-Cancel bug, where a
+    // reconnect under a new client_id orphaned the thread's session + in-flight
+    // handle.
+    assert_eq!(key_for("thread-abc"), "thread-abc");
+    assert_eq!(key_for(""), "");
+    // The same thread resolves to the same key no matter which socket asks.
+    assert_eq!(key_for("thread-xyz"), key_for("thread-xyz"));
 }
 
 #[test]