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

1 files changed, 41 insertions, 11 deletions

diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx
index d3ee0f5f..a009f3d8 100644
--- a/packages/cli/src/ui/components/SessionSummaryDisplay.tsx
+++ b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx
@@ -9,31 +9,57 @@ import { Box, Text } from 'ink';
 import Gradient from 'ink-gradient';
 import { Colors } from '../colors.js';
 import { formatDuration } from '../utils/formatters.js';
-import { CumulativeStats } from '../contexts/SessionContext.js';
+import { useSessionStats } from '../contexts/SessionContext.js';
+import { computeSessionStats } from '../utils/computeStats.js';
 import { FormattedStats, StatRow, StatsColumn } from './Stats.js';
 
 // --- Prop and Data Structures ---
 
 interface SessionSummaryDisplayProps {
-  stats: CumulativeStats;
   duration: string;
 }
 
 // --- Main Component ---
 
 export const SessionSummaryDisplay: React.FC<SessionSummaryDisplayProps> = ({
-  stats,
   duration,
 }) => {
+  const { stats } = useSessionStats();
+  const { metrics } = stats;
+  const computed = computeSessionStats(metrics);
+
   const cumulativeFormatted: FormattedStats = {
-    inputTokens: stats.promptTokenCount,
-    outputTokens: stats.candidatesTokenCount,
-    toolUseTokens: stats.toolUsePromptTokenCount,
-    thoughtsTokens: stats.thoughtsTokenCount,
-    cachedTokens: stats.cachedContentTokenCount,
-    totalTokens: stats.totalTokenCount,
+    inputTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.prompt,
+      0,
+    ),
+    outputTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.candidates,
+      0,
+    ),
+    toolUseTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.tool,
+      0,
+    ),
+    thoughtsTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.thoughts,
+      0,
+    ),
+    cachedTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.cached,
+      0,
+    ),
+    totalTokens: Object.values(metrics.models).reduce(
+      (acc, model) => acc + model.tokens.total,
+      0,
+    ),
   };
 
+  const totalRequests = Object.values(metrics.models).reduce(
+    (acc, model) => acc + model.api.totalRequests,
+    0,
+  );
+
   const title = 'Agent powering down. Goodbye!';
 
   return (
@@ -57,14 +83,18 @@ export const SessionSummaryDisplay: React.FC<SessionSummaryDisplayProps> = ({
 
       <Box marginTop={1}>
         <StatsColumn
-          title={`Cumulative Stats (${stats.turnCount} Turns)`}
+          title={`Cumulative Stats (${totalRequests} API calls)`}
           stats={cumulativeFormatted}
           isCumulative={true}
         >
           <Box marginTop={1} flexDirection="column">
             <StatRow
               label="Total duration (API)"
-              value={formatDuration(stats.apiTimeMs)}
+              value={formatDuration(computed.totalApiTime)}
+            />
+            <StatRow
+              label="Total duration (Tools)"
+              value={formatDuration(computed.totalToolTime)}
             />
             <StatRow label="Total duration (wall)" value={duration} />
           </Box>