@assistant-ui/react-ag-ui 0.0.65 → 0.0.66

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -38,17 +38,16 @@ import jsonpatch, { type Operation } from "fast-json-patch";
 import type { Logger } from "./logger";
 import { readMcpAppResourceUri } from "./mcp-tool-result";
 import type {
+  AgUiCustomMetadata,
   AgUiEvent,
   AgUiInterrupt,
   AgUiResumeEntry,
   AgUiResumeTranscript,
 } from "./types";
+import { AG_UI_METADATA_NAMESPACE, MCP_APPS_ACTIVITY_TYPE } from "./types";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import {
-  AG_UI_METADATA_NAMESPACE,
-  type AgUiCustomMetadata,
   isPlainObject,
-  MCP_APPS_ACTIVITY_TYPE,
   RunAggregator,
   tryParseJSON,
 } from "./adapter/run-aggregator";
@@ -83,6 +82,9 @@ const isResolvedToolCall = (
 ): boolean =>
   part.type === "tool-call" && "result" in part && part.result !== undefined;
 
+const isActivityPart = (part: ThreadAssistantMessage["content"][number]) =>
+  part.type === "data" && part.name.startsWith("agui-activity/");
+
 type RunConfig = NonNullable<AppendMessage["runConfig"]>;
 type ResumeStream = (
   options: ChatModelRunOptions,
@@ -113,7 +115,7 @@ type CoreOptions = {
   onError?: (error: Error) => void;
   onCancel?: () => void;
   history?: ThreadHistoryAdapter;
-  notifyUpdate: () => void;
+  isThreadSwitching?: () => boolean;
 };
 
 const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
@@ -157,7 +159,16 @@ export class AgUiThreadRuntimeCore {
   private autoCancelPendingToolCalls: boolean | undefined;
   private onError: ((error: Error) => void) | undefined;
   private onCancel: (() => void) | undefined;
-  private readonly notifyUpdate: () => void;
+  private readonly listeners = new Set<() => void>();
+  private snapshot:
+    | {
+        isLoading: boolean;
+        isRunning: boolean;
+        messageRepository: ExportedMessageRepository;
+        state: ReadonlyJSONValue | undefined;
+      }
+    | undefined;
+  private readonly isThreadSwitching: (() => boolean) | undefined;
   private readonly reportedErrors = new WeakSet<object>();
 
   private runtime: AssistantRuntime | undefined;
@@ -190,6 +201,8 @@ export class AgUiThreadRuntimeCore {
   private readonly persistedHistoryIds = new Set<string>();
   private readonly historyWrites = new Map<string, Promise<void>>();
   private _isLoading = false;
+  private _historyImportPromise: Promise<void> | undefined;
+  private historyMutationGeneration = 0;
   private _loadPromise: Promise<void> | undefined;
   private _loadRequested = false;
   private pendingResume: { owner: AbortController; messageId: string } | null =
@@ -206,10 +219,28 @@ export class AgUiThreadRuntimeCore {
     this.onError = options.onError;
     this.onCancel = options.onCancel;
     this.history = options.history;
-    this.notifyUpdate = options.notifyUpdate;
+    this.isThreadSwitching = options.isThreadSwitching;
   }
 
-  updateOptions(options: Omit<CoreOptions, "notifyUpdate">) {
+  subscribe = (listener: () => void): (() => void) => {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  };
+
+  getSnapshot = () =>
+    (this.snapshot ??= {
+      isLoading: this._isLoading,
+      isRunning: this.isRunningFlag,
+      messageRepository: this.session.export(),
+      state: this.stateSnapshot,
+    });
+
+  private publishUpdate(): void {
+    this.snapshot = undefined;
+    for (const listener of this.listeners) listener();
+  }
+
+  updateOptions(options: CoreOptions) {
     this.agent = options.agent;
     this.logger = options.logger;
     this.showThinking = options.showThinking;
@@ -266,28 +297,38 @@ export class AgUiThreadRuntimeCore {
     if (this._loadPromise) return this._loadPromise;
     if (!this.history) return Promise.resolve();
 
-    const promise = this.history.load();
+    const promise = this.history.load().then((repo) => {
+      if (!repo) return;
+
+      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.publishUpdate();
+
+      if (repo.state !== undefined) {
+        this.loadExternalState(repo.state);
+      }
+
+      return repo;
+    });
 
     this._isLoading = true;
+    this._historyImportPromise = promise.then(
+      () => {
+        this._historyImportPromise = undefined;
+      },
+      () => {
+        this._historyImportPromise = undefined;
+      },
+    );
 
     this._loadPromise = promise
       .then(async (repo) => {
-        if (!repo) return;
-
-        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) {
+        if (repo?.unstable_resume) {
           const parentId = repo.headId ?? this.session.headId;
           const resumeStream = this.history?.resume?.bind(this.history);
           await this.startRun(
@@ -304,14 +345,60 @@ export class AgUiThreadRuntimeCore {
       })
       .finally(() => {
         this._isLoading = false;
-        this.notifyUpdate();
+        this.publishUpdate();
       });
 
-    this.notifyUpdate();
+    this.publishUpdate();
     return this._loadPromise;
   }
 
+  private waitForHistoryImport(): Promise<boolean> | undefined {
+    const historyImport = this._historyImportPromise;
+    if (!historyImport) return undefined;
+    const generation = this.historyMutationGeneration;
+    return historyImport.then(
+      () =>
+        generation === this.historyMutationGeneration &&
+        !this.isThreadSwitching?.(),
+    );
+  }
+
+  private deferUntilHistoryImported(callback: () => void): boolean {
+    const historyImport = this._historyImportPromise;
+    if (!historyImport) return false;
+    const generation = this.historyMutationGeneration;
+    void historyImport.then(() => {
+      if (
+        generation !== this.historyMutationGeneration ||
+        this.isThreadSwitching?.()
+      ) {
+        return;
+      }
+      try {
+        callback();
+      } catch (error) {
+        this.reportError(error);
+      }
+    });
+    return true;
+  }
+
   async append(message: AppendMessage): Promise<void> {
+    const historyImport = this._historyImportPromise;
+    if (historyImport) {
+      const wasAtTail = message.parentId === this.session.headId;
+      const generation = this.historyMutationGeneration;
+      await historyImport;
+      if (
+        generation !== this.historyMutationGeneration ||
+        this.isThreadSwitching?.()
+      )
+        return;
+      if (wasAtTail) {
+        message = { ...message, parentId: this.session.headId };
+      }
+    }
+
     const startRun = message.startRun ?? message.role === "user";
     let ownsThread = true;
     if (startRun) {
@@ -331,7 +418,7 @@ export class AgUiThreadRuntimeCore {
     const parentId = this.session.headId;
     this.session.addOrUpdateMessage(parentId, message);
     this.session.switchToBranch(message.id);
-    this.notifyUpdate();
+    this.publishUpdate();
     this.recordHistoryEntry(parentId, message);
   }
 
@@ -353,7 +440,7 @@ export class AgUiThreadRuntimeCore {
           : this.session.headId;
     this.session.addOrUpdateMessage(parentId, threadMessage);
     this.session.switchToBranch(threadMessage.id);
-    this.notifyUpdate();
+    this.publishUpdate();
     this.recordHistoryEntry(parentId, threadMessage);
     return threadMessage.id;
   }
@@ -366,12 +453,15 @@ export class AgUiThreadRuntimeCore {
     parentId: string | null,
     config: { runConfig?: RunConfig } = {},
   ): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     this.assertNoPendingInterrupts();
     this.maybeAutoCancelPendingToolCalls();
     await this.startRun(parentId, config.runConfig);
   }
 
   async cancel(): Promise<void> {
+    this.historyMutationGeneration++;
     this.abortActiveRun();
   }
 
@@ -413,6 +503,8 @@ export class AgUiThreadRuntimeCore {
   }
 
   async resume(config: ResumeRunConfig): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     this.assertNoPendingInterrupts();
     await this.startRun(
       config.parentId,
@@ -423,6 +515,8 @@ export class AgUiThreadRuntimeCore {
   }
 
   async resumeInFlightRun(messages: readonly ThreadMessage[]): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     // Without a resume stream startRun would re-run the agent from scratch.
     const resumeStream = this.history?.resume?.bind(this.history);
     if (!resumeStream) {
@@ -502,6 +596,8 @@ export class AgUiThreadRuntimeCore {
   async submitInterruptResponses(
     responses: readonly AgUiResumeEntry[],
   ): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     const pending = this.getPendingInterrupts();
     if (!pending) {
       throw new Error(
@@ -597,6 +693,8 @@ export class AgUiThreadRuntimeCore {
   async respondToToolApproval(
     options: RespondToToolApprovalOptions,
   ): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     const pending = this.getPendingInterrupts();
     if (!pending) {
       throw new Error(
@@ -648,7 +746,7 @@ export class AgUiThreadRuntimeCore {
         `[agui] respondToToolApproval: approval "${options.approvalId}" is already decided`,
       );
     }
-    this.notifyUpdate();
+    this.publishUpdate();
 
     const assistant = this.session.tryGetMessage(pending.messageId)?.message as
       | ThreadAssistantMessage
@@ -670,6 +768,8 @@ export class AgUiThreadRuntimeCore {
     message: CreateAppendMessage,
     responses?: readonly AgUiResumeEntry[],
   ): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     const pending = this.getPendingInterrupts();
     if (!pending) {
       const pendingTools = this.getPendingToolCalls();
@@ -822,7 +922,7 @@ export class AgUiThreadRuntimeCore {
       };
     });
     if (touched) {
-      this.notifyUpdate();
+      this.publishUpdate();
     }
   }
 
@@ -857,10 +957,12 @@ export class AgUiThreadRuntimeCore {
       });
       return { ...assistant, content };
     });
-    if (updated) this.notifyUpdate();
+    if (updated) this.publishUpdate();
   }
 
   addToolResult(options: AddToolResultOptions): void {
+    if (this.deferUntilHistoryImported(() => this.addToolResult(options)))
+      return;
     // 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.
@@ -893,7 +995,7 @@ export class AgUiThreadRuntimeCore {
     });
 
     if (!updated) return;
-    this.notifyUpdate();
+    this.publishUpdate();
     this.maybeResumeAfterToolResults(sessionMessageId);
   }
 
@@ -902,6 +1004,8 @@ export class AgUiThreadRuntimeCore {
     toolCallId: string;
     interaction: Unstable_ToolInteraction;
   }): Promise<void> {
+    const historyImport = this.waitForHistoryImport();
+    if (historyImport && !(await historyImport)) return;
     const sessionMessageId = this.session.tryGetMessage(options.messageId)
       ? options.messageId
       : this.findMessageIdForToolCall(options.toolCallId);
@@ -941,7 +1045,7 @@ export class AgUiThreadRuntimeCore {
       );
     }
 
-    this.notifyUpdate();
+    this.publishUpdate();
 
     const message = this.session.tryGetMessage(sessionMessageId)?.message;
     const history = this.history;
@@ -959,6 +1063,8 @@ export class AgUiThreadRuntimeCore {
   }
 
   sendA2uiAction(action: Record<string, unknown>): void {
+    if (this.deferUntilHistoryImported(() => this.sendA2uiAction(action)))
+      return;
     this.assertNoPendingInterrupts();
     this.maybeAutoCancelPendingToolCalls();
     const parentId = this.session.headId;
@@ -1028,7 +1134,7 @@ export class AgUiThreadRuntimeCore {
         : current,
     );
     if (!updated) return false;
-    this.notifyUpdate();
+    this.publishUpdate();
     this.persistAssistantHistory(messageId);
     return true;
   }
@@ -1108,12 +1214,12 @@ export class AgUiThreadRuntimeCore {
     ) {
       this.pendingResume = null;
     }
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   loadExternalState(state: ReadonlyJSONValue): void {
     this.stateSnapshot = state;
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   setState(
@@ -1123,15 +1229,16 @@ export class AgUiThreadRuntimeCore {
   ): void {
     this.stateSnapshot =
       typeof next === "function" ? next(this.stateSnapshot) : next;
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   resetState(): void {
     this.stateSnapshot = undefined;
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   resetThreadState(): void {
+    this.historyMutationGeneration++;
     const controller = this.abortController;
     const activeRunAgent = this.activeRunAgent;
 
@@ -1152,7 +1259,7 @@ export class AgUiThreadRuntimeCore {
         controller.abort();
       }
     } else {
-      this.notifyUpdate();
+      this.publishUpdate();
     }
   }
 
@@ -1162,6 +1269,7 @@ export class AgUiThreadRuntimeCore {
     resume?: ResumeDispatch,
     resumeStream?: ResumeStream,
   ): Promise<void> {
+    if (this.isThreadSwitching?.()) return;
     // A default AG-UI run supersedes the active run; the hook's opt-in message
     // queue serializes sends instead. append supersedes earlier, before it links
     // its message; this covers the entry points that start a run without one.
@@ -1553,7 +1661,7 @@ export class AgUiThreadRuntimeCore {
 
   private setRunning(running: boolean) {
     this.isRunningFlag = running;
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private finishRun(controller: AbortController | null) {
@@ -1584,7 +1692,7 @@ export class AgUiThreadRuntimeCore {
     };
     this.session.addOrUpdateMessage(parentId ?? this.session.headId, assistant);
     this.session.switchToBranch(id);
-    this.notifyUpdate();
+    this.publishUpdate();
     return id;
   }
 
@@ -1650,7 +1758,7 @@ export class AgUiThreadRuntimeCore {
       }
     }
 
-    this.notifyUpdate();
+    this.publishUpdate();
     return !collidesWithExisting;
   }
 
@@ -1690,7 +1798,7 @@ export class AgUiThreadRuntimeCore {
       this.reassignAssistantId(messageId, stableId);
       resolvedMessageId = stableId;
     } else {
-      this.notifyUpdate();
+      this.publishUpdate();
     }
     if (this.isPersistableStatus(latestStatus)) {
       this.persistAssistantHistory(resolvedMessageId);
@@ -1760,6 +1868,128 @@ export class AgUiThreadRuntimeCore {
     return changed ? content : next;
   }
 
+  // A snapshot built from user, assistant, and tool records alone carries no
+  // reasoning or activity of its own, so the parts streamed locally would be
+  // lost on every refresh; keep the local ones where they streamed when the
+  // snapshot has none of that kind, matching the merge rule of @ag-ui/client.
+  private mergeSnapshotAssistantContent(
+    previous: ThreadAssistantMessage["content"],
+    next: ThreadAssistantMessage["content"],
+    snapshotHasReasoning: boolean,
+    snapshotHasActivity: boolean,
+  ): ThreadAssistantMessage["content"] {
+    const shouldKeep = (part: ThreadAssistantMessage["content"][number]) =>
+      part.type === "reasoning"
+        ? this.showThinking && !snapshotHasReasoning
+        : isActivityPart(part) && !snapshotHasActivity;
+    const kept = previous.filter(shouldKeep);
+    const merged = this.reconcileSnapshotNestedContent(
+      previous,
+      this.preserveToolInteractions(previous, next),
+    );
+    if (kept.length === 0) return merged;
+
+    const mergedByType = new Map<string, number[]>();
+    const mergedTools = new Map<string, number>();
+    for (const [index, part] of merged.entries()) {
+      if (part.type === "tool-call") {
+        mergedTools.set(part.toolCallId, index);
+      } else {
+        const indexes = mergedByType.get(part.type) ?? [];
+        indexes.push(index);
+        mergedByType.set(part.type, indexes);
+      }
+    }
+
+    const previousOrdinals = new Map<string, number>();
+    const insertions: (typeof kept)[] = Array.from(
+      { length: merged.length + 1 },
+      () => [],
+    );
+    let predecessorIndex = -1;
+    for (const part of previous) {
+      if (shouldKeep(part)) {
+        insertions[predecessorIndex + 1]!.push(part);
+        continue;
+      }
+
+      let matchedIndex: number | undefined;
+      if (part.type === "tool-call") {
+        matchedIndex = mergedTools.get(part.toolCallId);
+      } else {
+        const ordinal = previousOrdinals.get(part.type) ?? 0;
+        matchedIndex = mergedByType.get(part.type)?.[ordinal];
+        previousOrdinals.set(part.type, ordinal + 1);
+      }
+      if (matchedIndex !== undefined) predecessorIndex = matchedIndex;
+    }
+
+    const result: ThreadAssistantMessage["content"][number][] = [];
+    for (const [index, part] of merged.entries()) {
+      result.push(...insertions[index]!, part);
+    }
+    result.push(...insertions[merged.length]!);
+    return result;
+  }
+
+  private reconcileSnapshotNestedContent(
+    previous: ThreadAssistantMessage["content"],
+    next: ThreadAssistantMessage["content"],
+  ): ThreadAssistantMessage["content"] {
+    const nestedIds = new Set<string>();
+    for (const part of previous) {
+      if (part.type !== "tool-call" || !part.messages) continue;
+      for (const nested of part.messages) {
+        if (nested.role !== "assistant") continue;
+        for (const call of iterateToolCallParts(nested.content)) {
+          nestedIds.add(call.toolCallId);
+        }
+      }
+    }
+
+    const flattened = new Map<string, ToolCallMessagePart>();
+    for (const part of next) {
+      if (part.type === "tool-call" && nestedIds.has(part.toolCallId)) {
+        flattened.set(part.toolCallId, part);
+      }
+    }
+    const carried =
+      flattened.size === 0
+        ? previous
+        : mapToolCallPartsDeep(previous, (part) => {
+            const snapshot = flattened.get(part.toolCallId);
+            if (!snapshot) return part;
+            return {
+              ...part,
+              ...(snapshot.result !== undefined
+                ? { result: snapshot.result }
+                : {}),
+              ...(snapshot.isError !== undefined
+                ? { isError: snapshot.isError }
+                : {}),
+              ...(snapshot.artifact !== undefined
+                ? { artifact: snapshot.artifact }
+                : {}),
+            };
+          }).content;
+    const previousCalls = new Map<string, ToolCallMessagePart>();
+    for (const part of carried) {
+      if (part.type !== "tool-call") continue;
+      previousCalls.set(part.toolCallId, part);
+    }
+
+    return next.flatMap<ThreadAssistantMessage["content"][number]>((part) => {
+      if (part.type === "tool-call" && flattened.has(part.toolCallId)) {
+        return [];
+      }
+      if (part.type !== "tool-call") return [part];
+      const prior = previousCalls.get(part.toolCallId);
+      return prior?.messages === undefined
+        ? [part]
+        : [{ ...part, messages: prior.messages }];
+    });
+  }
+
   private mergeAssistantMetadata(
     current: ThreadAssistantMessage["metadata"],
     incoming: NonNullable<ChatModelRunResult["metadata"]>,
@@ -1894,7 +2124,7 @@ export class AgUiThreadRuntimeCore {
         };
       });
     }
-    this.notifyUpdate();
+    this.publishUpdate();
   }
 
   private applyCrossRunToolResult(
@@ -1965,7 +2195,7 @@ export class AgUiThreadRuntimeCore {
     });
 
     if (!updated) return;
-    this.notifyUpdate();
+    this.publishUpdate();
     // Not maybeResumeAfterToolResults: the delivering run is already in
     // flight, and a resume from the owner would reset the head past it.
     this.maybeCompleteAfterToolResults(messageId);
@@ -2028,7 +2258,7 @@ export class AgUiThreadRuntimeCore {
     });
 
     if (!updated) return;
-    this.notifyUpdate();
+    this.publishUpdate();
     this.maybeCompleteAfterToolResults(messageId);
   }
 
@@ -2037,9 +2267,11 @@ export class AgUiThreadRuntimeCore {
     activeAssistantId: string | undefined,
   ) {
     try {
-      const activeMessage = activeAssistantId
-        ? this.session.tryGetMessage(activeAssistantId)?.message
+      const activeAssistantItem = activeAssistantId
+        ? this.session.tryGetMessage(activeAssistantId)
         : undefined;
+      const activeMessage = activeAssistantItem?.message;
+      const activeAssistantParentId = activeAssistantItem?.parentId;
       const activeAssistant =
         activeMessage?.role === "assistant" ? activeMessage : undefined;
       const normalized = fromAgUiMessages(rawMessages, {
@@ -2048,28 +2280,9 @@ export class AgUiThreadRuntimeCore {
       const converted: ThreadMessage[] = [];
       for (const message of normalized) {
         try {
-          const convertedMessage = fromThreadMessageLike(
-            message,
-            generateId(),
-            FALLBACK_USER_STATUS,
+          converted.push(
+            fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS),
           );
-          const existing = this.session.tryGetMessage(
-            convertedMessage.id,
-          )?.message;
-          if (
-            convertedMessage.role === "assistant" &&
-            existing?.role === "assistant"
-          ) {
-            converted.push({
-              ...convertedMessage,
-              content: this.preserveToolInteractions(
-                existing.content,
-                convertedMessage.content,
-              ),
-            });
-          } else {
-            converted.push(convertedMessage);
-          }
         } catch (error) {
           this.logger.error?.(
             "[agui] failed to import message from snapshot",
@@ -2077,6 +2290,35 @@ export class AgUiThreadRuntimeCore {
           );
         }
       }
+      const snapshotHasReasoning = converted.some(
+        (message) =>
+          message.role === "assistant" &&
+          message.content.some((part) => part.type === "reasoning"),
+      );
+      const snapshotHasActivity = converted.some(
+        (message) =>
+          message.role === "assistant" && message.content.some(isActivityPart),
+      );
+      for (let index = 0; index < converted.length; index++) {
+        const convertedMessage = converted[index]!;
+        const existing = this.session.tryGetMessage(
+          convertedMessage.id,
+        )?.message;
+        if (
+          convertedMessage.role === "assistant" &&
+          existing?.role === "assistant"
+        ) {
+          converted[index] = {
+            ...convertedMessage,
+            content: this.mergeSnapshotAssistantContent(
+              existing.content,
+              convertedMessage.content,
+              snapshotHasReasoning,
+              snapshotHasActivity,
+            ),
+          };
+        }
+      }
       const snapshotContainsActiveAssistant = converted.some(
         (message) => message.id === activeAssistant?.id,
       );
@@ -2084,7 +2326,9 @@ export class AgUiThreadRuntimeCore {
         activeAssistant !== undefined &&
         !snapshotContainsActiveAssistant &&
         (activeAssistant.metadata.isOptimistic !== true ||
-          converted.at(-1)?.role !== "assistant");
+          converted.at(-1)?.role !== "assistant" ||
+          (activeAssistantParentId !== undefined &&
+            converted.at(-1)?.id === activeAssistantParentId));
       if (preservesActiveAssistant) {
         converted.push(activeAssistant);
       }