diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx new file mode 100644 index 000000000000..a403bde23800 --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx @@ -0,0 +1,82 @@ +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 { $compactingSessions, setSessionCompacting } from '@/store/compaction' +import type { RpcEvent } from '@/types/hermes' + +import { useMessageStream } from './index' + +const SID = 'session-1' +const OTHER_SID = 'session-2' +let handleEvent: ((event: RpcEvent) => void) | null = null + +function Harness() { + const activeSessionIdRef = useRef(SID) + const sessionStateByRuntimeIdRef = useRef(new Map()) + 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 current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState() + const next = updater(current) + sessionStateByRuntimeIdRef.current.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'] = {}) { + act(() => handleEvent!({ payload, session_id: SID, type })) +} + +describe('useMessageStream compaction lifecycle', () => { + beforeEach(() => { + handleEvent = null + $compactingSessions.set({}) + }) + + afterEach(() => { + cleanup() + $compactingSessions.set({}) + vi.restoreAllMocks() + }) + + it.each([ + ['message.delta', { text: 'resumed' }], + ['thinking.delta', { text: 'still working' }], + ['reasoning.delta', { text: 'thinking again' }], + ['tool.start', { name: 'terminal', tool_id: 'tool-1' }] + ] as const)('clears the stale compaction phase when %s resumes the turn', async (type, payload) => { + await mountStream() + setSessionCompacting(OTHER_SID, true) + + emit('status.update', { kind: 'compacting' }) + expect($compactingSessions.get()).toEqual({ [OTHER_SID]: true, [SID]: true }) + + emit(type, payload) + + expect($compactingSessions.get()).toEqual({ [OTHER_SID]: true }) + }) +}) 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 77104ba6295d..c50db6ef1b1d 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 @@ -50,6 +50,19 @@ import type { ClientSessionState } from '../../../types' import { hasSessionInfoStatePatch, sessionInfoStatePatch, SUBAGENT_EVENT_TYPES, toTodoPayload } from './utils' +const COMPACTION_RESUME_EVENT_TYPES = new Set([ + 'message.delta', + 'thinking.delta', + 'reasoning.delta', + 'reasoning.available', + 'moa.reference', + 'moa.aggregating', + 'tool.start', + 'tool.progress', + 'tool.generating', + 'tool.complete' +]) + interface GatewayEventDeps { activeSessionIdRef: MutableRefObject compactedTurnRef: MutableRefObject> @@ -118,6 +131,18 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { const sessionId = route.sessionId const isActiveEvent = !!sessionId && sessionId === activeSessionIdRef.current + // Mid-turn compaction does not emit another message.start. The first + // model output or tool event proves summarization has finished and the + // turn has resumed, so retire the phase label without waiting for the + // whole turn to complete. + if ( + sessionId && + COMPACTION_RESUME_EVENT_TYPES.has(event.type) && + compactedTurnRef.current.has(sessionId) + ) { + setSessionCompacting(sessionId, false) + } + if (event.type === 'gateway.ready') { return } else if (event.type === 'session.info') { diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx index 3e725b1480d3..9c5c269eb27b 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx @@ -5,7 +5,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { textPart } from '@/lib/chat-messages' import { $composerAttachments, $composerDraft, type ComposerAttachment, setComposerDraft } from '@/store/composer' -import { $busy, $connection, $messages, $sessions, setSessions } from '@/store/session' +import { $busy, $connection, $messages, $sessions, $turnStartedAt, setSessions } from '@/store/session' import type { SessionInfo } from '@/types/hermes' import { uploadComposerAttachment, usePromptActions } from '.' @@ -1075,6 +1075,7 @@ describe('usePromptActions sleep/wake session recovery', () => { afterEach(() => { cleanup() + $turnStartedAt.set(null) vi.restoreAllMocks() }) @@ -1168,6 +1169,32 @@ describe('usePromptActions sleep/wake session recovery', () => { expect(calls[2]?.params).toEqual({ session_id: RECOVERED_SESSION_ID }) }) + it('clears the active and cached turn clocks when stopping a turn', async () => { + const states: Record[] = [] + const requestGateway = vi.fn(async () => ({}) as never) + $turnStartedAt.set(1_700_000_000_000) + + let handle: HarnessHandle | null = null + await actRender( + (handle = h)} + onSeedState={state => states.push(state)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + /> + ) + + await handle!.cancelRun() + + expect($turnStartedAt.get()).toBeNull() + expect(states.at(-1)).toMatchObject({ + awaitingResponse: false, + busy: false, + interrupted: true, + turnStartedAt: null + }) + }) + it('surfaces the original error (no resume) when the failure is not "session not found"', async () => { const calls: string[] = [] const states: Record[] = [] 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 51a3689ae4e1..89f978cd69f8 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 @@ -21,7 +21,15 @@ import { resetSessionBackground } from '@/store/composer-status' import { clearNotifications, notify, notifyError } from '@/store/notifications' import { clearPreviewArtifacts } from '@/store/preview-status' import { clearAllPrompts } from '@/store/prompts' -import { $busy, $connection, $messages, setAwaitingResponse, setBusy, setMessages } from '@/store/session' +import { + $busy, + $connection, + $messages, + setAwaitingResponse, + setBusy, + setMessages, + setTurnStartedAt +} from '@/store/session' import { clearSessionSubagents } from '@/store/subagents' import { clearSessionTodos } from '@/store/todos' @@ -501,6 +509,7 @@ export function usePromptActions({ } setAwaitingResponse(false) + setTurnStartedAt(null) const finalizeMessages = (messages: ChatMessage[], streamId?: string | null) => messages @@ -526,7 +535,8 @@ export function usePromptActions({ streamId: null, pendingBranchGroup: null, needsInput: false, - interrupted: true + interrupted: true, + turnStartedAt: null } }) diff --git a/apps/desktop/src/components/assistant-ui/thread/status.test.tsx b/apps/desktop/src/components/assistant-ui/thread/status.test.tsx new file mode 100644 index 000000000000..b4920e1ec483 --- /dev/null +++ b/apps/desktop/src/components/assistant-ui/thread/status.test.tsx @@ -0,0 +1,56 @@ +import { act, cleanup, render, screen } from '@testing-library/react' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { __resetElapsedTimerRegistryForTests } from '@/components/chat/activity-timer' +import { I18nProvider } from '@/i18n' +import { $activeSessionId, $turnStartedAt } from '@/store/session' + +import { ResponseLoadingIndicator } from './status' + +function renderIndicator() { + return render( + + + + ) +} + +describe('ResponseLoadingIndicator timer', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + __resetElapsedTimerRegistryForTests() + }) + + afterEach(() => { + cleanup() + $activeSessionId.set(null) + $turnStartedAt.set(null) + __resetElapsedTimerRegistryForTests() + vi.useRealTimers() + }) + + it('preserves each running session timer while switching between sessions', () => { + $activeSessionId.set('session-a') + $turnStartedAt.set(Date.now()) + const sessionA = renderIndicator() + + act(() => vi.advanceTimersByTime(5_000)) + expect(screen.getByText('5s')).toBeTruthy() + sessionA.unmount() + + $activeSessionId.set('session-b') + $turnStartedAt.set(Date.now()) + const sessionB = renderIndicator() + + act(() => vi.advanceTimersByTime(3_000)) + expect(screen.getByText('3s')).toBeTruthy() + sessionB.unmount() + + $activeSessionId.set('session-a') + $turnStartedAt.set(new Date('2026-01-01T00:00:00.000Z').getTime()) + renderIndicator() + + expect(screen.getByText('8s')).toBeTruthy() + }) +}) diff --git a/apps/desktop/src/components/assistant-ui/thread/status.tsx b/apps/desktop/src/components/assistant-ui/thread/status.tsx index 53c6f415a85f..f65793d18ca8 100644 --- a/apps/desktop/src/components/assistant-ui/thread/status.tsx +++ b/apps/desktop/src/components/assistant-ui/thread/status.tsx @@ -11,6 +11,7 @@ import { cn } from '@/lib/utils' import { $backgroundResume } from '@/store/background-delegation' import { $compactionActive } from '@/store/compaction' import { $activeSessionAwaitingInput } from '@/store/prompts' +import { $activeSessionId, $turnStartedAt } from '@/store/session' const StatusRow: FC<{ children: ReactNode; label: string } & React.ComponentPropsWithoutRef<'div'>> = ({ children, @@ -36,6 +37,13 @@ const CompactionHint: FC = () => ( {COMPACTION_LABEL} ) +function useActiveTurnTimerKey(): string | undefined { + const activeSessionId = useStore($activeSessionId) + const turnStartedAt = useStore($turnStartedAt) + + return activeSessionId && turnStartedAt ? `turn:${activeSessionId}:${turnStartedAt}` : undefined +} + export const CenteredThreadSpinner: FC = () => { const { t } = useI18n() @@ -59,7 +67,8 @@ export const CenteredThreadSpinner: FC = () => { export const ResponseLoadingIndicator: FC = () => { const { t } = useI18n() - const elapsed = useElapsedSeconds() + const timerKey = useActiveTurnTimerKey() + const elapsed = useElapsedSeconds(true, timerKey) const compacting = useStore($compactionActive) return ( @@ -134,6 +143,7 @@ export const StreamStallIndicator: FC = () => { const [stalled, setStalled] = useState(false) const compacting = useStore($compactionActive) + const turnTimerKey = useActiveTurnTimerKey() // A pending clarify / approval / sudo / secret means the turn is paused on the // user, not working — so don't resurrect the "thinking" timer while they // decide (matches the pet's awaitingInput pose taking priority over busy). @@ -147,7 +157,7 @@ export const StreamStallIndicator: FC = () => { }, [activity]) const active = (stalled || compacting) && !awaitingInput - const elapsed = useElapsedSeconds(active) + const elapsed = useElapsedSeconds(active, compacting ? turnTimerKey : undefined) if (!active) { return null