diff --git a/src/apps/desktop/src/runtime/session_application.rs b/src/apps/desktop/src/runtime/session_application.rs index 27bfbb11b..98a09d5dc 100644 --- a/src/apps/desktop/src/runtime/session_application.rs +++ b/src/apps/desktop/src/runtime/session_application.rs @@ -36,7 +36,7 @@ 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/store/FlowChatStore.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.ts index 625281c79..a522ee637 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.ts @@ -6303,6 +6303,9 @@ export class FlowChatStore { deepReviewRunManifest: metadata.deepReviewRunManifest, reviewTargetEvidence: metadata.reviewTargetEvidence, isTransient: false, + currentTokenUsage: !validatedAgentType.startsWith('acp:') && metadata.customMetadata?.lastRequestTokenUsage + ? { ...metadata.customMetadata.lastRequestTokenUsage, timestamp: Date.now() } + : undefined, }; const newSessions = new Map(prev.sessions); @@ -6678,6 +6681,9 @@ export class FlowChatStore { deepReviewRunManifest: metadata.deepReviewRunManifest, reviewTargetEvidence: metadata.reviewTargetEvidence, isTransient: false, + currentTokenUsage: !validatedAgentType.startsWith('acp:') && metadata.customMetadata?.lastRequestTokenUsage + ? { ...metadata.customMetadata.lastRequestTokenUsage, timestamp: Date.now() } + : undefined, }; const newSessions = new Map(prev.sessions); @@ -7306,6 +7312,20 @@ export class FlowChatStore { const session = prev.sessions.get(sessionId); if (!session) return prev; + // Fallback for sessions persisted before lastRequestTokenUsage metadata: + // backfill from the last completed turn only when it had a single model + // round, so the accumulated input-token sum equals one request. + let fallbackTokenUsage = session.currentTokenUsage; + if (!fallbackTokenUsage && !isAcpSession && dialogTurns.length > 0) { + for (let i = dialogTurns.length - 1; i >= 0; i--) { + const turn = dialogTurns[i]; + if (turn.status === 'completed' && turn.tokenUsage && turn.modelRounds.length === 1) { + fallbackTokenUsage = { ...turn.tokenUsage }; + break; + } + } + } + const updatedSession = { ...session, dialogTurns, @@ -7330,6 +7350,7 @@ export class FlowChatStore { lastUserDialogMode: restoredLastUserDialogMode, lastSubmittedMode: restoredSessionInfo?.lastSubmittedAgentType ?? session.lastSubmittedMode, + ...(fallbackTokenUsage ? { currentTokenUsage: fallbackTokenUsage } : {}), }; const newSessions = new Map(prev.sessions); diff --git a/src/web-ui/src/flow_chat/utils/sessionMetadata.test.ts b/src/web-ui/src/flow_chat/utils/sessionMetadata.test.ts index 2de140ccf..8bd5fc2cf 100644 --- a/src/web-ui/src/flow_chat/utils/sessionMetadata.test.ts +++ b/src/web-ui/src/flow_chat/utils/sessionMetadata.test.ts @@ -225,6 +225,73 @@ describe('sessionMetadata', () => { })).toBe(4321); }); + it('persists lastRequestTokenUsage from session currentTokenUsage', () => { + const session = createSession({ + currentTokenUsage: { + inputTokens: 12000, + outputTokens: 3400, + totalTokens: 15400, + timestamp: 9999, + }, + }); + + const metadata = buildSessionMetadata(session); + + expect(metadata.customMetadata?.lastRequestTokenUsage).toEqual({ + inputTokens: 12000, + outputTokens: 3400, + totalTokens: 15400, + }); + }); + + it('omits lastRequestTokenUsage when session has no currentTokenUsage', () => { + const session = createSession(); + + const metadata = buildSessionMetadata(session); + + expect(metadata.customMetadata?.lastRequestTokenUsage).toBeUndefined(); + }); + + it('replaces stale lastRequestTokenUsage from existing customMetadata', () => { + const session = createSession({ + currentTokenUsage: { + inputTokens: 500, + totalTokens: 600, + timestamp: 1, + }, + }); + + const existingMetadata: SessionMetadata = { + sessionId: 'session-1', + sessionName: 'Session Title', + agentType: 'agentic', + modelName: 'gpt-test', + createdAt: 1000, + lastActiveAt: 1000, + turnCount: 0, + messageCount: 0, + toolCallCount: 0, + status: 'active', + tags: [], + customMetadata: { + lastRequestTokenUsage: { + inputTokens: 999999, + totalTokens: 999999, + }, + }, + todos: [], + workspacePath: '/workspace', + }; + + const metadata = buildSessionMetadata(session, existingMetadata); + + expect(metadata.customMetadata?.lastRequestTokenUsage).toEqual({ + inputTokens: 500, + outputTokens: undefined, + totalTokens: 600, + }); + }); + it('persists locale-aware default title metadata before the first message', () => { const session = createSession({ title: 'flow-chat:session.newCodeWithIndex', diff --git a/src/web-ui/src/flow_chat/utils/sessionMetadata.ts b/src/web-ui/src/flow_chat/utils/sessionMetadata.ts index 642b8dab2..3f7ebd62e 100644 --- a/src/web-ui/src/flow_chat/utils/sessionMetadata.ts +++ b/src/web-ui/src/flow_chat/utils/sessionMetadata.ts @@ -230,6 +230,7 @@ function buildSessionCustomMetadata( | 'titleSource' | 'titleI18nKey' | 'titleI18nParams' + | 'currentTokenUsage' >, existingCustomMetadata?: SessionCustomMetadata ): SessionCustomMetadata { @@ -253,6 +254,14 @@ function buildSessionCustomMetadata( nextCustomMetadata.titleParams = session.titleI18nParams ?? null; } + if (session.currentTokenUsage) { + nextCustomMetadata.lastRequestTokenUsage = { + inputTokens: session.currentTokenUsage.inputTokens, + outputTokens: session.currentTokenUsage.outputTokens, + totalTokens: session.currentTokenUsage.totalTokens, + }; + } + return nextCustomMetadata; } @@ -362,6 +371,7 @@ export function buildSessionMetadata( | 'needsUserAttention' | 'deepReviewRunManifest' | 'reviewTargetEvidence' + | 'currentTokenUsage' >, existingMetadata?: SessionMetadata | null ): SessionMetadata { @@ -404,6 +414,7 @@ export function buildSessionMetadata( titleSource: session.titleSource, titleI18nKey: session.titleI18nKey, titleI18nParams: session.titleI18nParams, + currentTokenUsage: session.currentTokenUsage, }, existingMetadata?.customMetadata ), diff --git a/src/web-ui/src/shared/types/session-history.ts b/src/web-ui/src/shared/types/session-history.ts index d865f6108..0e7dc0351 100644 --- a/src/web-ui/src/shared/types/session-history.ts +++ b/src/web-ui/src/shared/types/session-history.ts @@ -41,6 +41,7 @@ export interface SessionCustomMetadata extends Record { titleSource?: SessionTitleSource | null; titleKey?: string | null; titleParams?: Record | null; + lastRequestTokenUsage?: { inputTokens: number; outputTokens?: number; totalTokens: number } | null; } export interface SessionMetadata {