From 770f862832dfef477705bee69bd2a84397d105a8 Mon Sep 17 00:00:00 2001
From: Abhi <43648792+abhipatel12@users.noreply.github.com>
Date: Sun, 29 Jun 2025 20:44:33 -0400
Subject: feat: Change /stats to include more detailed breakdowns (#2615)

---
 .../ui/components/SessionSummaryDisplay.test.tsx   | 106 +++++++++++++++------
 1 file changed, 75 insertions(+), 31 deletions(-)

(limited to 'packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx')

diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
index 14d8a277..afb822e5 100644
--- a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
+++ b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
@@ -5,48 +5,92 @@
  */
 
 import { render } from 'ink-testing-library';
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import { SessionSummaryDisplay } from './SessionSummaryDisplay.js';
-import { type CumulativeStats } from '../contexts/SessionContext.js';
+import * as SessionContext from '../contexts/SessionContext.js';
+import { SessionMetrics } from '../contexts/SessionContext.js';
 
-describe('<SessionSummaryDisplay />', () => {
-  const mockStats: CumulativeStats = {
-    turnCount: 10,
-    promptTokenCount: 1000,
-    candidatesTokenCount: 2000,
-    totalTokenCount: 3500,
-    cachedContentTokenCount: 500,
-    toolUsePromptTokenCount: 200,
-    thoughtsTokenCount: 300,
-    apiTimeMs: 50234,
+vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
+  const actual = await importOriginal<typeof SessionContext>();
+  return {
+    ...actual,
+    useSessionStats: vi.fn(),
   };
+});
 
-  const mockDuration = '1h 23m 45s';
+const useSessionStatsMock = vi.mocked(SessionContext.useSessionStats);
 
-  it('renders correctly with given stats and duration', () => {
-    const { lastFrame } = render(
-      <SessionSummaryDisplay stats={mockStats} duration={mockDuration} />,
-    );
+const renderWithMockedStats = (metrics: SessionMetrics) => {
+  useSessionStatsMock.mockReturnValue({
+    stats: {
+      sessionStartTime: new Date(),
+      metrics,
+      lastPromptTokenCount: 0,
+    },
+  });
 
-    expect(lastFrame()).toMatchSnapshot();
+  return render(<SessionSummaryDisplay duration="1h 23m 45s" />);
+};
+
+describe('<SessionSummaryDisplay />', () => {
+  it('correctly sums and displays stats from multiple models', () => {
+    const metrics: SessionMetrics = {
+      models: {
+        'gemini-2.5-pro': {
+          api: { totalRequests: 10, totalErrors: 1, totalLatencyMs: 50234 },
+          tokens: {
+            prompt: 1000,
+            candidates: 2000,
+            total: 3500,
+            cached: 500,
+            thoughts: 300,
+            tool: 200,
+          },
+        },
+        'gemini-2.5-flash': {
+          api: { totalRequests: 5, totalErrors: 0, totalLatencyMs: 12345 },
+          tokens: {
+            prompt: 500,
+            candidates: 1000,
+            total: 1500,
+            cached: 100,
+            thoughts: 50,
+            tool: 20,
+          },
+        },
+      },
+      tools: {
+        totalCalls: 0,
+        totalSuccess: 0,
+        totalFail: 0,
+        totalDurationMs: 0,
+        totalDecisions: { accept: 0, reject: 0, modify: 0 },
+        byName: {},
+      },
+    };
+
+    const { lastFrame } = renderWithMockedStats(metrics);
+    const output = lastFrame();
+
+    // Verify totals are summed correctly
+    expect(output).toContain('Cumulative Stats (15 API calls)');
+    expect(output).toMatchSnapshot();
   });
 
   it('renders zero state correctly', () => {
-    const zeroStats: CumulativeStats = {
-      turnCount: 0,
-      promptTokenCount: 0,
-      candidatesTokenCount: 0,
-      totalTokenCount: 0,
-      cachedContentTokenCount: 0,
-      toolUsePromptTokenCount: 0,
-      thoughtsTokenCount: 0,
-      apiTimeMs: 0,
+    const zeroMetrics: SessionMetrics = {
+      models: {},
+      tools: {
+        totalCalls: 0,
+        totalSuccess: 0,
+        totalFail: 0,
+        totalDurationMs: 0,
+        totalDecisions: { accept: 0, reject: 0, modify: 0 },
+        byName: {},
+      },
     };
 
-    const { lastFrame } = render(
-      <SessionSummaryDisplay stats={zeroStats} duration="0s" />,
-    );
-
+    const { lastFrame } = renderWithMockedStats(zeroMetrics);
     expect(lastFrame()).toMatchSnapshot();
   });
 });
-- 
cgit v1.2.3