diff --git a/desktop/src/pages/ActiveSession.test.tsx b/desktop/src/pages/ActiveSession.test.tsx index 52f17ff0..e5b99d4d 100644 --- a/desktop/src/pages/ActiveSession.test.tsx +++ b/desktop/src/pages/ActiveSession.test.tsx @@ -100,7 +100,7 @@ describe('ActiveSession task polling', () => { expect( fetchSessionTasks.mock.calls.filter(([currentSessionId]) => currentSessionId === sessionId), - ).toHaveLength(3) + ).toHaveLength(4) unmount() useCLITaskStore.setState(originalCliTaskState) diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index 9ac30215..03205ac1 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -8,6 +8,11 @@ const { handleTeamCreatedMock, handleTeamUpdateMock, handleTeamDeletedMock, + fetchSessionTasksMock, + clearTasksMock, + setTasksFromTodosMock, + markCompletedAndDismissedMock, + refreshTasksMock, } = vi.hoisted(() => ({ sendMock: vi.fn(), getMemberBySessionIdMock: vi.fn<(sessionId: string) => any>(() => null), @@ -15,6 +20,11 @@ const { handleTeamCreatedMock: vi.fn(), handleTeamUpdateMock: vi.fn(), handleTeamDeletedMock: vi.fn(), + fetchSessionTasksMock: vi.fn(), + clearTasksMock: vi.fn(), + setTasksFromTodosMock: vi.fn(), + markCompletedAndDismissedMock: vi.fn(), + refreshTasksMock: vi.fn(), })) vi.mock('../api/websocket', () => ({ @@ -66,12 +76,13 @@ vi.mock('./sessionStore', () => ({ vi.mock('./cliTaskStore', () => ({ useCLITaskStore: { getState: () => ({ - fetchSessionTasks: vi.fn(), + fetchSessionTasks: fetchSessionTasksMock, tasks: [], - clearTasks: vi.fn(), - setTasksFromTodos: vi.fn(), - markCompletedAndDismissed: vi.fn(), - refreshTasks: vi.fn(), + sessionId: null, + clearTasks: clearTasksMock, + setTasksFromTodos: setTasksFromTodosMock, + markCompletedAndDismissed: markCompletedAndDismissedMock, + refreshTasks: refreshTasksMock, }), }, })) @@ -87,6 +98,11 @@ describe('chatStore history mapping', () => { getMemberBySessionIdMock.mockReset() getMemberBySessionIdMock.mockReturnValue(null) sendMessageToMemberMock.mockReset() + fetchSessionTasksMock.mockReset() + clearTasksMock.mockReset() + setTasksFromTodosMock.mockReset() + markCompletedAndDismissedMock.mockReset() + refreshTasksMock.mockReset() useChatStore.setState({ ...initialState, sessions: {}, @@ -502,4 +518,33 @@ describe('chatStore history mapping', () => { pending: true, }) }) + + it('refreshes CLI tasks when switching to an already-connected session', () => { + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: { + messages: [], + chatState: 'idle', + 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, + }, + }, + }) + + useChatStore.getState().connectToSession(TEST_SESSION_ID) + + expect(fetchSessionTasksMock).toHaveBeenCalledWith(TEST_SESSION_ID) + }) }) diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index e8420753..64221662 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -120,6 +120,8 @@ export const useChatStore = create((set, get) => ({ getSession: (sessionId) => get().sessions[sessionId] ?? createDefaultSessionState(), connectToSession: (sessionId) => { + void useCLITaskStore.getState().fetchSessionTasks(sessionId) + const existing = get().sessions[sessionId] if (existing && existing.connectionState !== 'disconnected') return @@ -144,7 +146,6 @@ export const useChatStore = create((set, get) => ({ }) get().loadHistory(sessionId) - useCLITaskStore.getState().fetchSessionTasks(sessionId) sessionsApi.getSlashCommands(sessionId) .then(({ commands }) => { if (get().sessions[sessionId]) { diff --git a/desktop/src/stores/cliTaskStore.test.ts b/desktop/src/stores/cliTaskStore.test.ts new file mode 100644 index 00000000..8ec9aa2d --- /dev/null +++ b/desktop/src/stores/cliTaskStore.test.ts @@ -0,0 +1,70 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { cliTasksApi } from '../api/cliTasks' +import type { CLITask, TaskStatus } from '../types/cliTask' +import { useCLITaskStore } from './cliTaskStore' + +vi.mock('../api/cliTasks', () => ({ + cliTasksApi: { + getTasksForList: vi.fn(), + }, +})) + +function makeTask(taskListId: string, status: TaskStatus = 'in_progress'): CLITask { + return { + id: '1', + subject: 'Keep current session isolated', + description: '', + status, + blocks: [], + blockedBy: [], + taskListId, + } +} + +describe('cliTaskStore', () => { + beforeEach(() => { + useCLITaskStore.getState().clearTasks() + vi.clearAllMocks() + }) + + afterEach(() => { + useCLITaskStore.getState().clearTasks() + }) + + it('clears stale tasks immediately when switching tracked sessions', async () => { + let resolveRequest: ((value: { tasks: ReturnType[] }) => void) | null = null + + vi.mocked(cliTasksApi.getTasksForList).mockImplementation( + (sessionId: string) => + new Promise<{ tasks: ReturnType[] }>((resolve) => { + if (sessionId === 'session-2') resolveRequest = resolve + }), + ) + + useCLITaskStore.setState({ + sessionId: 'session-1', + tasks: [makeTask('session-1')], + expanded: true, + completedAndDismissed: true, + dismissedCompletionKey: 'session-1::done', + }) + + const fetchPromise = useCLITaskStore.getState().fetchSessionTasks('session-2') + + expect(useCLITaskStore.getState()).toMatchObject({ + sessionId: 'session-2', + tasks: [], + expanded: false, + completedAndDismissed: false, + dismissedCompletionKey: null, + }) + + expect(resolveRequest).not.toBeNull() + resolveRequest!({ tasks: [makeTask('session-2', 'completed')] }) + await fetchPromise + + expect(useCLITaskStore.getState().tasks).toMatchObject([ + { taskListId: 'session-2', status: 'completed' }, + ]) + }) +}) diff --git a/desktop/src/stores/cliTaskStore.ts b/desktop/src/stores/cliTaskStore.ts index 960a38a7..979b145e 100644 --- a/desktop/src/stores/cliTaskStore.ts +++ b/desktop/src/stores/cliTaskStore.ts @@ -83,7 +83,16 @@ export const useCLITaskStore = create((set, get) => ({ dismissedCompletionKey: null, fetchSessionTasks: async (sessionId) => { - set({ sessionId }) + if (get().sessionId !== sessionId) { + set({ + sessionId, + tasks: [], + completedAndDismissed: false, + dismissedCompletionKey: null, + expanded: false, + }) + } + try { const { tasks } = await cliTasksApi.getTasksForList(sessionId) // Only update if still tracking the same session @@ -96,7 +105,7 @@ export const useCLITaskStore = create((set, get) => ({ } catch { // No tasks for this session — that's fine if (get().sessionId === sessionId) { - set({ tasks: [], completedAndDismissed: false, dismissedCompletionKey: null }) + set({ tasks: [], completedAndDismissed: false, dismissedCompletionKey: null, expanded: false }) } } },