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

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);