diff --git a/apps/desktop/src/app/chat/session-tile-actions.ts b/apps/desktop/src/app/chat/session-tile-actions.ts index 912878e53b0..46969138459 100644 --- a/apps/desktop/src/app/chat/session-tile-actions.ts +++ b/apps/desktop/src/app/chat/session-tile-actions.ts @@ -29,6 +29,7 @@ import { $sessionStates, sessionTileDelegate } from '@/store/session-states' import { broadcastSessionsChanged } from '@/store/session-sync' import { clearSessionSubagents } from '@/store/subagents' import { clearSessionTodos } from '@/store/todos' +import { setSessionDraftingTool } from '@/store/tool-drafting' import type { SessionInfo } from '@/types/hermes' import { uploadComposerAttachment } from '../session/hooks/use-prompt-actions' @@ -253,6 +254,7 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses clearSessionTodos(sessionId) clearSessionSubagents(sessionId) resetSessionBackground(sessionId) + setSessionDraftingTool(sessionId, '') clearAllPrompts(sessionId) clearClarifyRequest(undefined, sessionId) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts index 04aa94522a0..500540b6ccf 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts @@ -106,6 +106,28 @@ function surfaceBillingBlock(sessionId: string, raw: unknown): void { }) } +/** + * Events that retire a "drafting a tool call" claim. + * + * `tool.generating` opens the claim and nothing closes it — a draft can be + * abandoned without ever reaching `tool.start`, so enumerating the ways one + * *ends* left the label on screen for the rest of the turn. Inverted: the + * claim only covers what the model is emitting right now, and any other output + * from the session proves it moved on. Same rule the TUI applies to its + * transient trail lines (`turnController.pruneTransient`). + */ +const DRAFT_SUPERSEDING_EVENT_TYPES = new Set([ + 'error', + 'message.complete', + 'message.delta', + 'message.start', + 'reasoning.delta', + 'thinking.delta', + 'tool.complete', + 'tool.progress', + 'tool.start' +]) + const COMPACTION_RESUME_EVENT_TYPES = new Set([ 'message.delta', 'message.interim', @@ -244,6 +266,10 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { setSessionCompacting(sessionId, false) } + if (sessionId && DRAFT_SUPERSEDING_EVENT_TYPES.has(event.type)) { + setSessionDraftingTool(sessionId, '') + } + if (event.type === 'gateway.ready') { // Seed the active skin into the desktop theme registry without applying, // so a fresh connect never overrides the user's persisted desktop theme. @@ -442,7 +468,6 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { flushQueuedDeltas(sessionId) pruneFinishedSessionSubagents(sessionId) setSessionCompacting(sessionId, false) - setSessionDraftingTool(sessionId, '') compactedTurnRef.current.delete(sessionId) nativeSubagentSessionsRef.current.delete(sessionId) // A fresh turn on this session optimistically clears its billing wall; @@ -613,7 +638,6 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { // last item stuck pending/in_progress. Finished lists keep their linger. clearActiveSessionTodos(sessionId) setSessionCompacting(sessionId, false) - setSessionDraftingTool(sessionId, '') flushQueuedDeltas(sessionId) @@ -686,7 +710,11 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { // from it strands an argless placeholder whenever the bubble is sealed // before the real `tool.start` arrives, because the two can no longer be // reconciled across the boundary. It's a status, so say it as one. - if (!sessionId) { + // A stopped turn can still emit a frame or two before the backend + // notices, and naming a tool we will never run leaves the label up + // until something else retires it. `mutateStream` drops late tool rows + // on the same condition; the status line has to agree with it. + if (!sessionId || sessionInterrupted(sessionId)) { return } @@ -701,7 +729,6 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { } flushQueuedDeltas(sessionId) - setSessionDraftingTool(sessionId, '') upsertToolCall(sessionId, toTodoPayload(payload) ?? payload, 'running', event.type) if (isActiveEvent) { @@ -710,7 +737,6 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { } else if (event.type === 'tool.complete') { if (sessionId) { flushQueuedDeltas(sessionId) - setSessionDraftingTool(sessionId, '') upsertToolCall(sessionId, toTodoPayload(payload) ?? payload, 'complete', event.type) if (isActiveEvent) { @@ -1002,7 +1028,6 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { clearClarifyRequest(undefined, sessionId) clearActiveSessionTodos(sessionId) setSessionCompacting(sessionId, false) - setSessionDraftingTool(sessionId, '') compactedTurnRef.current.delete(sessionId) } diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/tool-drafting-event.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/tool-drafting-event.test.tsx new file mode 100644 index 00000000000..57d5c5b1ced --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/tool-drafting-event.test.tsx @@ -0,0 +1,124 @@ +import { QueryClient } from '@tanstack/react-query' +import { act, cleanup, render, waitFor } from '@testing-library/react' +import { useEffect, useRef } from 'react' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import type { ClientSessionState } from '@/app/types' +import { createClientSessionState } from '@/lib/chat-runtime' +import { $draftingToolSessions } from '@/store/tool-drafting' +import type { RpcEvent } from '@/types/hermes' + +import { useMessageStream } from './index' + +const SID = 'session-1' +const OTHER_SID = 'session-2' + +// Module-scoped so a test can seed session state (e.g. interrupted) before the +// handler reads it — `sessionInterrupted` resolves against this map. +const sessionStates = new Map() +let handleEvent: ((event: RpcEvent) => void) | null = null + +function Harness() { + const activeSessionIdRef = useRef(SID) + const sessionStateByRuntimeIdRef = useRef(sessionStates) + const queryClientRef = useRef(new QueryClient()) + + const stream = useMessageStream({ + activeSessionIdRef, + hydrateFromStoredSession: vi.fn(async () => undefined), + queryClient: queryClientRef.current, + refreshHermesConfig: vi.fn(async () => undefined), + refreshSessions: vi.fn(async () => undefined), + sessionStateByRuntimeIdRef, + updateSessionState: (sessionId, updater) => { + const next = updater(sessionStates.get(sessionId) ?? createClientSessionState()) + sessionStates.set(sessionId, next) + + return next + } + }) + + useEffect(() => { + handleEvent = stream.handleGatewayEvent + }, [stream.handleGatewayEvent]) + + return null +} + +async function mountStream() { + render() + await waitFor(() => expect(handleEvent).not.toBeNull()) +} + +function emit(type: RpcEvent['type'], payload: RpcEvent['payload'] = {}, sessionId = SID) { + act(() => handleEvent!({ payload, session_id: sessionId, type })) +} + +function draftedTool(sessionId = SID) { + return $draftingToolSessions.get()[sessionId]?.name +} + +describe('drafting-tool label lifecycle', () => { + beforeEach(() => { + handleEvent = null + sessionStates.clear() + $draftingToolSessions.set({}) + }) + + afterEach(() => { + cleanup() + sessionStates.clear() + $draftingToolSessions.set({}) + vi.restoreAllMocks() + }) + + it('names the tool the model is drafting', async () => { + await mountStream() + + emit('tool.generating', { name: 'write_file' }) + + expect(draftedTool()).toBe('write_file') + }) + + // The label used to be retired only by the events that mean "this tool ran". + // A tool can be abandoned without ever reaching `tool.start` — a mid-stream + // retry drops a partial call, a guardrail-blocked tool skips the lifecycle + // callbacks — and the name then sat on screen for the rest of the turn. + it.each([ + ['message.delta', { text: 'never mind' }], + ['reasoning.delta', { text: 'reconsidering' }], + ['thinking.delta', { text: 'reconsidering' }], + ['tool.start', { name: 'terminal', tool_id: 'tool-1' }], + ['tool.complete', { name: 'terminal', tool_id: 'tool-1' }], + ['message.complete', { text: 'done' }], + ['error', { message: 'boom' }] + ] as const)('retires the label when %s proves the model moved on', async (type, payload) => { + await mountStream() + emit('tool.generating', { name: 'write_file' }) + + emit(type, payload) + + expect(draftedTool()).toBeUndefined() + }) + + it('leaves another session’s label alone', async () => { + await mountStream() + emit('tool.generating', { name: 'patch' }, OTHER_SID) + emit('tool.generating', { name: 'write_file' }) + + emit('message.delta', { text: 'moving on' }) + + expect(draftedTool()).toBeUndefined() + expect(draftedTool(OTHER_SID)).toBe('patch') + }) + + // A stopped turn can still emit a frame or two before the backend notices. + it('ignores a tool announced after the user hit stop', async () => { + sessionStates.set(SID, { ...createClientSessionState(), interrupted: true }) + await mountStream() + + emit('tool.generating', { name: 'write_file' }) + + expect(draftedTool()).toBeUndefined() + }) +}) diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts index bee6192894c..681ef8d4199 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts @@ -33,6 +33,7 @@ import { } from '@/store/session' import { clearSessionSubagents } from '@/store/subagents' import { clearSessionTodos } from '@/store/todos' +import { setSessionDraftingTool } from '@/store/tool-drafting' import type { ClientSessionState, @@ -592,6 +593,7 @@ export function usePromptActions({ clearSessionTodos(sessionId) clearSessionSubagents(sessionId) resetSessionBackground(sessionId) + setSessionDraftingTool(sessionId, '') // Stop ends the turn, so the gateway is no longer blocked on any prompt it // raised. Drop this session's pending clarify / approval / sudo / secret so // a dead panel (and the sidebar "needs input" dot) can't linger and accept