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

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,
@@ -85,6 +92,39 @@ const isResolvedToolCall = (
 ): boolean =>
   part.type === "tool-call" && "result" in part && part.result !== undefined;
 
+const listedPendingToolCallIds = (
+  assistant: ThreadAssistantMessage,
+): ReadonlySet<string> | undefined => {
+  const ids = (
+    assistant.metadata.custom[AG_UI_METADATA_NAMESPACE] as
+      | AgUiCustomMetadata
+      | undefined
+  )?.pendingToolCallIds;
+  return ids?.length ? new Set(ids) : undefined;
+};
+
+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 isActivityPart = (part: ThreadAssistantMessage["content"][number]) =>
   part.type === "data" && part.name.startsWith("agui-activity/");
 
@@ -155,7 +195,7 @@ const mergeAgUiNamespace = (
   return merged;
 };
 
-export class AgUiThreadRuntimeCore {
+export class AgUiThreadController {
   private agent: AbstractAgent;
   private logger: Logger;
   private showThinking: boolean;
@@ -194,27 +234,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 +268,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 +296,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 +326,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 +361,7 @@ export class AgUiThreadRuntimeCore {
       return repo;
     });
 
-    this._isLoading = true;
+    this.dispatch({ type: "loadStarted" });
     this._historyImportPromise = promise.then(
       () => {
         this._historyImportPromise = undefined;
@@ -343,7 +378,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 +389,9 @@ export class AgUiThreadRuntimeCore {
         this.reportError(error);
       })
       .finally(() => {
-        this._isLoading = false;
+        this.dispatch({
+          type: "loadFinished",
+        });
         this.publishUpdate();
       });
 
@@ -365,10 +402,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 +413,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 +434,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 +471,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 +508,9 @@ export class AgUiThreadRuntimeCore {
   }
 
   async cancel(): Promise<void> {
-    this.historyMutationGeneration++;
+    this.dispatch({
+      type: "historyMutationAdvanced",
+    });
     this.abortActiveRun();
   }
 
@@ -518,7 +557,7 @@ export class AgUiThreadRuntimeCore {
     this.assertNoPendingInterrupts();
     await this.startRun(
       config.parentId,
-      config.runConfig ?? this.lastRunConfig,
+      config.runConfig ?? this.state.lastRunConfig,
       undefined,
       config.stream,
     );
@@ -541,7 +580,7 @@ export class AgUiThreadRuntimeCore {
     try {
       await this.startRun(
         parentId,
-        this.lastRunConfig,
+        this.state.lastRunConfig,
         undefined,
         resumeStream,
       );
@@ -573,6 +612,25 @@ export class AgUiThreadRuntimeCore {
     return assistant;
   }
 
+  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 turn;
+  }
+
+  private getPendingToolCallMessages(): ThreadAssistantMessage[] {
+    return this.getTrailingAssistantTurn().filter(
+      (message) =>
+        message.status?.type === "requires-action" &&
+        message.status.reason === "tool-calls",
+    );
+  }
+
   getPendingInterrupts(): {
     messageId: string;
     interrupts: readonly AgUiInterrupt[];
@@ -592,15 +650,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 +696,7 @@ export class AgUiThreadRuntimeCore {
     this.assertInterruptsAnswerable(
       "submitInterruptResponses",
       pending.interrupts,
+      responsesById,
     );
 
     await this.resumeWithResponses(
@@ -651,7 +716,7 @@ export class AgUiThreadRuntimeCore {
     resume: AgUiResumeEntry[],
   ): Promise<void> {
     this.clearPendingInterrupts(messageId, resume);
-    await this.startRun(messageId, this.lastRunConfig, {
+    await this.startRun(messageId, this.state.lastRunConfig, {
       entries: resume,
       interruptedMessageId: messageId,
     });
@@ -679,6 +744,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 +755,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 +858,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 +881,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 +1138,7 @@ export class AgUiThreadRuntimeCore {
       return;
     }
 
-    await this.chainHistoryWrite(sessionMessageId, () =>
+    await this.historyWrites.chain(sessionMessageId, () =>
       history.update!({ parentId: item.parentId, message }),
     );
   }
@@ -1090,11 +1161,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);
@@ -1105,15 +1182,25 @@ export class AgUiThreadRuntimeCore {
   // run is still draining) or after it.
   private maybeResumeAfterToolResults(messageId: string): void {
     if (!this.maybeCompleteAfterToolResults(messageId)) 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
       // closed). Defer until startRun's tail so we never start two runs.
-      this.pendingResume = { owner, messageId };
+      this.dispatch({
+        type: "resumeDeferred",
+        owner,
+        messageId: headId,
+      });
       return;
     }
-    this.startResumeRun(messageId);
+    this.startResumeRun(headId);
   }
 
   private maybeCompleteAfterToolResults(messageId: string): boolean {
@@ -1127,7 +1214,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 +1229,7 @@ export class AgUiThreadRuntimeCore {
         ? {
             ...(current as ThreadAssistantMessage),
             status: { type: "complete" as const, reason: "unknown" as const },
+            metadata: withoutPendingToolCallIds(current.metadata),
           }
         : current,
     );
@@ -1150,7 +1240,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 +1257,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) {
@@ -1217,18 +1308,21 @@ export class AgUiThreadRuntimeCore {
     // 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;
+    const resumeTarget = this.state.pendingResume?.messageId;
     if (
       resumeTarget !== undefined &&
       !this.session.getMessages().some((message) => message.id === resumeTarget)
     ) {
-      this.pendingResume = null;
+      this.dispatch({ type: "resumeCleared" });
     }
     this.publishUpdate();
   }
 
   loadExternalState(state: ReadonlyJSONValue): void {
-    this.stateSnapshot = state;
+    this.dispatch({
+      type: "stateSnapshot",
+      snapshot: state,
+    });
     this.publishUpdate();
   }
 
@@ -1237,25 +1331,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;
@@ -1286,7 +1383,10 @@ export class AgUiThreadRuntimeCore {
     if (!this.supersedeActiveRun()) return;
 
     const normalizedRunConfig = runConfig ?? {};
-    this.lastRunConfig = normalizedRunConfig;
+    this.dispatch({
+      type: "runConfigRecorded",
+      runConfig: normalizedRunConfig,
+    });
     const parent =
       parentId === null ? undefined : this.session.tryGetMessage(parentId);
     const shouldEagerlyInsertAssistant =
@@ -1315,6 +1415,9 @@ 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 restrictedPendingToolCallIds: ReadonlySet<string> | undefined;
+    let settlingRun = false;
     const ensureAssistant = (allowRecreate = false): string => {
       const cached = assistantMessageId;
       if (cached !== undefined && this.session.tryGetMessage(cached))
@@ -1344,7 +1447,66 @@ export class AgUiThreadRuntimeCore {
 
     if (shouldEagerlyInsertAssistant) ensureAssistant();
 
+    const settleHeldMessages = (terminalStatus: MessageStatus): string[] => {
+      const pendingMessages: string[] = [];
+      for (const [messageId, openToolCallIds] of heldMessages) {
+        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 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
+                ? {
+                    ...metadata,
+                    custom: {
+                      ...metadata.custom,
+                      [AG_UI_METADATA_NAMESPACE]: {
+                        ...(metadata.custom[AG_UI_METADATA_NAMESPACE] as
+                          | AgUiCustomMetadata
+                          | undefined),
+                        pendingToolCallIds,
+                      },
+                    },
+                  }
+                : metadata,
+          };
+        });
+        if (updated) {
+          this.publishUpdate();
+          this.persistAssistantHistory(messageId);
+          if (requiresToolResult) {
+            pendingMessages.push(messageId);
+          }
+        }
+      }
+      heldMessages.clear();
+      return pendingMessages;
+    };
+
     const applyUpdate = (update: ChatModelRunResult) => {
+      const terminalStatus = update.status;
+      const pendingMessages =
+        settlingRun && terminalStatus && terminalStatus.type !== "running"
+          ? settleHeldMessages(terminalStatus)
+          : [];
       const hasStreamContent =
         Array.isArray(update.content) &&
         update.content.some((part) => part.type !== "data");
@@ -1355,14 +1517,22 @@ export class AgUiThreadRuntimeCore {
       if (resolved !== assistantMessageId) {
         assistantMessageId = resolved;
       }
+      for (const messageId of pendingMessages) {
+        this.maybeResumeAfterToolResults(messageId);
+      }
     };
 
     const adoptServerMessageId = (
       serverId: string,
       startNewMessage: boolean,
+      openToolCallIds: string[] = [],
     ) => {
       if (startNewMessage && assistantMessageId !== undefined) {
-        applyUpdate({ status: { type: "complete", reason: "unknown" } });
+        if (openToolCallIds.length > 0) {
+          heldMessages.set(assistantMessageId, new Set(openToolCallIds));
+        } else {
+          applyUpdate({ status: { type: "complete", reason: "unknown" } });
+        }
         const previousId = assistantMessageId;
         assistantMessageId = undefined;
         if (this.session.tryGetMessage(previousId)) {
@@ -1404,22 +1574,43 @@ export class AgUiThreadRuntimeCore {
       onServerMessageId: (serverId) => {
         adoptServerMessageId(serverId, false);
       },
-      onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true),
+      onTextMessageStart: (serverId, openToolCallIds) =>
+        adoptServerMessageId(serverId, true, openToolCallIds),
     });
     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;
+        const ids =
+          event.outcome?.type === "success"
+            ? event.outcome.pendingToolCallIds
+            : undefined;
+        restrictedPendingToolCallIds = ids?.length ? new Set(ids) : undefined;
+      }
+      if (event.type === "TOOL_CALL_RESULT") {
+        for (const openToolCallIds of heldMessages.values()) {
+          openToolCallIds.delete(event.toolCallId);
+        }
+      }
+      settlingRun =
+        event.type === "RUN_FINISHED" ||
+        event.type === "RUN_ERROR" ||
+        event.type === "RUN_CANCELLED";
+      try {
+        const nextAssistantMessageId = this.handleEvent(
+          aggregator,
+          event,
+          assistantMessageId,
+          joinStrategy,
+          runStartMessageIds,
+          hasRunUserBoundary,
+        );
+        if (event.type === "MESSAGES_SNAPSHOT") {
+          assistantMessageId = nextAssistantMessageId;
+        }
+      } finally {
+        settlingRun = false;
       }
     };
 
@@ -1482,9 +1673,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
         }
@@ -1514,19 +1705,23 @@ export class AgUiThreadRuntimeCore {
 
     // 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 (this.state.pendingResume?.owner === abortController) {
+      const { messageId } = this.state.pendingResume;
+      this.dispatch({ type: "resumeCleared" });
       if (!abortSignal.aborted) {
         this.startResumeRun(messageId);
       }
     }
 
-    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 +1732,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 +1761,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 +1814,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 +1837,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 +1880,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 +1915,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 +1973,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();
@@ -1798,10 +1999,15 @@ export class AgUiThreadRuntimeCore {
     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 metadata =
+        latestStatus?.type === "requires-action" &&
+        latestStatus.reason === "tool-calls"
+          ? mergedMetadata
+          : withoutPendingToolCallIds(mergedMetadata);
       const content =
         update.content !== undefined
           ? this.preserveToolResults(
@@ -2057,21 +2263,26 @@ export class AgUiThreadRuntimeCore {
   ): 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);
@@ -2156,7 +2367,7 @@ export class AgUiThreadRuntimeCore {
           ...assistant,
           metadata: {
             ...assistant.metadata,
-            unstable_state: this.stateSnapshot ?? null,
+            unstable_state: this.state.stateSnapshot ?? null,
           },
         };
       });
@@ -2172,6 +2383,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 +2400,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 +2417,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 }
@@ -2531,7 +2747,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 +2756,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 +2777,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 +2791,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 +2806,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;
-  }
 }