From 5cc11ac0bf2fc55acbd81c372de487e3fc156558 Mon Sep 17 00:00:00 2001
From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com>
Date: Thu, 2 Oct 2025 17:15:54 +0530
Subject: [PATCH] feat(message-handling): fix repeated first token, think tag
 handling

closes #889
---
 src/lib/hooks/useChat.tsx | 38 +++++++++++++++++++++-----------------
 1 file changed, 21 insertions(+), 17 deletions(-)

diff --git a/src/lib/hooks/useChat.tsx b/src/lib/hooks/useChat.tsx
index a28a53d..479fc79 100644
--- a/src/lib/hooks/useChat.tsx
+++ b/src/lib/hooks/useChat.tsx
@@ -399,6 +399,7 @@ export const ChatProvider = ({
           (m, j) =>
             j > i &&
             m.role === 'source' &&
+            m.sources &&
             (nextUserMessageIndex === -1 || j < nextUserMessageIndex),
         ) as SourceMessage | undefined;
 
@@ -417,7 +418,7 @@ export const ChatProvider = ({
             const closeThinkTag =
               processedMessage.match(/<\/think>/g)?.length || 0;
 
-            if (openThinkTag > closeThinkTag) {
+            if (openThinkTag && !closeThinkTag) {
               processedMessage += '</think> <a> </a>';
             }
           }
@@ -426,7 +427,11 @@ export const ChatProvider = ({
             thinkingEnded = true;
           }
 
-          if (sourceMessage && sourceMessage.sources.length > 0) {
+          if (
+            sourceMessage &&
+            sourceMessage.sources &&
+            sourceMessage.sources.length > 0
+          ) {
             processedMessage = processedMessage.replace(
               citationRegex,
               (_, capturedContent: string) => {
@@ -635,23 +640,22 @@ export const ChatProvider = ({
             },
           ]);
           added = true;
+          setMessageAppeared(true);
+        } else {
+          setMessages((prev) =>
+            prev.map((message) => {
+              if (
+                message.messageId === data.messageId &&
+                message.role === 'assistant'
+              ) {
+                return { ...message, content: message.content + data.data };
+              }
+
+              return message;
+            }),
+          );
         }
-
-        setMessages((prev) =>
-          prev.map((message) => {
-            if (
-              message.messageId === data.messageId &&
-              message.role === 'assistant'
-            ) {
-              return { ...message, content: message.content + data.data };
-            }
-
-            return message;
-          }),
-        );
-
         recievedMessage += data.data;
-        setMessageAppeared(true);
       }
 
       if (data.type === 'messageEnd') {