Show token usage in chats

This commit is contained in:
Ramnique Singh 2026-07-09 22:47:25 +05:30
parent 1d16d16f15
commit 9df3630423
10 changed files with 297 additions and 3 deletions

View file

@ -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<typeof workspace.DirEntry>
type RunEventType = z.infer<typeof RunEvent>
@ -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 (
<div key={item.id} className="-mt-6 flex justify-start px-1" data-message-id={item.id}>
<TokenUsageMenu
usage={item.usage}
scope="turn"
modelCallCount={item.modelCallCount}
align="start"
/>
</div>
)
}
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}

View file

@ -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({
<span className="truncate">{activeTitle}</span>
</div>
)}
{showUsage && (
<TokenUsageMenu
usage={sessionUsage}
scope="session"
className="titlebar-no-drag my-1 shrink-0"
align="end"
/>
)}
<Tooltip>
<TooltipTrigger asChild>
<Button

View file

@ -47,6 +47,7 @@ import {
type ChatTabViewState,
type ConversationItem,
type PermissionResponse,
type TokenUsage,
createEmptyChatTabViewState,
getWebSearchCardData,
getComposioConnectCardData,
@ -57,12 +58,14 @@ import {
isErrorMessage,
isToolCall,
isToolGroup,
isTurnUsageMessage,
normalizeToolInput,
normalizeToolOutput,
parseAttachedFiles,
toToolState,
} from '@/lib/chat-conversation'
import { matchBillingError } from '@/lib/billing-error'
import { TokenUsageMenu } from '@/components/token-usage-menu'
const streamdownComponents = { pre: MarkdownPreOverride }
@ -140,6 +143,7 @@ interface ChatSidebarProps {
onOpenFullScreen?: () => void
conversation: ConversationItem[]
currentAssistantMessage: string
sessionUsage?: TokenUsage
chatTabStates?: Record<string, ChatTabViewState>
viewportAnchors?: Record<string, ChatViewportAnchorState>
isProcessing: boolean
@ -210,6 +214,7 @@ export function ChatSidebar({
onOpenFullScreen,
conversation,
currentAssistantMessage,
sessionUsage = {},
chatTabStates = {},
viewportAnchors = {},
isProcessing,
@ -350,6 +355,7 @@ export function ChatSidebar({
runId: runId ?? null,
conversation,
currentAssistantMessage,
sessionUsage,
pendingAskHumanRequests,
allPermissionRequests,
permissionResponses,
@ -358,6 +364,7 @@ export function ChatSidebar({
runId,
conversation,
currentAssistantMessage,
sessionUsage,
pendingAskHumanRequests,
allPermissionRequests,
permissionResponses,
@ -508,6 +515,19 @@ export function ChatSidebar({
)
}
if (isTurnUsageMessage(item)) {
return (
<div key={item.id} className="-mt-6 flex justify-start px-1" data-message-id={item.id}>
<TokenUsageMenu
usage={item.usage}
scope="turn"
modelCallCount={item.modelCallCount}
align="start"
/>
</div>
)
}
if (isErrorMessage(item)) {
if (matchBillingError(item.message)) {
return null
@ -600,6 +620,7 @@ export function ChatSidebar({
onNewChatTab={onNewChatTab}
recentRuns={recentRuns}
activeRunId={runId}
sessionUsage={activeTabState.sessionUsage}
onSelectRun={onSelectRun}
onOpenChatHistory={onOpenChatHistory}
/>

View file

@ -6,12 +6,14 @@ import {
isChatMessage,
isErrorMessage,
isToolCall,
isTurnUsageMessage,
getToolDisplayName,
getToolErrorText,
toToolState,
normalizeToolOutput,
} from '@/lib/chat-conversation'
import { Tool, ToolHeader, ToolContent, ToolTabbedContent } from '@/components/ai-elements/tool'
import { TokenUsageMenu } from '@/components/token-usage-menu'
/**
* Compact rendering of a run's conversation log — used by the live-note panel's
@ -34,6 +36,19 @@ export function CompactConversation({ items }: { items: ConversationItem[] }) {
</div>
)
}
if (isTurnUsageMessage(item)) {
return (
<div key={item.id} className="flex justify-start px-1">
<TokenUsageMenu
usage={item.usage}
scope="turn"
modelCallCount={item.modelCallCount}
className="size-5 border-transparent bg-transparent hover:bg-transparent"
align="start"
/>
</div>
)
}
if (isToolCall(item)) return <CompactToolRow key={item.id} tool={item} />
if (isChatMessage(item)) {
const isUser = item.role === 'user'

View file

@ -0,0 +1,96 @@
import { useState } from 'react'
import { BarChart3, MoreHorizontal } from 'lucide-react'
import {
Dialog,
DialogContent,
DialogDescription,
DialogHeader,
DialogTitle,
} from '@/components/ui/dialog'
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuTrigger,
} from '@/components/ui/dropdown-menu'
import { cn } from '@/lib/utils'
import type { TokenUsage } from '@/lib/chat-conversation'
import { formatTokenCount, totalTokensOf } from '@/lib/token-usage'
type TokenUsageMenuProps = {
usage: TokenUsage
scope: 'turn' | 'session'
modelCallCount?: number
className?: string
align?: 'start' | 'center' | 'end'
}
export function TokenUsageMenu({
usage,
scope,
modelCallCount,
className,
align = 'center',
}: TokenUsageMenuProps) {
const [dialogOpen, setDialogOpen] = useState(false)
const total = totalTokensOf(usage)
const totalText = `${formatTokenCount(total)} tokens`
const title = `Token usage for this ${scope}`
return (
<>
<DropdownMenu>
<DropdownMenuTrigger asChild>
<button
type="button"
className={cn(
'inline-flex size-6 items-center justify-center rounded-md border border-border/60 bg-background text-muted-foreground transition-colors hover:bg-accent hover:text-foreground',
className,
)}
aria-label={`${title} options`}
>
<MoreHorizontal className="size-3.5" strokeWidth={1.8} />
</button>
</DropdownMenuTrigger>
<DropdownMenuContent align={align} className="w-44">
<DropdownMenuItem onSelect={() => setDialogOpen(true)}>
<BarChart3 className="size-4" />
View token usage
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
<Dialog open={dialogOpen} onOpenChange={setDialogOpen}>
<DialogContent className="sm:max-w-sm">
<DialogHeader>
<DialogTitle>{title}</DialogTitle>
<DialogDescription>{totalText}</DialogDescription>
</DialogHeader>
<div className="space-y-2">
<UsageRow label="Input tokens" value={usage.inputTokens} />
<UsageRow label="Output tokens" value={usage.outputTokens} />
<UsageRow label="Cached input tokens" value={usage.cachedInputTokens} />
<UsageRow label="Reasoning tokens" value={usage.reasoningTokens} />
{modelCallCount !== undefined && modelCallCount > 0 && (
<div className="flex items-center justify-between gap-3 border-t border-border pt-2 text-sm">
<span className="text-muted-foreground">Model calls</span>
<span className="tabular-nums">{modelCallCount}</span>
</div>
)}
</div>
</DialogContent>
</Dialog>
</>
)
}
function UsageRow({ label, value }: { label: string; value: number | undefined }) {
if (!value) return null
return (
<div className="flex items-center justify-between gap-3 text-sm">
<span className="text-muted-foreground">{label}</span>
<span className="tabular-nums">{formatTokenCount(value)}</span>
</div>
)
}

View file

@ -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<string, z.infer<typeof AskHumanRequestEvent>>
allPermissionRequests: Map<string, z.infer<typeof ToolPermissionRequestEvent>>
permissionResponses: Map<string, PermissionResponse>
@ -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) {

View file

@ -67,6 +67,7 @@ export function completed(
turnId: string,
index: number,
message: { role: 'assistant'; content: unknown },
usage: Extract<TEvent, { type: 'model_call_completed' }>['usage'] = {},
): TEvent {
return {
type: 'model_call_completed',
@ -75,7 +76,7 @@ export function completed(
modelCallIndex: index,
message: message as never,
finishReason: 'stop',
usage: {},
usage,
}
}

View file

@ -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())

View file

@ -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<typeof ToolPermissionMetadata>
export type SessionChatState = {
conversation: ConversationItem[]
currentAssistantMessage: string
sessionUsage: TokenUsage
// See LiveOverlay.voiceSegments.
voiceSegments: string[]
pendingAskHumanRequests: Map<string, z.infer<typeof AskHumanRequestEvent>>
@ -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,

View file

@ -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)
}