From 10158ac16457affd793f5adb1d39058969e29de6 Mon Sep 17 00:00:00 2001 From: Chris Scott <> Date: Mon, 12 Oct 2026 01:12:55 +0000 Subject: [PATCH] fix(sessions): match OpenCode TUI context usage The header context count differed from the OpenCode TUI and dipped then jumped back up mid-session: - cache write tokens were left out of the total, so context stored as a cache write was dropped and then reappeared as a cache read on the next step - the newest assistant message without token data was skipped only one message back, so a step without usage showed 0 or a stale count The total now includes cache write tokens and the newest assistant message that reports tokens is counted, matching the TUI: messages before the latest completed compaction are ignored, messages at or after a revert point are ignored, and the indicator is hidden when no message qualifies or the total is zero. --- .../session/ContextUsageIndicator.tsx | 25 +- frontend/src/hooks/useContextUsage.test.tsx | 285 ++++++++++++++++++ frontend/src/hooks/useContextUsage.ts | 95 +++--- frontend/src/pages/SessionDetail.tsx | 1 + 4 files changed, 354 insertions(+), 52 deletions(-) create mode 100644 frontend/src/hooks/useContextUsage.test.tsx diff --git a/frontend/src/components/session/ContextUsageIndicator.tsx b/frontend/src/components/session/ContextUsageIndicator.tsx index 36bd4cfa9..c11bd1f96 100644 --- a/frontend/src/components/session/ContextUsageIndicator.tsx +++ b/frontend/src/components/session/ContextUsageIndicator.tsx @@ -5,6 +5,7 @@ interface ContextUsageIndicatorProps { sessionID: string | undefined isConnected: boolean isReconnecting?: boolean + revertMessageID?: string } const getUsageTextColor = (percentage: number) => { @@ -13,8 +14,14 @@ const getUsageTextColor = (percentage: number) => { return 'text-destructive' } -export function ContextUsageIndicator({ directory, sessionID, isConnected, isReconnecting }: ContextUsageIndicatorProps) { - const { totalTokens, contextLimit, usagePercentage, isLoading } = useContextUsage(sessionID, directory) +export function ContextUsageIndicator({ + directory, + sessionID, + isConnected, + isReconnecting, + revertMessageID, +}: ContextUsageIndicatorProps) { + const { usage, isLoading } = useContextUsage(sessionID, directory, revertMessageID) if (isLoading) { return ( @@ -32,15 +39,19 @@ export function ContextUsageIndicator({ directory, sessionID, isConnected, isRec return Disconnected } - const tokenText = contextLimit - ? `${totalTokens.toLocaleString()} (${Math.round(usagePercentage || 0)}%)` - : totalTokens.toLocaleString() + if (!usage) { + return null + } + + const tokenText = usage.usagePercentage === null + ? usage.totalTokens.toLocaleString() + : `${usage.totalTokens.toLocaleString()} (${usage.usagePercentage}%)` return (
- + {tokenText}
) -} \ No newline at end of file +} diff --git a/frontend/src/hooks/useContextUsage.test.tsx b/frontend/src/hooks/useContextUsage.test.tsx new file mode 100644 index 000000000..4935aaf13 --- /dev/null +++ b/frontend/src/hooks/useContextUsage.test.tsx @@ -0,0 +1,285 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { renderHook } from '@testing-library/react' +import type { ReactNode } from 'react' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { + SessionMessageAssistant, + SessionMessageCompaction, + SessionMessageInfo, +} from '@opencode-manager/shared/opencode' +import { useContextUsage } from './useContextUsage' +import { emptySessionTranscript } from '@/lib/session-projection' +import { sessionTranscriptQueryKey } from '@/lib/queryInvalidation' + +const mocks = vi.hoisted(() => ({ + useProviders: vi.fn(), +})) + +vi.mock('@/hooks/useProviders', () => ({ + useProviders: mocks.useProviders, +})) + +const SESSION_ID = 'session-1' +const DIRECTORY = '/repo' + +const model = { id: 'claude-sonnet', providerID: 'anthropic' } + +function assistantMessage(id: string, tokens: SessionMessageAssistant['tokens']): SessionMessageAssistant { + return { + id, + type: 'assistant', + agent: 'build', + model, + content: [], + time: { created: 1 }, + tokens, + } +} + +function compactionMessage(id: string, status: 'running' | 'completed'): SessionMessageCompaction { + const base = { + id, + type: 'compaction' as const, + reason: 'auto' as const, + summary: '', + recent: '', + time: { created: 1 }, + } + return status === 'running' + ? { ...base, status: 'running' } + : { ...base, status: 'completed' } +} + +function seedTranscript(queryClient: QueryClient, messages: SessionMessageInfo[]) { + queryClient.setQueryData(sessionTranscriptQueryKey(SESSION_ID), { + transcript: { ...emptySessionTranscript, messages }, + }) +} + +function createWrapper(queryClient: QueryClient) { + return function Wrapper({ children }: { children: ReactNode }) { + return {children} + } +} + +function createQueryClient() { + return new QueryClient({ defaultOptions: { queries: { retry: false } } }) +} + +describe('useContextUsage', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.useProviders.mockReturnValue({ data: { models: [] } }) + }) + + it('counts cache.write in the total token usage', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.totalTokens).toBe(1515) + }) + + it('rounds the usage percentage against the model context limit', () => { + mocks.useProviders.mockReturnValue({ + data: { models: [{ providerID: 'anthropic', id: 'claude-sonnet', limit: { context: 1000 } }] }, + }) + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.usagePercentage).toBe(152) + }) + + it('hides usage when the latest assistant has all-zero tokens even if an earlier one had usage', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + assistantMessage('assistant-2', { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage).toBeNull() + }) + + it('skips the latest assistant when tokens are undefined and uses the earlier assistant', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + assistantMessage('assistant-2', undefined), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.totalTokens).toBe(1515) + }) + + it('ignores assistants before the latest completed compaction', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + compactionMessage('compaction-1', 'completed'), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage).toBeNull() + }) + + it('uses an assistant that appears after the latest completed compaction', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + compactionMessage('compaction-1', 'completed'), + assistantMessage('assistant-2', { + input: 20, + output: 10, + reasoning: 0, + cache: { read: 100, write: 50 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.totalTokens).toBe(180) + }) + + it('does not treat a running compaction as a boundary', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + compactionMessage('compaction-1', 'running'), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.totalTokens).toBe(1515) + }) + + it('ignores assistants at or after the revert boundary', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + assistantMessage('assistant-2', { + input: 20, + output: 10, + reasoning: 0, + cache: { read: 100, write: 50 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY, 'assistant-2'), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage?.totalTokens).toBe(1515) + }) + + it('hides usage when the revert boundary excludes every assistant', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + assistantMessage('assistant-2', { + input: 20, + output: 10, + reasoning: 0, + cache: { read: 100, write: 50 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY, 'assistant-1'), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage).toBeNull() + }) + + it('hides usage when the revert boundary is unknown', () => { + const queryClient = createQueryClient() + seedTranscript(queryClient, [ + assistantMessage('assistant-1', { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 1000, write: 500 }, + }), + ]) + + const { result } = renderHook(() => useContextUsage(SESSION_ID, DIRECTORY, 'missing-message'), { + wrapper: createWrapper(queryClient), + }) + + expect(result.current.usage).toBeNull() + }) +}) diff --git a/frontend/src/hooks/useContextUsage.ts b/frontend/src/hooks/useContextUsage.ts index 519096fa3..a12ba5f1c 100644 --- a/frontend/src/hooks/useContextUsage.ts +++ b/frontend/src/hooks/useContextUsage.ts @@ -1,19 +1,42 @@ import { useMemo } from 'react' import { skipToken, useQuery } from '@tanstack/react-query' -import { formatOpenCodeModelRef, findModelInfo } from '@opencode-manager/shared/opencode' +import { findModelInfo } from '@opencode-manager/shared/opencode' +import type { SessionMessageInfo } from '@opencode-manager/shared/opencode' import { useProviders } from '@/hooks/useProviders' import { sessionTranscriptQueryKey } from '@/lib/queryInvalidation' import type { TranscriptCache } from '@/lib/session-projection' -interface ContextUsage { +export interface ContextUsage { totalTokens: number - contextLimit: number | null usagePercentage: number | null - currentModel: string | null +} + +interface ContextUsageResult { + usage: ContextUsage | null isLoading: boolean } -export const useContextUsage = (sessionID: string | undefined, directory?: string): ContextUsage => { +function lastCompletedCompactionIndex(messages: SessionMessageInfo[], end: number): number { + for (let index = end - 1; index >= 0; index -= 1) { + const message = messages[index] + if (message?.type === 'compaction' && message.status === 'completed') return index + } + return -1 +} + +function lastAssistantWithUsageIndex(messages: SessionMessageInfo[], start: number, end: number): number { + for (let index = end - 1; index > start; index -= 1) { + const message = messages[index] + if (message?.type === 'assistant' && message.tokens !== undefined) return index + } + return -1 +} + +export const useContextUsage = ( + sessionID: string | undefined, + directory?: string, + revertMessageID?: string, +): ContextUsageResult => { const transcriptQuery = useQuery({ queryKey: sessionTranscriptQueryKey(sessionID ?? ''), queryFn: skipToken, @@ -26,52 +49,34 @@ export const useContextUsage = (sessionID: string | undefined, directory?: strin return useMemo(() => { const messages = transcriptData?.transcript.messages ?? [] - const assistantMessages = messages.filter(msg => msg.type === 'assistant') || [] - let latestAssistantMessage = assistantMessages[assistantMessages.length - 1] - const sumTokens = (msg: typeof latestAssistantMessage) => { - if (msg?.type !== 'assistant' || !msg.tokens) return 0 - return msg.tokens.input + msg.tokens.output + msg.tokens.reasoning + msg.tokens.cache.read + const boundaryIndex = revertMessageID + ? messages.findIndex((message) => message.id === revertMessageID) + : -1 + if (revertMessageID && boundaryIndex === -1) { + return { usage: null, isLoading: messagesLoading } } - if (sumTokens(latestAssistantMessage) === 0 && assistantMessages.length > 1) { - latestAssistantMessage = assistantMessages[assistantMessages.length - 2] + const end = boundaryIndex === -1 ? messages.length : boundaryIndex + const compactionIndex = lastCompletedCompactionIndex(messages, end) + const assistantIndex = lastAssistantWithUsageIndex(messages, compactionIndex, end) + const assistant = assistantIndex === -1 ? undefined : messages[assistantIndex] + if (assistant?.type !== 'assistant' || assistant.tokens === undefined) { + return { usage: null, isLoading: messagesLoading } } - const activeModel = latestAssistantMessage?.type === 'assistant' ? latestAssistantMessage.model : undefined - - const currentModel = activeModel - ? formatOpenCodeModelRef(activeModel) - : null - - let contextLimit: number | null = null - if (activeModel && models) { - const model = findModelInfo(models, { providerID: activeModel.providerID, modelID: activeModel.id }) - if (model?.limit) { - contextLimit = model.limit.context - } - } - - if (messages.length === 0) { - return { - totalTokens: 0, - contextLimit, - usagePercentage: contextLimit ? 0 : null, - currentModel, - isLoading: messagesLoading - } + const { input, output, reasoning, cache } = assistant.tokens + const totalTokens = input + output + reasoning + cache.read + cache.write + if (totalTokens <= 0) { + return { usage: null, isLoading: messagesLoading } } - const totalTokens = sumTokens(latestAssistantMessage) - - const usagePercentage = contextLimit ? (totalTokens / contextLimit) * 100 : null + const model = models + ? findModelInfo(models, { providerID: assistant.model.providerID, modelID: assistant.model.id }) + : undefined + const contextLimit = model?.limit.context + const usagePercentage = contextLimit ? Math.round((totalTokens / contextLimit) * 100) : null - return { - totalTokens, - contextLimit, - usagePercentage, - currentModel, - isLoading: false - } - }, [transcriptData, messagesLoading, models]) + return { usage: { totalTokens, usagePercentage }, isLoading: messagesLoading } + }, [transcriptData, messagesLoading, models, revertMessageID]) } diff --git a/frontend/src/pages/SessionDetail.tsx b/frontend/src/pages/SessionDetail.tsx index 18cb7c94a..c32f9745f 100644 --- a/frontend/src/pages/SessionDetail.tsx +++ b/frontend/src/pages/SessionDetail.tsx @@ -726,6 +726,7 @@ export function SessionDetail() { sessionID={sessionId} isConnected={isConnected} isReconnecting={isReconnecting} + revertMessageID={revertMessageID} />