@tanstack/ai 0.6.3 → 0.8.1

package/dist/esm/activities/chat/stream/processor.js

@@ -5,13 +5,11 @@ import { ImmediateStrategy } from "./strategies.js";
 class StreamProcessor {
   constructor(options = {}) {
     this.messages = [];
-    this.currentAssistantMessageId = null;
-    this.totalTextContent = "";
-    this.currentSegmentText = "";
-    this.lastEmittedText = "";
-    this.thinkingContent = "";
-    this.toolCalls = /* @__PURE__ */ new Map();
-    this.toolCallOrder = [];
+    this.messageStates = /* @__PURE__ */ new Map();
+    this.activeMessageIds = /* @__PURE__ */ new Set();
+    this.toolCallToMessage = /* @__PURE__ */ new Map();
+    this.pendingManualMessageId = null;
+    this.activeRuns = /* @__PURE__ */ new Set();
     this.finishReason = null;
     this.hasError = false;
     this.isDone = false;
@@ -81,15 +79,16 @@ class StreamProcessor {
    */
   prepareAssistantMessage() {
     this.resetStreamState();
-    this.currentAssistantMessageId = null;
   }
   /**
    * @deprecated Use prepareAssistantMessage() instead. This eagerly creates
    * an assistant message which can cause empty message flicker.
    */
-  startAssistantMessage() {
+  startAssistantMessage(messageId) {
     this.prepareAssistantMessage();
-    return this.ensureAssistantMessage();
+    const { messageId: id } = this.ensureAssistantMessage(messageId);
+    this.pendingManualMessageId = id;
+    return id;
   }
   /**
    * Get the current assistant message ID (if one has been created).
@@ -97,33 +96,13 @@ class StreamProcessor {
    * has arrived yet.
    */
   getCurrentAssistantMessageId() {
-    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
-   */
-  ensureAssistantMessage() {
-    if (this.currentAssistantMessageId) {
-      return this.currentAssistantMessageId;
+    let lastId = null;
+    for (const [id, state] of this.messageStates) {
+      if (state.role === "assistant") {
+        lastId = id;
+      }
     }
-    const assistantMessage = {
-      id: generateMessageId(),
-      role: "assistant",
-      parts: [],
-      createdAt: /* @__PURE__ */ new Date()
-    };
-    this.currentAssistantMessageId = assistantMessage.id;
-    this.messages = [...this.messages, assistantMessage];
-    this.events.onStreamStart?.();
-    this.emitMessagesChange();
-    return assistantMessage.id;
+    return lastId;
   }
   /**
    * Add a tool result (called by client after handling onToolCall)
@@ -219,7 +198,10 @@ class StreamProcessor {
    */
   clearMessages() {
     this.messages = [];
-    this.currentAssistantMessageId = null;
+    this.messageStates.clear();
+    this.activeMessageIds.clear();
+    this.toolCallToMessage.clear();
+    this.pendingManualMessageId = null;
     this.emitMessagesChange();
   }
   // ============================================
@@ -247,7 +229,7 @@ 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
    */
@@ -262,11 +244,14 @@ 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;
@@ -285,28 +270,190 @@ class StreamProcessor {
       case "STEP_FINISHED":
         this.handleStepFinishedEvent(chunk);
         break;
+      case "MESSAGES_SNAPSHOT":
+        this.handleMessagesSnapshotEvent(chunk);
+        break;
       case "CUSTOM":
         this.handleCustomEvent(chunk);
         break;
+      case "RUN_STARTED":
+        this.handleRunStartedEvent(chunk);
+        break;
     }
   }
+  // ============================================
+  // Per-Message State Helpers
+  // ============================================
   /**
-   * 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.
+   * Create a new MessageStreamState for a message
+   */
+  createMessageState(messageId, role) {
+    const state = {
+      id: messageId,
+      role,
+      totalTextContent: "",
+      currentSegmentText: "",
+      lastEmittedText: "",
+      thinkingContent: "",
+      toolCalls: /* @__PURE__ */ new Map(),
+      toolCallOrder: [],
+      hasToolCallsSinceTextStart: false,
+      isComplete: false
+    };
+    this.messageStates.set(messageId, state);
+    return state;
+  }
+  /**
+   * Get the MessageStreamState for a message
+   */
+  getMessageState(messageId) {
+    return this.messageStates.get(messageId);
+  }
+  /**
+   * Get the most recent active assistant message ID.
+   * Used as fallback for events that don't include a messageId.
+   */
+  getActiveAssistantMessageId() {
+    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.
    */
-  handleTextMessageStartEvent() {
-    if (this.currentSegmentText !== this.lastEmittedText) {
-      this.emitTextUpdate();
+  ensureAssistantMessage(preferredId) {
+    if (preferredId) {
+      const state2 = this.getMessageState(preferredId);
+      if (state2) return { messageId: preferredId, state: state2 };
     }
-    this.currentSegmentText = "";
-    this.lastEmittedText = "";
+    const activeId = this.getActiveAssistantMessageId();
+    if (activeId) {
+      const state2 = this.getMessageState(activeId);
+      return { messageId: activeId, state: state2 };
+    }
+    if (preferredId) {
+      const existingMsg = this.messages.find((m) => m.id === preferredId);
+      if (existingMsg) {
+        const state2 = this.createMessageState(preferredId, existingMsg.role);
+        this.activeMessageIds.add(preferredId);
+        const lastPart = existingMsg.parts.length > 0 ? existingMsg.parts[existingMsg.parts.length - 1] : null;
+        if (lastPart && lastPart.type === "text") {
+          state2.currentSegmentText = lastPart.content;
+          state2.lastEmittedText = lastPart.content;
+          state2.totalTextContent = lastPart.content;
+        }
+        return { messageId: preferredId, state: state2 };
+      }
+    }
+    const id = preferredId || generateMessageId();
+    const assistantMessage = {
+      id,
+      role: "assistant",
+      parts: [],
+      createdAt: /* @__PURE__ */ 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
+   */
+  handleTextMessageStartEvent(chunk) {
+    const { messageId, role } = chunk;
+    const uiRole = role === "tool" ? "assistant" : role;
+    if (this.pendingManualMessageId) {
+      const pendingId = this.pendingManualMessageId;
+      this.pendingManualMessageId = null;
+      if (pendingId !== messageId) {
+        this.messages = this.messages.map(
+          (msg) => msg.id === pendingId ? { ...msg, id: messageId } : msg
+        );
+        const existingState = this.messageStates.get(pendingId);
+        if (existingState) {
+          existingState.id = messageId;
+          this.messageStates.delete(pendingId);
+          this.messageStates.set(messageId, existingState);
+        }
+        this.activeMessageIds.delete(pendingId);
+        this.activeMessageIds.add(messageId);
+      }
+      if (!this.messageStates.has(messageId)) {
+        this.createMessageState(messageId, uiRole);
+        this.activeMessageIds.add(messageId);
+      }
+      this.emitMessagesChange();
+      return;
+    }
+    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 (existingState.hasToolCallsSinceTextStart) {
+          if (existingState.currentSegmentText !== existingState.lastEmittedText) {
+            this.emitTextUpdateForMessage(messageId);
+          }
+          existingState.currentSegmentText = "";
+          existingState.lastEmittedText = "";
+          existingState.hasToolCallsSinceTextStart = false;
+        }
+      }
+      return;
+    }
+    const newMessage = {
+      id: messageId,
+      role: uiRole,
+      parts: [],
+      createdAt: /* @__PURE__ */ 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
+   */
+  handleTextMessageEndEvent(chunk) {
+    const { messageId } = chunk;
+    const state = this.getMessageState(messageId);
+    if (!state) return;
+    if (state.isComplete) return;
+    if (state.currentSegmentText !== state.lastEmittedText) {
+      this.emitTextUpdateForMessage(messageId);
+    }
+    this.completeAllToolCallsForMessage(messageId);
+  }
+  /**
+   * Handle MESSAGES_SNAPSHOT event
+   */
+  handleMessagesSnapshotEvent(chunk) {
+    this.resetStreamState();
+    this.messages = [...chunk.messages];
+    this.emitMessagesChange();
   }
   /**
    * Handle TEXT_MESSAGE_CONTENT event.
@@ -320,15 +467,42 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#uimessage-part-ordering-invariants — Replace vs. push logic
    */
   handleTextMessageContentEvent(chunk) {
-    this.ensureAssistantMessage();
-    this.currentSegmentText += chunk.delta;
-    this.totalTextContent += chunk.delta;
+    const { messageId, state } = this.ensureAssistantMessage(chunk.messageId);
+    this.completeAllToolCallsForMessage(messageId);
+    const previousSegment = state.currentSegmentText;
+    const isNewSegment = state.hasToolCallsSinceTextStart && previousSegment.length > 0 && this.isNewTextSegment(chunk, previousSegment);
+    if (isNewSegment) {
+      if (previousSegment !== state.lastEmittedText) {
+        this.emitTextUpdateForMessage(messageId);
+      }
+      state.currentSegmentText = "";
+      state.lastEmittedText = "";
+      state.hasToolCallsSinceTextStart = false;
+    }
+    const currentText = state.currentSegmentText;
+    let nextText = currentText;
+    const delta = chunk.delta || "";
+    if (delta !== "") {
+      nextText = currentText + delta;
+    } else if (chunk.content !== void 0 && chunk.content !== "") {
+      if (chunk.content.startsWith(currentText)) {
+        nextText = chunk.content;
+      } else if (currentText.startsWith(chunk.content)) {
+        nextText = currentText;
+      } else {
+        nextText = currentText + chunk.content;
+      }
+    }
+    const textDelta = nextText.slice(currentText.length);
+    state.currentSegmentText = nextText;
+    state.totalTextContent += textDelta;
+    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);
     }
   }
   /**
@@ -345,9 +519,13 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#adapter-contract — Ordering requirements
    */
   handleToolCallStartEvent(chunk) {
-    this.ensureAssistantMessage();
+    const targetMessageId = chunk.parentMessageId ?? this.getActiveAssistantMessageId();
+    const { messageId, state } = this.ensureAssistantMessage(
+      targetMessageId ?? void 0
+    );
+    state.hasToolCallsSinceTextStart = true;
     const toolCallId = chunk.toolCallId;
-    const existingToolCall = this.toolCalls.get(toolCallId);
+    const existingToolCall = state.toolCalls.get(toolCallId);
     if (!existingToolCall) {
       const initialState = "awaiting-input";
       const newToolCall = {
@@ -356,29 +534,24 @@ class StreamProcessor {
         arguments: "",
         state: initialState,
         parsedArguments: void 0,
-        index: chunk.index ?? this.toolCalls.size
+        index: chunk.index ?? state.toolCalls.size
       };
-      this.toolCalls.set(toolCallId, newToolCall);
-      this.toolCallOrder.push(toolCallId);
-      if (this.currentAssistantMessageId) {
-        this.messages = updateToolCallPart(
-          this.messages,
-          this.currentAssistantMessageId,
-          {
-            id: chunk.toolCallId,
-            name: chunk.toolName,
-            arguments: "",
-            state: initialState
-          }
-        );
-        this.emitMessagesChange();
-        this.events.onToolCallStateChange?.(
-          this.currentAssistantMessageId,
-          chunk.toolCallId,
-          initialState,
-          ""
-        );
-      }
+      state.toolCalls.set(toolCallId, newToolCall);
+      state.toolCallOrder.push(toolCallId);
+      this.toolCallToMessage.set(toolCallId, messageId);
+      this.messages = updateToolCallPart(this.messages, messageId, {
+        id: chunk.toolCallId,
+        name: chunk.toolName,
+        arguments: "",
+        state: initialState
+      });
+      this.emitMessagesChange();
+      this.events.onToolCallStateChange?.(
+        messageId,
+        chunk.toolCallId,
+        initialState,
+        ""
+      );
     }
   }
   /**
@@ -395,36 +568,33 @@ class StreamProcessor {
    */
   handleToolCallArgsEvent(chunk) {
     const toolCallId = chunk.toolCallId;
-    const existingToolCall = this.toolCalls.get(toolCallId);
-    if (existingToolCall) {
-      const wasAwaitingInput = existingToolCall.state === "awaiting-input";
-      existingToolCall.arguments += chunk.delta || "";
-      if (wasAwaitingInput && chunk.delta) {
-        existingToolCall.state = "input-streaming";
-      }
-      existingToolCall.parsedArguments = this.jsonParser.parse(
-        existingToolCall.arguments
-      );
-      if (this.currentAssistantMessageId) {
-        this.messages = updateToolCallPart(
-          this.messages,
-          this.currentAssistantMessageId,
-          {
-            id: existingToolCall.id,
-            name: existingToolCall.name,
-            arguments: existingToolCall.arguments,
-            state: existingToolCall.state
-          }
-        );
-        this.emitMessagesChange();
-        this.events.onToolCallStateChange?.(
-          this.currentAssistantMessageId,
-          existingToolCall.id,
-          existingToolCall.state,
-          existingToolCall.arguments
-        );
-      }
+    const messageId = this.toolCallToMessage.get(toolCallId);
+    if (!messageId) return;
+    const state = this.getMessageState(messageId);
+    if (!state) return;
+    const existingToolCall = state.toolCalls.get(toolCallId);
+    if (!existingToolCall) return;
+    const wasAwaitingInput = existingToolCall.state === "awaiting-input";
+    existingToolCall.arguments += chunk.delta || "";
+    if (wasAwaitingInput && chunk.delta) {
+      existingToolCall.state = "input-streaming";
     }
+    existingToolCall.parsedArguments = this.jsonParser.parse(
+      existingToolCall.arguments
+    );
+    this.messages = updateToolCallPart(this.messages, messageId, {
+      id: existingToolCall.id,
+      name: existingToolCall.name,
+      arguments: existingToolCall.arguments,
+      state: existingToolCall.state
+    });
+    this.emitMessagesChange();
+    this.events.onToolCallStateChange?.(
+      messageId,
+      existingToolCall.id,
+      existingToolCall.state,
+      existingToolCall.arguments
+    );
   }
   /**
    * Handle TOOL_CALL_END event — authoritative signal that a tool call's input is finalized.
@@ -441,16 +611,19 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#single-shot-tool-call-response — End-to-end flow
    */
   handleToolCallEndEvent(chunk) {
-    const existingToolCall = this.toolCalls.get(chunk.toolCallId);
+    const messageId = this.toolCallToMessage.get(chunk.toolCallId);
+    if (!messageId) return;
+    const msgState = this.getMessageState(messageId);
+    if (!msgState) return;
+    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 (chunk.input !== void 0) {
         existingToolCall.parsedArguments = chunk.input;
       }
     }
-    if (this.currentAssistantMessageId && chunk.result) {
-      const state = "complete";
+    if (chunk.result) {
       let output;
       try {
         output = JSON.parse(chunk.result);
@@ -462,36 +635,55 @@ class StreamProcessor {
         chunk.toolCallId,
         output
       );
+      const resultState = "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.
+   */
+  handleRunStartedEvent(chunk) {
+    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
    */
   handleRunFinishedEvent(chunk) {
     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();
+    }
   }
   /**
    * Handle RUN_ERROR event
    */
   handleRunErrorEvent(chunk) {
     this.hasError = true;
+    if (chunk.runId) {
+      this.activeRuns.delete(chunk.runId);
+    } else {
+      this.activeRuns.clear();
+    }
     this.ensureAssistantMessage();
     this.events.onError?.(new Error(chunk.error.message || "An error occurred"));
   }
@@ -504,20 +696,30 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#thinkingreasoning-content — Thinking flow
    */
   handleStepFinishedEvent(chunk) {
-    this.ensureAssistantMessage();
-    this.thinkingContent += chunk.delta;
-    if (this.currentAssistantMessageId) {
-      this.messages = updateThinkingPart(
-        this.messages,
-        this.currentAssistantMessageId,
-        this.thinkingContent
-      );
-      this.emitMessagesChange();
-      this.events.onThinkingUpdate?.(
-        this.currentAssistantMessageId,
-        this.thinkingContent
-      );
+    const { messageId, state } = this.ensureAssistantMessage(
+      this.getActiveAssistantMessageId() ?? void 0
+    );
+    const previous = state.thinkingContent;
+    let nextThinking = previous;
+    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;
+      }
     }
+    state.thinkingContent = nextThinking;
+    this.messages = updateThinkingPart(
+      this.messages,
+      messageId,
+      state.thinkingContent
+    );
+    this.emitMessagesChange();
+    this.events.onThinkingUpdate?.(messageId, state.thinkingContent);
   }
   /**
    * Handle CUSTOM event.
@@ -530,6 +732,7 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#client-tools-and-approval-flows — Full flow details
    */
   handleCustomEvent(chunk) {
+    const messageId = this.getActiveAssistantMessageId();
     if (chunk.name === "tool-input-available" && chunk.value) {
       const { toolCallId, toolName, input } = chunk.value;
       this.events.onToolCall?.({
@@ -541,10 +744,11 @@ class StreamProcessor {
     }
     if (chunk.name === "approval-requested" && chunk.value) {
       const { toolCallId, toolName, input, approval } = chunk.value;
-      if (this.currentAssistantMessageId) {
+      const resolvedMessageId = messageId ?? this.toolCallToMessage.get(toolCallId);
+      if (resolvedMessageId) {
         this.messages = updateToolCallApproval(
           this.messages,
-          this.currentAssistantMessageId,
+          resolvedMessageId,
           toolCallId,
           approval.id
         );
@@ -563,8 +767,25 @@ class StreamProcessor {
       this.events.onCustomEvent(chunk.name, chunk.value, { toolCallId });
     }
   }
+  // ============================================
+  // Internal Helpers
+  // ============================================
   /**
-   * Complete all tool calls — safety net for stream termination.
+   * Detect if an incoming content chunk represents a NEW text segment
+   */
+  isNewTextSegment(chunk, previous) {
+    if (chunk.content !== void 0) {
+      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
@@ -573,41 +794,45 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#single-shot-tool-call-response — Safety net behavior
    */
   completeAllToolCalls() {
-    this.toolCalls.forEach((toolCall, id) => {
+    for (const messageId of this.activeMessageIds) {
+      this.completeAllToolCallsForMessage(messageId);
+    }
+  }
+  /**
+   * Complete all tool calls for a specific message
+   */
+  completeAllToolCallsForMessage(messageId) {
+    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);
       }
     });
   }
   /**
    * Mark a tool call as complete and emit event
    */
-  completeToolCall(_index, toolCall) {
+  completeToolCall(messageId, _index, toolCall) {
     toolCall.state = "input-complete";
     toolCall.parsedArguments = this.jsonParser.parse(toolCall.arguments);
-    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.events.onToolCallStateChange?.(
-        this.currentAssistantMessageId,
-        toolCall.id,
-        "input-complete",
-        toolCall.arguments
-      );
-    }
+    this.messages = updateToolCallPart(this.messages, messageId, {
+      id: toolCall.id,
+      name: toolCall.name,
+      arguments: toolCall.arguments,
+      state: "input-complete"
+    });
+    this.emitMessagesChange();
+    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).
@@ -615,20 +840,17 @@ class StreamProcessor {
    *
    * @see docs/chat-architecture.md#uimessage-part-ordering-invariants — Replace vs. push logic
    */
-  emitTextUpdate() {
-    this.lastEmittedText = this.currentSegmentText;
-    if (this.currentAssistantMessageId) {
-      this.messages = updateTextPart(
-        this.messages,
-        this.currentAssistantMessageId,
-        this.currentSegmentText
-      );
-      this.emitMessagesChange();
-      this.events.onTextUpdate?.(
-        this.currentAssistantMessageId,
-        this.currentSegmentText
-      );
-    }
+  emitTextUpdateForMessage(messageId) {
+    const state = this.getMessageState(messageId);
+    if (!state) return;
+    state.lastEmittedText = state.currentSegmentText;
+    this.messages = updateTextPart(
+      this.messages,
+      messageId,
+      state.currentSegmentText
+    );
+    this.emitMessagesChange();
+    this.events.onTextUpdate?.(messageId, state.currentSegmentText);
   }
   /**
    * Emit messages change event
@@ -646,66 +868,94 @@ class StreamProcessor {
    * @see docs/chat-architecture.md#single-shot-text-response — Finalization step
    */
   finalizeStream() {
-    this.completeAllToolCalls();
-    if (this.currentSegmentText !== this.lastEmittedText) {
-      this.emitTextUpdate();
+    let lastAssistantMessage;
+    for (const messageId of this.activeMessageIds) {
+      const state = this.getMessageState(messageId);
+      if (!state) continue;
+      this.completeAllToolCallsForMessage(messageId);
+      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;
+      }
     }
-    if (this.currentAssistantMessageId && !this.hasError) {
-      const assistantMessage = this.messages.find(
-        (m) => m.id === this.currentAssistantMessageId
-      );
-      if (assistantMessage && this.isWhitespaceOnlyMessage(assistantMessage)) {
+    this.activeMessageIds.clear();
+    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;
       }
     }
-    if (this.currentAssistantMessageId) {
-      const assistantMessage = this.messages.find(
-        (m) => m.id === this.currentAssistantMessageId
-      );
-      if (assistantMessage) {
-        this.events.onStreamEnd?.(assistantMessage);
-      }
+    if (lastAssistantMessage) {
+      this.events.onStreamEnd?.(lastAssistantMessage);
     }
   }
   /**
-   * Get completed tool calls in API format
+   * Get completed tool calls in API format (aggregated across all messages)
    */
   getCompletedToolCalls() {
-    return Array.from(this.toolCalls.values()).filter((tc) => tc.state === "input-complete").map((tc) => ({
-      id: tc.id,
-      type: "function",
-      function: {
-        name: tc.name,
-        arguments: tc.arguments
+    const result = [];
+    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",
+            function: {
+              name: tc.name,
+              arguments: tc.arguments
+            }
+          });
+        }
       }
-    }));
+    }
+    return result;
   }
   /**
-   * Get current result
+   * Get current result (aggregated across all messages)
    */
   getResult() {
     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 || void 0,
+      content,
+      thinking: thinking || void 0,
       toolCalls: toolCalls.length > 0 ? toolCalls : void 0,
       finishReason: this.finishReason
     };
   }
   /**
-   * Get current processor state
+   * Get current processor state (aggregated across all messages)
    */
   getState() {
+    let content = "";
+    let thinking = "";
+    const toolCalls = /* @__PURE__ */ new Map();
+    const toolCallOrder = [];
+    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
     };
@@ -732,12 +982,11 @@ class StreamProcessor {
    * Reset stream state (but keep messages)
    */
   resetStreamState() {
-    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;
@@ -749,7 +998,6 @@ class StreamProcessor {
   reset() {
     this.resetStreamState();
     this.messages = [];
-    this.currentAssistantMessageId = null;
   }
   /**
    * Check if a message contains only whitespace text and no other meaningful parts