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

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -13,7 +13,7 @@ import type {
   ThreadMessage,
   ToolCallMessagePart,
 } from "@assistant-ui/core";
-import type { HttpAgent } from "@ag-ui/client";
+import type { AbstractAgent } from "@ag-ui/client";
 import jsonpatch, { type Operation } from "fast-json-patch";
 import type { Logger } from "./logger";
 import type { AgUiEvent, AgUiInterrupt, AgUiResumeEntry } from "./types";
@@ -22,6 +22,7 @@ import {
   AG_UI_METADATA_NAMESPACE,
   type AgUiCustomMetadata,
   RunAggregator,
+  tryParseJSON,
 } from "./adapter/run-aggregator";
 import {
   fromAgUiMessages,
@@ -48,7 +49,7 @@ type ResumeRunConfig = {
 };
 
 type CoreOptions = {
-  agent: HttpAgent;
+  agent: AbstractAgent;
   logger: Logger;
   showThinking: boolean;
   onError?: (error: Error) => void;
@@ -60,7 +61,7 @@ type CoreOptions = {
 const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
 
 export class AgUiThreadRuntimeCore {
-  private agent: HttpAgent;
+  private agent: AbstractAgent;
   private logger: Logger;
   private showThinking: boolean;
   private onError: ((error: Error) => void) | undefined;
@@ -214,6 +215,30 @@ export class AgUiThreadRuntimeCore {
     );
   }
 
+  async resumeInFlightRun(messages: readonly ThreadMessage[]): Promise<void> {
+    // Without a resume stream startRun would re-run the agent from scratch.
+    const resumeStream = this.history?.resume?.bind(this.history);
+    if (!resumeStream) {
+      const error = new Error(
+        "[agui] unstable_resume requires a ThreadHistoryAdapter with a resume() method; skipping resume after thread switch",
+      );
+      this.logger.error?.(error.message);
+      this.onError?.(error);
+      return;
+    }
+    const parentId = messages.at(-1)?.id ?? null;
+    try {
+      await this.startRun(
+        parentId,
+        this.lastRunConfig,
+        undefined,
+        resumeStream,
+      );
+    } catch {
+      // startRun already reported via onError; don't reject the switch.
+    }
+  }
+
   private assertNoPendingInterrupts(): void {
     if (!this.getPendingInterrupts()) return;
     throw new Error(
@@ -396,27 +421,37 @@ export class AgUiThreadRuntimeCore {
     this.maybeResumeAfterToolResults(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.
+  // The continuation fires whether the frontend result lands before
+  // RUN_FINISHED (the status flips to requires-action only later, while the
+  // run is still draining) or after it.
   private maybeResumeAfterToolResults(messageId: string): void {
+    if (!this.maybeCompleteAfterToolResults(messageId)) return;
+
+    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 maybeCompleteAfterToolResults(messageId: string): boolean {
     const message = this.messages.find((m) => m.id === messageId);
-    if (!message || message.role !== "assistant") return;
+    if (!message || message.role !== "assistant") return false;
     const assistant = message as ThreadAssistantMessage;
     if (
       assistant.status?.type !== "requires-action" ||
       assistant.status.reason !== "tool-calls"
     ) {
-      return;
+      return false;
     }
     const allResolved = assistant.content.every(
       (part) =>
         part.type !== "tool-call" ||
         ("result" in part && part.result !== undefined),
     );
-    if (!allResolved) return;
+    if (!allResolved) return false;
 
     this.messages = this.messages.map((m) =>
       m.id === messageId && m.role === "assistant"
@@ -428,14 +463,7 @@ export class AgUiThreadRuntimeCore {
     );
     this.notifyUpdate();
     this.persistAssistantHistory(messageId);
-
-    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);
+    return true;
   }
 
   private startResumeRun(messageId: string): void {
@@ -722,6 +750,7 @@ export class AgUiThreadRuntimeCore {
         unstable_annotations: [],
         unstable_data: [],
         steps: [],
+        isOptimistic: true,
         custom: {},
       },
     };
@@ -742,9 +771,11 @@ export class AgUiThreadRuntimeCore {
       );
       this.messages = this.messages.filter((m) => m.id !== oldId);
     } else {
-      this.messages = this.messages.map((m) =>
-        m.id === oldId ? { ...m, id: newId } : m,
-      );
+      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 pendingParent = this.assistantHistoryParents.get(oldId);
@@ -869,6 +900,7 @@ export class AgUiThreadRuntimeCore {
       unstable_annotations: annotations,
       unstable_data: data,
       steps,
+      ...(current.isOptimistic ? { isOptimistic: true } : {}),
       ...(incoming.timing ? { timing: incoming.timing } : {}),
       custom: incoming.custom
         ? { ...current.custom, ...incoming.custom }
@@ -904,11 +936,57 @@ export class AgUiThreadRuntimeCore {
         this.importMessagesSnapshot(event.messages);
         return;
       }
+      case "TOOL_CALL_RESULT": {
+        if (!aggregator.hasToolCall(event.toolCallId)) {
+          const messageId = this.findMessageIdForToolCall(event.toolCallId);
+          if (messageId !== undefined) {
+            this.applyCrossRunToolResult(messageId, event);
+            return;
+          }
+        }
+        aggregator.handle(event);
+        return;
+      }
       default:
         aggregator.handle(event);
     }
   }
 
+  private applyCrossRunToolResult(
+    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 assistant = message as ThreadAssistantMessage;
+      let matchedToolCall = false;
+      const content = assistant.content.map((part) => {
+        if (part.type !== "tool-call" || part.toolCallId !== event.toolCallId)
+          return part;
+        matchedToolCall = true;
+        return {
+          ...part,
+          result: tryParseJSON(event.content ?? "") as ReadonlyJSONValue,
+          ...(event.role === "tool" ? { isError: false } : {}),
+          ...(event.messageId
+            ? { unstable_toolMessageId: event.messageId }
+            : {}),
+        };
+      });
+      if (!matchedToolCall) return message;
+      updated = true;
+      return { ...assistant, content };
+    });
+
+    if (!updated) return;
+    this.notifyUpdate();
+    // Not maybeResumeAfterToolResults: the delivering run is already in
+    // flight, and a resume from the owner would reset the head past it.
+    this.maybeCompleteAfterToolResults(messageId);
+  }
+
   private importMessagesSnapshot(rawMessages: readonly unknown[]) {
     try {
       const normalized = fromAgUiMessages(rawMessages, {

package/src/runtime/adapter/conversions.ts

@@ -204,6 +204,115 @@ function toInputContent(
   return null;
 }
 
+type SnapshotAttachment = NonNullable<
+  CoreThreadMessageLike["attachments"]
+>[number];
+
+const mediaInputTypes = new Set(["image", "audio", "video", "document"]);
+
+// Inverse of buildInputSource.
+function inputSourceToString(
+  value: unknown,
+): { value: string; mimeType?: string } | null {
+  if (!isObject(value)) return null;
+  const sourceValue = getString(value, "value");
+  if (sourceValue === undefined) return null;
+  const mimeType = getString(value, "mimeType");
+  const type = getString(value, "type");
+  if (type === "url") {
+    return { value: sourceValue, ...(mimeType !== undefined && { mimeType }) };
+  }
+  if (type === "data") {
+    const resolvedMimeType = mimeType ?? "application/octet-stream";
+    return {
+      value: `data:${resolvedMimeType};base64,${sourceValue}`,
+      mimeType: resolvedMimeType,
+    };
+  }
+  return null;
+}
+
+// Mirrors @ag-ui/client's BackwardCompatibility_0_0_47 middleware.
+function upgradeBinaryInputPart(
+  part: Record<string, unknown>,
+): Record<string, unknown> | null {
+  const mimeType = getString(part, "mimeType");
+  if (mimeType === undefined) return null;
+  const data = getString(part, "data");
+  const url = getString(part, "url");
+  const source = data
+    ? { type: "data", value: data, mimeType }
+    : url
+      ? { type: "url", value: url, mimeType }
+      : null;
+  if (!source) return null;
+  const filename = getString(part, "filename");
+  return {
+    type: mediaTypeForMime(mimeType),
+    source,
+    ...(filename && { metadata: { filename } }),
+  };
+}
+
+function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
+  if (!Array.isArray(content)) return [];
+
+  const attachments: SnapshotAttachment[] = [];
+  for (const rawPart of content) {
+    if (!isObject(rawPart)) continue;
+    const part =
+      getString(rawPart, "type") === "binary"
+        ? upgradeBinaryInputPart(rawPart)
+        : rawPart;
+    if (!part) continue;
+    const type = getString(part, "type");
+    if (type === undefined || !mediaInputTypes.has(type)) continue;
+    const source = inputSourceToString(part.source);
+    if (!source) continue;
+
+    const filename = isObject(part.metadata)
+      ? getString(part.metadata, "filename")
+      : undefined;
+    const id = attachments.length.toString();
+
+    if (type === "image") {
+      attachments.push({
+        id,
+        type: "image",
+        name: filename ?? "image",
+        ...(source.mimeType !== undefined && { contentType: source.mimeType }),
+        status: { type: "complete" },
+        content: [
+          {
+            type: "image",
+            image: source.value,
+            ...(filename !== undefined && { filename }),
+          },
+        ],
+      });
+      continue;
+    }
+
+    const mimeType = source.mimeType ?? "application/octet-stream";
+    attachments.push({
+      id,
+      type: type === "document" ? "document" : "file",
+      name: filename ?? "file",
+      contentType: mimeType,
+      status: { type: "complete" },
+      content: [
+        {
+          type: "file",
+          data: source.value,
+          mimeType,
+          ...(filename !== undefined && { filename }),
+        },
+      ],
+    });
+  }
+  return attachments;
+}
+
 function buildUserContent(message: ThreadMessageLike): string | InputContent[] {
   // File parts in message.content are intentionally skipped: the canonical
   // binary payload for files always flows through message.attachments.
@@ -353,11 +462,14 @@ function toUserOrSystemSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const messageName = getString(rawMessage, "name");
+  const attachments =
+    role === "user" ? toSnapshotAttachments(rawMessage.content) : [];
   return {
     id: getString(rawMessage, "id") ?? generateId(),
     role,
     content: extractText(rawMessage.content),
     ...(messageName !== undefined ? { name: messageName } : {}),
+    ...(attachments.length > 0 ? { attachments } : {}),
   };
 }
 

package/src/runtime/adapter/run-aggregator.ts

@@ -82,6 +82,10 @@ export class RunAggregator {
     this.onServerMessageId = options.onServerMessageId;
   }
 
+  hasToolCall(toolCallId: string): boolean {
+    return this.toolCalls.has(toolCallId);
+  }
+
   handle(event: AgUiEvent): void {
     switch (event.type) {
       case "RUN_STARTED": {
@@ -388,7 +392,7 @@ export class RunAggregator {
     ) {
       this.partOrder.push({ kind: "tool-call", toolCallId: id });
     }
-    entry.result = this.tryParseJSON(content);
+    entry.result = tryParseJSON(content);
     entry.isError = isError;
     if (toolMessageId) {
       entry.toolMessageId = toolMessageId;
@@ -396,15 +400,6 @@ export class RunAggregator {
     this.lastResolvedToolCallId = id;
   }
 
-  private tryParseJSON(value: string): unknown {
-    if (!value) return value;
-    try {
-      return JSON.parse(value);
-    } catch {
-      return value;
-    }
-  }
-
   private emit(): void {
     const snapshot: ThreadAssistantMessagePart[] = [];
 
@@ -542,3 +537,12 @@ export class RunAggregator {
     this.emit();
   }
 }
+
+export function tryParseJSON(value: string): unknown {
+  if (!value) return value;
+  try {
+    return JSON.parse(value);
+  } catch {
+    return value;
+  }
+}

package/src/runtime/adapter/subscriber.ts

@@ -38,6 +38,16 @@ type Subscriber = {
   onRunFailed?: (payload: { error: Error }) => void;
 };
 
+const isAbortError = (error: Error): boolean => {
+  if (error.name === "AbortError") return true;
+  const message = typeof error.message === "string" ? error.message : "";
+  return (
+    message === "Fetch is aborted" ||
+    message === "signal is aborted without reason" ||
+    message === "component unmounted"
+  );
+};
+
 const ensureEvent = (
   raw: unknown,
   type: AgUiEvent["type"],
@@ -143,7 +153,12 @@ export const createAgUiSubscriber = (
       dispatch({ type: "RUN_FINISHED", runId });
     },
     onRunFailed: ({ error }) => {
+      runFinishedDispatched = true;
       onRunFailed?.(error);
+      if (isAbortError(error)) {
+        dispatch({ type: "RUN_CANCELLED" } satisfies AgUiEvent);
+        return;
+      }
       const message =
         typeof error.message === "string" ? error.message : "Run failed";
       const code =

package/src/runtime/types.ts

@@ -2,29 +2,42 @@ import type {
   AttachmentAdapter,
   DictationAdapter,
   ExternalStoreSharedOptions,
+  ExternalStoreThreadListAdapter,
   FeedbackAdapter,
   RealtimeVoiceAdapter,
   SpeechSynthesisAdapter,
   ThreadHistoryAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
-import type { HttpAgent } from "@ag-ui/client";
+import type { AbstractAgent } from "@ag-ui/client";
 import type { Logger } from "./logger";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 
 /**
  * @experimental This API is still under active development and might change without notice.
+ *
+ * Same as ExternalStoreThreadListAdapter, except `onSwitchToThread` returns
+ * the messages (and optional state) to hydrate the thread with.
  */
-export type UseAgUiThreadListAdapter = {
-  threadId?: string | undefined;
-  onSwitchToNewThread?: (() => Promise<void> | void) | undefined;
+type SwitchToThreadResult = {
+  messages: readonly ThreadMessage[];
+  state?: ReadonlyJSONValue;
+  /**
+   * Set when the thread has a run in flight. The runtime resumes the run
+   * after hydrating, the same way `ThreadHistoryAdapter.load()` does when it
+   * returns `unstable_resume: true`.
+   */
+  unstable_resume?: boolean;
+};
+
+export type UseAgUiThreadListAdapter = Omit<
+  ExternalStoreThreadListAdapter,
+  "onSwitchToThread"
+> & {
   onSwitchToThread?:
-    | ((threadId: string) =>
-        | Promise<{
-            messages: readonly ThreadMessage[];
-            state?: ReadonlyJSONValue;
-          }>
-        | { messages: readonly ThreadMessage[]; state?: ReadonlyJSONValue })
+    | ((
+        threadId: string,
+      ) => Promise<SwitchToThreadResult> | SwitchToThreadResult)
     | undefined;
 };
 
@@ -42,7 +55,7 @@ export type UseAgUiRuntimeAdapters = {
 };
 
 export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
-  agent: HttpAgent;
+  agent: AbstractAgent;
   logger?: Partial<Logger>;
   showThinking?: boolean;
   onError?: (e: Error) => void;

package/src/useAgUiRuntime.ts

@@ -74,10 +74,11 @@ export function useAgUiRuntime(
   const threadList = useMemo(() => {
     if (!threadListAdapter) return undefined;
 
-    const { onSwitchToNewThread, onSwitchToThread } = threadListAdapter;
+    const { onSwitchToNewThread, onSwitchToThread, ...rest } =
+      threadListAdapter;
 
     return {
-      threadId: threadListAdapter.threadId,
+      ...rest,
       onSwitchToNewThread: onSwitchToNewThread
         ? async () => {
             await onSwitchToNewThread();
@@ -86,11 +87,17 @@ export function useAgUiRuntime(
         : undefined,
       onSwitchToThread: onSwitchToThread
         ? async (threadId: string) => {
+            // Clear before the thread id flips, or the old messages leak
+            // into the new thread as a sibling branch.
+            core.applyExternalMessages([]);
             const result = await onSwitchToThread(threadId);
             core.applyExternalMessages(result.messages);
             if (result.state) {
               core.loadExternalState(result.state);
             }
+            if (result.unstable_resume) {
+              void core.resumeInFlightRun(result.messages);
+            }
           }
         : undefined,
     };