fix(ai): reset segment state when a tool-first message's TEXT_MESSAGE_START arrives

.changeset/tool-first-text-drop.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai': patch
+---
+
+Fix `StreamProcessor` dropping the first `TEXT_MESSAGE_CONTENT` delta when a `TOOL_CALL_START` event's `parentMessageId` precedes that message's `TEXT_MESSAGE_START` — the normal AG-UI shape for "call a tool, then explain the result" as one assistant turn (#1247).

packages/ai/src/activities/chat/stream/processor.ts

@@ -900,9 +900,21 @@ export class StreamProcessor {
       }
 
       // Ensure state exists
-      if (!this.messageStates.has(messageId)) {
-        this.createMessageState(messageId, uiRole)
+      let pendingState = this.messageStates.get(messageId)
+      if (!pendingState) {
+        pendingState = this.createMessageState(messageId, uiRole)
         this.activeMessageIds.add(messageId)
+      } else if (pendingState.hasToolCallsSinceTextStart) {
+        // A tool call (e.g. TOOL_CALL_START with parentMessageId) marked
+        // this message before its "real" TEXT_MESSAGE_START arrived — same
+        // reset Case 2 performs, so the segment accumulator doesn't carry
+        // stale tool-call state into the text that follows.
+        if (pendingState.currentSegmentText !== pendingState.lastEmittedText) {
+          this.emitTextUpdateForMessage(messageId)
+        }
+        pendingState.currentSegmentText = ''
+        pendingState.lastEmittedText = ''
+        pendingState.hasToolCallsSinceTextStart = false
       }
 
       this.mergeMessageMetadata(messageId, chunk.metadata)

packages/ai/tests/stream-processor.test.ts

@@ -4109,6 +4109,38 @@ describe('StreamProcessor', () => {
         },
       ])
     })
+
+    it('should not drop the first TEXT_MESSAGE_CONTENT delta in a tool-first flow (#1247)', () => {
+      const processor = new StreamProcessor()
+
+      processor.processChunk(
+        chunk(EventType.TOOL_CALL_START, {
+          toolCallId: 'tc-1',
+          toolCallName: 'lookupWeather',
+          toolName: 'lookupWeather',
+          parentMessageId: 'anthropic-msg-1',
+        }),
+      )
+      processor.processChunk(ev.toolArgs('tc-1', '{"location":"Berlin"}'))
+      processor.processChunk(ev.toolEnd('tc-1', 'lookupWeather'))
+
+      processor.processChunk(
+        chunk(EventType.TEXT_MESSAGE_START, {
+          messageId: 'anthropic-msg-1',
+          role: 'assistant' as const,
+        }),
+      )
+      // Two deltas: the bug only surfaces once a second delta arrives after
+      // the tool-first message's real TEXT_MESSAGE_START.
+      processor.processChunk(ev.textContent('It is ', 'anthropic-msg-1'))
+      processor.processChunk(ev.textContent('sunny.', 'anthropic-msg-1'))
+      processor.processChunk(ev.textEnd('anthropic-msg-1'))
+      processor.finalizeStream()
+
+      const messages = processor.getMessages()
+      const textParts = messages[0]?.parts.filter((p) => p.type === 'text')
+      expect(textParts).toEqual([{ type: 'text', content: 'It is sunny.' }])
+    })
   })
 
   describe('double onStreamEnd guard', () => {