feat: Add UI for /stats slash command (#883)

2 files changed, 85 insertions, 43 deletions

diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx
index fedf3d74..b00a5d75 100644
--- a/packages/cli/src/ui/contexts/SessionContext.test.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx
@@ -59,7 +59,7 @@ describe('SessionStatsContext', () => {
     const stats = contextRef.current?.stats;
 
     expect(stats?.sessionStartTime).toBeInstanceOf(Date);
-    expect(stats?.lastTurn).toBeNull();
+    expect(stats?.currentTurn).toBeDefined();
     expect(stats?.cumulative.turnCount).toBe(0);
     expect(stats?.cumulative.totalTokenCount).toBe(0);
     expect(stats?.cumulative.promptTokenCount).toBe(0);
@@ -81,6 +81,7 @@ describe('SessionStatsContext', () => {
     });
 
     const stats = contextRef.current?.stats;
+    expect(stats?.currentTurn.totalTokenCount).toBe(0);
     expect(stats?.cumulative.turnCount).toBe(1);
     // Ensure token counts are unaffected
     expect(stats?.cumulative.totalTokenCount).toBe(0);
@@ -98,7 +99,7 @@ describe('SessionStatsContext', () => {
     );
 
     act(() => {
-      contextRef.current?.addUsage(mockMetadata1);
+      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 });
     });
 
     const stats = contextRef.current?.stats;
@@ -110,12 +111,16 @@ describe('SessionStatsContext', () => {
     expect(stats?.cumulative.promptTokenCount).toBe(
       mockMetadata1.promptTokenCount ?? 0,
     );
+    expect(stats?.cumulative.apiTimeMs).toBe(123);
 
     // Check that turn count is NOT incremented
     expect(stats?.cumulative.turnCount).toBe(0);
 
-    // Check that lastTurn is updated
-    expect(stats?.lastTurn?.metadata).toEqual(mockMetadata1);
+    // Check that currentTurn is updated
+    expect(stats?.currentTurn?.totalTokenCount).toEqual(
+      mockMetadata1.totalTokenCount,
+    );
+    expect(stats?.currentTurn?.apiTimeMs).toBe(123);
   });
 
   it('should correctly track a full logical turn with multiple API calls', () => {
@@ -136,12 +141,12 @@ describe('SessionStatsContext', () => {
 
     // 2. First API call (e.g., prompt with a tool request)
     act(() => {
-      contextRef.current?.addUsage(mockMetadata1);
+      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
     });
 
     // 3. Second API call (e.g., sending tool response back)
     act(() => {
-      contextRef.current?.addUsage(mockMetadata2);
+      contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
     });
 
     const stats = contextRef.current?.stats;
@@ -149,18 +154,27 @@ describe('SessionStatsContext', () => {
     // Turn count should only be 1
     expect(stats?.cumulative.turnCount).toBe(1);
 
+    // --- Check Cumulative Stats ---
     // These fields should be the SUM of both calls
-    expect(stats?.cumulative.totalTokenCount).toBe(330); // 300 + 30
-    expect(stats?.cumulative.candidatesTokenCount).toBe(220); // 200 + 20
-    expect(stats?.cumulative.thoughtsTokenCount).toBe(22); // 20 + 2
+    expect(stats?.cumulative.totalTokenCount).toBe(300 + 30);
+    expect(stats?.cumulative.candidatesTokenCount).toBe(200 + 20);
+    expect(stats?.cumulative.thoughtsTokenCount).toBe(20 + 2);
+    expect(stats?.cumulative.apiTimeMs).toBe(100 + 50);
 
-    // These fields should ONLY be from the FIRST call, because isNewTurnForAggregation was true
-    expect(stats?.cumulative.promptTokenCount).toBe(100);
-    expect(stats?.cumulative.cachedContentTokenCount).toBe(50);
-    expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10);
+    // These fields should be the SUM of both calls
+    expect(stats?.cumulative.promptTokenCount).toBe(100 + 10);
+    expect(stats?.cumulative.cachedContentTokenCount).toBe(50 + 5);
+    expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10 + 1);
 
-    // Last turn should hold the metadata from the most recent call
-    expect(stats?.lastTurn?.metadata).toEqual(mockMetadata2);
+    // --- Check Current Turn Stats ---
+    // All fields should be the SUM of both calls for the turn
+    expect(stats?.currentTurn.totalTokenCount).toBe(300 + 30);
+    expect(stats?.currentTurn.candidatesTokenCount).toBe(200 + 20);
+    expect(stats?.currentTurn.thoughtsTokenCount).toBe(20 + 2);
+    expect(stats?.currentTurn.promptTokenCount).toBe(100 + 10);
+    expect(stats?.currentTurn.cachedContentTokenCount).toBe(50 + 5);
+    expect(stats?.currentTurn.toolUsePromptTokenCount).toBe(10 + 1);
+    expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
   });
 
   it('should throw an error when useSessionStats is used outside of a provider', () => {
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,
+      },
     }));
   }, []);