1 files changed, 56 insertions, 28 deletions

diff --git a/packages/cli/src/ui/contexts/SessionContext.tsx b/packages/cli/src/ui/contexts/SessionContext.tsx
index 0549e3e1..0d574e75 100644
--- a/packages/cli/src/ui/contexts/SessionContext.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.tsx
@@ -16,7 +16,7 @@ import { type GenerateContentResponseUsageMetadata } from '@google/genai';
 
 // --- Interface Definitions ---
 
-interface CumulativeStats {
+export interface CumulativeStats {
   turnCount: number;
   promptTokenCount: number;
   candidatesTokenCount: number;
@@ -24,18 +24,13 @@ interface CumulativeStats {
   cachedContentTokenCount: number;
   toolUsePromptTokenCount: number;
   thoughtsTokenCount: number;
-}
-
-interface LastTurnStats {
-  metadata: GenerateContentResponseUsageMetadata;
-  // TODO(abhipatel12): Add apiTime, etc. here in a future step.
+  apiTimeMs: number;
 }
 
 interface SessionStatsState {
   sessionStartTime: Date;
   cumulative: CumulativeStats;
-  lastTurn: LastTurnStats | null;
-  isNewTurnForAggregation: boolean;
+  currentTurn: CumulativeStats;
 }
 
 // Defines the final "value" of our context, including the state
@@ -43,7 +38,9 @@ interface SessionStatsState {
 interface SessionStatsContextValue {
   stats: SessionStatsState;
   startNewTurn: () => void;
-  addUsage: (metadata: GenerateContentResponseUsageMetadata) => void;
+  addUsage: (
+    metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
+  ) => void;
 }
 
 // --- Context Definition ---
@@ -52,6 +49,27 @@ const SessionStatsContext = createContext<SessionStatsContextValue | undefined>(
   undefined,
 );
 
+// --- Helper Functions ---
+
+/**
+ * A small, reusable helper function to sum token counts.
+ * It unconditionally adds all token values from the source to the target.
+ * @param target The object to add the tokens to (e.g., cumulative, currentTurn).
+ * @param source The metadata object from the API response.
+ */
+const addTokens = (
+  target: CumulativeStats,
+  source: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
+) => {
+  target.candidatesTokenCount += source.candidatesTokenCount ?? 0;
+  target.thoughtsTokenCount += source.thoughtsTokenCount ?? 0;
+  target.totalTokenCount += source.totalTokenCount ?? 0;
+  target.apiTimeMs += source.apiTimeMs ?? 0;
+  target.promptTokenCount += source.promptTokenCount ?? 0;
+  target.cachedContentTokenCount += source.cachedContentTokenCount ?? 0;
+  target.toolUsePromptTokenCount += source.toolUsePromptTokenCount ?? 0;
+};
+
 // --- Provider Component ---
 
 export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
@@ -67,36 +85,37 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
       cachedContentTokenCount: 0,
       toolUsePromptTokenCount: 0,
       thoughtsTokenCount: 0,
+      apiTimeMs: 0,
+    },
+    currentTurn: {
+      turnCount: 0,
+      promptTokenCount: 0,
+      candidatesTokenCount: 0,
+      totalTokenCount: 0,
+      cachedContentTokenCount: 0,
+      toolUsePromptTokenCount: 0,
+      thoughtsTokenCount: 0,
+      apiTimeMs: 0,
     },
-    lastTurn: null,
-    isNewTurnForAggregation: true,
   });
 
   // A single, internal worker function to handle all metadata aggregation.
   const aggregateTokens = useCallback(
-    (metadata: GenerateContentResponseUsageMetadata) => {
+    (
+      metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
+    ) => {
       setStats((prevState) => {
-        const { isNewTurnForAggregation } = prevState;
         const newCumulative = { ...prevState.cumulative };
+        const newCurrentTurn = { ...prevState.currentTurn };
 
-        newCumulative.candidatesTokenCount +=
-          metadata.candidatesTokenCount ?? 0;
-        newCumulative.thoughtsTokenCount += metadata.thoughtsTokenCount ?? 0;
-        newCumulative.totalTokenCount += metadata.totalTokenCount ?? 0;
-
-        if (isNewTurnForAggregation) {
-          newCumulative.promptTokenCount += metadata.promptTokenCount ?? 0;
-          newCumulative.cachedContentTokenCount +=
-            metadata.cachedContentTokenCount ?? 0;
-          newCumulative.toolUsePromptTokenCount +=
-            metadata.toolUsePromptTokenCount ?? 0;
-        }
+        // Add all tokens to the current turn's stats as well as cumulative stats.
+        addTokens(newCurrentTurn, metadata);
+        addTokens(newCumulative, metadata);
 
         return {
           ...prevState,
           cumulative: newCumulative,
-          lastTurn: { metadata },
-          isNewTurnForAggregation: false,
+          currentTurn: newCurrentTurn,
         };
       });
     },
@@ -110,7 +129,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
         ...prevState.cumulative,
         turnCount: prevState.cumulative.turnCount + 1,
       },
-      isNewTurnForAggregation: true,
+      currentTurn: {
+        turnCount: 0, // Reset for the new turn's accumulation.
+        promptTokenCount: 0,
+        candidatesTokenCount: 0,
+        totalTokenCount: 0,
+        cachedContentTokenCount: 0,
+        toolUsePromptTokenCount: 0,
+        thoughtsTokenCount: 0,
+        apiTimeMs: 0,
+      },
     }));
   }, []);