fix(desktop): preserve interim assistant text

apps/desktop/src/app/session/hooks/use-message-stream.ts

@@ -8,6 +8,7 @@ import {
   type ChatMessage,
   type ChatMessagePart,
   chatMessageText,
+  finalizeAssistantTextParts,
   type GatewayEventPayload,
   reasoningPart,
   renderMediaTags,
@@ -458,26 +459,6 @@ export function useMessageStream({
         const streamId = state.streamId
         const finalText = renderMediaTags(text).trim()
         const completionError = completionErrorText(finalText)
-        const normalize = (value: string) => value.replace(/\s+/g, ' ').trim()
-        const dedupeReference = normalize(finalText)
-
-        const replaceTextPart = (parts: ChatMessagePart[]) => {
-          const kept = parts.filter(part => {
-            if (part.type === 'text') {
-              return false
-            }
-
-            if (part.type !== 'reasoning' || !dedupeReference) {
-              return true
-            }
-
-            const r = normalize(part.text)
-
-            return !(r && (dedupeReference.startsWith(r) || r.startsWith(dedupeReference)))
-          })
-
-          return finalText ? [...kept, assistantTextPart(finalText)] : kept
-        }
 
         const completeMessage = (message: ChatMessage): ChatMessage =>
           completionError
@@ -489,7 +470,7 @@ export function useMessageStream({
               }
             : {
                 ...message,
-                parts: replaceTextPart(message.parts),
+                parts: finalizeAssistantTextParts(message.parts, finalText),
                 pending: false
               }
 

apps/desktop/src/lib/chat-messages.test.ts

@@ -3,13 +3,27 @@ import { describe, expect, it } from 'vitest'
 import type { ChatMessage, ChatMessagePart } from './chat-messages'
 import {
   appendAssistantTextPart,
+  assistantTextPart,
   chatMessageText,
+  finalizeAssistantTextParts,
   preserveLocalAssistantErrors,
   renderMediaTags,
   toChatMessages,
   upsertToolPart
 } from './chat-messages'
 
+function completedToolPart(id: string): ChatMessagePart {
+  return {
+    args: {} as never,
+    argsText: '{}',
+    isError: false,
+    result: {},
+    toolCallId: id,
+    toolName: 'terminal',
+    type: 'tool-call'
+  }
+}
+
 describe('toChatMessages', () => {
   it('keeps a turn with interleaved tool-only rows in a single bubble', () => {
     const messages = toChatMessages([
@@ -143,6 +157,34 @@ describe('renderMediaTags', () => {
   })
 })
 
+describe('finalizeAssistantTextParts', () => {
+  it('preserves assistant text before tool calls when completion text is only the final segment', () => {
+    const parts = [
+      assistantTextPart('First, I will inspect the tree.'),
+      completedToolPart('tool-1'),
+      assistantTextPart('Now I will inspect the tests.'),
+      completedToolPart('tool-2'),
+      assistantTextPart('Partial final')
+    ]
+
+    const finalized = finalizeAssistantTextParts(parts, 'Final answer.')
+
+    expect(finalized.map(part => part.type)).toEqual(['text', 'tool-call', 'text', 'tool-call', 'text'])
+    expect(chatMessageText({ id: 'a', parts: finalized, role: 'assistant' })).toBe(
+      'First, I will inspect the tree.Now I will inspect the tests.Final answer.'
+    )
+  })
+
+  it('does not duplicate prior text when completion text contains the full streamed assistant text', () => {
+    const parts = [assistantTextPart('Planning.'), completedToolPart('tool-1'), assistantTextPart('Done.')]
+
+    const finalized = finalizeAssistantTextParts(parts, 'Planning.Done.')
+
+    expect(finalized.map(part => part.type)).toEqual(['text', 'tool-call', 'text'])
+    expect(chatMessageText({ id: 'a', parts: finalized, role: 'assistant' })).toBe('Planning.Done.')
+  })
+})
+
 describe('preserveLocalAssistantErrors', () => {
   it('preserves a local user+error pair when hydration omits the failed turn', () => {
     const nextMessages: ChatMessage[] = [

apps/desktop/src/lib/chat-messages.ts

@@ -110,6 +110,65 @@ export function chatMessageText(message: ChatMessage): string {
     .join('')
 }
 
+function normalizeTextForDedupe(value: string): string {
+  return value.replace(/\s+/g, ' ').trim()
+}
+
+export function finalizeAssistantTextParts(parts: ChatMessagePart[], finalText: string): ChatMessagePart[] {
+  const dedupeReference = normalizeTextForDedupe(finalText)
+
+  const withoutDuplicateReasoning = parts.filter(part => {
+    if (part.type !== 'reasoning' || !dedupeReference) {
+      return true
+    }
+
+    const reasoning = normalizeTextForDedupe(part.text)
+
+    return !(reasoning && (dedupeReference.startsWith(reasoning) || reasoning.startsWith(dedupeReference)))
+  })
+
+  if (!finalText) {
+    return withoutDuplicateReasoning
+  }
+
+  const textIndexes = withoutDuplicateReasoning
+    .map((part, index) => (part.type === 'text' ? index : -1))
+    .filter(index => index >= 0)
+
+  if (!textIndexes.length) {
+    return [...withoutDuplicateReasoning, assistantTextPart(finalText)]
+  }
+
+  const lastTextIndex = textIndexes.at(-1)!
+
+  const textBeforeLast = textIndexes
+    .slice(0, -1)
+    .map(index => (withoutDuplicateReasoning[index] as Extract<ChatMessagePart, { type: 'text' }>).text)
+    .join('')
+
+  const lastText = (withoutDuplicateReasoning[lastTextIndex] as Extract<ChatMessagePart, { type: 'text' }>).text
+  const existingText = `${textBeforeLast}${lastText}`
+
+  if (normalizeTextForDedupe(existingText) === dedupeReference) {
+    return withoutDuplicateReasoning
+  }
+
+  const next = [...withoutDuplicateReasoning]
+
+  const currentSegmentText =
+    textBeforeLast && finalText.startsWith(textBeforeLast) ? finalText.slice(textBeforeLast.length) : finalText
+
+  if (!currentSegmentText) {
+    next.splice(lastTextIndex, 1)
+
+    return next
+  }
+
+  next[lastTextIndex] = assistantTextPart(currentSegmentText)
+
+  return next
+}
+
 const ATTACHED_CONTEXT_MARKER_RE = /(?:^|\n)--- Attached Context ---\s*\n/
 const CONTEXT_WARNINGS_MARKER_RE = /(?:^|\n)--- Context Warnings ---[\s\S]*$/
 const CONTEXT_REF_RE = /@(file|folder|url|image|tool|terminal):(?:"[^"\n]+"|'[^'\n]+'|`[^`\n]+`|\S+)/g