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

@@ -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(<ActiveSession />)
+
+    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'
 

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 ''

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<WebSocketData> & { 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()

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<string, unknown>
+    : {}
+  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