feat: Display initial token usage metrics in /stats (#879)

4 files changed, 158 insertions, 23 deletions

diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
index aa1e701f..cc6be49e 100644
--- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
+++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
@@ -61,13 +61,13 @@ import {
   MCPServerStatus,
   getMCPServerStatus,
 } from '@gemini-cli/core';
-import { useSession } from '../contexts/SessionContext.js';
+import { useSessionStats } from '../contexts/SessionContext.js';
 
 import * as ShowMemoryCommandModule from './useShowMemoryCommand.js';
 import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
 
 vi.mock('../contexts/SessionContext.js', () => ({
-  useSession: vi.fn(),
+  useSessionStats: vi.fn(),
 }));
 
 vi.mock('./useShowMemoryCommand.js', () => ({
@@ -89,7 +89,7 @@ describe('useSlashCommandProcessor', () => {
   let mockPerformMemoryRefresh: ReturnType<typeof vi.fn>;
   let mockConfig: Config;
   let mockCorgiMode: ReturnType<typeof vi.fn>;
-  const mockUseSession = useSession as Mock;
+  const mockUseSessionStats = useSessionStats as Mock;
 
   beforeEach(() => {
     mockAddItem = vi.fn();
@@ -105,8 +105,19 @@ describe('useSlashCommandProcessor', () => {
       getModel: vi.fn(() => 'test-model'),
     } as unknown as Config;
     mockCorgiMode = vi.fn();
-    mockUseSession.mockReturnValue({
-      startTime: new Date('2025-01-01T00:00:00.000Z'),
+    mockUseSessionStats.mockReturnValue({
+      stats: {
+        sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
+        cumulative: {
+          turnCount: 0,
+          promptTokenCount: 0,
+          candidatesTokenCount: 0,
+          totalTokenCount: 0,
+          cachedContentTokenCount: 0,
+          toolUsePromptTokenCount: 0,
+          thoughtsTokenCount: 0,
+        },
+      },
     });
 
     (open as Mock).mockClear();
@@ -240,29 +251,55 @@ describe('useSlashCommandProcessor', () => {
   });
 
   describe('/stats command', () => {
-    it('should show the session duration', async () => {
-      const { handleSlashCommand } = getProcessor();
-      let commandResult: SlashCommandActionReturn | boolean = false;
+    it('should show detailed session statistics', async () => {
+      // Arrange
+      mockUseSessionStats.mockReturnValue({
+        stats: {
+          sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
+          cumulative: {
+            totalTokenCount: 900,
+            promptTokenCount: 200,
+            candidatesTokenCount: 400,
+            cachedContentTokenCount: 100,
+            turnCount: 1,
+            toolUsePromptTokenCount: 50,
+            thoughtsTokenCount: 150,
+          },
+        },
+      });
 
-      // Mock current time
-      const mockDate = new Date('2025-01-01T00:01:05.000Z');
+      const { handleSlashCommand } = getProcessor();
+      const mockDate = new Date('2025-01-01T01:02:03.000Z'); // 1h 2m 3s duration
       vi.setSystemTime(mockDate);
 
+      // Act
       await act(async () => {
-        commandResult = handleSlashCommand('/stats');
+        handleSlashCommand('/stats');
       });
 
+      // Assert
+      const expectedContent = [
+        `  ⎿ Total duration (wall): 1h 2m 3s`,
+        `    Total Token usage:`,
+        `         Turns: 1`,
+        `         Total: 900`,
+        `             ├─ Input: 200`,
+        `             ├─ Output: 400`,
+        `             ├─ Cached: 100`,
+        `             └─ Overhead: 200`,
+        `                  ├─ Model thoughts: 150`,
+        `                  └─ Tool-use prompts: 50`,
+      ].join('\n');
+
       expect(mockAddItem).toHaveBeenNthCalledWith(
-        2,
+        2, // Called after the user message
         expect.objectContaining({
           type: MessageType.INFO,
-          text: 'Session duration: 1m 5s',
+          text: expectedContent,
         }),
         expect.any(Number),
       );
-      expect(commandResult).toBe(true);
 
-      // Restore system time
       vi.useRealTimers();
     });
   });
diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts
index daec0379..6159fe89 100644
--- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts
@@ -11,7 +11,7 @@ import process from 'node:process';
 import { UseHistoryManagerReturn } from './useHistoryManager.js';
 import { Config, MCPServerStatus, getMCPServerStatus } from '@gemini-cli/core';
 import { Message, MessageType, HistoryItemWithoutId } from '../types.js';
-import { useSession } from '../contexts/SessionContext.js';
+import { useSessionStats } from '../contexts/SessionContext.js';
 import { createShowMemoryAction } from './useShowMemoryCommand.js';
 import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
 import { formatMemoryUsage } from '../utils/formatters.js';
@@ -50,8 +50,7 @@ export const useSlashCommandProcessor = (
   toggleCorgiMode: () => void,
   showToolDescriptions: boolean = false,
 ) => {
-  const session = useSession();
-
+  const session = useSessionStats();
   const addMessage = useCallback(
     (message: Message) => {
       // Convert Message to HistoryItemWithoutId
@@ -147,7 +146,9 @@ export const useSlashCommandProcessor = (
         description: 'check session stats',
         action: (_mainCommand, _subCommand, _args) => {
           const now = new Date();
-          const duration = now.getTime() - session.startTime.getTime();
+          const { sessionStartTime, cumulative } = session.stats;
+
+          const duration = now.getTime() - sessionStartTime.getTime();
           const durationInSeconds = Math.floor(duration / 1000);
           const hours = Math.floor(durationInSeconds / 3600);
           const minutes = Math.floor((durationInSeconds % 3600) / 60);
@@ -161,9 +162,25 @@ export const useSlashCommandProcessor = (
             .filter(Boolean)
             .join(' ');
 
+          const overheadTotal =
+            cumulative.thoughtsTokenCount + cumulative.toolUsePromptTokenCount;
+
+          const statsContent = [
+            `  ⎿ Total duration (wall): ${durationString}`,
+            `    Total Token usage:`,
+            `         Turns: ${cumulative.turnCount.toLocaleString()}`,
+            `         Total: ${cumulative.totalTokenCount.toLocaleString()}`,
+            `             ├─ Input: ${cumulative.promptTokenCount.toLocaleString()}`,
+            `             ├─ Output: ${cumulative.candidatesTokenCount.toLocaleString()}`,
+            `             ├─ Cached: ${cumulative.cachedContentTokenCount.toLocaleString()}`,
+            `             └─ Overhead: ${overheadTotal.toLocaleString()}`,
+            `                  ├─ Model thoughts: ${cumulative.thoughtsTokenCount.toLocaleString()}`,
+            `                  └─ Tool-use prompts: ${cumulative.toolUsePromptTokenCount.toLocaleString()}`,
+          ].join('\n');
+
           addMessage({
             type: MessageType.INFO,
-            content: `Session duration: ${durationString}`,
+            content: statsContent,
             timestamp: new Date(),
           });
         },
@@ -477,7 +494,7 @@ Add any other context about the problem here.
       toggleCorgiMode,
       config,
       showToolDescriptions,
-      session.startTime,
+      session,
     ],
   );
 
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index f41f7f9c..ed0f2aac 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -96,6 +96,15 @@ vi.mock('./useLogger.js', () => ({
   }),
 }));
 
+const mockStartNewTurn = vi.fn();
+const mockAddUsage = vi.fn();
+vi.mock('../contexts/SessionContext.js', () => ({
+  useSessionStats: vi.fn(() => ({
+    startNewTurn: mockStartNewTurn,
+    addUsage: mockAddUsage,
+  })),
+}));
+
 vi.mock('./slashCommandProcessor.js', () => ({
   handleSlashCommand: vi.fn().mockReturnValue(false),
 }));
@@ -531,4 +540,63 @@ 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();
+    });
+  });
 });
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 2b47ae6f..bad9f78a 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -42,6 +42,7 @@ import {
   TrackedCompletedToolCall,
   TrackedCancelledToolCall,
 } from './useReactToolScheduler.js';
+import { useSessionStats } from '../contexts/SessionContext.js';
 
 export function mergePartListUnions(list: PartListUnion[]): PartListUnion {
   const resultParts: PartListUnion = [];
@@ -82,6 +83,7 @@ export const useGeminiStream = (
   const [pendingHistoryItemRef, setPendingHistoryItem] =
     useStateAndRef<HistoryItemWithoutId | null>(null);
   const logger = useLogger();
+  const { startNewTurn, addUsage } = useSessionStats();
 
   const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] =
     useReactToolScheduler(
@@ -390,6 +392,9 @@ export const useGeminiStream = (
           case ServerGeminiEventType.ChatCompressed:
             handleChatCompressionEvent();
             break;
+          case ServerGeminiEventType.UsageMetadata:
+            addUsage(event.value);
+            break;
           case ServerGeminiEventType.ToolCallConfirmation:
           case ServerGeminiEventType.ToolCallResponse:
             // do nothing
@@ -412,11 +417,12 @@ export const useGeminiStream = (
       handleErrorEvent,
       scheduleToolCalls,
       handleChatCompressionEvent,
+      addUsage,
     ],
   );
 
   const submitQuery = useCallback(
-    async (query: PartListUnion) => {
+    async (query: PartListUnion, options?: { isContinuation: boolean }) => {
       if (
         streamingState === StreamingState.Responding ||
         streamingState === StreamingState.WaitingForConfirmation
@@ -426,6 +432,10 @@ export const useGeminiStream = (
       const userMessageTimestamp = Date.now();
       setShowHelp(false);
 
+      if (!options?.isContinuation) {
+        startNewTurn();
+      }
+
       abortControllerRef.current = new AbortController();
       const abortSignal = abortControllerRef.current.signal;
 
@@ -491,6 +501,7 @@ export const useGeminiStream = (
       setPendingHistoryItem,
       setInitError,
       geminiClient,
+      startNewTurn,
     ],
   );
 
@@ -576,7 +587,9 @@ export const useGeminiStream = (
       );
 
       markToolsAsSubmitted(callIdsToMarkAsSubmitted);
-      submitQuery(mergePartListUnions(responsesToSend));
+      submitQuery(mergePartListUnions(responsesToSend), {
+        isContinuation: true,
+      });
     }
   }, [
     toolCalls,