fix: merge split assistant text blocks in desktop chat

desktop/src/stores/chatStore.test.ts

@@ -157,6 +157,30 @@ describe('chatStore history mapping', () => {
     expect(mapped[3]).toMatchObject({ parentToolUseId: 'agent-1' })
   })
 
+  it('merges consecutive assistant text blocks when restoring transcript history', () => {
+    const messages: MessageEntry[] = [
+      {
+        id: 'assistant-merge-1',
+        type: 'assistant',
+        timestamp: '2026-04-06T00:00:00.000Z',
+        model: 'opus',
+        content: [
+          { type: 'text', text: '第一段：Windows 下的桌面端输出。' },
+          { type: 'text', text: '\r\n第二段：刷新后也不应该被拆开。' },
+        ],
+      },
+    ]
+
+    const mapped = mapHistoryMessagesToUiMessages(messages)
+
+    expect(mapped).toMatchObject([
+      {
+        type: 'assistant_text',
+        content: '第一段：Windows 下的桌面端输出。\r\n第二段：刷新后也不应该被拆开。',
+      },
+    ])
+  })
+
   it('surfaces teammate prompt content when mapping member transcript history', () => {
     const messages: MessageEntry[] = [
       {
@@ -532,6 +556,68 @@ describe('chatStore history mapping', () => {
     ).toBe('permission_pending')
   })
 
+  it('keeps delayed text blocks from one streamed assistant turn in a single message', () => {
+    vi.useFakeTimers()
+
+    useChatStore.setState({
+      sessions: {
+        [TEST_SESSION_ID]: {
+          messages: [],
+          chatState: 'idle',
+          connectionState: 'connected',
+          streamingText: '',
+          streamingToolInput: '',
+          activeToolUseId: null,
+          activeToolName: null,
+          activeThinkingId: null,
+          pendingPermission: null,
+          pendingComputerUsePermission: null,
+          tokenUsage: { input_tokens: 0, output_tokens: 0 },
+          elapsedSeconds: 0,
+          statusVerb: '',
+          slashCommands: [],
+          agentTaskNotifications: {},
+          elapsedTimer: null,
+        },
+      },
+    })
+
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'content_start',
+      blockType: 'text',
+    })
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'content_delta',
+      text: '第一段：先到达。',
+    })
+    vi.advanceTimersByTime(60)
+
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'content_start',
+      blockType: 'text',
+    })
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'content_delta',
+      text: '\r\n第二段：稍后到达，但仍属于同一轮回复。',
+    })
+    vi.advanceTimersByTime(60)
+
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'message_complete',
+      usage: { input_tokens: 1, output_tokens: 2 },
+    })
+
+    expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([
+      {
+        type: 'assistant_text',
+        content: '第一段：先到达。\r\n第二段：稍后到达，但仍属于同一轮回复。',
+      },
+    ])
+
+    vi.runOnlyPendingTimers()
+    vi.useRealTimers()
+  })
+
   it('sends Computer Use approval payloads back over websocket', () => {
     useChatStore.setState({
       sessions: {

desktop/src/stores/chatStore.ts

@@ -128,6 +128,46 @@ const nextId = () => `msg-${++msgCounter}-${Date.now()}`
 let pendingDelta = ''
 let flushTimer: ReturnType<typeof setTimeout> | null = null
 
+function consumePendingDelta(): string {
+  if (flushTimer) {
+    clearTimeout(flushTimer)
+    flushTimer = null
+  }
+  const text = pendingDelta
+  pendingDelta = ''
+  return text
+}
+
+function appendAssistantTextMessage(
+  messages: UIMessage[],
+  content: string,
+  timestamp: number,
+  model?: string,
+): UIMessage[] {
+  if (!content.trim()) return messages
+
+  const last = messages[messages.length - 1]
+  if (last?.type === 'assistant_text') {
+    const merged: UIMessage = {
+      ...last,
+      content: last.content + content,
+      ...(model ?? last.model ? { model: model ?? last.model } : {}),
+    }
+    return [...messages.slice(0, -1), merged]
+  }
+
+  return [
+    ...messages,
+    {
+      id: nextId(),
+      type: 'assistant_text',
+      content,
+      timestamp,
+      ...(model ? { model } : {}),
+    },
+  ]
+}
+
 /** Helper: immutably update a specific session within the sessions record */
 function updateSessionIn(
   sessions: Record<string, PerSessionState>,
@@ -200,8 +240,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
     if (session?.elapsedTimer) clearInterval(session.elapsedTimer)
     if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
     if (pendingDelta) {
-      const text = pendingDelta
-      pendingDelta = ''
+      const text = consumePendingDelta()
       set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
     }
     wsManager.disconnect(sessionId)
@@ -241,20 +280,11 @@ export const useChatStore = create<ChatStore>((set, get) => ({
         clearTimeout(flushTimer)
         flushTimer = null
       }
-      const bufferedDelta = pendingDelta
-      pendingDelta = ''
-      const pendingAssistantText = `${session.streamingText}${bufferedDelta}`.trim()
+      const bufferedDelta = consumePendingDelta()
+      const pendingAssistantText = `${session.streamingText}${bufferedDelta}`
 
-      const newMessages = pendingAssistantText
-        ? [
-            ...session.messages,
-            {
-              id: nextId(),
-              type: 'assistant_text' as const,
-              content: pendingAssistantText,
-              timestamp: Date.now(),
-            },
-          ]
+      const newMessages = pendingAssistantText.trim()
+        ? appendAssistantTextMessage(session.messages, pendingAssistantText, Date.now())
         : [...session.messages]
       if (!isMemberSession && allTasksDone) {
         newMessages.push({
@@ -357,8 +387,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
     wsManager.send(sessionId, { type: 'stop_generation' })
     if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
     if (pendingDelta) {
-      const text = pendingDelta
-      pendingDelta = ''
+      const text = consumePendingDelta()
       set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
     }
     set((s) => {
@@ -481,17 +510,17 @@ export const useChatStore = create<ChatStore>((set, get) => ({
 
       case 'status':
         update((session) => {
-          const pendingText = session.streamingText.trim()
-          const shouldFlush = pendingText && session.chatState === 'streaming' && msg.state !== 'streaming'
+          const pendingText = `${session.streamingText}${consumePendingDelta()}`
+          const shouldFlush = pendingText.trim() && session.chatState === 'streaming' && msg.state !== 'streaming'
           return {
             chatState: msg.state,
             ...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}),
             ...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}),
             ...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}),
             ...(shouldFlush ? {
-              messages: [...session.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }],
+              messages: appendAssistantTextMessage(session.messages, pendingText, Date.now()),
               streamingText: '',
-            } : {}),
+            } : pendingText !== session.streamingText ? { streamingText: pendingText } : {}),
           }
         })
         if (msg.state === 'idle') {
@@ -508,15 +537,19 @@ export const useChatStore = create<ChatStore>((set, get) => ({
       case 'content_start': {
         const session = get().sessions[sessionId]
         if (!session) break
-        const pendingText = session.streamingText.trim()
-        if (pendingText) {
+        const pendingText = `${session.streamingText}${consumePendingDelta()}`
+        if (msg.blockType !== 'text' && pendingText.trim()) {
           update((s) => ({
-            messages: [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }],
+            messages: appendAssistantTextMessage(s.messages, pendingText, Date.now()),
             streamingText: '',
           }))
         }
         if (msg.blockType === 'text') {
-          update(() => ({ streamingText: '', chatState: 'streaming', activeThinkingId: null }))
+          update((s) => ({
+            ...(pendingText !== s.streamingText ? { streamingText: pendingText } : {}),
+            chatState: 'streaming',
+            activeThinkingId: null,
+          }))
         } else if (msg.blockType === 'tool_use') {
           update(() => ({
             activeToolUseId: msg.toolUseId ?? null,
@@ -546,9 +579,9 @@ export const useChatStore = create<ChatStore>((set, get) => ({
 
       case 'thinking':
         update((s) => {
-          const pendingText = s.streamingText.trim()
-          const base = pendingText
-            ? [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }]
+          const pendingText = `${s.streamingText}${consumePendingDelta()}`
+          const base = pendingText.trim()
+            ? appendAssistantTextMessage(s.messages, pendingText, Date.now())
             : s.messages
           const last = base[base.length - 1]
           if (last && last.type === 'thinking') {
@@ -643,12 +676,14 @@ export const useChatStore = create<ChatStore>((set, get) => ({
       case 'message_complete': {
         const session = get().sessions[sessionId]
         if (!session) break
-        const text = session.streamingText
-        if (text) {
+        const text = `${session.streamingText}${consumePendingDelta()}`
+        if (text.trim()) {
           update((s) => ({
-            messages: [...s.messages, { id: nextId(), type: 'assistant_text', content: text, timestamp: Date.now() }],
+            messages: appendAssistantTextMessage(s.messages, text, Date.now()),
             streamingText: '',
           }))
+        } else if (text !== session.streamingText) {
+          update(() => ({ streamingText: text }))
         }
         if (session.elapsedTimer) clearInterval(session.elapsedTimer)
         update(() => ({
@@ -664,12 +699,12 @@ export const useChatStore = create<ChatStore>((set, get) => ({
 
       case 'error':
         update((s) => {
-          const pendingText = s.streamingText.trim()
-          const newMessages = [...s.messages]
-          if (pendingText) {
-            newMessages.push({ id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() })
+          const pendingText = `${s.streamingText}${consumePendingDelta()}`
+          let newMessages = s.messages
+          if (pendingText.trim()) {
+            newMessages = appendAssistantTextMessage(newMessages, pendingText, Date.now())
           }
-          newMessages.push({ id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() })
+          newMessages = [...newMessages, { id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }]
           return {
             messages: newMessages,
             chatState: 'idle',
@@ -791,6 +826,30 @@ function extractVisibleTeammateMessageContents(text: string): string[] {
   return contents
 }
 
+function pushAssistantHistoryText(
+  messages: UIMessage[],
+  content: string,
+  timestamp: number,
+  model?: string,
+): void {
+  if (!content.trim()) return
+
+  const last = messages[messages.length - 1]
+  if (last?.type === 'assistant_text') {
+    last.content += content
+    if (model && !last.model) last.model = model
+    return
+  }
+
+  messages.push({
+    id: nextId(),
+    type: 'assistant_text',
+    content,
+    timestamp,
+    ...(model ? { model } : {}),
+  })
+}
+
 type HistoryMappingOptions = {
   includeTeammateMessages?: boolean
 }
@@ -900,7 +959,7 @@ export function mapHistoryMessagesToUiMessages(
     if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
       for (const block of msg.content as AssistantHistoryBlock[]) {
         if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp })
-        else if (block.type === 'text' && block.text) uiMessages.push({ id: nextId(), type: 'assistant_text', content: block.text, timestamp, model: msg.model })
+        else if (block.type === 'text' && block.text) pushAssistantHistoryText(uiMessages, block.text, timestamp, msg.model)
         else if (block.type === 'tool_use') uiMessages.push({ id: nextId(), type: 'tool_use', toolName: block.name ?? 'unknown', toolUseId: block.id ?? '', input: block.input, timestamp, parentToolUseId: msg.parentToolUseId })
       }
       continue