@assistant-ui/react-ag-ui 0.0.54 → 0.0.57

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -14,12 +14,16 @@ import type {
   ChatModelRunResult,
   ExportedMessageRepository,
   MessageStatus,
+  RespondToToolApprovalOptions,
   ThreadAssistantMessage,
   ThreadHistoryAdapter,
   ThreadMessage,
   ToolCallMessagePart,
 } from "@assistant-ui/core";
-import { MessageRepository } from "@assistant-ui/core/internal";
+import {
+  createMessageRepositorySession,
+  invokeUserCallback,
+} from "@assistant-ui/core/internal";
 import type {
   AbstractAgent,
   AgentSubscriber,
@@ -45,6 +49,12 @@ import {
   toAgUiTools,
 } from "./adapter/conversions";
 import { createAgUiSubscriber } from "./adapter/subscriber";
+import {
+  buildToolApprovalResume,
+  projectAgUiToolApprovals,
+  withSettledToolApprovals,
+  withToolApprovalDecision,
+} from "./adapter/tool-approval";
 
 // AbstractAgent.runAgent declares two parameters. HttpAgent ignores a third and
 // is cancelled through agent.abortRun(); the run options stay for subclasses
@@ -92,36 +102,17 @@ const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
 
 type AgUiRuntimeCallbackName = "onError" | "onCancel";
 
-const reportCallbackError = (name: AgUiRuntimeCallbackName, error: unknown) => {
-  console.error(`[react-ag-ui] ${name} callback threw an error`, error);
-};
-
-const invokeRuntimeCallback = <TArgs extends unknown[]>(
+const invokeRuntimeCallback = <TArgs extends readonly unknown[]>(
   name: AgUiRuntimeCallbackName,
-  callback: ((...args: TArgs) => void) | undefined,
+  callback: ((...args: TArgs) => unknown) | undefined,
   ...args: TArgs
-) => {
-  if (!callback) return;
-
-  try {
-    const result = callback(...args) as unknown;
-    if (
-      result !== null &&
-      (typeof result === "object" || typeof result === "function") &&
-      "then" in result &&
-      typeof result.then === "function"
-    ) {
-      void Promise.resolve(result).catch((error) => {
-        reportCallbackError(name, error);
-      });
-    }
-  } catch (error) {
-    reportCallbackError(name, error);
-  }
+): void => {
+  void invokeUserCallback("react-ag-ui", name, callback, ...args);
 };
 
-// The aggregator only ever sends the interrupts it owns, so a shallow spread at
-// the custom level would drop sibling agui state such as opaqueReasoning.
+// The aggregator sends only the agui keys it owns (interrupts and
+// opaqueReasoning, each as a full replacement), so a shallow spread at the
+// custom level would drop any sibling agui state a snapshot import attached.
 const mergeAgUiNamespace = (
   current: Record<string, unknown> | undefined,
   incoming: Record<string, unknown>,
@@ -148,10 +139,25 @@ export class AgUiThreadRuntimeCore {
   private onError: ((error: Error) => void) | undefined;
   private onCancel: (() => void) | undefined;
   private readonly notifyUpdate: () => void;
+  private readonly reportedErrors = new WeakSet<object>();
 
   private runtime: AssistantRuntime | undefined;
-  private readonly repository = new MessageRepository();
-  private exportedRepository: ExportedMessageRepository | undefined;
+  private readonly session = createMessageRepositorySession({
+    decorateExport: (exported, repository) => {
+      let parentId: string | null = null;
+      for (const message of repository.getMessages()) {
+        if (message.metadata.isOptimistic) {
+          exported.messages.push({ parentId, message });
+        }
+        parentId = message.id;
+      }
+
+      return {
+        ...exported,
+        headId: repository.headId,
+      };
+    },
+  });
   private isRunningFlag = false;
   private abortController: AbortController | null = null;
   // The agent that started the active run. updateOptions can swap this.agent
@@ -162,9 +168,12 @@ export class AgUiThreadRuntimeCore {
   private history: ThreadHistoryAdapter | undefined;
   private lastRunConfig: RunConfig | undefined;
   private readonly assistantHistoryParents = new Map<string, string | null>();
-  private readonly recordedHistoryIds = new Set<string>();
+  private readonly snapshotHistoryIds = new Set<string>();
+  private readonly persistedHistoryIds = new Set<string>();
+  private readonly historyWrites = new Map<string, Promise<void>>();
   private _isLoading = false;
   private _loadPromise: Promise<void> | undefined;
+  private _loadRequested = false;
   private pendingResumeMessageId: string | null = null;
   private pendingA2uiResume = false;
   private pendingA2uiAction: Record<string, unknown> | undefined;
@@ -188,8 +197,19 @@ export class AgUiThreadRuntimeCore {
     this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
     this.onError = options.onError;
     this.onCancel = options.onCancel;
+    const previousHistory = this.history;
     this.history = options.history;
     this.installResumeShim();
+
+    if (
+      this._loadRequested &&
+      !this._loadPromise &&
+      !previousHistory &&
+      options.history &&
+      this.session.getMessages().length === 0
+    ) {
+      void this.__internal_load();
+    }
   }
 
   attachRuntime(runtime: AssistantRuntime) {
@@ -198,100 +218,17 @@ export class AgUiThreadRuntimeCore {
 
   detachRuntime() {
     this.runtime = undefined;
+    void this.cancel().catch((error) => {
+      this.logger.error("[agui] failed to cancel run during teardown", error);
+    });
   }
 
   getMessages(): readonly ThreadMessage[] {
-    return this.repository.getMessages();
+    return this.session.getMessages();
   }
 
   getMessageRepository(): ExportedMessageRepository {
-    if (this.exportedRepository) return this.exportedRepository;
-
-    const exportedRepository = this.repository.export();
-    let parentId: string | null = null;
-    for (const message of this.repository.getMessages()) {
-      if (message.metadata.isOptimistic) {
-        exportedRepository.messages.push({ parentId, message });
-      }
-      parentId = message.id;
-    }
-
-    this.exportedRepository = {
-      ...exportedRepository,
-      headId: this.repository.headId,
-    };
-    return this.exportedRepository;
-  }
-
-  private tryGetMessage(messageId: string) {
-    try {
-      return this.repository.getMessage(messageId);
-    } catch {
-      return undefined;
-    }
-  }
-
-  private tryGetMessages(
-    messageId: string,
-  ): readonly ThreadMessage[] | undefined {
-    try {
-      return this.repository.getMessages(messageId);
-    } catch {
-      return undefined;
-    }
-  }
-
-  private hasMessage(messageId: string): boolean {
-    return this.tryGetMessage(messageId) !== undefined;
-  }
-
-  private addOrUpdateMessage(
-    parentId: string | null,
-    message: ThreadMessage,
-  ): void {
-    this.repository.addOrUpdateMessage(parentId, message);
-    this.exportedRepository = undefined;
-  }
-
-  private deleteMessage(
-    messageId: string,
-    replacementId?: string | null,
-  ): void {
-    this.repository.deleteMessage(messageId, replacementId);
-    this.exportedRepository = undefined;
-  }
-
-  private tryDeleteMessage(messageId: string): boolean {
-    if (!this.hasMessage(messageId)) return false;
-    this.deleteMessage(messageId);
-    return true;
-  }
-
-  private switchToBranch(messageId: string): void {
-    this.repository.switchToBranch(messageId);
-    this.exportedRepository = undefined;
-  }
-
-  private resetRepositoryHead(messageId: string | null): void {
-    this.repository.resetHead(messageId);
-    this.exportedRepository = undefined;
-  }
-
-  private clearRepository(): void {
-    this.repository.clear();
-    this.exportedRepository = undefined;
-  }
-
-  private updateMessage(
-    messageId: string,
-    updater: (message: ThreadMessage) => ThreadMessage,
-  ): boolean {
-    const item = this.tryGetMessage(messageId);
-    if (!item) return false;
-    const message = updater(item.message);
-    if (message === item.message) return false;
-    this.addOrUpdateMessage(item.parentId, message);
-    return true;
+    return this.session.export();
   }
 
   getState(): ReadonlyJSONValue | undefined {
@@ -307,9 +244,11 @@ export class AgUiThreadRuntimeCore {
   }
 
   __internal_load(): Promise<void> {
+    this._loadRequested = true;
     if (this._loadPromise) return this._loadPromise;
+    if (!this.history) return Promise.resolve();
 
-    const promise = this.history?.load() ?? Promise.resolve(null);
+    const promise = this.history.load();
 
     this._isLoading = true;
 
@@ -317,14 +256,21 @@ export class AgUiThreadRuntimeCore {
       .then(async (repo) => {
         if (!repo) return;
 
-        this.applyExternalMessageRepository(repo);
+        this.session.applyExternalMessageRepository(repo);
+        this.assistantHistoryParents.clear();
+        this.snapshotHistoryIds.clear();
+        this.persistedHistoryIds.clear();
+        for (const { message } of repo.messages) {
+          this.persistedHistoryIds.add(message.id);
+        }
+        this.notifyUpdate();
 
         if (repo.state !== undefined) {
           this.loadExternalState(repo.state);
         }
 
         if (repo.unstable_resume) {
-          const parentId = repo.headId ?? this.repository.headId;
+          const parentId = repo.headId ?? this.session.headId;
           const resumeStream = this.history?.resume?.bind(this.history);
           await this.startRun(
             parentId,
@@ -371,17 +317,17 @@ export class AgUiThreadRuntimeCore {
   }
 
   private appendEntry(message: AppendMessage): string {
-    if (message.sourceId) this.tryDeleteMessage(message.sourceId);
+    if (message.sourceId) this.session.tryDeleteMessage(message.sourceId);
 
     const threadMessage = this.toThreadMessage(message);
     const parentId =
       message.parentId === null
         ? null
-        : message.parentId && this.hasMessage(message.parentId)
+        : message.parentId && this.session.hasMessage(message.parentId)
           ? message.parentId
-          : this.repository.headId;
-    this.addOrUpdateMessage(parentId, threadMessage);
-    this.switchToBranch(threadMessage.id);
+          : this.session.headId;
+    this.session.addOrUpdateMessage(parentId, threadMessage);
+    this.session.switchToBranch(threadMessage.id);
     this.notifyUpdate();
     this.recordHistoryEntry(parentId, threadMessage);
     return threadMessage.id;
@@ -526,34 +472,144 @@ export class AgUiThreadRuntimeCore {
       );
     }
 
+    this.assertInterruptsAnswerable(
+      "submitInterruptResponses",
+      pending.interrupts,
+    );
+
+    await this.resumeWithResponses(
+      pending.messageId,
+      openIds.map((id) => responsesById.get(id)!),
+    );
+  }
+
+  /**
+   * Resumes with an already validated resume array. Answerability is checked
+   * once per submission: re-checking here would let a clock crossing reject a
+   * decision this runtime has already recorded, stranding the gate decided and
+   * unretryable.
+   */
+  private async resumeWithResponses(
+    messageId: string,
+    resume: AgUiResumeEntry[],
+  ): Promise<void> {
+    this.clearPendingInterrupts(messageId, resume);
+    await this.startRun(messageId, this.lastRunConfig, resume);
+  }
+
+  /**
+   * The core seams that reach `respondToToolApproval` discard the promise they
+   * receive, so a rejected decision would otherwise surface only as an
+   * unhandled rejection. A failure raised by the resumed run itself was already
+   * reported by `startRun` before it rethrew, so reporting it here again would
+   * give the consumer two notifications for one failure.
+   */
+  reportError(error: unknown): void {
+    if (this.reportedErrors.has(error as object)) {
+      this.reportedErrors.delete(error as object);
+      return;
+    }
+    invokeRuntimeCallback(
+      "onError",
+      this.onError,
+      error instanceof Error ? error : new Error(String(error)),
+    );
+  }
+
+  private assertInterruptsAnswerable(
+    method: "submitInterruptResponses" | "respondToToolApproval",
+    interrupts: readonly AgUiInterrupt[],
+  ): void {
     const now = Date.now();
-    for (const interrupt of pending.interrupts) {
+    for (const interrupt of interrupts) {
       if (!interrupt.expiresAt) continue;
       const expiry = new Date(interrupt.expiresAt).getTime();
       if (Number.isNaN(expiry)) {
         throw new Error(
-          `[agui] submitInterruptResponses: interrupt ${interrupt.id} has malformed expiresAt "${interrupt.expiresAt}"`,
+          `[agui] ${method}: interrupt ${interrupt.id} has malformed expiresAt "${interrupt.expiresAt}"`,
         );
       }
       if (expiry <= now) {
         throw new Error(
-          `[agui] submitInterruptResponses: interrupt ${interrupt.id} expired at ${interrupt.expiresAt}`,
+          `[agui] ${method}: interrupt ${interrupt.id} expired at ${interrupt.expiresAt}`,
         );
       }
     }
+    if (this.isRunningFlag) {
+      throw new Error(`[agui] ${method}: a run is already in progress`);
+    }
+  }
 
-    const resume: AgUiResumeEntry[] = openIds.map((id) =>
-      responsesById.get(id)!,
+  async respondToToolApproval(
+    options: RespondToToolApprovalOptions,
+  ): Promise<void> {
+    const pending = this.getPendingInterrupts();
+    if (!pending) {
+      throw new Error(
+        "[agui] respondToToolApproval: no pending interrupts on this thread",
+      );
+    }
+
+    // Bound against the message the gates landed on, so this check claims a
+    // batch only where the projection did.
+    const gatedMessage = this.session.tryGetMessage(pending.messageId)
+      ?.message as ThreadAssistantMessage | undefined;
+    const gated = projectAgUiToolApprovals(
+      pending.interrupts,
+      new Set(
+        (gatedMessage?.content ?? [])
+          .filter((part) => part.type === "tool-call")
+          .map((part) => part.toolCallId),
+      ),
+    );
+    const isGated = [...gated.values()].some(
+      (approval) => approval.id === options.approvalId,
     );
+    if (!isGated) {
+      throw new Error(
+        `[agui] respondToToolApproval: no pending tool-call interrupt for approval id "${options.approvalId}"`,
+      );
+    }
 
-    if (this.isRunningFlag) {
+    // The decision is recorded only once the batch is known to be answerable:
+    // a rejected submission would otherwise leave the gate decided and
+    // unretryable, because a second click reports it as already decided.
+    this.assertInterruptsAnswerable(
+      "respondToToolApproval",
+      pending.interrupts,
+    );
+
+    const recorded = this.session.updateMessage(
+      pending.messageId,
+      (message) => {
+        if (message.role !== "assistant") return message;
+        const assistant = message as ThreadAssistantMessage;
+        const content = withToolApprovalDecision(assistant.content, options);
+        if (content === assistant.content) return assistant;
+        return { ...assistant, content };
+      },
+    );
+    if (!recorded) {
       throw new Error(
-        "[agui] submitInterruptResponses: a run is already in progress",
+        `[agui] respondToToolApproval: approval "${options.approvalId}" is already decided`,
       );
     }
+    this.notifyUpdate();
+
+    const assistant = this.session.tryGetMessage(pending.messageId)?.message as
+      | ThreadAssistantMessage
+      | undefined;
+    if (!assistant) return;
+
+    // AG-UI resumes a run with one response per open interrupt, so the run
+    // stays paused until every gate in the batch has been answered.
+    const resume = buildToolApprovalResume(
+      assistant.content,
+      pending.interrupts,
+    );
+    if (!resume) return;
 
-    this.clearPendingInterrupts(pending.messageId);
-    await this.startRun(pending.messageId, this.lastRunConfig, resume);
+    await this.resumeWithResponses(pending.messageId, resume);
   }
 
   async steerAway(
@@ -595,7 +651,7 @@ export class AgUiThreadRuntimeCore {
     }
 
     const normalized = this.toAppendMessage(message);
-    this.clearPendingInterrupts(pending.messageId);
+    this.clearPendingInterrupts(pending.messageId, resume);
     const threadMessageId = this.appendEntry(normalized);
     await this.startRun(threadMessageId, normalized.runConfig, resume);
   }
@@ -655,7 +711,7 @@ export class AgUiThreadRuntimeCore {
     if (typeof message === "string") {
       return {
         createdAt: new Date(),
-        parentId: this.repository.headId,
+        parentId: this.session.headId,
         sourceId: null,
         runConfig: {},
         role: "user",
@@ -666,7 +722,7 @@ export class AgUiThreadRuntimeCore {
     }
     return {
       createdAt: message.createdAt ?? new Date(),
-      parentId: message.parentId ?? this.repository.headId,
+      parentId: message.parentId ?? this.session.headId,
       sourceId: message.sourceId ?? null,
       role: message.role ?? "user",
       content: message.content,
@@ -677,8 +733,11 @@ export class AgUiThreadRuntimeCore {
     } as AppendMessage;
   }
 
-  private clearPendingInterrupts(messageId: string): void {
-    const touched = this.updateMessage(messageId, (message) => {
+  private clearPendingInterrupts(
+    messageId: string,
+    resume: readonly AgUiResumeEntry[],
+  ): void {
+    const touched = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       if (
@@ -699,6 +758,7 @@ export class AgUiThreadRuntimeCore {
       }
       return {
         ...assistant,
+        content: withSettledToolApprovals(assistant.content, resume),
         status: { type: "complete" as const, reason: "unknown" as const },
         metadata: { ...assistant.metadata, custom: newCustom },
       };
@@ -727,7 +787,7 @@ export class AgUiThreadRuntimeCore {
   }
 
   private cancelUnresolvedToolCalls(messageId: string): void {
-    const updated = this.updateMessage(messageId, (message) => {
+    const updated = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       const content = assistant.content.map((part) => {
@@ -744,7 +804,7 @@ export class AgUiThreadRuntimeCore {
   }
 
   addToolResult(options: AddToolResultOptions): void {
-    const updated = this.updateMessage(options.messageId, (message) => {
+    const updated = this.session.updateMessage(options.messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       let matchedToolCall = false;
@@ -771,7 +831,7 @@ export class AgUiThreadRuntimeCore {
   sendA2uiAction(action: Record<string, unknown>): void {
     this.assertNoPendingInterrupts();
     this.maybeAutoCancelPendingToolCalls();
-    const parentId = this.repository.headId;
+    const parentId = this.session.headId;
     if (parentId === null) {
       this.logger.debug(
         "[agui] sendA2uiAction: no messages to resume, dropping action",
@@ -809,7 +869,7 @@ export class AgUiThreadRuntimeCore {
   }
 
   private maybeCompleteAfterToolResults(messageId: string): boolean {
-    const message = this.tryGetMessage(messageId)?.message;
+    const message = this.session.tryGetMessage(messageId)?.message;
     if (!message || message.role !== "assistant") return false;
     const assistant = message as ThreadAssistantMessage;
     if (
@@ -823,7 +883,7 @@ export class AgUiThreadRuntimeCore {
     );
     if (!allResolved) return false;
 
-    const updated = this.updateMessage(messageId, (current) =>
+    const updated = this.session.updateMessage(messageId, (current) =>
       current.role === "assistant"
         ? {
             ...(current as ThreadAssistantMessage),
@@ -853,7 +913,7 @@ export class AgUiThreadRuntimeCore {
     this.assistantHistoryParents.clear();
 
     if (messages.length === 0) {
-      this.clearRepository();
+      this.session.clear();
     } else {
       let expectedParentId: string | null = null;
       let lastAppliedId: string | null = null;
@@ -863,100 +923,36 @@ export class AgUiThreadRuntimeCore {
       for (const message of messages) {
         if (seen.has(message.id)) continue;
         seen.add(message.id);
-        const existing = this.tryGetMessage(message.id);
+        const existing = this.session.tryGetMessage(message.id);
         if (existing && existing.parentId !== expectedParentId) {
           hardReplace = true;
           break;
         }
-        this.addOrUpdateMessage(expectedParentId, message);
+        this.session.addOrUpdateMessage(expectedParentId, message);
         expectedParentId = message.id;
         lastAppliedId = message.id;
       }
 
       if (hardReplace) {
-        this.clearRepository();
+        this.session.clear();
         expectedParentId = null;
         lastAppliedId = null;
         seen.clear();
         for (const message of messages) {
           if (seen.has(message.id)) continue;
           seen.add(message.id);
-          this.addOrUpdateMessage(expectedParentId, message);
+          this.session.addOrUpdateMessage(expectedParentId, message);
           expectedParentId = message.id;
           lastAppliedId = message.id;
         }
       }
 
-      this.resetRepositoryHead(lastAppliedId);
+      this.session.resetHead(lastAppliedId);
     }
 
-    this.recordedHistoryIds.clear();
+    this.snapshotHistoryIds.clear();
     for (const { message } of this.getMessageRepository().messages) {
-      this.recordedHistoryIds.add(message.id);
-    }
-    this.notifyUpdate();
-  }
-
-  private applyExternalMessageRepository(
-    loaded: ExportedMessageRepository,
-  ): void {
-    const headId = loaded.headId ?? loaded.messages.at(-1)?.message.id ?? null;
-    const ids = new Set<string>();
-    let degenerate = false;
-    for (const { message } of loaded.messages) {
-      if (ids.has(message.id)) {
-        degenerate = true;
-        break;
-      }
-      ids.add(message.id);
-    }
-    if (headId !== null && !ids.has(headId)) degenerate = true;
-
-    if (!degenerate) {
-      this.clearRepository();
-      let pending = [...loaded.messages];
-      const importedIds = new Set<string>();
-
-      while (pending.length > 0) {
-        const unresolved: typeof pending = [];
-        let progressed = false;
-        for (const item of pending) {
-          if (item.parentId !== null && !importedIds.has(item.parentId)) {
-            unresolved.push(item);
-            continue;
-          }
-          this.addOrUpdateMessage(item.parentId, item.message);
-          importedIds.add(item.message.id);
-          progressed = true;
-        }
-        if (!progressed) {
-          degenerate = true;
-          break;
-        }
-        pending = unresolved;
-      }
-    }
-
-    if (degenerate) {
-      this.clearRepository();
-      let previousId: string | null = null;
-      for (const { message } of loaded.messages) {
-        const existing = this.tryGetMessage(message.id);
-        this.addOrUpdateMessage(
-          existing ? existing.parentId : previousId,
-          message,
-        );
-        previousId = message.id;
-      }
-      this.resetRepositoryHead(previousId);
-    } else {
-      this.resetRepositoryHead(headId);
-    }
-
-    this.assistantHistoryParents.clear();
-    this.recordedHistoryIds.clear();
-    for (const { message } of loaded.messages) {
-      this.recordedHistoryIds.add(message.id);
+      this.snapshotHistoryIds.add(message.id);
     }
     this.notifyUpdate();
   }
@@ -989,37 +985,53 @@ export class AgUiThreadRuntimeCore {
   ): Promise<void> {
     const normalizedRunConfig = runConfig ?? {};
     this.lastRunConfig = normalizedRunConfig;
-    const parent = parentId === null ? undefined : this.tryGetMessage(parentId);
+    const parent =
+      parentId === null ? undefined : this.session.tryGetMessage(parentId);
     const shouldEagerlyInsertAssistant =
       parentId !== null &&
       parent !== undefined &&
-      parentId !== this.repository.headId;
+      parentId !== this.session.headId;
     const historicalMessages = [
       ...(shouldEagerlyInsertAssistant
-        ? (this.tryGetMessages(parentId) ?? this.repository.getMessages())
-        : this.repository.getMessages()),
+        ? (this.session.tryGetMessages(parentId) ?? this.session.getMessages())
+        : this.session.getMessages()),
     ];
+    const runStartMessageIds = new Set(
+      this.getMessageRepository().messages.map(({ message }) => message.id),
+    );
 
     this.pendingError = null;
-    const assistantParentId = parent ? parentId : this.repository.headId;
+    const assistantParentId = parent ? parentId : this.session.headId;
     let assistantMessageId: 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
-    // and server-id collisions must not recreate.
+    // emit keeps both the stream and the message identity. Status-only emits,
+    // data-only content, and off-branch server-id collisions must not
+    // recreate: the snapshot already carries this turn's assistant, so
+    // resurrecting for custom-event data would leave a trailing data-only
+    // duplicate.
     let assistantCollided = false;
     const ensureAssistant = (allowRecreate = false): string => {
       const cached = assistantMessageId;
-      if (cached !== undefined && this.tryGetMessage(cached)) return cached;
+      if (cached !== undefined && this.session.tryGetMessage(cached))
+        return cached;
       if (cached !== undefined && (assistantCollided || !allowRecreate)) {
         return cached;
       }
+      const repositoryHeadId = this.session.headId;
+      const shouldUseSelectedParent =
+        cached === undefined ||
+        (shouldEagerlyInsertAssistant &&
+          (repositoryHeadId === null ||
+            runStartMessageIds.has(repositoryHeadId)));
+      // Branch runs keep their selected parent unless the snapshot advanced to
+      // a message introduced during this run.
       const parentId =
-        cached === undefined &&
+        shouldUseSelectedParent &&
         assistantParentId &&
-        this.hasMessage(assistantParentId)
+        this.session.hasMessage(assistantParentId)
           ? assistantParentId
-          : this.repository.headId;
+          : repositoryHeadId;
       const created = this.insertAssistantPlaceholder(parentId, cached);
       assistantMessageId = created;
       this.markPendingAssistantHistory(created, parentId);
@@ -1029,10 +1041,11 @@ export class AgUiThreadRuntimeCore {
     if (shouldEagerlyInsertAssistant) ensureAssistant();
 
     const applyUpdate = (update: ChatModelRunResult) => {
-      const hasContent =
-        Array.isArray(update.content) && update.content.length > 0;
+      const hasStreamContent =
+        Array.isArray(update.content) &&
+        update.content.some((part) => part.type !== "data");
       const resolved = this.updateAssistantMessage(
-        ensureAssistant(hasContent),
+        ensureAssistant(hasStreamContent),
         update,
       );
       if (resolved !== assistantMessageId) {
@@ -1040,19 +1053,51 @@ export class AgUiThreadRuntimeCore {
       }
     };
 
+    const adoptServerMessageId = (
+      serverId: string,
+      startNewMessage: boolean,
+    ) => {
+      if (startNewMessage && assistantMessageId !== undefined) {
+        applyUpdate({ status: { type: "complete", reason: "unknown" } });
+        const previousId = assistantMessageId;
+        assistantMessageId = undefined;
+        if (this.session.tryGetMessage(previousId)) {
+          const created = this.insertAssistantPlaceholder(previousId);
+          assistantMessageId = created;
+          this.markPendingAssistantHistory(created, previousId);
+        }
+      }
+      const placeholder = ensureAssistant(true);
+      if (placeholder === serverId) return;
+      const reassigned = this.reassignAssistantId(placeholder, serverId);
+      // A collision drops the placeholder before revealing the existing
+      // server message as the current head. Only messages introduced during
+      // this run can replace the placeholder; regeneration must not rewrite
+      // a previous branch when the server incorrectly reuses its id.
+      const adoptsVisibleCollision =
+        !reassigned &&
+        !runStartMessageIds.has(serverId) &&
+        this.session.headId === serverId;
+      if (reassigned || adoptsVisibleCollision) {
+        assistantMessageId = serverId;
+        if (adoptsVisibleCollision) {
+          const parentId =
+            this.session.tryGetMessage(serverId)?.parentId ?? null;
+          this.markPendingAssistantHistory(serverId, parentId);
+        }
+      } else {
+        assistantCollided = true;
+      }
+    };
+
     const aggregator = new RunAggregator({
       showThinking: this.showThinking,
       logger: this.logger,
       emit: applyUpdate,
       onServerMessageId: (serverId) => {
-        const placeholder = ensureAssistant(true);
-        if (placeholder === serverId) return;
-        if (this.reassignAssistantId(placeholder, serverId)) {
-          assistantMessageId = serverId;
-        } else {
-          assistantCollided = true;
-        }
+        adoptServerMessageId(serverId, false);
       },
+      onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true),
     });
     const dispatch = (event: AgUiEvent) =>
       this.handleEvent(aggregator, event, assistantMessageId);
@@ -1134,6 +1179,7 @@ export class AgUiThreadRuntimeCore {
 
     if (this.pendingError) {
       const err = this.pendingError;
+      this.reportedErrors.add(err);
       this.pendingError = null;
       this.pendingResumeMessageId = null;
       this.pendingA2uiResume = false;
@@ -1164,7 +1210,7 @@ export class AgUiThreadRuntimeCore {
           return;
         }
         this.maybeAutoCancelPendingToolCalls();
-        const parentId = this.repository.headId;
+        const parentId = this.session.headId;
         if (parentId !== null) {
           this.startResumeRun(parentId);
         } else {
@@ -1206,7 +1252,7 @@ export class AgUiThreadRuntimeCore {
       unstable_threadId: ctx.threadId,
       unstable_parentId: ctx.parentId,
       unstable_getMessage: () => {
-        const message = this.tryGetMessage(currentId())?.message;
+        const message = this.session.tryGetMessage(currentId())?.message;
         if (!message) {
           throw new Error(
             "[agui] resume stream requested the assistant message before it existed",
@@ -1233,7 +1279,7 @@ export class AgUiThreadRuntimeCore {
     }
 
     if (ctx.abortSignal.aborted) return;
-    const current = this.tryGetMessage(currentId())?.message;
+    const current = this.session.tryGetMessage(currentId())?.message;
     if (!current || current.status?.type === "running") {
       ctx.applyUpdate({ status: { type: "complete", reason: "unknown" } });
     }
@@ -1247,7 +1293,7 @@ export class AgUiThreadRuntimeCore {
   ) {
     const threadId = this.agent.threadId || "main";
     const messages = toAgUiMessages(
-      historyMessages ?? this.repository.getMessages(),
+      historyMessages ?? this.session.getMessages(),
     );
     const context = this.runtime?.thread.getModelContext();
     const input = {
@@ -1326,34 +1372,34 @@ export class AgUiThreadRuntimeCore {
         custom: {},
       },
     };
-    this.addOrUpdateMessage(parentId ?? this.repository.headId, assistant);
-    this.switchToBranch(id);
+    this.session.addOrUpdateMessage(parentId ?? this.session.headId, assistant);
+    this.session.switchToBranch(id);
     this.notifyUpdate();
     return id;
   }
 
   private reassignAssistantId(oldId: string, newId: string): boolean {
     if (oldId === newId) return true;
-    const oldItem = this.tryGetMessage(oldId);
+    const oldItem = this.session.tryGetMessage(oldId);
     if (!oldItem) return false;
 
-    const collidesWithExisting = this.hasMessage(newId);
+    const collidesWithExisting = this.session.hasMessage(newId);
 
     if (collidesWithExisting) {
       this.logger.debug?.(
         "[agui] reassignAssistantId: server id already present in messages or repository, dropping placeholder",
         { oldId, newId },
       );
-      this.deleteMessage(oldId, oldItem.parentId ?? null);
+      this.session.deleteMessage(oldId, oldItem.parentId ?? null);
     } else {
       const { isOptimistic: _, ...metadata } = oldItem.message.metadata;
-      this.addOrUpdateMessage(oldItem.parentId, {
+      this.session.addOrUpdateMessage(oldItem.parentId, {
         ...oldItem.message,
         id: newId,
         metadata,
       } as ThreadMessage);
-      this.switchToBranch(newId);
-      this.tryDeleteMessage(oldId);
+      this.session.switchToBranch(newId);
+      this.session.tryDeleteMessage(oldId);
     }
 
     const pendingParent = this.assistantHistoryParents.get(oldId);
@@ -1364,10 +1410,33 @@ export class AgUiThreadRuntimeCore {
       }
     }
 
-    if (this.recordedHistoryIds.has(oldId)) {
-      this.recordedHistoryIds.delete(oldId);
+    for (const ids of [this.snapshotHistoryIds, this.persistedHistoryIds]) {
+      if (ids.has(oldId)) {
+        ids.delete(oldId);
+        if (!collidesWithExisting) {
+          ids.add(newId);
+        }
+      }
+    }
+
+    // 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.recordedHistoryIds.add(newId);
+        this.historyWrites.set(newId, pendingWrite);
+        const settle = () => {
+          if (this.historyWrites.get(newId) === pendingWrite) {
+            this.historyWrites.delete(newId);
+          }
+        };
+        void pendingWrite.then(() => {
+          this.persistedHistoryIds.delete(oldId);
+          this.persistedHistoryIds.add(newId);
+          settle();
+        }, settle);
       }
     }
 
@@ -1380,7 +1449,7 @@ export class AgUiThreadRuntimeCore {
     update: ChatModelRunResult,
   ): string {
     let latestStatus: MessageStatus | undefined;
-    const touched = this.updateMessage(messageId, (message) => {
+    const touched = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       const metadata = update.metadata
@@ -1549,7 +1618,7 @@ export class AgUiThreadRuntimeCore {
     messageId: string,
     event: Extract<AgUiEvent, { type: "TOOL_CALL_RESULT" }>,
   ): void {
-    const updated = this.updateMessage(messageId, (message) => {
+    const updated = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       const mcpAppUri = readMcpAppResourceUri(event.mcpResult?._meta);
@@ -1636,7 +1705,7 @@ export class AgUiThreadRuntimeCore {
           ? serverHash
           : undefined;
     const result = event.content["result"];
-    const updated = this.updateMessage(messageId, (message) => {
+    const updated = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       let matchedToolCall = false;
@@ -1688,7 +1757,7 @@ export class AgUiThreadRuntimeCore {
   ) {
     try {
       const activeMessage = activeAssistantId
-        ? this.tryGetMessage(activeAssistantId)?.message
+        ? this.session.tryGetMessage(activeAssistantId)?.message
         : undefined;
       const activeAssistant =
         activeMessage?.role === "assistant" ? activeMessage : undefined;
@@ -1708,7 +1777,6 @@ export class AgUiThreadRuntimeCore {
           );
         }
       }
-      const snapshotHeadId = converted.at(-1)?.id ?? null;
       const snapshotContainsActiveAssistant = converted.some(
         (message) => message.id === activeAssistant?.id,
       );
@@ -1721,9 +1789,17 @@ export class AgUiThreadRuntimeCore {
         converted.push(activeAssistant);
       }
       this.applyExternalMessages(converted);
-      if (preservesActiveAssistant) {
-        this.recordedHistoryIds.delete(activeAssistant.id);
-        this.markPendingAssistantHistory(activeAssistant.id, snapshotHeadId);
+      if (activeAssistant !== undefined) {
+        const activeItem = this.session.tryGetMessage(activeAssistant.id);
+        if (activeItem) {
+          if (preservesActiveAssistant) {
+            this.snapshotHistoryIds.delete(activeAssistant.id);
+          }
+          this.markPendingAssistantHistory(
+            activeAssistant.id,
+            activeItem.parentId,
+          );
+        }
       }
     } catch (error) {
       this.logger.error?.("[agui] failed to import messages snapshot", error);
@@ -1748,7 +1824,9 @@ export class AgUiThreadRuntimeCore {
   }
 
   private recordHistoryEntry(parentId: string | null, message: ThreadMessage) {
-    this.appendHistoryItem(parentId, message);
+    void this.appendHistoryItem(parentId, message)?.catch((error) => {
+      this.logger.error?.("[agui] failed to append history entry", error);
+    });
   }
 
   private markPendingAssistantHistory(
@@ -1760,22 +1838,108 @@ export class AgUiThreadRuntimeCore {
   }
 
   private persistAssistantHistory(messageId: string) {
-    if (!this.history) return;
+    const history = this.history;
+    if (!history) return;
     const parentId = this.assistantHistoryParents.get(messageId);
     if (parentId === undefined) return;
-    const message = this.tryGetMessage(messageId)?.message;
+    const message = this.session.tryGetMessage(messageId)?.message;
     if (!message || message.role !== "assistant") return;
     if (!this.isPersistableStatus(message.status)) return;
+    const wasPersisted = this.persistedHistoryIds.has(messageId);
+    const update = history.update;
+    const shouldUpdate =
+      update !== undefined &&
+      (wasPersisted || this.snapshotHistoryIds.has(messageId));
+
+    if (shouldUpdate) {
+      const write = this.chainHistoryWrite(messageId, () =>
+        update.call(history, { parentId, message }),
+      );
+      this.assistantHistoryParents.delete(messageId);
+      void write.then(
+        () => {
+          this.persistedHistoryIds.add(messageId);
+        },
+        (error) => {
+          const pending = this.historyWrites.get(messageId);
+          if (pending === undefined || pending === write) {
+            this.assistantHistoryParents.set(messageId, parentId);
+          }
+          this.logger.error?.("[agui] failed to update history entry", error);
+        },
+      );
+      return;
+    }
+
+    if (wasPersisted) {
+      this.assistantHistoryParents.delete(messageId);
+      return;
+    }
+
+    const write = this.appendHistoryItem(parentId, message);
+    if (!write) return;
     this.assistantHistoryParents.delete(messageId);
-    this.appendHistoryItem(parentId, message);
+    void write.then(
+      () => {},
+      (error) => {
+        const pending = this.historyWrites.get(messageId);
+        if (pending === undefined || pending === write) {
+          this.assistantHistoryParents.set(messageId, parentId);
+        }
+        this.logger.error?.("[agui] failed to append history entry", error);
+      },
+    );
   }
 
-  private appendHistoryItem(parentId: string | null, message: ThreadMessage) {
-    if (!this.history || this.recordedHistoryIds.has(message.id)) return;
-    this.recordedHistoryIds.add(message.id);
-    void this.history.append({ parentId, message }).catch((error) => {
-      this.recordedHistoryIds.delete(message.id);
-      this.logger.error?.("[agui] failed to append history entry", error);
-    });
+  private appendHistoryItem(
+    parentId: string | null,
+    message: ThreadMessage,
+  ): Promise<void> | undefined {
+    if (!this.history || this.persistedHistoryIds.has(message.id)) return;
+    const pending = this.historyWrites.get(message.id);
+    if (pending) return pending;
+
+    const append = this.history.append.bind(this.history);
+    const write = this.chainHistoryWrite(message.id, () =>
+      append({ parentId, message }),
+    );
+    void write.then(
+      () => {
+        this.persistedHistoryIds.add(message.id);
+      },
+      () => {},
+    );
+    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;
   }
 }