Token state not showing on load, or after message is finished. (#5606)

ui/desktop/src/api/types.gen.ts

@@ -374,6 +374,7 @@ export type MessageEvent = {
     type: 'Error';
 } | {
     reason: string;
+    token_state: TokenState;
     type: 'Finish';
 } | {
     mode: string;

ui/desktop/src/components/BaseChat.tsx

@@ -443,12 +443,12 @@ function BaseChatContent({
             commandHistory={commandHistory}
             initialValue={input || ''}
             setView={setView}
-            totalTokens={tokenState?.totalTokens ?? sessionTokenCount}
+            totalTokens={tokenState?.totalTokens || sessionTokenCount}
             accumulatedInputTokens={
-              tokenState?.accumulatedInputTokens ?? sessionInputTokens ?? localInputTokens
+              tokenState?.accumulatedInputTokens || sessionInputTokens || localInputTokens
             }
             accumulatedOutputTokens={
-              tokenState?.accumulatedOutputTokens ?? sessionOutputTokens ?? localOutputTokens
+              tokenState?.accumulatedOutputTokens || sessionOutputTokens || localOutputTokens
             }
             droppedFiles={droppedFiles}
             onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing

ui/desktop/src/hooks/useChatEngine.ts

@@ -213,7 +213,6 @@ export const useChatEngine = ({
 
   // Update token counts when session changes from the message stream
   useEffect(() => {
-    console.log('Session received:', session);
     if (session) {
       setSessionTokenCount(session.total_tokens || 0);
       setSessionInputTokens(session.accumulated_input_tokens || 0);

ui/desktop/src/hooks/useMessageStream.ts

@@ -37,7 +37,7 @@ export interface NotificationEvent {
 type MessageEvent =
   | { type: 'Message'; message: Message; token_state: TokenState }
   | { type: 'Error'; error: string }
-  | { type: 'Finish'; reason: string }
+  | { type: 'Finish'; reason: string; token_state: TokenState }
   | { type: 'ModelChange'; model: string; mode: string }
   | { type: 'UpdateConversation'; conversation: Conversation }
   | NotificationEvent;
@@ -368,6 +368,8 @@ export function useMessageStream({
                   }
 
                   case 'Finish': {
+                    setTokenState(parsedEvent.token_state);
+
                     if (onFinish && currentMessages.length > 0) {
                       const lastMessage = currentMessages[currentMessages.length - 1];
                       onFinish(lastMessage, parsedEvent.reason);