From f72a7e6d8717eb51c3664b6b47f451a7de369a24 Mon Sep 17 00:00:00 2001
From: chenchen <chenchen@xinghanlab.com>
Date: Mon, 18 May 2026 23:03:17 +0800
Subject: [PATCH] fix: prevent cross-turn assistant message merging in chat UI

appendAssistantTextMessage was merging consecutive assistant_text messages
regardless of whether they belonged to different conversation turns. Added
turnComplete flag that gets set on message_complete, error, and status(idle)
events, preventing subsequent turns from merging into the previous response.
---
 cc-haha/desktop/src/stores/chatStore.ts | 22 ++++++++++++++++++----
 cc-haha/desktop/src/types/chat.ts       |  2 +-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/cc-haha/desktop/src/stores/chatStore.ts b/cc-haha/desktop/src/stores/chatStore.ts
index 9516b79..8ba71a6 100644
--- a/cc-haha/desktop/src/stores/chatStore.ts
+++ b/cc-haha/desktop/src/stores/chatStore.ts
@@ -174,7 +174,7 @@ function appendAssistantTextMessage(
   if (!content.trim()) return messages
 
   const last = messages[messages.length - 1]
-  if (last?.type === 'assistant_text') {
+  if (last?.type === 'assistant_text' && !last.turnComplete) {
     const merged: UIMessage = {
       ...last,
       content: last.content + content,
@@ -195,6 +195,17 @@ function appendAssistantTextMessage(
   ]
 }
 
+function markLastAssistantComplete(messages: UIMessage[]): UIMessage[] {
+  const lastIdx = messages.length - 1
+  const last = messages[lastIdx]
+  if (last?.type === 'assistant_text' && !last.turnComplete) {
+    const updated = [...messages]
+    updated[lastIdx] = { ...last, turnComplete: true }
+    return updated
+  }
+  return messages
+}
+
 /** Helper: immutably update a specific session within the sessions record */
 function updateSessionIn(
   sessions: Record<string, PerSessionState>,
@@ -562,7 +573,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
             ...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}),
             ...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}),
             ...(shouldFlush ? {
-              messages: appendAssistantTextMessage(session.messages, pendingText, Date.now()),
+              messages: markLastAssistantComplete(appendAssistantTextMessage(session.messages, pendingText, Date.now())),
               streamingText: '',
             } : pendingText !== session.streamingText ? { streamingText: pendingText } : {}),
           }
@@ -727,11 +738,14 @@ export const useChatStore = create<ChatStore>((set, get) => ({
         const text = `${session.streamingText}${consumePendingDelta(sessionId)}`
         if (text.trim()) {
           update((s) => ({
-            messages: appendAssistantTextMessage(s.messages, text, Date.now()),
+            messages: markLastAssistantComplete(appendAssistantTextMessage(s.messages, text, Date.now())),
             streamingText: '',
           }))
         } else if (text !== session.streamingText) {
           update(() => ({ streamingText: text }))
+          update((s) => ({ messages: markLastAssistantComplete(s.messages) }))
+        } else {
+          update((s) => ({ messages: markLastAssistantComplete(s.messages) }))
         }
         if (session.elapsedTimer) clearInterval(session.elapsedTimer)
         update(() => ({
@@ -750,7 +764,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
           const pendingText = `${s.streamingText}${consumePendingDelta(sessionId)}`
           let newMessages = s.messages
           if (pendingText.trim()) {
-            newMessages = appendAssistantTextMessage(newMessages, pendingText, Date.now())
+            newMessages = markLastAssistantComplete(appendAssistantTextMessage(newMessages, pendingText, Date.now()))
           }
           newMessages = [...newMessages, { id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }]
           return {
diff --git a/cc-haha/desktop/src/types/chat.ts b/cc-haha/desktop/src/types/chat.ts
index c0c0247..6aa1eb1 100644
--- a/cc-haha/desktop/src/types/chat.ts
+++ b/cc-haha/desktop/src/types/chat.ts
@@ -158,7 +158,7 @@ export type TaskSummaryItem = {
 
 export type UIMessage =
   | { id: string; type: 'user_text'; content: string; timestamp: number; attachments?: UIAttachment[]; pending?: boolean }
-  | { id: string; type: 'assistant_text'; content: string; timestamp: number; model?: string }
+  | { id: string; type: 'assistant_text'; content: string; timestamp: number; model?: string; turnComplete?: boolean }
   | { id: string; type: 'thinking'; content: string; timestamp: number }
   | { id: string; type: 'tool_use'; toolName: string; toolUseId: string; input: unknown; timestamp: number; parentToolUseId?: string }
   | { id: string; type: 'tool_result'; toolUseId: string; content: unknown; isError: boolean; timestamp: number; parentToolUseId?: string }