@assistant-ui/react-ag-ui 0.0.34 → 0.0.36

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -5,13 +5,16 @@ import type {
   AddToolResultOptions,
   AppendMessage,
   AssistantRuntime,
+  ChatModelRunOptions,
   ChatModelRunResult,
   MessageStatus,
   ThreadAssistantMessage,
   ThreadHistoryAdapter,
   ThreadMessage,
+  ToolCallMessagePart,
 } from "@assistant-ui/core";
 import type { HttpAgent } from "@ag-ui/client";
+import jsonpatch, { type Operation } from "fast-json-patch";
 import type { Logger } from "./logger";
 import type { AgUiEvent, AgUiInterrupt, AgUiResumeEntry } from "./types";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
@@ -34,11 +37,14 @@ const isOptimisticId = (id: string) => id.startsWith(optimisticPrefix);
 const symbolResumeShim = Symbol("agui-resume-shim");
 
 type RunConfig = NonNullable<AppendMessage["runConfig"]>;
+type ResumeStream = (
+  options: ChatModelRunOptions,
+) => AsyncGenerator<ChatModelRunResult, void, unknown>;
 type ResumeRunConfig = {
   parentId: string | null;
   sourceId: string | null;
   runConfig: RunConfig;
-  stream?: unknown;
+  stream?: ResumeStream;
 };
 
 type CoreOptions = {
@@ -73,6 +79,7 @@ export class AgUiThreadRuntimeCore {
   private readonly recordedHistoryIds = new Set<string>();
   private _isLoading = false;
   private _loadPromise: Promise<void> | undefined;
+  private pendingResumeMessageId: string | null = null;
 
   constructor(options: CoreOptions) {
     this.agent = options.agent;
@@ -135,7 +142,13 @@ export class AgUiThreadRuntimeCore {
 
         if (repo.unstable_resume) {
           const parentId = repo.headId ?? messages.at(-1)?.id ?? null;
-          await this.startRun(parentId, this.lastRunConfig);
+          const resumeStream = this.history?.resume?.bind(this.history);
+          await this.startRun(
+            parentId,
+            this.lastRunConfig,
+            undefined,
+            resumeStream,
+          );
         }
       })
       .catch((error) => {
@@ -193,14 +206,11 @@ export class AgUiThreadRuntimeCore {
 
   async resume(config: ResumeRunConfig): Promise<void> {
     this.assertNoPendingInterrupts();
-    if (config.stream) {
-      this.logger.debug?.(
-        "[agui] resume stream is not supported, falling back to regular run",
-      );
-    }
     await this.startRun(
       config.parentId,
       config.runConfig ?? this.lastRunConfig,
+      undefined,
+      config.stream,
     );
   }
 
@@ -360,7 +370,6 @@ export class AgUiThreadRuntimeCore {
 
   addToolResult(options: AddToolResultOptions): void {
     let updated = false;
-    let shouldResume = false;
     this.messages = this.messages.map((message) => {
       if (message.id !== options.messageId || message.role !== "assistant")
         return message;
@@ -379,47 +388,60 @@ export class AgUiThreadRuntimeCore {
       });
       if (!matchedToolCall) return message;
       updated = true;
-
-      if (
-        assistant.status?.type === "requires-action" &&
-        assistant.status.reason === "tool-calls" &&
-        content.every(
-          (part) =>
-            part.type !== "tool-call" ||
-            ("result" in part && part.result !== undefined),
-        )
-      ) {
-        shouldResume = true;
-        return {
-          ...assistant,
-          content,
-          status: { type: "complete" as const, reason: "unknown" as const },
-        };
-      }
-
-      return {
-        ...assistant,
-        content,
-      };
+      return { ...assistant, content };
     });
 
-    if (updated) {
-      this.notifyUpdate();
+    if (!updated) return;
+    this.notifyUpdate();
+    this.maybeResumeAfterToolResults(options.messageId);
+  }
 
-      if (shouldResume) {
-        this.persistAssistantHistory(options.messageId);
+  // Re-evaluate whether every tool call on the assistant message now has a
+  // result and, if so, drive the follow-up run. Centralizing this lets the
+  // continuation fire whether the frontend result lands before RUN_FINISHED
+  // (status flips to requires-action only later, while a run is still draining)
+  // or after it. Idempotent: once the message is marked complete it no-ops.
+  private maybeResumeAfterToolResults(messageId: string): void {
+    const message = this.messages.find((m) => m.id === messageId);
+    if (!message || message.role !== "assistant") return;
+    const assistant = message as ThreadAssistantMessage;
+    if (
+      assistant.status?.type !== "requires-action" ||
+      assistant.status.reason !== "tool-calls"
+    ) {
+      return;
+    }
+    const allResolved = assistant.content.every(
+      (part) =>
+        part.type !== "tool-call" ||
+        ("result" in part && part.result !== undefined),
+    );
+    if (!allResolved) return;
+
+    this.messages = this.messages.map((m) =>
+      m.id === messageId && m.role === "assistant"
+        ? {
+            ...(m as ThreadAssistantMessage),
+            status: { type: "complete" as const, reason: "unknown" as const },
+          }
+        : m,
+    );
+    this.notifyUpdate();
+    this.persistAssistantHistory(messageId);
 
-        if (!this.isRunningFlag) {
-          void this.startRun(options.messageId, this.lastRunConfig).catch(
-            (error) => {
-              this.onError?.(
-                error instanceof Error ? error : new Error(String(error)),
-              );
-            },
-          );
-        }
-      }
+    if (this.isRunningFlag) {
+      // A run is still draining (RUN_FINISHED arrived but the stream has not
+      // closed). Defer until startRun's tail so we never start two runs.
+      this.pendingResumeMessageId = messageId;
+      return;
     }
+    this.startResumeRun(messageId);
+  }
+
+  private startResumeRun(messageId: string): void {
+    void this.startRun(messageId, this.lastRunConfig).catch((error) => {
+      this.onError?.(error instanceof Error ? error : new Error(String(error)));
+    });
   }
 
   applyExternalMessages(messages: readonly ThreadMessage[]): void {
@@ -441,20 +463,14 @@ export class AgUiThreadRuntimeCore {
     parentId: string | null,
     runConfig?: RunConfig,
     resume?: AgUiResumeEntry[],
+    resumeStream?: ResumeStream,
   ): Promise<void> {
     const normalizedRunConfig = runConfig ?? {};
     this.lastRunConfig = normalizedRunConfig;
     this.resetHead(parentId);
     const historicalMessages = [...this.messages];
 
-    const runId = generateId();
     this.pendingError = null;
-    const input = this.buildRunInput(
-      runId,
-      normalizedRunConfig,
-      historicalMessages,
-      resume,
-    );
     const assistantParentId = parentId ?? this.messages.at(-1)?.id ?? null;
     let assistantMessageId: string | undefined;
     const ensureAssistant = () => {
@@ -465,15 +481,17 @@ export class AgUiThreadRuntimeCore {
       return created;
     };
 
+    const applyUpdate = (update: ChatModelRunResult) => {
+      const resolved = this.updateAssistantMessage(ensureAssistant(), update);
+      if (resolved !== assistantMessageId) {
+        assistantMessageId = resolved;
+      }
+    };
+
     const aggregator = new RunAggregator({
       showThinking: this.showThinking,
       logger: this.logger,
-      emit: (update) => {
-        const resolved = this.updateAssistantMessage(ensureAssistant(), update);
-        if (resolved !== assistantMessageId) {
-          assistantMessageId = resolved;
-        }
-      },
+      emit: applyUpdate,
       onServerMessageId: (serverId) => {
         const placeholder = ensureAssistant();
         if (placeholder === serverId) return;
@@ -487,40 +505,63 @@ export class AgUiThreadRuntimeCore {
     const abortSignal = abortController.signal;
     this.abortController = abortController;
 
+    let cancelRun = () => dispatch({ type: "RUN_CANCELLED" });
     abortSignal.addEventListener(
       "abort",
       () => {
-        dispatch({ type: "RUN_CANCELLED" });
+        cancelRun();
         this.finishRun(abortController);
         this.onCancel?.();
       },
       { once: true },
     );
 
-    const subscriber = createAgUiSubscriber({
-      dispatch,
-      runId,
-      logger: this.logger,
-      onRunFailed: (error) => {
-        this.pendingError = error;
-        this.onError?.(error);
-      },
-    });
-
-    aggregator.handle({ type: "RUN_STARTED", runId });
     this.setRunning(true);
 
     try {
-      try {
-        (this.agent as any).messages = input.messages;
-        (this.agent as any).threadId = input.threadId;
-        (this.agent as any).state = input.state ?? null;
-      } catch {
-        // ignore
+      if (resumeStream) {
+        // Cancel flips only the status; an aggregator RUN_CANCELLED would emit an empty snapshot and wipe the replayed content.
+        cancelRun = () =>
+          applyUpdate({ status: { type: "incomplete", reason: "cancelled" } });
+        await this.consumeResumeStream(resumeStream, {
+          runConfig: normalizedRunConfig,
+          threadId: this.agent.threadId || "main",
+          parentId: assistantParentId,
+          historicalMessages,
+          abortSignal,
+          ensureAssistant,
+          applyUpdate,
+          getAssistantMessageId: () => assistantMessageId,
+        });
+      } else {
+        const runId = generateId();
+        aggregator.handle({ type: "RUN_STARTED", runId });
+        const input = this.buildRunInput(
+          runId,
+          normalizedRunConfig,
+          historicalMessages,
+          resume,
+        );
+        const subscriber = createAgUiSubscriber({
+          dispatch,
+          runId,
+          logger: this.logger,
+          onRunFailed: (error) => {
+            this.pendingError = error;
+            this.onError?.(error);
+          },
+        });
+        try {
+          (this.agent as any).messages = input.messages;
+          (this.agent as any).threadId = input.threadId;
+          (this.agent as any).state = input.state ?? null;
+        } catch {
+          // ignore
+        }
+        await (this.agent as any).runAgent(input, subscriber, {
+          signal: abortSignal,
+        });
       }
-      await (this.agent as any).runAgent(input, subscriber, {
-        signal: abortSignal,
-      });
     } catch (error) {
       if (!abortSignal.aborted) {
         const err = error instanceof Error ? error : new Error(String(error));
@@ -535,8 +576,77 @@ export class AgUiThreadRuntimeCore {
     if (this.pendingError) {
       const err = this.pendingError;
       this.pendingError = null;
+      this.pendingResumeMessageId = null;
       throw err;
     }
+
+    // A tool result that landed before the run settled deferred its
+    // continuation here so a second run never overlaps the first.
+    if (this.pendingResumeMessageId !== null) {
+      const resumeMessageId = this.pendingResumeMessageId;
+      this.pendingResumeMessageId = null;
+      if (!abortSignal.aborted) {
+        this.startResumeRun(resumeMessageId);
+      }
+    }
+  }
+
+  // Replays a persisted run's snapshots into the existing assistant message, bypassing agent.runAgent so it is not re-invoked.
+  private async consumeResumeStream(
+    stream: ResumeStream,
+    ctx: {
+      runConfig: RunConfig;
+      threadId: string;
+      parentId: string | null;
+      historicalMessages: readonly ThreadMessage[];
+      abortSignal: AbortSignal;
+      ensureAssistant: () => string;
+      applyUpdate: (update: ChatModelRunResult) => void;
+      getAssistantMessageId: () => string | undefined;
+    },
+  ): Promise<void> {
+    const assistantId = ctx.ensureAssistant();
+    const currentId = () => ctx.getAssistantMessageId() ?? assistantId;
+    const options: ChatModelRunOptions = {
+      messages: ctx.historicalMessages,
+      runConfig: ctx.runConfig,
+      abortSignal: ctx.abortSignal,
+      context: this.runtime?.thread.getModelContext() ?? {},
+      unstable_assistantMessageId: assistantId,
+      unstable_threadId: ctx.threadId,
+      unstable_parentId: ctx.parentId,
+      unstable_getMessage: () => {
+        const message = this.messages.find((m) => m.id === currentId());
+        if (!message) {
+          throw new Error(
+            "[agui] resume stream requested the assistant message before it existed",
+          );
+        }
+        return message;
+      },
+    };
+
+    try {
+      for await (const result of stream(options)) {
+        if (ctx.abortSignal.aborted) return;
+        ctx.applyUpdate(result);
+      }
+    } catch (error) {
+      if (ctx.abortSignal.aborted) return;
+      const err = error instanceof Error ? error : new Error(String(error));
+      ctx.applyUpdate({
+        status: { type: "incomplete", reason: "error", error: err.message },
+      });
+      this.onError?.(err);
+      this.pendingError = this.pendingError ?? err;
+      return;
+    }
+
+    if (ctx.abortSignal.aborted) return;
+    const current = this.messages.find((m) => m.id === currentId());
+    if (!current || current.status?.type === "running") {
+      ctx.applyUpdate({ status: { type: "complete", reason: "unknown" } });
+    }
   }
 
   private buildRunInput(
@@ -668,10 +778,16 @@ export class AgUiThreadRuntimeCore {
         ? this.mergeAssistantMetadata(assistant.metadata, update.metadata)
         : assistant.metadata;
       latestStatus = update.status ?? assistant.status;
+      const content =
+        update.content !== undefined
+          ? this.preserveToolResults(
+              assistant.content,
+              update.content as ThreadAssistantMessage["content"],
+            )
+          : assistant.content;
       return {
         ...assistant,
-        content: (update.content ??
-          assistant.content) as ThreadAssistantMessage["content"],
+        content,
         status: latestStatus,
         metadata,
       };
@@ -691,9 +807,47 @@ export class AgUiThreadRuntimeCore {
     if (this.isPersistableStatus(latestStatus)) {
       this.persistAssistantHistory(resolvedMessageId);
     }
+    this.maybeResumeAfterToolResults(resolvedMessageId);
     return resolvedMessageId;
   }
 
+  // The RunAggregator rebuilds the assistant content from stream events only,
+  // so a fresh snapshot omits results injected via addToolResult (frontend tool
+  // execution). Carry those results forward so the aggregator never clobbers
+  // them. Results are only ever added in this flow, so preserving is safe.
+  private preserveToolResults(
+    previous: ThreadAssistantMessage["content"],
+    next: ThreadAssistantMessage["content"],
+  ): ThreadAssistantMessage["content"] {
+    const resolved = new Map<string, ToolCallMessagePart>();
+    for (const part of previous) {
+      if (
+        part.type === "tool-call" &&
+        "result" in part &&
+        part.result !== undefined
+      ) {
+        resolved.set(part.toolCallId, part);
+      }
+    }
+    if (resolved.size === 0) return next;
+
+    let changed = false;
+    const merged = next.map((part) => {
+      if (part.type !== "tool-call") return part;
+      if ("result" in part && part.result !== undefined) return part;
+      const prior = resolved.get(part.toolCallId);
+      if (!prior) return part;
+      changed = true;
+      return {
+        ...part,
+        result: prior.result,
+        ...(prior.artifact !== undefined ? { artifact: prior.artifact } : {}),
+        ...(prior.isError !== undefined ? { isError: prior.isError } : {}),
+      };
+    });
+    return changed ? (merged as ThreadAssistantMessage["content"]) : next;
+  }
+
   private mergeAssistantMetadata(
     current: ThreadAssistantMessage["metadata"],
     incoming: NonNullable<ChatModelRunResult["metadata"]>,
@@ -730,7 +884,20 @@ export class AgUiThreadRuntimeCore {
         return;
       }
       case "STATE_DELTA": {
-        this.logger.debug?.("[agui] state delta event ignored", event.delta);
+        if (event.delta.length === 0) return;
+        try {
+          const state = this.stateSnapshot ?? {};
+          const result = jsonpatch.applyPatch(
+            state,
+            event.delta as Operation[],
+            /* validateOperation */ true,
+            /* mutateDocument */ false,
+          );
+          this.stateSnapshot = result.newDocument as ReadonlyJSONValue;
+          this.notifyUpdate();
+        } catch (error) {
+          this.logger.error?.("[agui] failed to apply state delta", error);
+        }
         return;
       }
       case "MESSAGES_SNAPSHOT": {
@@ -744,16 +911,14 @@ export class AgUiThreadRuntimeCore {
 
   private importMessagesSnapshot(rawMessages: readonly unknown[]) {
     try {
-      const normalized = fromAgUiMessages(rawMessages);
+      const normalized = fromAgUiMessages(rawMessages, {
+        showThinking: this.showThinking,
+      });
       const converted: ThreadMessage[] = [];
       for (const message of normalized) {
         try {
           converted.push(
-            fromThreadMessageLike(
-              message as any,
-              generateId(),
-              FALLBACK_USER_STATUS,
-            ),
+            fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS),
           );
         } catch (error) {
           this.logger.error?.(