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

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