@assistant-ui/react-ag-ui 0.0.71 → 0.1.1

package/src/runtime/{AgUiThreadRuntimeCore.ts → AgUiThreadController.ts}

@@ -24,6 +24,7 @@ import type {
 } from "@assistant-ui/core";
 import {
   appendToolInteraction,
+  createHistoryWriteChain,
   createMessageRepositorySession,
   invokeUserCallback,
   iterateToolCallParts,
@@ -37,8 +38,14 @@ import type {
   RunAgentResult,
 } from "@ag-ui/client";
 import jsonpatch, { type Operation } from "fast-json-patch";
+import {
+  createAgUiThreadState,
+  reduceAgUiThreadState,
+  type AgUiThreadAction,
+  type AgUiThreadState,
+} from "./agUiThreadState";
 import type { Logger } from "./logger";
-import { readMcpAppResourceUri } from "./mcp-tool-result";
+import { readMcpAppResourceUri, toolContentText } from "./mcp-tool-result";
 import type {
   AgUiCustomMetadata,
   AgUiEvent,
@@ -60,9 +67,16 @@ import {
   toAgUiTools,
 } 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";
@@ -85,6 +99,97 @@ 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: Pick<ThreadAssistantMessage, "metadata">,
+): ReadonlySet<string> | undefined => {
+  const ids = (
+    assistant.metadata.custom[AG_UI_METADATA_NAMESPACE] as
+      | AgUiCustomMetadata
+      | undefined
+  )?.pendingToolCallIds;
+  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"] => {
+  const namespace = metadata.custom[AG_UI_METADATA_NAMESPACE] as
+    | AgUiCustomMetadata
+    | undefined;
+  if (!namespace || !("pendingToolCallIds" in namespace)) return metadata;
+  const { pendingToolCallIds: _pendingToolCallIds, ...remainingNamespace } =
+    namespace;
+  const { [AG_UI_METADATA_NAMESPACE]: _namespace, ...remainingCustom } =
+    metadata.custom;
+  return {
+    ...metadata,
+    custom: {
+      ...remainingCustom,
+      ...(Object.keys(remainingNamespace).length > 0
+        ? { [AG_UI_METADATA_NAMESPACE]: remainingNamespace }
+        : {}),
+    },
+  };
+};
+
+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/");
 
@@ -155,7 +260,7 @@ const mergeAgUiNamespace = (
   return merged;
 };
 
-export class AgUiThreadRuntimeCore {
+export class AgUiThreadController {
   private agent: AbstractAgent;
   private logger: Logger;
   private showThinking: boolean;
@@ -194,27 +299,18 @@ export class AgUiThreadRuntimeCore {
       };
     },
   });
-  private isRunningFlag = false;
+  private state: AgUiThreadState = createAgUiThreadState();
   private abortController: AbortController | null = null;
   // The agent that started the active run. updateOptions can swap this.agent
   // mid-run, and cancelling has to reach the agent holding the live request.
   private activeRunAgent: AbstractAgent | null = null;
-  private stateSnapshot: ReadonlyJSONValue | undefined;
   private history: ThreadHistoryAdapter | undefined;
-  private lastRunConfig: RunConfig | undefined;
   private readonly assistantHistoryParents = new Map<string, string | null>();
   private readonly snapshotHistoryIds = new Set<string>();
   private readonly persistedHistoryIds = new Set<string>();
-  private readonly historyWrites = new Map<string, Promise<void>>();
-  private _isLoading = false;
+  private readonly historyWrites = createHistoryWriteChain();
   private _historyImportPromise: Promise<void> | undefined;
-  private historyMutationGeneration = 0;
   private _loadPromise: Promise<void> | undefined;
-  private _loadRequested = false;
-  private pendingResume: { owner: AbortController; messageId: string } | null =
-    null;
-  private pendingA2uiResumeOwner: AbortController | null = null;
-  private pendingA2uiAction: Record<string, unknown> | undefined;
 
   constructor(options: CoreOptions) {
     this.agent = options.agent;
@@ -237,12 +333,16 @@ export class AgUiThreadRuntimeCore {
 
   getSnapshot = () =>
     (this.snapshot ??= {
-      isLoading: this._isLoading,
-      isRunning: this.isRunningFlag,
+      isLoading: this.state.isLoading,
+      isRunning: this.state.isRunning,
       messageRepository: this.session.export(),
-      state: this.stateSnapshot,
+      state: this.state.stateSnapshot,
     });
 
+  private dispatch(action: AgUiThreadAction): void {
+    this.state = reduceAgUiThreadState(this.state, action);
+  }
+
   private publishUpdate(): void {
     this.snapshot = undefined;
     for (const listener of this.listeners) listener();
@@ -261,7 +361,7 @@ export class AgUiThreadRuntimeCore {
     const previousHistory = this.history;
     this.history = options.history;
     if (
-      this._loadRequested &&
+      this.state.loadRequested &&
       !this._loadPromise &&
       !previousHistory &&
       options.history &&
@@ -291,19 +391,19 @@ export class AgUiThreadRuntimeCore {
   }
 
   getState(): ReadonlyJSONValue | undefined {
-    return this.stateSnapshot;
+    return this.state.stateSnapshot;
   }
 
   isRunning(): boolean {
-    return this.isRunningFlag;
+    return this.state.isRunning;
   }
 
   get isLoading(): boolean {
-    return this._isLoading;
+    return this.state.isLoading;
   }
 
   __internal_load(): Promise<void> {
-    this._loadRequested = true;
+    this.dispatch({ type: "loadRequested" });
     if (this._loadPromise) return this._loadPromise;
     if (!this.history) return Promise.resolve();
 
@@ -326,7 +426,7 @@ export class AgUiThreadRuntimeCore {
       return repo;
     });
 
-    this._isLoading = true;
+    this.dispatch({ type: "loadStarted" });
     this._historyImportPromise = promise.then(
       () => {
         this._historyImportPromise = undefined;
@@ -343,7 +443,7 @@ export class AgUiThreadRuntimeCore {
           const resumeStream = this.history?.resume?.bind(this.history);
           await this.startRun(
             parentId,
-            this.lastRunConfig,
+            this.state.lastRunConfig,
             undefined,
             resumeStream,
           );
@@ -354,7 +454,9 @@ export class AgUiThreadRuntimeCore {
         this.reportError(error);
       })
       .finally(() => {
-        this._isLoading = false;
+        this.dispatch({
+          type: "loadFinished",
+        });
         this.publishUpdate();
       });
 
@@ -365,10 +467,10 @@ export class AgUiThreadRuntimeCore {
   private waitForHistoryImport(): Promise<boolean> | undefined {
     const historyImport = this._historyImportPromise;
     if (!historyImport) return undefined;
-    const generation = this.historyMutationGeneration;
+    const generation = this.state.historyMutationGeneration;
     return historyImport.then(
       () =>
-        generation === this.historyMutationGeneration &&
+        generation === this.state.historyMutationGeneration &&
         !this.isThreadSwitching?.(),
     );
   }
@@ -376,10 +478,10 @@ export class AgUiThreadRuntimeCore {
   private deferUntilHistoryImported(callback: () => void): boolean {
     const historyImport = this._historyImportPromise;
     if (!historyImport) return false;
-    const generation = this.historyMutationGeneration;
+    const generation = this.state.historyMutationGeneration;
     void historyImport.then(() => {
       if (
-        generation !== this.historyMutationGeneration ||
+        generation !== this.state.historyMutationGeneration ||
         this.isThreadSwitching?.()
       ) {
         return;
@@ -397,10 +499,10 @@ export class AgUiThreadRuntimeCore {
     const historyImport = this._historyImportPromise;
     if (historyImport) {
       const wasAtTail = message.parentId === this.session.headId;
-      const generation = this.historyMutationGeneration;
+      const generation = this.state.historyMutationGeneration;
       await historyImport;
       if (
-        generation !== this.historyMutationGeneration ||
+        generation !== this.state.historyMutationGeneration ||
         this.isThreadSwitching?.()
       )
         return;
@@ -434,10 +536,10 @@ export class AgUiThreadRuntimeCore {
 
   private maybeAutoCancelPendingToolCalls(): void {
     if (this.autoCancelPendingToolCalls === false) return;
-    const pending = this.getPendingToolCalls();
-    if (!pending) return;
-    this.cancelUnresolvedToolCalls(pending.messageId);
-    this.maybeCompleteAfterToolResults(pending.messageId);
+    for (const assistant of this.getPendingToolCallMessages()) {
+      this.cancelUnresolvedToolCalls(assistant.id);
+      this.maybeCompleteAfterToolResults(assistant.id);
+    }
   }
 
   private appendEntry(message: AppendMessage): string {
@@ -471,7 +573,9 @@ export class AgUiThreadRuntimeCore {
   }
 
   async cancel(): Promise<void> {
-    this.historyMutationGeneration++;
+    this.dispatch({
+      type: "historyMutationAdvanced",
+    });
     this.abortActiveRun();
   }
 
@@ -518,9 +622,10 @@ export class AgUiThreadRuntimeCore {
     this.assertNoPendingInterrupts();
     await this.startRun(
       config.parentId,
-      config.runConfig ?? this.lastRunConfig,
+      config.runConfig ?? this.state.lastRunConfig,
       undefined,
       config.stream,
+      true,
     );
   }
 
@@ -541,7 +646,7 @@ export class AgUiThreadRuntimeCore {
     try {
       await this.startRun(
         parentId,
-        this.lastRunConfig,
+        this.state.lastRunConfig,
         undefined,
         resumeStream,
       );
@@ -560,17 +665,32 @@ export class AgUiThreadRuntimeCore {
   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 (
+      this.getTrailingAssistantTurn().findLast(
+        (message) =>
+          message.status?.type === "requires-action" &&
+          message.status.reason === reason,
+      ) ?? null
+    );
+  }
+
+  private getTrailingAssistantTurn(): ThreadAssistantMessage[] {
+    const messages = this.getMessages();
+    const turn: ThreadAssistantMessage[] = [];
+    for (let index = messages.length - 1; index >= 0; index--) {
+      const message = messages[index];
+      if (message?.role !== "assistant") break;
+      turn.unshift(message as ThreadAssistantMessage);
     }
-    return assistant;
+    return turn;
+  }
+
+  private getPendingToolCallMessages(): ThreadAssistantMessage[] {
+    return this.getTrailingAssistantTurn().filter(
+      (message) =>
+        message.status?.type === "requires-action" &&
+        message.status.reason === "tool-calls",
+    );
   }
 
   getPendingInterrupts(): {
@@ -592,15 +712,21 @@ export class AgUiThreadRuntimeCore {
     messageId: string;
     toolCallIds: string[];
   } | null {
-    const assistant = this.findRequiresActionAssistant("tool-calls");
-    if (!assistant) return null;
     const toolCallIds: string[] = [];
-    for (const part of iterateToolCallParts(assistant.content)) {
-      if (isResolvedToolCall(part)) continue;
-      toolCallIds.push(part.toolCallId);
+    let messageId: string | undefined;
+    for (const assistant of this.getPendingToolCallMessages()) {
+      const listed = listedPendingToolCallIds(assistant);
+      for (const part of iterateToolCallParts(assistant.content)) {
+        if (
+          isResolvedToolCall(part) ||
+          (listed && !listed.has(part.toolCallId))
+        )
+          continue;
+        messageId ??= assistant.id;
+        toolCallIds.push(part.toolCallId);
+      }
     }
-    if (toolCallIds.length === 0) return null;
-    return { messageId: assistant.id, toolCallIds };
+    return messageId === undefined ? null : { messageId, toolCallIds };
   }
 
   async submitInterruptResponses(
@@ -632,6 +758,7 @@ export class AgUiThreadRuntimeCore {
     this.assertInterruptsAnswerable(
       "submitInterruptResponses",
       pending.interrupts,
+      responsesById,
     );
 
     await this.resumeWithResponses(
@@ -650,8 +777,9 @@ export class AgUiThreadRuntimeCore {
     messageId: string,
     resume: AgUiResumeEntry[],
   ): Promise<void> {
+    const parentId = this.session.headId ?? messageId;
     this.clearPendingInterrupts(messageId, resume);
-    await this.startRun(messageId, this.lastRunConfig, {
+    await this.startRun(parentId, this.state.lastRunConfig, {
       entries: resume,
       interruptedMessageId: messageId,
     });
@@ -679,6 +807,7 @@ export class AgUiThreadRuntimeCore {
   private assertInterruptsAnswerable(
     method: "submitInterruptResponses" | "respondToToolApproval",
     interrupts: readonly AgUiInterrupt[],
+    responsesById?: ReadonlyMap<string, AgUiResumeEntry>,
   ): void {
     const now = Date.now();
     for (const interrupt of interrupts) {
@@ -689,13 +818,16 @@ export class AgUiThreadRuntimeCore {
           `[agui] ${method}: interrupt ${interrupt.id} has malformed expiresAt "${interrupt.expiresAt}"`,
         );
       }
-      if (expiry <= now) {
+      if (
+        expiry <= now &&
+        responsesById?.get(interrupt.id)?.status !== "cancelled"
+      ) {
         throw new Error(
           `[agui] ${method}: interrupt ${interrupt.id} expired at ${interrupt.expiresAt}`,
         );
       }
     }
-    if (this.isRunningFlag) {
+    if (this.state.isRunning) {
       throw new Error(`[agui] ${method}: a run is already in progress`);
     }
   }
@@ -789,11 +921,13 @@ export class AgUiThreadRuntimeCore {
             "[agui] steerAway: responses are only valid for pending interrupts",
           );
         }
-        if (this.isRunningFlag) {
+        if (this.state.isRunning) {
           throw new Error("[agui] steerAway: a run is already in progress");
         }
-        this.cancelUnresolvedToolCalls(pendingTools.messageId);
-        this.maybeCompleteAfterToolResults(pendingTools.messageId);
+        for (const assistant of this.getPendingToolCallMessages()) {
+          this.cancelUnresolvedToolCalls(assistant.id);
+          this.maybeCompleteAfterToolResults(assistant.id);
+        }
         const normalized = this.toAppendMessage(message);
         const threadMessageId = this.appendEntry(normalized);
         await this.startRun(threadMessageId, normalized.runConfig);
@@ -810,7 +944,7 @@ export class AgUiThreadRuntimeCore {
 
     const resume = this.resolveSteerAwayResume(pending.interrupts, responses);
 
-    if (this.isRunningFlag) {
+    if (this.state.isRunning) {
       throw new Error("[agui] steerAway: a run is already in progress");
     }
 
@@ -1067,7 +1201,7 @@ export class AgUiThreadRuntimeCore {
       return;
     }
 
-    await this.chainHistoryWrite(sessionMessageId, () =>
+    await this.historyWrites.chain(sessionMessageId, () =>
       history.update!({ parentId: item.parentId, message }),
     );
   }
@@ -1090,11 +1224,17 @@ export class AgUiThreadRuntimeCore {
     if (!("timestamp" in userAction)) {
       userAction.timestamp = new Date().toISOString();
     }
-    this.pendingA2uiAction = userAction;
+    this.dispatch({
+      type: "a2uiActionDeferred",
+      action: userAction,
+    });
 
     const owner = this.abortController;
     if (owner) {
-      this.pendingA2uiResumeOwner = owner;
+      this.dispatch({
+        type: "a2uiResumeDeferred",
+        owner,
+      });
       return;
     }
     this.startResumeRun(parentId);
@@ -1103,17 +1243,41 @@ export class AgUiThreadRuntimeCore {
   // 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;
     const owner = this.abortController;
     if (owner) {
       // A run is still draining (RUN_FINISHED arrived but the stream has not
       // closed). Defer until startRun's tail so we never start two runs.
-      this.pendingResume = { owner, messageId };
+      this.dispatch({
+        type: "resumeDeferred",
+        owner,
+        messageId,
+      });
       return;
     }
-    this.startResumeRun(messageId);
+    if (this.getPendingToolCalls()) return;
+    const headId = turn.at(-1)?.id;
+    if (headId === undefined) return;
+    this.startResumeRun(headId);
   }
 
   private maybeCompleteAfterToolResults(messageId: string): boolean {
@@ -1127,7 +1291,9 @@ export class AgUiThreadRuntimeCore {
       return false;
     }
     let allResolved = true;
+    const listed = listedPendingToolCallIds(assistant);
     for (const part of iterateToolCallParts(assistant.content)) {
+      if (listed && !listed.has(part.toolCallId)) continue;
       if (!isResolvedToolCall(part)) {
         allResolved = false;
         break;
@@ -1140,6 +1306,7 @@ export class AgUiThreadRuntimeCore {
         ? {
             ...(current as ThreadAssistantMessage),
             status: { type: "complete" as const, reason: "unknown" as const },
+            metadata: withoutPendingToolCallIds(current.metadata),
           }
         : current,
     );
@@ -1150,7 +1317,7 @@ export class AgUiThreadRuntimeCore {
   }
 
   private startResumeRun(messageId: string): void {
-    void this.startRun(messageId, this.lastRunConfig).catch((error) => {
+    void this.startRun(messageId, this.state.lastRunConfig).catch((error) => {
       this.reportError(error);
     });
   }
@@ -1167,8 +1334,9 @@ export class AgUiThreadRuntimeCore {
         metadata: { ...message.metadata, modality },
       } as ThreadMessage;
     });
-    this.pendingA2uiResumeOwner = null;
-    this.pendingA2uiAction = undefined;
+    this.dispatch({
+      type: "a2uiContinuationsCleared",
+    });
     this.assistantHistoryParents.clear();
 
     if (messages.length === 0) {
@@ -1178,6 +1346,19 @@ export class AgUiThreadRuntimeCore {
       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;
@@ -1187,12 +1368,21 @@ export class AgUiThreadRuntimeCore {
           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;
@@ -1200,35 +1390,57 @@ export class AgUiThreadRuntimeCore {
         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.pendingResume?.messageId;
-    if (
-      resumeTarget !== undefined &&
-      !this.session.getMessages().some((message) => message.id === resumeTarget)
-    ) {
-      this.pendingResume = null;
+    // 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();
   }
 
   loadExternalState(state: ReadonlyJSONValue): void {
-    this.stateSnapshot = state;
+    this.dispatch({
+      type: "stateSnapshot",
+      snapshot: state,
+    });
     this.publishUpdate();
   }
 
@@ -1237,25 +1449,28 @@ export class AgUiThreadRuntimeCore {
       | ReadonlyJSONValue
       | ((prev: ReadonlyJSONValue | undefined) => ReadonlyJSONValue),
   ): void {
-    this.stateSnapshot =
-      typeof next === "function" ? next(this.stateSnapshot) : next;
+    const snapshot =
+      typeof next === "function" ? next(this.state.stateSnapshot) : next;
+    this.dispatch({
+      type: "stateSnapshot",
+      snapshot,
+    });
     this.publishUpdate();
   }
 
   resetState(): void {
-    this.stateSnapshot = undefined;
+    this.dispatch({ type: "resetState" });
     this.publishUpdate();
   }
 
   resetThreadState(): void {
-    this.historyMutationGeneration++;
+    this.dispatch({
+      type: "historyMutationAdvanced",
+    });
     const controller = this.abortController;
     const activeRunAgent = this.activeRunAgent;
 
-    this.stateSnapshot = undefined;
-    this.pendingResume = null;
-    this.pendingA2uiResumeOwner = null;
-    this.pendingA2uiAction = undefined;
+    this.dispatch({ type: "resetThread" });
 
     if (controller) {
       this.abortController = null;
@@ -1278,6 +1493,7 @@ export class AgUiThreadRuntimeCore {
     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
@@ -1285,8 +1501,12 @@ export class AgUiThreadRuntimeCore {
     // 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.lastRunConfig = normalizedRunConfig;
+    this.dispatch({
+      type: "runConfigRecorded",
+      runConfig: normalizedRunConfig,
+    });
     const parent =
       parentId === null ? undefined : this.session.tryGetMessage(parentId);
     const shouldEagerlyInsertAssistant =
@@ -1301,12 +1521,62 @@ export class AgUiThreadRuntimeCore {
     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,
@@ -1315,6 +1585,10 @@ export class AgUiThreadRuntimeCore {
     // resurrecting for custom-event data would leave a trailing data-only
     // duplicate.
     let assistantCollided = false;
+    const heldMessages = new Map<string, Set<string>>();
+    let frameOwnerIds = new Set<string>();
+    let restrictedPendingToolCallIds: ReadonlySet<string> | undefined;
+    let settlingRun = false;
     const ensureAssistant = (allowRecreate = false): string => {
       const cached = assistantMessageId;
       if (cached !== undefined && this.session.tryGetMessage(cached))
@@ -1324,10 +1598,16 @@ export class AgUiThreadRuntimeCore {
       }
       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 =
@@ -1342,35 +1622,348 @@ export class AgUiThreadRuntimeCore {
       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 = 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" }
+              : { type: "complete", reason: "unknown" }
+            : terminalStatus;
+          requiresToolResult = status.type === "requires-action";
+          const metadata = withoutPendingToolCallIds(assistant.metadata);
+          return {
+            ...assistant,
+            status,
+            metadata:
+              status.type === "requires-action" &&
+              (restrictedPendingToolCallIds || replayContexts.has(messageId))
+                ? withPendingToolCallIds(metadata, pendingToolCallIds)
+                : metadata,
+          };
+        });
+        if (updated) {
+          this.publishUpdate();
+          this.persistAssistantHistory(messageId);
+          if (requiresToolResult) {
+            pendingMessages.push(messageId);
+          }
+        }
+      }
+      heldMessages.clear();
+      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 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,
+          replayMessageIds.includes(messageId) ? accountedResultIds : undefined,
+        );
+      }
     };
 
     const adoptServerMessageId = (
       serverId: string,
       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) {
-        applyUpdate({ status: { type: "complete", reason: "unknown" } });
+        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);
@@ -1404,22 +1997,85 @@ export class AgUiThreadRuntimeCore {
       onServerMessageId: (serverId) => {
         adoptServerMessageId(serverId, false);
       },
-      onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true),
+      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;
-      const nextAssistantMessageId = this.handleEvent(
-        aggregator,
-        event,
-        assistantMessageId,
-        joinStrategy,
-        runStartMessageIds,
-        hasRunUserBoundary,
-      );
-      if (event.type === "MESSAGES_SNAPSHOT") {
-        assistantMessageId = nextAssistantMessageId;
+      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 === 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);
+        }
+      }
+      settlingRun =
+        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,
+          event,
+          assistantMessageId,
+          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;
+          if (!replayEnded && assistantMessageId === undefined) {
+            const tail = this.session.getMessages().at(-1);
+            snapshotAssistantId =
+              tail?.role === "assistant" ? tail.id : undefined;
+          }
+        }
+      } finally {
+        settlingRun = false;
       }
     };
 
@@ -1482,9 +2138,9 @@ export class AgUiThreadRuntimeCore {
           },
         });
         try {
-          (runAgentInstance as any).messages = input.messages;
-          (runAgentInstance as any).threadId = input.threadId;
-          (runAgentInstance as any).state = input.state ?? null;
+          runAgentInstance.messages = input.messages;
+          runAgentInstance.threadId = input.threadId;
+          runAgentInstance.state = input.state ?? null;
         } catch {
           // ignore
         }
@@ -1512,21 +2168,78 @@ export class AgUiThreadRuntimeCore {
       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.pendingResume?.owner === abortController) {
-      const { messageId } = this.pendingResume;
-      this.pendingResume = null;
-      if (!abortSignal.aborted) {
-        this.startResumeRun(messageId);
+    if (this.state.pendingResume?.owner === abortController) {
+      const { messageIds } = this.state.pendingResume;
+      this.dispatch({ type: "resumeCleared" });
+      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);
+        }
       }
     }
 
-    if (this.pendingA2uiResumeOwner === abortController) {
-      this.pendingA2uiResumeOwner = null;
-      if (!abortSignal.aborted && this.pendingA2uiAction !== undefined) {
+    if (this.state.pendingA2uiResumeOwner === abortController) {
+      this.dispatch({
+        type: "a2uiResumeCleared",
+      });
+      if (!abortSignal.aborted && this.state.pendingA2uiAction !== undefined) {
         if (this.getPendingInterrupts()) {
-          this.pendingA2uiAction = undefined;
+          this.dispatch({
+            type: "a2uiActionCleared",
+          });
           this.logger.debug(
             "[agui] sendA2uiAction: pending interrupts, dropping action",
           );
@@ -1537,13 +2250,17 @@ export class AgUiThreadRuntimeCore {
         if (parentId !== null) {
           this.startResumeRun(parentId);
         } else {
-          this.pendingA2uiAction = undefined;
+          this.dispatch({
+            type: "a2uiActionCleared",
+          });
           this.logger.debug(
             "[agui] sendA2uiAction: no messages to resume, dropping action",
           );
         }
       } else {
-        this.pendingA2uiAction = undefined;
+        this.dispatch({
+          type: "a2uiActionCleared",
+        });
       }
     }
   }
@@ -1562,8 +2279,9 @@ export class AgUiThreadRuntimeCore {
       getAssistantMessageId: () => string | undefined;
     },
   ): Promise<Error | undefined> {
-    this.pendingA2uiAction = undefined;
-    this.pendingA2uiResumeOwner = null;
+    this.dispatch({
+      type: "a2uiContinuationsCleared",
+    });
     const assistantId = ctx.ensureAssistant();
     const currentId = () => ctx.getAssistantMessageId() ?? assistantId;
     const options: ChatModelRunOptions = {
@@ -1614,8 +2332,10 @@ export class AgUiThreadRuntimeCore {
     historyMessages: readonly ThreadMessage[] | undefined,
     resume?: ResumeDispatch,
   ) {
-    const pendingA2uiAction = this.pendingA2uiAction;
-    this.pendingA2uiAction = undefined;
+    const pendingA2uiAction = this.state.pendingA2uiAction;
+    this.dispatch({
+      type: "a2uiActionCleared",
+    });
     const threadId = this.agent.threadId || "main";
     const messages = toAgUiMessages(
       this.selectRunMessages(
@@ -1635,7 +2355,7 @@ export class AgUiThreadRuntimeCore {
     const input = {
       threadId,
       runId,
-      state: this.stateSnapshot ?? null,
+      state: this.state.stateSnapshot ?? null,
       messages,
       tools: toAgUiTools(context?.tools),
       context: [
@@ -1678,17 +2398,20 @@ export class AgUiThreadRuntimeCore {
   }
 
   private clearDeferredContinuations(owner: AbortController): void {
-    if (this.pendingResume?.owner === owner) {
-      this.pendingResume = null;
+    if (this.state.pendingResume?.owner === owner) {
+      this.dispatch({ type: "resumeCleared" });
     }
-    if (this.pendingA2uiResumeOwner === owner) {
-      this.pendingA2uiResumeOwner = null;
-      this.pendingA2uiAction = undefined;
+    if (this.state.pendingA2uiResumeOwner === owner) {
+      this.dispatch({
+        type: "a2uiContinuationsCleared",
+      });
     }
   }
 
   private setRunning(running: boolean) {
-    this.isRunningFlag = running;
+    this.dispatch({
+      type: running ? "runStarted" : "runFinished",
+    });
     this.publishUpdate();
   }
 
@@ -1710,7 +2433,7 @@ export class AgUiThreadRuntimeCore {
       status: { type: "running" },
       content: [],
       metadata: {
-        unstable_state: this.stateSnapshot ?? null,
+        unstable_state: this.state.stateSnapshot ?? null,
         unstable_annotations: [],
         unstable_data: [],
         steps: [],
@@ -1768,22 +2491,18 @@ export class AgUiThreadRuntimeCore {
     // An in-flight write completes under the id it was started with, so the
     // rename moves the chain entry and transfers the completion mark; without
     // this the resolve records the dead id and the live one appends again.
-    const pendingWrite = this.historyWrites.get(oldId);
-    if (pendingWrite) {
-      this.historyWrites.delete(oldId);
-      if (!collidesWithExisting) {
-        this.historyWrites.set(newId, pendingWrite);
-        const settle = () => {
-          if (this.historyWrites.get(newId) === pendingWrite) {
-            this.historyWrites.delete(newId);
-          }
-        };
-        void pendingWrite.then(() => {
+    const pendingWrite = this.historyWrites.move(
+      oldId,
+      collidesWithExisting ? undefined : newId,
+    );
+    if (pendingWrite && !collidesWithExisting) {
+      void pendingWrite.then(
+        () => {
           this.persistedHistoryIds.delete(oldId);
           this.persistedHistoryIds.add(newId);
-          settle();
-        }, settle);
-      }
+        },
+        () => {},
+      );
     }
 
     this.publishUpdate();
@@ -1793,22 +2512,113 @@ export class AgUiThreadRuntimeCore {
   private updateAssistantMessage(
     messageId: string,
     update: ChatModelRunResult,
+    replay?: ReplayContext,
+    pendingToolCallIds?: ReadonlySet<string>,
   ): string {
     let latestStatus: MessageStatus | undefined;
     const touched = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
-      const metadata = update.metadata
+      const mergedMetadata = update.metadata
         ? this.mergeAssistantMetadata(assistant.metadata, update.metadata)
         : assistant.metadata;
       latestStatus = update.status ?? assistant.status;
-      const content =
+      let content =
         update.content !== undefined
-          ? this.preserveToolResults(
-              assistant.content,
-              update.content as ThreadAssistantMessage["content"],
-            )
+          ? 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"
+          ? admission
+          : replay && latestStatus?.type === "running"
+            ? mergedMetadata
+            : withoutPendingToolCallIds(admission);
       return {
         ...assistant,
         content,
@@ -1831,7 +2641,10 @@ export class AgUiThreadRuntimeCore {
     if (this.isPersistableStatus(latestStatus)) {
       this.persistAssistantHistory(resolvedMessageId);
     }
-    this.maybeResumeAfterToolResults(resolvedMessageId);
+    this.maybeResumeAfterToolResults(
+      resolvedMessageId,
+      replay?.accountedResultIds,
+    );
     return resolvedMessageId;
   }
 
@@ -1842,6 +2655,7 @@ export class AgUiThreadRuntimeCore {
   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)) {
@@ -1849,25 +2663,27 @@ export class AgUiThreadRuntimeCore {
         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);
   }
@@ -2054,24 +2870,30 @@ export class AgUiThreadRuntimeCore {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    reconcileSnapshot?: SnapshotMessageReconciler,
   ): string | undefined {
     switch (event.type) {
       case "STATE_SNAPSHOT": {
-        this.stateSnapshot = event.snapshot as ReadonlyJSONValue;
+        this.dispatch({
+          type: "stateSnapshot",
+          snapshot: event.snapshot as ReadonlyJSONValue,
+        });
         this.updateActiveAssistantState(activeAssistantId);
         break;
       }
       case "STATE_DELTA": {
         if (event.delta.length === 0) break;
         try {
-          const state = this.stateSnapshot ?? {};
           const result = jsonpatch.applyPatch(
-            state,
+            this.state.stateSnapshot ?? {},
             event.delta as Operation[],
             /* validateOperation */ true,
             /* mutateDocument */ false,
           );
-          this.stateSnapshot = result.newDocument as ReadonlyJSONValue;
+          this.dispatch({
+            type: "stateSnapshot",
+            snapshot: result.newDocument as ReadonlyJSONValue,
+          });
           this.updateActiveAssistantState(activeAssistantId);
         } catch (error) {
           this.logger.error?.("[agui] failed to apply state delta", error);
@@ -2093,6 +2915,7 @@ export class AgUiThreadRuntimeCore {
           joinStrategy,
           runStartMessageIds,
           hasRunUserBoundary,
+          reconcileSnapshot,
         );
         const headId = this.session.headId;
         if (hasActiveText && headId !== previousHeadId) {
@@ -2156,7 +2979,7 @@ export class AgUiThreadRuntimeCore {
           ...assistant,
           metadata: {
             ...assistant.metadata,
-            unstable_state: this.stateSnapshot ?? null,
+            unstable_state: this.state.stateSnapshot ?? null,
           },
         };
       });
@@ -2172,6 +2995,7 @@ export class AgUiThreadRuntimeCore {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       const mcpAppUri = readMcpAppResourceUri(event.mcpResult?._meta);
+      const modelText = toolContentText(event.content);
       let matchedToolCall = false;
       const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
         if (part.toolCallId !== event.toolCallId) return part;
@@ -2188,11 +3012,9 @@ export class AgUiThreadRuntimeCore {
         if (snapshotResultApplied) {
           return {
             ...part,
-            ...(part.modelContent === undefined && event.content
+            ...(part.modelContent === undefined
               ? {
-                  modelContent: [
-                    { type: "text" as const, text: event.content },
-                  ],
+                  modelContent: [{ type: "text" as const, text: modelText }],
                 }
               : {}),
             ...(part.isError === undefined
@@ -2207,12 +3029,18 @@ export class AgUiThreadRuntimeCore {
               : {}),
           };
         }
+        const { modelContent: _modelContent, ...partWithoutModelContent } =
+          part;
         return {
-          ...part,
+          ...(event.mcpResult === undefined && Array.isArray(event.content)
+            ? partWithoutModelContent
+            : part),
           result: (event.mcpResult ??
-            tryParseJSON(event.content ?? "")) as ReadonlyJSONValue,
+            (Array.isArray(event.content)
+              ? event.content
+              : tryParseJSON(event.content))) as ReadonlyJSONValue,
           ...(event.mcpResult !== undefined
-            ? { modelContent: [{ type: "text" as const, text: event.content }] }
+            ? { modelContent: [{ type: "text" as const, text: modelText }] }
             : {}),
           ...(typeof event.mcpResult?.isError === "boolean"
             ? { isError: event.mcpResult.isError }
@@ -2305,6 +3133,7 @@ export class AgUiThreadRuntimeCore {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    reconcileSnapshot?: SnapshotMessageReconciler,
   ) {
     try {
       const activeAssistantItem = activeAssistantId
@@ -2344,19 +3173,22 @@ export class AgUiThreadRuntimeCore {
         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;
+          converted[index] = previous
+            ? {
+                ...snapshot,
+                content: this.mergeSnapshotAssistantContent(
+                  previous.content,
+                  snapshot.content,
+                  snapshotHasReasoning,
+                  snapshotHasActivity,
+                ),
+              }
+            : snapshot;
         }
       }
       if (joinStrategy === "concat-content") {
@@ -2531,7 +3363,7 @@ export class AgUiThreadRuntimeCore {
       (wasPersisted || this.snapshotHistoryIds.has(messageId));
 
     if (shouldUpdate) {
-      const write = this.chainHistoryWrite(messageId, () =>
+      const write = this.historyWrites.chain(messageId, () =>
         update.call(history, { parentId, message }),
       );
       this.assistantHistoryParents.delete(messageId);
@@ -2540,7 +3372,7 @@ export class AgUiThreadRuntimeCore {
           this.persistedHistoryIds.add(messageId);
         },
         (error) => {
-          const pending = this.historyWrites.get(messageId);
+          const pending = this.historyWrites.pending(messageId);
           if (pending === undefined || pending === write) {
             this.assistantHistoryParents.set(messageId, parentId);
           }
@@ -2561,7 +3393,7 @@ export class AgUiThreadRuntimeCore {
     void write.then(
       () => {},
       (error) => {
-        const pending = this.historyWrites.get(messageId);
+        const pending = this.historyWrites.pending(messageId);
         if (pending === undefined || pending === write) {
           this.assistantHistoryParents.set(messageId, parentId);
         }
@@ -2575,11 +3407,11 @@ export class AgUiThreadRuntimeCore {
     message: ThreadMessage,
   ): Promise<void> | undefined {
     if (!this.history || this.persistedHistoryIds.has(message.id)) return;
-    const pending = this.historyWrites.get(message.id);
+    const pending = this.historyWrites.pending(message.id);
     if (pending) return pending;
 
     const append = this.history.append.bind(this.history);
-    const write = this.chainHistoryWrite(message.id, () =>
+    const write = this.historyWrites.chain(message.id, () =>
       append({ parentId, message }),
     );
     void write.then(
@@ -2590,35 +3422,4 @@ export class AgUiThreadRuntimeCore {
     );
     return write;
   }
-
-  private chainHistoryWrite(
-    id: string,
-    write: () => Promise<void>,
-  ): Promise<void> {
-    const pending = this.historyWrites.get(id);
-    let next: Promise<void>;
-    if (pending) {
-      next = pending.then(write, write);
-    } else {
-      try {
-        next = Promise.resolve(write());
-      } catch (error) {
-        next = Promise.reject(error);
-      }
-    }
-    this.historyWrites.set(id, next);
-    void next.then(
-      () => {
-        if (this.historyWrites.get(id) === next) {
-          this.historyWrites.delete(id);
-        }
-      },
-      () => {
-        if (this.historyWrites.get(id) === next) {
-          this.historyWrites.delete(id);
-        }
-      },
-    );
-    return next;
-  }
 }