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

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

diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
index f31a4297a..7f3503f36 100644
--- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
+++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
@@ -157,19 +157,9 @@ export const useChatStreamMessage = () => {
       if (!response_message?.id) return;
 
       const responseMessageId = response_message.id;
-      const foundResponseMessage: BusterChatMessageResponse | undefined =
+      const existingMessage =
         chatRef.current[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId];
-      const isNewMessage = !foundResponseMessage;
-
-      if (response_message.type === 'text') {
-        const existingMessage =
-          (foundResponseMessage as BusterChatResponseMessage_text)?.message || '';
-        const isStreaming =
-          response_message.message_chunk !== undefined || response_message.message_chunk !== null;
-        if (isStreaming) {
-          response_message.message = existingMessage + response_message.message_chunk;
-        }
-      }
+      const isNewMessage = !existingMessage;
 
       if (isNewMessage) {
         initializeOrUpdateMessage(chat_id, message_id, (draft) => {
@@ -182,6 +172,28 @@ export const useChatStreamMessage = () => {
         });
       }
 
+      if (response_message.type === 'text') {
+        const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text;
+        const isStreaming =
+          response_message.message_chunk !== undefined || response_message.message_chunk !== null;
+
+        initializeOrUpdateMessage(chat_id, message_id, (draft) => {
+          const responseMessage =
+            draft[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId];
+          if (!responseMessage) return;
+          const messageText = responseMessage as BusterChatMessageReasoning_text;
+
+          Object.assign(messageText, {
+            ...existingResponseMessageText,
+            ...response_message,
+            message: isStreaming
+              ? (existingResponseMessageText?.message || '') +
+                (response_message.message_chunk || '')
+              : response_message.message
+          });
+        });
+      }
+
       const response_messages = chatRef.current[chat_id]?.messages?.[message_id]?.response_messages;
       const response_message_ids =
         chatRef.current[chat_id]?.messages?.[message_id]?.response_message_ids;