diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index b075eaa0..17f1ba71 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -807,6 +807,49 @@ describe('chatStore history mapping', () => { }) }) + it('restores token usage from transcript history after reopening a session', async () => { + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: 'build the docs', + }, + { + id: 'assistant-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:01.000Z', + content: 'done', + usage: { input_tokens: 1200, output_tokens: 80 }, + }, + { + id: 'assistant-2', + type: 'assistant', + timestamp: '2026-04-06T00:00:02.000Z', + content: [{ type: 'text', text: 'follow-up done' }], + usage: { input_tokens: 3400, output_tokens: 120 }, + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [], + tokenUsage: { input_tokens: 0, output_tokens: 0 }, + }), + }, + }) + + await useChatStore.getState().loadHistory(TEST_SESSION_ID) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.tokenUsage).toEqual({ + input_tokens: 4600, + output_tokens: 200, + }) + }) + it('uses transcript terminal events to repair stale live goal and background task state', async () => { vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ messages: [ @@ -2954,6 +2997,40 @@ describe('chatStore history mapping', () => { expect(updateTabStatusMock).toHaveBeenLastCalledWith(TEST_SESSION_ID, 'running') }) + it('starts an elapsed timer when a reconnected session reports running status', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + chatState: 'idle', + elapsedSeconds: 0, + elapsedTimer: null, + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'status', + state: 'thinking', + verb: 'Thinking', + }) + + vi.advanceTimersByTime(2100) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.elapsedSeconds).toBe(2) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'status', + state: 'idle', + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.elapsedTimer).toBeNull() + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + it('tracks API retry status until the request finishes', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 68551830..07c28a5c 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -840,6 +840,37 @@ function mergeSlashCommandUpdates( return [...merged.values()] } +function readUsageToken(value: unknown): number { + return typeof value === 'number' && Number.isFinite(value) ? value : 0 +} + +function summarizeTokenUsageFromHistory(messages: MessageEntry[]): TokenUsage | null { + let inputTokens = 0 + let outputTokens = 0 + let cacheReadTokens = 0 + let cacheCreationTokens = 0 + + for (const message of messages) { + const usage = message.usage + if (!usage) continue + inputTokens += readUsageToken(usage.input_tokens) + outputTokens += readUsageToken(usage.output_tokens) + cacheReadTokens += readUsageToken(usage.cache_read_input_tokens) + cacheCreationTokens += readUsageToken(usage.cache_creation_input_tokens) + } + + if (inputTokens === 0 && outputTokens === 0 && cacheReadTokens === 0 && cacheCreationTokens === 0) { + return null + } + + return { + input_tokens: inputTokens, + output_tokens: outputTokens, + ...(cacheReadTokens > 0 ? { cache_read_tokens: cacheReadTokens } : {}), + ...(cacheCreationTokens > 0 ? { cache_creation_tokens: cacheCreationTokens } : {}), + } +} + async function fetchAndMapSessionHistory(sessionId: string) { const { messages, taskNotifications } = await sessionsApi.getMessages(sessionId) const uiMessages = mapHistoryMessagesToUiMessages(messages) @@ -855,6 +886,7 @@ async function fetchAndMapSessionHistory(sessionId: string) { restoredBackgroundTasks: backgroundTaskRecordFromNotifications(Object.values(restoredNotifications)), lastTodos: extractLastTodoWriteFromHistory(messages), hasMessagesAfterTaskCompletion: hasUserMessagesAfterTaskCompletion(messages), + tokenUsage: summarizeTokenUsageFromHistory(messages), } } @@ -1162,6 +1194,7 @@ export const useChatStore = create((set, get) => ({ restoredBackgroundTasks, lastTodos, hasMessagesAfterTaskCompletion, + tokenUsage, } = await fetchAndMapSessionHistory(sessionId) set((state) => { const session = state.sessions[sessionId] @@ -1176,6 +1209,7 @@ export const useChatStore = create((set, get) => ({ s.backgroundAgentTasks ?? {}, restoredBackgroundTasks, ), + tokenUsage: tokenUsage ?? s.tokenUsage, messages: mergeRestoredHistoryIntoLiveMessages( mergeBackgroundTaskMessages(s.messages, restoredBackgroundTasks), uiMessages, @@ -1192,6 +1226,7 @@ export const useChatStore = create((set, get) => ({ s.backgroundAgentTasks ?? {}, restoredBackgroundTasks, ), + tokenUsage: tokenUsage ?? s.tokenUsage, })) } }) if (lastTodos && lastTodos.length > 0) { @@ -1235,6 +1270,7 @@ export const useChatStore = create((set, get) => ({ restoredBackgroundTasks, lastTodos, hasMessagesAfterTaskCompletion, + tokenUsage, } = await fetchAndMapSessionHistory(sessionId) set((state) => { @@ -1249,6 +1285,7 @@ export const useChatStore = create((set, get) => ({ activeGoal, agentTaskNotifications: restoredNotifications, backgroundAgentTasks: restoredBackgroundTasks, + tokenUsage: tokenUsage ?? session.tokenUsage, chatState: 'idle', activeThinkingId: null, activeToolUseId: null, @@ -1500,6 +1537,19 @@ export const useChatStore = create((set, get) => ({ } : pendingText !== session.streamingText ? { streamingText: pendingText } : {}), } }) + if (msg.state !== 'idle') { + const session = get().sessions[sessionId] + if (session && !session.elapsedTimer) { + const timer = setInterval(() => { + set((st) => ({ + sessions: updateSessionIn(st.sessions, sessionId, (sess) => ({ + elapsedSeconds: sess.elapsedSeconds + 1, + })), + })) + }, 1000) + update(() => ({ elapsedTimer: timer })) + } + } if (msg.state === 'idle') { const session = get().sessions[sessionId] if (session?.elapsedTimer) { diff --git a/desktop/src/types/session.ts b/desktop/src/types/session.ts index 6c83366e..f4de90e9 100644 --- a/desktop/src/types/session.ts +++ b/desktop/src/types/session.ts @@ -13,6 +13,13 @@ export type SessionListItem = { permissionMode?: string } +export type MessageUsage = { + input_tokens?: number + output_tokens?: number + cache_read_input_tokens?: number + cache_creation_input_tokens?: number +} + export type MessageEntry = { id: string type: 'user' | 'assistant' | 'system' | 'tool_use' | 'tool_result' @@ -20,6 +27,7 @@ export type MessageEntry = { toolUseResult?: unknown timestamp: string model?: string + usage?: MessageUsage parentUuid?: string parentToolUseId?: string isSidechain?: boolean diff --git a/src/server/__tests__/sessions.test.ts b/src/server/__tests__/sessions.test.ts index 0aec97cd..901db6f4 100644 --- a/src/server/__tests__/sessions.test.ts +++ b/src/server/__tests__/sessions.test.ts @@ -225,7 +225,16 @@ function makeUserEntry(content: string, uuid?: string): Record } } -function makeAssistantEntry(content: string, parentUuid?: string): Record { +function makeAssistantEntry( + content: string, + parentUuid?: string, + usage?: { + input_tokens?: number + output_tokens?: number + cache_read_input_tokens?: number + cache_creation_input_tokens?: number + }, +): Record { return { parentUuid: parentUuid || null, isSidechain: false, @@ -236,6 +245,7 @@ function makeAssistantEntry(content: string, parentUuid?: string): Record { await writeSessionFile('-tmp-api-test', sessionId, [ makeSnapshotEntry(), makeUserEntry('Hello'), - makeAssistantEntry('World'), + makeAssistantEntry('World', undefined, { input_tokens: 1234, output_tokens: 56 }), makeUserEntry( '\nbg-1\ntoolu_bg\nfailed\nBackground command failed & stopped\nStack trace & failed assertion\nC:\\Temp\\bg.output\n', crypto.randomUUID(), @@ -2073,6 +2083,10 @@ describe('Sessions API', () => { taskNotifications: unknown[] } expect(body.messages).toHaveLength(2) + expect(body.messages[1]).toMatchObject({ + type: 'assistant', + usage: { input_tokens: 1234, output_tokens: 56 }, + }) expect(JSON.stringify(body.messages)).not.toContain('') expect(body.taskNotifications).toEqual([ { diff --git a/src/server/services/sessionService.ts b/src/server/services/sessionService.ts index 44799070..59f2f30f 100644 --- a/src/server/services/sessionService.ts +++ b/src/server/services/sessionService.ts @@ -94,6 +94,13 @@ export type TrimSessionResult = { removedMessageIds: string[] } +export type MessageUsage = { + input_tokens?: number + output_tokens?: number + cache_read_input_tokens?: number + cache_creation_input_tokens?: number +} + export type MessageEntry = { id: string type: 'user' | 'assistant' | 'system' | 'tool_use' | 'tool_result' @@ -101,6 +108,7 @@ export type MessageEntry = { toolUseResult?: unknown timestamp: string model?: string + usage?: MessageUsage parentUuid?: string parentToolUseId?: string isSidechain?: boolean @@ -225,6 +233,34 @@ type RawEntry = { [key: string]: unknown } +type RawMessageUsage = NonNullable['usage'] + +function normalizeMessageUsage(usage: RawMessageUsage): MessageUsage | undefined { + if (!usage) return undefined + + const normalized: MessageUsage = {} + if (typeof usage.input_tokens === 'number' && Number.isFinite(usage.input_tokens)) { + normalized.input_tokens = usage.input_tokens + } + if (typeof usage.output_tokens === 'number' && Number.isFinite(usage.output_tokens)) { + normalized.output_tokens = usage.output_tokens + } + if ( + typeof usage.cache_read_input_tokens === 'number' && + Number.isFinite(usage.cache_read_input_tokens) + ) { + normalized.cache_read_input_tokens = usage.cache_read_input_tokens + } + if ( + typeof usage.cache_creation_input_tokens === 'number' && + Number.isFinite(usage.cache_creation_input_tokens) + ) { + normalized.cache_creation_input_tokens = usage.cache_creation_input_tokens + } + + return Object.keys(normalized).length > 0 ? normalized : undefined +} + type PersistedWorktreeSession = { originalCwd: string worktreePath: string @@ -700,6 +736,8 @@ export class SessionService { type = 'system' } + const usage = normalizeMessageUsage(msg.usage) + return { id: entry.uuid || crypto.randomUUID(), type, @@ -707,6 +745,7 @@ export class SessionService { ...(entry.toolUseResult !== undefined ? { toolUseResult: entry.toolUseResult } : {}), timestamp: entry.timestamp || new Date().toISOString(), model: msg.model, + ...(usage ? { usage } : {}), parentUuid: entry.parentUuid ?? undefined, parentToolUseId, isSidechain: entry.isSidechain,