From 15aa8187f16ec232d1617fc4c5d2118865b8845b Mon Sep 17 00:00:00 2001
From: David Katz <dkatz@squareup.com>
Date: Tue, 16 Dec 2025 11:35:19 -0500
Subject: [PATCH] Fix tokenState loading on new sessions (#6129)

---
 ui/desktop/src/hooks/useChatStream.ts | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts
index a7260e96..01555361 100644
--- a/ui/desktop/src/hooks/useChatStream.ts
+++ b/ui/desktop/src/hooks/useChatStream.ts
@@ -212,6 +212,14 @@ export function useChatStream({
     if (cached) {
       setSession(cached.session);
       updateMessages(cached.messages);
+      setTokenState({
+        inputTokens: cached.session?.input_tokens ?? 0,
+        outputTokens: cached.session?.output_tokens ?? 0,
+        totalTokens: cached.session?.total_tokens ?? 0,
+        accumulatedInputTokens: cached.session?.accumulated_input_tokens ?? 0,
+        accumulatedOutputTokens: cached.session?.accumulated_output_tokens ?? 0,
+        accumulatedTotalTokens: cached.session?.accumulated_total_tokens ?? 0,
+      });
       setChatState(ChatState.Idle);
       return;
     }
@@ -241,6 +249,14 @@ export function useChatStream({
         const session = response.data;
         setSession(session);
         updateMessages(session?.conversation || []);
+        setTokenState({
+          inputTokens: session?.input_tokens ?? 0,
+          outputTokens: session?.output_tokens ?? 0,
+          totalTokens: session?.total_tokens ?? 0,
+          accumulatedInputTokens: session?.accumulated_input_tokens ?? 0,
+          accumulatedOutputTokens: session?.accumulated_output_tokens ?? 0,
+          accumulatedTotalTokens: session?.accumulated_total_tokens ?? 0,
+        });
         setChatState(ChatState.Idle);
         onSessionLoaded?.();
       } catch (error) {