Stream token usage on every agent message (#5342)

ui/desktop/src/components/BaseChat.tsx

@@ -132,6 +132,7 @@ function BaseChatContent({
     sessionOutputTokens,
     localInputTokens,
     localOutputTokens,
+    tokenState,
     commandHistory,
     toolCallNotifications,
     sessionMetadata,
@@ -442,9 +443,13 @@ function BaseChatContent({
             commandHistory={commandHistory}
             initialValue={input || ''}
             setView={setView}
-            numTokens={sessionTokenCount}
-            inputTokens={sessionInputTokens || localInputTokens}
-            outputTokens={sessionOutputTokens || localOutputTokens}
+            totalTokens={tokenState?.totalTokens ?? sessionTokenCount}
+            accumulatedInputTokens={
+              tokenState?.accumulatedInputTokens ?? sessionInputTokens ?? localInputTokens
+            }
+            accumulatedOutputTokens={
+              tokenState?.accumulatedOutputTokens ?? sessionOutputTokens ?? localOutputTokens
+            }
             droppedFiles={droppedFiles}
             onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
             messages={messages}

ui/desktop/src/components/BaseChat2.tsx

@@ -72,6 +72,7 @@ function BaseChatContent({
     stopStreaming,
     sessionLoadError,
     setRecipeUserParams,
+    tokenState,
   } = useChatStream({
     sessionId,
     onStreamFinish,
@@ -281,9 +282,13 @@ function BaseChatContent({
             //commandHistory={commandHistory}
             initialValue={initialPrompt}
             setView={setView}
-            numTokens={session?.total_tokens || undefined}
-            inputTokens={session?.input_tokens || undefined}
-            outputTokens={session?.output_tokens || undefined}
+            totalTokens={tokenState?.totalTokens ?? session?.total_tokens ?? undefined}
+            accumulatedInputTokens={
+              tokenState?.accumulatedInputTokens ?? session?.accumulated_input_tokens ?? undefined
+            }
+            accumulatedOutputTokens={
+              tokenState?.accumulatedOutputTokens ?? session?.accumulated_output_tokens ?? undefined
+            }
             droppedFiles={droppedFiles}
             onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
             messages={messages}

ui/desktop/src/components/ChatInput.tsx

@@ -70,9 +70,9 @@ interface ChatInputProps {
   droppedFiles?: DroppedFile[];
   onFilesProcessed?: () => void; // Callback to clear dropped files after processing
   setView: (view: View) => void;
-  numTokens?: number;
-  inputTokens?: number;
-  outputTokens?: number;
+  totalTokens?: number;
+  accumulatedInputTokens?: number;
+  accumulatedOutputTokens?: number;
   messages?: Message[];
   sessionCosts?: {
     [key: string]: {
@@ -103,9 +103,9 @@ export default function ChatInput({
   droppedFiles = [],
   onFilesProcessed,
   setView,
-  numTokens,
-  inputTokens,
-  outputTokens,
+  totalTokens,
+  accumulatedInputTokens,
+  accumulatedOutputTokens,
   messages = [],
   disableAnimation = false,
   sessionCosts,
@@ -505,16 +505,16 @@ export default function ChatInput({
     clearAlerts();
 
     // Show alert when either there is registered token usage, or we know the limit
-    if ((numTokens && numTokens > 0) || (isTokenLimitLoaded && tokenLimit)) {
+    if ((totalTokens && totalTokens > 0) || (isTokenLimitLoaded && tokenLimit)) {
       addAlert({
         type: AlertType.Info,
         message: 'Context window',
         progress: {
-          current: numTokens || 0,
+          current: totalTokens || 0,
           total: tokenLimit,
         },
         showCompactButton: true,
-        compactButtonDisabled: !numTokens,
+        compactButtonDisabled: !totalTokens,
         onCompact: () => {
           window.dispatchEvent(new CustomEvent('hide-alert-popover'));
 
@@ -542,7 +542,7 @@ export default function ChatInput({
     }
     // We intentionally omit setView as it shouldn't trigger a re-render of alerts
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [numTokens, toolCount, tokenLimit, isTokenLimitLoaded, addAlert, clearAlerts]);
+  }, [totalTokens, toolCount, tokenLimit, isTokenLimitLoaded, addAlert, clearAlerts]);
 
   // Cleanup effect for component unmount - prevent memory leaks
   useEffect(() => {
@@ -1540,8 +1540,8 @@ export default function ChatInput({
             <>
               <div className="flex items-center h-full ml-1 mr-1">
                 <CostTracker
-                  inputTokens={inputTokens}
-                  outputTokens={outputTokens}
+                  inputTokens={accumulatedInputTokens}
+                  outputTokens={accumulatedOutputTokens}
                   sessionCosts={sessionCosts}
                 />
               </div>

ui/desktop/src/components/hub.tsx

@@ -78,9 +78,9 @@ export default function Hub({
         commandHistory={[]}
         initialValue=""
         setView={setView}
-        numTokens={0}
-        inputTokens={0}
-        outputTokens={0}
+        totalTokens={0}
+        accumulatedInputTokens={0}
+        accumulatedOutputTokens={0}
         droppedFiles={[]}
         onFilesProcessed={() => {}}
         messages={[]}