From 3ba8b44a23e126316f3fa440bcd0f0e8543de9eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Tue, 23 Jun 2026 16:31:51 +0800 Subject: [PATCH] fix: preserve cache token usage in live sessions Map CLI result cache_read_input_tokens/cache_creation_input_tokens into the desktop websocket usage payload, and count cache read/write tokens in the active session header token badge. Fixes locally: #842 Tested: bun test src/server/__tests__/websocket-handler.test.ts Tested: cd desktop && bun run test -- src/pages/ActiveSession.test.tsx --run Tested: bun run check:server Tested: bun run check:desktop Not-tested: live Windows v0.4.3 reproduction; GitHub issues remain open until release and retest. Confidence: high Scope-risk: narrow --- desktop/src/pages/ActiveSession.test.tsx | 57 +++++++++++++++++++ desktop/src/pages/ActiveSession.tsx | 13 ++++- .../__tests__/websocket-handler.test.ts | 33 +++++++++++ src/server/ws/handler.ts | 26 +++++++-- 4 files changed, 122 insertions(+), 7 deletions(-) diff --git a/desktop/src/pages/ActiveSession.test.tsx b/desktop/src/pages/ActiveSession.test.tsx index 00a9a4cb..7ead7b0e 100644 --- a/desktop/src/pages/ActiveSession.test.tsx +++ b/desktop/src/pages/ActiveSession.test.tsx @@ -148,6 +148,63 @@ describe('ActiveSession task polling', () => { expect(screen.getByTestId('chat-input')).toHaveAttribute('data-variant', 'default') }) + it('shows the session token badge when usage is cache-only', () => { + const sessionId = 'cache-only-token-session' + + useSessionStore.setState({ + sessions: [{ + id: sessionId, + title: 'Cache Only Token Session', + createdAt: '2026-05-07T00:00:00.000Z', + modifiedAt: '2026-05-07T00:00:00.000Z', + messageCount: 1, + projectPath: '/workspace/project', + workDir: '/workspace/project', + workDirExists: true, + }], + activeSessionId: sessionId, + isLoading: false, + error: null, + }) + useTabStore.setState({ + tabs: [{ sessionId, title: 'Cache Only Token Session', type: 'session', status: 'idle' }], + activeTabId: sessionId, + }) + useChatStore.setState({ + sessions: { + [sessionId]: { + messages: [], + chatState: 'idle', + connectionState: 'connected', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + pendingPermission: null, + pendingComputerUsePermission: null, + tokenUsage: { + input_tokens: 0, + output_tokens: 0, + cache_read_tokens: 1200, + cache_creation_tokens: 300, + }, + streamingResponseChars: 0, + elapsedSeconds: 0, + statusVerb: '', + slashCommands: [], + agentTaskNotifications: {}, + elapsedTimer: null, + }, + }, + }) + + render() + + const tokenBadge = screen.getByTitle(/1,500/) + expect(tokenBadge).toHaveTextContent('1.5k') + }) + it('shows a loading state for historical sessions while messages are loading', () => { const sessionId = 'history-visible-loading-session' diff --git a/desktop/src/pages/ActiveSession.tsx b/desktop/src/pages/ActiveSession.tsx index 64f8e5ae..59fefdd3 100644 --- a/desktop/src/pages/ActiveSession.tsx +++ b/desktop/src/pages/ActiveSession.tsx @@ -29,7 +29,7 @@ import { WorkbenchPanel } from '../components/workbench/WorkbenchPanel' import { TeamStatusBar } from '../components/teams/TeamStatusBar' import { TerminalSettings } from './TerminalSettings' import type { SessionListItem } from '../types/session' -import type { ActiveGoalState } from '../types/chat' +import type { ActiveGoalState, TokenUsage } from '../types/chat' import { useMobileViewport } from '../hooks/useMobileViewport' import { isDesktopRuntime } from '../lib/desktopRuntime' import { formatTokenCount } from '../lib/formatTokenCount' @@ -52,6 +52,15 @@ function isSessionTabState(activeTabId: string | null, activeTabType: TabType | !activeTabId.startsWith(WORKBENCH_TAB_PREFIX) } +function getTokenUsageTotal(usage: TokenUsage): number { + return ( + usage.input_tokens + + usage.output_tokens + + (usage.cache_read_tokens ?? 0) + + (usage.cache_creation_tokens ?? 0) + ) +} + function getSessionTerminalCwd(session: SessionListItem | undefined) { if (!session) return undefined if (session.workDir && session.workDirExists !== false) return session.workDir @@ -355,7 +364,7 @@ export function ActiveSession() { const isActive = chatState !== 'idle' || (trackedTaskSessionId === activeTabId && hasRunningTasks) || hasRunningBackgroundTasks - const totalTokens = tokenUsage.input_tokens + tokenUsage.output_tokens + const totalTokens = getTokenUsageTotal(tokenUsage) const lastUpdated = useMemo(() => { if (!session?.modifiedAt) return '' diff --git a/src/server/__tests__/websocket-handler.test.ts b/src/server/__tests__/websocket-handler.test.ts index dcd20d0f..6d0fed5c 100644 --- a/src/server/__tests__/websocket-handler.test.ts +++ b/src/server/__tests__/websocket-handler.test.ts @@ -9,6 +9,7 @@ import { closeSessionConnection, getActiveSessionIds, handleWebSocket, + translateCliMessage, type WebSocketData, } from '../ws/handler.js' import { @@ -37,6 +38,38 @@ function makeClientSocket(sessionId: string) { } as unknown as ServerWebSocket & { sent: string[] } } +describe('translateCliMessage usage mapping', () => { + afterEach(() => { + __resetWebSocketHandlerStateForTests() + mock.restore() + }) + + it('keeps cache token counts on result completion events', () => { + const sessionId = `usage-${crypto.randomUUID()}` + + const messages = translateCliMessage({ + type: 'result', + subtype: 'success', + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 3456, + cache_creation_input_tokens: 789, + }, + }, sessionId) + + expect(messages).toEqual([{ + type: 'message_complete', + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_tokens: 3456, + cache_creation_tokens: 789, + }, + }]) + }) +}) + describe('WebSocket handler session isolation', () => { afterEach(() => { __resetWebSocketHandlerStateForTests() diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 26d10310..f4bf3fc8 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -7,7 +7,7 @@ */ import type { ServerWebSocket } from 'bun' -import type { ClientMessage, ServerMessage, StreamingFallbackCause } from './events.js' +import type { ClientMessage, ServerMessage, StreamingFallbackCause, TokenUsage } from './events.js' import * as os from 'node:os' import { ConversationStartupError, @@ -132,6 +132,25 @@ export function getSlashCommands(sessionId: string): SessionSlashCommand[] { return sessionSlashCommands.get(sessionId) || [] } +function usageNumber(value: unknown): number { + return typeof value === 'number' && Number.isFinite(value) ? value : 0 +} + +function translateCliUsage(usage: unknown): TokenUsage { + const record = usage && typeof usage === 'object' + ? usage as Record + : {} + const cacheReadTokens = usageNumber(record.cache_read_input_tokens ?? record.cache_read_tokens) + const cacheCreationTokens = usageNumber(record.cache_creation_input_tokens ?? record.cache_creation_tokens) + + return { + input_tokens: usageNumber(record.input_tokens), + output_tokens: usageNumber(record.output_tokens), + ...(cacheReadTokens > 0 ? { cache_read_tokens: cacheReadTokens } : {}), + ...(cacheCreationTokens > 0 ? { cache_creation_tokens: cacheCreationTokens } : {}), + } +} + export type WebSocketData = { sessionId: string connectedAt: number @@ -1691,10 +1710,7 @@ export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessa case 'result': { // 对话结果（成功或错误） - const usage = { - input_tokens: cliMsg.usage?.input_tokens || 0, - output_tokens: cliMsg.usage?.output_tokens || 0, - } + const usage = translateCliUsage(cliMsg.usage) if (cliMsg.is_error) { // If the user requested stop, this "error" is just the interrupt