@assistant-ui/react-ag-ui 0.1.0 → 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/adapter/tool-approval.ts

@@ -228,6 +228,40 @@ export const projectAgUiToolApprovals = (
   );
 };
 
+/** A recorded decision authorizes one exact call and gate, not a reused id. */
+export const preserveToolApprovalDecision = (
+  previous: ToolCallMessagePart | undefined,
+  next: ToolCallMessagePart,
+): ToolApproval | undefined => {
+  const prior = previous?.approval;
+  let approval = next.approval;
+  if (!approval && prior) {
+    const {
+      approved: _approved,
+      reason: _reason,
+      resolution: _resolution,
+      ...pending
+    } = prior;
+    approval = pending;
+  }
+  if (!approval) return undefined;
+  // AG-UI may omit argument events for a no-argument call; both forms mean {}.
+  if (
+    !previous ||
+    prior?.id !== approval.id ||
+    previous.toolCallId !== next.toolCallId ||
+    previous.toolName !== next.toolName ||
+    (previous.argsText || "{}") !== (next.argsText || "{}")
+  )
+    return approval;
+  return {
+    ...approval,
+    ...(prior.approved !== undefined ? { approved: prior.approved } : {}),
+    ...(prior.reason !== undefined ? { reason: prior.reason } : {}),
+    ...(prior.resolution !== undefined ? { resolution: prior.resolution } : {}),
+  };
+};
+
 /**
  * Records a decision on the gated tool call. The decision is held on the part
  * until every gate in the batch is answered, because AG-UI resumes a run with
@@ -266,16 +300,20 @@ export const buildToolApprovalResume = (
   content: readonly ThreadAssistantMessagePart[],
   interrupts: readonly AgUiInterrupt[],
 ): AgUiResumeEntry[] | null => {
-  const open = new Set(interrupts.map((interrupt) => interrupt.id));
+  const calls = Array.from(iterateToolCallParts(content));
+  const gates = projectAgUiToolApprovals(
+    interrupts,
+    new Set(calls.map((part) => part.toolCallId)),
+  );
   // Decisions accumulate in a `Map` rather than in an object literal, because an
   // interrupt id of `__proto__` would otherwise write to the prototype chain
   // instead of recording a response, and `constructor` or `toString` would
   // report themselves as answered while undecided.
   const responses = new Map<string, { status: "resolved"; payload: unknown }>();
-  for (const part of iterateToolCallParts(content)) {
+  for (const part of calls) {
     const approval = part.approval;
     if (!approval || approval.approved === undefined) continue;
-    if (!open.has(approval.id)) continue;
+    if (gates.get(part.toolCallId)?.id !== approval.id) continue;
     responses.set(approval.id, {
       status: "resolved",
       payload: {

package/src/runtime/agUiThreadState.test.ts

@@ -115,7 +115,10 @@ describe("reduceAgUiThreadState", () => {
       owner,
       messageId: "assistant",
     });
-    expect(withResume.pendingResume).toEqual({ owner, messageId: "assistant" });
+    expect(withResume.pendingResume).toEqual({
+      owner,
+      messageIds: ["assistant"],
+    });
     expect(
       reduceAgUiThreadState(withResume, {
         type: "resumeDeferred",
@@ -123,6 +126,22 @@ describe("reduceAgUiThreadState", () => {
         messageId: "assistant",
       }),
     ).toBe(withResume);
+    const withTwoOwners = reduceAgUiThreadState(withResume, {
+      type: "resumeDeferred",
+      owner,
+      messageId: "assistant-next",
+    });
+    expect(withTwoOwners.pendingResume).toEqual({
+      owner,
+      messageIds: ["assistant", "assistant-next"],
+    });
+    expect(
+      reduceAgUiThreadState(withTwoOwners, {
+        type: "resumeDeferred",
+        owner: new AbortController(),
+        messageId: "new-run",
+      }).pendingResume?.messageIds,
+    ).toEqual(["new-run"]);
     const withOwner = reduceAgUiThreadState(withResume, {
       type: "a2uiResumeDeferred",
       owner,
@@ -189,7 +208,7 @@ describe("reduceAgUiThreadState", () => {
     });
     expect(cleared.pendingA2uiResumeOwner).toBeNull();
     expect(cleared.pendingA2uiAction).toBeUndefined();
-    expect(cleared.pendingResume).toEqual({ owner, messageId: "assistant" });
+    expect(cleared.pendingResume).toEqual({ owner, messageIds: ["assistant"] });
     expect(
       reduceAgUiThreadState(cleared, { type: "a2uiContinuationsCleared" }),
     ).toBe(cleared);

package/src/runtime/agUiThreadState.ts

@@ -8,7 +8,10 @@ export type AgUiThreadState = {
   historyMutationGeneration: number;
   stateSnapshot: ReadonlyJSONValue | undefined;
   lastRunConfig: RunConfig | undefined;
-  pendingResume: { owner: AbortController; messageId: string } | null;
+  pendingResume: {
+    owner: AbortController;
+    messageIds: readonly string[];
+  } | null;
   pendingA2uiResumeOwner: AbortController | null;
   pendingA2uiAction: Record<string, unknown> | undefined;
 };
@@ -78,11 +81,17 @@ export const reduceAgUiThreadState = (
         : { ...state, lastRunConfig: action.runConfig };
     case "resumeDeferred":
       return state.pendingResume?.owner === action.owner &&
-        state.pendingResume.messageId === action.messageId
+        state.pendingResume.messageIds.includes(action.messageId)
         ? state
         : {
             ...state,
-            pendingResume: { owner: action.owner, messageId: action.messageId },
+            pendingResume: {
+              owner: action.owner,
+              messageIds:
+                state.pendingResume?.owner === action.owner
+                  ? [...state.pendingResume.messageIds, action.messageId]
+                  : [action.messageId],
+            },
           };
     case "resumeCleared":
       return state.pendingResume === null

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