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

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

@@ -15,13 +15,21 @@ 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";
 
 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;
@@ -61,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;
@@ -77,6 +89,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>();
@@ -87,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;
@@ -111,6 +127,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();
@@ -194,14 +213,42 @@ 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":
       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 +258,7 @@ export class RunAggregator {
         this.handleReasoningContent(
           event.delta,
           "messageId" in event ? event.messageId : undefined,
+          true,
         );
         this.totalChunks++;
         this.recordFirstToken();
@@ -583,13 +631,33 @@ 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") {
         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;
@@ -603,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 }
@@ -698,12 +777,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 +799,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;

package/src/runtime/adapter/subscriber.ts

@@ -22,6 +22,7 @@ type Subscriber = {
   onReasoningMessageStartEvent?: (payload: { event: unknown }) => void;
   onReasoningMessageContentEvent?: (payload: { event: unknown }) => void;
   onReasoningMessageEndEvent?: (payload: { event: unknown }) => void;
+  onReasoningEncryptedValueEvent?: (payload: { event: unknown }) => void;
   onToolCallStartEvent?: (payload: { event: unknown }) => void;
   onToolCallArgsEvent?: (payload: { event: unknown }) => void;
   onToolCallEndEvent?: (payload: { event: unknown }) => void;
@@ -34,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;
 };
@@ -124,6 +126,8 @@ export const createAgUiSubscriber = (
       dispatchIfValid(event, "REASONING_MESSAGE_CONTENT"),
     onReasoningMessageEndEvent: ({ event }) =>
       dispatchIfValid(event, "REASONING_MESSAGE_END"),
+    onReasoningEncryptedValueEvent: ({ event }) =>
+      dispatchIfValid(event, "REASONING_ENCRYPTED_VALUE"),
     onToolCallStartEvent: ({ event }) =>
       dispatchIfValid(event, "TOOL_CALL_START"),
     onToolCallArgsEvent: ({ event }) =>
@@ -143,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 });