@assistant-ui/react-ag-ui 0.1.1 → 0.1.2

package/src/runtime/adapter/conversions.ts

@@ -5,8 +5,10 @@ import type {
   PartSource,
   RunAgentParameters,
 } from "@ag-ui/client";
-import { generateId } from "@assistant-ui/core";
+import { fromThreadMessageLike, generateId } from "@assistant-ui/core";
 import type {
+  MessageStatus,
+  ThreadMessage,
   ThreadMessageLike as CoreThreadMessageLike,
   PartProviderMetadata,
   ReasoningMessagePart,
@@ -26,6 +28,7 @@ import {
   MCP_APPS_ACTIVITY_TYPE,
   type AgUiCustomMetadata,
   type AgUiOpaqueReasoning,
+  MAX_SUBAGENT_DEPTH,
 } from "../types";
 import {
   applyA2uiOperations,
@@ -212,7 +215,7 @@ function withAgUiMessageMetadata(metadata: unknown): {
   return messageMetadata !== undefined ? { metadata: messageMetadata } : {};
 }
 
-function withAgUiPartMessageMetadata(
+export function withAgUiPartMessageMetadata(
   providerMetadata: PartProviderMetadata | undefined,
   messageMetadata: Record<string, unknown> | undefined,
 ): PartProviderMetadata {
@@ -222,11 +225,17 @@ function withAgUiPartMessageMetadata(
   const namespaced = isObject(existing[AG_UI_METADATA_NAMESPACE])
     ? existing[AG_UI_METADATA_NAMESPACE]
     : {};
+  const previousMessageMetadata = isObject(namespaced.messageMetadata)
+    ? namespaced.messageMetadata
+    : {};
   return {
     ...existing,
     [AG_UI_METADATA_NAMESPACE]: {
       ...namespaced,
-      messageMetadata: messageMetadata as ReadonlyJSONObject,
+      messageMetadata: {
+        ...previousMessageMetadata,
+        ...messageMetadata,
+      } as ReadonlyJSONObject,
     },
   };
 }
@@ -342,8 +351,6 @@ function extractText(content: unknown): string {
     .join("\n");
 }
 
-type EncodedPartSource = Extract<PartSource, { type: "data" | "url" }>;
-
 type MediaInputType = "image" | "audio" | "video" | "document";
 
 function mediaTypeForMime(mimeType: string | undefined): MediaInputType {
@@ -363,7 +370,14 @@ function buildPartSource(
   value: string,
   declaredMimeType: string | undefined,
   sourceType?: string,
-): EncodedPartSource {
+): PartSource {
+  if (sourceType === "id") {
+    return {
+      type: "file",
+      value,
+      ...(declaredMimeType !== undefined && { mimeType: declaredMimeType }),
+    };
+  }
   const source = resolveFilePartSource({
     data: value,
     mimeType: declaredMimeType ?? "application/octet-stream",
@@ -460,9 +474,12 @@ type SnapshotAttachment = NonNullable<
 
 const mediaInputTypes = new Set(["image", "audio", "video", "document"]);
 
-function partSourceToString(
-  value: unknown,
-): { value: string; mimeType?: string; isUrl?: boolean } | null {
+function partSourceToString(value: unknown): {
+  value: string;
+  mimeType?: string;
+  isUrl?: boolean;
+  isId?: boolean;
+} | null {
   if (!isObject(value)) return null;
   const sourceValue = getString(value, "value");
   if (sourceValue === undefined) return null;
@@ -475,6 +492,13 @@ function partSourceToString(
       ...(mimeType !== undefined && { mimeType }),
     };
   }
+  if (type === "file") {
+    return {
+      value: sourceValue,
+      isId: true,
+      ...(mimeType !== undefined && { mimeType }),
+    };
+  }
   if (type === "data") {
     const resolvedMimeType = mimeType ?? "application/octet-stream";
     return {
@@ -541,15 +565,11 @@ function toUserContentLayout(content: unknown): {
     }
     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++;
+    if (!partSourceToString(part.source)) continue;
+    if (isObject(part.source) && part.source.type === "file") {
+      fileSourceParts.push({ index: mediaIndex, part });
     }
+    mediaIndex++;
   }
 
   const leadingText =
@@ -576,7 +596,7 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
     const { filename, providerMetadata } = readPartMetadata(part.metadata);
     const id = attachments.length.toString();
 
-    if (type === "image") {
+    if (type === "image" && !source.isId) {
       attachments.push({
         id,
         type: "image",
@@ -608,6 +628,7 @@ function toSnapshotAttachments(content: unknown): SnapshotAttachment[] {
           data: source.value,
           mimeType,
           ...(source.isUrl && { sourceType: "url" as const }),
+          ...(source.isId && { sourceType: "id" as const }),
           ...(filename !== undefined && { filename }),
           ...(providerMetadata && { providerMetadata }),
         },
@@ -679,15 +700,28 @@ function buildUserContent(message: ThreadMessageLike): string | ContentPart[] {
       isObject(part.source) &&
       part.source.type === "file" &&
       typeof part.source.value === "string"
-        ? [{ index, part: part as ContentPart }]
+        ? [{ index, part: part as ContentPart, value: part.source.value }]
         : [],
   );
   const ordered = restoresText
     ? converted.filter((part) => part.type !== "text")
     : converted;
-  for (const { index, part } of fileSourceSlots.sort(
+  for (const { index, part, value } of fileSourceSlots.sort(
     (a, b) => a.index - b.index,
   )) {
+    // A snapshot also imports the handle as an attachment; the slot replaces
+    // it in place to restore fields the attachment cannot hold.
+    const attached = ordered.findIndex(
+      (contentPart) =>
+        contentPart.type !== "text" &&
+        contentPart.source.type === "file" &&
+        contentPart.source.value === value &&
+        !fileSourceSlots.some((slot) => slot.part === contentPart),
+    );
+    if (attached !== -1) {
+      ordered[attached] = part;
+      continue;
+    }
     const mediaPositions = ordered.flatMap((contentPart, position) =>
       contentPart.type !== "text" ? [position] : [],
     );
@@ -1102,9 +1136,286 @@ export type FromAgUiMessagesOptions = {
   showThinking?: boolean;
 };
 
+const readSubagentRunId = (record: Record<string, unknown>) => {
+  const runId = getString(record, "subagentRunId");
+  return runId?.trim() ? runId : undefined;
+};
+
+const readToolCallIds = (record: Record<string, unknown>) =>
+  extractAssistantToolCalls(record).flatMap((part) =>
+    part.toolCallId ? [part.toolCallId] : [],
+  );
+
 export function fromAgUiMessages(
   messages: readonly unknown[],
   options?: FromAgUiMessagesOptions,
+): CoreThreadMessageLike[] {
+  return fromAgUiSnapshotMessages(messages, options);
+}
+
+/**
+ * Converts AG-UI messages, nesting each subagent run's messages under the tool
+ * call `SUBAGENT_STARTED.parentToolCallId` announced for it. A message's
+ * `subagentRunId` names its run; a run with no announced spawning call in
+ * `messages` stays in the parent thread.
+ */
+export function fromAgUiSnapshotMessages(
+  messages: readonly unknown[],
+  options?: FromAgUiMessagesOptions,
+  subagentParentToolCallIds?: ReadonlyMap<string, string | undefined>,
+): CoreThreadMessageLike[] {
+  const records = messages.filter(isObject);
+  const callOwners = new Map<string, string | undefined>();
+  for (const record of records) {
+    if (getString(record, "role") !== "assistant") continue;
+    for (const toolCallId of readToolCallIds(record)) {
+      callOwners.set(toolCallId, readSubagentRunId(record));
+    }
+  }
+
+  const spawningCalls = new Map<string, string>();
+  for (const record of records) {
+    const runId = readSubagentRunId(record);
+    if (runId === undefined || spawningCalls.has(runId)) continue;
+    const toolCallId = subagentParentToolCallIds?.get(runId);
+    if (toolCallId !== undefined && callOwners.has(toolCallId))
+      spawningCalls.set(runId, toolCallId);
+  }
+  if (spawningCalls.size === 0) return convertAgUiMessages(messages, options);
+
+  const reachesRoot = (runId: string) => {
+    let current: string | undefined = runId;
+    for (let depth = 0; current !== undefined; depth++) {
+      const toolCallId = spawningCalls.get(current);
+      if (toolCallId === undefined || depth > MAX_SUBAGENT_DEPTH) return false;
+      current = callOwners.get(toolCallId);
+    }
+    return true;
+  };
+  const nestedRuns = new Set([...spawningCalls.keys()].filter(reachesRoot));
+
+  const groups = new Map<string | undefined, unknown[]>();
+  for (const record of records) {
+    const ownRunId = readSubagentRunId(record);
+    const answered =
+      getString(record, "role") === "tool" ? getToolCallId(record) : undefined;
+    const runId =
+      answered !== undefined && callOwners.has(answered)
+        ? callOwners.get(answered)
+        : ownRunId;
+    const role = getString(record, "role");
+    const group =
+      role !== "user" &&
+      role !== "system" &&
+      role !== "developer" &&
+      runId !== undefined &&
+      nestedRuns.has(runId)
+        ? runId
+        : undefined;
+    const members = groups.get(group) ?? [];
+    members.push(record);
+    groups.set(group, members);
+  }
+
+  const runsBySpawningCall = new Map<string, string[]>();
+  for (const runId of nestedRuns) {
+    const toolCallId = spawningCalls.get(runId)!;
+    runsBySpawningCall.set(toolCallId, [
+      ...(runsBySpawningCall.get(toolCallId) ?? []),
+      runId,
+    ]);
+  }
+
+  const convertGroup = (group: string | undefined): CoreThreadMessageLike[] =>
+    convertAgUiMessages(groups.get(group) ?? [], options).map((message) => {
+      if (message.role !== "assistant" || !Array.isArray(message.content))
+        return message;
+      const nested: ThreadMessage[] = [];
+      const content = message.content.map((part) => {
+        if (part.type !== "tool-call") return part;
+        const runIds = runsBySpawningCall.get(part.toolCallId);
+        if (!runIds) return part;
+        const subagents = runIds.map(toSubagentMessage);
+        nested.push(...subagents);
+        return { ...part, messages: [...(part.messages ?? []), ...subagents] };
+      });
+      return liftSubagentState({ ...message, content }, nested);
+    });
+
+  const toSubagentMessage = (runId: string): ThreadMessage => {
+    const messages = convertGroup(runId).filter(
+      (message) => message.role === "assistant",
+    );
+    const folded = messages
+      .map((message) =>
+        typeof message.content === "string"
+          ? {
+              ...message,
+              content: message.content
+                ? [{ type: "text" as const, text: message.content }]
+                : [],
+            }
+          : message,
+      )
+      .reduce<CoreThreadMessageLike | undefined>(
+        (merged, message) =>
+          merged ? foldTurnRecords(merged, message) : message,
+        undefined,
+      );
+    const content = Array.isArray(folded?.content) ? folded.content : [];
+    const pendingToolCallIds = new Set(
+      messages
+        .filter(
+          (message) =>
+            message.status?.type === "requires-action" &&
+            message.status.reason === "tool-calls",
+        )
+        .flatMap((message) => admittedUnresolvedToolCallIds(message, true)),
+    );
+    const metadata = pendingToolCallIds.size
+      ? withPendingToolCallIds(folded?.metadata, pendingToolCallIds)
+      : folded?.metadata;
+    return fromThreadMessageLike(
+      {
+        id: runId,
+        role: "assistant",
+        content,
+        ...(metadata ? { metadata } : {}),
+      },
+      runId,
+      subagentStatus(folded?.metadata, messages),
+    );
+  };
+
+  return convertGroup(undefined);
+}
+
+const hasUnresolvedToolCall = (content: readonly unknown[]) =>
+  content.some(
+    (part) =>
+      isObject(part) && part.type === "tool-call" && part.result === undefined,
+  );
+
+function admittedUnresolvedToolCallIds(
+  message: CoreThreadMessageLike,
+  includeNested: boolean,
+): string[] {
+  if (message.role !== "assistant" || !Array.isArray(message.content))
+    return [];
+  const listed = readPersistedPendingToolCallIds(message.metadata);
+  if (listed) return listed;
+  return message.content.flatMap((part) => {
+    if (part.type !== "tool-call") return [];
+    const calls = includeNested
+      ? [
+          part,
+          ...Array.from(
+            walkToolCallTree(part.messages ?? []),
+            ({ part }) => part,
+          ),
+        ]
+      : [part];
+    return calls.flatMap((call) =>
+      call.result === undefined && call.toolCallId ? [call.toolCallId] : [],
+    );
+  });
+}
+
+function withPendingToolCallIds(
+  metadata: CoreThreadMessageLike["metadata"],
+  ids: ReadonlySet<string>,
+): CoreThreadMessageLike["metadata"] {
+  return {
+    ...(isObject(metadata) ? metadata : {}),
+    custom: {
+      ...readCustomMetadata(metadata),
+      [AG_UI_METADATA_NAMESPACE]: {
+        ...readNamespacedMetadata(metadata),
+        pendingToolCallIds: [...ids],
+      },
+    },
+  };
+}
+
+function subagentStatus(
+  metadata: unknown,
+  messages: readonly CoreThreadMessageLike[],
+): MessageStatus {
+  const hasInterrupt = readPersistedInterrupts(metadata) !== undefined;
+  const hasPendingToolCall = messages.some(
+    (message) =>
+      message.status?.type === "requires-action" &&
+      message.status.reason === "tool-calls",
+  );
+  return hasInterrupt || hasPendingToolCall
+    ? getAutoStatus(false, false, hasInterrupt, hasPendingToolCall, undefined)
+    : { type: "complete", reason: "unknown" };
+}
+
+/**
+ * A nested run's open interrupts and unanswered calls surface on the message
+ * that spawned it, as the live run does.
+ */
+function liftSubagentState(
+  message: CoreThreadMessageLike,
+  nested: readonly ThreadMessage[],
+): CoreThreadMessageLike {
+  const waiting = nested.filter(
+    (subagent) => subagent.status?.type === "requires-action",
+  );
+  if (waiting.length === 0) return message;
+  const interrupts = [
+    ...(readPersistedInterrupts(message.metadata) ?? []),
+    ...waiting.flatMap(
+      (subagent) => readPersistedInterrupts(subagent.metadata) ?? [],
+    ),
+  ];
+  const custom = readCustomMetadata(message.metadata);
+  const metadata = {
+    ...(isObject(message.metadata) ? message.metadata : {}),
+    custom: {
+      ...custom,
+      [AG_UI_METADATA_NAMESPACE]: {
+        ...readNamespacedMetadata(message.metadata),
+        ...(interrupts.length > 0 ? { interrupts } : {}),
+      },
+    },
+  };
+  const hasInterrupt = interrupts.length > 0;
+  const hasPendingToolCall =
+    waiting.some(
+      (subagent) =>
+        subagent.status?.type === "requires-action" &&
+        subagent.status.reason === "tool-calls",
+    ) ||
+    (Array.isArray(message.content) && hasUnresolvedToolCall(message.content));
+  const pendingToolCallIds = new Set([
+    ...admittedUnresolvedToolCallIds(message, false),
+    ...waiting.flatMap((subagent) =>
+      subagent.status?.type === "requires-action" &&
+      subagent.status.reason === "tool-calls"
+        ? admittedUnresolvedToolCallIds(subagent, true)
+        : [],
+    ),
+  ]);
+  return {
+    ...message,
+    metadata: pendingToolCallIds.size
+      ? withPendingToolCallIds(metadata, pendingToolCallIds)
+      : metadata,
+    status: getAutoStatus(
+      false,
+      false,
+      hasInterrupt,
+      hasPendingToolCall,
+      undefined,
+    ),
+  };
+}
+
+function convertAgUiMessages(
+  messages: readonly unknown[],
+  options?: FromAgUiMessagesOptions,
 ): CoreThreadMessageLike[] {
   const showThinking = options?.showThinking ?? true;
   const converted: CoreThreadMessageLike[] = [];

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

@@ -0,0 +1,151 @@
+import { describe, expect, it } from "vitest";
+import type { ChatModelRunResult } from "@assistant-ui/core";
+import { RunAggregator } from "./run-aggregator";
+import { fromAgUiMessages, toAgUiMessages } from "./conversions";
+import { parseAgUiEvent } from "../event-parser";
+
+describe("AG-UI tool result metadata", () => {
+  it("keeps interleaved tool result metadata through finish and message export/import", () => {
+    const expectedSecondMetadata = {
+      trace: "second-final",
+      attempts: 0,
+      optional: null,
+      nested: { value: "after" },
+    };
+    const updates: ChatModelRunResult[] = [];
+    const aggregator = new RunAggregator({
+      showThinking: false,
+      logger: { debug: () => {}, error: () => {} },
+      emit: (update) => updates.push(update),
+    });
+    aggregator.handle({ type: "RUN_STARTED", runId: "run" });
+    for (const toolCallId of ["first", "second", "empty"]) {
+      aggregator.handle({
+        type: "TOOL_CALL_START",
+        toolCallId,
+        toolCallName: "lookup",
+      });
+    }
+
+    for (const event of [
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "second",
+        messageId: "result-second",
+        content: "0",
+        metadata: {
+          trace: "second",
+          attempts: 0,
+          optional: null,
+          nested: { old: true, value: "before" },
+        },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "first",
+        messageId: "result-first",
+        content: "false",
+        metadata: { trace: "first", labels: [] },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "empty",
+        messageId: "result-empty",
+        content: "",
+        metadata: { trace: "empty" },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "second",
+        messageId: "result-second",
+        content: "0",
+        metadata: { trace: "second-final", nested: { value: "after" } },
+      },
+    ]) {
+      const parsed = parseAgUiEvent(event);
+      if (!parsed) throw new Error("Expected a parsed result event");
+      aggregator.handle(parsed);
+    }
+
+    const liveSecond = updates
+      .at(-1)
+      ?.content?.find(
+        (part) => part.type === "tool-call" && part.toolCallId === "second",
+      );
+    if (liveSecond?.type !== "tool-call") {
+      throw new Error("Expected the streamed second tool call");
+    }
+    expect(liveSecond).toMatchObject({ result: 0 });
+    expect(liveSecond.providerMetadata).toEqual({
+      agui: { messageMetadata: expectedSecondMetadata },
+    });
+    aggregator.handle({ type: "RUN_FINISHED", runId: "run" });
+    const content = updates.at(-1)?.content;
+    expect(content).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          toolCallId: "first",
+          result: false,
+          providerMetadata: {
+            agui: { messageMetadata: { trace: "first", labels: [] } },
+          },
+        }),
+        expect.objectContaining({
+          toolCallId: "second",
+          result: 0,
+          providerMetadata: {
+            agui: {
+              messageMetadata: expectedSecondMetadata,
+            },
+          },
+        }),
+        expect.objectContaining({
+          toolCallId: "empty",
+          result: "",
+          providerMetadata: {
+            agui: { messageMetadata: { trace: "empty" } },
+          },
+        }),
+      ]),
+    );
+    const exported = toAgUiMessages([
+      { id: "assistant", role: "assistant", content: content ?? [] },
+    ]);
+    expect(exported).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "first",
+          metadata: { trace: "first", labels: [] },
+        }),
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "second",
+          metadata: expectedSecondMetadata,
+        }),
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "empty",
+          metadata: { trace: "empty" },
+        }),
+      ]),
+    );
+    expect(fromAgUiMessages(exported)).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          content: expect.arrayContaining([
+            expect.objectContaining({
+              toolCallId: "second",
+              result: 0,
+              providerMetadata: {
+                agui: {
+                  messageMetadata: expectedSecondMetadata,
+                },
+              },
+            }),
+          ]),
+        }),
+      ]),
+    );
+  });
+});