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

package/src/runtime/adapter/conversions.ts

@@ -27,6 +27,7 @@ import {
   type A2uiSurfaceState,
 } from "@assistant-ui/react-generative-ui/a2ui";
 import type { AgUiInterrupt } from "../types";
+import { projectAgUiToolApprovals } from "./tool-approval";
 import {
   parseMcpToolCallResult,
   readMcpAppResourceUri,
@@ -41,7 +42,7 @@ type AttachmentLike = {
 };
 
 type ThreadMessageLike = {
-  id: string;
+  id?: string;
   role: string;
   content: unknown;
   metadata?: unknown;
@@ -51,6 +52,8 @@ type ThreadMessageLike = {
   attachments?: readonly AttachmentLike[];
 };
 
+type NormalizedThreadMessageLike = ThreadMessageLike & { id: string };
+
 type AgUiToolCall = {
   id: string;
   type: "function";
@@ -60,11 +63,23 @@ type AgUiToolCall = {
 export type AgUiMessage =
   | {
       id: string;
-      role: string;
+      role: "user";
       content: string | InputContent[];
       name?: string;
+    }
+  | {
+      id: string;
+      role: "assistant";
+      content: string;
+      name?: string;
       toolCalls?: AgUiToolCall[];
     }
+  | {
+      id: string;
+      role: "system" | "developer";
+      content: string;
+      name?: string;
+    }
   | {
       id: string;
       role: "reasoning";
@@ -606,13 +621,27 @@ function toAssistantSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
   const text = extractText(rawMessage.content);
-  const toolCallParts = extractAssistantToolCalls(rawMessage);
+  const interrupts = readPersistedInterrupts(rawMessage.metadata);
+  const restoredToolCalls = extractAssistantToolCalls(rawMessage);
+  const approvals = projectAgUiToolApprovals(
+    interrupts,
+    new Set(
+      restoredToolCalls
+        .map((part) => part.toolCallId)
+        .filter((id): id is string => !!id),
+    ),
+  );
+  const toolCallParts = restoredToolCalls.map((part) => {
+    const approval = part.toolCallId
+      ? approvals.get(part.toolCallId)
+      : undefined;
+    return approval ? { ...part, approval } : part;
+  });
   const assistantContent = [
     ...(text.length > 0 ? [{ type: "text" as const, text }] : []),
     ...toolCallParts,
   ];
   const messageName = getString(rawMessage, "name");
-  const interrupts = readPersistedInterrupts(rawMessage.metadata);
   return {
     id: getString(rawMessage, "id") ?? generateId(),
     role: "assistant",
@@ -960,7 +989,7 @@ export function fromAgUiMessages(
 }
 
 function convertAssistantMessage(
-  message: ThreadMessageLike,
+  message: NormalizedThreadMessageLike,
   converted: AgUiMessage[],
 ): void {
   const content = extractText(message.content);
@@ -1012,18 +1041,15 @@ function convertAssistantMessage(
     return;
   }
 
-  const assistantMessage: AgUiMessage = {
+  converted.push({
     id: message.id,
     role: "assistant",
     content,
-  };
-  if (message.name) {
-    assistantMessage.name = message.name;
-  }
-  if (toolCalls.length > 0) {
-    assistantMessage.toolCalls = toolCalls.map((entry) => entry.call);
-  }
-  converted.push(assistantMessage);
+    ...(message.name ? { name: message.name } : {}),
+    ...(toolCalls.length > 0
+      ? { toolCalls: toolCalls.map((entry) => entry.call) }
+      : {}),
+  });
 
   for (const { id: toolCallId, part } of toolCalls) {
     if (part.result === undefined) continue;
@@ -1035,36 +1061,30 @@ function convertAssistantMessage(
           ? part.result
           : JSON.stringify(part.result);
 
-    const toolMessage: AgUiMessage = {
+    converted.push({
       id: part.unstable_toolMessageId ?? `${toolCallId}:tool`,
       role: "tool",
       content: resultContent,
       toolCallId,
-    };
-    if (part.isError) {
-      toolMessage.error = resultContent;
-    }
-    converted.push(toolMessage);
+      ...(part.isError ? { error: resultContent } : {}),
+    });
   }
 }
 
 function convertToolMessage(
-  message: ThreadMessageLike,
+  message: NormalizedThreadMessageLike,
   converted: AgUiMessage[],
 ): void {
   const content = extractText(message.content);
   const toolCallId = message.toolCallId ?? generateId();
 
-  const toolMessage: AgUiMessage = {
+  converted.push({
     id: message.id,
     role: "tool",
     content,
     toolCallId,
-  };
-  if (typeof message.error === "string") {
-    toolMessage.error = message.error;
-  }
-  converted.push(toolMessage);
+    ...(typeof message.error === "string" ? { error: message.error } : {}),
+  });
 }
 
 export function toAgUiMessages(
@@ -1072,7 +1092,11 @@ export function toAgUiMessages(
 ): AgUiMessage[] {
   const converted: AgUiMessage[] = [];
 
-  for (const message of messages) {
+  for (const rawMessage of messages) {
+    const message: NormalizedThreadMessageLike = {
+      ...rawMessage,
+      id: rawMessage.id ?? generateId(),
+    };
     const opaqueReasoning = readOpaqueReasoning(message.metadata);
     const toOpaqueRecord = (entry: AgUiOpaqueReasoning): AgUiMessage => ({
       id: entry.id,
@@ -1101,18 +1125,38 @@ export function toAgUiMessages(
       continue;
     }
 
-    const genericMessage: AgUiMessage = {
-      id: message.id,
-      role: message.role,
-      content:
-        message.role === "user"
-          ? buildUserContent(message)
-          : extractText(message.content),
-    };
-    if (message.name) {
-      genericMessage.name = message.name;
+    if (message.role === "user") {
+      converted.push({
+        id: message.id,
+        role: "user",
+        content: buildUserContent(message),
+        ...(message.name ? { name: message.name } : {}),
+      });
+      flushTrailingOpaqueReasoning();
+      continue;
+    }
+
+    if (message.role === "system" || message.role === "developer") {
+      converted.push({
+        id: message.id,
+        role: message.role,
+        content: extractText(message.content),
+        ...(message.name ? { name: message.name } : {}),
+      });
+      flushTrailingOpaqueReasoning();
+      continue;
+    }
+
+    if (message.role === "reasoning") {
+      converted.push({
+        id: message.id,
+        role: "reasoning",
+        content: extractText(message.content),
+      });
+      flushTrailingOpaqueReasoning();
+      continue;
     }
-    converted.push(genericMessage);
+
     flushTrailingOpaqueReasoning();
   }
 

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

@@ -0,0 +1,103 @@
+import { describe, it, expect, vi } from "vitest";
+import type { ChatModelRunResult } from "@assistant-ui/core";
+import { RunAggregator } from "./run-aggregator";
+import type { AgUiInterrupt } from "../types";
+
+const GATE: AgUiInterrupt = {
+  id: "int-1",
+  reason: "tool_call",
+  toolCallId: "tc-1",
+  message: "Delete /tmp/a?",
+};
+
+const noopLogger = {
+  debug: vi.fn(),
+  info: vi.fn(),
+  warn: vi.fn(),
+  error: vi.fn(),
+} as any;
+
+const gatedAggregator = () => {
+  const emitted: ChatModelRunResult[] = [];
+  const aggregator = new RunAggregator({
+    showThinking: false,
+    logger: noopLogger,
+    emit: (result) => {
+      emitted.push(result);
+    },
+  });
+
+  aggregator.handle({ type: "RUN_STARTED", runId: "run-1" });
+  aggregator.handle({
+    type: "TOOL_CALL_START",
+    toolCallId: "tc-1",
+    toolCallName: "delete_file",
+  });
+  aggregator.handle({
+    type: "TOOL_CALL_ARGS",
+    toolCallId: "tc-1",
+    delta: '{"path":"/tmp/a"}',
+  });
+  aggregator.handle({ type: "TOOL_CALL_END", toolCallId: "tc-1" });
+  aggregator.handle({
+    type: "RUN_FINISHED",
+    runId: "run-1",
+    outcome: { type: "interrupt", interrupts: [GATE] },
+  } as any);
+
+  return { aggregator, emitted, last: () => emitted[emitted.length - 1]! };
+};
+
+const approvalOf = (result: ChatModelRunResult) => {
+  const part = (result.content ?? []).find((p) => p.type === "tool-call");
+  return (part as { approval?: unknown } | undefined)?.approval;
+};
+
+describe("RunAggregator tool approval projection", () => {
+  it("projects the gate while the run awaits the interrupt", () => {
+    const { last } = gatedAggregator();
+
+    expect(last().status).toMatchObject({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+    expect(approvalOf(last())).toMatchObject({ id: "int-1" });
+  });
+
+  /**
+   * The gate cannot be clicked once the run is unresumable, but the interrupts
+   * stay on the message: the bespoke hooks read that payload, and dropping it
+   * would take a shipped surface away with the projection.
+   */
+  it("drops the gate when the run then errors, keeping the interrupts", () => {
+    const { aggregator, last } = gatedAggregator();
+
+    aggregator.handle({ type: "RUN_ERROR", message: "stream dropped" } as any);
+
+    expect(last().status).toMatchObject({
+      type: "incomplete",
+      reason: "error",
+    });
+    expect(approvalOf(last())).toBeUndefined();
+    expect(
+      (last().metadata?.custom as Record<string, any> | undefined)?.agui
+        ?.interrupts,
+    ).toEqual([GATE]);
+  });
+
+  it("drops the gate when the run is then cancelled, keeping the interrupts", () => {
+    const { aggregator, last } = gatedAggregator();
+
+    aggregator.handle({ type: "RUN_CANCELLED" } as any);
+
+    expect(last().status).toMatchObject({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+    expect(approvalOf(last())).toBeUndefined();
+    expect(
+      (last().metadata?.custom as Record<string, any> | undefined)?.agui
+        ?.interrupts,
+    ).toEqual([GATE]);
+  });
+});

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

@@ -15,6 +15,7 @@ import {
   type A2uiSurfaceState,
 } from "@assistant-ui/react-generative-ui/a2ui";
 import { readMcpAppResourceUri } from "../mcp-tool-result";
+import { projectAgUiToolApprovals } from "./tool-approval";
 import type { AgUiEvent, AgUiInterrupt } from "../types";
 import type { Logger } from "../logger";
 
@@ -68,7 +69,11 @@ export type RunAggregatorOptions = {
  * Collects AG-UI events into assistant-ui run snapshots that can be yielded from a ChatModelAdapter.
  *
  * The aggregator keeps a single assistant message worth of parts. Each incoming event updates the parts and
- * emits a fresh snapshot through the provided `emit` callback.
+ * emits a fresh snapshot through the provided `emit` callback. `CUSTOM` events
+ * become canonical data parts; integration plumbing names such as
+ * `on_interrupt`, `PredictState`, `Exit`, `hook_error`, `state_update_error`,
+ * `system:*`, and `MultiAgentHandoff` are forwarded the same way, so apps only
+ * need data renderers for names they own.
  */
 export class RunAggregator {
   private readonly emitUpdate: Emit;
@@ -97,6 +102,7 @@ export class RunAggregator {
     | { kind: "text"; key: string }
     | { kind: "reasoning"; key: string }
     | { kind: "tool-call"; toolCallId: string }
+    | { kind: "data"; name: string; value: unknown }
   )[] = [];
   private textPartCounter = 0;
   private serverMessageIdReported = false;
@@ -207,6 +213,17 @@ export class RunAggregator {
         break;
       }
 
+      case "CUSTOM": {
+        this.activeTextMessageId = undefined;
+        this.partOrder.push({
+          kind: "data",
+          name: event.name,
+          value: event.value,
+        });
+        this.emit();
+        break;
+      }
+
       case "THINKING_START":
       case "THINKING_TEXT_MESSAGE_START":
       case "REASONING_START":
@@ -614,6 +631,14 @@ export class RunAggregator {
 
   private emit(): void {
     const snapshot: ThreadAssistantMessagePart[] = [];
+    // A run that ended incomplete can no longer be resumed, so a gate left over
+    // from an earlier interrupt outcome is unanswerable and must not stay
+    // projected. The interrupts themselves are kept on the message, since the
+    // bespoke hooks read that payload.
+    const approvals = projectAgUiToolApprovals(
+      this.status?.type === "requires-action" ? this.interrupts : undefined,
+      new Set(this.toolCalls.keys()),
+    );
 
     for (const part of this.partOrder) {
       if (part.kind === "reasoning") {
@@ -646,14 +671,25 @@ export class RunAggregator {
         continue;
       }
 
+      if (part.kind === "data") {
+        snapshot.push({
+          type: "data",
+          name: part.name,
+          data: part.value,
+        });
+        continue;
+      }
+
       const entry = this.toolCalls.get(part.toolCallId);
       if (!entry) continue;
+      const approval = approvals.get(entry.toolCallId);
       const toolPart: ToolCallMessagePart = {
         type: "tool-call",
         toolCallId: entry.toolCallId,
         toolName: entry.toolCallName,
         args: (entry.parsedArgs ?? {}) as any,
         argsText: entry.argsText,
+        ...(approval ? { approval } : {}),
         ...(entry.result !== undefined ? { result: entry.result } : {}),
         ...(entry.modelContent !== undefined
           ? { modelContent: entry.modelContent }

package/src/runtime/adapter/subscriber.ts

@@ -35,6 +35,7 @@ type Subscriber = {
   onCustomEvent?: (payload: { event: unknown }) => void;
   onRawEvent?: (payload: { event: unknown }) => void;
   onRunFinishedEvent?: (payload: { event: unknown }) => void;
+  onRunErrorEvent?: (payload: { event: unknown }) => void;
   onRunFinalized?: () => void;
   onRunFailed?: (payload: { error: Error }) => void;
 };
@@ -146,11 +147,24 @@ export const createAgUiSubscriber = (
     onCustomEvent: ({ event }) => dispatchIfValid(event, "CUSTOM"),
     onRawEvent: ({ event }) => dispatchIfValid(event, "RAW"),
     onRunFinishedEvent: ({ event }) => {
+      if (runFinishedDispatched) return;
       const parsed = ensureEvent(event, "RUN_FINISHED", logger);
       if (!parsed) return;
       runFinishedDispatched = true;
       dispatch(parsed);
     },
+    onRunErrorEvent: ({ event }) => {
+      const parsed = ensureEvent(event, "RUN_ERROR", logger);
+      if (!parsed || parsed.type !== "RUN_ERROR") return;
+      runFinishedDispatched = true;
+      dispatch(parsed);
+
+      const error = Object.assign(
+        new Error(parsed.message ?? "Run failed"),
+        parsed.code === undefined ? {} : { code: parsed.code },
+      );
+      onRunFailed?.(error);
+    },
     onRunFinalized: () => {
       if (runFinishedDispatched) return;
       dispatch({ type: "RUN_FINISHED", runId });