fix(flow-chat): persist exact last request usage for startup display

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

src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts

@@ -57,6 +57,7 @@ const pendingImageAnalysisTurns = new Map<string, string>();
 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);
   }

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<typeof setTimeout>
+>();
+
+/**
+ * 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<void> {
+  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
+  );
+}

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();
+  });
 });