From 9df36304230cc9c0722e9397b2224d0c059c839d Mon Sep 17 00:00:00 2001 From: Ramnique Singh <30795890+ramnique@users.noreply.github.com> Date: Thu, 9 Jul 2026 22:47:25 +0530 Subject: [PATCH] Show token usage in chats --- apps/x/apps/renderer/src/App.tsx | 19 ++++ .../renderer/src/components/chat-header.tsx | 14 +++ .../renderer/src/components/chat-sidebar.tsx | 21 ++++ .../src/components/compact-conversation.tsx | 15 +++ .../src/components/token-usage-menu.tsx | 96 +++++++++++++++++++ .../renderer/src/lib/chat-conversation.ts | 22 ++++- .../src/lib/session-chat/test-fixtures.ts | 3 +- .../src/lib/session-chat/turn-view.test.ts | 61 +++++++++++- .../src/lib/session-chat/turn-view.ts | 16 ++++ apps/x/apps/renderer/src/lib/token-usage.ts | 33 +++++++ 10 files changed, 297 insertions(+), 3 deletions(-) create mode 100644 apps/x/apps/renderer/src/components/token-usage-menu.tsx create mode 100644 apps/x/apps/renderer/src/lib/token-usage.ts diff --git a/apps/x/apps/renderer/src/App.tsx b/apps/x/apps/renderer/src/App.tsx index 2952b7d0..5aa4e411 100644 --- a/apps/x/apps/renderer/src/App.tsx +++ b/apps/x/apps/renderer/src/App.tsx @@ -114,6 +114,7 @@ import { isErrorMessage, isToolCall, isToolGroup, + isTurnUsageMessage, normalizeToolInput, normalizeToolOutput, parseAttachedFiles, @@ -133,6 +134,7 @@ import { useAnalyticsIdentity } from '@/hooks/useAnalyticsIdentity' import * as analytics from '@/lib/analytics' import { playAckCue } from '@/lib/call-sounds' import { useTheme } from '@/contexts/theme-context' +import { TokenUsageMenu } from '@/components/token-usage-menu' type DirEntry = z.infer type RunEventType = z.infer @@ -1596,6 +1598,7 @@ function App() { runId, conversation, currentAssistantMessage, + sessionUsage: {}, pendingAskHumanRequests: new Map(pendingAskHumanRequests), allPermissionRequests: new Map(allPermissionRequests), permissionResponses: new Map(permissionResponses), @@ -6021,6 +6024,19 @@ function App() { ) } + if (isTurnUsageMessage(item)) { + return ( +
+ +
+ ) + } + if (isErrorMessage(item)) { if (matchBillingError(item.message)) { return null @@ -6044,6 +6060,7 @@ function App() { ? { runId, ...sessionChat.chatState } : { runId, + sessionUsage: {}, conversation: sessionLoadErrorItems.length > 0 ? sessionLoadErrorItems : conversation, currentAssistantMessage, pendingAskHumanRequests, @@ -6206,6 +6223,7 @@ function App() { onNewChatTab={handleNewChatTab} recentRuns={runs} activeRunId={runId} + sessionUsage={activeChatTabState.sessionUsage} onSelectRun={(rid) => void navigateToView({ type: 'chat', runId: rid })} onOpenChatHistory={() => void navigateToView({ type: 'chat-history' })} /> @@ -6929,6 +6947,7 @@ function App() { onOpenFullScreen={toggleRightPaneMaximize} conversation={activeChatTabState.conversation} currentAssistantMessage={activeChatTabState.currentAssistantMessage} + sessionUsage={activeChatTabState.sessionUsage} chatTabStates={chatTabStatesForRender} viewportAnchors={chatViewportAnchorByTab} isProcessing={activeIsProcessing} diff --git a/apps/x/apps/renderer/src/components/chat-header.tsx b/apps/x/apps/renderer/src/components/chat-header.tsx index 0350f895..186ca5ef 100644 --- a/apps/x/apps/renderer/src/components/chat-header.tsx +++ b/apps/x/apps/renderer/src/components/chat-header.tsx @@ -3,6 +3,9 @@ import { ArrowUpRight, ChevronDown, MessageSquare, Plus } from 'lucide-react' import { Button } from '@/components/ui/button' import { cn } from '@/lib/utils' import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip' +import { TokenUsageMenu } from '@/components/token-usage-menu' +import type { TokenUsage } from '@/lib/chat-conversation' +import { hasTokenUsage } from '@/lib/token-usage' import { DropdownMenu, DropdownMenuContent, @@ -24,6 +27,7 @@ export interface ChatHeaderProps { onNewChatTab: () => void recentRuns?: ChatHeaderRecentRun[] activeRunId?: string | null + sessionUsage?: TokenUsage onSelectRun?: (runId: string) => void onOpenChatHistory?: () => void } @@ -40,10 +44,12 @@ export function ChatHeader({ onNewChatTab, recentRuns = [], activeRunId, + sessionUsage, onSelectRun, onOpenChatHistory, }: ChatHeaderProps) { const hasHistory = recentRuns.length > 0 || Boolean(onOpenChatHistory) + const showUsage = hasTokenUsage(sessionUsage) return ( <> @@ -95,6 +101,14 @@ export function ChatHeader({ {activeTitle} )} + {showUsage && ( + + )} + + + setDialogOpen(true)}> + + View token usage + + + + + + + + {title} + {totalText} + +
+ + + + + {modelCallCount !== undefined && modelCallCount > 0 && ( +
+ Model calls + {modelCallCount} +
+ )} +
+
+
+ + ) +} + +function UsageRow({ label, value }: { label: string; value: number | undefined }) { + if (!value) return null + return ( +
+ {label} + {formatTokenCount(value)} +
+ ) +} diff --git a/apps/x/apps/renderer/src/lib/chat-conversation.ts b/apps/x/apps/renderer/src/lib/chat-conversation.ts index ff3206a1..6397ded2 100644 --- a/apps/x/apps/renderer/src/lib/chat-conversation.ts +++ b/apps/x/apps/renderer/src/lib/chat-conversation.ts @@ -4,6 +4,14 @@ import { AskHumanRequestEvent, ToolPermissionAutoDecisionEvent, ToolPermissionRe import { COMPOSIO_DISPLAY_NAMES } from '@x/shared/src/composio.js' import type { CodeRunEvent, PermissionAsk } from '@x/shared/src/code-mode.js' +export interface TokenUsage { + inputTokens?: number + outputTokens?: number + totalTokens?: number + reasoningTokens?: number + cachedInputTokens?: number +} + export interface MessageAttachment { path: string filename: string @@ -45,13 +53,22 @@ export interface ErrorMessage { timestamp: number } -export type ConversationItem = ChatMessage | ToolCall | ErrorMessage +export interface TurnUsageMessage { + id: string + kind: 'turn-usage' + usage: TokenUsage + modelCallCount: number + timestamp: number +} + +export type ConversationItem = ChatMessage | ToolCall | ErrorMessage | TurnUsageMessage export type PermissionResponse = 'approve' | 'deny' export type ChatTabViewState = { runId: string | null conversation: ConversationItem[] currentAssistantMessage: string + sessionUsage: TokenUsage pendingAskHumanRequests: Map> allPermissionRequests: Map> permissionResponses: Map @@ -67,6 +84,7 @@ export const createEmptyChatTabViewState = (): ChatTabViewState => ({ runId: null, conversation: [], currentAssistantMessage: '', + sessionUsage: {}, pendingAskHumanRequests: new Map(), allPermissionRequests: new Map(), permissionResponses: new Map(), @@ -79,6 +97,8 @@ export const isChatMessage = (item: ConversationItem): item is ChatMessage => 'r export const isToolCall = (item: ConversationItem): item is ToolCall => 'name' in item export const isErrorMessage = (item: ConversationItem): item is ErrorMessage => 'kind' in item && item.kind === 'error' +export const isTurnUsageMessage = (item: ConversationItem): item is TurnUsageMessage => + 'kind' in item && item.kind === 'turn-usage' export const toToolState = (status: ToolCall['status']): ToolState => { switch (status) { diff --git a/apps/x/apps/renderer/src/lib/session-chat/test-fixtures.ts b/apps/x/apps/renderer/src/lib/session-chat/test-fixtures.ts index fa5bd107..a58a5a1c 100644 --- a/apps/x/apps/renderer/src/lib/session-chat/test-fixtures.ts +++ b/apps/x/apps/renderer/src/lib/session-chat/test-fixtures.ts @@ -67,6 +67,7 @@ export function completed( turnId: string, index: number, message: { role: 'assistant'; content: unknown }, + usage: Extract['usage'] = {}, ): TEvent { return { type: 'model_call_completed', @@ -75,7 +76,7 @@ export function completed( modelCallIndex: index, message: message as never, finishReason: 'stop', - usage: {}, + usage, } } diff --git a/apps/x/apps/renderer/src/lib/session-chat/turn-view.test.ts b/apps/x/apps/renderer/src/lib/session-chat/turn-view.test.ts index d4fc78a6..87e5da8b 100644 --- a/apps/x/apps/renderer/src/lib/session-chat/turn-view.test.ts +++ b/apps/x/apps/renderer/src/lib/session-chat/turn-view.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' import { reduceTurn } from '@x/shared/src/turns.js' -import { isChatMessage, isErrorMessage, isToolCall } from '@/lib/chat-conversation' +import { isChatMessage, isErrorMessage, isToolCall, isTurnUsageMessage } from '@/lib/chat-conversation' import { applyOverlay, buildSessionChatState, @@ -179,6 +179,34 @@ describe('buildTurnConversation', () => { expect(tool?.input).toEqual({ x: 1 }) }) + it('adds one turn usage row from accumulated model-call usage', () => { + const state = reduceTurn([ + created(T1, S1, user('run it')), + requested(T1, 0), + completed(T1, 0, assistantText('checking'), { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + }), + requested(T1, 1, ['assistant:0']), + completed(T1, 1, assistantText('done'), { + inputTokens: 12, + outputTokens: 3, + totalTokens: 15, + cachedInputTokens: 4, + }), + turnCompleted(T1, 'done'), + ]) + const usage = buildTurnConversation(state).find(isTurnUsageMessage) + expect(usage?.usage).toEqual({ + inputTokens: 22, + outputTokens: 8, + totalTokens: 30, + cachedInputTokens: 4, + }) + expect(usage?.modelCallCount).toBe(2) + }) + it('marks permission-pending tools pending and running tools running', () => { const state = reduceTurn([ created(T1, S1), @@ -388,6 +416,37 @@ describe('buildSessionChatState', () => { expect(state.isWaitingOnHuman).toBe(false) }) + it('aggregates usage across session turns', () => { + const first = reduceTurn([ + created('turn-1', S1, user('first?')), + requested('turn-1', 0), + completed('turn-1', 0, assistantText('first answer'), { + inputTokens: 20, + outputTokens: 5, + totalTokens: 25, + }), + turnCompleted('turn-1', 'first answer'), + ]) + const second = reduceTurn([ + created('turn-2', S1, user('second?')), + requested('turn-2', 0), + completed('turn-2', 0, assistantText('second answer'), { + inputTokens: 30, + outputTokens: 7, + totalTokens: 37, + reasoningTokens: 2, + }), + turnCompleted('turn-2', 'second answer'), + ]) + const state = buildSessionChatState([first, second], emptyOverlay()) + expect(state.sessionUsage).toEqual({ + inputTokens: 50, + outputTokens: 12, + totalTokens: 62, + reasoningTokens: 2, + }) + }) + it('is settled (not processing) when the latest turn is terminal', () => { const turn = reduceTurn(completedTurnLog(T1, S1, 'q', 'a')) const state = buildSessionChatState([turn], emptyOverlay()) diff --git a/apps/x/apps/renderer/src/lib/session-chat/turn-view.ts b/apps/x/apps/renderer/src/lib/session-chat/turn-view.ts index cfef1249..f5d46295 100644 --- a/apps/x/apps/renderer/src/lib/session-chat/turn-view.ts +++ b/apps/x/apps/renderer/src/lib/session-chat/turn-view.ts @@ -20,8 +20,10 @@ import type { ErrorMessage, MessageAttachment, PermissionResponse, + TokenUsage, ToolCall, } from '@/lib/chat-conversation' +import { addTokenUsage, hasTokenUsage } from '@/lib/token-usage' // Pure derivations from reduced turn state (+ the ephemeral live overlay) to // the view shapes the existing chat components already consume. No IPC, no @@ -337,6 +339,16 @@ export function buildTurnConversation(state: TurnState): ConversationItem[] { } satisfies ErrorMessage) } + if (hasTokenUsage(state.usage)) { + items.push({ + id: `${turnId}:usage`, + kind: 'turn-usage', + usage: state.usage, + modelCallCount: state.modelCalls.filter((call) => call.usage !== undefined).length, + timestamp: ts(), + }) + } + return items } @@ -349,6 +361,7 @@ type PermMeta = z.infer export type SessionChatState = { conversation: ConversationItem[] currentAssistantMessage: string + sessionUsage: TokenUsage // See LiveOverlay.voiceSegments. voiceSegments: string[] pendingAskHumanRequests: Map> @@ -410,8 +423,10 @@ export function buildSessionChatState( overlay: LiveOverlay, ): SessionChatState { const conversation: ConversationItem[] = [] + let sessionUsage: TokenUsage = {} for (const turn of turns) { conversation.push(...buildTurnConversation(turn)) + sessionUsage = addTokenUsage(sessionUsage, turn.usage) } for (let i = 0; i < conversation.length; i++) { const item = conversation[i] @@ -482,6 +497,7 @@ export function buildSessionChatState( return { conversation, currentAssistantMessage: stripVoiceTags(overlay.text), + sessionUsage, voiceSegments: overlay.voiceSegments, pendingAskHumanRequests, allPermissionRequests, diff --git a/apps/x/apps/renderer/src/lib/token-usage.ts b/apps/x/apps/renderer/src/lib/token-usage.ts new file mode 100644 index 00000000..0b616f4b --- /dev/null +++ b/apps/x/apps/renderer/src/lib/token-usage.ts @@ -0,0 +1,33 @@ +import type { TokenUsage } from '@/lib/chat-conversation' + +const USAGE_KEYS = [ + 'inputTokens', + 'outputTokens', + 'totalTokens', + 'reasoningTokens', + 'cachedInputTokens', +] as const + +export function addTokenUsage(total: TokenUsage, usage: TokenUsage): TokenUsage { + const next = { ...total } + for (const key of USAGE_KEYS) { + const value = usage[key] + if (value !== undefined) { + next[key] = (next[key] ?? 0) + value + } + } + return next +} + +export function hasTokenUsage(usage: TokenUsage | undefined): usage is TokenUsage { + return !!usage && USAGE_KEYS.some((key) => (usage[key] ?? 0) > 0) +} + +export function totalTokensOf(usage: TokenUsage): number { + return usage.totalTokens ?? (usage.inputTokens ?? 0) + (usage.outputTokens ?? 0) +} + +export function formatTokenCount(tokens: number | undefined): string { + if (tokens === undefined) return '-' + return new Intl.NumberFormat('en-US', { notation: 'compact' }).format(tokens) +}