@assistant-ui/react-ag-ui 0.0.56 → 0.0.58

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -20,7 +20,12 @@ import type {
   ThreadMessage,
   ToolCallMessagePart,
 } from "@assistant-ui/core";
-import { MessageRepository } from "@assistant-ui/core/internal";
+import {
+  createMessageRepositorySession,
+  invokeUserCallback,
+  iterateToolCallParts,
+  mapToolCallPartsDeep,
+} from "@assistant-ui/core/internal";
 import type {
   AbstractAgent,
   AgentSubscriber,
@@ -71,8 +76,6 @@ const isResolvedToolCall = (
 ): boolean =>
   part.type === "tool-call" && "result" in part && part.result !== undefined;
 
-const symbolResumeShim = Symbol("agui-resume-shim");
-
 type RunConfig = NonNullable<AppendMessage["runConfig"]>;
 type ResumeStream = (
   options: ChatModelRunOptions,
@@ -99,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>,
@@ -158,8 +142,22 @@ export class AgUiThreadRuntimeCore {
   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
@@ -175,6 +173,7 @@ export class AgUiThreadRuntimeCore {
   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,7 +187,6 @@ export class AgUiThreadRuntimeCore {
     this.onCancel = options.onCancel;
     this.history = options.history;
     this.notifyUpdate = options.notifyUpdate;
-    this.installResumeShim();
   }
 
   updateOptions(options: Omit<CoreOptions, "notifyUpdate">) {
@@ -198,8 +196,17 @@ 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) {
@@ -214,97 +221,11 @@ export class AgUiThreadRuntimeCore {
   }
 
   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 {
@@ -320,9 +241,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;
 
@@ -330,14 +253,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,
@@ -384,17 +314,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;
@@ -506,8 +436,7 @@ export class AgUiThreadRuntimeCore {
     const assistant = this.findRequiresActionAssistant("tool-calls");
     if (!assistant) return null;
     const toolCallIds: string[] = [];
-    for (const part of assistant.content) {
-      if (part.type !== "tool-call") continue;
+    for (const part of iterateToolCallParts(assistant.content)) {
       if (isResolvedToolCall(part)) continue;
       toolCallIds.push(part.toolCallId);
     }
@@ -565,9 +494,9 @@ export class AgUiThreadRuntimeCore {
   }
 
   /**
-   * 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
+   * A rejected decision reaches the caller through the approval seam, and
+   * `onError` in addition, because a consumer watching only the runtime would
+   * otherwise miss it. 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.
    */
@@ -619,15 +548,14 @@ export class AgUiThreadRuntimeCore {
 
     // Bound against the message the gates landed on, so this check claims a
     // batch only where the projection did.
-    const gatedMessage = this.tryGetMessage(pending.messageId)?.message as
-      | ThreadAssistantMessage
-      | undefined;
+    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),
+        Array.from(iterateToolCallParts(gatedMessage?.content ?? [])).map(
+          (part) => part.toolCallId,
+        ),
       ),
     );
     const isGated = [...gated.values()].some(
@@ -647,13 +575,16 @@ export class AgUiThreadRuntimeCore {
       pending.interrupts,
     );
 
-    const recorded = this.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 };
-    });
+    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] respondToToolApproval: approval "${options.approvalId}" is already decided`,
@@ -661,7 +592,7 @@ export class AgUiThreadRuntimeCore {
     }
     this.notifyUpdate();
 
-    const assistant = this.tryGetMessage(pending.messageId)?.message as
+    const assistant = this.session.tryGetMessage(pending.messageId)?.message as
       | ThreadAssistantMessage
       | undefined;
     if (!assistant) return;
@@ -776,7 +707,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",
@@ -787,7 +718,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,
@@ -802,7 +733,7 @@ export class AgUiThreadRuntimeCore {
     messageId: string,
     resume: readonly AgUiResumeEntry[],
   ): void {
-    const touched = this.updateMessage(messageId, (message) => {
+    const touched = this.session.updateMessage(messageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       if (
@@ -839,9 +770,8 @@ export class AgUiThreadRuntimeCore {
     for (let index = messages.length - 1; index >= 0; index--) {
       const message = messages[index];
       if (!message || message.role !== "assistant") continue;
-      for (const part of message.content) {
-        if (part.type !== "tool-call" || part.toolCallId !== toolCallId)
-          continue;
+      for (const part of iterateToolCallParts(message.content)) {
+        if (part.toolCallId !== toolCallId) continue;
         if (!isResolvedToolCall(part)) {
           return message.id;
         }
@@ -852,11 +782,11 @@ 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) => {
-        if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
+      const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+        if (isResolvedToolCall(part)) return part;
         return {
           ...part,
           result: { error: "Tool call cancelled by user" },
@@ -869,13 +799,19 @@ export class AgUiThreadRuntimeCore {
   }
 
   addToolResult(options: AddToolResultOptions): void {
-    const updated = this.updateMessage(options.messageId, (message) => {
+    // Core's ToolInvocationTracker resolves a nested call to the nested
+    // subagent message's id, which is not a session message; re-anchor on the
+    // top-level message that owns the tree so the update can land.
+    const sessionMessageId = this.session.tryGetMessage(options.messageId)
+      ? options.messageId
+      : this.findMessageIdForToolCall(options.toolCallId);
+    if (sessionMessageId === undefined) return;
+    const updated = this.session.updateMessage(sessionMessageId, (message) => {
       if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       let matchedToolCall = false;
-      const content = assistant.content.map((part) => {
-        if (part.type !== "tool-call" || part.toolCallId !== options.toolCallId)
-          return part;
+      const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+        if (part.toolCallId !== options.toolCallId) return part;
         matchedToolCall = true;
         return {
           ...part,
@@ -890,13 +826,13 @@ export class AgUiThreadRuntimeCore {
 
     if (!updated) return;
     this.notifyUpdate();
-    this.maybeResumeAfterToolResults(options.messageId);
+    this.maybeResumeAfterToolResults(sessionMessageId);
   }
 
   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",
@@ -934,7 +870,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 (
@@ -943,12 +879,16 @@ export class AgUiThreadRuntimeCore {
     ) {
       return false;
     }
-    const allResolved = assistant.content.every(
-      (part) => part.type !== "tool-call" || isResolvedToolCall(part),
-    );
+    let allResolved = true;
+    for (const part of iterateToolCallParts(assistant.content)) {
+      if (!isResolvedToolCall(part)) {
+        allResolved = false;
+        break;
+      }
+    }
     if (!allResolved) return false;
 
-    const updated = this.updateMessage(messageId, (current) =>
+    const updated = this.session.updateMessage(messageId, (current) =>
       current.role === "assistant"
         ? {
             ...(current as ThreadAssistantMessage),
@@ -978,7 +918,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;
@@ -988,31 +928,31 @@ 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.snapshotHistoryIds.clear();
@@ -1022,71 +962,6 @@ export class AgUiThreadRuntimeCore {
     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.snapshotHistoryIds.clear();
-    this.persistedHistoryIds.clear();
-    for (const { message } of loaded.messages) {
-      this.persistedHistoryIds.add(message.id);
-    }
-    this.notifyUpdate();
-  }
-
   loadExternalState(state: ReadonlyJSONValue): void {
     this.stateSnapshot = state;
     this.notifyUpdate();
@@ -1115,22 +990,23 @@ 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
@@ -1142,11 +1018,12 @@ export class AgUiThreadRuntimeCore {
     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.repository.headId;
+      const repositoryHeadId = this.session.headId;
       const shouldUseSelectedParent =
         cached === undefined ||
         (shouldEagerlyInsertAssistant &&
@@ -1157,7 +1034,7 @@ export class AgUiThreadRuntimeCore {
       const parentId =
         shouldUseSelectedParent &&
         assistantParentId &&
-        this.hasMessage(assistantParentId)
+        this.session.hasMessage(assistantParentId)
           ? assistantParentId
           : repositoryHeadId;
       const created = this.insertAssistantPlaceholder(parentId, cached);
@@ -1181,32 +1058,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;
-        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.repository.headId === serverId;
-        if (reassigned || adoptsVisibleCollision) {
-          assistantMessageId = serverId;
-          if (adoptsVisibleCollision) {
-            const parentId = this.tryGetMessage(serverId)?.parentId ?? null;
-            this.markPendingAssistantHistory(serverId, parentId);
-          }
-        } else {
-          assistantCollided = true;
-        }
+        adoptServerMessageId(serverId, false);
       },
+      onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true),
     });
     const dispatch = (event: AgUiEvent) =>
       this.handleEvent(aggregator, event, assistantMessageId);
@@ -1319,7 +1215,7 @@ export class AgUiThreadRuntimeCore {
           return;
         }
         this.maybeAutoCancelPendingToolCalls();
-        const parentId = this.repository.headId;
+        const parentId = this.session.headId;
         if (parentId !== null) {
           this.startResumeRun(parentId);
         } else {
@@ -1361,7 +1257,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",
@@ -1388,7 +1284,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" } });
     }
@@ -1402,7 +1298,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 = {
@@ -1428,27 +1324,6 @@ export class AgUiThreadRuntimeCore {
     return input;
   }
 
-  private installResumeShim(): void {
-    const agent = this.agent as any;
-    if (agent[symbolResumeShim]) return;
-    agent[symbolResumeShim] = true;
-    const onInstance = Object.hasOwn(agent, "prepareRunAgentInput");
-    const original = onInstance
-      ? agent.prepareRunAgentInput
-      : Object.getPrototypeOf(agent)?.prepareRunAgentInput;
-    if (typeof original !== "function") return;
-    agent.prepareRunAgentInput = function (
-      this: unknown,
-      params: { resume?: unknown } | undefined,
-    ) {
-      const input = original.call(this, params);
-      if (params?.resume !== undefined && input && typeof input === "object") {
-        return { ...(input as object), resume: params.resume };
-      }
-      return input;
-    };
-  }
-
   private setRunning(running: boolean) {
     this.isRunningFlag = running;
     this.notifyUpdate();
@@ -1481,34 +1356,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);
@@ -1558,7 +1433,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
@@ -1607,19 +1482,17 @@ export class AgUiThreadRuntimeCore {
     next: ThreadAssistantMessage["content"],
   ): ThreadAssistantMessage["content"] {
     const resolved = new Map<string, ToolCallMessagePart>();
-    for (const part of previous) {
-      if (part.type === "tool-call" && isResolvedToolCall(part)) {
+    for (const part of iterateToolCallParts(previous)) {
+      if (isResolvedToolCall(part)) {
         resolved.set(part.toolCallId, part);
       }
     }
     if (resolved.size === 0) return next;
 
-    let changed = false;
-    const merged = next.map((part) => {
-      if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
+    const { content: merged, changed } = mapToolCallPartsDeep(next, (part) => {
+      if (isResolvedToolCall(part)) return part;
       const prior = resolved.get(part.toolCallId);
       if (!prior) return part;
-      changed = true;
       return {
         ...part,
         result: prior.result,
@@ -1627,7 +1500,7 @@ export class AgUiThreadRuntimeCore {
         ...(prior.isError !== undefined ? { isError: prior.isError } : {}),
       };
     });
-    return changed ? (merged as ThreadAssistantMessage["content"]) : next;
+    return changed ? merged : next;
   }
 
   private mergeAssistantMetadata(
@@ -1727,14 +1600,13 @@ 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);
       let matchedToolCall = false;
-      const content = assistant.content.map((part) => {
-        if (part.type !== "tool-call" || part.toolCallId !== event.toolCallId)
-          return part;
+      const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+        if (part.toolCallId !== event.toolCallId) return part;
         matchedToolCall = true;
         // An applied activity snapshot owns part.result; a later result only
         // fills what is missing, mirroring the aggregator's finishToolCall.
@@ -1814,13 +1686,12 @@ 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;
-      const content = assistant.content.map((part) => {
-        if (part.type !== "tool-call" || part.toolCallId !== toolCallId)
-          return part;
+      const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+        if (part.toolCallId !== toolCallId) return part;
         matchedToolCall = true;
         return {
           ...part,
@@ -1866,7 +1737,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;
@@ -1899,7 +1770,7 @@ export class AgUiThreadRuntimeCore {
       }
       this.applyExternalMessages(converted);
       if (activeAssistant !== undefined) {
-        const activeItem = this.tryGetMessage(activeAssistant.id);
+        const activeItem = this.session.tryGetMessage(activeAssistant.id);
         if (activeItem) {
           if (preservesActiveAssistant) {
             this.snapshotHistoryIds.delete(activeAssistant.id);
@@ -1951,7 +1822,7 @@ export class AgUiThreadRuntimeCore {
     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);