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

package/src/runtime/adapter/conversions.ts

@@ -69,6 +69,7 @@ export type AgUiMessage =
       role: "user";
       content: string | InputContent[];
       name?: string;
+      metadata?: Record<string, unknown>;
     }
   | {
       id: string;
@@ -76,18 +77,21 @@ export type AgUiMessage =
       content: string;
       name?: string;
       toolCalls?: AgUiToolCall[];
+      metadata?: Record<string, unknown>;
     }
   | {
       id: string;
       role: "system" | "developer";
       content: string;
       name?: string;
+      metadata?: Record<string, unknown>;
     }
   | {
       id: string;
       role: "reasoning";
       content: string;
       encryptedValue?: string;
+      metadata?: Record<string, unknown>;
     }
   | {
       id: string;
@@ -95,6 +99,7 @@ export type AgUiMessage =
       content: string;
       toolCallId: string;
       error?: string;
+      metadata?: Record<string, unknown>;
     };
 
 type CoreToolCallPart = Extract<
@@ -117,6 +122,108 @@ const getString = (record: Record<string, unknown>, key: string) => {
   return typeof value === "string" ? value : undefined;
 };
 
+function toAgUiWireMetadata(
+  metadata: unknown,
+): Record<string, unknown> | undefined {
+  return isObject(metadata) ? { ...metadata } : undefined;
+}
+
+function readLegacyAgUiNamespace(metadata: unknown): Record<string, unknown> {
+  if (!isObject(metadata) || !isObject(metadata.custom)) return {};
+  const namespaced = metadata.custom[AG_UI_METADATA_NAMESPACE];
+  if (!isObject(namespaced)) return {};
+  return "interrupts" in namespaced
+    ? { interrupts: namespaced.interrupts }
+    : {};
+}
+
+function withoutPersistedInterrupts(
+  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;
+  const { [AG_UI_METADATA_NAMESPACE]: _namespace, ...remainingCustom } =
+    wire.custom;
+  if (Object.keys(remainingNamespace).length > 0) {
+    remainingCustom[AG_UI_METADATA_NAMESPACE] = remainingNamespace;
+  }
+  const { custom: _custom, ...remainingWire } = wire;
+  return Object.keys(remainingCustom).length > 0
+    ? { ...remainingWire, custom: remainingCustom }
+    : remainingWire;
+}
+
+function toSnapshotCustom(
+  metadata: unknown,
+  state: Record<string, unknown> = {},
+): Record<string, unknown> {
+  const wire = withoutPersistedInterrupts(toAgUiWireMetadata(metadata) ?? {});
+  const needsWireSlot =
+    isObject(metadata) &&
+    (Object.keys(metadata).length === 0 || AG_UI_METADATA_NAMESPACE in wire);
+  const namespaced = {
+    ...readLegacyAgUiNamespace(metadata),
+    ...state,
+    ...(needsWireSlot ? { messageMetadata: wire } : {}),
+  };
+  return {
+    ...wire,
+    ...(Object.keys(namespaced).length > 0
+      ? { [AG_UI_METADATA_NAMESPACE]: namespaced }
+      : {}),
+  };
+}
+
+function toAgUiMessageMetadata(
+  metadata: unknown,
+): Record<string, unknown> | undefined {
+  const custom = readCustomMetadata(metadata);
+  const namespaced = readNamespacedMetadata(metadata);
+  if (isObject(namespaced.messageMetadata)) {
+    return { ...namespaced.messageMetadata };
+  }
+  const { [AG_UI_METADATA_NAMESPACE]: _state, ...wire } = custom;
+  return Object.keys(wire).length > 0 ? wire : undefined;
+}
+
+function withAgUiMessageMetadata(metadata: unknown): {
+  metadata?: Record<string, unknown>;
+} {
+  const messageMetadata = toAgUiMessageMetadata(metadata);
+  return messageMetadata !== undefined ? { metadata: messageMetadata } : {};
+}
+
+function withAgUiPartMessageMetadata(
+  providerMetadata: PartProviderMetadata | undefined,
+  messageMetadata: Record<string, unknown> | undefined,
+): PartProviderMetadata {
+  if (!messageMetadata) return providerMetadata ?? {};
+
+  const existing = providerMetadata ?? {};
+  const namespaced = isObject(existing[AG_UI_METADATA_NAMESPACE])
+    ? existing[AG_UI_METADATA_NAMESPACE]
+    : {};
+  return {
+    ...existing,
+    [AG_UI_METADATA_NAMESPACE]: {
+      ...namespaced,
+      messageMetadata: messageMetadata as ReadonlyJSONObject,
+    },
+  };
+}
+
+function readAgUiPartMessageMetadata(
+  providerMetadata: unknown,
+): Record<string, unknown> | undefined {
+  const namespaced = readAgUiNamespace(providerMetadata);
+  return namespaced && isObject(namespaced.messageMetadata)
+    ? toAgUiWireMetadata(namespaced.messageMetadata)
+    : undefined;
+}
+
 const getToolCallId = (record: Record<string, unknown>) =>
   getString(record, "toolCallId") ?? getString(record, "tool_call_id");
 
@@ -132,6 +239,7 @@ const readAgUiReasoningMeta = (providerMetadata: unknown) => {
   return {
     encryptedValue: getString(namespaced, "encryptedValue"),
     reasoningId: getString(namespaced, "reasoningId"),
+    messageMetadata: toAgUiWireMetadata(namespaced.messageMetadata),
   };
 };
 
@@ -600,8 +708,14 @@ function readOpaqueReasoning(metadata: unknown): AgUiOpaqueReasoning[] {
     const id = getString(entry, "id");
     const encryptedValue = getString(entry, "encryptedValue");
     if (!id?.trim() || !encryptedValue?.trim()) return [];
+    const metadata = toAgUiWireMetadata(entry.metadata);
     return [
-      { id, encryptedValue, ...(entry.after === true ? { after: true } : {}) },
+      {
+        id,
+        encryptedValue,
+        ...(metadata ? { metadata } : {}),
+        ...(entry.after === true ? { after: true } : {}),
+      },
     ];
   });
 }
@@ -631,7 +745,8 @@ function toAssistantSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const text = extractText(rawMessage.content);
-  const interrupts = readPersistedInterrupts(rawMessage.metadata);
+  const custom = toSnapshotCustom(rawMessage.metadata);
+  const interrupts = readPersistedInterrupts({ custom });
   const restoredToolCalls = extractAssistantToolCalls(rawMessage);
   const approvals = projectAgUiToolApprovals(
     interrupts,
@@ -652,21 +767,24 @@ function toAssistantSnapshotMessage(
     ...toolCallParts,
   ];
   const messageName = getString(rawMessage, "name");
+  const namespaced = isObject(custom[AG_UI_METADATA_NAMESPACE])
+    ? { ...custom[AG_UI_METADATA_NAMESPACE] }
+    : {};
+  if (interrupts) namespaced.interrupts = interrupts;
+  else delete namespaced.interrupts;
+  const messageCustom = { ...custom };
+  if (Object.keys(namespaced).length > 0) {
+    messageCustom[AG_UI_METADATA_NAMESPACE] = namespaced;
+  } else {
+    delete messageCustom[AG_UI_METADATA_NAMESPACE];
+  }
   return {
     id: getString(rawMessage, "id") ?? generateId(),
     role: "assistant",
     content: assistantContent.length > 0 ? assistantContent : "",
     ...(messageName !== undefined ? { name: messageName } : {}),
-    ...(interrupts
-      ? {
-          metadata: {
-            custom: {
-              [AG_UI_METADATA_NAMESPACE]: {
-                interrupts,
-              } satisfies AgUiCustomMetadata,
-            },
-          },
-        }
+    ...(Object.keys(messageCustom).length > 0
+      ? { metadata: { custom: messageCustom } }
       : {}),
   };
 }
@@ -676,6 +794,10 @@ function toUserOrSystemSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const messageName = getString(rawMessage, "name");
+  const custom = toSnapshotCustom(
+    rawMessage.metadata,
+    role === "developer" ? { role: "developer" } : {},
+  );
   const attachments =
     role === "user" ? toSnapshotAttachments(rawMessage.content) : [];
   return {
@@ -686,17 +808,7 @@ function toUserOrSystemSnapshotMessage(
     content: extractText(rawMessage.content),
     ...(messageName !== undefined ? { name: messageName } : {}),
     ...(attachments.length > 0 ? { attachments } : {}),
-    ...(role === "developer"
-      ? {
-          metadata: {
-            custom: {
-              [AG_UI_METADATA_NAMESPACE]: {
-                role: "developer",
-              } satisfies AgUiCustomMetadata,
-            },
-          },
-        }
-      : {}),
+    ...(Object.keys(custom).length > 0 ? { metadata: { custom } } : {}),
   };
 }
 
@@ -941,6 +1053,7 @@ export function fromAgUiMessages(
       const mcpAppUri = readMcpAppResourceUri(mcpResult?._meta);
       const mcpApp =
         mcpAppUri !== undefined ? { resourceUri: mcpAppUri } : undefined;
+      const messageMetadata = toAgUiWireMetadata(rawMessage.metadata);
 
       let updated = false;
       for (
@@ -974,6 +1087,14 @@ export function fromAgUiMessages(
               ? { unstable_toolMessageId: toolMessageId }
               : {}),
             ...(mcpApp ? { mcp: { app: mcpApp } } : {}),
+            ...(messageMetadata
+              ? {
+                  providerMetadata: withAgUiPartMessageMetadata(
+                    (part as ToolCallPart).providerMetadata,
+                    messageMetadata,
+                  ),
+                }
+              : {}),
           };
           const updatedContent = message.content.map((contentPart, index) =>
             index === partIndex ? updatedPart : contentPart,
@@ -1010,6 +1131,14 @@ export function fromAgUiMessages(
               ? { unstable_toolMessageId: toolMessageId }
               : {}),
             ...(mcpApp ? { mcp: { app: mcpApp } } : {}),
+            ...(messageMetadata
+              ? {
+                  providerMetadata: withAgUiPartMessageMetadata(
+                    undefined,
+                    messageMetadata,
+                  ),
+                }
+              : {}),
           },
         ],
       });
@@ -1108,6 +1237,7 @@ export function fromAgUiMessages(
     if (role === "reasoning") {
       const text = extractText(rawMessage.content);
       const encryptedValue = getString(rawMessage, "encryptedValue");
+      const messageMetadata = toAgUiWireMetadata(rawMessage.metadata);
       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
@@ -1120,6 +1250,7 @@ export function fromAgUiMessages(
           opaqueReasoning.push({
             id: opaqueId,
             encryptedValue,
+            ...(messageMetadata ? { metadata: messageMetadata } : {}),
             anchor: converted.length,
           });
         }
@@ -1135,6 +1266,7 @@ export function fromAgUiMessages(
           opaqueReasoning.push({
             id: hiddenId,
             encryptedValue,
+            ...(messageMetadata ? { metadata: messageMetadata } : {}),
             anchor: converted.length,
           });
         }
@@ -1149,6 +1281,9 @@ export function fromAgUiMessages(
       const meta = {
         ...(reasoningId !== undefined ? { reasoningId } : {}),
         ...(encryptedValue !== undefined ? { encryptedValue } : {}),
+        ...(messageMetadata
+          ? { messageMetadata: messageMetadata as ReadonlyJSONObject }
+          : {}),
       };
       pendingReasoning.push({
         id: reasoningId,
@@ -1265,9 +1400,8 @@ function convertAssistantMessage(
     if (!isObject(part) || part.type !== "reasoning") continue;
     const text = getString(part, "text") ?? "";
     if (text.trim().length === 0) continue;
-    const { encryptedValue, reasoningId } = readAgUiReasoningMeta(
-      part.providerMetadata,
-    );
+    const { encryptedValue, reasoningId, messageMetadata } =
+      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.
@@ -1279,12 +1413,12 @@ function convertAssistantMessage(
       role: "reasoning",
       content: text,
       ...(encryptedValue !== undefined ? { encryptedValue } : {}),
+      ...(messageMetadata ? { metadata: messageMetadata } : {}),
     });
     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.
+  // Some providers reject a replayed blank assistant turn, even when it carries metadata.
   if (shellIsDropped) {
     return;
   }
@@ -1319,6 +1453,7 @@ function convertAssistantMessage(
     role: "assistant",
     content,
     ...(message.name ? { name: message.name } : {}),
+    ...withAgUiMessageMetadata(message.metadata),
     ...(toolCalls.length + nestedToolCalls.length > 0
       ? {
           toolCalls: [...toolCalls, ...nestedToolCalls].map(
@@ -1356,6 +1491,7 @@ function emitToolResult(
       : typeof part.result === "string"
         ? part.result
         : JSON.stringify(part.result);
+  const messageMetadata = readAgUiPartMessageMetadata(part.providerMetadata);
 
   converted.push({
     id: part.unstable_toolMessageId ?? `${toolCallId}:tool`,
@@ -1363,6 +1499,7 @@ function emitToolResult(
     content: resultContent,
     toolCallId,
     ...(part.isError ? { error: resultContent } : {}),
+    ...(messageMetadata ? { metadata: messageMetadata } : {}),
   });
 }
 
@@ -1379,6 +1516,7 @@ function convertToolMessage(
     content,
     toolCallId,
     ...(typeof message.error === "string" ? { error: message.error } : {}),
+    ...withAgUiMessageMetadata(message.metadata),
   });
 }
 
@@ -1398,6 +1536,7 @@ export function toAgUiMessages(
       role: "reasoning",
       content: "",
       encryptedValue: entry.encryptedValue,
+      ...(entry.metadata ? { metadata: entry.metadata } : {}),
     });
     for (const entry of opaqueReasoning) {
       if (entry.after !== true) converted.push(toOpaqueRecord(entry));
@@ -1426,6 +1565,7 @@ export function toAgUiMessages(
         role: "user",
         content: buildUserContent(message),
         ...(message.name ? { name: message.name } : {}),
+        ...withAgUiMessageMetadata(message.metadata),
       });
       flushTrailingOpaqueReasoning();
       continue;
@@ -1449,6 +1589,7 @@ export function toAgUiMessages(
         role: wireRole,
         content: extractText(message.content),
         ...(message.name ? { name: message.name } : {}),
+        ...withAgUiMessageMetadata(message.metadata),
       });
       flushTrailingOpaqueReasoning();
       continue;
@@ -1459,6 +1600,7 @@ export function toAgUiMessages(
         id: message.id,
         role: "reasoning",
         content: extractText(message.content),
+        ...withAgUiMessageMetadata(message.metadata),
       });
       flushTrailingOpaqueReasoning();
       continue;

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

@@ -85,16 +85,16 @@ describe("RunAggregator tool approval projection", () => {
     ).toEqual([GATE]);
   });
 
-  it("drops the gate when the run is then cancelled, keeping the interrupts", () => {
+  it("keeps the gate when the finished run is then cancelled", () => {
     const { aggregator, last } = gatedAggregator();
 
     aggregator.handle({ type: "RUN_CANCELLED" } as any);
 
     expect(last().status).toMatchObject({
-      type: "incomplete",
-      reason: "cancelled",
+      type: "requires-action",
+      reason: "interrupt",
     });
-    expect(approvalOf(last())).toBeUndefined();
+    expect(approvalOf(last())).toBeDefined();
     expect(
       (last().metadata?.custom as Record<string, any> | undefined)?.agui
         ?.interrupts,

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

@@ -11,6 +11,7 @@ import {
   type ToolCallMessagePart,
   type ToolModelContentPart,
 } from "@assistant-ui/core";
+import type { JoinStrategy } from "@assistant-ui/core/react";
 import {
   applyA2uiOperations,
   surfaceToPresentToolCall,
@@ -176,6 +177,7 @@ export const isPlainObject = (
 
 export type RunAggregatorOptions = {
   showThinking: boolean;
+  joinStrategy?: JoinStrategy | undefined;
   logger: Logger;
   emit: Emit;
   onServerMessageId?: (messageId: string) => void;
@@ -195,6 +197,7 @@ export type RunAggregatorOptions = {
 export class RunAggregator {
   private readonly emitUpdate: Emit;
   private readonly showThinking: boolean;
+  private readonly joinStrategy: JoinStrategy | undefined;
   private readonly logger: Logger;
   private readonly onServerMessageId: ((messageId: string) => void) | undefined;
   private readonly onTextMessageStart:
@@ -259,6 +262,7 @@ export class RunAggregator {
   constructor(options: RunAggregatorOptions) {
     this.emitUpdate = options.emit;
     this.showThinking = options.showThinking;
+    this.joinStrategy = options.joinStrategy;
     this.logger = options.logger;
     this.onServerMessageId = options.onServerMessageId;
     this.onTextMessageStart = options.onTextMessageStart;
@@ -319,8 +323,14 @@ export class RunAggregator {
       }
       case "RUN_CANCELLED": {
         // A cancel can land after RUN_FINISHED while the response is still
-        // closing, and the answer it finished is already persisted as complete.
-        if (this.status?.type === "complete") break;
+        // closing: the finished answer stays complete, and an interrupt the run
+        // already delivered stays answerable.
+        if (
+          this.status?.type === "complete" ||
+          (this.status?.type === "requires-action" &&
+            this.status.reason === "interrupt")
+        )
+          break;
         this.status = { type: "incomplete", reason: "cancelled" };
         this.closeOpenSubagentRuns(this.status);
         this.emit();
@@ -845,6 +855,7 @@ export class RunAggregator {
 
   private beginDistinctTextMessage(messageId: string | undefined): void {
     if (
+      this.joinStrategy === "concat-content" ||
       !messageId ||
       this.lastTextMessageId === undefined ||
       this.lastTextMessageId === messageId ||

package/src/runtime/types.ts

@@ -6,11 +6,13 @@ import type {
   ExternalStoreThreadListAdapter,
   FeedbackAdapter,
   RealtimeVoiceAdapter,
+  RunConfig,
   SpeechSynthesisAdapter,
   ThreadHistoryAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
-import type { AbstractAgent } from "@ag-ui/client";
+import type { JoinStrategy } from "@assistant-ui/core/react";
+import type { AbstractAgent, Context, Message } from "@ag-ui/client";
 import type { Logger } from "./logger";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { McpToolCallResult } from "./mcp-tool-result";
@@ -22,6 +24,7 @@ export const A2UI_SURFACE_ACTIVITY_TYPE = "a2ui-surface";
 export type AgUiOpaqueReasoning = {
   id: string;
   encryptedValue: string;
+  metadata?: Record<string, unknown>;
   after?: boolean;
 };
 
@@ -80,6 +83,22 @@ export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
   agent: AbstractAgent;
   logger?: Partial<Logger>;
   showThinking?: boolean;
+  joinStrategy?: JoinStrategy;
+  /**
+   * Return context entries for each run. `messages` is the exact AG-UI
+   * transcript sent in the run input, after transcript selection and message
+   * conversion. On resume, it follows `resumeTranscript`: "full" sends the
+   * whole thread, while "appended" sends messages after the interrupted
+   * assistant message, falling back to the whole thread if that message is
+   * missing. `runConfig` is the run config passed while building that input.
+   * If the getter throws, the error is reported and the run continues without host entries.
+   */
+  getRunContext?:
+    | ((input: {
+        messages: readonly Message[];
+        runConfig: RunConfig | undefined;
+      }) => readonly Context[])
+    | undefined;
   /**
    * What `messages` carries on a resume run, meaning a `RunAgentInput` that
    * also carries `resume`. The AG-UI interrupt spec leaves this undefined: its