@assistant-ui/react-ag-ui 0.0.71 → 0.1.0

package/src/runtime/adapter/conversions.content.test.ts

@@ -1,9 +1,6 @@
 import { describe, expect, expectTypeOf, it } from "vitest";
-import {
-  MessageSchema,
-  UserMessageSchema,
-  type Message as AgUiWireMessage,
-} from "@ag-ui/client";
+import type { Message as AgUiWireMessage } from "@ag-ui/client";
+import { MessageSchema, UserMessageSchema } from "@ag-ui/core/schemas";
 import type { AppendMessage } from "@assistant-ui/core";
 import {
   fromAgUiMessages,
@@ -82,6 +79,41 @@ describe("toAgUiMessages data URLs", () => {
 });
 
 describe("toAgUiMessages content metadata", () => {
+  it.each(["image", "document"] as const)(
+    "preserves a %s file source and its metadata through a snapshot round trip",
+    (type) => {
+      const part = {
+        type,
+        source: {
+          type: "file",
+          value: `${type}-handle`,
+          provider: "openai",
+          mimeType: type === "image" ? "image/png" : "application/pdf",
+        },
+        metadata: { filename: `${type}.bin`, file_id: `${type}-1` },
+      };
+      const original = {
+        id: "u-file",
+        role: "user",
+        content: [
+          { type: "text", text: "label" },
+          {
+            type: "image",
+            source: { type: "url", value: "https://example.com/other.png" },
+          },
+          part,
+        ],
+        metadata: { trace: "kept" },
+      };
+
+      const imported = fromAgUiMessages([original]);
+      const exported = toAgUiMessages(imported);
+
+      expect(MessageSchema.parse(exported[0])).toEqual(original);
+      expect(exported).toEqual([original]);
+    },
+  );
+
   it("emits AgUiMessage values assignable to AG-UI Message", () => {
     expectTypeOf<AgUiMessage>().toExtend<AgUiWireMessage>();
   });

package/src/runtime/adapter/conversions.tools.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, expectTypeOf, it } from "vitest";
-import { ToolSchema, type RunAgentParameters } from "@ag-ui/client";
+import type { RunAgentParameters } from "@ag-ui/client";
+import { ToolSchema } from "@ag-ui/core/schemas";
 import { toAgUiTools } from "./conversions";
 
 describe("toAgUiTools", () => {

package/src/runtime/adapter/conversions.ts

@@ -1,6 +1,10 @@
 "use client";
 
-import type { InputContent, RunAgentParameters } from "@ag-ui/client";
+import type {
+  ContentPart,
+  PartSource,
+  RunAgentParameters,
+} from "@ag-ui/client";
 import { generateId } from "@assistant-ui/core";
 import type {
   ThreadMessageLike as CoreThreadMessageLike,
@@ -34,10 +38,9 @@ import { projectAgUiToolApprovals } from "./tool-approval";
 import {
   parseMcpToolCallResult,
   readMcpAppResourceUri,
+  toolContentText,
 } from "../mcp-tool-result";
 
-export type { InputContent };
-
 type AttachmentLike = {
   name?: string | undefined;
   contentType?: string | undefined;
@@ -67,7 +70,7 @@ export type AgUiMessage =
   | {
       id: string;
       role: "user";
-      content: string | InputContent[];
+      content: string | ContentPart[];
       name?: string;
       metadata?: Record<string, unknown>;
     }
@@ -96,7 +99,7 @@ export type AgUiMessage =
   | {
       id: string;
       role: "tool";
-      content: string;
+      content: string | ContentPart[];
       toolCallId: string;
       error?: string;
       metadata?: Record<string, unknown>;
@@ -137,14 +140,22 @@ function readLegacyAgUiNamespace(metadata: unknown): Record<string, unknown> {
     : {};
 }
 
-function withoutPersistedInterrupts(
+function withoutPersistedStatusMetadata(
   wire: Record<string, unknown>,
 ): Record<string, unknown> {
   if (!isObject(wire.custom)) return wire;
   const namespaced = wire.custom[AG_UI_METADATA_NAMESPACE];
-  if (!isObject(namespaced) || !("interrupts" in namespaced)) return wire;
-
-  const { interrupts: _interrupts, ...remainingNamespace } = namespaced;
+  if (
+    !isObject(namespaced) ||
+    (!("interrupts" in namespaced) && !("pendingToolCallIds" in namespaced))
+  )
+    return wire;
+
+  const {
+    interrupts: _interrupts,
+    pendingToolCallIds: _pendingToolCallIds,
+    ...remainingNamespace
+  } = namespaced;
   const { [AG_UI_METADATA_NAMESPACE]: _namespace, ...remainingCustom } =
     wire.custom;
   if (Object.keys(remainingNamespace).length > 0) {
@@ -160,12 +171,17 @@ function toSnapshotCustom(
   metadata: unknown,
   state: Record<string, unknown> = {},
 ): Record<string, unknown> {
-  const wire = withoutPersistedInterrupts(toAgUiWireMetadata(metadata) ?? {});
+  const wire = withoutPersistedStatusMetadata(
+    toAgUiWireMetadata(metadata) ?? {},
+  );
   const needsWireSlot =
     isObject(metadata) &&
     (Object.keys(metadata).length === 0 || AG_UI_METADATA_NAMESPACE in wire);
   const namespaced = {
     ...readLegacyAgUiNamespace(metadata),
+    ...(readPersistedPendingToolCallIds(metadata) !== undefined
+      ? { pendingToolCallIds: readPersistedPendingToolCallIds(metadata) }
+      : {}),
     ...state,
     ...(needsWireSlot ? { messageMetadata: wire } : {}),
   };
@@ -246,7 +262,7 @@ const readAgUiReasoningMeta = (providerMetadata: unknown) => {
 // AG-UI carries per-item extras in `metadata`; anything else on the item is
 // stripped by its schema. The part's own `filename` wins over a key of the
 // same name in the bag.
-const buildInputMetadata = (
+const buildPartMetadata = (
   part: Record<string, unknown>,
   filename?: string | undefined,
 ) => {
@@ -260,7 +276,7 @@ const buildInputMetadata = (
 // The inverse: `filename` lands on the part's own field, everything else the
 // item carried goes back into the namespace it came from, so a snapshot echo
 // resends what the host attached.
-const readInputMetadata = (
+const readPartMetadata = (
   metadata: unknown,
 ): {
   filename?: string | undefined;
@@ -326,9 +342,7 @@ function extractText(content: unknown): string {
     .join("\n");
 }
 
-type InputContentSource =
-  | { type: "data"; value: string; mimeType: string }
-  | { type: "url"; value: string; mimeType?: string };
+type EncodedPartSource = Extract<PartSource, { type: "data" | "url" }>;
 
 type MediaInputType = "image" | "audio" | "video" | "document";
 
@@ -345,11 +359,11 @@ function mediaTypeForMime(mimeType: string | undefined): MediaInputType {
 // `sourceType: "url"` on the part forces the url leg for non-http references,
 // and a data URL that is not base64 takes it too, since a `data` source
 // carries base64 bytes.
-function buildInputSource(
+function buildPartSource(
   value: string,
   declaredMimeType: string | undefined,
   sourceType?: string,
-): InputContentSource {
+): EncodedPartSource {
   const source = resolveFilePartSource({
     data: value,
     mimeType: declaredMimeType ?? "application/octet-stream",
@@ -368,10 +382,10 @@ function buildInputSource(
   };
 }
 
-function toInputContent(
+function toContentPart(
   part: unknown,
   fallbackMimeType: string | undefined,
-): InputContent | null {
+): ContentPart | null {
   if (!isObject(part)) return null;
   const type = getString(part, "type");
 
@@ -384,8 +398,8 @@ function toInputContent(
   if (type === "image") {
     const image = getString(part, "image");
     if (image === undefined) return null;
-    const metadata = buildInputMetadata(part, getString(part, "filename"));
-    const source = buildInputSource(image, fallbackMimeType);
+    const metadata = buildPartMetadata(part, getString(part, "filename"));
+    const source = buildPartSource(image, fallbackMimeType);
     return {
       type: "image",
       source:
@@ -404,12 +418,12 @@ function toInputContent(
     if (data === undefined) return null;
     const declaredMimeType = getString(part, "mimeType") || fallbackMimeType;
     const filename = getString(part, "filename");
-    const source = buildInputSource(
+    const source = buildPartSource(
       data,
       declaredMimeType,
       getString(part, "sourceType"),
     );
-    const metadata = buildInputMetadata(part, filename);
+    const metadata = buildPartMetadata(part, filename);
     switch (mediaTypeForMime(source.mimeType)) {
       case "image":
         return { type: "image", source, ...(metadata && { metadata }) };
@@ -430,7 +444,7 @@ function toInputContent(
     if (data === undefined || format === undefined) return null;
     return {
       type: "audio",
-      source: buildInputSource(
+      source: buildPartSource(
         parseDataUrl(data)?.data ?? data,
         `audio/${format}`,
       ),
@@ -446,8 +460,7 @@ type SnapshotAttachment = NonNullable<
 
 const mediaInputTypes = new Set(["image", "audio", "video", "document"]);
 
-// Inverse of buildInputSource.
-function inputSourceToString(
+function partSourceToString(
   value: unknown,
 ): { value: string; mimeType?: string; isUrl?: boolean } | null {
   if (!isObject(value)) return null;
@@ -494,23 +507,73 @@ function upgradeBinaryInputPart(
   };
 }
 
+function normalizeInputPart(
+  rawPart: Record<string, unknown>,
+): Record<string, unknown> | null {
+  return getString(rawPart, "type") === "binary"
+    ? upgradeBinaryInputPart(rawPart)
+    : rawPart;
+}
+
+type ContentSlot<T> = { index: number } & T;
+
+// Text imports as one string and encoded media as attachments, so export
+// needs the wire layout to restore the original order: each file source part
+// keeps its rank among media, and each text part its position among all kept
+// parts unless the text is one non-empty leading part.
+function toUserContentLayout(content: unknown): {
+  fileSourceParts: ContentSlot<{ part: Record<string, unknown> }>[];
+  textParts: ContentSlot<{ text: string }>[];
+} {
+  const fileSourceParts: ContentSlot<{ part: Record<string, unknown> }>[] = [];
+  const textParts: ContentSlot<{ text: string }>[] = [];
+  if (!Array.isArray(content)) return { fileSourceParts, textParts };
+
+  let mediaIndex = 0;
+  for (const rawPart of content) {
+    if (!isObject(rawPart)) continue;
+    if (rawPart.type === "text" && typeof rawPart.text === "string") {
+      textParts.push({
+        index: textParts.length + mediaIndex,
+        text: rawPart.text,
+      });
+      continue;
+    }
+    const part = normalizeInputPart(rawPart);
+    if (!part || !mediaInputTypes.has(getString(part, "type") ?? "")) continue;
+    if (
+      isObject(part.source) &&
+      part.source.type === "file" &&
+      typeof part.source.value === "string"
+    ) {
+      fileSourceParts.push({ index: mediaIndex++, part });
+    } else if (partSourceToString(part.source)) {
+      mediaIndex++;
+    }
+  }
+
+  const leadingText =
+    textParts.length === 0 ||
+    (textParts.length === 1 &&
+      textParts[0]!.index === 0 &&
+      textParts[0]!.text !== "");
+  return { fileSourceParts, textParts: leadingText ? [] : textParts };
+}
+
 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;
+    const part = normalizeInputPart(rawPart);
     if (!part) continue;
     const type = getString(part, "type");
     if (type === undefined || !mediaInputTypes.has(type)) continue;
-    const source = inputSourceToString(part.source);
+    const source = partSourceToString(part.source);
     if (!source) continue;
 
-    const { filename, providerMetadata } = readInputMetadata(part.metadata);
+    const { filename, providerMetadata } = readPartMetadata(part.metadata);
     const id = attachments.length.toString();
 
     if (type === "image") {
@@ -554,12 +617,26 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
   return attachments;
 }
 
-function buildUserContent(message: ThreadMessageLike): string | InputContent[] {
+function readContentSlots(
+  value: unknown,
+): { index: number; entry: Record<string, unknown> }[] {
+  if (!Array.isArray(value)) return [];
+  return value.flatMap((entry) =>
+    isObject(entry) &&
+    typeof entry.index === "number" &&
+    Number.isInteger(entry.index) &&
+    entry.index >= 0
+      ? [{ index: entry.index, entry }]
+      : [],
+  );
+}
+
+function buildUserContent(message: ThreadMessageLike): string | ContentPart[] {
   const contentParts = Array.isArray(message.content) ? message.content : [];
 
   const attachments = message.attachments ?? [];
 
-  const converted: InputContent[] = [];
+  const converted: ContentPart[] = [];
 
   // Promote string-form content to a leading text part so it survives when
   // non-text attachments are present (fromAgUiMessages emits string content).
@@ -568,7 +645,7 @@ function buildUserContent(message: ThreadMessageLike): string | InputContent[] {
   }
 
   for (const part of contentParts) {
-    const input = toInputContent(part, undefined);
+    const input = toContentPart(part, undefined);
     if (input) converted.push(input);
   }
   for (const attachment of attachments) {
@@ -577,13 +654,53 @@ function buildUserContent(message: ThreadMessageLike): string | InputContent[] {
     if (!Array.isArray(attachmentContent)) continue;
     const fallbackMime = getString(attachment, "contentType");
     for (const part of attachmentContent) {
-      const input = toInputContent(part, fallbackMime);
+      const input = toContentPart(part, fallbackMime);
       if (input) converted.push(input);
     }
   }
 
-  const hasNonText = converted.some((part) => part.type !== "text");
-  if (hasNonText) return converted;
+  const namespaced = readNamespacedMetadata(message.metadata);
+  const textSlots = readContentSlots(namespaced.textParts).flatMap(
+    ({ index, entry }) =>
+      typeof entry.text === "string"
+        ? [{ index, part: { type: "text" as const, text: entry.text } }]
+        : [],
+  );
+  const mergedText = converted
+    .flatMap((part) => (part.type === "text" ? [part.text] : []))
+    .join("\n");
+  const restoresText =
+    textSlots.length > 0 &&
+    textSlots.map(({ part }) => part.text).join("\n") === mergedText;
+  const fileSourceSlots = readContentSlots(namespaced.fileSourceParts).flatMap(
+    ({ index, entry: { part } }) =>
+      isObject(part) &&
+      mediaInputTypes.has(getString(part, "type") ?? "") &&
+      isObject(part.source) &&
+      part.source.type === "file" &&
+      typeof part.source.value === "string"
+        ? [{ index, part: part as ContentPart }]
+        : [],
+  );
+  const ordered = restoresText
+    ? converted.filter((part) => part.type !== "text")
+    : converted;
+  for (const { index, part } of fileSourceSlots.sort(
+    (a, b) => a.index - b.index,
+  )) {
+    const mediaPositions = ordered.flatMap((contentPart, position) =>
+      contentPart.type !== "text" ? [position] : [],
+    );
+    ordered.splice(mediaPositions[index] ?? ordered.length, 0, part);
+  }
+  if (restoresText) {
+    for (const { index, part } of textSlots.sort((a, b) => a.index - b.index)) {
+      ordered.splice(Math.min(index, ordered.length), 0, part);
+    }
+  }
+
+  if (restoresText || ordered.some((part) => part.type !== "text"))
+    return ordered;
 
   // All-text path: collapse to plain string. Join text parts collected from
   // both content and attachments so attachment-sourced text is not dropped.
@@ -695,6 +812,17 @@ function readPersistedInterrupts(
   return validateInterrupts(namespaced.interrupts);
 }
 
+function readPersistedPendingToolCallIds(
+  metadata: unknown,
+): string[] | undefined {
+  const value = readNamespacedMetadata(metadata).pendingToolCallIds;
+  return Array.isArray(value) &&
+    value.length > 0 &&
+    value.every((id): id is string => typeof id === "string")
+    ? value
+    : undefined;
+}
+
 function readOpaqueReasoning(metadata: unknown): AgUiOpaqueReasoning[] {
   if (!isObject(metadata)) return [];
   const custom = metadata.custom;
@@ -794,10 +922,15 @@ function toUserOrSystemSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const messageName = getString(rawMessage, "name");
-  const custom = toSnapshotCustom(
-    rawMessage.metadata,
-    role === "developer" ? { role: "developer" } : {},
-  );
+  const { fileSourceParts, textParts } =
+    role === "user"
+      ? toUserContentLayout(rawMessage.content)
+      : { fileSourceParts: [], textParts: [] };
+  const custom = toSnapshotCustom(rawMessage.metadata, {
+    ...(role === "developer" ? { role: "developer" } : {}),
+    ...(fileSourceParts.length > 0 ? { fileSourceParts } : {}),
+    ...(textParts.length > 0 ? { textParts } : {}),
+  });
   const attachments =
     role === "user" ? toSnapshotAttachments(rawMessage.content) : [];
   return {
@@ -1030,18 +1163,27 @@ export function fromAgUiMessages(
       flushPendingReasoning();
       const toolCallId = getToolCallId(rawMessage) ?? `tool-${generateId()}`;
       const toolMessageId = getString(rawMessage, "id");
-      const modelContent = extractText(rawMessage.content);
-      const mcpResult = parseMcpToolCallResult(rawMessage, modelContent);
-      const mcpModelContent = mcpResult
+      const modelContent = Array.isArray(rawMessage.content)
+        ? toolContentText(rawMessage.content as ContentPart[])
+        : extractText(rawMessage.content);
+      const mcpResult = parseMcpToolCallResult(
+        rawMessage,
+        Array.isArray(rawMessage.content)
+          ? (rawMessage.content as ContentPart[])
+          : modelContent,
+      );
+      const toolModelContent = mcpResult
         ? [{ type: "text" as const, text: modelContent }]
         : undefined;
       const result =
         mcpResult ??
-        (rawMessage.result !== undefined
-          ? rawMessage.result
-          : typeof rawMessage.content === "string"
-            ? parseJSONText(rawMessage.content)
-            : rawMessage.content);
+        (Array.isArray(rawMessage.content)
+          ? rawMessage.content
+          : rawMessage.result !== undefined
+            ? rawMessage.result
+            : typeof rawMessage.content === "string"
+              ? parseJSONText(rawMessage.content)
+              : rawMessage.content);
       const isError =
         typeof rawMessage.error === "string" ||
         rawMessage.isError === true ||
@@ -1078,10 +1220,14 @@ export function fromAgUiMessages(
           if (!isObject(part) || part.type !== "tool-call") continue;
           if (getString(part, "toolCallId") !== toolCallId) continue;
 
+          const { modelContent: _modelContent, ...partWithoutModelContent } =
+            part as ToolCallPart;
           const updatedPart: ToolCallPart = {
-            ...(part as ToolCallPart),
+            ...(Array.isArray(rawMessage.content) && !mcpResult
+              ? partWithoutModelContent
+              : (part as ToolCallPart)),
             result,
-            ...(mcpModelContent ? { modelContent: mcpModelContent } : {}),
+            ...(toolModelContent ? { modelContent: toolModelContent } : {}),
             ...(isError !== undefined ? { isError } : {}),
             ...(toolMessageId !== undefined
               ? { unstable_toolMessageId: toolMessageId }
@@ -1125,7 +1271,7 @@ export function fromAgUiMessages(
             args: {},
             argsText: "{}",
             result,
-            ...(mcpModelContent ? { modelContent: mcpModelContent } : {}),
+            ...(toolModelContent ? { modelContent: toolModelContent } : {}),
             ...(isError !== undefined ? { isError } : {}),
             ...(toolMessageId !== undefined
               ? { unstable_toolMessageId: toolMessageId }
@@ -1345,13 +1491,18 @@ export function fromAgUiMessages(
 
     const hasInterrupt =
       readPersistedInterrupts(message.metadata) !== undefined;
+    const pendingToolCallIds = readPersistedPendingToolCallIds(
+      message.metadata,
+    );
     const hasPendingToolCall =
       Array.isArray(message.content) &&
       message.content.some(
         (part) =>
           isObject(part) &&
           part.type === "tool-call" &&
-          part.result === undefined,
+          part.result === undefined &&
+          (pendingToolCallIds === undefined ||
+            pendingToolCallIds.includes(getString(part, "toolCallId") ?? "")),
       );
 
     if (hasInterrupt || hasPendingToolCall) {
@@ -1478,6 +1629,22 @@ function convertAssistantMessage(
   }
 }
 
+function isToolContentParts(value: unknown): value is ContentPart[] {
+  return (
+    Array.isArray(value) &&
+    value.length > 0 &&
+    value.every(
+      (part) =>
+        isObject(part) &&
+        (part.type === "text"
+          ? typeof part.text === "string"
+          : mediaInputTypes.has(getString(part, "type") ?? "") &&
+            isObject(part.source) &&
+            typeof part.source.value === "string"),
+    )
+  );
+}
+
 function emitToolResult(
   toolCallId: string,
   part: ToolCallPart,
@@ -1488,9 +1655,11 @@ function emitToolResult(
   const resultContent =
     part.modelContent !== undefined
       ? extractText(part.modelContent)
-      : typeof part.result === "string"
+      : isToolContentParts(part.result)
         ? part.result
-        : JSON.stringify(part.result);
+        : typeof part.result === "string"
+          ? part.result
+          : JSON.stringify(part.result);
   const messageMetadata = readAgUiPartMessageMetadata(part.providerMetadata);
 
   converted.push({
@@ -1498,7 +1667,13 @@ function emitToolResult(
     role: "tool",
     content: resultContent,
     toolCallId,
-    ...(part.isError ? { error: resultContent } : {}),
+    ...(part.isError
+      ? {
+          error: Array.isArray(resultContent)
+            ? toolContentText(resultContent)
+            : resultContent,
+        }
+      : {}),
     ...(messageMetadata ? { metadata: messageMetadata } : {}),
   });
 }
@@ -1507,7 +1682,9 @@ function convertToolMessage(
   message: NormalizedThreadMessageLike,
   converted: AgUiMessage[],
 ): void {
-  const content = extractText(message.content);
+  const content = isToolContentParts(message.content)
+    ? message.content
+    : extractText(message.content);
   const toolCallId = message.toolCallId ?? generateId();
 
   converted.push({