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

package/src/runtime/types.ts

@@ -127,6 +127,10 @@ export type AgUiRunFinishedOutcome =
   | { type: "success" }
   | { type: "interrupt"; interrupts: AgUiInterrupt[] };
 
+export type AgUiSubagentFinishedOutcome =
+  | { type: "success" }
+  | { type: "suspended"; interruptIds?: string[] };
+
 export type AgUiEvent =
   | { type: "RUN_STARTED"; runId: string }
   | {
@@ -136,40 +140,71 @@ export type AgUiEvent =
     }
   | { type: "RUN_CANCELLED"; runId?: string }
   | { type: "RUN_ERROR"; message?: string; code?: string }
-  | { type: "TEXT_MESSAGE_START"; messageId?: string }
-  | { type: "TEXT_MESSAGE_CONTENT"; messageId?: string; delta: string }
-  | { type: "TEXT_MESSAGE_END"; messageId?: string }
-  | { type: "TEXT_MESSAGE_CHUNK"; delta: string }
+  | { type: "TEXT_MESSAGE_START"; messageId?: string; subagentRunId?: string }
+  | {
+      type: "TEXT_MESSAGE_CONTENT";
+      messageId?: string;
+      delta: string;
+      subagentRunId?: string;
+    }
+  | { type: "TEXT_MESSAGE_END"; messageId?: string; subagentRunId?: string }
+  | {
+      type: "TEXT_MESSAGE_CHUNK";
+      messageId?: string;
+      delta: string;
+      subagentRunId?: string;
+    }
   | { type: "THINKING_START"; title?: string }
   | { type: "THINKING_TEXT_MESSAGE_START" }
   | { type: "THINKING_TEXT_MESSAGE_CONTENT"; delta: string }
   | { type: "THINKING_TEXT_MESSAGE_END" }
   | { type: "THINKING_END" }
-  | { type: "REASONING_START"; messageId?: string }
-  | { type: "REASONING_MESSAGE_START"; messageId?: string }
-  | { type: "REASONING_MESSAGE_CONTENT"; messageId?: string; delta: string }
-  | { type: "REASONING_MESSAGE_END"; messageId?: string }
+  | { type: "REASONING_START"; messageId?: string; subagentRunId?: string }
+  | {
+      type: "REASONING_MESSAGE_START";
+      messageId?: string;
+      subagentRunId?: string;
+    }
+  | {
+      type: "REASONING_MESSAGE_CONTENT";
+      messageId?: string;
+      delta: string;
+      subagentRunId?: string;
+    }
+  | {
+      type: "REASONING_MESSAGE_END";
+      messageId?: string;
+      subagentRunId?: string;
+    }
   | {
       type: "REASONING_ENCRYPTED_VALUE";
       subtype: "message" | "tool-call";
       entityId: string;
       encryptedValue: string;
+      subagentRunId?: string;
     }
-  | { type: "REASONING_END"; messageId?: string }
+  | { type: "REASONING_END"; messageId?: string; subagentRunId?: string }
   | {
       type: "TOOL_CALL_START";
       toolCallId: string;
       toolCallName?: string;
       parentMessageId?: string;
+      subagentRunId?: string;
+    }
+  | {
+      type: "TOOL_CALL_ARGS";
+      toolCallId: string;
+      delta: string;
+      subagentRunId?: string;
     }
-  | { type: "TOOL_CALL_ARGS"; toolCallId: string; delta: string }
-  | { type: "TOOL_CALL_END"; toolCallId: string }
+  | { type: "TOOL_CALL_END"; toolCallId: string; subagentRunId?: string }
   | {
       type: "TOOL_CALL_CHUNK";
       toolCallId?: string;
       toolCallName?: string;
       parentMessageId?: string;
       delta?: string;
+      subagentRunId?: string;
     }
   | {
       type: "TOOL_CALL_RESULT";
@@ -178,6 +213,7 @@ export type AgUiEvent =
       content: string;
       role?: "tool";
       mcpResult?: McpToolCallResult;
+      subagentRunId?: string;
     }
   | {
       type: "ACTIVITY_SNAPSHOT";
@@ -185,9 +221,31 @@ export type AgUiEvent =
       content: Record<string, unknown>;
       messageId?: string;
       replace?: boolean;
+      subagentRunId?: string;
     }
   | { type: "RAW"; event: any; source?: string }
   | { type: "CUSTOM"; name: string; value: any }
   | { type: "STATE_SNAPSHOT"; snapshot: any }
   | { type: "STATE_DELTA"; delta: any[] }
-  | { type: "MESSAGES_SNAPSHOT"; messages: any[] };
+  | { type: "MESSAGES_SNAPSHOT"; messages: any[] }
+  | {
+      type: "SUBAGENT_STARTED";
+      subagentRunId: string;
+      name: string;
+      description?: string;
+      parentSubagentRunId?: string;
+      parentToolCallId?: string;
+      parentMessageId?: string;
+    }
+  | {
+      type: "SUBAGENT_FINISHED";
+      subagentRunId: string;
+      result?: unknown;
+      outcome?: AgUiSubagentFinishedOutcome;
+    }
+  | {
+      type: "SUBAGENT_ERROR";
+      subagentRunId: string;
+      message: string;
+      code?: string;
+    };

package/src/useAgUiRuntime.approval.test.tsx

@@ -30,13 +30,23 @@ import { readRawResponseSchema } from "./runtime/interrupt-internals";
 
 type Subscriber = Record<string, ((payload: any) => void) | undefined>;
 
+const GATE_PROMPT = "Delete /tmp/a?";
+
 const GATE: AgUiInterrupt = {
   id: "int-1",
   reason: "tool_call",
   toolCallId: "tc-1",
-  message: "Delete /tmp/a?",
+  message: GATE_PROMPT,
 };
 
+// The gate's message reaches the renderer as the approval's prompt, so every
+// projection of GATE carries it.
+const gateApproval = (rest: Record<string, unknown> = {}) => ({
+  id: "int-1",
+  prompt: GATE_PROMPT,
+  ...rest,
+});
+
 // Gates the first run on an interrupt and lets every later run settle, so a
 // resumed run is observable as a second call rather than a second gate.
 const gatingAgent = (interrupts: readonly AgUiInterrupt[] = [GATE]) => {
@@ -145,7 +155,7 @@ describe("useAgUiRuntime tool approvals", () => {
   it("gates the tool call a tool_call interrupt names", async () => {
     const { runtime } = await gatedThread();
 
-    expect(allToolCalls(runtime.current)[0]!.approval).toEqual({ id: "int-1" });
+    expect(allToolCalls(runtime.current)[0]!.approval).toEqual(gateApproval());
     expect(gatedPart(runtime.current).getState().status).toMatchObject({
       type: "requires-action",
     });
@@ -200,10 +210,9 @@ describe("useAgUiRuntime tool approvals", () => {
     });
     expect(runAgent).toHaveBeenCalledTimes(1);
     // The first decision survives the rerender the second gate triggers.
-    expect(allToolCalls(runtime.current)[0]!.approval).toEqual({
-      id: "int-1",
-      approved: true,
-    });
+    expect(allToolCalls(runtime.current)[0]!.approval).toEqual(
+      gateApproval({ approved: true }),
+    );
 
     await act(async () => {
       gatedPart(runtime.current, "tc-2").respondToToolApproval({
@@ -222,10 +231,9 @@ describe("useAgUiRuntime tool approvals", () => {
       },
     ]);
     // Settlement consumes the resume array, so the approved gate stays approved.
-    expect(allToolCalls(runtime.current)[0]!.approval).toEqual({
-      id: "int-1",
-      approved: true,
-    });
+    expect(allToolCalls(runtime.current)[0]!.approval).toEqual(
+      gateApproval({ approved: true }),
+    );
     expect(
       allToolCalls(runtime.current).every((part) => part.result === undefined),
     ).toBe(true);
@@ -255,7 +263,7 @@ describe("useAgUiRuntime tool approvals", () => {
     // The wire cancelled both, so the gate approved while its sibling was
     // still open cannot keep displaying an approval that was never sent.
     expect(allToolCalls(runtime.current).map((part) => part.approval)).toEqual([
-      { id: "int-1", resolution: "cancelled" },
+      gateApproval({ resolution: "cancelled" }),
       { id: "int-2", resolution: "cancelled" },
     ]);
   });
@@ -286,7 +294,7 @@ describe("useAgUiRuntime tool approvals", () => {
     expect(
       allToolCalls(runtime.current).find((p) => p.approval?.id === "int-1")
         ?.approval,
-    ).toEqual({ id: "int-1", resolution: "cancelled" });
+    ).toEqual(gateApproval({ resolution: "cancelled" }));
   });
 
   it("shows a bespoke resolved override as the decision it carried", async () => {
@@ -305,7 +313,7 @@ describe("useAgUiRuntime tool approvals", () => {
     expect(
       allToolCalls(runtime.current).find((p) => p.approval?.id === "int-1")
         ?.approval,
-    ).toEqual({ id: "int-1", approved: true });
+    ).toEqual(gateApproval({ approved: true }));
   });
 
   it("leaves the tool call ungated when the live interrupt carries a rejecting schema", async () => {
@@ -344,7 +352,7 @@ describe("useAgUiRuntime tool approvals", () => {
     ]);
     const [interrupt] = interruptsRef.current;
 
-    expect(allToolCalls(runtime.current)[0]!.approval).toEqual({ id: "int-1" });
+    expect(allToolCalls(runtime.current)[0]!.approval).toEqual(gateApproval());
     expect(interrupt?.responseSchema).toBeUndefined();
     expect(interrupt && readRawResponseSchema(interrupt)).toBeUndefined();
   });
@@ -412,7 +420,7 @@ describe("useAgUiRuntime tool approvals", () => {
     // `{answer:42}` went out for the first gate, so the approval recorded while
     // its sibling was still open cannot keep displaying a decision never sent.
     expect(allToolCalls(runtime.current).map((part) => part.approval)).toEqual([
-      { id: "int-1", resolution: "cancelled" },
+      gateApproval({ resolution: "cancelled" }),
       { id: "int-2", resolution: "cancelled" },
     ]);
   });
@@ -469,7 +477,7 @@ describe("useAgUiRuntime tool approvals", () => {
       type: "requires-action",
       reason: "interrupt",
     });
-    expect(allToolCalls(runtime.current)[0]!.approval).toEqual({ id: "int-1" });
+    expect(allToolCalls(runtime.current)[0]!.approval).toEqual(gateApproval());
     expect(execute).not.toHaveBeenCalled();
 
     await act(async () => {
@@ -542,14 +550,14 @@ describe("useAgUiRuntime tool approvals under an unsettled run", () => {
       content: [{ type: "text", text: "delete it" }],
     });
     await waitFor(() =>
-      expect(allToolCalls(result.current)[0]?.approval).toEqual({
-        id: "int-1",
-      }),
+      expect(allToolCalls(result.current)[0]?.approval).toEqual(gateApproval()),
     );
     expect(result.current.thread.getState().isRunning).toBe(true);
 
     await act(async () => {
-      gatedPart(result.current).respondToToolApproval({ approved: true });
+      await expect(
+        gatedPart(result.current).respondToToolApproval({ approved: true }),
+      ).rejects.toThrow("a run is already in progress");
     });
     await waitFor(() => expect(onError).toHaveBeenCalledTimes(1));
     expect((onError.mock.calls[0]![0] as Error).message).toContain(
@@ -557,7 +565,7 @@ describe("useAgUiRuntime tool approvals under an unsettled run", () => {
     );
     // The rejected click must not have recorded a decision, or the retry
     // below would report the approval as already decided.
-    expect(allToolCalls(result.current)[0]!.approval).toEqual({ id: "int-1" });
+    expect(allToolCalls(result.current)[0]!.approval).toEqual(gateApproval());
 
     await act(async () => {
       release();
@@ -662,7 +670,7 @@ describe("useAgUiRuntime tool approvals across a messages snapshot", () => {
     const gated = allToolCalls(result.current).filter(
       (part) => part.approval !== undefined,
     );
-    expect(gated.map((part) => part.approval)).toEqual([{ id: "int-1" }]);
+    expect(gated.map((part) => part.approval)).toEqual([gateApproval()]);
 
     await act(async () => {
       await gatedPart(result.current).respondToToolApproval({ approved: true });
@@ -795,7 +803,9 @@ describe("useAgUiRuntime approvals on a run that fails after the interrupt", ()
     await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
 
     await act(async () => {
-      gatedPart(result.current).respondToToolApproval({ approved: true });
+      await expect(
+        gatedPart(result.current).respondToToolApproval({ approved: true }),
+      ).rejects.toThrow("resume failed");
     });
     await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
     await waitFor(() => expect(onError).toHaveBeenCalled());
@@ -879,7 +889,7 @@ describe("useAgUiRuntime approvals that cannot be completed through the seam", (
     // The interrupt is closed, so the gate must not stay actionable: a click
     // would find no pending interrupt on the thread.
     expect(allToolCalls(runtime.current).map((part) => part.approval)).toEqual([
-      { id: "int-1", resolution: "cancelled" },
+      gateApproval({ resolution: "cancelled" }),
     ]);
   });
 });

package/src/useAgUiRuntime.history.test.tsx

@@ -0,0 +1,76 @@
+// @vitest-environment jsdom
+
+import { describe, it, expect, vi, afterEach } from "vitest";
+import { cleanup, renderHook, waitFor } from "@testing-library/react";
+import type { ThreadHistoryAdapter } from "@assistant-ui/core";
+import type { HttpAgent } from "@ag-ui/client";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+
+const agent = { runAgent: vi.fn(), abortRun: vi.fn() } as unknown as HttpAgent;
+
+function createHistory(): ThreadHistoryAdapter {
+  return {
+    load: vi.fn().mockResolvedValue({
+      headId: "restored",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "restored",
+            role: "user" as const,
+            content: [{ type: "text" as const, text: "hello" }],
+            createdAt: new Date(0),
+            metadata: { custom: {} },
+          },
+        },
+      ],
+    }),
+    append: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("useAgUiRuntime history", () => {
+  it("loads a history adapter that arrives on a later render", async () => {
+    const history = createHistory();
+    const { result, rerender } = renderHook(
+      ({ history }: { history?: ThreadHistoryAdapter }) =>
+        useAgUiRuntime({ agent, adapters: history ? { history } : {} }),
+      { initialProps: {} },
+    );
+
+    await waitFor(() =>
+      expect(result.current.thread.getState().isLoading).toBe(false),
+    );
+    expect(history.load).not.toHaveBeenCalled();
+    expect(result.current.thread.getState().messages).toEqual([]);
+
+    rerender({ history });
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    expect(history.load).toHaveBeenCalledOnce();
+  });
+
+  it("loads once when the adapter is present from the first render", async () => {
+    const history = createHistory();
+    const { result, rerender } = renderHook(() =>
+      useAgUiRuntime({ agent, adapters: { history } }),
+    );
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(["restored"]),
+    );
+    rerender();
+    rerender();
+    expect(history.load).toHaveBeenCalledOnce();
+  });
+});

package/src/useAgUiRuntime.toolDeferral.test.tsx

@@ -0,0 +1,186 @@
+// @vitest-environment jsdom
+
+import { describe, it, expect, vi, afterEach } from "vitest";
+import { act, cleanup, render, renderHook } from "@testing-library/react";
+import type { AssistantRuntime, ToolCallMessagePart } from "@assistant-ui/core";
+import {
+  AssistantRuntimeProvider,
+  useAssistantTool,
+} from "@assistant-ui/core/react";
+import type { HttpAgent } from "@ag-ui/client";
+import { z } from "zod";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+import type { AgUiInterrupt } from "./runtime/types";
+
+type Subscriber = Record<string, ((payload: any) => void) | undefined>;
+
+const GATE: AgUiInterrupt = {
+  id: "int-1",
+  reason: "tool_call",
+  toolCallId: "tc-1",
+  message: "Delete /tmp/a?",
+};
+
+const toolCalls = (runtime: AssistantRuntime): ToolCallMessagePart[] =>
+  runtime.thread
+    .getState()
+    .messages.flatMap((message) => message.content as readonly unknown[])
+    .filter(
+      (part): part is ToolCallMessagePart =>
+        (part as ToolCallMessagePart).type === "tool-call",
+    );
+
+const flush = () =>
+  act(async () => void (await new Promise((r) => setTimeout(r, 0))));
+
+/**
+ * Holds the run open after the tool call's arguments complete, so the window
+ * between complete args and the run's outcome is observable as its own render
+ * rather than coalescing with the tail.
+ */
+const renderHeldRun = async (tail: (subscriber: Subscriber) => void) => {
+  let releaseTail!: () => void;
+  const held = new Promise<void>((r) => (releaseTail = r));
+  let argsDelivered!: () => void;
+  const argsDeliveredP = new Promise<void>((r) => (argsDelivered = r));
+
+  const runAgent = vi.fn(async (_input: unknown, subscriber: Subscriber) => {
+    if (runAgent.mock.calls.length > 1) {
+      subscriber.onRunFinalized?.(undefined);
+      return;
+    }
+    subscriber.onToolCallStartEvent?.({
+      event: {
+        type: "TOOL_CALL_START",
+        toolCallId: "tc-1",
+        toolCallName: "delete_file",
+      },
+    });
+    subscriber.onToolCallArgsEvent?.({
+      event: {
+        type: "TOOL_CALL_ARGS",
+        toolCallId: "tc-1",
+        delta: '{"path":"/tmp/a"}',
+      },
+    });
+    subscriber.onToolCallEndEvent?.({
+      event: { type: "TOOL_CALL_END", toolCallId: "tc-1" },
+    });
+    argsDelivered();
+    await held;
+    tail(subscriber);
+    subscriber.onRunFinalized?.(undefined);
+  });
+  const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+
+  const execute = vi.fn(async () => ({ deleted: true }));
+  const DeleteFileTool = () => {
+    useAssistantTool({
+      toolName: "delete_file",
+      description: "delete a file",
+      parameters: z.object({ path: z.string() }),
+      execute,
+    });
+    return null;
+  };
+
+  const { result } = renderHook(() => useAgUiRuntime({ agent }));
+  render(
+    <AssistantRuntimeProvider runtime={result.current}>
+      <DeleteFileTool />
+    </AssistantRuntimeProvider>,
+  );
+  await flush();
+
+  await act(async () => {
+    void result.current.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "delete it" }],
+    });
+    await argsDeliveredP;
+  });
+  await flush();
+
+  return {
+    runtime: result,
+    execute,
+    settle: async () => {
+      await act(async () => {
+        releaseTail();
+        await new Promise((r) => setTimeout(r, 0));
+      });
+      await flush();
+    },
+  };
+};
+
+afterEach(() => cleanup());
+
+describe("useAgUiRuntime frontend tool deferral", () => {
+  it("does not run a frontend execute on a call the provider is about to gate", async () => {
+    const { runtime, execute, settle } = await renderHeldRun((subscriber) => {
+      subscriber.onRunFinishedEvent?.({
+        event: {
+          type: "RUN_FINISHED",
+          runId: "run-1",
+          outcome: { type: "interrupt", interrupts: [GATE] },
+        },
+      });
+    });
+
+    const inWindow = toolCalls(runtime.current)[0]!;
+    expect(inWindow.argsText).toBe('{"path":"/tmp/a"}');
+    expect(inWindow.approval).toBeUndefined();
+    expect(inWindow.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    const gated = toolCalls(runtime.current)[0]!;
+    expect(gated.approval).toEqual({
+      id: "int-1",
+      prompt: "Delete /tmp/a?",
+    });
+    expect(gated.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("does not run a frontend execute on a call the agent answers itself", async () => {
+    const { runtime, execute, settle } = await renderHeldRun((subscriber) => {
+      subscriber.onToolCallResultEvent?.({
+        event: {
+          type: "TOOL_CALL_RESULT",
+          messageId: "m-result",
+          toolCallId: "tc-1",
+          content: "server did the delete",
+        },
+      });
+      subscriber.onRunFinishedEvent?.({
+        event: { type: "RUN_FINISHED", runId: "run-1" },
+      });
+    });
+
+    expect(toolCalls(runtime.current)[0]!.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    expect(toolCalls(runtime.current)[0]!.result).toBe("server did the delete");
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("runs a frontend execute once the run settles with no outcome for the call", async () => {
+    const { runtime, execute, settle } = await renderHeldRun((subscriber) => {
+      subscriber.onRunFinishedEvent?.({
+        event: { type: "RUN_FINISHED", runId: "run-1" },
+      });
+    });
+
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    expect(execute).toHaveBeenCalledTimes(1);
+    expect(toolCalls(runtime.current)[0]!.result).toEqual({ deleted: true });
+  });
+});

package/src/useAgUiRuntime.ts

@@ -81,14 +81,16 @@ export function useAgUiRuntime(
   }
 
   const core = coreRef.current;
-  core.updateOptions({
-    agent: options.agent,
-    logger,
-    showThinking: options.showThinking ?? true,
-    autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
-    ...(options.onError && { onError: options.onError }),
-    ...(options.onCancel && { onCancel: options.onCancel }),
-    ...(historyAdapter && { history: historyAdapter }),
+  useEffect(() => {
+    core.updateOptions({
+      agent: options.agent,
+      logger,
+      showThinking: options.showThinking ?? true,
+      autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
+      ...(options.onError && { onError: options.onError }),
+      ...(options.onCancel && { onCancel: options.onCancel }),
+      ...(historyAdapter && { history: historyAdapter }),
+    });
   });
 
   const [toolStatuses, setToolStatuses] = useState<
@@ -228,7 +230,7 @@ export function useAgUiRuntime(
         isLoading: core.isLoading,
         messageRepository: core.getMessageRepository(),
         state: core.getState(),
-        isRunning,
+        isRunning: core.isRunning(),
         extras: agUiExtras.provide({
           interrupts:
             core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,
@@ -258,6 +260,7 @@ export function useAgUiRuntime(
         onRespondToToolApproval: (options) =>
           core.respondToToolApproval(options).catch((error: unknown) => {
             core.reportError(error);
+            throw error;
           }),
         onResume: (config) => core.resume(config),
         setMessages: (messages: readonly ThreadMessage[]) =>