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

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

@@ -22,6 +22,7 @@ import {
 import jsonpatch, { type Operation } from "fast-json-patch";
 import { readMcpAppResourceUri, toolContentText } from "../mcp-tool-result";
 import { projectAgUiToolApprovals } from "./tool-approval";
+import { withAgUiPartMessageMetadata } from "./conversions";
 import {
   AG_UI_METADATA_NAMESPACE,
   A2UI_SURFACE_ACTIVITY_TYPE,
@@ -30,15 +31,12 @@ import {
   type AgUiEvent,
   type AgUiInterrupt,
   type AgUiOpaqueReasoning,
+  MAX_SUBAGENT_DEPTH,
 } from "../types";
 import type { Logger } from "../logger";
 
 const ROOT_SCOPE = "";
 
-// Defensive guard against a malformed or cyclic parentToolCallId /
-// parentSubagentRunId chain, not a realistic nesting depth.
-const MAX_SUBAGENT_DEPTH = 16;
-
 type SubagentRunState = {
   subagentRunId: string;
   name: string;
@@ -93,6 +91,7 @@ type ToolCallState = {
   mcpAppServerId?: string;
   artifact?: unknown;
   modelContent?: ToolModelContentPart[];
+  resultMetadata?: Record<string, unknown>;
   snapshotResultApplied: boolean;
   subagentRunId?: string;
 };
@@ -209,6 +208,7 @@ export class RunAggregator {
   private interrupts: AgUiInterrupt[] | undefined;
   private pendingToolCallIds: string[] | undefined;
   private readonly textParts = new Map<string, { buffer: string }>();
+  private readonly pendingSnapshotText = new Map<string, string>();
   private readonly activeTextMessageIdByScope = new Map<string, string>();
   private readonly reasoningParts = new Map<
     string,
@@ -282,6 +282,20 @@ export class RunAggregator {
         this.emit();
         break;
       }
+      case "MESSAGES_SNAPSHOT": {
+        this.pendingSnapshotText.clear();
+        for (const message of event.messages) {
+          if (
+            message.role !== "assistant" ||
+            typeof message.content !== "string"
+          )
+            continue;
+          const part = this.textParts.get(message.id);
+          if (part) part.buffer = message.content;
+          else this.pendingSnapshotText.set(message.id, message.content);
+        }
+        break;
+      }
       case "RUN_FINISHED": {
         if (event.outcome?.type === "cancelled") {
           this.interrupts = undefined;
@@ -368,7 +382,7 @@ export class RunAggregator {
         const incomingId = event.messageId;
         const scope = this.scopeOf(event);
         if (scope === ROOT_SCOPE) {
-          this.beginDistinctTextMessage(incomingId);
+          this.beginDistinctTextMessage(incomingId, true);
           this.reportServerMessageId(incomingId);
           if (incomingId) this.lastTextMessageId = incomingId;
         }
@@ -494,6 +508,7 @@ export class RunAggregator {
               : undefined,
           event.messageId,
           event.mcpResult,
+          event.metadata,
         );
         this.emit();
         break;
@@ -793,6 +808,7 @@ export class RunAggregator {
 
   private resetMessageParts(): void {
     this.clearTextParts();
+    this.pendingSnapshotText.clear();
     this.reasoningParts.clear();
     this.reasoningSignatures.clear();
     this.reasoningSignatureIds.clear();
@@ -837,7 +853,10 @@ export class RunAggregator {
     return scope === ROOT_SCOPE ? id : `${scope}\u0000${id}`;
   }
 
-  private beginDistinctTextMessage(messageId: string | undefined): void {
+  private beginDistinctTextMessage(
+    messageId: string | undefined,
+    preserveSnapshotText = false,
+  ): void {
     if (
       this.joinStrategy === "concat-content" ||
       !messageId ||
@@ -854,7 +873,13 @@ export class RunAggregator {
     const openToolCallIds = Array.from(this.toolCalls.values())
       .filter((call) => call.result === undefined)
       .map((call) => call.toolCallId);
+    const snapshotText = preserveSnapshotText
+      ? this.pendingSnapshotText.get(messageId)
+      : undefined;
     this.resetMessageParts();
+    if (snapshotText !== undefined) {
+      this.pendingSnapshotText.set(messageId, snapshotText);
+    }
     this.onTextMessageStart(messageId, openToolCallIds);
     this.restoreReasoning(carried);
   }
@@ -946,15 +971,21 @@ export class RunAggregator {
   private startTextMessage(scope: string, messageId?: string): string {
     const id = messageId ?? this.generateTextKey();
     const key = this.ensureTextPart(scope, id);
+    if (scope === ROOT_SCOPE && messageId) {
+      const snapshotText = this.pendingSnapshotText.get(messageId);
+      const part = this.textParts.get(key);
+      if (snapshotText !== undefined && part) {
+        part.buffer = snapshotText;
+        this.pendingSnapshotText.delete(messageId);
+      }
+    }
     this.activeTextMessageIdByScope.set(scope, key);
     return key;
   }
 
   private resolveTextMessageId(scope: string, messageId?: string): string {
     if (messageId) {
-      const key = this.ensureTextPart(scope, messageId);
-      this.activeTextMessageIdByScope.set(scope, key);
-      return key;
+      return this.startTextMessage(scope, messageId);
     }
 
     const active = this.activeTextMessageIdByScope.get(scope);
@@ -1098,6 +1129,7 @@ export class RunAggregator {
     isError?: boolean,
     toolMessageId?: string,
     mcpResult?: Record<string, unknown>,
+    metadata?: Record<string, unknown>,
   ) {
     if (!id) return;
     const modelText = toolContentText(content);
@@ -1149,6 +1181,9 @@ export class RunAggregator {
     if (toolMessageId) {
       entry.toolMessageId = toolMessageId;
     }
+    if (metadata !== undefined) {
+      entry.resultMetadata = { ...entry.resultMetadata, ...metadata };
+    }
     this.lastResolvedToolCallIdByScope.set(scope, id);
   }
 
@@ -1299,6 +1334,14 @@ export class RunAggregator {
           ...(entry.toolMessageId
             ? { unstable_toolMessageId: entry.toolMessageId }
             : {}),
+          ...(entry.resultMetadata !== undefined
+            ? {
+                providerMetadata: withAgUiPartMessageMetadata(
+                  undefined,
+                  entry.resultMetadata,
+                ),
+              }
+            : {}),
           ...(nestedMessages.length > 0 ? { messages: nestedMessages } : {}),
         } as ToolCallMessagePart & { unstable_toolMessageId?: string };
         snapshot.push(toolPart);

package/src/runtime/adapter/subscriber.test.ts

@@ -4,7 +4,7 @@ import { createAgUiSubscriber } from "./subscriber";
 const createSubscriber = () => {
   const dispatch = vi.fn();
   const logger = { debug: vi.fn(), error: vi.fn() };
-  const subscriber = createAgUiSubscriber({ dispatch, runId: "run-1", logger });
+  const subscriber = createAgUiSubscriber({ dispatch, logger });
   return { subscriber, dispatch, logger };
 };
 

package/src/runtime/adapter/subscriber.ts

@@ -60,7 +60,6 @@ const ensureEvent = (
 
 type SubscriberOptions = {
   dispatch: Dispatch;
-  runId: string;
   logger?: Logger;
   onRunFailed?: (error: Error) => void;
 };
@@ -68,7 +67,7 @@ type SubscriberOptions = {
 export const createAgUiSubscriber = (
   options: SubscriberOptions,
 ): Subscriber => {
-  const { dispatch, runId, logger, onRunFailed } = options;
+  const { dispatch, logger, onRunFailed } = options;
   let runFinishedDispatched = false;
   const dispatchIfValid = (raw: unknown, type: AgUiEvent["type"]) => {
     const event = ensureEvent(raw, type, logger);
@@ -140,10 +139,6 @@ export const createAgUiSubscriber = (
       logger?.debug?.("[agui] step boundary ignored", event),
     onStepFinishedEvent: ({ event }) =>
       logger?.debug?.("[agui] step boundary ignored", event),
-    onRunFinalized: () => {
-      if (runFinishedDispatched) return;
-      dispatch({ type: "RUN_FINISHED", runId });
-    },
     onRunFailed: ({ error }) => {
       if (runFinishedDispatched) return;
       runFinishedDispatched = true;

package/src/runtime/event-parser.ts

@@ -344,6 +344,9 @@ export const parseAgUiEvent = (
           messageId: getString("messageId"),
           role: payload.role === "tool" ? "tool" : undefined,
           mcpResult: parseMcpToolCallResult(payload, content),
+          metadata: isPlainObject(payload.metadata)
+            ? payload.metadata
+            : undefined,
           subagentRunId: getString("subagentRunId"),
         },
       );

package/src/runtime/types.ts

@@ -26,6 +26,10 @@ export const AG_UI_METADATA_NAMESPACE = "agui";
 export const MCP_APPS_ACTIVITY_TYPE = "mcp-apps";
 export const A2UI_SURFACE_ACTIVITY_TYPE = "a2ui-surface";
 
+// Defensive guard against a malformed or cyclic parentToolCallId /
+// parentSubagentRunId chain, not a realistic nesting depth.
+export const MAX_SUBAGENT_DEPTH = 16;
+
 export type AgUiOpaqueReasoning = {
   id: string;
   encryptedValue: string;
@@ -278,6 +282,7 @@ export type AgUiEvent =
       content: string | ContentPart[];
       role?: "tool";
       mcpResult?: McpToolCallResult;
+      metadata?: Record<string, unknown>;
       subagentRunId?: string;
     }
   | {