feat(ui): implement per-task token tracking in LoadingIndicator

packages/cli/src/ui/AppContainer.tsx

@@ -1022,10 +1022,16 @@ export const AppContainer = (props: AppContainerProps) => {
     [historyManager, setShowCommandMigrationNudge, config.storage],
   );
 
-  const { elapsedTime, currentLoadingPhrase } = useLoadingIndicator(
-    streamingState,
-    settings.merged.ui?.customWittyPhrases,
-  );
+  const currentCandidatesTokens = Object.values(
+    sessionStats.metrics?.models ?? {},
+  ).reduce((acc, model) => acc + (model.tokens?.candidates ?? 0), 0);
+
+  const { elapsedTime, currentLoadingPhrase, taskStartTokens } =
+    useLoadingIndicator(
+      streamingState,
+      settings.merged.ui?.customWittyPhrases,
+      currentCandidatesTokens,
+    );
 
   useAttentionNotifications({
     isFocused,
@@ -1430,6 +1436,8 @@ export const AppContainer = (props: AppContainerProps) => {
       isMcpDialogOpen,
       // Feedback dialog
       isFeedbackDialogOpen,
+      // Per-task token tracking
+      taskStartTokens,
     }),
     [
       isThemeDialogOpen,
@@ -1524,6 +1532,8 @@ export const AppContainer = (props: AppContainerProps) => {
       isMcpDialogOpen,
       // Feedback dialog
       isFeedbackDialogOpen,
+      // Per-task token tracking
+      taskStartTokens,
     ],
   );
 

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

@@ -111,6 +111,7 @@ const createMockUIState = (overrides: Partial<UIState> = {}): UIState =>
     debugMessage: '',
     nightly: false,
     isTrustedFolder: true,
+    taskStartTokens: 0,
     ...overrides,
   }) as UIState;
 

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

@@ -27,7 +27,7 @@ export const Composer = () => {
   const uiActions = useUIActions();
   const { vimEnabled } = useVimMode();
 
-  const { showAutoAcceptIndicator, sessionStats } = uiState;
+  const { showAutoAcceptIndicator, sessionStats, taskStartTokens } = uiState;
 
   const tokens = Object.values(sessionStats.metrics?.models ?? {}).reduce(
     (acc, model) => ({
@@ -37,6 +37,8 @@ export const Composer = () => {
     { prompt: 0, candidates: 0 },
   );
 
+  const taskTokens = tokens.candidates - taskStartTokens;
+
   // State for keyboard shortcuts display toggle
   const [showShortcuts, setShowShortcuts] = useState(false);
   const handleToggleShortcuts = useCallback(() => {
@@ -72,7 +74,7 @@ export const Composer = () => {
               : uiState.currentLoadingPhrase
           }
           elapsedTime={uiState.elapsedTime}
-          candidatesTokens={tokens.candidates}
+          candidatesTokens={taskTokens}
         />
       )}
 

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

@@ -131,6 +131,8 @@ export interface UIState {
   isMcpDialogOpen: boolean;
   // Feedback dialog
   isFeedbackDialogOpen: boolean;
+  // Per-task token tracking
+  taskStartTokens: number;
 }
 
 export const UIStateContext = createContext<UIState | null>(null);

packages/cli/src/ui/hooks/useLoadingIndicator.test.ts

@@ -133,4 +133,119 @@ describe('useLoadingIndicator', () => {
     });
     expect(result.current.elapsedTime).toBe(0);
   });
+
+  describe('token tracking', () => {
+    it('should capture token snapshot when task starts', () => {
+      const { result, rerender } = renderHook(
+        ({ streamingState, currentCandidatesTokens }) =>
+          useLoadingIndicator(
+            streamingState,
+            undefined,
+            currentCandidatesTokens,
+          ),
+        {
+          initialProps: {
+            streamingState: StreamingState.Idle,
+            currentCandidatesTokens: 100,
+          },
+        },
+      );
+
+      expect(result.current.taskStartTokens).toBe(0);
+
+      act(() => {
+        rerender({
+          streamingState: StreamingState.Responding,
+          currentCandidatesTokens: 100,
+        });
+      });
+
+      expect(result.current.taskStartTokens).toBe(100);
+    });
+
+    it('should reset token snapshot when transitioning from Responding to Idle', async () => {
+      const { result, rerender } = renderHook(
+        ({ streamingState, currentCandidatesTokens }) =>
+          useLoadingIndicator(
+            streamingState,
+            undefined,
+            currentCandidatesTokens,
+          ),
+        {
+          initialProps: {
+            streamingState: StreamingState.Idle,
+            currentCandidatesTokens: 0,
+          },
+        },
+      );
+
+      act(() => {
+        rerender({
+          streamingState: StreamingState.Responding,
+          currentCandidatesTokens: 0,
+        });
+      });
+      expect(result.current.taskStartTokens).toBe(0);
+
+      await act(async () => {
+        await vi.advanceTimersByTimeAsync(1000);
+        rerender({
+          streamingState: StreamingState.Responding,
+          currentCandidatesTokens: 500,
+        });
+      });
+
+      act(() => {
+        rerender({
+          streamingState: StreamingState.Idle,
+          currentCandidatesTokens: 500,
+        });
+      });
+
+      expect(result.current.taskStartTokens).toBe(0);
+    });
+
+    it('should reset token snapshot when transitioning from WaitingForConfirmation to Responding', async () => {
+      const { result, rerender } = renderHook(
+        ({ streamingState, currentCandidatesTokens }) =>
+          useLoadingIndicator(
+            streamingState,
+            undefined,
+            currentCandidatesTokens,
+          ),
+        {
+          initialProps: {
+            streamingState: StreamingState.Responding,
+            currentCandidatesTokens: 100,
+          },
+        },
+      );
+
+      expect(result.current.taskStartTokens).toBe(100);
+
+      await act(async () => {
+        await vi.advanceTimersByTimeAsync(5000);
+        rerender({
+          streamingState: StreamingState.Responding,
+          currentCandidatesTokens: 500,
+        });
+      });
+
+      act(() => {
+        rerender({
+          streamingState: StreamingState.WaitingForConfirmation,
+          currentCandidatesTokens: 500,
+        });
+      });
+
+      act(() => {
+        rerender({
+          streamingState: StreamingState.Responding,
+          currentCandidatesTokens: 500,
+        });
+      });
+
+      expect(result.current.taskStartTokens).toBe(500);
+    });
+  });
 });

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

@@ -7,11 +7,12 @@
 import { StreamingState } from '../types.js';
 import { useTimer } from './useTimer.js';
 import { usePhraseCycler } from './usePhraseCycler.js';
-import { useState, useEffect, useRef } from 'react'; // Added useRef
+import { useState, useEffect, useRef } from 'react';
 
 export const useLoadingIndicator = (
   streamingState: StreamingState,
   customWittyPhrases?: string[],
+  currentCandidatesTokens?: number,
 ) => {
   const [timerResetKey, setTimerResetKey] = useState(0);
   const isTimerActive = streamingState === StreamingState.Responding;
@@ -27,6 +28,7 @@ export const useLoadingIndicator = (
   );
 
   const [retainedElapsedTime, setRetainedElapsedTime] = useState(0);
+  const [taskStartTokens, setTaskStartTokens] = useState(0);
   const prevStreamingStateRef = useRef<StreamingState | null>(null);
 
   useEffect(() => {
@@ -35,21 +37,26 @@ export const useLoadingIndicator = (
       streamingState === StreamingState.Responding
     ) {
       setTimerResetKey((prevKey) => prevKey + 1);
-      setRetainedElapsedTime(0); // Clear retained time when going back to responding
+      setRetainedElapsedTime(0);
+      setTaskStartTokens(currentCandidatesTokens ?? 0);
     } else if (
       streamingState === StreamingState.Idle &&
       prevStreamingStateRef.current === StreamingState.Responding
     ) {
-      setTimerResetKey((prevKey) => prevKey + 1); // Reset timer when becoming idle from responding
+      setTimerResetKey((prevKey) => prevKey + 1);
       setRetainedElapsedTime(0);
+      setTaskStartTokens(0);
+    } else if (
+      streamingState === StreamingState.Responding &&
+      prevStreamingStateRef.current !== StreamingState.Responding
+    ) {
+      setTaskStartTokens(currentCandidatesTokens ?? 0);
     } else if (streamingState === StreamingState.WaitingForConfirmation) {
-      // Capture the time when entering WaitingForConfirmation
-      // elapsedTimeFromTimer will hold the last value from when isTimerActive was true.
       setRetainedElapsedTime(elapsedTimeFromTimer);
     }
 
     prevStreamingStateRef.current = streamingState;
-  }, [streamingState, elapsedTimeFromTimer]);
+  }, [streamingState, elapsedTimeFromTimer, currentCandidatesTokens]);
 
   return {
     elapsedTime:
@@ -57,5 +64,6 @@ export const useLoadingIndicator = (
         ? retainedElapsedTime
         : elapsedTimeFromTimer,
     currentLoadingPhrase,
+    taskStartTokens,
   };
 };