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

1 files changed, 15 insertions, 2 deletions

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,