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