@assistant-ui/react-ag-ui 0.0.43 → 0.0.45

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -8,12 +8,14 @@ import type {
   AssistantRuntime,
   ChatModelRunOptions,
   ChatModelRunResult,
+  ExportedMessageRepository,
   MessageStatus,
   ThreadAssistantMessage,
   ThreadHistoryAdapter,
   ThreadMessage,
   ToolCallMessagePart,
 } from "@assistant-ui/core";
+import { MessageRepository } from "@assistant-ui/core/internal";
 import type { AbstractAgent } from "@ag-ui/client";
 import jsonpatch, { type Operation } from "fast-json-patch";
 import type { Logger } from "./logger";
@@ -36,6 +38,11 @@ const optimisticPrefix = "__optimistic__";
 const generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
 const isOptimisticId = (id: string) => id.startsWith(optimisticPrefix);
 
+const isResolvedToolCall = (
+  part: ThreadAssistantMessage["content"][number],
+): boolean =>
+  part.type === "tool-call" && "result" in part && part.result !== undefined;
+
 const symbolResumeShim = Symbol("agui-resume-shim");
 
 type RunConfig = NonNullable<AppendMessage["runConfig"]>;
@@ -53,6 +60,7 @@ type CoreOptions = {
   agent: AbstractAgent;
   logger: Logger;
   showThinking: boolean;
+  autoCancelPendingToolCalls?: boolean | undefined;
   onError?: (error: Error) => void;
   onCancel?: () => void;
   history?: ThreadHistoryAdapter;
@@ -65,12 +73,14 @@ export class AgUiThreadRuntimeCore {
   private agent: AbstractAgent;
   private logger: Logger;
   private showThinking: boolean;
+  private autoCancelPendingToolCalls: boolean | undefined;
   private onError: ((error: Error) => void) | undefined;
   private onCancel: (() => void) | undefined;
   private readonly notifyUpdate: () => void;
 
   private runtime: AssistantRuntime | undefined;
-  private messages: ThreadMessage[] = [];
+  private readonly repository = new MessageRepository();
+  private exportedRepository: ExportedMessageRepository | undefined;
   private isRunningFlag = false;
   private abortController: AbortController | null = null;
   private stateSnapshot: ReadonlyJSONValue | undefined;
@@ -87,6 +97,7 @@ export class AgUiThreadRuntimeCore {
     this.agent = options.agent;
     this.logger = options.logger;
     this.showThinking = options.showThinking;
+    this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
     this.onError = options.onError;
     this.onCancel = options.onCancel;
     this.history = options.history;
@@ -98,6 +109,7 @@ export class AgUiThreadRuntimeCore {
     this.agent = options.agent;
     this.logger = options.logger;
     this.showThinking = options.showThinking;
+    this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
     this.onError = options.onError;
     this.onCancel = options.onCancel;
     this.history = options.history;
@@ -113,7 +125,97 @@ export class AgUiThreadRuntimeCore {
   }
 
   getMessages(): readonly ThreadMessage[] {
-    return this.messages;
+    return this.repository.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;
   }
 
   getState(): ReadonlyJSONValue | undefined {
@@ -139,15 +241,14 @@ export class AgUiThreadRuntimeCore {
       .then(async (repo) => {
         if (!repo) return;
 
-        const messages = repo.messages.map((item) => item.message);
-        this.applyExternalMessages(messages);
+        this.applyExternalMessageRepository(repo);
 
         if (repo.state !== undefined) {
           this.loadExternalState(repo.state);
         }
 
         if (repo.unstable_resume) {
-          const parentId = repo.headId ?? messages.at(-1)?.id ?? null;
+          const parentId = repo.headId ?? this.repository.headId;
           const resumeStream = this.history?.resume?.bind(this.history);
           await this.startRun(
             parentId,
@@ -174,24 +275,37 @@ export class AgUiThreadRuntimeCore {
 
   async append(message: AppendMessage): Promise<void> {
     const startRun = message.startRun ?? message.role === "user";
-    if (startRun) this.assertNoPendingInterrupts();
+    if (startRun) {
+      this.assertNoPendingInterrupts();
+      this.maybeAutoCancelPendingToolCalls();
+    }
     const threadMessageId = this.appendEntry(message);
     if (!startRun) return;
     await this.startRun(threadMessageId, message.runConfig);
   }
 
+  private maybeAutoCancelPendingToolCalls(): void {
+    if (this.autoCancelPendingToolCalls === false) return;
+    const pending = this.getPendingToolCalls();
+    if (!pending) return;
+    this.cancelUnresolvedToolCalls(pending.messageId);
+    this.maybeCompleteAfterToolResults(pending.messageId);
+  }
+
   private appendEntry(message: AppendMessage): string {
-    if (message.sourceId) {
-      this.messages = this.messages.filter(
-        (entry) => entry.id !== message.sourceId,
-      );
-    }
-    this.resetHead(message.parentId);
+    if (message.sourceId) this.tryDeleteMessage(message.sourceId);
 
     const threadMessage = this.toThreadMessage(message);
-    this.messages = [...this.messages, threadMessage];
+    const parentId =
+      message.parentId === null
+        ? null
+        : message.parentId && this.hasMessage(message.parentId)
+          ? message.parentId
+          : this.repository.headId;
+    this.addOrUpdateMessage(parentId, threadMessage);
+    this.switchToBranch(threadMessage.id);
     this.notifyUpdate();
-    this.recordHistoryEntry(message.parentId ?? null, threadMessage);
+    this.recordHistoryEntry(parentId, threadMessage);
     return threadMessage.id;
   }
 
@@ -204,8 +318,7 @@ export class AgUiThreadRuntimeCore {
     config: { runConfig?: RunConfig } = {},
   ): Promise<void> {
     this.assertNoPendingInterrupts();
-    this.resetHead(parentId);
-    this.notifyUpdate();
+    this.maybeAutoCancelPendingToolCalls();
     await this.startRun(parentId, config.runConfig);
   }
 
@@ -255,20 +368,28 @@ export class AgUiThreadRuntimeCore {
     );
   }
 
-  getPendingInterrupts(): {
-    messageId: string;
-    interrupts: readonly AgUiInterrupt[];
-  } | null {
-    const assistant = this.messages.findLast((m) => m.role === "assistant") as
-      | ThreadAssistantMessage
-      | undefined;
+  private findRequiresActionAssistant(
+    reason: "interrupt" | "tool-calls",
+  ): ThreadAssistantMessage | null {
+    const assistant = this.getMessages().findLast(
+      (message) => message.role === "assistant",
+    ) as ThreadAssistantMessage | undefined;
     if (
       !assistant ||
       assistant.status?.type !== "requires-action" ||
-      assistant.status.reason !== "interrupt"
+      assistant.status.reason !== reason
     ) {
       return null;
     }
+    return assistant;
+  }
+
+  getPendingInterrupts(): {
+    messageId: string;
+    interrupts: readonly AgUiInterrupt[];
+  } | null {
+    const assistant = this.findRequiresActionAssistant("interrupt");
+    if (!assistant) return null;
     const stored = (
       assistant.metadata.custom[AG_UI_METADATA_NAMESPACE] as
         | AgUiCustomMetadata
@@ -278,6 +399,22 @@ export class AgUiThreadRuntimeCore {
     return { messageId: assistant.id, interrupts: stored };
   }
 
+  getPendingToolCalls(): {
+    messageId: string;
+    toolCallIds: string[];
+  } | null {
+    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;
+      if (isResolvedToolCall(part)) continue;
+      toolCallIds.push(part.toolCallId);
+    }
+    if (toolCallIds.length === 0) return null;
+    return { messageId: assistant.id, toolCallIds };
+  }
+
   async submitInterruptResponses(
     responses: readonly AgUiResumeEntry[],
   ): Promise<void> {
@@ -339,8 +476,8 @@ export class AgUiThreadRuntimeCore {
       }
     }
 
-    const resume: AgUiResumeEntry[] = openIds.map(
-      (id) => responsesById.get(id)!,
+    const resume: AgUiResumeEntry[] = openIds.map((id) =>
+      responsesById.get(id)!,
     );
 
     if (this.isRunningFlag) {
@@ -359,6 +496,23 @@ export class AgUiThreadRuntimeCore {
   ): Promise<void> {
     const pending = this.getPendingInterrupts();
     if (!pending) {
+      const pendingTools = this.getPendingToolCalls();
+      if (pendingTools) {
+        if (responses?.length) {
+          throw new Error(
+            "[agui] steerAway: responses are only valid for pending interrupts",
+          );
+        }
+        if (this.isRunningFlag) {
+          throw new Error("[agui] steerAway: a run is already in progress");
+        }
+        this.cancelUnresolvedToolCalls(pendingTools.messageId);
+        this.maybeCompleteAfterToolResults(pendingTools.messageId);
+        const normalized = this.toAppendMessage(message);
+        const threadMessageId = this.appendEntry(normalized);
+        await this.startRun(threadMessageId, normalized.runConfig);
+        return;
+      }
       if (responses?.length) {
         throw new Error(
           "[agui] steerAway: no pending interrupts on this thread",
@@ -375,9 +529,6 @@ export class AgUiThreadRuntimeCore {
     }
 
     const normalized = this.toAppendMessage(message);
-    // Clear before appendEntry so the interrupted assistant is still in
-    // this.messages when its status is flipped; a parentId that points at an
-    // ancestor would otherwise truncate it away before it can be cleared.
     this.clearPendingInterrupts(pending.messageId);
     const threadMessageId = this.appendEntry(normalized);
     await this.startRun(threadMessageId, normalized.runConfig, resume);
@@ -422,7 +573,7 @@ export class AgUiThreadRuntimeCore {
     if (typeof message === "string") {
       return {
         createdAt: new Date(),
-        parentId: this.messages.at(-1)?.id ?? null,
+        parentId: this.repository.headId,
         sourceId: null,
         runConfig: {},
         role: "user",
@@ -433,7 +584,7 @@ export class AgUiThreadRuntimeCore {
     }
     return {
       createdAt: message.createdAt ?? new Date(),
-      parentId: message.parentId ?? this.messages.at(-1)?.id ?? null,
+      parentId: message.parentId ?? this.repository.headId,
       sourceId: message.sourceId ?? null,
       role: message.role ?? "user",
       content: message.content,
@@ -445,10 +596,8 @@ export class AgUiThreadRuntimeCore {
   }
 
   private clearPendingInterrupts(messageId: string): void {
-    let touched = false;
-    this.messages = this.messages.map((message) => {
-      if (message.id !== messageId || message.role !== "assistant")
-        return message;
+    const touched = this.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       if (
         assistant.status?.type !== "requires-action" ||
@@ -456,7 +605,6 @@ export class AgUiThreadRuntimeCore {
       ) {
         return assistant;
       }
-      touched = true;
       const aguiMeta = assistant.metadata.custom[AG_UI_METADATA_NAMESPACE] as
         | AgUiCustomMetadata
         | undefined;
@@ -473,18 +621,21 @@ export class AgUiThreadRuntimeCore {
         metadata: { ...assistant.metadata, custom: newCustom },
       };
     });
-    if (touched) this.notifyUpdate();
+    if (touched) {
+      this.notifyUpdate();
+    }
   }
 
   findMessageIdForToolCall(toolCallId: string): string | undefined {
     let fallbackMessageId: string | undefined;
-    for (let index = this.messages.length - 1; index >= 0; index--) {
-      const message = this.messages[index];
+    const messages = this.getMessages();
+    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;
-        if (!("result" in part) || part.result === undefined) {
+        if (!isResolvedToolCall(part)) {
           return message.id;
         }
         fallbackMessageId ??= message.id;
@@ -493,11 +644,26 @@ export class AgUiThreadRuntimeCore {
     return fallbackMessageId;
   }
 
+  private cancelUnresolvedToolCalls(messageId: string): void {
+    const updated = this.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;
+        return {
+          ...part,
+          result: { error: "Tool call cancelled by user" },
+          isError: true,
+        };
+      });
+      return { ...assistant, content };
+    });
+    if (updated) this.notifyUpdate();
+  }
+
   addToolResult(options: AddToolResultOptions): void {
-    let updated = false;
-    this.messages = this.messages.map((message) => {
-      if (message.id !== options.messageId || message.role !== "assistant")
-        return message;
+    const updated = this.updateMessage(options.messageId, (message) => {
+      if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       let matchedToolCall = false;
       const content = assistant.content.map((part) => {
@@ -512,7 +678,6 @@ export class AgUiThreadRuntimeCore {
         };
       });
       if (!matchedToolCall) return message;
-      updated = true;
       return { ...assistant, content };
     });
 
@@ -537,7 +702,7 @@ export class AgUiThreadRuntimeCore {
   }
 
   private maybeCompleteAfterToolResults(messageId: string): boolean {
-    const message = this.messages.find((m) => m.id === messageId);
+    const message = this.tryGetMessage(messageId)?.message;
     if (!message || message.role !== "assistant") return false;
     const assistant = message as ThreadAssistantMessage;
     if (
@@ -547,20 +712,19 @@ export class AgUiThreadRuntimeCore {
       return false;
     }
     const allResolved = assistant.content.every(
-      (part) =>
-        part.type !== "tool-call" ||
-        ("result" in part && part.result !== undefined),
+      (part) => part.type !== "tool-call" || isResolvedToolCall(part),
     );
     if (!allResolved) return false;
 
-    this.messages = this.messages.map((m) =>
-      m.id === messageId && m.role === "assistant"
+    const updated = this.updateMessage(messageId, (current) =>
+      current.role === "assistant"
         ? {
-            ...(m as ThreadAssistantMessage),
+            ...(current as ThreadAssistantMessage),
             status: { type: "complete" as const, reason: "unknown" as const },
           }
-        : m,
+        : current,
     );
+    if (!updated) return false;
     this.notifyUpdate();
     this.persistAssistantHistory(messageId);
     return true;
@@ -574,9 +738,111 @@ export class AgUiThreadRuntimeCore {
 
   applyExternalMessages(messages: readonly ThreadMessage[]): void {
     this.assistantHistoryParents.clear();
-    this.messages = [...messages];
+
+    if (messages.length === 0) {
+      this.clearRepository();
+    } else {
+      let expectedParentId: string | null = null;
+      let lastAppliedId: string | null = null;
+      let hardReplace = false;
+      const seen = new Set<string>();
+
+      for (const message of messages) {
+        if (seen.has(message.id)) continue;
+        seen.add(message.id);
+        const existing = this.tryGetMessage(message.id);
+        if (existing && existing.parentId !== expectedParentId) {
+          hardReplace = true;
+          break;
+        }
+        this.addOrUpdateMessage(expectedParentId, message);
+        expectedParentId = message.id;
+        lastAppliedId = message.id;
+      }
+
+      if (hardReplace) {
+        this.clearRepository();
+        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);
+          expectedParentId = message.id;
+          lastAppliedId = message.id;
+        }
+      }
+
+      this.resetRepositoryHead(lastAppliedId);
+    }
+
+    this.recordedHistoryIds.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 this.messages) {
+    for (const { message } of loaded.messages) {
       this.recordedHistoryIds.add(message.id);
     }
     this.notifyUpdate();
@@ -587,6 +853,21 @@ export class AgUiThreadRuntimeCore {
     this.notifyUpdate();
   }
 
+  setState(
+    next:
+      | ReadonlyJSONValue
+      | ((prev: ReadonlyJSONValue | undefined) => ReadonlyJSONValue),
+  ): void {
+    this.stateSnapshot =
+      typeof next === "function" ? next(this.stateSnapshot) : next;
+    this.notifyUpdate();
+  }
+
+  resetState(): void {
+    this.stateSnapshot = undefined;
+    this.notifyUpdate();
+  }
+
   private async startRun(
     parentId: string | null,
     runConfig?: RunConfig,
@@ -595,20 +876,32 @@ export class AgUiThreadRuntimeCore {
   ): Promise<void> {
     const normalizedRunConfig = runConfig ?? {};
     this.lastRunConfig = normalizedRunConfig;
-    this.resetHead(parentId);
-    const historicalMessages = [...this.messages];
+    const parent = parentId === null ? undefined : this.tryGetMessage(parentId);
+    const shouldEagerlyInsertAssistant =
+      parentId !== null &&
+      parent !== undefined &&
+      parentId !== this.repository.headId;
+    const historicalMessages = [
+      ...(shouldEagerlyInsertAssistant
+        ? (this.tryGetMessages(parentId) ?? this.repository.getMessages())
+        : this.repository.getMessages()),
+    ];
 
     this.pendingError = null;
-    const assistantParentId = parentId ?? this.messages.at(-1)?.id ?? null;
+    const assistantParentId = parent ? parentId : this.repository.headId;
     let assistantMessageId: string | undefined;
     const ensureAssistant = () => {
       if (assistantMessageId) return assistantMessageId;
-      const created = this.insertAssistantPlaceholder();
+      const created = this.insertAssistantPlaceholder(
+        assistantParentId ?? null,
+      );
       assistantMessageId = created;
       this.markPendingAssistantHistory(created, assistantParentId ?? null);
       return created;
     };
 
+    if (shouldEagerlyInsertAssistant) ensureAssistant();
+
     const applyUpdate = (update: ChatModelRunResult) => {
       const resolved = this.updateAssistantMessage(ensureAssistant(), update);
       if (resolved !== assistantMessageId) {
@@ -623,8 +916,9 @@ export class AgUiThreadRuntimeCore {
       onServerMessageId: (serverId) => {
         const placeholder = ensureAssistant();
         if (placeholder === serverId) return;
-        this.reassignAssistantId(placeholder, serverId);
-        assistantMessageId = serverId;
+        if (this.reassignAssistantId(placeholder, serverId)) {
+          assistantMessageId = serverId;
+        }
       },
     });
     const dispatch = (event: AgUiEvent) => this.handleEvent(aggregator, event);
@@ -744,7 +1038,7 @@ export class AgUiThreadRuntimeCore {
       unstable_threadId: ctx.threadId,
       unstable_parentId: ctx.parentId,
       unstable_getMessage: () => {
-        const message = this.messages.find((m) => m.id === currentId());
+        const message = this.tryGetMessage(currentId())?.message;
         if (!message) {
           throw new Error(
             "[agui] resume stream requested the assistant message before it existed",
@@ -771,7 +1065,7 @@ export class AgUiThreadRuntimeCore {
     }
 
     if (ctx.abortSignal.aborted) return;
-    const current = this.messages.find((m) => m.id === currentId());
+    const current = this.tryGetMessage(currentId())?.message;
     if (!current || current.status?.type === "running") {
       ctx.applyUpdate({ status: { type: "complete", reason: "unknown" } });
     }
@@ -784,7 +1078,9 @@ export class AgUiThreadRuntimeCore {
     resume?: AgUiResumeEntry[],
   ) {
     const threadId = this.agent.threadId || "main";
-    const messages = toAgUiMessages(historyMessages ?? this.messages);
+    const messages = toAgUiMessages(
+      historyMessages ?? this.repository.getMessages(),
+    );
     const context = this.runtime?.thread.getModelContext();
     return {
       threadId,
@@ -837,7 +1133,7 @@ export class AgUiThreadRuntimeCore {
     this.setRunning(false);
   }
 
-  private insertAssistantPlaceholder(): string {
+  private insertAssistantPlaceholder(parentId: string | null): string {
     const id = generateOptimisticId();
     const assistant: ThreadAssistantMessage = {
       id,
@@ -854,56 +1150,62 @@ export class AgUiThreadRuntimeCore {
         custom: {},
       },
     };
-    this.messages = [...this.messages, assistant];
+    this.addOrUpdateMessage(parentId ?? this.repository.headId, assistant);
+    this.switchToBranch(id);
     this.notifyUpdate();
     return id;
   }
 
-  private reassignAssistantId(oldId: string, newId: string): void {
-    if (oldId === newId) return;
+  private reassignAssistantId(oldId: string, newId: string): boolean {
+    if (oldId === newId) return true;
+    const oldItem = this.tryGetMessage(oldId);
+    if (!oldItem) return false;
 
-    const collidesWithExisting = this.messages.some((m) => m.id === newId);
+    const collidesWithExisting = this.hasMessage(newId);
 
     if (collidesWithExisting) {
       this.logger.debug?.(
-        "[agui] reassignAssistantId: server id already present in messages, dropping placeholder",
+        "[agui] reassignAssistantId: server id already present in messages or repository, dropping placeholder",
         { oldId, newId },
       );
-      this.messages = this.messages.filter((m) => m.id !== oldId);
+      this.deleteMessage(oldId, oldItem.parentId ?? null);
     } else {
-      this.messages = this.messages.map((m) => {
-        if (m.id !== oldId) return m;
-        const { isOptimistic: _, ...metadata } = m.metadata;
-        return { ...m, id: newId, metadata } as ThreadMessage;
-      });
+      const { isOptimistic: _, ...metadata } = oldItem.message.metadata;
+      this.addOrUpdateMessage(oldItem.parentId, {
+        ...oldItem.message,
+        id: newId,
+        metadata,
+      } as ThreadMessage);
+      this.switchToBranch(newId);
+      this.tryDeleteMessage(oldId);
     }
 
     const pendingParent = this.assistantHistoryParents.get(oldId);
     if (pendingParent !== undefined) {
       this.assistantHistoryParents.delete(oldId);
-      if (!this.assistantHistoryParents.has(newId)) {
+      if (!collidesWithExisting && !this.assistantHistoryParents.has(newId)) {
         this.assistantHistoryParents.set(newId, pendingParent);
       }
     }
 
     if (this.recordedHistoryIds.has(oldId)) {
       this.recordedHistoryIds.delete(oldId);
-      this.recordedHistoryIds.add(newId);
+      if (!collidesWithExisting) {
+        this.recordedHistoryIds.add(newId);
+      }
     }
 
     this.notifyUpdate();
+    return !collidesWithExisting;
   }
 
   private updateAssistantMessage(
     messageId: string,
     update: ChatModelRunResult,
   ): string {
-    let touched = false;
     let latestStatus: MessageStatus | undefined;
-    this.messages = this.messages.map((message) => {
-      if (message.id !== messageId || message.role !== "assistant")
-        return message;
-      touched = true;
+    const touched = this.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       const metadata = update.metadata
         ? this.mergeAssistantMetadata(assistant.metadata, update.metadata)
@@ -952,11 +1254,7 @@ export class AgUiThreadRuntimeCore {
   ): ThreadAssistantMessage["content"] {
     const resolved = new Map<string, ToolCallMessagePart>();
     for (const part of previous) {
-      if (
-        part.type === "tool-call" &&
-        "result" in part &&
-        part.result !== undefined
-      ) {
+      if (part.type === "tool-call" && isResolvedToolCall(part)) {
         resolved.set(part.toolCallId, part);
       }
     }
@@ -964,8 +1262,7 @@ export class AgUiThreadRuntimeCore {
 
     let changed = false;
     const merged = next.map((part) => {
-      if (part.type !== "tool-call") return part;
-      if ("result" in part && part.result !== undefined) return part;
+      if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
       const prior = resolved.get(part.toolCallId);
       if (!prior) return part;
       changed = true;
@@ -1056,10 +1353,8 @@ export class AgUiThreadRuntimeCore {
     messageId: string,
     event: Extract<AgUiEvent, { type: "TOOL_CALL_RESULT" }>,
   ): void {
-    let updated = false;
-    this.messages = this.messages.map((message) => {
-      if (message.id !== messageId || message.role !== "assistant")
-        return message;
+    const updated = this.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
       const assistant = message as ThreadAssistantMessage;
       let matchedToolCall = false;
       const content = assistant.content.map((part) => {
@@ -1076,7 +1371,6 @@ export class AgUiThreadRuntimeCore {
         };
       });
       if (!matchedToolCall) return message;
-      updated = true;
       return { ...assistant, content };
     });
 
@@ -1119,18 +1413,6 @@ export class AgUiThreadRuntimeCore {
     );
   }
 
-  private resetHead(parentId: string | null | undefined) {
-    if (!parentId) {
-      if (this.messages.length) {
-        this.messages = [];
-      }
-      return;
-    }
-    const idx = this.messages.findIndex((message) => message.id === parentId);
-    if (idx === -1) return;
-    this.messages = this.messages.slice(0, idx + 1);
-  }
-
   private isTerminalStatus(status?: MessageStatus): boolean {
     return status?.type === "complete" || status?.type === "incomplete";
   }
@@ -1156,7 +1438,7 @@ export class AgUiThreadRuntimeCore {
     if (!this.history) return;
     const parentId = this.assistantHistoryParents.get(messageId);
     if (parentId === undefined) return;
-    const message = this.messages.find((m) => m.id === messageId);
+    const message = this.tryGetMessage(messageId)?.message;
     if (!message || message.role !== "assistant") return;
     if (!this.isPersistableStatus(message.status)) return;
     this.assistantHistoryParents.delete(messageId);