@assistant-ui/react-ag-ui 0.1.0 → 0.1.2

package/src/runtime/AgUiThreadController.ts

@@ -62,14 +62,22 @@ import {
   tryParseJSON,
 } from "./adapter/run-aggregator";
 import {
-  fromAgUiMessages,
+  fromAgUiSnapshotMessages,
   toAgUiMessages,
   toAgUiTools,
+  withAgUiPartMessageMetadata,
 } from "./adapter/conversions";
 import { createAgUiSubscriber } from "./adapter/subscriber";
+import {
+  createReplayContentMerger,
+  groupReplayToolContent,
+  indexReplayToolOwners,
+  preserveToolResult,
+} from "./adapter/replay-content";
 import {
   buildToolApprovalResume,
   projectAgUiToolApprovals,
+  preserveToolApprovalDecision,
   withSettledToolApprovals,
   withToolApprovalDecision,
 } from "./adapter/tool-approval";
@@ -92,17 +100,38 @@ const isResolvedToolCall = (
 ): boolean =>
   part.type === "tool-call" && "result" in part && part.result !== undefined;
 
+const hasResolvedToolCalls = (message: ThreadMessage): boolean => {
+  if (message.role !== "assistant") return false;
+  for (const part of iterateToolCallParts(message.content)) {
+    if (isResolvedToolCall(part)) return true;
+  }
+  return false;
+};
+
 const listedPendingToolCallIds = (
-  assistant: ThreadAssistantMessage,
+  assistant: Pick<ThreadAssistantMessage, "metadata">,
 ): ReadonlySet<string> | undefined => {
   const ids = (
     assistant.metadata.custom[AG_UI_METADATA_NAMESPACE] as
       | AgUiCustomMetadata
       | undefined
   )?.pendingToolCallIds;
-  return ids?.length ? new Set(ids) : undefined;
+  return ids === undefined ? undefined : new Set(ids);
 };
 
+type ReplayContext = {
+  // Immutable history reference: partial argument frames cannot redefine consent.
+  readonly originalContent: ThreadAssistantMessage["content"];
+  accountedResultIds: ReadonlySet<string>;
+  admittedToolCallIds: ReadonlySet<string>;
+  merge: ReturnType<typeof createReplayContentMerger>;
+};
+
+type SnapshotMessageReconciler = (
+  snapshot: ThreadAssistantMessage,
+  previous: ThreadAssistantMessage | undefined,
+) => ThreadAssistantMessage;
+
 const withoutPendingToolCallIds = (
   metadata: ThreadAssistantMessage["metadata"],
 ): ThreadAssistantMessage["metadata"] => {
@@ -125,6 +154,43 @@ const withoutPendingToolCallIds = (
   };
 };
 
+const withPendingToolCallIds = (
+  metadata: ThreadAssistantMessage["metadata"],
+  ids: Iterable<string>,
+): ThreadAssistantMessage["metadata"] => ({
+  ...metadata,
+  custom: {
+    ...metadata.custom,
+    [AG_UI_METADATA_NAMESPACE]: {
+      ...(metadata.custom[AG_UI_METADATA_NAMESPACE] as
+        | AgUiCustomMetadata
+        | undefined),
+      pendingToolCallIds: [...ids],
+    },
+  },
+});
+
+const actionableToolCallIds = (
+  message: ThreadMessage | undefined,
+): ReadonlySet<string> => {
+  if (
+    message?.role !== "assistant" ||
+    message.status?.type !== "requires-action" ||
+    message.status.reason !== "tool-calls"
+  )
+    return new Set();
+  const permitted = listedPendingToolCallIds(message);
+  return new Set(
+    Array.from(iterateToolCallParts(message.content))
+      .filter(
+        (part) =>
+          !isResolvedToolCall(part) &&
+          (!permitted || permitted.has(part.toolCallId)),
+      )
+      .map((part) => part.toolCallId),
+  );
+};
+
 const isActivityPart = (part: ThreadAssistantMessage["content"][number]) =>
   part.type === "data" && part.name.startsWith("agui-activity/");
 
@@ -164,6 +230,73 @@ type CoreOptions = {
 
 const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
 
+const textOf = (content: ThreadAssistantMessage["content"]) =>
+  content.flatMap((part) => (part.type === "text" ? [part.text] : [])).join("");
+
+/**
+ * Swaps `content`'s text parts, in order, for `source`'s. A text part `source`
+ * has beyond `content`'s text slots goes after the nearest earlier part of
+ * `source` that `content` also has: a tool call by id, another part by its
+ * type and position among parts of that type. With none, it goes first.
+ */
+const replaceTextParts = (
+  content: ThreadAssistantMessage["content"],
+  source: ThreadAssistantMessage["content"],
+): ThreadAssistantMessage["content"] => {
+  const texts = source.filter((part) => part.type === "text");
+  const result: ThreadAssistantMessage["content"][number][] = [];
+  let used = 0;
+  for (const part of content) {
+    if (part.type !== "text") result.push(part);
+    else if (used < texts.length) result.push(texts[used++]!);
+  }
+  const positionIn = (
+    parts: readonly ThreadAssistantMessage["content"][number][],
+    part: ThreadAssistantMessage["content"][number],
+  ) => parts.filter((other) => other.type === part.type).indexOf(part);
+  for (const text of texts.slice(used)) {
+    let anchor = -1;
+    for (
+      let index = source.indexOf(text) - 1;
+      index >= 0 && anchor < 0;
+      index--
+    ) {
+      const before = source[index]!;
+      if (before.type === "tool-call") {
+        anchor = result.findIndex(
+          (part) =>
+            part.type === "tool-call" && part.toolCallId === before.toolCallId,
+        );
+      } else if (before.type === "text") {
+        anchor = result.indexOf(before);
+      } else {
+        const sameType = result.filter((part) => part.type === before.type);
+        const match = sameType[positionIn(source, before)];
+        anchor = match ? result.indexOf(match) : -1;
+      }
+    }
+    result.splice(anchor + 1, 0, text);
+  }
+  return result;
+};
+
+const settleNestedContent = (
+  content: ThreadAssistantMessage["content"],
+  status: MessageStatus,
+) =>
+  mapToolCallPartsDeep(content, (part) =>
+    part.messages
+      ? {
+          ...part,
+          messages: part.messages.map((message) =>
+            message.role === "assistant" && message.status?.type === "running"
+              ? { ...message, status }
+              : message,
+          ),
+        }
+      : part,
+  ).content;
+
 type AgUiRuntimeCallbackName = "onError" | "onCancel";
 
 const invokeRuntimeCallback = <TArgs extends readonly unknown[]>(
@@ -241,6 +374,10 @@ export class AgUiThreadController {
   private activeRunAgent: AbstractAgent | null = null;
   private history: ThreadHistoryAdapter | undefined;
   private readonly assistantHistoryParents = new Map<string, string | null>();
+  private readonly subagentParentToolCallIds = new Map<
+    string,
+    string | undefined
+  >();
   private readonly snapshotHistoryIds = new Set<string>();
   private readonly persistedHistoryIds = new Set<string>();
   private readonly historyWrites = createHistoryWriteChain();
@@ -347,6 +484,7 @@ export class AgUiThreadController {
 
       this.session.applyExternalMessageRepository(repo);
       this.assistantHistoryParents.clear();
+      this.subagentParentToolCallIds.clear();
       this.snapshotHistoryIds.clear();
       this.persistedHistoryIds.clear();
       for (const { message } of repo.messages) {
@@ -560,6 +698,7 @@ export class AgUiThreadController {
       config.runConfig ?? this.state.lastRunConfig,
       undefined,
       config.stream,
+      true,
     );
   }
 
@@ -599,17 +738,13 @@ export class AgUiThreadController {
   private findRequiresActionAssistant(
     reason: "interrupt" | "tool-calls",
   ): ThreadAssistantMessage | null {
-    const assistant = this.getMessages().findLast(
-      (message) => message.role === "assistant",
-    ) as ThreadAssistantMessage | undefined;
-    if (
-      !assistant ||
-      assistant.status?.type !== "requires-action" ||
-      assistant.status.reason !== reason
-    ) {
-      return null;
-    }
-    return assistant;
+    return (
+      this.getTrailingAssistantTurn().findLast(
+        (message) =>
+          message.status?.type === "requires-action" &&
+          message.status.reason === reason,
+      ) ?? null
+    );
   }
 
   private getTrailingAssistantTurn(): ThreadAssistantMessage[] {
@@ -715,8 +850,9 @@ export class AgUiThreadController {
     messageId: string,
     resume: AgUiResumeEntry[],
   ): Promise<void> {
+    const parentId = this.session.headId ?? messageId;
     this.clearPendingInterrupts(messageId, resume);
-    await this.startRun(messageId, this.state.lastRunConfig, {
+    await this.startRun(parentId, this.state.lastRunConfig, {
       entries: resume,
       interruptedMessageId: messageId,
     });
@@ -1004,6 +1140,7 @@ export class AgUiThreadController {
     });
     if (touched) {
       this.publishUpdate();
+      this.persistSettledAssistant(messageId);
     }
   }
 
@@ -1072,7 +1209,35 @@ export class AgUiThreadController {
         };
       });
       if (!matchedToolCall) return message;
-      return { ...assistant, content };
+      const settled = mapToolCallPartsDeep(content, (part) => {
+        if (!part.messages) return part;
+        let changed = false;
+        const messages = part.messages.map((nested) => {
+          if (
+            nested.role !== "assistant" ||
+            nested.status.type !== "requires-action" ||
+            nested.status.reason !== "tool-calls"
+          )
+            return nested;
+          const listed = listedPendingToolCallIds(nested);
+          if (
+            Array.from(iterateToolCallParts(nested.content)).some(
+              (call) =>
+                !isResolvedToolCall(call) &&
+                (!listed || listed.has(call.toolCallId)),
+            )
+          )
+            return nested;
+          changed = true;
+          return {
+            ...nested,
+            status: { type: "complete" as const, reason: "unknown" as const },
+            metadata: withoutPendingToolCallIds(nested.metadata),
+          };
+        });
+        return changed ? { ...part, messages } : part;
+      }).content;
+      return { ...assistant, content: settled };
     });
 
     if (!updated) return;
@@ -1180,15 +1345,26 @@ export class AgUiThreadController {
   // The continuation fires whether the frontend result lands before
   // RUN_FINISHED (the status flips to requires-action only later, while the
   // run is still draining) or after it.
-  private maybeResumeAfterToolResults(messageId: string): void {
+  private maybeResumeAfterToolResults(
+    messageId: string,
+    accountedResultIds?: ReadonlySet<string>,
+  ): void {
     if (!this.maybeCompleteAfterToolResults(messageId)) return;
+    if (accountedResultIds) {
+      const message = this.session.tryGetMessage(messageId)?.message;
+      if (
+        message?.role !== "assistant" ||
+        !Array.from(iterateToolCallParts(message.content)).some(
+          (part) =>
+            isResolvedToolCall(part) &&
+            !accountedResultIds.has(part.toolCallId),
+        )
+      )
+        return;
+    }
     if (this.getPendingInterrupts()) return;
     const turn = this.getTrailingAssistantTurn();
     if (!turn.some((message) => message.id === messageId)) return;
-    if (this.getPendingToolCalls()) return;
-    const headId = turn.at(-1)?.id;
-    if (headId === undefined) return;
-
     const owner = this.abortController;
     if (owner) {
       // A run is still draining (RUN_FINISHED arrived but the stream has not
@@ -1196,10 +1372,13 @@ export class AgUiThreadController {
       this.dispatch({
         type: "resumeDeferred",
         owner,
-        messageId: headId,
+        messageId,
       });
       return;
     }
+    if (this.getPendingToolCalls()) return;
+    const headId = turn.at(-1)?.id;
+    if (headId === undefined) return;
     this.startResumeRun(headId);
   }
 
@@ -1235,7 +1414,11 @@ export class AgUiThreadController {
     );
     if (!updated) return false;
     this.publishUpdate();
-    this.persistAssistantHistory(messageId);
+    if (this.history?.update) {
+      this.persistSettledAssistant(messageId);
+    } else {
+      this.persistAssistantHistory(messageId);
+    }
     return true;
   }
 
@@ -1269,6 +1452,19 @@ export class AgUiThreadController {
       let lastAppliedId: string | null = null;
       let hardReplace = false;
       const seen = new Set<string>();
+      const retainToolResults = (
+        existing: ThreadMessage | undefined,
+        message: ThreadMessage,
+      ): ThreadMessage =>
+        existing?.role === "assistant" && message.role === "assistant"
+          ? {
+              ...message,
+              content: this.preserveToolResults(
+                existing.content,
+                message.content,
+              ),
+            }
+          : message;
 
       for (const message of messages) {
         if (seen.has(message.id)) continue;
@@ -1278,12 +1474,21 @@ export class AgUiThreadController {
           hardReplace = true;
           break;
         }
-        this.session.addOrUpdateMessage(expectedParentId, message);
+        this.session.addOrUpdateMessage(
+          expectedParentId,
+          retainToolResults(existing?.message, message),
+        );
         expectedParentId = message.id;
         lastAppliedId = message.id;
       }
 
       if (hardReplace) {
+        const previous = new Map(
+          this.getMessageRepository().messages.map(({ message }) => [
+            message.id,
+            message,
+          ]),
+        );
         this.session.clear();
         expectedParentId = null;
         lastAppliedId = null;
@@ -1291,29 +1496,48 @@ export class AgUiThreadController {
         for (const message of messages) {
           if (seen.has(message.id)) continue;
           seen.add(message.id);
-          this.session.addOrUpdateMessage(expectedParentId, message);
+          this.session.addOrUpdateMessage(
+            expectedParentId,
+            retainToolResults(previous.get(message.id), message),
+          );
           expectedParentId = message.id;
           lastAppliedId = message.id;
         }
       }
 
       this.session.resetHead(lastAppliedId);
+      for (const messageId of seen) {
+        this.maybeCompleteAfterToolResults(messageId);
+      }
     }
 
     this.snapshotHistoryIds.clear();
     for (const { message } of this.getMessageRepository().messages) {
       this.snapshotHistoryIds.add(message.id);
     }
-    // MESSAGES_SNAPSHOT re-appends the active assistant, so a parked
-    // continuation whose target survived the snapshot is still answerable.
-    // Membership is the head branch, not the repository: a soft merge leaves
-    // the messages it dropped behind as off-branch nodes.
-    const resumeTarget = this.state.pendingResume?.messageId;
-    if (
-      resumeTarget !== undefined &&
-      !this.session.getMessages().some((message) => message.id === resumeTarget)
-    ) {
-      this.dispatch({ type: "resumeCleared" });
+    // A soft merge leaves dropped messages off branch; only selected
+    // messages with resolved tool calls can still own a continuation.
+    const pendingResume = this.state.pendingResume;
+    if (pendingResume) {
+      const selectedResultOwnerIds = new Set(
+        this.session
+          .getMessages()
+          .filter(hasResolvedToolCalls)
+          .map((message) => message.id),
+      );
+      const retained = pendingResume.messageIds.filter((id) =>
+        selectedResultOwnerIds.has(id),
+      );
+      if (retained.length !== pendingResume.messageIds.length) {
+        this.dispatch({ type: "resumeCleared" });
+        for (const messageId of retained) {
+          this.dispatch({
+            type: "resumeDeferred",
+            owner: pendingResume.owner,
+            messageId,
+          });
+        }
+      }
     }
     this.publishUpdate();
   }
@@ -1353,6 +1577,7 @@ export class AgUiThreadController {
     const activeRunAgent = this.activeRunAgent;
 
     this.dispatch({ type: "resetThread" });
+    this.subagentParentToolCallIds.clear();
 
     if (controller) {
       this.abortController = null;
@@ -1375,6 +1600,7 @@ export class AgUiThreadController {
     runConfig?: RunConfig,
     resume?: ResumeDispatch,
     resumeStream?: ResumeStream,
+    replayExistingMessages = false,
   ): Promise<void> {
     if (this.isThreadSwitching?.()) return;
     // A default AG-UI run supersedes the active run; the hook's opt-in message
@@ -1382,6 +1608,7 @@ export class AgUiThreadController {
     // its message; this covers the entry points that start a run without one.
     if (!this.supersedeActiveRun()) return;
 
+    const generation = this.state.historyMutationGeneration;
     const normalizedRunConfig = runConfig ?? {};
     this.dispatch({
       type: "runConfigRecorded",
@@ -1401,12 +1628,62 @@ export class AgUiThreadController {
     const runStartMessageIds = new Set(
       this.getMessageRepository().messages.map(({ message }) => message.id),
     );
+    const selectedMessages = this.session.getMessages();
+    const parentIndex = selectedMessages.findIndex(
+      (message) => message.id === parentId,
+    );
+    let replayStart = parentIndex + 1;
+    if (selectedMessages[parentIndex]?.role === "assistant") {
+      replayStart = parentIndex;
+      while (
+        replayStart > 0 &&
+        selectedMessages[replayStart - 1]?.role === "assistant"
+      )
+        replayStart--;
+    }
+    const replayMessageIds: string[] = [];
+    if (replayExistingMessages && !resumeStream && parentIndex >= 0) {
+      for (const message of selectedMessages.slice(replayStart)) {
+        if (message.role !== "assistant") break;
+        replayMessageIds.push(message.id);
+      }
+    }
+    // Only request-delivered and producer-observed results suppress continuation.
+    const accountedResultIds = new Set(
+      historicalMessages.flatMap((message) =>
+        message.role === "assistant"
+          ? Array.from(iterateToolCallParts(message.content))
+              .filter(isResolvedToolCall)
+              .map((part) => part.toolCallId)
+          : [],
+      ),
+    );
+    const candidateContexts = new Map<string, ReplayContext>();
+    for (const message of selectedMessages) {
+      if (
+        message.role !== "assistant" ||
+        !replayMessageIds.includes(message.id)
+      )
+        continue;
+      candidateContexts.set(message.id, {
+        originalContent: message.content,
+        accountedResultIds,
+        admittedToolCallIds: actionableToolCallIds(message),
+        merge: createReplayContentMerger(),
+      });
+    }
+    let replayContexts = new Map<string, ReplayContext>();
+    let lastReplayedIndex = -1;
+    let replayEnded = replayMessageIds.length === 0;
+    const replayTail = replayMessageIds.at(-1);
     const joinStrategy = this.joinStrategy;
     const hasRunUserBoundary = parent?.message.role === "user" && !resume;
 
     let pendingError: Error | null = null;
     const assistantParentId = parent ? parentId : this.session.headId;
     let assistantMessageId: string | undefined;
+    let snapshotAssistantId: string | undefined;
+    let replayInterruptMessageId: string | undefined;
     // A snapshot the preserve gate declines still evicts the in-flight
     // assistant; recreating under the cached id on the next content-bearing
     // emit keeps both the stream and the message identity. Status-only emits,
@@ -1416,8 +1693,11 @@ export class AgUiThreadController {
     // duplicate.
     let assistantCollided = false;
     const heldMessages = new Map<string, Set<string>>();
+    let frameOwnerIds = new Set<string>();
     let restrictedPendingToolCallIds: ReadonlySet<string> | undefined;
     let settlingRun = false;
+    let latestSnapshot: ThreadAssistantMessage | undefined;
+    let snapshotToolCallIds = new Set<string>();
     const ensureAssistant = (allowRecreate = false): string => {
       const cached = assistantMessageId;
       if (cached !== undefined && this.session.tryGetMessage(cached))
@@ -1427,10 +1707,16 @@ export class AgUiThreadController {
       }
       const repositoryHeadId = this.session.headId;
       const shouldUseSelectedParent =
-        cached === undefined ||
-        (shouldEagerlyInsertAssistant &&
-          (repositoryHeadId === null ||
-            runStartMessageIds.has(repositoryHeadId)));
+        replayContexts.size === 0 &&
+        (cached === undefined ||
+          (shouldEagerlyInsertAssistant &&
+            (repositoryHeadId === null ||
+              runStartMessageIds.has(repositoryHeadId)))) &&
+        // Replay candidates defer the placeholder; honor snapshots that advanced
+        // the head before the producer identified its first assistant message.
+        (replayMessageIds.length === 0 ||
+          repositoryHeadId === null ||
+          runStartMessageIds.has(repositoryHeadId));
       // Branch runs keep their selected parent unless the snapshot advanced to
       // a message introduced during this run.
       const parentId =
@@ -1445,23 +1731,59 @@ export class AgUiThreadController {
       return created;
     };
 
-    if (shouldEagerlyInsertAssistant) ensureAssistant();
+    if (shouldEagerlyInsertAssistant && replayEnded) ensureAssistant();
 
     const settleHeldMessages = (terminalStatus: MessageStatus): string[] => {
       const pendingMessages: string[] = [];
+      const producerPending = restrictedPendingToolCallIds;
+      const canWaitForTools =
+        terminalStatus.type === "complete" ||
+        terminalStatus.type === "requires-action";
+      if (producerPending && canWaitForTools) {
+        for (const message of this.getTrailingAssistantTurn()) {
+          const replay = replayContexts.get(message.id);
+          if (
+            message.id === assistantMessageId ||
+            heldMessages.has(message.id) ||
+            !replay ||
+            (message.status?.type === "requires-action" &&
+              message.status.reason === "interrupt")
+          )
+            continue;
+          if (
+            replay.admittedToolCallIds.size > 0 ||
+            actionableToolCallIds(message).size > 0 ||
+            Array.from(iterateToolCallParts(message.content)).some(
+              (part) =>
+                producerPending.has(part.toolCallId) &&
+                !accountedResultIds.has(part.toolCallId),
+            )
+          )
+            heldMessages.set(message.id, new Set());
+        }
+      }
       for (const [messageId, openToolCallIds] of heldMessages) {
+        const existing = this.session.tryGetMessage(messageId);
+        if (existing)
+          this.markPendingAssistantHistory(messageId, existing.parentId);
         let requiresToolResult = false;
         const updated = this.session.updateMessage(messageId, (message) => {
           if (message.role !== "assistant") return message;
           const assistant = message as ThreadAssistantMessage;
-          const pendingToolCallIds = Array.from(openToolCallIds).filter(
-            (id) =>
-              !restrictedPendingToolCallIds ||
-              restrictedPendingToolCallIds.has(id),
-          );
-          const canWaitForTools =
-            terminalStatus.type === "complete" ||
-            terminalStatus.type === "requires-action";
+          const pendingToolCallIds = producerPending
+            ? Array.from(iterateToolCallParts(assistant.content))
+                .filter(
+                  (part) =>
+                    producerPending.has(part.toolCallId) &&
+                    (!isResolvedToolCall(part) ||
+                      openToolCallIds.has(part.toolCallId) ||
+                      (replayContexts.has(messageId) &&
+                        !accountedResultIds.has(part.toolCallId))),
+                )
+                .map((part) => part.toolCallId)
+            : Array.from(openToolCallIds);
+          const replay = replayContexts.get(messageId);
+          if (replay) replay.admittedToolCallIds = new Set(pendingToolCallIds);
           const status: MessageStatus = canWaitForTools
             ? pendingToolCallIds.length > 0
               ? { type: "requires-action", reason: "tool-calls" }
@@ -1473,19 +1795,9 @@ export class AgUiThreadController {
             ...assistant,
             status,
             metadata:
-              status.type === "requires-action" && restrictedPendingToolCallIds
-                ? {
-                    ...metadata,
-                    custom: {
-                      ...metadata.custom,
-                      [AG_UI_METADATA_NAMESPACE]: {
-                        ...(metadata.custom[AG_UI_METADATA_NAMESPACE] as
-                          | AgUiCustomMetadata
-                          | undefined),
-                        pendingToolCallIds,
-                      },
-                    },
-                  }
+              status.type === "requires-action" &&
+              (restrictedPendingToolCallIds || replayContexts.has(messageId))
+                ? withPendingToolCallIds(metadata, pendingToolCallIds)
                 : metadata,
           };
         });
@@ -1501,24 +1813,190 @@ export class AgUiThreadController {
       return pendingMessages;
     };
 
+    const accountProducerResults = (
+      content: ThreadAssistantMessage["content"],
+    ) => {
+      for (const part of iterateToolCallParts(content))
+        if (isResolvedToolCall(part)) accountedResultIds.add(part.toolCallId);
+    };
+    const reconcileReplaySnapshot: SnapshotMessageReconciler = (
+      snapshot,
+      previous,
+    ) => {
+      // Observe the producer before preserving any local frontend results.
+      accountProducerResults(snapshot.content);
+      const replay = candidateContexts.get(snapshot.id);
+      if (!replay) return snapshot;
+      replay.admittedToolCallIds = actionableToolCallIds(snapshot);
+      if (heldMessages.has(snapshot.id))
+        heldMessages.set(snapshot.id, new Set(replay.admittedToolCallIds));
+      if (!replayContexts.has(snapshot.id)) return snapshot;
+      return previous
+        ? {
+            ...snapshot,
+            content: this.preserveToolResults(
+              previous.content,
+              snapshot.content,
+              replay.originalContent,
+            ),
+          }
+        : snapshot;
+    };
+
     const applyUpdate = (update: ChatModelRunResult) => {
+      if (
+        !assistantCollided &&
+        replayContexts.size &&
+        update.content !== undefined
+      )
+        accountProducerResults(
+          update.content as ThreadAssistantMessage["content"],
+        );
+      let ownerParts:
+        | Map<string, ThreadAssistantMessage["content"][number][]>
+        | undefined;
+      let contentOwnerIds: Set<string> | undefined;
+      const hasStreamContent =
+        Array.isArray(update.content) &&
+        update.content.some((part) => part.type !== "data");
+      if (!replayEnded && assistantMessageId === undefined) {
+        if (!update.status || update.status.type === "running") return;
+        // No matching producer identity arrived: retain native regeneration.
+        replayEnded = true;
+        if (
+          replayContexts.size === 0 &&
+          snapshotAssistantId !== undefined &&
+          this.session.headId === snapshotAssistantId &&
+          !hasStreamContent
+        ) {
+          if (update.status.type !== "complete") {
+            this.updateAssistantMessage(snapshotAssistantId, {
+              status: update.status,
+              ...(update.metadata ? { metadata: update.metadata } : {}),
+            });
+          }
+          return;
+        }
+      }
+      const streamMessageId = assistantMessageId;
+      let primaryMessageId = streamMessageId;
       const terminalStatus = update.status;
+      const replayInterrupt =
+        replayContexts.size > 0 &&
+        terminalStatus?.type === "requires-action" &&
+        terminalStatus.reason === "interrupt";
+      if (
+        streamMessageId !== undefined &&
+        !assistantCollided &&
+        replayContexts.size > 0 &&
+        (update.content !== undefined || replayInterrupt)
+      ) {
+        const owners = this.getTrailingAssistantTurn().filter((message) =>
+          replayContexts.has(message.id),
+        );
+        // Buffered calls keep their owners when a later text frame arrives.
+        const grouped =
+          update.content !== undefined
+            ? groupReplayToolContent(owners, update.content, streamMessageId)
+            : new Map([[streamMessageId, [...(update.content ?? [])]]]);
+        if (update.content !== undefined) {
+          const currentOwnerIds = new Set(grouped.keys());
+          for (const ownerId of frameOwnerIds)
+            grouped.set(ownerId, grouped.get(ownerId) ?? []);
+          contentOwnerIds = new Set(grouped.keys());
+          frameOwnerIds = currentOwnerIds;
+        }
+        ownerParts = grouped;
+        if (replayInterrupt) {
+          const interrupts = (
+            update.metadata?.custom?.[AG_UI_METADATA_NAMESPACE] as
+              | AgUiCustomMetadata
+              | undefined
+          )?.interrupts;
+          const approvals = projectAgUiToolApprovals(interrupts);
+          for (const owner of owners) {
+            if (
+              Array.from(iterateToolCallParts(owner.content)).some((part) =>
+                approvals.has(part.toolCallId),
+              )
+            )
+              grouped.set(owner.id, grouped.get(owner.id) ?? []);
+          }
+          const anchor = owners.find(
+            (message) =>
+              projectAgUiToolApprovals(
+                interrupts,
+                new Set(
+                  Array.from(
+                    iterateToolCallParts([
+                      ...message.content,
+                      ...(grouped.get(message.id) ?? []),
+                    ]),
+                  ).map((part) => part.toolCallId),
+                ),
+              ).size > 0,
+          );
+          if (anchor) {
+            primaryMessageId = anchor.id;
+            replayInterruptMessageId = anchor.id;
+            grouped.set(anchor.id, grouped.get(anchor.id) ?? []);
+          }
+        }
+      }
       const pendingMessages =
         settlingRun && terminalStatus && terminalStatus.type !== "running"
           ? settleHeldMessages(terminalStatus)
           : [];
-      const hasStreamContent =
-        Array.isArray(update.content) &&
-        update.content.some((part) => part.type !== "data");
+      const frameMessageId = ensureAssistant(hasStreamContent);
+      const messageId = primaryMessageId ?? frameMessageId;
+      let primaryUpdate = update;
+      if (ownerParts) {
+        const interrupt =
+          terminalStatus?.type === "requires-action" &&
+          terminalStatus.reason === "interrupt";
+        const ownerStatus: MessageStatus | undefined =
+          terminalStatus?.type === "requires-action"
+            ? { type: "complete", reason: "unknown" }
+            : terminalStatus;
+        for (const [ownerId, content] of ownerParts) {
+          const existing = this.session.tryGetMessage(ownerId);
+          if (existing)
+            this.markPendingAssistantHistory(ownerId, existing.parentId);
+          if (ownerId === messageId) continue;
+          this.updateAssistantMessage(
+            ownerId,
+            {
+              ...(contentOwnerIds?.has(ownerId) ? { content } : {}),
+              ...(ownerStatus ? { status: ownerStatus } : {}),
+            },
+            replayContexts.get(ownerId),
+            interrupt ? new Set<string>() : restrictedPendingToolCallIds,
+          );
+        }
+        const { content: _content, ...controlUpdate } = update;
+        const content = ownerParts.get(messageId) ?? [];
+        primaryUpdate = {
+          ...controlUpdate,
+          ...(contentOwnerIds?.has(messageId) ? { content } : {}),
+          ...(ownerStatus && !interrupt && replayContexts.has(messageId)
+            ? { status: ownerStatus }
+            : {}),
+        };
+      }
       const resolved = this.updateAssistantMessage(
-        ensureAssistant(hasStreamContent),
-        update,
+        messageId,
+        primaryUpdate,
+        replayContexts.get(messageId),
+        restrictedPendingToolCallIds,
       );
-      if (resolved !== assistantMessageId) {
+      if (messageId === frameMessageId && resolved !== assistantMessageId) {
         assistantMessageId = resolved;
       }
       for (const messageId of pendingMessages) {
-        this.maybeResumeAfterToolResults(messageId);
+        this.maybeResumeAfterToolResults(
+          messageId,
+          replayMessageIds.includes(messageId) ? accountedResultIds : undefined,
+        );
       }
     };
 
@@ -1527,20 +2005,74 @@ export class AgUiThreadController {
       startNewMessage: boolean,
       openToolCallIds: string[] = [],
     ) => {
+      const existing = this.session.tryGetMessage(serverId);
+      const replayIndex = replayMessageIds.indexOf(serverId);
+      const adoptsReplay =
+        !replayEnded &&
+        !assistantCollided &&
+        replayIndex > lastReplayedIndex &&
+        existing?.message.role === "assistant" &&
+        this.session.getMessages().some((message) => message.id === serverId);
+      if (adoptsReplay && replayContexts.size === 0) {
+        replayContexts = candidateContexts;
+        const selected = this.session.getMessages();
+        let tailId = serverId;
+        for (const message of selected.slice(
+          selected.findIndex((message) => message.id === serverId) + 1,
+        )) {
+          if (message.role !== "assistant") break;
+          tailId = message.id;
+        }
+        this.session.selectHead(tailId);
+      }
       if (startNewMessage && assistantMessageId !== undefined) {
-        if (openToolCallIds.length > 0) {
+        let hasReplayCallsToHold = false;
+        const replay = replayContexts.get(assistantMessageId);
+        if (replay) {
+          const current =
+            this.session.tryGetMessage(assistantMessageId)?.message;
+          if (current?.role === "assistant") {
+            const retainedPending = Array.from(
+              iterateToolCallParts(current.content),
+            ).filter(
+              (part) =>
+                !isResolvedToolCall(part) ||
+                (replay.admittedToolCallIds.has(part.toolCallId) &&
+                  !accountedResultIds.has(part.toolCallId)),
+            );
+            hasReplayCallsToHold = retainedPending.length > 0;
+            openToolCallIds = [
+              ...new Set([
+                ...openToolCallIds,
+                ...retainedPending
+                  .filter((part) =>
+                    replay.admittedToolCallIds.has(part.toolCallId),
+                  )
+                  .map((part) => part.toolCallId),
+              ]),
+            ];
+          }
+        }
+        if (openToolCallIds.length > 0 || hasReplayCallsToHold) {
           heldMessages.set(assistantMessageId, new Set(openToolCallIds));
         } else {
           applyUpdate({ status: { type: "complete", reason: "unknown" } });
         }
         const previousId = assistantMessageId;
         assistantMessageId = undefined;
-        if (this.session.tryGetMessage(previousId)) {
+        if (!adoptsReplay && this.session.tryGetMessage(previousId)) {
           const created = this.insertAssistantPlaceholder(previousId);
           assistantMessageId = created;
           this.markPendingAssistantHistory(created, previousId);
         }
       }
+      if (adoptsReplay) {
+        assistantMessageId = serverId;
+        lastReplayedIndex = replayIndex;
+        this.markPendingAssistantHistory(serverId, existing.parentId);
+        return;
+      }
+      replayEnded = true;
       const placeholder = ensureAssistant(true);
       if (placeholder === serverId) return;
       const reassigned = this.reassignAssistantId(placeholder, serverId);
@@ -1570,25 +2102,146 @@ export class AgUiThreadController {
       showThinking: this.showThinking,
       joinStrategy,
       logger: this.logger,
-      emit: applyUpdate,
+      emit: (update) => {
+        const currentSnapshotAssistant =
+          assistantMessageId === undefined
+            ? undefined
+            : this.session.tryGetMessage(assistantMessageId)?.message;
+        const listedSnapshotToolCallIds =
+          currentSnapshotAssistant?.role === "assistant"
+            ? listedPendingToolCallIds(currentSnapshotAssistant)
+            : undefined;
+        const hasPendingSnapshotToolCall =
+          latestSnapshot?.id === assistantMessageId &&
+          currentSnapshotAssistant?.role === "assistant" &&
+          Array.from(
+            iterateToolCallParts(currentSnapshotAssistant.content),
+          ).some(
+            (part) =>
+              !isResolvedToolCall(part) &&
+              (!listedSnapshotToolCallIds ||
+                listedSnapshotToolCallIds.has(part.toolCallId)) &&
+              (!restrictedPendingToolCallIds ||
+                restrictedPendingToolCallIds.has(part.toolCallId)),
+          );
+        const reconciledUpdate =
+          settlingRun &&
+          update.status?.type === "complete" &&
+          hasPendingSnapshotToolCall
+            ? {
+                ...update,
+                status: {
+                  type: "requires-action",
+                  reason: "tool-calls",
+                } as const,
+                ...(restrictedPendingToolCallIds
+                  ? {
+                      metadata: {
+                        ...update.metadata,
+                        custom: {
+                          ...update.metadata?.custom,
+                          [AG_UI_METADATA_NAMESPACE]: {
+                            ...(update.metadata?.custom?.[
+                              AG_UI_METADATA_NAMESPACE
+                            ] as AgUiCustomMetadata | undefined),
+                            pendingToolCallIds: [
+                              ...restrictedPendingToolCallIds,
+                            ],
+                          },
+                        },
+                      },
+                    }
+                  : {}),
+              }
+            : update;
+        if (
+          settlingRun &&
+          latestSnapshot &&
+          latestSnapshot.id === assistantMessageId &&
+          reconciledUpdate.status &&
+          reconciledUpdate.status.type !== "running" &&
+          reconciledUpdate.status.type !== "requires-action"
+        ) {
+          const status = reconciledUpdate.status;
+          const current =
+            this.session.tryGetMessage(assistantMessageId)?.message;
+          const content = settleNestedContent(
+            current?.role === "assistant"
+              ? current.content
+              : latestSnapshot.content,
+            status,
+          );
+          applyUpdate({ ...reconciledUpdate, content });
+        } else if (
+          latestSnapshot &&
+          latestSnapshot.id === assistantMessageId &&
+          reconciledUpdate.content !== undefined
+        ) {
+          const content = this.mergeSnapshotAssistantContent(
+            latestSnapshot.content,
+            reconciledUpdate.content as ThreadAssistantMessage["content"],
+            reconciledUpdate.content.some((part) => part.type === "reasoning"),
+            reconciledUpdate.content.some(isActivityPart),
+            snapshotToolCallIds,
+          );
+          applyUpdate({ ...reconciledUpdate, content });
+        } else {
+          applyUpdate(reconciledUpdate);
+        }
+        const current = assistantMessageId
+          ? this.session.tryGetMessage(assistantMessageId)?.message
+          : undefined;
+        latestSnapshot =
+          !settlingRun &&
+          current?.role === "assistant" &&
+          current.id === latestSnapshot?.id
+            ? current
+            : undefined;
+      },
       onServerMessageId: (serverId) => {
         adoptServerMessageId(serverId, false);
       },
-      onTextMessageStart: (serverId, openToolCallIds) =>
-        adoptServerMessageId(serverId, true, openToolCallIds),
+      onTextMessageStart: (serverId, openToolCallIds) => {
+        const owners = this.getTrailingAssistantTurn().filter((message) =>
+          replayContexts.has(message.id),
+        );
+        const toolOwners = indexReplayToolOwners(owners);
+        const frameOpenIds = openToolCallIds.filter((id) => {
+          const ownerId = toolOwners.get(id);
+          if (!ownerId || ownerId === assistantMessageId || assistantCollided)
+            return true;
+          const held = heldMessages.get(ownerId) ?? new Set<string>();
+          held.add(id);
+          heldMessages.set(ownerId, held);
+          return false;
+        });
+        // The native accumulator just reset; prior frame contributions are
+        // retained history, while removals within this frame remain authoritative.
+        for (const ownerId of frameOwnerIds) {
+          const replay = replayContexts.get(ownerId);
+          if (replay) replay.merge = createReplayContentMerger();
+        }
+        frameOwnerIds.clear();
+        adoptServerMessageId(serverId, true, frameOpenIds);
+      },
     });
     let runFinished = false;
     const dispatch = (event: AgUiEvent) => {
       if (this.abortController !== abortController) return;
       if (event.type === "RUN_FINISHED") {
         runFinished = true;
+        if (event.outcome?.type === "interrupt")
+          replayInterruptMessageId = assistantMessageId ?? replayTail;
         const ids =
           event.outcome?.type === "success"
             ? event.outcome.pendingToolCallIds
             : undefined;
-        restrictedPendingToolCallIds = ids?.length ? new Set(ids) : undefined;
+        restrictedPendingToolCallIds =
+          ids === undefined ? undefined : new Set(ids);
       }
       if (event.type === "TOOL_CALL_RESULT") {
+        if (!assistantCollided && candidateContexts.size)
+          accountedResultIds.add(event.toolCallId);
         for (const openToolCallIds of heldMessages.values()) {
           openToolCallIds.delete(event.toolCallId);
         }
@@ -1597,6 +2250,8 @@ export class AgUiThreadController {
         event.type === "RUN_FINISHED" ||
         event.type === "RUN_ERROR" ||
         event.type === "RUN_CANCELLED";
+      const beforeSnapshot =
+        event.type === "MESSAGES_SNAPSHOT" ? this.session.getMessages() : [];
       try {
         const nextAssistantMessageId = this.handleEvent(
           aggregator,
@@ -1605,9 +2260,39 @@ export class AgUiThreadController {
           joinStrategy,
           runStartMessageIds,
           hasRunUserBoundary,
+          candidateContexts.size ? reconcileReplaySnapshot : undefined,
         );
         if (event.type === "MESSAGES_SNAPSHOT") {
+          for (const message of beforeSnapshot) {
+            const context = replayContexts.get(message.id);
+            const current = this.session.tryGetMessage(message.id)?.message;
+            if (context && current?.content !== message.content) {
+              context.merge = createReplayContentMerger();
+              frameOwnerIds.delete(message.id);
+            }
+          }
           assistantMessageId = nextAssistantMessageId;
+          const imported = assistantMessageId
+            ? this.session.tryGetMessage(assistantMessageId)?.message
+            : undefined;
+          const source = event.messages.find(
+            (message) =>
+              message?.role === "assistant" && message.id === imported?.id,
+          );
+          latestSnapshot =
+            imported?.role === "assistant" && source ? imported : undefined;
+          snapshotToolCallIds = new Set(
+            Array.isArray(source?.toolCalls)
+              ? source.toolCalls
+                  .map((call: { id?: unknown }) => call.id)
+                  .filter((id: unknown): id is string => typeof id === "string")
+              : [],
+          );
+          if (!replayEnded && assistantMessageId === undefined) {
+            const tail = this.session.getMessages().at(-1);
+            snapshotAssistantId =
+              tail?.role === "assistant" ? tail.id : undefined;
+          }
         }
       } finally {
         settlingRun = false;
@@ -1664,7 +2349,6 @@ export class AgUiThreadController {
         );
         const subscriber = createAgUiSubscriber({
           dispatch,
-          runId,
           logger: this.logger,
           onRunFailed: (error) => {
             if (abortSignal.aborted) return;
@@ -1675,7 +2359,7 @@ export class AgUiThreadController {
         try {
           runAgentInstance.messages = input.messages;
           runAgentInstance.threadId = input.threadId;
-          runAgentInstance.state = input.state ?? null;
+          runAgentInstance.state = input.state;
         } catch {
           // ignore
         }
@@ -1693,6 +2377,26 @@ export class AgUiThreadController {
         pendingError = err;
       }
     } finally {
+      const current = assistantMessageId
+        ? this.session.tryGetMessage(assistantMessageId)?.message
+        : undefined;
+      if (
+        this.abortController === abortController &&
+        current?.role === "assistant" &&
+        current.status?.type === "running"
+      ) {
+        // Finalization is transport cleanup, not proof of a terminal event.
+        const status: MessageStatus = {
+          type: "incomplete",
+          reason: "error",
+          error: "The connection ended before the answer finished.",
+        };
+        settleHeldMessages(status);
+        applyUpdate({
+          status,
+          content: settleNestedContent(current.content, status),
+        });
+      }
       this.finishRun(abortController);
     }
 
@@ -1703,13 +2407,66 @@ export class AgUiThreadController {
       throw err;
     }
 
+    const pending = replayInterruptMessageId
+      ? this.getPendingInterrupts()
+      : null;
+    if (
+      pending &&
+      pending.messageId === replayInterruptMessageId &&
+      replayContexts.has(pending.messageId) &&
+      !abortSignal.aborted &&
+      this.abortController === null &&
+      generation === this.state.historyMutationGeneration &&
+      !this.isThreadSwitching?.()
+    ) {
+      const message = this.session.tryGetMessage(pending.messageId)?.message;
+      const responses =
+        message?.role === "assistant"
+          ? buildToolApprovalResume(message.content, pending.interrupts)
+          : null;
+      if (responses) {
+        this.clearDeferredContinuations(abortController);
+        try {
+          await this.submitInterruptResponses(responses);
+        } catch (error) {
+          const err = error instanceof Error ? error : new Error(String(error));
+          if (!this.reportedErrors.has(err)) {
+            this.reportError(err);
+            this.reportedErrors.add(err);
+          }
+          throw err;
+        }
+        return;
+      }
+    }
+
     // A tool result that landed before the run settled deferred its
     // continuation here so a second run never overlaps the first.
     if (this.state.pendingResume?.owner === abortController) {
-      const { messageId } = this.state.pendingResume;
+      const { messageIds } = this.state.pendingResume;
       this.dispatch({ type: "resumeCleared" });
-      if (!abortSignal.aborted) {
-        this.startResumeRun(messageId);
+      if (
+        !abortSignal.aborted &&
+        !this.getPendingInterrupts() &&
+        !this.getPendingToolCalls()
+      ) {
+        const turn = this.getTrailingAssistantTurn();
+        if (
+          turn.some((message) => {
+            if (!messageIds.includes(message.id)) return false;
+            const replay = replayContexts.get(message.id);
+            if (!replay) return hasResolvedToolCalls(message);
+            return Array.from(iterateToolCallParts(message.content)).some(
+              (part) =>
+                isResolvedToolCall(part) &&
+                replay.admittedToolCallIds.has(part.toolCallId) &&
+                !accountedResultIds.has(part.toolCallId),
+            );
+          })
+        ) {
+          const headId = turn.at(-1)?.id;
+          if (headId !== undefined) this.startResumeRun(headId);
+        }
       }
     }
 
@@ -1837,7 +2594,8 @@ export class AgUiThreadController {
     const input = {
       threadId,
       runId,
-      state: this.state.stateSnapshot ?? null,
+      state:
+        this.state.stateSnapshot === undefined ? {} : this.state.stateSnapshot,
       messages,
       tools: toAgUiTools(context?.tools),
       context: [
@@ -1994,6 +2752,8 @@ export class AgUiThreadController {
   private updateAssistantMessage(
     messageId: string,
     update: ChatModelRunResult,
+    replay?: ReplayContext,
+    pendingToolCallIds?: ReadonlySet<string>,
   ): string {
     let latestStatus: MessageStatus | undefined;
     const touched = this.session.updateMessage(messageId, (message) => {
@@ -2003,18 +2763,102 @@ export class AgUiThreadController {
         ? this.mergeAssistantMetadata(assistant.metadata, update.metadata)
         : assistant.metadata;
       latestStatus = update.status ?? assistant.status;
+      let content =
+        update.content !== undefined
+          ? replay
+            ? replay.merge(
+                assistant.content,
+                update.content as ThreadAssistantMessage["content"],
+              )
+            : this.preserveToolResults(
+                assistant.content,
+                update.content as ThreadAssistantMessage["content"],
+              )
+          : assistant.content;
+      const calls = replay ? Array.from(iterateToolCallParts(content)) : [];
+      // A retained historical call needs existing or current stream authority.
+      const pending =
+        pendingToolCallIds ??
+        (replay
+          ? new Set([
+              ...replay.admittedToolCallIds,
+              ...Array.from(iterateToolCallParts(update.content ?? []))
+                .filter((part) => !isResolvedToolCall(part))
+                .map((part) => part.toolCallId),
+            ])
+          : undefined);
+      if (replay && latestStatus?.type !== "running" && pending)
+        replay.admittedToolCallIds = pending;
+      if (replay) {
+        // Project gates against the content the user can actually answer.
+        const interrupts =
+          latestStatus?.type === "requires-action" &&
+          latestStatus.reason === "interrupt"
+            ? (
+                mergedMetadata.custom[AG_UI_METADATA_NAMESPACE] as
+                  | AgUiCustomMetadata
+                  | undefined
+              )?.interrupts
+            : undefined;
+        const approvals = projectAgUiToolApprovals(
+          interrupts,
+          new Set(calls.map((part) => part.toolCallId)),
+        );
+        const originalCalls = new Map(
+          Array.from(iterateToolCallParts(replay.originalContent), (part) => [
+            part.toolCallId,
+            part,
+          ]),
+        );
+        content = mapToolCallPartsDeep(content, (part) => {
+          const approval = approvals.get(part.toolCallId);
+          if (approval)
+            return {
+              ...part,
+              approval:
+                preserveToolApprovalDecision(
+                  originalCalls.get(part.toolCallId),
+                  { ...part, approval },
+                ) ?? approval,
+            };
+          if (
+            !part.approval ||
+            part.approval.approved !== undefined ||
+            part.approval.resolution
+          )
+            return part;
+          const { approval: _approval, ...withoutApproval } = part;
+          return withoutApproval;
+        }).content;
+      }
+      if (replay && latestStatus?.type === "complete") {
+        const listed =
+          pending ??
+          listedPendingToolCallIds({
+            ...assistant,
+            metadata: mergedMetadata,
+          });
+        if (
+          calls.some(
+            (part) =>
+              (!isResolvedToolCall(part) ||
+                !replay.accountedResultIds.has(part.toolCallId)) &&
+              (!listed || listed.has(part.toolCallId)),
+          )
+        )
+          latestStatus = { type: "requires-action", reason: "tool-calls" };
+      }
+      const admission =
+        pending && replay
+          ? withPendingToolCallIds(mergedMetadata, pending)
+          : mergedMetadata;
       const metadata =
         latestStatus?.type === "requires-action" &&
         latestStatus.reason === "tool-calls"
-          ? mergedMetadata
-          : withoutPendingToolCallIds(mergedMetadata);
-      const content =
-        update.content !== undefined
-          ? this.preserveToolResults(
-              assistant.content,
-              update.content as ThreadAssistantMessage["content"],
-            )
-          : assistant.content;
+          ? admission
+          : replay && latestStatus?.type === "running"
+            ? mergedMetadata
+            : withoutPendingToolCallIds(admission);
       return {
         ...assistant,
         content,
@@ -2037,7 +2881,10 @@ export class AgUiThreadController {
     if (this.isPersistableStatus(latestStatus)) {
       this.persistAssistantHistory(resolvedMessageId);
     }
-    this.maybeResumeAfterToolResults(resolvedMessageId);
+    this.maybeResumeAfterToolResults(
+      resolvedMessageId,
+      replay?.accountedResultIds,
+    );
     return resolvedMessageId;
   }
 
@@ -2048,6 +2895,7 @@ export class AgUiThreadController {
   private preserveToolResults(
     previous: ThreadAssistantMessage["content"],
     next: ThreadAssistantMessage["content"],
+    decisionSource?: ThreadAssistantMessage["content"],
   ): ThreadAssistantMessage["content"] {
     const preserved = new Map<string, ToolCallMessagePart>();
     for (const part of iterateToolCallParts(previous)) {
@@ -2055,25 +2903,27 @@ export class AgUiThreadController {
         preserved.set(part.toolCallId, part);
       }
     }
-    if (preserved.size === 0) {
+    const decisions = new Map(
+      Array.from(iterateToolCallParts(decisionSource ?? []))
+        .filter((part) => part.approval?.approved !== undefined)
+        .map((part) => [part.toolCallId, part]),
+    );
+    if (preserved.size === 0 && decisions.size === 0) {
       return this.preserveToolInteractions(previous, next);
     }
 
     const { content: merged, changed } = mapToolCallPartsDeep(next, (part) => {
       const prior = preserved.get(part.toolCallId);
-      if (!prior) return part;
-      if (!isResolvedToolCall(part) && isResolvedToolCall(prior)) {
-        return {
-          ...part,
-          result: prior.result,
-          ...(prior.artifact !== undefined ? { artifact: prior.artifact } : {}),
-          ...(prior.isError !== undefined ? { isError: prior.isError } : {}),
-          ...(prior.modelContent !== undefined
-            ? { modelContent: prior.modelContent }
-            : {}),
-        };
+      const result = prior ? preserveToolResult(prior, part) : part;
+      if (result.approval && decisions.has(part.toolCallId)) {
+        const approval = preserveToolApprovalDecision(
+          decisions.get(part.toolCallId),
+          result,
+        );
+        if (approval && approval !== result.approval)
+          return { ...result, approval };
       }
-      return part;
+      return result;
     });
     return this.preserveToolInteractions(previous, changed ? merged : next);
   }
@@ -2103,7 +2953,7 @@ export class AgUiThreadController {
   }
 
   // A snapshot built from user, assistant, and tool records alone carries no
-  // reasoning or activity of its own, so the parts streamed locally would be
+  // reasoning, activity, or custom data of its own, so locally streamed parts would be
   // lost on every refresh; keep the local ones where they streamed when the
   // snapshot has none of that kind, matching the merge rule of @ag-ui/client.
   private mergeSnapshotAssistantContent(
@@ -2111,15 +2961,30 @@ export class AgUiThreadController {
     next: ThreadAssistantMessage["content"],
     snapshotHasReasoning: boolean,
     snapshotHasActivity: boolean,
+    keepMissingToolCalls: ReadonlySet<string> = new Set(),
+    keepPriorTranscripts = false,
   ): ThreadAssistantMessage["content"] {
     const shouldKeep = (part: ThreadAssistantMessage["content"][number]) =>
-      part.type === "reasoning"
-        ? this.showThinking && !snapshotHasReasoning
-        : isActivityPart(part) && !snapshotHasActivity;
+      part.type === "tool-call"
+        ? keepMissingToolCalls.has(part.toolCallId) &&
+          !next.some(
+            (incoming) =>
+              incoming.type === "tool-call" &&
+              incoming.toolCallId === part.toolCallId,
+          )
+        : part.type === "reasoning"
+          ? this.showThinking && !snapshotHasReasoning
+          : part.type === "data" && !isActivityPart(part)
+            ? !next.some(
+                (incoming) =>
+                  incoming.type === "data" && incoming.name === part.name,
+              )
+            : isActivityPart(part) && !snapshotHasActivity;
     const kept = previous.filter(shouldKeep);
     const merged = this.reconcileSnapshotNestedContent(
       previous,
       this.preserveToolInteractions(previous, next),
+      keepPriorTranscripts,
     );
     if (kept.length === 0) return merged;
 
@@ -2169,6 +3034,7 @@ export class AgUiThreadController {
   private reconcileSnapshotNestedContent(
     previous: ThreadAssistantMessage["content"],
     next: ThreadAssistantMessage["content"],
+    keepPriorTranscripts: boolean,
   ): ThreadAssistantMessage["content"] {
     const nestedIds = new Set<string>();
     for (const part of previous) {
@@ -2218,10 +3084,92 @@ export class AgUiThreadController {
       }
       if (part.type !== "tool-call") return [part];
       const prior = previousCalls.get(part.toolCallId);
-      return prior?.messages === undefined
-        ? [part]
-        : [{ ...part, messages: prior.messages }];
+      if (prior?.messages === undefined) return [part];
+      if (part.messages === undefined)
+        return [{ ...part, messages: prior.messages }];
+      if (part.messages.length === 0) return [part];
+      return [
+        {
+          ...part,
+          messages: this.mergeNestedTranscripts(
+            prior.messages,
+            part.messages,
+            keepPriorTranscripts,
+          ),
+        },
+      ];
+    });
+  }
+
+  /**
+   * Merges one tool call's nested transcripts message by message. The incoming
+   * copy is the base for tool calls and reasoning; resolved results and calls
+   * only the prior copy has carry over, and the prior text wins when it extends
+   * the incoming text. Importing a snapshot keeps prior messages it lacks, a
+   * prior status that is running or incomplete, and the prior status when the
+   * snapshot waits on calls the merged content has already resolved.
+   */
+  private mergeNestedTranscripts(
+    prior: readonly ThreadMessage[],
+    next: readonly ThreadMessage[],
+    importingSnapshot: boolean,
+  ): readonly ThreadMessage[] {
+    const priorById = new Map(prior.map((message) => [message.id, message]));
+    const merged = next.map((message) => {
+      const previous = priorById.get(message.id);
+      if (previous?.role !== "assistant" || message.role !== "assistant")
+        return message;
+      let content = this.mergeSnapshotAssistantContent(
+        previous.content,
+        this.preserveToolResults(previous.content, message.content),
+        message.content.some((part) => part.type === "reasoning"),
+        message.content.some(isActivityPart),
+        new Set(
+          Array.from(
+            iterateToolCallParts(previous.content),
+            (call) => call.toolCallId,
+          ),
+        ),
+        importingSnapshot,
+      );
+      const priorText = textOf(previous.content);
+      const incomingText = textOf(message.content);
+      if (
+        priorText.length > incomingText.length &&
+        priorText.startsWith(incomingText)
+      ) {
+        content = replaceTextParts(content, previous.content);
+      }
+      const status = !importingSnapshot
+        ? message.status
+        : previous.status.type === "running" ||
+            previous.status.type === "incomplete" ||
+            (message.status.type === "requires-action" &&
+              message.status.reason === "tool-calls" &&
+              !Array.from(iterateToolCallParts(content)).some(
+                (call) => !isResolvedToolCall(call),
+              ))
+          ? previous.status
+          : message.status;
+      return {
+        ...message,
+        status,
+        content,
+        metadata: {
+          ...message.metadata,
+          custom: mergeAgUiNamespace(
+            previous.metadata.custom,
+            message.metadata.custom,
+          ),
+        },
+      };
     });
+    if (!importingSnapshot) return merged;
+    const mergedById = new Map(merged.map((message) => [message.id, message]));
+    return [
+      ...prior.map((message) => mergedById.get(message.id) ?? message),
+      ...merged.filter((message) => !priorById.has(message.id)),
+    ];
   }
 
   private mergeAssistantMetadata(
@@ -2260,7 +3208,17 @@ export class AgUiThreadController {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    reconcileSnapshot?: SnapshotMessageReconciler,
   ): string | undefined {
+    if (
+      event.type === "SUBAGENT_STARTED" &&
+      !this.subagentParentToolCallIds.has(event.subagentRunId)
+    ) {
+      this.subagentParentToolCallIds.set(
+        event.subagentRunId,
+        event.parentToolCallId,
+      );
+    }
     switch (event.type) {
       case "STATE_SNAPSHOT": {
         this.dispatch({
@@ -2290,6 +3248,7 @@ export class AgUiThreadController {
         break;
       }
       case "MESSAGES_SNAPSHOT": {
+        aggregator.handle(event);
         const previousHeadId = this.session.headId;
         const activeAssistant =
           activeAssistantId === undefined
@@ -2304,15 +3263,49 @@ export class AgUiThreadController {
           joinStrategy,
           runStartMessageIds,
           hasRunUserBoundary,
+          reconcileSnapshot,
         );
         const headId = this.session.headId;
-        if (hasActiveText && headId !== previousHeadId) {
-          const head =
-            headId === null
-              ? undefined
-              : this.session.tryGetMessage(headId)?.message;
-          activeAssistantId =
+        const head =
+          headId === null
+            ? undefined
+            : this.session.tryGetMessage(headId)?.message;
+        const hasNewPendingToolCallAnswer =
+          head?.role === "assistant" &&
+          !runStartMessageIds.has(head.id) &&
+          head.status?.type === "requires-action" &&
+          head.status.reason === "tool-calls";
+        if (
+          headId !== previousHeadId &&
+          (hasActiveText || hasNewPendingToolCallAnswer)
+        ) {
+          const nextActiveAssistantId =
             head?.role === "assistant" ? (headId ?? undefined) : undefined;
+          if (
+            activeAssistantId !== nextActiveAssistantId &&
+            activeAssistantId
+          ) {
+            const previous =
+              this.session.tryGetMessage(activeAssistantId)?.message;
+            if (
+              previous?.role === "assistant" &&
+              previous.status?.type === "running"
+            ) {
+              this.updateAssistantMessage(activeAssistantId, {
+                status: FALLBACK_USER_STATUS,
+                content: settleNestedContent(
+                  previous.content,
+                  FALLBACK_USER_STATUS,
+                ),
+              });
+            }
+          }
+          activeAssistantId = nextActiveAssistantId;
+          if (hasNewPendingToolCallAnswer && head) {
+            this.updateAssistantMessage(head.id, {
+              status: { type: "running" },
+            });
+          }
         }
         this.updateActiveAssistantState(activeAssistantId);
         break;
@@ -2415,6 +3408,14 @@ export class AgUiThreadController {
             ...(event.messageId
               ? { unstable_toolMessageId: event.messageId }
               : {}),
+            ...(event.metadata !== undefined
+              ? {
+                  providerMetadata: withAgUiPartMessageMetadata(
+                    part.providerMetadata,
+                    event.metadata,
+                  ),
+                }
+              : {}),
           };
         }
         const { modelContent: _modelContent, ...partWithoutModelContent } =
@@ -2441,6 +3442,14 @@ export class AgUiThreadController {
           ...(event.messageId
             ? { unstable_toolMessageId: event.messageId }
             : {}),
+          ...(event.metadata !== undefined
+            ? {
+                providerMetadata: withAgUiPartMessageMetadata(
+                  part.providerMetadata,
+                  event.metadata,
+                ),
+              }
+            : {}),
         };
       });
       if (!matchedToolCall) return message;
@@ -2451,7 +3460,8 @@ export class AgUiThreadController {
     this.publishUpdate();
     // Not maybeResumeAfterToolResults: the delivering run is already in
     // flight, and a resume from the owner would reset the head past it.
-    this.maybeCompleteAfterToolResults(messageId);
+    if (this.maybeCompleteAfterToolResults(messageId)) return;
+    this.persistSettledAssistant(messageId);
   }
 
   private applyCrossRunActivitySnapshot(
@@ -2512,7 +3522,8 @@ export class AgUiThreadController {
 
     if (!updated) return;
     this.publishUpdate();
-    this.maybeCompleteAfterToolResults(messageId);
+    if (this.maybeCompleteAfterToolResults(messageId)) return;
+    this.persistSettledAssistant(messageId);
   }
 
   private importMessagesSnapshot(
@@ -2521,6 +3532,7 @@ export class AgUiThreadController {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    reconcileSnapshot?: SnapshotMessageReconciler,
   ) {
     try {
       const activeAssistantItem = activeAssistantId
@@ -2530,14 +3542,24 @@ export class AgUiThreadController {
       const activeAssistantParentId = activeAssistantItem?.parentId;
       const activeAssistant =
         activeMessage?.role === "assistant" ? activeMessage : undefined;
-      const normalized = fromAgUiMessages(rawMessages, {
-        showThinking: this.showThinking,
-      });
+      const normalized = fromAgUiSnapshotMessages(
+        rawMessages,
+        { showThinking: this.showThinking },
+        this.subagentParentToolCallIds,
+      );
       const converted: ThreadMessage[] = [];
       for (const message of normalized) {
         try {
+          const previous = message.id
+            ? this.session.tryGetMessage(message.id)?.message
+            : undefined;
+          const fallbackStatus =
+            previous?.role === "assistant" &&
+            previous.status?.type === "running"
+              ? previous.status
+              : FALLBACK_USER_STATUS;
           converted.push(
-            fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS),
+            fromThreadMessageLike(message, generateId(), fallbackStatus),
           );
         } catch (error) {
           this.logger.error?.(
@@ -2560,19 +3582,30 @@ export class AgUiThreadController {
         const existing = this.session.tryGetMessage(
           convertedMessage.id,
         )?.message;
-        if (
-          convertedMessage.role === "assistant" &&
-          existing?.role === "assistant"
-        ) {
-          converted[index] = {
-            ...convertedMessage,
-            content: this.mergeSnapshotAssistantContent(
-              existing.content,
-              convertedMessage.content,
-              snapshotHasReasoning,
-              snapshotHasActivity,
-            ),
-          };
+        if (convertedMessage.role === "assistant") {
+          const previous =
+            existing?.role === "assistant" ? existing : undefined;
+          const snapshot =
+            reconcileSnapshot?.(convertedMessage, previous) ?? convertedMessage;
+          const preservesRunningStatus =
+            convertedMessage.id === activeAssistantId &&
+            previous?.status?.type === "running" &&
+            snapshot.status?.type === "requires-action" &&
+            snapshot.status.reason === "tool-calls";
+          converted[index] = previous
+            ? {
+                ...snapshot,
+                ...(preservesRunningStatus ? { status: previous.status } : {}),
+                content: this.mergeSnapshotAssistantContent(
+                  previous.content,
+                  snapshot.content,
+                  snapshotHasReasoning,
+                  snapshotHasActivity,
+                  new Set(),
+                  true,
+                ),
+              }
+            : snapshot;
         }
       }
       if (joinStrategy === "concat-content") {
@@ -2732,6 +3765,19 @@ export class AgUiThreadController {
     this.assistantHistoryParents.set(messageId, parentId);
   }
 
+  private persistSettledAssistant(messageId: string): void {
+    const item = this.session.tryGetMessage(messageId);
+    if (!this.history?.update || !item) return;
+    if (!this.isTerminalStatus(item.message.status)) return;
+    if (
+      !this.persistedHistoryIds.has(messageId) &&
+      this.historyWrites.pending(messageId) === undefined
+    )
+      return;
+    this.markPendingAssistantHistory(messageId, item.parentId);
+    this.persistAssistantHistory(messageId);
+  }
+
   private persistAssistantHistory(messageId: string) {
     const history = this.history;
     if (!history) return;
@@ -2744,7 +3790,9 @@ export class AgUiThreadController {
     const update = history.update;
     const shouldUpdate =
       update !== undefined &&
-      (wasPersisted || this.snapshotHistoryIds.has(messageId));
+      (wasPersisted ||
+        this.snapshotHistoryIds.has(messageId) ||
+        this.historyWrites.pending(messageId) !== undefined);
 
     if (shouldUpdate) {
       const write = this.historyWrites.chain(messageId, () =>