diff --git a/src/apps/desktop/src/runtime/session_application.rs b/src/apps/desktop/src/runtime/session_application.rs index 27bfbb11b..17997b7ac 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 972d59dac..b3a22dc8c 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 fc6bffd29..daac1d78e 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 65756d828..1e45a2ff5 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,197 @@ 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), + modelRounds: [{ + id: 'round-0', + turnId: 'turn-0', + roundIndex: 0, + timestamp: 1, + textItems: [], + toolItems: [], + thinkingItems: [], + startTime: 1, + status: 'completed', + }], + endTime: 2, + tokenUsage: { + inputTokens: 1000, + outputTokens: 100, + totalTokens: 1100, + timestamp: 2, + }, + }, + { + ...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, + 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), + modelRounds: [{ + id: 'round-0', + turnId: 'turn-0', + roundIndex: 0, + timestamp: 1, + textItems: [], + toolItems: [], + thinkingItems: [], + startTime: 1, + status: 'completed', + }], + 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, + }); + }); + + 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 625281c79..bf3f9108b 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,31 @@ 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 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; @@ -6243,6 +6269,9 @@ export class FlowChatStore { remoteConnectionId, remoteSshHost, ); + const restoredCurrentTokenUsage = deriveRestoredCurrentTokenUsage( + metadata.customMetadata?.lastRequestTokenUsage, + ); this.setState(prev => { if (surfaceGeneration !== this.surfaceGeneration) { @@ -6284,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, @@ -6621,6 +6651,9 @@ export class FlowChatStore { remoteConnectionId, remoteSshHost, ); + const restoredCurrentTokenUsage = deriveRestoredCurrentTokenUsage( + metadata.customMetadata?.lastRequestTokenUsage, + ); this.setState(prev => { if (prev.sessions.has(metadata.sessionId)) { @@ -6659,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, @@ -6878,6 +6912,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 +7369,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 73c2691e7..7c50c66f0 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,106 @@ 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: [{ id: 'round-1' }], + 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('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 f77ef19b0..3f4b9961f 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,44 @@ function formatCompactNumber(value: number): string { return Number.isInteger(value) ? String(value) : value.toFixed(1).replace(/\.0$/, ''); } +/** + * Derive the last completed single-round turn's token usage as a + * context-usage approximation. + * + * 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) { + 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 ( + turn.modelRounds.length === 1 + && typeof usage.inputTokens === 'number' + && usage.inputTokens > 0 + ) { + return usage; + } + } + } + return undefined; +} + export function getSessionContextUsageDisplay(session?: Session): ContextUsageDisplay { if (!session) { return {