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

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";
 
@@ -27,6 +28,9 @@ export type AgUiOpaqueReasoning = {
 };
 
 export type AgUiCustomMetadata = {
+  /** Wire role restored on export for messages the internal model cannot
+   * represent (a developer record rides as a system message). */
+  role?: "developer";
   interrupts?: AgUiInterrupt[];
   opaqueReasoning?: AgUiOpaqueReasoning[];
 };
@@ -62,29 +66,47 @@ export type RunAggregatorOptions = {
   logger: Logger;
   emit: Emit;
   onServerMessageId?: (messageId: string) => void;
+  onTextMessageStart?: (messageId: string) => void;
 };
 
 /**
  * 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;
   private readonly showThinking: boolean;
   private readonly logger: Logger;
   private readonly onServerMessageId: ((messageId: string) => void) | undefined;
+  private readonly onTextMessageStart:
+    | ((messageId: string) => void)
+    | undefined;
 
   private status: ChatModelRunResult["status"] | undefined;
   private interrupts: AgUiInterrupt[] | undefined;
-  private readonly textParts = new Map<
-    string,
-    { buffer: string; touched: boolean }
-  >();
+  private readonly textParts = new Map<string, { buffer: string }>();
   private activeTextMessageId: string | undefined;
   private readonly reasoningParts = new Map<string, string>(); // key → buffer
   private readonly reasoningSignatures = new Map<string, string>();
+  private readonly reasoningSignatureIds = new Map<string, string>();
+  // Signatures captured while thinking is hidden have no block to live on;
+  // they are transport state that must survive to the next run input. The
+  // claim rules mirror the visible path: an id-carrying block, or an open
+  // anonymous one.
+  private readonly hiddenSignatures = new Map<string, string>();
+  private readonly hiddenSignatureAnchors = new Map<string, number>();
+  private readonly hiddenBlockAnchors = new Map<string, number>();
+  private hiddenAnonymousAnchor: number | undefined;
+  private readonly hiddenReasoningIds = new Set<string>();
+  private hiddenActiveReasoning: "none" | "anonymous" | "identified" = "none";
+  private hasEmittedOpaqueReasoning = false;
+  private readonly loggedDroppedOpaqueIds = new Set<string>();
   private readonly reasoningMessageIds = new Map<string, string>();
   private readonly anonymousReasoningKeys = new Set<string>();
   private activeReasoningKey: string | undefined;
@@ -97,9 +119,12 @@ 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;
+  private reportedServerMessageId: string | undefined;
+  private lastTextMessageId: string | undefined;
 
   private streamStartTime: number | undefined;
   private firstTokenTime: number | undefined;
@@ -110,6 +135,7 @@ export class RunAggregator {
     this.showThinking = options.showThinking;
     this.logger = options.logger;
     this.onServerMessageId = options.onServerMessageId;
+    this.onTextMessageStart = options.onTextMessageStart;
   }
 
   hasToolCall(toolCallId: string): boolean {
@@ -119,22 +145,11 @@ export class RunAggregator {
   handle(event: AgUiEvent): void {
     switch (event.type) {
       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();
-        this.a2uiBuckets.clear();
-        this.a2uiToolCallIds.clear();
-        this.lastResolvedToolCallId = undefined;
-        this.partOrder.length = 0;
-        this.textPartCounter = 0;
-        this.activeTextMessageId = undefined;
+        this.resetMessageParts();
         this.interrupts = undefined;
         this.serverMessageIdReported = false;
+        this.lastTextMessageId = undefined;
+        this.reportedServerMessageId = undefined;
         this.streamStartTime = Date.now();
         this.firstTokenTime = undefined;
         this.totalChunks = 0;
@@ -178,18 +193,19 @@ export class RunAggregator {
       }
 
       case "TEXT_MESSAGE_START": {
+        this.beginDistinctTextMessage(event.messageId);
         this.reportServerMessageId(event.messageId);
-        const id = this.startTextMessage(event.messageId);
-        if (id) {
-          this.markTextPartTouched(id);
-        }
+        this.startTextMessage(event.messageId);
+        if (event.messageId) this.lastTextMessageId = event.messageId;
         this.emit();
         break;
       }
       case "TEXT_MESSAGE_CONTENT":
       case "TEXT_MESSAGE_CHUNK": {
         const incomingId = "messageId" in event ? event.messageId : undefined;
+        this.beginDistinctTextMessage(incomingId);
         this.reportServerMessageId(incomingId);
+        if (incomingId) this.lastTextMessageId = incomingId;
         if (!event.delta) break;
         this.recordFirstToken();
         const id = this.resolveTextMessageId(incomingId);
@@ -207,6 +223,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":
@@ -221,13 +248,31 @@ export class RunAggregator {
           // 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;
+          const key = this.showThinking
+            ? this.reasoningParts.has(event.entityId)
+              ? event.entityId
+              : active !== undefined && this.anonymousReasoningKeys.has(active)
+                ? active
+                : undefined
+            : undefined;
           if (key !== undefined) {
             this.reasoningSignatures.set(key, event.encryptedValue);
+            this.reasoningSignatureIds.set(key, event.entityId);
+            this.emit();
+          } else if (
+            !this.showThinking &&
+            event.entityId.trim().length > 0 &&
+            event.encryptedValue.trim().length > 0 &&
+            (this.hiddenReasoningIds.has(event.entityId) ||
+              this.hiddenActiveReasoning === "anonymous")
+          ) {
+            this.hiddenSignatures.set(event.entityId, event.encryptedValue);
+            this.hiddenSignatureAnchors.set(
+              event.entityId,
+              this.hiddenBlockAnchors.get(event.entityId) ??
+                this.hiddenAnonymousAnchor ??
+                this.partOrder.length,
+            );
             this.emit();
           }
         }
@@ -419,8 +464,13 @@ export class RunAggregator {
   }
 
   private reportServerMessageId(messageId: string | undefined): void {
-    if (this.serverMessageIdReported || !messageId) return;
+    if (!messageId) return;
+    if (this.lastTextMessageId === undefined) {
+      this.lastTextMessageId = messageId;
+    }
+    if (this.serverMessageIdReported) return;
     this.serverMessageIdReported = true;
+    this.reportedServerMessageId = messageId;
     this.onServerMessageId?.(messageId);
   }
 
@@ -428,6 +478,45 @@ export class RunAggregator {
     this.textParts.clear();
   }
 
+  private resetMessageParts(): void {
+    this.clearTextParts();
+    this.reasoningParts.clear();
+    this.reasoningSignatures.clear();
+    this.reasoningSignatureIds.clear();
+    this.reasoningMessageIds.clear();
+    this.hiddenSignatures.clear();
+    this.hiddenSignatureAnchors.clear();
+    this.hiddenBlockAnchors.clear();
+    this.hiddenAnonymousAnchor = undefined;
+    this.hiddenReasoningIds.clear();
+    this.hiddenActiveReasoning = "none";
+    this.loggedDroppedOpaqueIds.clear();
+    this.anonymousReasoningKeys.clear();
+    this.activeReasoningKey = undefined;
+    this.reasoningPartCounter = 0;
+    this.toolCalls.clear();
+    this.a2uiBuckets.clear();
+    this.a2uiToolCallIds.clear();
+    this.lastResolvedToolCallId = undefined;
+    this.partOrder.length = 0;
+    this.textPartCounter = 0;
+    this.activeTextMessageId = undefined;
+    this.reportedServerMessageId = undefined;
+  }
+
+  private beginDistinctTextMessage(messageId: string | undefined): void {
+    if (
+      !messageId ||
+      this.lastTextMessageId === undefined ||
+      this.lastTextMessageId === messageId ||
+      !this.onTextMessageStart
+    ) {
+      return;
+    }
+    this.resetMessageParts();
+    this.onTextMessageStart(messageId);
+  }
+
   private generateTextKey(): string {
     this.textPartCounter += 1;
     return `text-${this.textPartCounter}`;
@@ -459,7 +548,7 @@ export class RunAggregator {
 
   private ensureTextPart(id: string): void {
     if (!this.textParts.has(id)) {
-      this.textParts.set(id, { buffer: "", touched: false });
+      this.textParts.set(id, { buffer: "" });
       if (
         !this.partOrder.some((part) => part.kind === "text" && part.key === id)
       ) {
@@ -468,18 +557,11 @@ export class RunAggregator {
     }
   }
 
-  private markTextPartTouched(id: string): void {
-    const entry = this.textParts.get(id);
-    if (!entry) return;
-    entry.touched = true;
-  }
-
   private appendText(id: string, delta: string): void {
     this.ensureTextPart(id);
     const entry = this.textParts.get(id);
     if (!entry) return;
     entry.buffer += delta;
-    entry.touched = true;
   }
 
   private startToolCall(
@@ -614,8 +696,31 @@ 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()),
+    );
+
+    const opaqueCandidates: (AgUiOpaqueReasoning & { anchor: number })[] =
+      Array.from(this.hiddenSignatures, ([id, encryptedValue]) => ({
+        id,
+        encryptedValue,
+        anchor: this.hiddenSignatureAnchors.get(id)!,
+      }));
+
+    let currentIndex = -1;
+    let lastMaterializedIndex = -1;
+    const pushSnapshotPart = (part: (typeof snapshot)[number]) => {
+      snapshot.push(part);
+      lastMaterializedIndex = currentIndex;
+    };
 
-    for (const part of this.partOrder) {
+    for (const [index, part] of this.partOrder.entries()) {
+      currentIndex = index;
       if (part.kind === "reasoning") {
         if (this.showThinking) {
           const buffer = this.reasoningParts.get(part.key) ?? "";
@@ -626,13 +731,22 @@ export class RunAggregator {
               ...(reasoningId !== undefined ? { reasoningId } : {}),
               ...(encryptedValue !== undefined ? { encryptedValue } : {}),
             };
-            snapshot.push({
+            pushSnapshotPart({
               type: "reasoning",
               text: buffer,
               ...(Object.keys(meta).length > 0
                 ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } }
                 : {}),
             } as const);
+          } else {
+            // A retracted empty block still carries transport state: without
+            // a part to live on, its signature rides the message metadata,
+            // matching the shape a snapshot reload produces.
+            const encryptedValue = this.reasoningSignatures.get(part.key);
+            const id = this.reasoningSignatureIds.get(part.key);
+            if (id?.trim() && encryptedValue?.trim()) {
+              opaqueCandidates.push({ id, encryptedValue, anchor: index + 1 });
+            }
           }
         }
         continue;
@@ -640,20 +754,31 @@ export class RunAggregator {
 
       if (part.kind === "text") {
         const entry = this.textParts.get(part.key);
-        if (entry?.touched) {
-          snapshot.push({ type: "text", text: entry.buffer } as const);
+        if (entry && entry.buffer.trim().length > 0) {
+          pushSnapshotPart({ type: "text", text: entry.buffer } as const);
         }
         continue;
       }
 
+      if (part.kind === "data") {
+        pushSnapshotPart({
+          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 }
@@ -676,17 +801,57 @@ export class RunAggregator {
           ? { unstable_toolMessageId: entry.toolMessageId }
           : {}),
       } as ToolCallMessagePart & { unstable_toolMessageId?: string };
-      snapshot.push(toolPart);
+      pushSnapshotPart(toolPart);
     }
 
+    // An anonymous claim promotes the signature's entityId to a wire message
+    // id; when that id actually names a text message or the adopted assistant
+    // message id (which can also arrive via TOOL_CALL_START.parentMessageId),
+    // replaying it as its own record would put two wire records under one id,
+    // so such entries are dropped instead — the same outcome the visible
+    // path's claim guard produces for a signature that is "not for this
+    // block". Entries with no materialized part after their anchor trail the
+    // assistant record, matching the snapshot path's anchor/after
+    // bookkeeping.
+    const toolMessageIds = new Set<string>();
+    for (const call of this.toolCalls.values()) {
+      if (call.toolMessageId) toolMessageIds.add(call.toolMessageId);
+    }
+    opaqueCandidates.sort((a, b) => a.anchor - b.anchor);
+    const publishableOpaqueReasoning = opaqueCandidates
+      .filter((entry) => {
+        const collides =
+          this.textParts.has(entry.id) ||
+          entry.id === this.reportedServerMessageId ||
+          toolMessageIds.has(entry.id);
+        if (collides && !this.loggedDroppedOpaqueIds.has(entry.id)) {
+          this.loggedDroppedOpaqueIds.add(entry.id);
+          this.logger.debug?.(
+            "[agui] aggregator dropped opaque reasoning signature: id collides with a message id",
+            entry.id,
+          );
+        }
+        return !collides;
+      })
+      .map(({ anchor, ...entry }) =>
+        lastMaterializedIndex < anchor ? { ...entry, after: true } : entry,
+      );
+    if (publishableOpaqueReasoning.length > 0)
+      this.hasEmittedOpaqueReasoning = true;
+    // Once an opaque entry has been published this run, the key keeps being
+    // emitted (empty when withdrawn) so the namespace merge can retract it.
+    const includeOpaque = this.hasEmittedOpaqueReasoning;
     const timing = this.getTiming();
     const metadata = {
       ...(timing ? { timing } : {}),
-      ...(this.interrupts
+      ...(this.interrupts || includeOpaque
         ? {
             custom: {
               [AG_UI_METADATA_NAMESPACE]: {
-                interrupts: this.interrupts,
+                ...(this.interrupts ? { interrupts: this.interrupts } : {}),
+                ...(includeOpaque
+                  ? { opaqueReasoning: publishableOpaqueReasoning }
+                  : {}),
               } satisfies AgUiCustomMetadata,
             },
           }
@@ -742,7 +907,18 @@ export class RunAggregator {
   }
 
   private handleReasoningStart(messageId?: string, isMessageId = false): void {
-    if (!this.showThinking) return;
+    if (!this.showThinking) {
+      const anchor = this.partOrder.length;
+      if (messageId === undefined) {
+        this.hiddenActiveReasoning = "anonymous";
+        this.hiddenAnonymousAnchor = anchor;
+      } else {
+        this.hiddenReasoningIds.add(messageId);
+        this.hiddenBlockAnchors.set(messageId, anchor);
+        this.hiddenActiveReasoning = "identified";
+      }
+      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;
@@ -768,7 +944,15 @@ export class RunAggregator {
     messageId?: string,
     isMessageId = false,
   ): void {
-    if (!this.showThinking || !delta) return;
+    if (!delta) return;
+    if (!this.showThinking) {
+      // Content without a preceding START still names the block (the visible
+      // path opens it lazily); register it so a signature can claim it.
+      if (this.hiddenActiveReasoning === "none") {
+        this.handleReasoningStart(messageId, isMessageId);
+      }
+      return;
+    }
     if (!this.activeReasoningKey) {
       // Content arrived without a preceding START — create the slot lazily.
       this.handleReasoningStart(messageId, isMessageId);
@@ -780,7 +964,10 @@ export class RunAggregator {
   }
 
   private handleReasoningEnd(): void {
-    if (!this.showThinking) return;
+    if (!this.showThinking) {
+      this.hiddenActiveReasoning = "none";
+      return;
+    }
     this.activeReasoningKey = undefined;
     this.emit();
   }

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 });