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

1 files changed, 49 insertions, 134 deletions

diff --git a/packages/cli/src/ui/contexts/SessionContext.tsx b/packages/cli/src/ui/contexts/SessionContext.tsx
index f59e17e1..b89d19e7 100644
--- a/packages/cli/src/ui/contexts/SessionContext.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.tsx
@@ -9,39 +9,43 @@ import React, {
   useContext,
   useState,
   useMemo,
-  useCallback,
+  useEffect,
 } from 'react';
 
-import { type GenerateContentResponseUsageMetadata } from '@google/genai';
+import {
+  uiTelemetryService,
+  SessionMetrics,
+  ModelMetrics,
+} from '@google/gemini-cli-core';
 
 // --- Interface Definitions ---
 
-export interface CumulativeStats {
-  turnCount: number;
-  promptTokenCount: number;
-  candidatesTokenCount: number;
-  totalTokenCount: number;
-  cachedContentTokenCount: number;
-  toolUsePromptTokenCount: number;
-  thoughtsTokenCount: number;
-  apiTimeMs: number;
-}
+export type { SessionMetrics, ModelMetrics };
 
 interface SessionStatsState {
   sessionStartTime: Date;
-  cumulative: CumulativeStats;
-  currentTurn: CumulativeStats;
-  currentResponse: CumulativeStats;
+  metrics: SessionMetrics;
+  lastPromptTokenCount: number;
+}
+
+export interface ComputedSessionStats {
+  totalApiTime: number;
+  totalToolTime: number;
+  agentActiveTime: number;
+  apiTimePercent: number;
+  toolTimePercent: number;
+  cacheEfficiency: number;
+  totalDecisions: number;
+  successRate: number;
+  agreementRate: number;
+  totalCachedTokens: number;
+  totalPromptTokens: number;
 }
 
 // Defines the final "value" of our context, including the state
 // and the functions to update it.
 interface SessionStatsContextValue {
   stats: SessionStatsState;
-  startNewTurn: () => void;
-  addUsage: (
-    metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
-  ) => void;
 }
 
 // --- Context Definition ---
@@ -50,27 +54,6 @@ const SessionStatsContext = createContext<SessionStatsContextValue | undefined>(
   undefined,
 );
 
-// --- Helper Functions ---
-
-/**
- * A small, reusable helper function to sum token counts.
- * It unconditionally adds all token values from the source to the target.
- * @param target The object to add the tokens to (e.g., cumulative, currentTurn).
- * @param source The metadata object from the API response.
- */
-const addTokens = (
-  target: CumulativeStats,
-  source: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
-) => {
-  target.candidatesTokenCount += source.candidatesTokenCount ?? 0;
-  target.thoughtsTokenCount += source.thoughtsTokenCount ?? 0;
-  target.totalTokenCount += source.totalTokenCount ?? 0;
-  target.apiTimeMs += source.apiTimeMs ?? 0;
-  target.promptTokenCount += source.promptTokenCount ?? 0;
-  target.cachedContentTokenCount += source.cachedContentTokenCount ?? 0;
-  target.toolUsePromptTokenCount += source.toolUsePromptTokenCount ?? 0;
-};
-
 // --- Provider Component ---
 
 export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
@@ -78,110 +61,42 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
 }) => {
   const [stats, setStats] = useState<SessionStatsState>({
     sessionStartTime: new Date(),
-    cumulative: {
-      turnCount: 0,
-      promptTokenCount: 0,
-      candidatesTokenCount: 0,
-      totalTokenCount: 0,
-      cachedContentTokenCount: 0,
-      toolUsePromptTokenCount: 0,
-      thoughtsTokenCount: 0,
-      apiTimeMs: 0,
-    },
-    currentTurn: {
-      turnCount: 0,
-      promptTokenCount: 0,
-      candidatesTokenCount: 0,
-      totalTokenCount: 0,
-      cachedContentTokenCount: 0,
-      toolUsePromptTokenCount: 0,
-      thoughtsTokenCount: 0,
-      apiTimeMs: 0,
-    },
-    currentResponse: {
-      turnCount: 0,
-      promptTokenCount: 0,
-      candidatesTokenCount: 0,
-      totalTokenCount: 0,
-      cachedContentTokenCount: 0,
-      toolUsePromptTokenCount: 0,
-      thoughtsTokenCount: 0,
-      apiTimeMs: 0,
-    },
+    metrics: uiTelemetryService.getMetrics(),
+    lastPromptTokenCount: 0,
   });
 
-  // A single, internal worker function to handle all metadata aggregation.
-  const aggregateTokens = useCallback(
-    (
-      metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
-    ) => {
-      setStats((prevState) => {
-        const newCumulative = { ...prevState.cumulative };
-        const newCurrentTurn = { ...prevState.currentTurn };
-        const newCurrentResponse = {
-          turnCount: 0,
-          promptTokenCount: 0,
-          candidatesTokenCount: 0,
-          totalTokenCount: 0,
-          cachedContentTokenCount: 0,
-          toolUsePromptTokenCount: 0,
-          thoughtsTokenCount: 0,
-          apiTimeMs: 0,
-        };
+  useEffect(() => {
+    const handleUpdate = ({
+      metrics,
+      lastPromptTokenCount,
+    }: {
+      metrics: SessionMetrics;
+      lastPromptTokenCount: number;
+    }) => {
+      setStats((prevState) => ({
+        ...prevState,
+        metrics,
+        lastPromptTokenCount,
+      }));
+    };
 
-        // Add all tokens to the current turn's stats as well as cumulative stats.
-        addTokens(newCurrentTurn, metadata);
-        addTokens(newCumulative, metadata);
-        addTokens(newCurrentResponse, metadata);
-
-        return {
-          ...prevState,
-          cumulative: newCumulative,
-          currentTurn: newCurrentTurn,
-          currentResponse: newCurrentResponse,
-        };
-      });
-    },
-    [],
-  );
+    uiTelemetryService.on('update', handleUpdate);
+    // Set initial state
+    handleUpdate({
+      metrics: uiTelemetryService.getMetrics(),
+      lastPromptTokenCount: uiTelemetryService.getLastPromptTokenCount(),
+    });
 
-  const startNewTurn = useCallback(() => {
-    setStats((prevState) => ({
-      ...prevState,
-      cumulative: {
-        ...prevState.cumulative,
-        turnCount: prevState.cumulative.turnCount + 1,
-      },
-      currentTurn: {
-        turnCount: 0, // Reset for the new turn's accumulation.
-        promptTokenCount: 0,
-        candidatesTokenCount: 0,
-        totalTokenCount: 0,
-        cachedContentTokenCount: 0,
-        toolUsePromptTokenCount: 0,
-        thoughtsTokenCount: 0,
-        apiTimeMs: 0,
-      },
-      currentResponse: {
-        turnCount: 0,
-        promptTokenCount: 0,
-        candidatesTokenCount: 0,
-        totalTokenCount: 0,
-        cachedContentTokenCount: 0,
-        toolUsePromptTokenCount: 0,
-        thoughtsTokenCount: 0,
-        apiTimeMs: 0,
-      },
-    }));
+    return () => {
+      uiTelemetryService.off('update', handleUpdate);
+    };
   }, []);
 
   const value = useMemo(
     () => ({
       stats,
-      startNewTurn,
-      addUsage: aggregateTokens,
     }),
-    [stats, startNewTurn, aggregateTokens],
+    [stats],
   );
 
   return (