1 files changed, 29 insertions, 15 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', () => {