From fd5dc097a74145ed7feeff29b07d79b1988e2b00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Thu, 7 May 2026 18:33:48 +0800 Subject: [PATCH] fix: isolate desktop session streaming state Desktop chat sessions can stream and reconnect concurrently, so shared in-process buffers must not be keyed only by process lifetime. This change scopes streaming delta buffers, task-tool refresh bookkeeping, and CLI task mutations to the session that produced the event. It also ignores stale client socket closes after a newer socket has replaced the active connection for the same session. Constraint: Desktop users can keep multiple sessions and tabs active at the same time. Rejected: Serialize all desktop streaming through one active session | would hide the race instead of preserving multi-session behavior. Confidence: high Scope-risk: moderate Directive: Do not reintroduce process-global chat/task pending state without session keys and multi-session regression tests. Tested: cd desktop && bun run test -- src/stores/chatStore.test.ts src/stores/cliTaskStore.test.ts Tested: bun test src/server/__tests__/websocket-handler.test.ts Tested: cd desktop && bun run lint Tested: bun run check:desktop Tested: bun run check:server Tested: bun run check:coverage | changed-lines 100%, fails only existing agent-utils baseline Tested: bun run verify | 8 passed, 1 failed on existing agent-utils coverage baseline, 2 skipped Not-tested: Real desktop dual-window manual smoke. Related: https://github.com/NanmiCoder/cc-haha/issues/302 Related: https://github.com/NanmiCoder/cc-haha/issues/303 --- desktop/src/stores/chatStore.test.ts | 480 +++++++++++++++++- desktop/src/stores/chatStore.ts | 126 +++-- desktop/src/stores/cliTaskStore.test.ts | 74 +++ desktop/src/stores/cliTaskStore.ts | 37 +- .../__tests__/websocket-handler.test.ts | 55 ++ src/server/ws/handler.ts | 12 + 6 files changed, 721 insertions(+), 63 deletions(-) create mode 100644 src/server/__tests__/websocket-handler.test.ts diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index ce232140..7a0de26e 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -122,11 +122,39 @@ vi.mock('./cliTaskStore', () => ({ }, })) -import { mapHistoryMessagesToUiMessages, reconstructAgentNotifications, useChatStore } from './chatStore' +import { sessionsApi } from '../api/sessions' +import { + mapHistoryMessagesToUiMessages, + reconstructAgentNotifications, + type PerSessionState, + useChatStore, +} from './chatStore' const TEST_SESSION_ID = 'test-session-1' const initialState = useChatStore.getState() +function makeSession(overrides: Partial = {}): PerSessionState { + return { + messages: [], + chatState: 'streaming', + connectionState: 'connected', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + pendingPermission: null, + pendingComputerUsePermission: null, + tokenUsage: { input_tokens: 0, output_tokens: 0 }, + elapsedSeconds: 0, + statusVerb: '', + slashCommands: [], + agentTaskNotifications: {}, + elapsedTimer: null, + ...overrides, + } +} + describe('chatStore history mapping', () => { beforeEach(() => { sendMock.mockReset() @@ -527,6 +555,124 @@ describe('chatStore history mapping', () => { ]) }) + it('hydrates TodoWrite history into the currently tracked task store only', async () => { + const todos = [{ content: 'Session task', status: 'in_progress' }] + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'assistant-todo', + type: 'assistant', + timestamp: '2026-04-06T00:00:00.000Z', + content: [ + { type: 'tool_use', name: 'TodoWrite', id: 'todo-1', input: { todos } }, + ], + }, + ], + }) + cliTaskStoreSnapshot.sessionId = TEST_SESSION_ID + cliTaskStoreSnapshot.tasks = [] + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ messages: [] }), + }, + }) + + await useChatStore.getState().loadHistory(TEST_SESSION_ID) + + expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID) + }) + + it('marks history task completion dismissed when the user already continued', async () => { + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'assistant-task', + type: 'assistant', + timestamp: '2026-04-06T00:00:00.000Z', + content: [ + { type: 'tool_use', name: 'TaskCreate', id: 'task-1', input: { subject: 'Done' } }, + ], + }, + { + id: 'user-next', + type: 'user', + timestamp: '2026-04-06T00:00:01.000Z', + content: '继续下一步', + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ messages: [] }), + }, + }) + + await useChatStore.getState().loadHistory(TEST_SESSION_ID) + + expect(setTasksFromTodosMock).toHaveBeenCalledWith([], TEST_SESSION_ID) + expect(markCompletedAndDismissedMock).toHaveBeenCalledWith(TEST_SESSION_ID) + }) + + it('reloads history task state for the requested session', async () => { + const todos = [{ content: 'Reloaded task', status: 'pending' }] + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'assistant-todo', + type: 'assistant', + timestamp: '2026-04-06T00:00:00.000Z', + content: [ + { type: 'tool_use', name: 'TodoWrite', id: 'todo-1', input: { todos } }, + ], + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ messages: [{ id: 'old', type: 'assistant_text', content: 'old', timestamp: 1 }] }), + }, + }) + + await useChatStore.getState().reloadHistory(TEST_SESSION_ID) + + expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID) + }) + + it('clears reloaded task state after completed history is followed by a user turn', async () => { + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'assistant-task', + type: 'assistant', + timestamp: '2026-04-06T00:00:00.000Z', + content: [ + { type: 'tool_use', name: 'TaskUpdate', id: 'task-1', input: { subject: 'Done' } }, + ], + }, + { + id: 'user-next', + type: 'user', + timestamp: '2026-04-06T00:00:01.000Z', + content: '新的问题', + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ messages: [{ id: 'old', type: 'assistant_text', content: 'old', timestamp: 1 }] }), + }, + }) + + await useChatStore.getState().reloadHistory(TEST_SESSION_ID) + + expect(setTasksFromTodosMock).toHaveBeenCalledWith([], TEST_SESSION_ID) + expect(markCompletedAndDismissedMock).toHaveBeenCalledWith(TEST_SESSION_ID) + }) + it('keeps parent tool linkage for live tool events', () => { // Initialize the session first useChatStore.setState({ @@ -582,6 +728,24 @@ describe('chatStore history mapping', () => { ]) }) + it('syncs live TodoWrite tool input into the task store for that session', () => { + const todos = [{ content: 'Live todo', status: 'in_progress' }] + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ chatState: 'tool_executing' }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'tool_use_complete', + toolName: 'TodoWrite', + toolUseId: 'todo-live', + input: { todos }, + }) + + expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID) + }) + it('replays saved runtime selection when reconnecting a session', () => { useSessionRuntimeStore.getState().setSelection(TEST_SESSION_ID, { providerId: 'provider-1', @@ -811,6 +975,8 @@ describe('chatStore history mapping', () => { }) it('clears local desktop chat state when the server confirms /clear', () => { + vi.useFakeTimers() + useChatStore.setState({ sessions: { [TEST_SESSION_ID]: { @@ -837,6 +1003,10 @@ describe('chatStore history mapping', () => { }, }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_delta', + text: 'stale throttled delta', + }) useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { type: 'system_notification', subtype: 'session_cleared', @@ -849,7 +1019,34 @@ describe('chatStore history mapping', () => { expect(session?.chatState).toBe('idle') expect(session?.tokenUsage).toEqual({ input_tokens: 0, output_tokens: 0 }) expect(session?.slashCommands).toEqual([]) - expect(clearTasksMock).toHaveBeenCalled() + expect(clearTasksMock).toHaveBeenCalledWith(TEST_SESSION_ID) + + vi.advanceTimersByTime(60) + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.streamingText).toBe('') + vi.useRealTimers() + }) + + it('clears local message state for only the requested session', () => { + useChatStore.setState({ + sessions: { + 'session-a': makeSession({ + messages: [{ id: 'a1', type: 'assistant_text', content: 'A old', timestamp: 1 }], + streamingText: 'A pending', + }), + 'session-b': makeSession({ + messages: [{ id: 'b1', type: 'assistant_text', content: 'B old', timestamp: 1 }], + streamingText: 'B pending', + }), + }, + }) + + useChatStore.getState().clearMessages('session-a') + + expect(useChatStore.getState().sessions['session-a']?.messages).toEqual([]) + expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('') + expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([ + { content: 'B old' }, + ]) }) it('renders compact boundary notifications as system messages', () => { @@ -958,7 +1155,7 @@ describe('chatStore history mapping', () => { useChatStore.getState().sendMessage(TEST_SESSION_ID, '继续下一轮') - expect(resetCompletedTasksMock).toHaveBeenCalledTimes(1) + expect(resetCompletedTasksMock).toHaveBeenCalledWith(TEST_SESSION_ID) expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ { type: 'task_summary', @@ -974,6 +1171,71 @@ describe('chatStore history mapping', () => { ]) }) + it('does not attach completed tasks from another tracked session to a new user turn', () => { + cliTaskStoreSnapshot.sessionId = 'session-b' + cliTaskStoreSnapshot.tasks = [ + { id: '1', subject: 'Session B completed task', status: 'completed' }, + ] + + useChatStore.setState({ + sessions: { + 'session-a': makeSession({ chatState: 'idle' }), + 'session-b': makeSession({ chatState: 'idle' }), + }, + }) + + useChatStore.getState().sendMessage('session-a', '继续 A 会话') + + expect(resetCompletedTasksMock).not.toHaveBeenCalled() + expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([ + { + type: 'user_text', + content: '继续 A 会话', + }, + ]) + }) + + it('tracks task tool results independently per session even when tool IDs collide', () => { + useChatStore.setState({ + sessions: { + 'session-a': makeSession({ + activeToolUseId: 'tool-same', + activeToolName: 'TaskCreate', + }), + 'session-b': makeSession({ + activeToolUseId: 'tool-same', + activeToolName: 'TaskCreate', + }), + }, + }) + + for (const sessionId of ['session-a', 'session-b']) { + useChatStore.getState().handleServerMessage(sessionId, { + type: 'tool_use_complete', + toolName: 'TaskCreate', + toolUseId: 'tool-same', + input: { subject: sessionId }, + }) + } + + useChatStore.getState().handleServerMessage('session-a', { + type: 'tool_result', + toolUseId: 'tool-same', + content: 'created A', + isError: false, + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'tool_result', + toolUseId: 'tool-same', + content: 'created B', + isError: false, + }) + + expect(refreshTasksMock).toHaveBeenCalledTimes(2) + expect(refreshTasksMock).toHaveBeenNthCalledWith(1, 'session-a') + expect(refreshTasksMock).toHaveBeenNthCalledWith(2, 'session-b') + }) + it('tracks Computer Use approval requests separately from generic tool permissions', () => { useChatStore.setState({ sessions: { @@ -1104,6 +1366,218 @@ describe('chatStore history mapping', () => { vi.useRealTimers() }) + it('keeps throttled streaming deltas isolated per session', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + 'session-a': makeSession(), + 'session-b': makeSession(), + }, + }) + + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_start', + blockType: 'text', + }) + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_delta', + text: 'A-only response', + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'content_start', + blockType: 'text', + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'content_delta', + text: 'B-only response', + }) + + useChatStore.getState().handleServerMessage('session-a', { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 1 }, + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 1 }, + }) + + expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([ + { type: 'assistant_text', content: 'A-only response' }, + ]) + expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([ + { type: 'assistant_text', content: 'B-only response' }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + + it('flushes pending text before appending a thinking block', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ chatState: 'streaming' }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_delta', + text: 'visible answer before thinking', + }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'thinking', + text: 'internal note', + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ + { type: 'assistant_text', content: 'visible answer before thinking' }, + { type: 'thinking', content: 'internal note' }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + + it('flushes pending text before appending an error message', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ chatState: 'streaming' }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_delta', + text: 'partial answer before error', + }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'error', + message: 'provider failed', + code: 'provider_error', + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ + { type: 'assistant_text', content: 'partial answer before error' }, + { type: 'error', message: 'provider failed' }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + + it('flushes throttled deltas only for the stopped session', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + 'session-a': makeSession(), + 'session-b': makeSession(), + }, + }) + + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_delta', + text: 'A-only response', + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'content_delta', + text: 'B-only response', + }) + + useChatStore.getState().stopGeneration('session-a') + + expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('A-only response') + expect(useChatStore.getState().sessions['session-b']?.streamingText).toBe('') + + useChatStore.getState().handleServerMessage('session-b', { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 1 }, + }) + + expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([ + { type: 'assistant_text', content: 'B-only response' }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + + it('does not flush one session throttled delta into another disconnected session', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + 'session-a': makeSession(), + 'session-b': makeSession(), + }, + }) + + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_delta', + text: 'A-only response', + }) + useChatStore.getState().handleServerMessage('session-b', { + type: 'content_delta', + text: 'B-only response', + }) + + useChatStore.getState().disconnectSession('session-a') + + expect(useChatStore.getState().sessions['session-a']).toBeUndefined() + expect(useChatStore.getState().sessions['session-b']?.streamingText).toBe('') + + useChatStore.getState().handleServerMessage('session-b', { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 1 }, + }) + + expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([ + { type: 'assistant_text', content: 'B-only response' }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + + it('ignores late throttled deltas after a session has disconnected', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + 'session-a': makeSession(), + }, + }) + + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_delta', + text: 'before disconnect', + }) + useChatStore.getState().disconnectSession('session-a') + + useChatStore.getState().handleServerMessage('session-a', { + type: 'content_delta', + text: 'late stale delta', + }) + useChatStore.setState({ + sessions: { + 'session-a': makeSession({ chatState: 'idle' }), + }, + }) + + useChatStore.getState().sendMessage('session-a', 'fresh turn') + + expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([ + { type: 'user_text', content: 'fresh turn' }, + ]) + + vi.runOnlyPendingTimers() + expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('') + vi.useRealTimers() + }) + it('does not split one streamed markdown reply when task progress arrives mid-stream', () => { vi.useFakeTimers() diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 00020f54..9292e662 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -124,26 +124,62 @@ type ChatStore = { } const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite']) -const pendingTaskToolUseIds = new Set() +const pendingTaskToolUseIdsBySession = new Map>() + +function addPendingTaskToolUseId(sessionId: string, toolUseId: string): void { + const ids = pendingTaskToolUseIdsBySession.get(sessionId) ?? new Set() + ids.add(toolUseId) + pendingTaskToolUseIdsBySession.set(sessionId, ids) +} + +function consumePendingTaskToolUseId(sessionId: string, toolUseId: string): boolean { + const ids = pendingTaskToolUseIdsBySession.get(sessionId) + if (!ids?.has(toolUseId)) return false + ids.delete(toolUseId) + if (ids.size === 0) pendingTaskToolUseIdsBySession.delete(sessionId) + return true +} + +function clearPendingTaskToolUseIds(sessionId: string): void { + pendingTaskToolUseIdsBySession.delete(sessionId) +} const AGENT_COMPLETION_NOTIFICATION_PREVIEW_CHARS = 160 let msgCounter = 0 const nextId = () => `msg-${++msgCounter}-${Date.now()}` -// Streaming throttle for content_delta -let pendingDelta = '' -let flushTimer: ReturnType | null = null +// Streaming throttle for content_delta. Buffers must be per-session because +// multiple desktop tabs can stream at the same time. +const pendingDeltaBySession = new Map() +const flushTimerBySession = new Map>() -function consumePendingDelta(): string { +function consumePendingDelta(sessionId: string): string { + const flushTimer = flushTimerBySession.get(sessionId) if (flushTimer) { clearTimeout(flushTimer) - flushTimer = null + flushTimerBySession.delete(sessionId) } - const text = pendingDelta - pendingDelta = '' + const text = pendingDeltaBySession.get(sessionId) ?? '' + pendingDeltaBySession.delete(sessionId) return text } +function appendPendingDelta(sessionId: string, text: string): void { + pendingDeltaBySession.set( + sessionId, + `${pendingDeltaBySession.get(sessionId) ?? ''}${text}`, + ) +} + +function clearPendingDelta(sessionId: string): void { + const flushTimer = flushTimerBySession.get(sessionId) + if (flushTimer) { + clearTimeout(flushTimer) + flushTimerBySession.delete(sessionId) + } + pendingDeltaBySession.delete(sessionId) +} + function appendAssistantTextMessage( messages: UIMessage[], content: string, @@ -283,11 +319,11 @@ export const useChatStore = create((set, get) => ({ disconnectSession: (sessionId) => { const session = get().sessions[sessionId] if (session?.elapsedTimer) clearInterval(session.elapsedTimer) - if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } - if (pendingDelta) { - const text = consumePendingDelta() + if (pendingDeltaBySession.has(sessionId)) { + const text = consumePendingDelta(sessionId) set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) } + clearPendingTaskToolUseIds(sessionId) wsManager.disconnect(sessionId) set((s) => { const { [sessionId]: _, ...rest } = s.sessions @@ -317,13 +353,14 @@ export const useChatStore = create((set, get) => ({ : undefined const taskStore = useCLITaskStore.getState() - const allTasksDone = taskStore.tasks.length > 0 && taskStore.tasks.every((t) => t.status === 'completed') + const sessionTasks = taskStore.sessionId === sessionId ? taskStore.tasks : [] + const allTasksDone = sessionTasks.length > 0 && sessionTasks.every((t) => t.status === 'completed') const completedTaskSummary = allTasksDone - ? taskStore.tasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm })) + ? sessionTasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm })) : [] if (!isMemberSession && allTasksDone) { - void taskStore.resetCompletedTasks() + void taskStore.resetCompletedTasks(sessionId) } if (!isMemberSession) { @@ -332,11 +369,7 @@ export const useChatStore = create((set, get) => ({ set((s) => { const session = s.sessions[sessionId] ?? createDefaultSessionState() - if (flushTimer) { - clearTimeout(flushTimer) - flushTimer = null - } - const bufferedDelta = consumePendingDelta() + const bufferedDelta = consumePendingDelta(sessionId) const pendingAssistantText = `${session.streamingText}${bufferedDelta}` const newMessages = pendingAssistantText.trim() @@ -449,9 +482,8 @@ export const useChatStore = create((set, get) => ({ stopGeneration: (sessionId) => { wsManager.send(sessionId, { type: 'stop_generation' }) - if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } - if (pendingDelta) { - const text = consumePendingDelta() + if (pendingDeltaBySession.has(sessionId)) { + const text = consumePendingDelta(sessionId) set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) } set((s) => { @@ -491,12 +523,12 @@ export const useChatStore = create((set, get) => ({ }) if (lastTodos && lastTodos.length > 0) { const taskStore = useCLITaskStore.getState() - if (taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos) + if (taskStore.sessionId === sessionId && taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos, sessionId) } else { - useCLITaskStore.getState().setTasksFromTodos([]) + useCLITaskStore.getState().setTasksFromTodos([], sessionId) } if (hasMessagesAfterTaskCompletion) { - useCLITaskStore.getState().markCompletedAndDismissed() + useCLITaskStore.getState().markCompletedAndDismissed(sessionId) } } catch { // Session may not have messages yet @@ -535,12 +567,12 @@ export const useChatStore = create((set, get) => ({ }) if (lastTodos && lastTodos.length > 0) { - useCLITaskStore.getState().setTasksFromTodos(lastTodos) + useCLITaskStore.getState().setTasksFromTodos(lastTodos, sessionId) } else { - useCLITaskStore.getState().setTasksFromTodos([]) + useCLITaskStore.getState().setTasksFromTodos([], sessionId) } if (hasMessagesAfterTaskCompletion) { - useCLITaskStore.getState().markCompletedAndDismissed() + useCLITaskStore.getState().markCompletedAndDismissed(sessionId) } } catch { // Session may not have messages yet @@ -560,6 +592,7 @@ export const useChatStore = create((set, get) => ({ }, clearMessages: (sessionId) => { + clearPendingTaskToolUseIds(sessionId) set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ messages: [], streamingText: '', chatState: 'idle' })) })) }, @@ -574,7 +607,7 @@ export const useChatStore = create((set, get) => ({ case 'status': update((session) => { - const pendingText = `${session.streamingText}${consumePendingDelta()}` + const pendingText = `${session.streamingText}${consumePendingDelta(sessionId)}` const hasPendingStreamText = session.chatState === 'streaming' && pendingText.trim().length > 0 // Background task progress can arrive while the assistant is still @@ -608,7 +641,7 @@ export const useChatStore = create((set, get) => ({ case 'content_start': { const session = get().sessions[sessionId] if (!session) break - const pendingText = `${session.streamingText}${consumePendingDelta()}` + const pendingText = `${session.streamingText}${consumePendingDelta(sessionId)}` if (msg.blockType !== 'text' && pendingText.trim()) { update((s) => ({ messages: appendAssistantTextMessage(s.messages, pendingText, Date.now()), @@ -635,14 +668,16 @@ export const useChatStore = create((set, get) => ({ case 'content_delta': if (msg.text !== undefined) { - pendingDelta += msg.text - if (!flushTimer) { - flushTimer = setTimeout(() => { - const text = pendingDelta - pendingDelta = '' - flushTimer = null + if (!get().sessions[sessionId]) break + appendPendingDelta(sessionId, msg.text) + if (!flushTimerBySession.has(sessionId)) { + const timer = setTimeout(() => { + const text = pendingDeltaBySession.get(sessionId) ?? '' + pendingDeltaBySession.delete(sessionId) + flushTimerBySession.delete(sessionId) update((s) => ({ streamingText: s.streamingText + text })) }, 50) + flushTimerBySession.set(sessionId, timer) } } if (msg.toolInput !== undefined) update((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput })) @@ -650,7 +685,7 @@ export const useChatStore = create((set, get) => ({ case 'thinking': update((s) => { - const pendingText = `${s.streamingText}${consumePendingDelta()}` + const pendingText = `${s.streamingText}${consumePendingDelta(sessionId)}` const base = pendingText.trim() ? appendAssistantTextMessage(s.messages, pendingText, Date.now()) : s.messages @@ -682,10 +717,10 @@ export const useChatStore = create((set, get) => ({ activeToolUseId: null, activeToolName: null, activeThinkingId: null, streamingToolInput: '', })) if (toolName === 'TodoWrite' && Array.isArray((msg.input as any)?.todos)) { - useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos) + useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos, sessionId) } else if (TASK_TOOL_NAMES.has(toolName)) { const useId = msg.toolUseId || session?.activeToolUseId - if (useId) pendingTaskToolUseIds.add(useId) + if (useId) addPendingTaskToolUseId(sessionId, useId) } break } @@ -698,9 +733,8 @@ export const useChatStore = create((set, get) => ({ }], chatState: 'thinking', activeThinkingId: null, })) - if (pendingTaskToolUseIds.has(msg.toolUseId)) { - pendingTaskToolUseIds.delete(msg.toolUseId) - useCLITaskStore.getState().refreshTasks() + if (consumePendingTaskToolUseId(sessionId, msg.toolUseId)) { + useCLITaskStore.getState().refreshTasks(sessionId) } break @@ -766,7 +800,7 @@ export const useChatStore = create((set, get) => ({ const session = get().sessions[sessionId] if (!session) break const wasAgentRunning = session.chatState !== 'idle' - const text = `${session.streamingText}${consumePendingDelta()}` + const text = `${session.streamingText}${consumePendingDelta(sessionId)}` let completionMessages = session.messages if (text.trim()) { completionMessages = appendAssistantTextMessage(session.messages, text, Date.now()) @@ -803,7 +837,7 @@ export const useChatStore = create((set, get) => ({ case 'error': update((s) => { - const pendingText = `${s.streamingText}${consumePendingDelta()}` + const pendingText = `${s.streamingText}${consumePendingDelta(sessionId)}` let newMessages = s.messages if (pendingText.trim()) { newMessages = appendAssistantTextMessage(newMessages, pendingText, Date.now()) @@ -866,7 +900,9 @@ export const useChatStore = create((set, get) => ({ tokenUsage: { input_tokens: 0, output_tokens: 0 }, slashCommands: [], })) - useCLITaskStore.getState().clearTasks() + clearPendingDelta(sessionId) + clearPendingTaskToolUseIds(sessionId) + useCLITaskStore.getState().clearTasks(sessionId) useSessionStore.getState().updateSessionTitle(sessionId, 'New Session') useTabStore.getState().updateTabTitle(sessionId, 'New Session') useTabStore.getState().updateTabStatus(sessionId, 'idle') diff --git a/desktop/src/stores/cliTaskStore.test.ts b/desktop/src/stores/cliTaskStore.test.ts index fa338e1e..f6146d23 100644 --- a/desktop/src/stores/cliTaskStore.test.ts +++ b/desktop/src/stores/cliTaskStore.test.ts @@ -106,4 +106,78 @@ describe('cliTaskStore', () => { expect(useCLITaskStore.getState().resetting).toBe(false) }) + + it('refreshes tasks for the currently tracked session by default', async () => { + vi.mocked(cliTasksApi.getTasksForList).mockResolvedValue({ + tasks: [makeTask('session-1', 'in_progress')], + }) + + useCLITaskStore.setState({ + sessionId: 'session-1', + tasks: [], + expanded: false, + completedAndDismissed: false, + dismissedCompletionKey: null, + }) + + await useCLITaskStore.getState().refreshTasks() + + expect(cliTasksApi.getTasksForList).toHaveBeenCalledWith('session-1') + expect(useCLITaskStore.getState().tasks).toMatchObject([ + { taskListId: 'session-1', status: 'in_progress' }, + ]) + }) + + it('marks completed tasks dismissed for the currently tracked session by default', () => { + useCLITaskStore.setState({ + sessionId: 'session-1', + tasks: [makeTask('session-1', 'completed')], + expanded: true, + completedAndDismissed: false, + dismissedCompletionKey: null, + }) + + useCLITaskStore.getState().markCompletedAndDismissed() + + expect(useCLITaskStore.getState()).toMatchObject({ + completedAndDismissed: true, + dismissedCompletionKey: 'session-1::1::Keep current session isolated::completed::::', + expanded: false, + }) + }) + + it('ignores TodoWrite updates for a session that is not currently tracked', () => { + useCLITaskStore.setState({ + sessionId: 'session-1', + tasks: [makeTask('session-1', 'in_progress')], + expanded: true, + completedAndDismissed: false, + dismissedCompletionKey: null, + }) + + useCLITaskStore.getState().setTasksFromTodos([ + { content: 'Session 2 task', status: 'completed' }, + ], 'session-2') + + expect(useCLITaskStore.getState().tasks).toMatchObject([ + { taskListId: 'session-1', subject: 'Keep current session isolated' }, + ]) + }) + + it('does not reset completed tasks for a different session', async () => { + useCLITaskStore.setState({ + sessionId: 'session-1', + tasks: [makeTask('session-1', 'completed')], + expanded: true, + completedAndDismissed: false, + dismissedCompletionKey: null, + }) + + await useCLITaskStore.getState().resetCompletedTasks('session-2') + + expect(vi.mocked(cliTasksApi.resetTaskList)).not.toHaveBeenCalled() + expect(useCLITaskStore.getState().tasks).toMatchObject([ + { taskListId: 'session-1', status: 'completed' }, + ]) + }) }) diff --git a/desktop/src/stores/cliTaskStore.ts b/desktop/src/stores/cliTaskStore.ts index 3b8db3a0..192f9892 100644 --- a/desktop/src/stores/cliTaskStore.ts +++ b/desktop/src/stores/cliTaskStore.ts @@ -25,16 +25,16 @@ type CLITaskStore = { /** Fetch tasks for a given session (uses sessionId as taskListId) */ fetchSessionTasks: (sessionId: string) => Promise - /** Refresh tasks for the currently tracked session */ - refreshTasks: () => Promise + /** Refresh tasks for the currently tracked session, or a specific session if provided */ + refreshTasks: (sessionId?: string) => Promise /** Update tasks from TodoWrite V1 tool input (in-memory, no disk read needed) */ - setTasksFromTodos: (todos: TodoItem[]) => void + setTasksFromTodos: (todos: TodoItem[], sessionId?: string) => void /** Mark that completed tasks were already dismissed (conversation continued) */ - markCompletedAndDismissed: () => void + markCompletedAndDismissed: (sessionId?: string) => void /** Clear a completed task list locally and remotely so the next cycle starts clean */ - resetCompletedTasks: () => Promise + resetCompletedTasks: (sessionId?: string) => Promise /** Clear task tracking state */ - clearTasks: () => void + clearTasks: (sessionId?: string) => void /** Toggle expanded state */ toggleExpanded: () => void } @@ -116,8 +116,8 @@ export const useCLITaskStore = create((set, get) => ({ } }, - refreshTasks: async () => { - const { sessionId } = get() + refreshTasks: async (targetSessionId) => { + const sessionId = targetSessionId ?? get().sessionId if (!sessionId) return try { const { tasks } = await cliTasksApi.getTasksForList(sessionId) @@ -132,15 +132,19 @@ export const useCLITaskStore = create((set, get) => ({ } }, - setTasksFromTodos: (todos) => { - const tasks = mapTodosToTasks(todos, get().sessionId) + setTasksFromTodos: (todos, targetSessionId) => { + const sessionId = targetSessionId ?? get().sessionId + if (!sessionId || get().sessionId !== sessionId) return + const tasks = mapTodosToTasks(todos, sessionId) set((state) => ({ tasks, ...resolveDismissState(tasks, state.dismissedCompletionKey), })) }, - markCompletedAndDismissed: () => { + markCompletedAndDismissed: (targetSessionId) => { + const sessionId = targetSessionId ?? get().sessionId + if (!sessionId || get().sessionId !== sessionId) return const completionKey = buildCompletedTaskKey(get().tasks) if (!completionKey) return @@ -151,10 +155,12 @@ export const useCLITaskStore = create((set, get) => ({ }) }, - resetCompletedTasks: async () => { - const { sessionId, tasks } = get() + resetCompletedTasks: async (targetSessionId) => { + const sessionId = targetSessionId ?? get().sessionId + if (!sessionId || get().sessionId !== sessionId) return + const { tasks } = get() const completionKey = buildCompletedTaskKey(tasks) - if (!sessionId || !completionKey) return + if (!completionKey) return set({ tasks: [], @@ -173,7 +179,8 @@ export const useCLITaskStore = create((set, get) => ({ } }, - clearTasks: () => { + clearTasks: (targetSessionId) => { + if (targetSessionId && get().sessionId !== targetSessionId) return set({ sessionId: null, tasks: [], diff --git a/src/server/__tests__/websocket-handler.test.ts b/src/server/__tests__/websocket-handler.test.ts new file mode 100644 index 00000000..33224b10 --- /dev/null +++ b/src/server/__tests__/websocket-handler.test.ts @@ -0,0 +1,55 @@ +import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test' +import type { ServerWebSocket } from 'bun' +import { + __resetWebSocketHandlerStateForTests, + getActiveSessionIds, + handleWebSocket, + type WebSocketData, +} from '../ws/handler.js' +import { conversationService } from '../services/conversationService.js' +import { computerUseApprovalService } from '../services/computerUseApprovalService.js' + +function makeClientSocket(sessionId: string) { + const sent: string[] = [] + return { + data: { + sessionId, + connectedAt: Date.now(), + channel: 'client', + sdkToken: null, + serverPort: 0, + serverHost: '127.0.0.1', + }, + send: mock((payload: string) => { + sent.push(payload) + }), + close: mock(() => {}), + sent, + } as unknown as ServerWebSocket & { sent: string[] } +} + +describe('WebSocket handler session isolation', () => { + afterEach(() => { + __resetWebSocketHandlerStateForTests() + mock.restore() + }) + + it('ignores stale disconnects from an older socket for the same session', () => { + const sessionId = `duplicate-${crypto.randomUUID()}` + const first = makeClientSocket(sessionId) + const second = makeClientSocket(sessionId) + const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks') + const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession') + + handleWebSocket.open(first) + handleWebSocket.open(second) + clearCallbacks.mockClear() + cancelComputerUse.mockClear() + + handleWebSocket.close(first, 1000, 'stale tab closed') + + expect(getActiveSessionIds()).toContain(sessionId) + expect(clearCallbacks).not.toHaveBeenCalled() + expect(cancelComputerUse).not.toHaveBeenCalled() + }) +}) diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index a49bf367..8dac662e 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -193,6 +193,10 @@ export const handleWebSocket = { } console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`) + if (activeSessions.get(sessionId) !== ws) { + console.log(`[WS] Ignoring stale client disconnect for session: ${sessionId}`) + return + } computerUseApprovalService.cancelSession(sessionId) activeSessions.delete(sessionId) conversationService.clearOutputCallbacks(sessionId) @@ -1498,3 +1502,11 @@ export function sendToSession(sessionId: string, message: ServerMessage): boolea export function getActiveSessionIds(): string[] { return Array.from(activeSessions.keys()) } + +export function __resetWebSocketHandlerStateForTests(): void { + for (const timer of sessionCleanupTimers.values()) clearTimeout(timer) + for (const timer of prewarmIdleTimers.values()) clearTimeout(timer) + activeSessions.clear() + sessionCleanupTimers.clear() + prewarmIdleTimers.clear() +}