diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index f287fe77..b10924ae 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -1641,6 +1641,64 @@ describe('chatStore history mapping', () => { vi.useRealTimers() }) + it('marks a background shell task stopped when TaskStop returns before a task notification', () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-05-19T13:34:19.000Z')) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession(), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'system_notification', + subtype: 'task_started', + data: { + task_id: 'shell-task-1', + tool_use_id: 'shell-tool-1', + description: 'Start tap proxy', + task_type: 'local_bash', + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'tool_use_complete', + toolName: 'TaskStop', + toolUseId: 'task-stop-1', + input: { task_id: 'shell-task-1' }, + }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'tool_result', + toolUseId: 'task-stop-1', + isError: false, + content: JSON.stringify({ + message: 'Successfully stopped task: shell-task-1 (tap proxy)', + task_id: 'shell-task-1', + task_type: 'local_bash', + command: 'tap proxy', + }), + }) + + const session = useChatStore.getState().sessions[TEST_SESSION_ID] + expect(session?.backgroundAgentTasks?.['shell-task-1']).toMatchObject({ + status: 'stopped', + taskType: 'local_bash', + description: 'tap proxy', + }) + expect(session?.messages.find((message) => message.type === 'background_task')).toMatchObject({ + type: 'background_task', + task: { + taskId: 'shell-task-1', + status: 'stopped', + taskType: 'local_bash', + description: 'tap proxy', + }, + }) + + vi.useRealTimers() + }) + it('removes stale agent task transcript cards by matching tool use id', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index f64d2350..79eb792e 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -32,6 +32,7 @@ import type { } from '../types/chat' type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'reconnecting' +type ToolCall = Extract export type ComposerDraftState = { input: string @@ -143,6 +144,7 @@ type ChatStore = { } const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite']) +const TASK_STOP_TOOL_NAMES = new Set(['TaskStop', 'KillShell']) const pendingTaskToolUseIdsBySession = new Map>() const pendingToolParentUseIdsBySession = new Map>() @@ -1059,15 +1061,32 @@ export const useChatStore = create((set, get) => ({ } case 'tool_result': { + const now = Date.now() const pendingParentToolUseId = consumePendingToolParentUseId(sessionId, msg.toolUseId) const parentToolUseId = msg.parentToolUseId ?? pendingParentToolUseId - update((s) => ({ - messages: [...s.messages, { + update((s) => { + let messages: UIMessage[] = [...s.messages, { id: nextId(), type: 'tool_result', toolUseId: msg.toolUseId, - content: msg.content, isError: msg.isError, timestamp: Date.now(), parentToolUseId, - }], - chatState: 'thinking', activeThinkingId: null, - })) + content: msg.content, isError: msg.isError, timestamp: now, parentToolUseId, + }] + let backgroundAgentTasks = s.backgroundAgentTasks ?? {} + const stoppedTask = msg.isError + ? null + : getStoppedBackgroundTaskFromToolResult(s.messages, msg.toolUseId, msg.content) + if (stoppedTask) { + backgroundAgentTasks = upsertBackgroundAgentTask(backgroundAgentTasks, stoppedTask, now) + const task = backgroundAgentTasks[stoppedTask.taskId] + if (task) { + messages = upsertBackgroundTaskMessage(messages, task, now) + } + } + return { + messages, + ...(stoppedTask ? { backgroundAgentTasks } : {}), + chatState: 'thinking', + activeThinkingId: null, + } + }) if (consumePendingTaskToolUseId(sessionId, msg.toolUseId)) { useCLITaskStore.getState().refreshTasks(sessionId) } @@ -1470,6 +1489,55 @@ function readNonEmptyString(record: Record, ...keys: string[]): return undefined } +function readRecord(value: unknown): Record | null { + if (!value || typeof value !== 'object' || Array.isArray(value)) return null + return value as Record +} + +function parseJsonRecord(value: unknown): Record | null { + const record = readRecord(value) + if (record) return record + if (typeof value !== 'string') return null + const trimmed = value.trim() + if (!trimmed.startsWith('{')) return null + try { + return readRecord(JSON.parse(trimmed)) + } catch { + return null + } +} + +function findToolUseMessage(messages: UIMessage[], toolUseId: string): ToolCall | null { + return messages.find(( + message, + ): message is ToolCall => + message.type === 'tool_use' && + message.toolUseId === toolUseId) ?? null +} + +function getStoppedBackgroundTaskFromToolResult( + messages: UIMessage[], + toolUseId: string, + content: unknown, +): (Partial & Pick) | null { + const toolUse = findToolUseMessage(messages, toolUseId) + if (!toolUse || !TASK_STOP_TOOL_NAMES.has(toolUse.toolName)) return null + + const input = readRecord(toolUse.input) ?? {} + const output = parseJsonRecord(content) ?? {} + const taskId = readNonEmptyString(output, 'task_id', 'taskId') ?? + readNonEmptyString(input, 'task_id', 'taskId', 'shell_id', 'shellId') + if (!taskId) return null + + return { + taskId, + status: 'stopped', + taskType: readNonEmptyString(output, 'task_type', 'taskType'), + description: readNonEmptyString(output, 'command', 'description', 'message'), + summary: readNonEmptyString(output, 'message'), + } +} + function normalizeBackgroundTaskUsage(value: unknown): BackgroundAgentTaskUsage | undefined { if (!value || typeof value !== 'object') return undefined const record = value as Record