diff --git a/app/src/services/__tests__/socketService.events.test.ts b/app/src/services/__tests__/socketService.events.test.ts index b5d6aa513..07c126d78 100644 --- a/app/src/services/__tests__/socketService.events.test.ts +++ b/app/src/services/__tests__/socketService.events.test.ts @@ -10,7 +10,15 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; type EventHandlerMap = Record 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 }).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 }).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(); diff --git a/app/src/services/socketService.ts b/app/src/services/socketService.ts index 57cde7af1..f56f2919a 100644 --- a/app/src/services/socketService.ts +++ b/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', () => { diff --git a/src/core/socketio.rs b/src/core/socketio.rs index b6b8d6900..258c4628d 100644 --- a/src/core/socketio.rs +++ b/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:`). 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| 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 = 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); } } diff --git a/src/openhuman/channels/providers/web.rs b/src/openhuman/channels/providers/web.rs index 7ce450964..94334f52c 100644 --- a/src/openhuman/channels/providers/web.rs +++ b/src/openhuman/channels/providers/web.rs @@ -131,8 +131,17 @@ static BUDGET_ERROR_PATTERNS: Lazy> = 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 = 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); diff --git a/src/openhuman/channels/providers/web_tests.rs b/src/openhuman/channels/providers/web_tests.rs index b18ad6355..e803a4c52 100644 --- a/src/openhuman/channels/providers/web_tests.rs +++ b/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]