From b3d89a10758462947546a7d9be43bf617c0787c6 Mon Sep 17 00:00:00 2001
From: Asad Memon <asad.lionpk@gmail.com>
Date: Sun, 15 Jun 2025 11:15:53 -0700
Subject: feat: Add token stats in footer (#909)

---
 .../cli/src/ui/contexts/SessionContext.test.tsx    | 45 ++++++++++++++++++++++
 1 file changed, 45 insertions(+)

(limited to 'packages/cli/src/ui/contexts/SessionContext.test.tsx')

diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx
index b00a5d75..e9fc33e6 100644
--- a/packages/cli/src/ui/contexts/SessionContext.test.tsx
+++ b/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(() => {});
-- 
cgit v1.2.3