@tanstack/ai 0.6.3 → 0.8.1

package/src/activities/chat/stream/processor.ts

@@ -12,6 +12,7 @@
  * - Thinking/reasoning content
  * - Recording/replay for testing
  * - Event-driven architecture for UI updates
+ * - Per-message stream state tracking for multi-message sessions
  *
  * @see docs/chat-architecture.md — Canonical reference for AG-UI chunk ordering,
  *   adapter contract, single-shot flows, and expected UIMessage output.
@@ -32,6 +33,7 @@ import type {
   ChunkRecording,
   ChunkStrategy,
   InternalToolCallState,
+  MessageStreamState,
   ProcessorResult,
   ProcessorState,
   ToolCallState,
@@ -116,9 +118,8 @@ export interface StreamProcessorOptions {
  *
  * State tracking:
  * - Full message array
- * - Current assistant message being streamed
- * - Text content accumulation (reset on TEXT_MESSAGE_START)
- * - Multiple parallel tool calls
+ * - Per-message stream state (text, tool calls, thinking)
+ * - Multiple concurrent message streams
  * - Tool call completion via TOOL_CALL_END events
  *
  * @see docs/chat-architecture.md#streamprocessor-internal-state — State field reference
@@ -132,17 +133,17 @@ export class StreamProcessor {
 
   // Message state
   private messages: Array<UIMessage> = []
-  private currentAssistantMessageId: string | null = null
-
-  // Stream state for current assistant message
-  // Total accumulated text across all segments (for the final result)
-  private totalTextContent = ''
-  // Current segment's text content (for onTextUpdate callbacks)
-  private currentSegmentText = ''
-  private lastEmittedText = ''
-  private thinkingContent = ''
-  private toolCalls: Map<string, InternalToolCallState> = new Map()
-  private toolCallOrder: Array<string> = []
+
+  // Per-message stream state
+  private messageStates: Map<string, MessageStreamState> = new Map()
+  private activeMessageIds: Set<string> = new Set()
+  private toolCallToMessage: Map<string, string> = new Map()
+  private pendingManualMessageId: string | null = null
+
+  // Run tracking (for concurrent run safety)
+  private activeRuns = new Set<string>()
+
+  // Shared stream state
   private finishReason: string | null = null
   private hasError = false
   private isDone = false
@@ -231,18 +232,17 @@ export class StreamProcessor {
   prepareAssistantMessage(): void {
     // Reset stream state for new message
     this.resetStreamState()
-    // Clear the current assistant message ID so ensureAssistantMessage()
-    // will create a fresh message on the first content chunk
-    this.currentAssistantMessageId = null
   }
 
   /**
    * @deprecated Use prepareAssistantMessage() instead. This eagerly creates
    * an assistant message which can cause empty message flicker.
    */
-  startAssistantMessage(): string {
+  startAssistantMessage(messageId?: string): string {
     this.prepareAssistantMessage()
-    return this.ensureAssistantMessage()
+    const { messageId: id } = this.ensureAssistantMessage(messageId)
+    this.pendingManualMessageId = id
+    return id
   }
 
   /**
@@ -251,39 +251,16 @@ export class StreamProcessor {
    * has arrived yet.
    */
   getCurrentAssistantMessageId(): string | null {
-    return this.currentAssistantMessageId
-  }
-
-  /**
-   * Lazily create the assistant message if it hasn't been created yet.
-   * Called by content handlers on the first content-bearing chunk.
-   * Returns the message ID.
-   *
-   * Content-bearing chunks that trigger this:
-   * TEXT_MESSAGE_CONTENT, TOOL_CALL_START, STEP_FINISHED, RUN_ERROR.
-   *
-   * @see docs/chat-architecture.md#streamprocessor-internal-state — Lazy creation pattern
-   */
-  private ensureAssistantMessage(): string {
-    if (this.currentAssistantMessageId) {
-      return this.currentAssistantMessageId
-    }
-
-    const assistantMessage: UIMessage = {
-      id: generateMessageId(),
-      role: 'assistant',
-      parts: [],
-      createdAt: new Date(),
+    // Scan all message states (not just active) for the last assistant.
+    // After finalizeStream() clears activeMessageIds, messageStates retains entries.
+    // After reset() / resetStreamState(), messageStates is cleared → returns null.
+    let lastId: string | null = null
+    for (const [id, state] of this.messageStates) {
+      if (state.role === 'assistant') {
+        lastId = id
+      }
     }
-
-    this.currentAssistantMessageId = assistantMessage.id
-    this.messages = [...this.messages, assistantMessage]
-
-    // Emit events
-    this.events.onStreamStart?.()
-    this.emitMessagesChange()
-
-    return assistantMessage.id
+    return lastId
   }
 
   /**
@@ -410,7 +387,10 @@ export class StreamProcessor {
    */
   clearMessages(): void {
     this.messages = []
-    this.currentAssistantMessageId = null
+    this.messageStates.clear()
+    this.activeMessageIds.clear()
+    this.toolCallToMessage.clear()
+    this.pendingManualMessageId = null
     this.emitMessagesChange()
   }
 
@@ -451,7 +431,7 @@ export class StreamProcessor {
    *
    * Central dispatch for all AG-UI events. Each event type maps to a specific
    * handler. Events not listed in the switch are intentionally ignored
-   * (RUN_STARTED, TEXT_MESSAGE_END, STEP_STARTED, STATE_SNAPSHOT, STATE_DELTA).
+   * (RUN_STARTED, STEP_STARTED, STATE_DELTA).
    *
    * @see docs/chat-architecture.md#adapter-contract — Expected event types and ordering
    */
@@ -468,13 +448,17 @@ export class StreamProcessor {
     switch (chunk.type) {
       // AG-UI Events
       case 'TEXT_MESSAGE_START':
-        this.handleTextMessageStartEvent()
+        this.handleTextMessageStartEvent(chunk)
         break
 
       case 'TEXT_MESSAGE_CONTENT':
         this.handleTextMessageContentEvent(chunk)
         break
 
+      case 'TEXT_MESSAGE_END':
+        this.handleTextMessageEndEvent(chunk)
+        break
+
       case 'TOOL_CALL_START':
         this.handleToolCallStartEvent(chunk)
         break
@@ -499,35 +483,264 @@ export class StreamProcessor {
         this.handleStepFinishedEvent(chunk)
         break
 
+      case 'MESSAGES_SNAPSHOT':
+        this.handleMessagesSnapshotEvent(chunk)
+        break
+
       case 'CUSTOM':
         this.handleCustomEvent(chunk)
         break
 
+      case 'RUN_STARTED':
+        this.handleRunStartedEvent(chunk)
+        break
+
       default:
-        // RUN_STARTED, TEXT_MESSAGE_END, STEP_STARTED,
-        // STATE_SNAPSHOT, STATE_DELTA - no special handling needed
+        // STEP_STARTED, STATE_SNAPSHOT, STATE_DELTA - no special handling needed
         break
     }
   }
 
+  // ============================================
+  // Per-Message State Helpers
+  // ============================================
+
+  /**
+   * Create a new MessageStreamState for a message
+   */
+  private createMessageState(
+    messageId: string,
+    role: 'user' | 'assistant' | 'system',
+  ): MessageStreamState {
+    const state: MessageStreamState = {
+      id: messageId,
+      role,
+      totalTextContent: '',
+      currentSegmentText: '',
+      lastEmittedText: '',
+      thinkingContent: '',
+      toolCalls: new Map(),
+      toolCallOrder: [],
+      hasToolCallsSinceTextStart: false,
+      isComplete: false,
+    }
+    this.messageStates.set(messageId, state)
+    return state
+  }
+
   /**
-   * Handle TEXT_MESSAGE_START event — marks the beginning of a new text segment.
-   * Resets segment accumulation so text after tool calls starts fresh.
-   *
-   * This is the key mechanism for multi-segment text (text before and after tool
-   * calls becoming separate TextParts). Without this reset, all text would merge
-   * into a single TextPart and tool-call interleaving would be lost.
+   * Get the MessageStreamState for a message
+   */
+  private getMessageState(messageId: string): MessageStreamState | undefined {
+    return this.messageStates.get(messageId)
+  }
+
+  /**
+   * Get the most recent active assistant message ID.
+   * Used as fallback for events that don't include a messageId.
+   */
+  private getActiveAssistantMessageId(): string | null {
+    // Set iteration is insertion-order; convert to array and search from the end
+    const ids = Array.from(this.activeMessageIds)
+    for (let i = ids.length - 1; i >= 0; i--) {
+      const id = ids[i]!
+      const state = this.messageStates.get(id)
+      if (state && state.role === 'assistant') {
+        return id
+      }
+    }
+    return null
+  }
+
+  /**
+   * Ensure an active assistant message exists, creating one if needed.
+   * Used for backward compat when events arrive without prior TEXT_MESSAGE_START.
    *
-   * @see docs/chat-architecture.md#single-shot-text-response — Step-by-step text processing
-   * @see docs/chat-architecture.md#text-then-tool-interleaving-single-shot — Multi-segment text
+   * On reconnect/resume, a TEXT_MESSAGE_CONTENT may arrive for a message that
+   * already exists in this.messages (e.g. from initialMessages or a prior
+   * MESSAGES_SNAPSHOT) but whose transient state was cleared. In that case we
+   * hydrate state from the existing message rather than creating a duplicate.
    */
-  private handleTextMessageStartEvent(): void {
-    // Emit any pending text from a previous segment before resetting
-    if (this.currentSegmentText !== this.lastEmittedText) {
-      this.emitTextUpdate()
+  private ensureAssistantMessage(preferredId?: string): {
+    messageId: string
+    state: MessageStreamState
+  } {
+    // Try to find state by preferred ID
+    if (preferredId) {
+      const state = this.getMessageState(preferredId)
+      if (state) return { messageId: preferredId, state }
+    }
+
+    // Try active assistant message
+    const activeId = this.getActiveAssistantMessageId()
+    if (activeId) {
+      const state = this.getMessageState(activeId)!
+      return { messageId: activeId, state }
+    }
+
+    // Check if a message with preferredId already exists (reconnect/resume case).
+    // Hydrate transient state from the existing message instead of duplicating it.
+    if (preferredId) {
+      const existingMsg = this.messages.find((m) => m.id === preferredId)
+      if (existingMsg) {
+        const state = this.createMessageState(preferredId, existingMsg.role)
+        this.activeMessageIds.add(preferredId)
+
+        // Seed segment text from the existing last text part so that
+        // incoming deltas append correctly and updateTextPart produces
+        // the right content (existing text + new delta).
+        const lastPart =
+          existingMsg.parts.length > 0
+            ? existingMsg.parts[existingMsg.parts.length - 1]
+            : null
+        if (lastPart && lastPart.type === 'text') {
+          state.currentSegmentText = lastPart.content
+          state.lastEmittedText = lastPart.content
+          state.totalTextContent = lastPart.content
+        }
+
+        return { messageId: preferredId, state }
+      }
     }
-    this.currentSegmentText = ''
-    this.lastEmittedText = ''
+
+    // Auto-create an assistant message (backward compat for process() without TEXT_MESSAGE_START)
+    const id = preferredId || generateMessageId()
+    const assistantMessage: UIMessage = {
+      id,
+      role: 'assistant',
+      parts: [],
+      createdAt: new Date(),
+    }
+    this.messages = [...this.messages, assistantMessage]
+    const state = this.createMessageState(id, 'assistant')
+    this.activeMessageIds.add(id)
+    this.pendingManualMessageId = id
+    this.events.onStreamStart?.()
+    this.emitMessagesChange()
+    return { messageId: id, state }
+  }
+
+  // ============================================
+  // Event Handlers
+  // ============================================
+
+  /**
+   * Handle TEXT_MESSAGE_START event
+   */
+  private handleTextMessageStartEvent(
+    chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_START' }>,
+  ): void {
+    const { messageId, role } = chunk
+
+    // Map 'tool' role to 'assistant' for both UIMessage and MessageStreamState
+    // (UIMessage doesn't support 'tool' role, and lookups like
+    // getActiveAssistantMessageId() check state.role === 'assistant')
+    const uiRole: 'system' | 'user' | 'assistant' =
+      role === 'tool' ? 'assistant' : role
+
+    // Case 1: A manual message was created via startAssistantMessage()
+    if (this.pendingManualMessageId) {
+      const pendingId = this.pendingManualMessageId
+      this.pendingManualMessageId = null
+
+      if (pendingId !== messageId) {
+        // Update the message's ID in the messages array
+        this.messages = this.messages.map((msg) =>
+          msg.id === pendingId ? { ...msg, id: messageId } : msg,
+        )
+
+        // Move state to the new key
+        const existingState = this.messageStates.get(pendingId)
+        if (existingState) {
+          existingState.id = messageId
+          this.messageStates.delete(pendingId)
+          this.messageStates.set(messageId, existingState)
+        }
+
+        // Update activeMessageIds
+        this.activeMessageIds.delete(pendingId)
+        this.activeMessageIds.add(messageId)
+      }
+
+      // Ensure state exists
+      if (!this.messageStates.has(messageId)) {
+        this.createMessageState(messageId, uiRole)
+        this.activeMessageIds.add(messageId)
+      }
+
+      this.emitMessagesChange()
+      return
+    }
+
+    // Case 2: Message already exists (dedup)
+    const existingMsg = this.messages.find((m) => m.id === messageId)
+    if (existingMsg) {
+      this.activeMessageIds.add(messageId)
+      if (!this.messageStates.has(messageId)) {
+        this.createMessageState(messageId, uiRole)
+      } else {
+        const existingState = this.messageStates.get(messageId)!
+        // If tool calls happened since last text, this TEXT_MESSAGE_START
+        // signals a new text segment — reset segment accumulation
+        if (existingState.hasToolCallsSinceTextStart) {
+          if (
+            existingState.currentSegmentText !== existingState.lastEmittedText
+          ) {
+            this.emitTextUpdateForMessage(messageId)
+          }
+          existingState.currentSegmentText = ''
+          existingState.lastEmittedText = ''
+          existingState.hasToolCallsSinceTextStart = false
+        }
+      }
+      return
+    }
+
+    // Case 3: New message from the stream
+    const newMessage: UIMessage = {
+      id: messageId,
+      role: uiRole,
+      parts: [],
+      createdAt: new Date(),
+    }
+
+    this.messages = [...this.messages, newMessage]
+    this.createMessageState(messageId, uiRole)
+    this.activeMessageIds.add(messageId)
+
+    this.events.onStreamStart?.()
+    this.emitMessagesChange()
+  }
+
+  /**
+   * Handle TEXT_MESSAGE_END event
+   */
+  private handleTextMessageEndEvent(
+    chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_END' }>,
+  ): void {
+    const { messageId } = chunk
+    const state = this.getMessageState(messageId)
+    if (!state) return
+    if (state.isComplete) return
+
+    // Emit any pending text for this message
+    if (state.currentSegmentText !== state.lastEmittedText) {
+      this.emitTextUpdateForMessage(messageId)
+    }
+
+    // Complete all tool calls for this message
+    this.completeAllToolCallsForMessage(messageId)
+  }
+
+  /**
+   * Handle MESSAGES_SNAPSHOT event
+   */
+  private handleMessagesSnapshotEvent(
+    chunk: Extract<StreamChunk, { type: 'MESSAGES_SNAPSHOT' }>,
+  ): void {
+    this.resetStreamState()
+    this.messages = [...chunk.messages]
+    this.emitMessagesChange()
   }
 
   /**
@@ -544,17 +757,62 @@ export class StreamProcessor {
   private handleTextMessageContentEvent(
     chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_CONTENT' }>,
   ): void {
-    this.ensureAssistantMessage()
+    const { messageId, state } = this.ensureAssistantMessage(chunk.messageId)
+
+    // Content arriving means all current tool calls for this message are complete
+    this.completeAllToolCallsForMessage(messageId)
+
+    const previousSegment = state.currentSegmentText
+
+    // Detect if this is a NEW text segment (after tool calls) vs continuation
+    const isNewSegment =
+      state.hasToolCallsSinceTextStart &&
+      previousSegment.length > 0 &&
+      this.isNewTextSegment(chunk, previousSegment)
+
+    if (isNewSegment) {
+      // Emit any accumulated text before starting new segment
+      if (previousSegment !== state.lastEmittedText) {
+        this.emitTextUpdateForMessage(messageId)
+      }
+      // Reset SEGMENT text accumulation for the new text segment after tool calls
+      state.currentSegmentText = ''
+      state.lastEmittedText = ''
+      state.hasToolCallsSinceTextStart = false
+    }
 
-    this.currentSegmentText += chunk.delta
-    this.totalTextContent += chunk.delta
+    const currentText = state.currentSegmentText
+    let nextText = currentText
+
+    // Prefer delta over content - delta is the incremental change
+    // Normalize to empty string to avoid "undefined" string concatenation
+    const delta = chunk.delta || ''
+    if (delta !== '') {
+      nextText = currentText + delta
+    } else if (chunk.content !== undefined && chunk.content !== '') {
+      // Fallback: use content if delta is not provided
+      if (chunk.content.startsWith(currentText)) {
+        nextText = chunk.content
+      } else if (currentText.startsWith(chunk.content)) {
+        nextText = currentText
+      } else {
+        nextText = currentText + chunk.content
+      }
+    }
 
+    // Calculate the delta for totalTextContent
+    const textDelta = nextText.slice(currentText.length)
+    state.currentSegmentText = nextText
+    state.totalTextContent += textDelta
+
+    // Use delta for chunk strategy if available
+    const chunkPortion = chunk.delta || chunk.content || ''
     const shouldEmit = this.chunkStrategy.shouldEmit(
-      chunk.delta,
-      this.currentSegmentText,
+      chunkPortion,
+      state.currentSegmentText,
     )
-    if (shouldEmit && this.currentSegmentText !== this.lastEmittedText) {
-      this.emitTextUpdate()
+    if (shouldEmit && state.currentSegmentText !== state.lastEmittedText) {
+      this.emitTextUpdateForMessage(messageId)
     }
   }
 
@@ -574,10 +832,18 @@ export class StreamProcessor {
   private handleToolCallStartEvent(
     chunk: Extract<StreamChunk, { type: 'TOOL_CALL_START' }>,
   ): void {
-    this.ensureAssistantMessage()
+    // Determine the message this tool call belongs to
+    const targetMessageId =
+      chunk.parentMessageId ?? this.getActiveAssistantMessageId()
+    const { messageId, state } = this.ensureAssistantMessage(
+      targetMessageId ?? undefined,
+    )
+
+    // Mark that we've seen tool calls since the last text segment
+    state.hasToolCallsSinceTextStart = true
 
     const toolCallId = chunk.toolCallId
-    const existingToolCall = this.toolCalls.get(toolCallId)
+    const existingToolCall = state.toolCalls.get(toolCallId)
 
     if (!existingToolCall) {
       // New tool call starting
@@ -589,34 +855,31 @@ export class StreamProcessor {
         arguments: '',
         state: initialState,
         parsedArguments: undefined,
-        index: chunk.index ?? this.toolCalls.size,
+        index: chunk.index ?? state.toolCalls.size,
       }
 
-      this.toolCalls.set(toolCallId, newToolCall)
-      this.toolCallOrder.push(toolCallId)
+      state.toolCalls.set(toolCallId, newToolCall)
+      state.toolCallOrder.push(toolCallId)
+
+      // Store mapping for TOOL_CALL_ARGS/END routing
+      this.toolCallToMessage.set(toolCallId, messageId)
 
       // Update UIMessage
-      if (this.currentAssistantMessageId) {
-        this.messages = updateToolCallPart(
-          this.messages,
-          this.currentAssistantMessageId,
-          {
-            id: chunk.toolCallId,
-            name: chunk.toolName,
-            arguments: '',
-            state: initialState,
-          },
-        )
-        this.emitMessagesChange()
+      this.messages = updateToolCallPart(this.messages, messageId, {
+        id: chunk.toolCallId,
+        name: chunk.toolName,
+        arguments: '',
+        state: initialState,
+      })
+      this.emitMessagesChange()
 
-        // Emit granular event
-        this.events.onToolCallStateChange?.(
-          this.currentAssistantMessageId,
-          chunk.toolCallId,
-          initialState,
-          '',
-        )
-      }
+      // Emit granular event
+      this.events.onToolCallStateChange?.(
+        messageId,
+        chunk.toolCallId,
+        initialState,
+        '',
+      )
     }
   }
 
@@ -636,47 +899,46 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'TOOL_CALL_ARGS' }>,
   ): void {
     const toolCallId = chunk.toolCallId
-    const existingToolCall = this.toolCalls.get(toolCallId)
+    const messageId = this.toolCallToMessage.get(toolCallId)
+    if (!messageId) return
 
-    if (existingToolCall) {
-      const wasAwaitingInput = existingToolCall.state === 'awaiting-input'
+    const state = this.getMessageState(messageId)
+    if (!state) return
 
-      // Accumulate arguments from delta
-      existingToolCall.arguments += chunk.delta || ''
+    const existingToolCall = state.toolCalls.get(toolCallId)
+    if (!existingToolCall) return
 
-      // Update state
-      if (wasAwaitingInput && chunk.delta) {
-        existingToolCall.state = 'input-streaming'
-      }
-
-      // Try to parse the updated arguments
-      existingToolCall.parsedArguments = this.jsonParser.parse(
-        existingToolCall.arguments,
-      )
+    const wasAwaitingInput = existingToolCall.state === 'awaiting-input'
 
-      // Update UIMessage
-      if (this.currentAssistantMessageId) {
-        this.messages = updateToolCallPart(
-          this.messages,
-          this.currentAssistantMessageId,
-          {
-            id: existingToolCall.id,
-            name: existingToolCall.name,
-            arguments: existingToolCall.arguments,
-            state: existingToolCall.state,
-          },
-        )
-        this.emitMessagesChange()
+    // Accumulate arguments from delta
+    existingToolCall.arguments += chunk.delta || ''
 
-        // Emit granular event
-        this.events.onToolCallStateChange?.(
-          this.currentAssistantMessageId,
-          existingToolCall.id,
-          existingToolCall.state,
-          existingToolCall.arguments,
-        )
-      }
+    // Update state
+    if (wasAwaitingInput && chunk.delta) {
+      existingToolCall.state = 'input-streaming'
     }
+
+    // Try to parse the updated arguments
+    existingToolCall.parsedArguments = this.jsonParser.parse(
+      existingToolCall.arguments,
+    )
+
+    // Update UIMessage
+    this.messages = updateToolCallPart(this.messages, messageId, {
+      id: existingToolCall.id,
+      name: existingToolCall.name,
+      arguments: existingToolCall.arguments,
+      state: existingToolCall.state,
+    })
+    this.emitMessagesChange()
+
+    // Emit granular event
+    this.events.onToolCallStateChange?.(
+      messageId,
+      existingToolCall.id,
+      existingToolCall.state,
+      existingToolCall.arguments,
+    )
   }
 
   /**
@@ -696,11 +958,17 @@ export class StreamProcessor {
   private handleToolCallEndEvent(
     chunk: Extract<StreamChunk, { type: 'TOOL_CALL_END' }>,
   ): void {
+    const messageId = this.toolCallToMessage.get(chunk.toolCallId)
+    if (!messageId) return
+
+    const msgState = this.getMessageState(messageId)
+    if (!msgState) return
+
     // Transition the tool call to input-complete (the authoritative completion signal)
-    const existingToolCall = this.toolCalls.get(chunk.toolCallId)
+    const existingToolCall = msgState.toolCalls.get(chunk.toolCallId)
     if (existingToolCall && existingToolCall.state !== 'input-complete') {
-      const index = this.toolCallOrder.indexOf(chunk.toolCallId)
-      this.completeToolCall(index, existingToolCall)
+      const index = msgState.toolCallOrder.indexOf(chunk.toolCallId)
+      this.completeToolCall(messageId, index, existingToolCall)
       // If TOOL_CALL_END provides parsed input, use it as the canonical parsed
       // arguments (overrides the accumulated string parse from completeToolCall)
       if (chunk.input !== undefined) {
@@ -708,10 +976,8 @@ export class StreamProcessor {
       }
     }
 
-    // Update UIMessage if we have a current assistant message and a result
-    if (this.currentAssistantMessageId && chunk.result) {
-      const state: ToolResultState = 'complete'
-
+    // Update UIMessage if there's a result
+    if (chunk.result) {
       // Step 1: Update the tool-call part's output field (for UI consistency
       // with client tools — see GitHub issue #176)
       let output: unknown
@@ -727,23 +993,36 @@ export class StreamProcessor {
       )
 
       // Step 2: Create/update the tool-result part (for LLM conversation history)
+      const resultState: ToolResultState = 'complete'
       this.messages = updateToolResultPart(
         this.messages,
-        this.currentAssistantMessageId,
+        messageId,
         chunk.toolCallId,
         chunk.result,
-        state,
+        resultState,
       )
       this.emitMessagesChange()
     }
   }
 
+  /**
+   * Handle RUN_STARTED event.
+   *
+   * Registers the run so that RUN_FINISHED can determine whether other
+   * runs are still active before finalizing.
+   */
+  private handleRunStartedEvent(
+    chunk: Extract<StreamChunk, { type: 'RUN_STARTED' }>,
+  ): void {
+    this.activeRuns.add(chunk.runId)
+  }
+
   /**
    * Handle RUN_FINISHED event.
    *
-   * Records the finishReason and calls completeAllToolCalls() as a safety net
-   * to force-complete any tool calls that didn't receive an explicit TOOL_CALL_END.
-   * This handles cases like aborted streams or adapter bugs.
+   * Records the finishReason and removes the run from activeRuns.
+   * Only finalizes when no more runs are active, so that concurrent
+   * runs don't interfere with each other.
    *
    * @see docs/chat-architecture.md#single-shot-tool-call-response — finishReason semantics
    * @see docs/chat-architecture.md#adapter-contract — Why RUN_FINISHED is mandatory
@@ -752,8 +1031,13 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'RUN_FINISHED' }>,
   ): void {
     this.finishReason = chunk.finishReason
-    this.isDone = true
-    this.completeAllToolCalls()
+    this.activeRuns.delete(chunk.runId)
+
+    if (this.activeRuns.size === 0) {
+      this.isDone = true
+      this.completeAllToolCalls()
+      this.finalizeStream()
+    }
   }
 
   /**
@@ -763,8 +1047,12 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
   ): void {
     this.hasError = true
+    if (chunk.runId) {
+      this.activeRuns.delete(chunk.runId)
+    } else {
+      this.activeRuns.clear()
+    }
     this.ensureAssistantMessage()
-    // Emit error event
     this.events.onError?.(new Error(chunk.error.message || 'An error occurred'))
   }
 
@@ -779,25 +1067,38 @@ export class StreamProcessor {
   private handleStepFinishedEvent(
     chunk: Extract<StreamChunk, { type: 'STEP_FINISHED' }>,
   ): void {
-    this.ensureAssistantMessage()
+    const { messageId, state } = this.ensureAssistantMessage(
+      this.getActiveAssistantMessageId() ?? undefined,
+    )
+
+    const previous = state.thinkingContent
+    let nextThinking = previous
+
+    // Prefer delta over content
+    if (chunk.delta && chunk.delta !== '') {
+      nextThinking = previous + chunk.delta
+    } else if (chunk.content && chunk.content !== '') {
+      if (chunk.content.startsWith(previous)) {
+        nextThinking = chunk.content
+      } else if (previous.startsWith(chunk.content)) {
+        nextThinking = previous
+      } else {
+        nextThinking = previous + chunk.content
+      }
+    }
 
-    this.thinkingContent += chunk.delta
+    state.thinkingContent = nextThinking
 
     // Update UIMessage
-    if (this.currentAssistantMessageId) {
-      this.messages = updateThinkingPart(
-        this.messages,
-        this.currentAssistantMessageId,
-        this.thinkingContent,
-      )
-      this.emitMessagesChange()
+    this.messages = updateThinkingPart(
+      this.messages,
+      messageId,
+      state.thinkingContent,
+    )
+    this.emitMessagesChange()
 
-      // Emit granular event
-      this.events.onThinkingUpdate?.(
-        this.currentAssistantMessageId,
-        this.thinkingContent,
-      )
-    }
+    // Emit granular event
+    this.events.onThinkingUpdate?.(messageId, state.thinkingContent)
   }
 
   /**
@@ -813,6 +1114,8 @@ export class StreamProcessor {
   private handleCustomEvent(
     chunk: Extract<StreamChunk, { type: 'CUSTOM' }>,
   ): void {
+    const messageId = this.getActiveAssistantMessageId()
+
     // Handle client tool input availability - trigger client-side execution
     if (chunk.name === 'tool-input-available' && chunk.value) {
       const { toolCallId, toolName, input } = chunk.value as {
@@ -839,11 +1142,15 @@ export class StreamProcessor {
         approval: { id: string; needsApproval: boolean }
       }
 
-      // Update the tool call part with approval state
-      if (this.currentAssistantMessageId) {
+      // Resolve the message containing this tool call. After RUN_FINISHED,
+      // activeMessageIds is cleared, so fall back to the toolCallToMessage map
+      // which is populated during TOOL_CALL_START and preserved across finalize.
+      const resolvedMessageId =
+        messageId ?? this.toolCallToMessage.get(toolCallId)
+      if (resolvedMessageId) {
         this.messages = updateToolCallApproval(
           this.messages,
-          this.currentAssistantMessageId,
+          resolvedMessageId,
           toolCallId,
           approval.id,
         )
@@ -870,8 +1177,34 @@ export class StreamProcessor {
     }
   }
 
+  // ============================================
+  // Internal Helpers
+  // ============================================
+
   /**
-   * Complete all tool calls — safety net for stream termination.
+   * Detect if an incoming content chunk represents a NEW text segment
+   */
+  private isNewTextSegment(
+    chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_CONTENT' }>,
+    previous: string,
+  ): boolean {
+    // Check if content is present (delta is always defined but may be empty string)
+    if (chunk.content !== undefined) {
+      if (chunk.content.length < previous.length) {
+        return true
+      }
+      if (
+        !chunk.content.startsWith(previous) &&
+        !previous.startsWith(chunk.content)
+      ) {
+        return true
+      }
+    }
+    return false
+  }
+
+  /**
+   * Complete all tool calls across all active messages — safety net for stream termination.
    *
    * Called by RUN_FINISHED and finalizeStream(). Force-transitions any tool call
    * not yet in input-complete state. Handles cases where TOOL_CALL_END was
@@ -880,10 +1213,22 @@ export class StreamProcessor {
    * @see docs/chat-architecture.md#single-shot-tool-call-response — Safety net behavior
    */
   private completeAllToolCalls(): void {
-    this.toolCalls.forEach((toolCall, id) => {
+    for (const messageId of this.activeMessageIds) {
+      this.completeAllToolCallsForMessage(messageId)
+    }
+  }
+
+  /**
+   * Complete all tool calls for a specific message
+   */
+  private completeAllToolCallsForMessage(messageId: string): void {
+    const state = this.getMessageState(messageId)
+    if (!state) return
+
+    state.toolCalls.forEach((toolCall, id) => {
       if (toolCall.state !== 'input-complete') {
-        const index = this.toolCallOrder.indexOf(id)
-        this.completeToolCall(index, toolCall)
+        const index = state.toolCallOrder.indexOf(id)
+        this.completeToolCall(messageId, index, toolCall)
       }
     })
   }
@@ -892,6 +1237,7 @@ export class StreamProcessor {
    * Mark a tool call as complete and emit event
    */
   private completeToolCall(
+    messageId: string,
     _index: number,
     toolCall: InternalToolCallState,
   ): void {
@@ -901,31 +1247,25 @@ export class StreamProcessor {
     toolCall.parsedArguments = this.jsonParser.parse(toolCall.arguments)
 
     // Update UIMessage
-    if (this.currentAssistantMessageId) {
-      this.messages = updateToolCallPart(
-        this.messages,
-        this.currentAssistantMessageId,
-        {
-          id: toolCall.id,
-          name: toolCall.name,
-          arguments: toolCall.arguments,
-          state: 'input-complete',
-        },
-      )
-      this.emitMessagesChange()
+    this.messages = updateToolCallPart(this.messages, messageId, {
+      id: toolCall.id,
+      name: toolCall.name,
+      arguments: toolCall.arguments,
+      state: 'input-complete',
+    })
+    this.emitMessagesChange()
 
-      // Emit granular event
-      this.events.onToolCallStateChange?.(
-        this.currentAssistantMessageId,
-        toolCall.id,
-        'input-complete',
-        toolCall.arguments,
-      )
-    }
+    // Emit granular event
+    this.events.onToolCallStateChange?.(
+      messageId,
+      toolCall.id,
+      'input-complete',
+      toolCall.arguments,
+    )
   }
 
   /**
-   * Emit pending text update.
+   * Emit pending text update for a specific message.
    *
    * Calls updateTextPart() which has critical append-vs-replace logic:
    * - If last UIMessage part is TextPart → replaces its content (same segment).
@@ -933,24 +1273,22 @@ export class StreamProcessor {
    *
    * @see docs/chat-architecture.md#uimessage-part-ordering-invariants — Replace vs. push logic
    */
-  private emitTextUpdate(): void {
-    this.lastEmittedText = this.currentSegmentText
+  private emitTextUpdateForMessage(messageId: string): void {
+    const state = this.getMessageState(messageId)
+    if (!state) return
+
+    state.lastEmittedText = state.currentSegmentText
 
     // Update UIMessage
-    if (this.currentAssistantMessageId) {
-      this.messages = updateTextPart(
-        this.messages,
-        this.currentAssistantMessageId,
-        this.currentSegmentText,
-      )
-      this.emitMessagesChange()
+    this.messages = updateTextPart(
+      this.messages,
+      messageId,
+      state.currentSegmentText,
+    )
+    this.emitMessagesChange()
 
-      // Emit granular event
-      this.events.onTextUpdate?.(
-        this.currentAssistantMessageId,
-        this.currentSegmentText,
-      )
-    }
+    // Emit granular event
+    this.events.onTextUpdate?.(messageId, state.currentSegmentText)
   }
 
   /**
@@ -970,81 +1308,116 @@ export class StreamProcessor {
    * @see docs/chat-architecture.md#single-shot-text-response — Finalization step
    */
   finalizeStream(): void {
-    // Safety net: complete any remaining tool calls (e.g. on network errors / aborted streams)
-    this.completeAllToolCalls()
+    let lastAssistantMessage: UIMessage | undefined
 
-    // Emit any pending text if not already emitted
-    if (this.currentSegmentText !== this.lastEmittedText) {
-      this.emitTextUpdate()
+    // Finalize ALL active messages
+    for (const messageId of this.activeMessageIds) {
+      const state = this.getMessageState(messageId)
+      if (!state) continue
+
+      // Complete any remaining tool calls
+      this.completeAllToolCallsForMessage(messageId)
+
+      // Emit any pending text if not already emitted
+      if (state.currentSegmentText !== state.lastEmittedText) {
+        this.emitTextUpdateForMessage(messageId)
+      }
+
+      state.isComplete = true
+
+      const msg = this.messages.find((m) => m.id === messageId)
+      if (msg && msg.role === 'assistant') {
+        lastAssistantMessage = msg
+      }
     }
 
-    // Remove the assistant message if it only contains whitespace text
-    // (no tool calls, no meaningful content). This handles models like Gemini
-    // that sometimes return just "\n" during auto-continuation.
+    this.activeMessageIds.clear()
+
+    // Remove whitespace-only assistant messages (handles models like Gemini
+    // that sometimes return just "\n" during auto-continuation).
     // Preserve the message on errors so the UI can show error state.
-    if (this.currentAssistantMessageId && !this.hasError) {
-      const assistantMessage = this.messages.find(
-        (m) => m.id === this.currentAssistantMessageId,
-      )
-      if (assistantMessage && this.isWhitespaceOnlyMessage(assistantMessage)) {
+    if (lastAssistantMessage && !this.hasError) {
+      if (this.isWhitespaceOnlyMessage(lastAssistantMessage)) {
         this.messages = this.messages.filter(
-          (m) => m.id !== this.currentAssistantMessageId,
+          (m) => m.id !== lastAssistantMessage.id,
         )
         this.emitMessagesChange()
-        this.currentAssistantMessageId = null
         return
       }
     }
 
-    // Emit stream end event (only if a message was actually created)
-    if (this.currentAssistantMessageId) {
-      const assistantMessage = this.messages.find(
-        (m) => m.id === this.currentAssistantMessageId,
-      )
-      if (assistantMessage) {
-        this.events.onStreamEnd?.(assistantMessage)
-      }
+    // Emit stream end for the last assistant message
+    if (lastAssistantMessage) {
+      this.events.onStreamEnd?.(lastAssistantMessage)
     }
   }
 
   /**
-   * Get completed tool calls in API format
+   * Get completed tool calls in API format (aggregated across all messages)
    */
   private getCompletedToolCalls(): Array<ToolCall> {
-    return Array.from(this.toolCalls.values())
-      .filter((tc) => tc.state === 'input-complete')
-      .map((tc) => ({
-        id: tc.id,
-        type: 'function' as const,
-        function: {
-          name: tc.name,
-          arguments: tc.arguments,
-        },
-      }))
+    const result: Array<ToolCall> = []
+    for (const state of this.messageStates.values()) {
+      for (const tc of state.toolCalls.values()) {
+        if (tc.state === 'input-complete') {
+          result.push({
+            id: tc.id,
+            type: 'function' as const,
+            function: {
+              name: tc.name,
+              arguments: tc.arguments,
+            },
+          })
+        }
+      }
+    }
+    return result
   }
 
   /**
-   * Get current result
+   * Get current result (aggregated across all messages)
    */
   private getResult(): ProcessorResult {
     const toolCalls = this.getCompletedToolCalls()
+    let content = ''
+    let thinking = ''
+
+    for (const state of this.messageStates.values()) {
+      content += state.totalTextContent
+      thinking += state.thinkingContent
+    }
+
     return {
-      content: this.totalTextContent,
-      thinking: this.thinkingContent || undefined,
+      content,
+      thinking: thinking || undefined,
       toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
       finishReason: this.finishReason,
     }
   }
 
   /**
-   * Get current processor state
+   * Get current processor state (aggregated across all messages)
    */
   getState(): ProcessorState {
+    let content = ''
+    let thinking = ''
+    const toolCalls = new Map<string, InternalToolCallState>()
+    const toolCallOrder: Array<string> = []
+
+    for (const state of this.messageStates.values()) {
+      content += state.totalTextContent
+      thinking += state.thinkingContent
+      for (const [id, tc] of state.toolCalls) {
+        toolCalls.set(id, tc)
+      }
+      toolCallOrder.push(...state.toolCallOrder)
+    }
+
     return {
-      content: this.totalTextContent,
-      thinking: this.thinkingContent,
-      toolCalls: new Map(this.toolCalls),
-      toolCallOrder: [...this.toolCallOrder],
+      content,
+      thinking,
+      toolCalls,
+      toolCallOrder,
       finishReason: this.finishReason,
       done: this.isDone,
     }
@@ -1074,12 +1447,11 @@ export class StreamProcessor {
    * Reset stream state (but keep messages)
    */
   private resetStreamState(): void {
-    this.totalTextContent = ''
-    this.currentSegmentText = ''
-    this.lastEmittedText = ''
-    this.thinkingContent = ''
-    this.toolCalls.clear()
-    this.toolCallOrder = []
+    this.messageStates.clear()
+    this.activeMessageIds.clear()
+    this.activeRuns.clear()
+    this.toolCallToMessage.clear()
+    this.pendingManualMessageId = null
     this.finishReason = null
     this.hasError = false
     this.isDone = false
@@ -1092,7 +1464,6 @@ export class StreamProcessor {
   reset(): void {
     this.resetStreamState()
     this.messages = []
-    this.currentAssistantMessageId = null
   }
 
   /**