fix(web-ui): persist and hydrate lastRequestTokenUsage for context display

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")]

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);

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',

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
     ),

src/web-ui/src/shared/types/session-history.ts

@@ -41,6 +41,7 @@ export interface SessionCustomMetadata extends Record<string, unknown> {
   titleSource?: SessionTitleSource | null;
   titleKey?: string | null;
   titleParams?: Record<string, unknown> | null;
+  lastRequestTokenUsage?: { inputTokens: number; outputTokens?: number; totalTokens: number } | null;
 }
 
 export interface SessionMetadata {