feat: Add token stats in footer (#909)

packages/cli/src/ui/contexts/SessionContext.test.tsx

@@ -177,6 +177,51 @@ describe('SessionStatsContext', () => {
     expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
   });
 
+  it('should overwrite currentResponse with each API call', () => {
+    const contextRef: MutableRefObject<
+      ReturnType<typeof useSessionStats> | undefined
+    > = { current: undefined };
+
+    render(
+      <SessionStatsProvider>
+        <TestHarness contextRef={contextRef} />
+      </SessionStatsProvider>,
+    );
+
+    // 1. First API call
+    act(() => {
+      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
+    });
+
+    let stats = contextRef.current?.stats;
+
+    // currentResponse should match the first call
+    expect(stats?.currentResponse.totalTokenCount).toBe(300);
+    expect(stats?.currentResponse.apiTimeMs).toBe(100);
+
+    // 2. Second API call
+    act(() => {
+      contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
+    });
+
+    stats = contextRef.current?.stats;
+
+    // currentResponse should now match the second call
+    expect(stats?.currentResponse.totalTokenCount).toBe(30);
+    expect(stats?.currentResponse.apiTimeMs).toBe(50);
+
+    // 3. Start a new turn
+    act(() => {
+      contextRef.current?.startNewTurn();
+    });
+
+    stats = contextRef.current?.stats;
+
+    // currentResponse should be reset
+    expect(stats?.currentResponse.totalTokenCount).toBe(0);
+    expect(stats?.currentResponse.apiTimeMs).toBe(0);
+  });
+
   it('should throw an error when useSessionStats is used outside of a provider', () => {
     // Suppress the expected console error during this test.
     const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});

packages/cli/src/ui/contexts/SessionContext.tsx

@@ -31,6 +31,7 @@ interface SessionStatsState {
   sessionStartTime: Date;
   cumulative: CumulativeStats;
   currentTurn: CumulativeStats;
+  currentResponse: CumulativeStats;
 }
 
 // Defines the final "value" of our context, including the state
@@ -97,6 +98,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
       thoughtsTokenCount: 0,
       apiTimeMs: 0,
     },
+    currentResponse: {
+      turnCount: 0,
+      promptTokenCount: 0,
+      candidatesTokenCount: 0,
+      totalTokenCount: 0,
+      cachedContentTokenCount: 0,
+      toolUsePromptTokenCount: 0,
+      thoughtsTokenCount: 0,
+      apiTimeMs: 0,
+    },
   });
 
   // A single, internal worker function to handle all metadata aggregation.
@@ -107,15 +118,27 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
       setStats((prevState) => {
         const newCumulative = { ...prevState.cumulative };
         const newCurrentTurn = { ...prevState.currentTurn };
+        const newCurrentResponse = {
+          turnCount: 0,
+          promptTokenCount: 0,
+          candidatesTokenCount: 0,
+          totalTokenCount: 0,
+          cachedContentTokenCount: 0,
+          toolUsePromptTokenCount: 0,
+          thoughtsTokenCount: 0,
+          apiTimeMs: 0,
+        };
 
         // Add all tokens to the current turn's stats as well as cumulative stats.
         addTokens(newCurrentTurn, metadata);
         addTokens(newCumulative, metadata);
+        addTokens(newCurrentResponse, metadata);
 
         return {
           ...prevState,
           cumulative: newCumulative,
           currentTurn: newCurrentTurn,
+          currentResponse: newCurrentResponse,
         };
       });
     },
@@ -139,6 +162,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
         thoughtsTokenCount: 0,
         apiTimeMs: 0,
       },
+      currentResponse: {
+        turnCount: 0,
+        promptTokenCount: 0,
+        candidatesTokenCount: 0,
+        totalTokenCount: 0,
+        cachedContentTokenCount: 0,
+        toolUsePromptTokenCount: 0,
+        thoughtsTokenCount: 0,
+        apiTimeMs: 0,
+      },
     }));
   }, []);