feat(cli): display real-time token consumption during streaming (#2742)

packages/cli/src/ui/AppContainer.tsx

@@ -727,6 +727,8 @@ export const AppContainer = (props: AppContainerProps) => {
     handleApprovalModeChange,
     activePtyId,
     loopDetectionConfirmationRequest,
+    streamingResponseLengthRef,
+    isReceivingContent,
   } = useGeminiStream(
     config.getGeminiClient(),
     historyManager.history,
@@ -1947,6 +1949,9 @@ export const AppContainer = (props: AppContainerProps) => {
       isFeedbackDialogOpen,
       // Per-task token tracking
       taskStartTokens,
+      // Real-time token display
+      streamingResponseLengthRef,
+      isReceivingContent,
       // Prompt suggestion
       promptSuggestion,
       dismissPromptSuggestion,
@@ -2054,6 +2059,9 @@ export const AppContainer = (props: AppContainerProps) => {
       isFeedbackDialogOpen,
       // Per-task token tracking
       taskStartTokens,
+      // Real-time token display
+      streamingResponseLengthRef,
+      isReceivingContent,
       // Prompt suggestion
       promptSuggestion,
       dismissPromptSuggestion,

packages/cli/src/ui/components/Composer.test.tsx

@@ -112,6 +112,9 @@ const createMockUIState = (overrides: Partial<UIState> = {}): UIState =>
     nightly: false,
     isTrustedFolder: true,
     taskStartTokens: 0,
+    streamingResponseLengthRef: { current: 0 },
+    isReceivingContent: false,
+    pendingGeminiHistoryItems: [],
     ...overrides,
   }) as UIState;
 

packages/cli/src/ui/components/Composer.tsx

@@ -15,10 +15,11 @@ import { useUIState } from '../contexts/UIStateContext.js';
 import { useUIActions } from '../contexts/UIActionsContext.js';
 import { useVimMode } from '../contexts/VimModeContext.js';
 import { useConfig } from '../contexts/ConfigContext.js';
-import { StreamingState } from '../types.js';
+import { StreamingState, type HistoryItemToolGroup } from '../types.js';
 import { ConfigInitDisplay } from '../components/ConfigInitDisplay.js';
 import { FeedbackDialog } from '../FeedbackDialog.js';
 import { t } from '../../i18n/index.js';
+import { useAnimationFrame } from '../hooks/useAnimationFrame.js';
 
 export const Composer = () => {
   const config = useConfig();
@@ -27,17 +28,46 @@ export const Composer = () => {
   const uiActions = useUIActions();
   const { vimEnabled } = useVimMode();
 
-  const { showAutoAcceptIndicator, sessionStats, taskStartTokens } = uiState;
+  const {
+    showAutoAcceptIndicator,
+    streamingResponseLengthRef,
+    isReceivingContent,
+  } = uiState;
 
-  const tokens = Object.values(sessionStats.metrics?.models ?? {}).reduce(
-    (acc, model) => ({
-      prompt: acc.prompt + (model.tokens?.prompt ?? 0),
-      candidates: acc.candidates + (model.tokens?.candidates ?? 0),
-    }),
-    { prompt: 0, candidates: 0 },
+  // --- Real-time token estimation ---
+  // Poll streamingResponseLengthRef at 50ms; only re-renders when the value
+  // actually changes, avoiding unnecessary empty renders.
+  const isStreaming =
+    uiState.streamingState === StreamingState.Responding ||
+    uiState.streamingState === StreamingState.WaitingForConfirmation;
+  const streamingChars = useAnimationFrame(
+    streamingResponseLengthRef,
+    isStreaming ? 50 : null,
   );
+  const estimatedStreamingTokens = Math.round(streamingChars / 4);
 
-  const taskTokens = tokens.candidates - taskStartTokens;
+  // Aggregate agent tool tokens from executing tool calls
+  let agentTokens = 0;
+  for (const item of uiState.pendingGeminiHistoryItems ?? []) {
+    if (item.type === 'tool_group') {
+      const toolGroup = item as HistoryItemToolGroup;
+      for (const tool of toolGroup.tools) {
+        const display = tool.resultDisplay;
+        if (
+          typeof display === 'object' &&
+          display !== null &&
+          'type' in display &&
+          display.type === 'task_execution' &&
+          'tokenCount' in display &&
+          typeof display.tokenCount === 'number'
+        ) {
+          agentTokens += display.tokenCount;
+        }
+      }
+    }
+  }
+
+  const taskTokens = estimatedStreamingTokens + agentTokens;
 
   // State for keyboard shortcuts display toggle
   const [showShortcuts, setShowShortcuts] = useState(false);
@@ -75,6 +105,7 @@ export const Composer = () => {
           }
           elapsedTime={uiState.elapsedTime}
           candidatesTokens={taskTokens}
+          isReceivingContent={isReceivingContent}
         />
       )}
 

packages/cli/src/ui/components/LoadingIndicator.test.tsx

@@ -374,5 +374,29 @@ describe('<LoadingIndicator />', () => {
       const output = lastFrame();
       expect(output).toContain('(5s · ↓ 5.4k tokens · esc to cancel)');
     });
+
+    it('should show ↑ arrow when waiting for API response', () => {
+      const { lastFrame } = renderWithContext(
+        <LoadingIndicator
+          {...defaultProps}
+          candidatesTokens={500}
+          isReceivingContent={false}
+        />,
+        StreamingState.Responding,
+      );
+      const output = lastFrame();
+      expect(output).toContain('↑ 500 tokens');
+      expect(output).not.toContain('↓');
+    });
+
+    it('should show ↓ arrow when receiving content (default)', () => {
+      const { lastFrame } = renderWithContext(
+        <LoadingIndicator {...defaultProps} candidatesTokens={500} />,
+        StreamingState.Responding,
+      );
+      const output = lastFrame();
+      expect(output).toContain('↓ 500 tokens');
+      expect(output).not.toContain('↑');
+    });
   });
 });

packages/cli/src/ui/components/LoadingIndicator.tsx

@@ -22,6 +22,12 @@ interface LoadingIndicatorProps {
   rightContent?: React.ReactNode;
   thought?: ThoughtSummary | null;
   candidatesTokens?: number;
+  /**
+   * True when receiving content (shows ↓ arrow), false when waiting for API
+   * response (shows ↑ arrow).
+   * @default true
+   */
+  isReceivingContent?: boolean;
 }
 
 export const LoadingIndicator: React.FC<LoadingIndicatorProps> = ({
@@ -30,6 +36,7 @@ export const LoadingIndicator: React.FC<LoadingIndicatorProps> = ({
   rightContent,
   thought,
   candidatesTokens,
+  isReceivingContent = true,
 }) => {
   const streamingState = useStreamingContext();
   const { columns: terminalWidth } = useTerminalSize();
@@ -43,12 +50,13 @@ export const LoadingIndicator: React.FC<LoadingIndicatorProps> = ({
 
   const outputTokens = candidatesTokens ?? 0;
   const showTokens = !isNarrow && outputTokens > 0;
+  const tokenArrow = isReceivingContent ? '↓' : '↑';
 
   const timeStr =
     elapsedTime < 60 ? `${elapsedTime}s` : formatDuration(elapsedTime * 1000);
 
   const tokenStr = showTokens
-    ? ` · ↓ ${formatTokenCount(outputTokens)} tokens`
+    ? ` · ${tokenArrow} ${formatTokenCount(outputTokens)} tokens`
     : '';
 
   const cancelAndTimerContent =

packages/cli/src/ui/contexts/UIStateContext.tsx

@@ -143,6 +143,10 @@ export interface UIState {
   isFeedbackDialogOpen: boolean;
   // Per-task token tracking
   taskStartTokens: number;
+  // Real-time token display: ref to streaming output char length (polled, not state)
+  streamingResponseLengthRef: React.RefObject<number>;
+  // True = receiving content (↓), false = waiting for API response (↑)
+  isReceivingContent: boolean;
   // Prompt suggestion
   promptSuggestion: string | null;
   /** Dismiss prompt suggestion (clears state, aborts speculation) */

packages/cli/src/ui/hooks/useAnimationFrame.ts

@@ -0,0 +1,50 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useEffect, useRef, useState } from 'react';
+
+/**
+ * Hook that polls a ref at a fixed interval and triggers a re-render only
+ * when the ref's value has changed. This avoids the cost of unconditional
+ * re-renders while still providing smooth animation-style updates.
+ *
+ * Pass `null` as intervalMs to pause polling entirely.
+ *
+ * @param watchRef - The ref to poll for changes.
+ * @param intervalMs - How often to check (ms), or null to pause.
+ * @returns The latest value read from the ref.
+ */
+export function useAnimationFrame(
+  watchRef: React.RefObject<number>,
+  intervalMs: number | null = 50,
+): number {
+  const [value, setValue] = useState(() => watchRef.current);
+  const lastSeen = useRef(watchRef.current);
+
+  useEffect(() => {
+    if (intervalMs === null) return;
+
+    // Re-sync when the interval resumes or when the ref value changed
+    // externally (e.g. ref reset to 0 at new turn start while paused).
+    const current = watchRef.current;
+    if (current !== lastSeen.current) {
+      lastSeen.current = current;
+      setValue(current);
+    }
+
+    const id = setInterval(() => {
+      const now = watchRef.current;
+      if (now !== lastSeen.current) {
+        lastSeen.current = now;
+        setValue(now);
+      }
+    }, intervalMs);
+
+    return () => clearInterval(id);
+  }, [watchRef, intervalMs]);
+
+  return value;
+}

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -238,6 +238,12 @@ export const useGeminiStream = (
   );
   const processedMemoryToolsRef = useRef<Set<string>>(new Set());
   const modelOverrideRef = useRef<string | undefined>(undefined);
+  // --- Real-time token display ---
+  // Accumulates output character count across the whole turn (not per API call).
+  // Uses a ref to avoid re-renders on every text_delta.
+  const streamingResponseLengthRef = useRef(0);
+  // Tracks whether we are receiving content (↓) or waiting for API (↑).
+  const [isReceivingContent, setIsReceivingContent] = useState(false);
   const {
     startNewPrompt,
     getPromptCount,
@@ -645,6 +651,9 @@ export const useGeminiStream = (
         // Prevents additional output after a user initiated cancel.
         return '';
       }
+      // Track output chars for real-time token estimation & mark as receiving.
+      streamingResponseLengthRef.current += eventValue.length;
+      setIsReceivingContent(true);
       let newGeminiMessageBuffer = currentGeminiMessageBuffer + eventValue;
       if (
         pendingHistoryItemRef.current?.type !== 'gemini' &&
@@ -1354,6 +1363,13 @@ export const useGeminiStream = (
 
         setIsResponding(true);
         setInitError(null);
+        // Entering "requesting" phase — no content yet for this API call.
+        setIsReceivingContent(false);
+        // Reset char counter only on new user queries; tool-result continuations
+        // keep accumulating so the token count only goes up within a turn.
+        if (submitType !== SendMessageType.ToolResult) {
+          streamingResponseLengthRef.current = 0;
+        }
 
         try {
           const stream = geminiClient.sendMessageStream(
@@ -1853,5 +1869,7 @@ export const useGeminiStream = (
     handleApprovalModeChange,
     activePtyId,
     loopDetectionConfirmationRequest,
+    streamingResponseLengthRef,
+    isReceivingContent,
   };
 };

packages/core/src/tools/agent.ts

@@ -32,6 +32,7 @@ import type {
   AgentFinishEvent,
   AgentErrorEvent,
   AgentApprovalRequestEvent,
+  AgentUsageEvent,
 } from '../agents/runtime/agent-events.js';
 import { BuiltinAgentRegistry } from '../subagents/builtin-agents.js';
 import { createDebugLogger } from '../utils/debugLogger.js';
@@ -484,6 +485,21 @@ class AgentToolInvocation extends BaseToolInvocation<AgentParams, ToolResult> {
       );
     });
 
+    // Track real-time token consumption from subagent API calls.
+    // Each USAGE_METADATA event carries the cumulative totalTokenCount for the
+    // subagent session, so we replace (not accumulate) on every event.
+    this.eventEmitter.on(
+      AgentEventType.USAGE_METADATA,
+      (...args: unknown[]) => {
+        const event = args[0] as AgentUsageEvent;
+        const total =
+          event.usage?.totalTokenCount ?? event.usage?.promptTokenCount ?? 0;
+        if (total > 0) {
+          this.updateDisplay({ tokenCount: total }, updateOutput);
+        }
+      },
+    );
+
     // Indicate when a tool call is waiting for approval
     this.eventEmitter.on(
       AgentEventType.TOOL_WAITING_APPROVAL,

packages/core/src/tools/tools.ts

@@ -499,6 +499,8 @@ export interface AgentResultDisplay {
   terminateReason?: string;
   result?: string;
   executionSummary?: AgentStatsSummary;
+  /** Real-time token count during execution (input + output). */
+  tokenCount?: number;
 
   // If the subagent is awaiting approval for a tool call,
   // this contains the confirmation details for inline UI rendering.