From 0bd28f96e07839385f33d0598c1d1332872c9a7b Mon Sep 17 00:00:00 2001 From: Tant Date: Thu, 6 Aug 2026 01:54:04 +0800 Subject: [PATCH 1/2] fix(flow-chat): restore context usage display after session hydration Startup or opening a historical session left session.currentTokenUsage undefined because the only writer was the TokenUsageUpdated event fired after a model response. The ModelSelector then hid the context percentage (tokenPercentage > 0 guard) and the tooltip omitted the last-request context line (current <= 0 guard). Backfill currentTokenUsage from the last completed dialog turn's persisted tokenUsage during hydrate commits in loadSessionHistory and refreshPeerSessionSnapshot. The backfill is idempotent (keeps any live value) and skipped for ACP sessions. The exact value is still overwritten by the next TokenUsageUpdated event. --- .../src/flow_chat/store/FlowChatStore.test.ts | 106 ++++++++++++++++++ .../src/flow_chat/store/FlowChatStore.ts | 18 +++ .../flow_chat/utils/tokenUsageDisplay.test.ts | 84 +++++++++++++- .../src/flow_chat/utils/tokenUsageDisplay.ts | 33 +++++- 4 files changed, 239 insertions(+), 2 deletions(-) diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts index 65756d8280..2c73a2717e 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts @@ -5227,4 +5227,110 @@ describe('FlowChatStore historical session hydration state', () => { }); expect(apiMocks.loadSessionTurnWindow).toHaveBeenCalledTimes(1); }); + + it('backfills currentTokenUsage from the last completed turn after hydration', async () => { + peerModeFlagMock.active = true; + apiMocks.restoreSessionView.mockResolvedValueOnce({ + session: { + sessionId: 'history-1', + sessionName: 'History 1', + agentType: 'agentic', + state: 'Idle', + turnCount: 2, + createdAt: 1, + }, + turns: [ + { + ...createPersistedTurn(0), + endTime: 2, + tokenUsage: { + inputTokens: 1000, + outputTokens: 100, + totalTokens: 1100, + timestamp: 2, + }, + }, + { + ...createPersistedTurn(1), + endTime: 4, + tokenUsage: { + inputTokens: 2400, + outputTokens: 300, + totalTokens: 2700, + timestamp: 4, + }, + }, + ], + contextRestoreState: 'ready', + }); + flowChatStore.setState(() => ({ + sessions: new Map([ + ['history-1', createSession({ + sessionId: 'history-1', + isHistorical: true, + historyState: 'metadata-only', + })], + ]), + activeSessionId: 'history-1', + })); + + await flowChatStore.loadSessionHistory('history-1', 'D:/workspace/BitFun'); + + expect(flowChatStore.getState().sessions.get('history-1')?.currentTokenUsage).toMatchObject({ + inputTokens: 2400, + outputTokens: 300, + totalTokens: 2700, + }); + }); + + it('keeps an existing currentTokenUsage when hydrating historical turns', async () => { + peerModeFlagMock.active = true; + apiMocks.restoreSessionView.mockResolvedValueOnce({ + session: { + sessionId: 'history-1', + sessionName: 'History 1', + agentType: 'agentic', + state: 'Idle', + turnCount: 1, + createdAt: 1, + }, + turns: [ + { + ...createPersistedTurn(0), + endTime: 2, + tokenUsage: { + inputTokens: 2400, + outputTokens: 300, + totalTokens: 2700, + timestamp: 2, + }, + }, + ], + contextRestoreState: 'ready', + }); + flowChatStore.setState(() => ({ + sessions: new Map([ + ['history-1', createSession({ + sessionId: 'history-1', + isHistorical: true, + historyState: 'metadata-only', + currentTokenUsage: { + inputTokens: 999, + outputTokens: 1, + totalTokens: 1000, + timestamp: 5, + }, + })], + ]), + activeSessionId: 'history-1', + })); + + await flowChatStore.loadSessionHistory('history-1', 'D:/workspace/BitFun'); + + expect(flowChatStore.getState().sessions.get('history-1')?.currentTokenUsage).toMatchObject({ + inputTokens: 999, + outputTokens: 1, + totalTokens: 1000, + }); + }); }); diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.ts index 625281c79e..f673c8eee9 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.ts @@ -56,6 +56,7 @@ import { } from '../utils/sessionMetadata'; import { sessionProjectWorkspacePath } from '../utils/sessionWorkspace'; import type { SessionTitleDescriptor } from '../utils/sessionTitle'; +import { deriveContextUsageFromTurns } from '../utils/tokenUsageDisplay'; import { deriveSessionTitleState, deriveSessionTitleStateFromMetadata, @@ -106,6 +107,13 @@ function firstNonEmptyString(...values: unknown[]): string | undefined { return undefined; } +function isAcpSessionForContextUsage(session: Session): boolean { + return Boolean( + session.mode?.startsWith('acp:') + || session.config.agentType?.startsWith('acp:'), + ); +} + function persistedSessionRemoteScope( metadata: { remoteConnectionId?: unknown; @@ -6878,6 +6886,11 @@ export class FlowChatStore { restored.session.lastUserDialogAgentType || session.lastUserDialogMode, lastSubmittedMode: restored.session.lastSubmittedAgentType ?? session.lastSubmittedMode, + currentTokenUsage: + session.currentTokenUsage + ?? (!isAcpSessionForContextUsage(session) + ? deriveContextUsageFromTurns(mergedTurns) + : undefined), }); applied = true; @@ -7330,6 +7343,11 @@ export class FlowChatStore { lastUserDialogMode: restoredLastUserDialogMode, lastSubmittedMode: restoredSessionInfo?.lastSubmittedAgentType ?? session.lastSubmittedMode, + currentTokenUsage: + session.currentTokenUsage + ?? (!isAcpSessionForContextUsage(session) + ? deriveContextUsageFromTurns(dialogTurns) + : undefined), }; const newSessions = new Map(prev.sessions); diff --git a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts index 73c2691e72..80726d0e4b 100644 --- a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts +++ b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts @@ -1,8 +1,9 @@ import { describe, expect, it } from 'vitest'; -import type { Session, TokenUsage } from '../types/flow-chat'; +import type { DialogTurn, Session, TokenUsage } from '../types/flow-chat'; import { buildContextUsageTooltip, buildModelRoundUsageMeta, + deriveContextUsageFromTurns, formatCompactTokenCount, getSessionContextUsageDisplay, } from './tokenUsageDisplay'; @@ -144,3 +145,84 @@ describe('tokenUsageDisplay', () => { expect(formatCompactTokenCount(4000)).toBe('4K'); }); }); + +describe('deriveContextUsageFromTurns', () => { + const makeTurn = ( + overrides: Partial & { + status: DialogTurn['status']; + tokenUsage?: TokenUsage; + }, + ): DialogTurn => ({ + id: 'turn-1', + sessionId: 'session-1', + userMessage: { + id: 'user-1', + content: 'hello', + timestamp: 1000, + }, + modelRounds: [], + status: 'completed', + startTime: 1000, + ...overrides, + }); + + const usage = (inputTokens: number): TokenUsage => ({ + inputTokens, + outputTokens: 100, + totalTokens: inputTokens + 100, + timestamp: 2000, + }); + + it('returns the last completed turn usage', () => { + const turns = [ + makeTurn({ id: 'turn-1', status: 'completed', tokenUsage: usage(1000) }), + makeTurn({ id: 'turn-2', status: 'completed', tokenUsage: usage(2000) }), + ]; + + expect(deriveContextUsageFromTurns(turns)).toEqual(usage(2000)); + }); + + it('skips unfinished turns and falls back to the last completed turn', () => { + const turns = [ + makeTurn({ id: 'turn-1', status: 'completed', tokenUsage: usage(1000) }), + makeTurn({ id: 'turn-2', status: 'processing', tokenUsage: usage(500) }), + makeTurn({ id: 'turn-3', status: 'pending', tokenUsage: usage(300) }), + ]; + + expect(deriveContextUsageFromTurns(turns)).toEqual(usage(1000)); + }); + + it('skips turns without usage and returns the last completed one that has it', () => { + const turns = [ + makeTurn({ id: 'turn-1', status: 'completed' }), + makeTurn({ id: 'turn-2', status: 'error', tokenUsage: usage(2500) }), + ]; + + expect(deriveContextUsageFromTurns(turns)).toEqual(usage(2500)); + }); + + it('skips completed turns with zero or invalid input tokens', () => { + const turns = [ + makeTurn({ + id: 'turn-1', + status: 'completed', + tokenUsage: { inputTokens: 0, totalTokens: 0, timestamp: 2000 }, + }), + makeTurn({ + id: 'turn-2', + status: 'cancelled', + tokenUsage: { inputTokens: 420, totalTokens: 500, timestamp: 3000 }, + }), + ]; + + expect(deriveContextUsageFromTurns(turns)).toMatchObject({ inputTokens: 420 }); + }); + + it('returns undefined for empty input or when no completed turn has usage', () => { + expect(deriveContextUsageFromTurns([])).toBeUndefined(); + expect(deriveContextUsageFromTurns(undefined)).toBeUndefined(); + expect(deriveContextUsageFromTurns([ + makeTurn({ id: 'turn-1', status: 'processing' }), + ])).toBeUndefined(); + }); +}); diff --git a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts index f77ef19b00..9f8ade61a5 100644 --- a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts +++ b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts @@ -1,4 +1,4 @@ -import type { Session, TokenUsage } from '../types/flow-chat'; +import type { DialogTurn, Session, TokenUsage } from '../types/flow-chat'; export const DEFAULT_MAX_CONTEXT_TOKENS = 128128; @@ -35,6 +35,37 @@ function formatCompactNumber(value: number): string { return Number.isInteger(value) ? String(value) : value.toFixed(1).replace(/\.0$/, ''); } +/** + * Derive the last completed turn's token usage as a context-usage approximation. + * + * Used to restore `session.currentTokenUsage` when a session is hydrated from + * persisted history (startup or opening a historical session): the exact value + * is only reported by the backend after the next model response. + */ +export function deriveContextUsageFromTurns(turns: DialogTurn[] | undefined): TokenUsage | undefined { + if (!turns) { + return undefined; + } + + for (let i = turns.length - 1; i >= 0; i--) { + const turn = turns[i]; + const usage = turn.tokenUsage; + if (!usage) { + continue; + } + if ( + turn.status === 'completed' + || turn.status === 'error' + || turn.status === 'cancelled' + ) { + if (typeof usage.inputTokens === 'number' && usage.inputTokens > 0) { + return usage; + } + } + } + return undefined; +} + export function getSessionContextUsageDisplay(session?: Session): ContextUsageDisplay { if (!session) { return { From 973f6a913260519d1d4e2d5088a884a995bce45e Mon Sep 17 00:00:00 2001 From: Tant Date: Thu, 6 Aug 2026 03:03:47 +0800 Subject: [PATCH 2/2] fix(flow-chat): persist exact last request usage for startup display The hydration backfill reused the last completed dialog turn's accumulated token usage, which sums input tokens across every model round of that turn. Long agentic sessions therefore showed an absurd context usage (e.g. 8.9M tokens) right after startup. Store the exact last request usage in session metadata (customMetadata.lastRequestTokenUsage, via the existing UI metadata whitelist) on every TokenUsageUpdated for non-ACP sessions, restore it during metadata hydration, and restrict the turn-based fallback to single-round turns where the accumulated value equals a single request. --- .../src/runtime/session_application.rs | 7 +- .../flow-chat-manager/EventHandlerModule.ts | 12 +++ .../flow-chat-manager/PersistenceModule.ts | 80 +++++++++++++++++ .../src/flow_chat/store/FlowChatStore.test.ts | 87 +++++++++++++++++++ .../src/flow_chat/store/FlowChatStore.ts | 26 ++++++ .../flow_chat/utils/tokenUsageDisplay.test.ts | 26 +++++- .../src/flow_chat/utils/tokenUsageDisplay.ts | 17 ++-- 7 files changed, 247 insertions(+), 8 deletions(-) diff --git a/src/apps/desktop/src/runtime/session_application.rs b/src/apps/desktop/src/runtime/session_application.rs index 27bfbb11bb..17997b7ac5 100644 --- a/src/apps/desktop/src/runtime/session_application.rs +++ b/src/apps/desktop/src/runtime/session_application.rs @@ -36,7 +36,12 @@ use bitfun_runtime_ports::{AgentContextReloadRequest, SessionTurnWindowRequest}; use serde::{Deserialize, Serialize}; use tokio::sync::RwLock; -const UI_CUSTOM_METADATA_KEYS: [&str; 3] = ["titleSource", "titleKey", "titleParams"]; +const UI_CUSTOM_METADATA_KEYS: [&str; 4] = [ + "titleSource", + "titleKey", + "titleParams", + "lastRequestTokenUsage", +]; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts index 972d59dac3..b3a22dc8c8 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts @@ -57,6 +57,7 @@ const pendingImageAnalysisTurns = new Map(); import { debouncedSaveDialogTurn, immediateSaveDialogTurn, + persistLastRequestTokenUsage, saveDialogTurnToDisk, cleanupSaveState, } from './PersistenceModule'; @@ -2136,6 +2137,17 @@ function handleTokenUsageUpdate(context: FlowChatContext, event: any): void { totalTokens }, turnId); + // Persist the exact last request usage so the context display survives a + // restart. Skip ACP sessions: their display is driven by + // currentAcpContextUsage instead. + if (!session.mode?.startsWith('acp:') && !session.config.agentType?.startsWith('acp:')) { + persistLastRequestTokenUsage(context, sessionId, { + inputTokens, + outputTokens: typeof outputTokens === 'number' ? outputTokens : undefined, + totalTokens, + }); + } + if (maxContextTokens !== undefined && maxContextTokens !== null) { store.updateSessionMaxContextTokens(sessionId, maxContextTokens); } diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts index fc6bffd29e..daac1d78e4 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts @@ -604,3 +604,83 @@ export async function touchSessionActivity( log.debug('Failed to touch session activity', { sessionId, error }); } } + +const lastRequestTokenUsageDebouncers = new Map< + string, + ReturnType +>(); + +/** + * Persist the exact last model request usage into session metadata so the + * input-box context display can be restored after an app restart. + * + * Only the session-level last request value is stored; the dialog turn usage + * stays accumulated per turn and must not be reused as a single-request + * approximation. The write is trailing-throttled because agentic sessions + * can emit one TokenUsageUpdated per model round. + */ +export function persistLastRequestTokenUsage( + context: FlowChatContext, + sessionId: string, + usage: { inputTokens: number; outputTokens?: number; totalTokens: number }, +): void { + const existingTimer = lastRequestTokenUsageDebouncers.get(sessionId); + if (existingTimer) { + clearTimeout(existingTimer); + } + const timer = setTimeout(() => { + lastRequestTokenUsageDebouncers.delete(sessionId); + void persistLastRequestTokenUsageNow(context, sessionId, usage).catch(error => { + log.warn('Failed to persist last request token usage', { sessionId, error }); + }); + }, COALESCED_IMMEDIATE_SAVE_DELAY_MS); + lastRequestTokenUsageDebouncers.set(sessionId, timer); +} + +async function persistLastRequestTokenUsageNow( + context: FlowChatContext, + sessionId: string, + usage: { inputTokens: number; outputTokens?: number; totalTokens: number }, +): Promise { + const { sessionAPI } = await import('@/infrastructure/api/service-api/SessionAPI'); + + const session = context.flowChatStore.getState().sessions.get(sessionId); + if (!session) return; + if (isTransientSession(session) || isObserverOnlyDispatchSession(sessionId, session)) return; + + const workspacePath = requireSessionProjectWorkspacePath(session, sessionId); + + let existingMetadata: any = null; + try { + existingMetadata = await sessionAPI.loadSessionMetadata( + sessionId, + workspacePath, + session.remoteConnectionId, + session.remoteSshHost + ); + } catch { + // Metadata may not exist yet for a fresh session; the patch below still works. + } + + const metadata = { + ...existingMetadata, + sessionId, + customMetadata: { + ...(existingMetadata?.customMetadata ?? {}), + lastRequestTokenUsage: { + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + totalTokens: usage.totalTokens, + timestamp: Date.now(), + }, + }, + }; + + await sessionAPI.saveSessionMetadata( + metadata, + workspacePath, + ['titleMetadata'], + session.remoteConnectionId, + session.remoteSshHost + ); +} diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts index 2c73a2717e..1e45a2ff51 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts @@ -5242,6 +5242,17 @@ describe('FlowChatStore historical session hydration state', () => { turns: [ { ...createPersistedTurn(0), + modelRounds: [{ + id: 'round-0', + turnId: 'turn-0', + roundIndex: 0, + timestamp: 1, + textItems: [], + toolItems: [], + thinkingItems: [], + startTime: 1, + status: 'completed', + }], endTime: 2, tokenUsage: { inputTokens: 1000, @@ -5252,6 +5263,17 @@ describe('FlowChatStore historical session hydration state', () => { }, { ...createPersistedTurn(1), + modelRounds: [{ + id: 'round-1', + turnId: 'turn-1', + roundIndex: 0, + timestamp: 3, + textItems: [], + toolItems: [], + thinkingItems: [], + startTime: 3, + status: 'completed', + }], endTime: 4, tokenUsage: { inputTokens: 2400, @@ -5297,6 +5319,17 @@ describe('FlowChatStore historical session hydration state', () => { turns: [ { ...createPersistedTurn(0), + modelRounds: [{ + id: 'round-0', + turnId: 'turn-0', + roundIndex: 0, + timestamp: 1, + textItems: [], + toolItems: [], + thinkingItems: [], + startTime: 1, + status: 'completed', + }], endTime: 2, tokenUsage: { inputTokens: 2400, @@ -5333,4 +5366,58 @@ describe('FlowChatStore historical session hydration state', () => { totalTokens: 1000, }); }); + + it('restores the exact last request token usage from persisted metadata', async () => { + apiMocks.listSessions.mockResolvedValueOnce([ + { + sessionId: 'history-1', + title: 'Saved session', + agentType: 'agentic', + modelName: 'auto', + createdAt: 10, + lastActiveAt: 20, + customMetadata: { + lastRequestTokenUsage: { + inputTokens: 42000, + outputTokens: 1500, + totalTokens: 43500, + timestamp: 21, + }, + }, + }, + ]); + + await flowChatStore.initializeFromDisk('D:/workspace/BitFun'); + + expect(flowChatStore.getState().sessions.get('history-1')?.currentTokenUsage).toMatchObject({ + inputTokens: 42000, + outputTokens: 1500, + totalTokens: 43500, + }); + }); + + it('ignores invalid persisted last request token usage', async () => { + apiMocks.listSessions.mockResolvedValueOnce([ + { + sessionId: 'history-1', + title: 'Saved session', + agentType: 'agentic', + modelName: 'auto', + createdAt: 10, + lastActiveAt: 20, + customMetadata: { + lastRequestTokenUsage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + timestamp: 21, + }, + }, + }, + ]); + + await flowChatStore.initializeFromDisk('D:/workspace/BitFun'); + + expect(flowChatStore.getState().sessions.get('history-1')?.currentTokenUsage).toBeUndefined(); + }); }); diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.ts index f673c8eee9..bf3f9108b8 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.ts @@ -114,6 +114,24 @@ function isAcpSessionForContextUsage(session: Session): boolean { ); } +function deriveRestoredCurrentTokenUsage(value: unknown): TokenUsage | undefined { + if (!value || typeof value !== 'object') { + return undefined; + } + const record = value as Record; + const inputTokens = record.inputTokens; + if (typeof inputTokens !== 'number' || !Number.isFinite(inputTokens) || inputTokens <= 0) { + return undefined; + } + const totalTokens = record.totalTokens; + return { + inputTokens, + outputTokens: typeof record.outputTokens === 'number' ? record.outputTokens : undefined, + totalTokens: typeof totalTokens === 'number' && Number.isFinite(totalTokens) ? totalTokens : inputTokens, + timestamp: typeof record.timestamp === 'number' ? record.timestamp : Date.now(), + }; +} + function persistedSessionRemoteScope( metadata: { remoteConnectionId?: unknown; @@ -6251,6 +6269,9 @@ export class FlowChatStore { remoteConnectionId, remoteSshHost, ); + const restoredCurrentTokenUsage = deriveRestoredCurrentTokenUsage( + metadata.customMetadata?.lastRequestTokenUsage, + ); this.setState(prev => { if (surfaceGeneration !== this.surfaceGeneration) { @@ -6292,6 +6313,7 @@ export class FlowChatStore { historyState: 'metadata-only', todos: (metadata as any).todos || [], maxContextTokens, + currentTokenUsage: restoredCurrentTokenUsage, mode: validatedAgentType, lastUserDialogMode: metadata.lastUserDialogAgentType, lastSubmittedMode: metadata.lastSubmittedAgentType, @@ -6629,6 +6651,9 @@ export class FlowChatStore { remoteConnectionId, remoteSshHost, ); + const restoredCurrentTokenUsage = deriveRestoredCurrentTokenUsage( + metadata.customMetadata?.lastRequestTokenUsage, + ); this.setState(prev => { if (prev.sessions.has(metadata.sessionId)) { @@ -6667,6 +6692,7 @@ export class FlowChatStore { historyState: 'metadata-only', todos: (metadata as any).todos || [], maxContextTokens, + currentTokenUsage: restoredCurrentTokenUsage, mode: validatedAgentType, lastUserDialogMode: metadata.lastUserDialogAgentType, lastSubmittedMode: metadata.lastSubmittedAgentType, diff --git a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts index 80726d0e4b..7c50c66f06 100644 --- a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts +++ b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.test.ts @@ -160,7 +160,7 @@ describe('deriveContextUsageFromTurns', () => { content: 'hello', timestamp: 1000, }, - modelRounds: [], + modelRounds: [{ id: 'round-1' }], status: 'completed', startTime: 1000, ...overrides, @@ -218,11 +218,33 @@ describe('deriveContextUsageFromTurns', () => { expect(deriveContextUsageFromTurns(turns)).toMatchObject({ inputTokens: 420 }); }); - it('returns undefined for empty input or when no completed turn has usage', () => { + it('skips multi-round turns because accumulated usage would overestimate context', () => { + const turns = [ + makeTurn({ id: 'turn-1', status: 'completed', tokenUsage: usage(1000) }), + makeTurn({ + id: 'turn-2', + status: 'completed', + modelRounds: [{ id: 'round-1' }, { id: 'round-2' }], + tokenUsage: usage(8_900_000), + }), + ]; + + expect(deriveContextUsageFromTurns(turns)).toEqual(usage(1000)); + }); + + it('returns undefined for empty input or when no completed single-round turn has usage', () => { expect(deriveContextUsageFromTurns([])).toBeUndefined(); expect(deriveContextUsageFromTurns(undefined)).toBeUndefined(); expect(deriveContextUsageFromTurns([ makeTurn({ id: 'turn-1', status: 'processing' }), ])).toBeUndefined(); + expect(deriveContextUsageFromTurns([ + makeTurn({ + id: 'turn-1', + status: 'completed', + modelRounds: [{ id: 'round-1' }, { id: 'round-2' }], + tokenUsage: usage(9000), + }), + ])).toBeUndefined(); }); }); diff --git a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts index 9f8ade61a5..3f4b9961f2 100644 --- a/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts +++ b/src/web-ui/src/flow_chat/utils/tokenUsageDisplay.ts @@ -36,11 +36,14 @@ function formatCompactNumber(value: number): string { } /** - * Derive the last completed turn's token usage as a context-usage approximation. + * Derive the last completed single-round turn's token usage as a + * context-usage approximation. * - * Used to restore `session.currentTokenUsage` when a session is hydrated from - * persisted history (startup or opening a historical session): the exact value - * is only reported by the backend after the next model response. + * Used as a fallback to restore `session.currentTokenUsage` when a session is + * hydrated from persisted history and no exact last-request usage was stored + * in session metadata. Only single-round turns are used: dialog turn usage + * accumulates across model rounds, so a multi-round turn's input total would + * badly overestimate the current context. */ export function deriveContextUsageFromTurns(turns: DialogTurn[] | undefined): TokenUsage | undefined { if (!turns) { @@ -58,7 +61,11 @@ export function deriveContextUsageFromTurns(turns: DialogTurn[] | undefined): To || turn.status === 'error' || turn.status === 'cancelled' ) { - if (typeof usage.inputTokens === 'number' && usage.inputTokens > 0) { + if ( + turn.modelRounds.length === 1 + && typeof usage.inputTokens === 'number' + && usage.inputTokens > 0 + ) { return usage; } }