@assistant-ui/react-ag-ui 0.0.53 → 0.0.54

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

@@ -0,0 +1,220 @@
+import { describe, expect, it } from "vitest";
+import { UserMessageSchema } from "@ag-ui/client";
+import { fromAgUiMessages, toAgUiMessages } from "./conversions";
+
+type Message = Parameters<typeof toAgUiMessages>[0][number];
+
+const userMessage = (
+  content: unknown,
+  attachments?: Message["attachments"],
+): Message =>
+  ({
+    id: "u1",
+    role: "user",
+    content,
+    ...(attachments && { attachments }),
+  }) as Message;
+
+const contentOf = (message: Message) => {
+  const [converted] = toAgUiMessages([message]);
+  return (converted as { content: unknown }).content;
+};
+
+describe("toAgUiMessages content metadata", () => {
+  it("carries a part's agui provider metadata onto an image item", () => {
+    expect(
+      contentOf(
+        userMessage([
+          {
+            type: "image",
+            image: "https://example.com/cat.png",
+            providerMetadata: { agui: { file_id: "f_1" } },
+          },
+        ]),
+      ),
+    ).toEqual([
+      {
+        type: "image",
+        source: { type: "url", value: "https://example.com/cat.png" },
+        metadata: { file_id: "f_1" },
+      },
+    ]);
+  });
+
+  it("preserves a direct image part's filename through a snapshot round trip", () => {
+    const sent = toAgUiMessages([
+      userMessage([
+        {
+          type: "image",
+          image: "https://example.com/cat.png",
+          filename: "cat.png",
+        },
+      ]),
+    ]);
+
+    expect((sent[0] as { content: unknown }).content).toEqual([
+      {
+        type: "image",
+        source: { type: "url", value: "https://example.com/cat.png" },
+        metadata: { filename: "cat.png" },
+      },
+    ]);
+
+    const rebuilt = fromAgUiMessages(sent as never);
+    const attachment = (rebuilt[0] as unknown as { attachments: unknown[] })
+      .attachments[0];
+
+    expect(attachment).toMatchObject({
+      name: "cat.png",
+      content: [{ type: "image", filename: "cat.png" }],
+    });
+  });
+
+  it("merges a file part's metadata with its filename", () => {
+    expect(
+      contentOf(
+        userMessage([
+          {
+            type: "file",
+            data: "https://example.com/spec.pdf",
+            mimeType: "application/pdf",
+            filename: "spec.pdf",
+            providerMetadata: { agui: { file_id: "f_2" } },
+          },
+        ]),
+      ),
+    ).toEqual([
+      {
+        type: "document",
+        source: {
+          type: "url",
+          value: "https://example.com/spec.pdf",
+          mimeType: "application/pdf",
+        },
+        metadata: { file_id: "f_2", filename: "spec.pdf" },
+      },
+    ]);
+  });
+
+  it("keeps the part's own filename over one in the metadata", () => {
+    expect(
+      contentOf(
+        userMessage([
+          {
+            type: "file",
+            data: "https://example.com/spec.pdf",
+            mimeType: "application/pdf",
+            filename: "spec.pdf",
+            providerMetadata: { agui: { filename: "renamed.pdf" } },
+          },
+        ]),
+      ),
+    ).toMatchObject([{ metadata: { filename: "spec.pdf" } }]);
+  });
+
+  it("carries metadata from an attachment's content part", () => {
+    expect(
+      contentOf(
+        userMessage([], [
+          {
+            name: "spec.pdf",
+            contentType: "application/pdf",
+            content: [
+              {
+                type: "file",
+                data: "https://example.com/spec.pdf",
+                mimeType: "application/pdf",
+                providerMetadata: { agui: { file_id: "f_3" } },
+              },
+            ],
+          },
+        ] as Message["attachments"]),
+      ),
+    ).toMatchObject([{ type: "document", metadata: { file_id: "f_3" } }]);
+  });
+
+  it("ignores metadata namespaced for another provider", () => {
+    expect(
+      contentOf(
+        userMessage([
+          {
+            type: "image",
+            image: "https://example.com/cat.png",
+            providerMetadata: { openai: { file_id: "f_4" } },
+          },
+        ]),
+      ),
+    ).toEqual([
+      {
+        type: "image",
+        source: { type: "url", value: "https://example.com/cat.png" },
+      },
+    ]);
+  });
+
+  it("emits no metadata channel for text, which the schema has no field for", () => {
+    expect(
+      contentOf(
+        userMessage([
+          {
+            type: "text",
+            text: "hello",
+            providerMetadata: { agui: { file_id: "f_5" } },
+          },
+          { type: "image", image: "https://example.com/cat.png" },
+        ]),
+      ),
+    ).toEqual([
+      { type: "text", text: "hello" },
+      {
+        type: "image",
+        source: { type: "url", value: "https://example.com/cat.png" },
+      },
+    ]);
+  });
+
+  it("stays valid against the upstream user message schema", () => {
+    const [converted] = toAgUiMessages([
+      userMessage([
+        {
+          type: "file",
+          data: "https://example.com/spec.pdf",
+          mimeType: "application/pdf",
+          filename: "spec.pdf",
+          providerMetadata: { agui: { file_id: "f_6", nested: { a: 1 } } },
+        },
+      ]),
+    ]);
+
+    expect(
+      UserMessageSchema.safeParse(JSON.parse(JSON.stringify(converted)))
+        .success,
+    ).toBe(true);
+  });
+
+  it("survives a snapshot round trip", () => {
+    const sent = toAgUiMessages([
+      userMessage([
+        {
+          type: "file",
+          data: "https://example.com/spec.pdf",
+          mimeType: "application/pdf",
+          filename: "spec.pdf",
+          providerMetadata: { agui: { file_id: "f_7" } },
+        },
+      ]),
+    ]);
+
+    const rebuilt = fromAgUiMessages(sent as never);
+    const attachment = (rebuilt[0] as unknown as { attachments: unknown[] })
+      .attachments[0] as { content: unknown[] };
+
+    expect(attachment.content[0]).toMatchObject({
+      type: "file",
+      filename: "spec.pdf",
+      providerMetadata: { agui: { file_id: "f_7" } },
+    });
+
+    expect(toAgUiMessages(rebuilt as never)).toEqual(sent);
+  });
+});

package/src/runtime/adapter/conversions.ts

@@ -3,6 +3,7 @@
 import type { InputContent, RunAgentParameters } from "@ag-ui/client";
 import type {
   ThreadMessageLike as CoreThreadMessageLike,
+  PartProviderMetadata,
   ToolCallMessagePartMcpMetadata,
   ToolModelContentPart,
 } from "@assistant-ui/core";
@@ -17,6 +18,7 @@ import {
   AG_UI_METADATA_NAMESPACE,
   A2UI_SURFACE_ACTIVITY_TYPE,
   type AgUiCustomMetadata,
+  type AgUiOpaqueReasoning,
 } from "./run-aggregator";
 import {
   applyA2uiOperations,
@@ -42,6 +44,7 @@ type ThreadMessageLike = {
   id: string;
   role: string;
   content: unknown;
+  metadata?: unknown;
   name?: string;
   toolCallId?: string;
   error?: string;
@@ -62,6 +65,12 @@ export type AgUiMessage =
       name?: string;
       toolCalls?: AgUiToolCall[];
     }
+  | {
+      id: string;
+      role: "reasoning";
+      content: string;
+      encryptedValue?: string;
+    }
   | {
       id: string;
       role: "tool";
@@ -94,6 +103,56 @@ const getString = (record: Record<string, unknown>, key: string) => {
 const getToolCallId = (record: Record<string, unknown>) =>
   getString(record, "toolCallId") ?? getString(record, "tool_call_id");
 
+const readAgUiNamespace = (providerMetadata: unknown) => {
+  if (!isObject(providerMetadata)) return undefined;
+  const namespaced = providerMetadata[AG_UI_METADATA_NAMESPACE];
+  return isObject(namespaced) ? namespaced : undefined;
+};
+
+const readAgUiReasoningMeta = (providerMetadata: unknown) => {
+  const namespaced = readAgUiNamespace(providerMetadata);
+  if (!namespaced) return {};
+  return {
+    encryptedValue: getString(namespaced, "encryptedValue"),
+    reasoningId: getString(namespaced, "reasoningId"),
+  };
+};
+
+// 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 = (
+  part: Record<string, unknown>,
+  filename?: string | undefined,
+) => {
+  const metadata = {
+    ...readAgUiNamespace(part.providerMetadata),
+    ...(filename !== undefined && { filename }),
+  };
+  return Object.keys(metadata).length > 0 ? metadata : undefined;
+};
+
+// 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 = (
+  metadata: unknown,
+): {
+  filename?: string | undefined;
+  providerMetadata?: PartProviderMetadata | undefined;
+} => {
+  if (!isObject(metadata)) return {};
+  const { filename: _filename, ...rest } = metadata;
+  return {
+    filename: getString(metadata, "filename"),
+    ...(Object.keys(rest).length > 0 && {
+      providerMetadata: {
+        [AG_UI_METADATA_NAMESPACE]: rest as PartProviderMetadata[string],
+      },
+    }),
+  };
+};
+
 function parseJSONText(value: string): unknown {
   if (!value) return value;
   try {
@@ -198,7 +257,12 @@ function toInputContent(
   if (type === "image") {
     const image = getString(part, "image");
     if (image === undefined) return null;
-    return { type: "image", source: buildInputSource(image, fallbackMimeType) };
+    const metadata = buildInputMetadata(part, getString(part, "filename"));
+    return {
+      type: "image",
+      source: buildInputSource(image, fallbackMimeType),
+      ...(metadata && { metadata }),
+    };
   }
 
   if (type === "file") {
@@ -211,7 +275,7 @@ function toInputContent(
       declaredMimeType,
       getString(part, "sourceType"),
     );
-    const metadata = filename !== undefined ? { filename } : undefined;
+    const metadata = buildInputMetadata(part, filename);
     switch (mediaTypeForMime(source.mimeType)) {
       case "image":
         return { type: "image", source, ...(metadata && { metadata }) };
@@ -312,9 +376,7 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
     const source = inputSourceToString(part.source);
     if (!source) continue;
 
-    const filename = isObject(part.metadata)
-      ? getString(part.metadata, "filename")
-      : undefined;
+    const { filename, providerMetadata } = readInputMetadata(part.metadata);
     const id = attachments.length.toString();
 
     if (type === "image") {
@@ -329,6 +391,7 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
             type: "image",
             image: source.value,
             ...(filename !== undefined && { filename }),
+            ...(providerMetadata && { providerMetadata }),
           },
         ],
       });
@@ -349,6 +412,7 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
           mimeType,
           ...(source.isUrl && { sourceType: "url" as const }),
           ...(filename !== undefined && { filename }),
+          ...(providerMetadata && { providerMetadata }),
         },
       ],
     });
@@ -498,6 +562,46 @@ function readPersistedInterrupts(
   return validateInterrupts(namespaced.interrupts);
 }
 
+function readOpaqueReasoning(metadata: unknown): AgUiOpaqueReasoning[] {
+  if (!isObject(metadata)) return [];
+  const custom = metadata.custom;
+  if (!isObject(custom)) return [];
+  const namespaced = custom[AG_UI_METADATA_NAMESPACE];
+  if (!isObject(namespaced)) return [];
+  const entries = namespaced.opaqueReasoning;
+  if (!Array.isArray(entries)) return [];
+  return entries.flatMap((entry) => {
+    if (!isObject(entry)) return [];
+    const id = getString(entry, "id");
+    const encryptedValue = getString(entry, "encryptedValue");
+    if (!id?.trim() || !encryptedValue?.trim()) return [];
+    return [
+      { id, encryptedValue, ...(entry.after === true ? { after: true } : {}) },
+    ];
+  });
+}
+
+function withOpaqueReasoning(
+  message: CoreThreadMessageLike,
+  opaqueReasoning: AgUiOpaqueReasoning[],
+): CoreThreadMessageLike {
+  const metadata = isObject(message.metadata) ? message.metadata : {};
+  const custom = isObject(metadata.custom) ? metadata.custom : {};
+  const namespaced = isObject(custom[AG_UI_METADATA_NAMESPACE])
+    ? (custom[AG_UI_METADATA_NAMESPACE] as Record<string, unknown>)
+    : {};
+  return {
+    ...message,
+    metadata: {
+      ...metadata,
+      custom: {
+        ...custom,
+        [AG_UI_METADATA_NAMESPACE]: { ...namespaced, opaqueReasoning },
+      },
+    },
+  } as CoreThreadMessageLike;
+}
+
 function toAssistantSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
@@ -596,6 +700,11 @@ export function fromAgUiMessages(
   const converted: CoreThreadMessageLike[] = [];
   const a2uiBuckets = new Map<string, A2uiState>();
   const a2uiBucketOwnerIndices = new Map<string, number>();
+  // A zero-data-retention run carries its payload in encryptedValue with no
+  // readable content, so the record has no part to become and rides on the
+  // neighbouring message instead of being dropped. The anchor is the index the
+  // next pushed message will occupy, which is where it sat on the wire.
+  const opaqueReasoning: (AgUiOpaqueReasoning & { anchor: number })[] = [];
 
   for (const rawMessage of messages) {
     if (!isObject(rawMessage)) continue;
@@ -750,15 +859,53 @@ export function fromAgUiMessages(
     }
 
     if (role === "reasoning") {
-      // Gate on showThinking so a cold reload matches the live run: the
-      // aggregator never stores reasoning parts when showThinking is false.
-      if (!showThinking) continue;
       const text = extractText(rawMessage.content);
-      if (text.trim().length === 0) continue;
+      const encryptedValue = getString(rawMessage, "encryptedValue");
+      if (text.trim().length === 0) {
+        // showThinking hides reasoning; this record is never rendered anyway,
+        // and dropping it would deny the provider the payload it needs to
+        // accept the next run.
+        const opaqueId = getString(rawMessage, "id");
+        if (opaqueId?.trim() && encryptedValue?.trim()) {
+          opaqueReasoning.push({
+            id: opaqueId,
+            encryptedValue,
+            anchor: converted.length,
+          });
+        }
+        continue;
+      }
+      // Gate visible reasoning on showThinking so a cold reload matches the
+      // live run: the aggregator never stores reasoning parts when it is off.
+      // A signature on a hidden record is still transport state, so it is kept
+      // opaque rather than discarded with the text.
+      if (!showThinking) {
+        const hiddenId = getString(rawMessage, "id");
+        if (hiddenId?.trim() && encryptedValue?.trim()) {
+          opaqueReasoning.push({
+            id: hiddenId,
+            encryptedValue,
+            anchor: converted.length,
+          });
+        }
+        continue;
+      }
       converted.push({
         id: getString(rawMessage, "id") ?? generateId(),
         role: "assistant",
-        content: [{ type: "reasoning", text }],
+        content: [
+          {
+            type: "reasoning",
+            text,
+            ...(encryptedValue !== undefined
+              ? {
+                  providerMetadata: {
+                    [AG_UI_METADATA_NAMESPACE]: { encryptedValue },
+                  },
+                }
+              : {}),
+          },
+        ],
       });
       continue;
     }
@@ -768,6 +915,18 @@ export function fromAgUiMessages(
     }
   }
 
+  for (const { anchor, ...entry } of opaqueReasoning) {
+    // Nothing followed it on the wire, so it trails the last message instead.
+    const trailing = anchor >= converted.length;
+    const index = trailing ? converted.length - 1 : anchor;
+    const target = converted[index];
+    if (!target) continue;
+    converted[index] = withOpaqueReasoning(target, [
+      ...readOpaqueReasoning(target.metadata),
+      ...(trailing ? [{ ...entry, after: true }] : [entry]),
+    ]);
+  }
+
   for (let i = 0; i < converted.length; i++) {
     const message = converted[i]!;
     if (message.role !== "assistant") continue;
@@ -821,9 +980,35 @@ function convertAssistantMessage(
     part,
   }));
 
+  // An AG-UI assistant record has no reasoning field, so reasoning leaves as
+  // the standalone record it arrived as, ahead of the assistant it belongs to.
+  const shellIsDropped = content.length === 0 && toolCalls.length === 0;
+  let reasoningIndex = 0;
+  for (const part of contentArray) {
+    if (!isObject(part) || part.type !== "reasoning") continue;
+    const text = getString(part, "text") ?? "";
+    if (text.trim().length === 0) continue;
+    const { encryptedValue, reasoningId } = readAgUiReasoningMeta(
+      part.providerMetadata,
+    );
+    converted.push({
+      // The wire id is what a signature was issued against, so it wins over a
+      // synthesized one whenever the run carried it.
+      id: reasoningId?.trim()
+        ? reasoningId
+        : shellIsDropped && reasoningIndex === 0
+          ? message.id
+          : `${message.id}:reasoning-${reasoningIndex}`,
+      role: "reasoning",
+      content: text,
+      ...(encryptedValue !== undefined ? { encryptedValue } : {}),
+    });
+    reasoningIndex++;
+  }
+
   // Drop assistant messages with no text or tool calls (e.g. an imported
   // reasoning-only entry) so they are not re-sent as a blank assistant turn.
-  if (content.length === 0 && toolCalls.length === 0) {
+  if (shellIsDropped) {
     return;
   }
 
@@ -888,13 +1073,31 @@ export function toAgUiMessages(
   const converted: AgUiMessage[] = [];
 
   for (const message of messages) {
+    const opaqueReasoning = readOpaqueReasoning(message.metadata);
+    const toOpaqueRecord = (entry: AgUiOpaqueReasoning): AgUiMessage => ({
+      id: entry.id,
+      role: "reasoning",
+      content: "",
+      encryptedValue: entry.encryptedValue,
+    });
+    for (const entry of opaqueReasoning) {
+      if (entry.after !== true) converted.push(toOpaqueRecord(entry));
+    }
+    const flushTrailingOpaqueReasoning = () => {
+      for (const entry of opaqueReasoning) {
+        if (entry.after === true) converted.push(toOpaqueRecord(entry));
+      }
+    };
+
     if (message.role === "assistant") {
       convertAssistantMessage(message, converted);
+      flushTrailingOpaqueReasoning();
       continue;
     }
 
     if (message.role === "tool") {
       convertToolMessage(message, converted);
+      flushTrailingOpaqueReasoning();
       continue;
     }
 
@@ -910,6 +1113,7 @@ export function toAgUiMessages(
       genericMessage.name = message.name;
     }
     converted.push(genericMessage);
+    flushTrailingOpaqueReasoning();
   }
 
   return converted;

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

@@ -20,8 +20,15 @@ import type { Logger } from "../logger";
 
 export const AG_UI_METADATA_NAMESPACE = "agui";
 
+export type AgUiOpaqueReasoning = {
+  id: string;
+  encryptedValue: string;
+  after?: boolean;
+};
+
 export type AgUiCustomMetadata = {
   interrupts?: AgUiInterrupt[];
+  opaqueReasoning?: AgUiOpaqueReasoning[];
 };
 
 type Emit = (update: ChatModelRunResult) => void;
@@ -77,6 +84,9 @@ export class RunAggregator {
   >();
   private activeTextMessageId: string | undefined;
   private readonly reasoningParts = new Map<string, string>(); // key → buffer
+  private readonly reasoningSignatures = new Map<string, string>();
+  private readonly reasoningMessageIds = new Map<string, string>();
+  private readonly anonymousReasoningKeys = new Set<string>();
   private activeReasoningKey: string | undefined;
   private reasoningPartCounter = 0;
   private readonly toolCalls = new Map<string, ToolCallState>();
@@ -111,6 +121,9 @@ export class RunAggregator {
       case "RUN_STARTED": {
         this.clearTextParts();
         this.reasoningParts.clear();
+        this.reasoningSignatures.clear();
+        this.reasoningMessageIds.clear();
+        this.anonymousReasoningKeys.clear();
         this.activeReasoningKey = undefined;
         this.reasoningPartCounter = 0;
         this.toolCalls.clear();
@@ -200,8 +213,25 @@ export class RunAggregator {
       case "REASONING_MESSAGE_START":
         this.handleReasoningStart(
           "messageId" in event ? event.messageId : undefined,
+          event.type === "REASONING_MESSAGE_START",
         );
         break;
+      case "REASONING_ENCRYPTED_VALUE":
+        if (event.subtype === "message") {
+          // entityId names any message, not necessarily a reasoning one, so an
+          // unmatched id may only claim a block that has no id to contradict it.
+          const active = this.activeReasoningKey;
+          const key = this.reasoningParts.has(event.entityId)
+            ? event.entityId
+            : active !== undefined && this.anonymousReasoningKeys.has(active)
+              ? active
+              : undefined;
+          if (key !== undefined) {
+            this.reasoningSignatures.set(key, event.encryptedValue);
+            this.emit();
+          }
+        }
+        break;
       case "THINKING_TEXT_MESSAGE_CONTENT":
         this.handleReasoningContent(event.delta);
         this.totalChunks++;
@@ -211,6 +241,7 @@ export class RunAggregator {
         this.handleReasoningContent(
           event.delta,
           "messageId" in event ? event.messageId : undefined,
+          true,
         );
         this.totalChunks++;
         this.recordFirstToken();
@@ -589,7 +620,19 @@ export class RunAggregator {
         if (this.showThinking) {
           const buffer = this.reasoningParts.get(part.key) ?? "";
           if (buffer.length > 0 || this.activeReasoningKey === part.key) {
-            snapshot.push({ type: "reasoning", text: buffer } as const);
+            const encryptedValue = this.reasoningSignatures.get(part.key);
+            const reasoningId = this.reasoningMessageIds.get(part.key);
+            const meta = {
+              ...(reasoningId !== undefined ? { reasoningId } : {}),
+              ...(encryptedValue !== undefined ? { encryptedValue } : {}),
+            };
+            snapshot.push({
+              type: "reasoning",
+              text: buffer,
+              ...(Object.keys(meta).length > 0
+                ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } }
+                : {}),
+            } as const);
           }
         }
         continue;
@@ -698,12 +741,20 @@ export class RunAggregator {
     };
   }
 
-  private handleReasoningStart(messageId?: string): void {
+  private handleReasoningStart(messageId?: string, isMessageId = false): void {
     if (!this.showThinking) return;
     // A reasoning block acts as a boundary: anonymous text arriving after it
     // should be a new part, not appended to any pre-reasoning text.
     this.activeTextMessageId = undefined;
     const key = messageId ?? `__auto-reasoning-${++this.reasoningPartCounter}`;
+    // Two different questions: which id may be replayed as a ReasoningMessage.id
+    // (only the message-scoped aliases carry one), and which block an unmatched
+    // signature may claim (any block opened without an id at all).
+    if (messageId === undefined) {
+      this.anonymousReasoningKeys.add(key);
+    } else if (isMessageId) {
+      this.reasoningMessageIds.set(key, messageId);
+    }
     if (!this.reasoningParts.has(key)) {
       this.reasoningParts.set(key, "");
       this.partOrder.push({ kind: "reasoning", key });
@@ -712,11 +763,15 @@ export class RunAggregator {
     this.emit();
   }
 
-  private handleReasoningContent(delta: string, messageId?: string): void {
+  private handleReasoningContent(
+    delta: string,
+    messageId?: string,
+    isMessageId = false,
+  ): void {
     if (!this.showThinking || !delta) return;
     if (!this.activeReasoningKey) {
       // Content arrived without a preceding START — create the slot lazily.
-      this.handleReasoningStart(messageId);
+      this.handleReasoningStart(messageId, isMessageId);
     }
     const key = this.activeReasoningKey;
     if (!key) return;