feat: Change /stats to include more detailed breakdowns (#2615)

4 files changed, 64 insertions, 105 deletions

diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
index 01954670..d10ae22b 100644
--- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
+++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
@@ -296,19 +296,9 @@ describe('useSlashCommandProcessor', () => {
   describe('/stats command', () => {
     it('should show detailed session statistics', async () => {
       // Arrange
-      const cumulativeStats = {
-        totalTokenCount: 900,
-        promptTokenCount: 200,
-        candidatesTokenCount: 400,
-        cachedContentTokenCount: 100,
-        turnCount: 1,
-        toolUsePromptTokenCount: 50,
-        thoughtsTokenCount: 150,
-      };
       mockUseSessionStats.mockReturnValue({
         stats: {
           sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
-          cumulative: cumulativeStats,
         },
       });
 
@@ -326,7 +316,6 @@ describe('useSlashCommandProcessor', () => {
         2, // Called after the user message
         expect.objectContaining({
           type: MessageType.STATS,
-          stats: cumulativeStats,
           duration: '1h 2m 3s',
         }),
         expect.any(Number),
@@ -334,6 +323,44 @@ describe('useSlashCommandProcessor', () => {
 
       vi.useRealTimers();
     });
+
+    it('should show model-specific statistics when using /stats model', async () => {
+      // Arrange
+      const { handleSlashCommand } = getProcessor();
+
+      // Act
+      await act(async () => {
+        handleSlashCommand('/stats model');
+      });
+
+      // Assert
+      expect(mockAddItem).toHaveBeenNthCalledWith(
+        2, // Called after the user message
+        expect.objectContaining({
+          type: MessageType.MODEL_STATS,
+        }),
+        expect.any(Number),
+      );
+    });
+
+    it('should show tool-specific statistics when using /stats tools', async () => {
+      // Arrange
+      const { handleSlashCommand } = getProcessor();
+
+      // Act
+      await act(async () => {
+        handleSlashCommand('/stats tools');
+      });
+
+      // Assert
+      expect(mockAddItem).toHaveBeenNthCalledWith(
+        2, // Called after the user message
+        expect.objectContaining({
+          type: MessageType.TOOL_STATS,
+        }),
+        expect.any(Number),
+      );
+    });
   });
 
   describe('/about command', () => {
@@ -598,7 +625,6 @@ describe('useSlashCommandProcessor', () => {
           },
           {
             type: 'quit',
-            stats: expect.any(Object),
             duration: '1h 2m 3s',
             id: expect.any(Number),
           },
diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts
index b7dcbdcb..ffc3d7d1 100644
--- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts
@@ -110,14 +110,19 @@ export const useSlashCommandProcessor = (
       } else if (message.type === MessageType.STATS) {
         historyItemContent = {
           type: 'stats',
-          stats: message.stats,
-          lastTurnStats: message.lastTurnStats,
           duration: message.duration,
         };
+      } else if (message.type === MessageType.MODEL_STATS) {
+        historyItemContent = {
+          type: 'model_stats',
+        };
+      } else if (message.type === MessageType.TOOL_STATS) {
+        historyItemContent = {
+          type: 'tool_stats',
+        };
       } else if (message.type === MessageType.QUIT) {
         historyItemContent = {
           type: 'quit',
-          stats: message.stats,
           duration: message.duration,
         };
       } else if (message.type === MessageType.COMPRESSION) {
@@ -262,16 +267,28 @@ export const useSlashCommandProcessor = (
       {
         name: 'stats',
         altName: 'usage',
-        description: 'check session stats',
-        action: (_mainCommand, _subCommand, _args) => {
+        description: 'check session stats. Usage: /stats [model|tools]',
+        action: (_mainCommand, subCommand, _args) => {
+          if (subCommand === 'model') {
+            addMessage({
+              type: MessageType.MODEL_STATS,
+              timestamp: new Date(),
+            });
+            return;
+          } else if (subCommand === 'tools') {
+            addMessage({
+              type: MessageType.TOOL_STATS,
+              timestamp: new Date(),
+            });
+            return;
+          }
+
           const now = new Date();
-          const { sessionStartTime, cumulative, currentTurn } = session.stats;
+          const { sessionStartTime } = session.stats;
           const wallDuration = now.getTime() - sessionStartTime.getTime();
 
           addMessage({
             type: MessageType.STATS,
-            stats: cumulative,
-            lastTurnStats: currentTurn,
             duration: formatDuration(wallDuration),
             timestamp: new Date(),
           });
@@ -805,7 +822,7 @@ export const useSlashCommandProcessor = (
         description: 'exit the cli',
         action: async (mainCommand, _subCommand, _args) => {
           const now = new Date();
-          const { sessionStartTime, cumulative } = session.stats;
+          const { sessionStartTime } = session.stats;
           const wallDuration = now.getTime() - sessionStartTime.getTime();
 
           setQuittingMessages([
@@ -816,7 +833,6 @@ export const useSlashCommandProcessor = (
             },
             {
               type: 'quit',
-              stats: cumulative,
               duration: formatDuration(wallDuration),
               id: now.getTime(),
             },
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index 0c8b261e..9751f470 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -604,78 +604,6 @@ describe('useGeminiStream', () => {
     });
   });
 
-  describe('Session Stats Integration', () => {
-    it('should call startNewTurn and addUsage for a simple prompt', async () => {
-      const mockMetadata = { totalTokenCount: 123 };
-      const mockStream = (async function* () {
-        yield { type: 'content', value: 'Response' };
-        yield { type: 'usage_metadata', value: mockMetadata };
-      })();
-      mockSendMessageStream.mockReturnValue(mockStream);
-
-      const { result } = renderTestHook();
-
-      await act(async () => {
-        await result.current.submitQuery('Hello, world!');
-      });
-
-      expect(mockStartNewTurn).toHaveBeenCalledTimes(1);
-      expect(mockAddUsage).toHaveBeenCalledTimes(1);
-      expect(mockAddUsage).toHaveBeenCalledWith(mockMetadata);
-    });
-
-    it('should only call addUsage for a tool continuation prompt', async () => {
-      const mockMetadata = { totalTokenCount: 456 };
-      const mockStream = (async function* () {
-        yield { type: 'content', value: 'Final Answer' };
-        yield { type: 'usage_metadata', value: mockMetadata };
-      })();
-      mockSendMessageStream.mockReturnValue(mockStream);
-
-      const { result } = renderTestHook();
-
-      await act(async () => {
-        await result.current.submitQuery([{ text: 'tool response' }], {
-          isContinuation: true,
-        });
-      });
-
-      expect(mockStartNewTurn).not.toHaveBeenCalled();
-      expect(mockAddUsage).toHaveBeenCalledTimes(1);
-      expect(mockAddUsage).toHaveBeenCalledWith(mockMetadata);
-    });
-
-    it('should not call addUsage if the stream contains no usage metadata', async () => {
-      // Arrange: A stream that yields content but never a usage_metadata event
-      const mockStream = (async function* () {
-        yield { type: 'content', value: 'Some response text' };
-      })();
-      mockSendMessageStream.mockReturnValue(mockStream);
-
-      const { result } = renderTestHook();
-
-      await act(async () => {
-        await result.current.submitQuery('Query with no usage data');
-      });
-
-      expect(mockStartNewTurn).toHaveBeenCalledTimes(1);
-      expect(mockAddUsage).not.toHaveBeenCalled();
-    });
-
-    it('should not call startNewTurn for a slash command', async () => {
-      mockHandleSlashCommand.mockReturnValue(true);
-
-      const { result } = renderTestHook();
-
-      await act(async () => {
-        await result.current.submitQuery('/stats');
-      });
-
-      expect(mockStartNewTurn).not.toHaveBeenCalled();
-      expect(mockSendMessageStream).not.toHaveBeenCalled();
-    });
-  });
-
   it('should not flicker streaming state to Idle between tool completion and submission', async () => {
     const toolCallResponseParts: PartListUnion = [
       { text: 'tool 1 final response' },
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 3d24ede7..e2226761 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -51,7 +51,6 @@ import {
   TrackedCompletedToolCall,
   TrackedCancelledToolCall,
 } from './useReactToolScheduler.js';
-import { useSessionStats } from '../contexts/SessionContext.js';
 
 export function mergePartListUnions(list: PartListUnion[]): PartListUnion {
   const resultParts: PartListUnion = [];
@@ -101,7 +100,6 @@ export const useGeminiStream = (
     useStateAndRef<HistoryItemWithoutId | null>(null);
   const processedMemoryToolsRef = useRef<Set<string>>(new Set());
   const logger = useLogger();
-  const { startNewTurn, addUsage } = useSessionStats();
   const gitService = useMemo(() => {
     if (!config.getProjectRoot()) {
       return;
@@ -461,9 +459,6 @@ export const useGeminiStream = (
           case ServerGeminiEventType.ChatCompressed:
             handleChatCompressionEvent(event.value);
             break;
-          case ServerGeminiEventType.UsageMetadata:
-            addUsage(event.value);
-            break;
           case ServerGeminiEventType.ToolCallConfirmation:
           case ServerGeminiEventType.ToolCallResponse:
             // do nothing
@@ -486,7 +481,6 @@ export const useGeminiStream = (
       handleErrorEvent,
       scheduleToolCalls,
       handleChatCompressionEvent,
-      addUsage,
     ],
   );
 
@@ -516,10 +510,6 @@ export const useGeminiStream = (
         return;
       }
 
-      if (!options?.isContinuation) {
-        startNewTurn();
-      }
-
       setIsResponding(true);
       setInitError(null);
 
@@ -568,7 +558,6 @@ export const useGeminiStream = (
       setPendingHistoryItem,
       setInitError,
       geminiClient,
-      startNewTurn,
       onAuthError,
       config,
     ],