From d7430f15941c5a93b9b7fcaedaa0af9af0720633 Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Wed, 5 Mar 2025 16:22:45 -0700
Subject: [PATCH] Update useChatStreamMessage.ts

---
 .../NewChatProvider/useChatStreamMessage.ts   | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
index 6d60689e8..9d3b6066e 100644
--- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
+++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
@@ -219,15 +219,18 @@ export const useChatStreamMessage = () => {
             reasoning.message_chunk !== null || reasoning.message_chunk !== undefined;
 
           initializeOrUpdateMessage(chat_id, message_id, (draft) => {
-            if (!draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId])
-              return;
-            const messageText = draft[chat_id].messages[message_id].reasoning_messages[
-              reasoningMessageId
-            ] as BusterChatMessageReasoning_text;
-            Object.assign(messageText, existingReasoningMessageText);
-            messageText.message = isStreaming
-              ? (existingReasoningMessageText?.message || '') + (reasoning.message_chunk || '')
-              : reasoning.message;
+            const reasoningMessage =
+              draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId];
+            if (!reasoningMessage) return;
+            const messageText = reasoningMessage as BusterChatMessageReasoning_text;
+
+            Object.assign(messageText, {
+              ...existingReasoningMessageText,
+              ...reasoning,
+              message: isStreaming
+                ? (existingReasoningMessageText?.message || '') + (reasoning.message_chunk || '')
+                : reasoning.message
+            });
           });
 
           break;