@assistant-ui/react-ag-ui 0.0.62 → 0.0.64

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

@@ -17,6 +17,7 @@ import {
   type A2uiState,
   type A2uiSurfaceState,
 } from "@assistant-ui/react-generative-ui/a2ui";
+import jsonpatch, { type Operation } from "fast-json-patch";
 import { readMcpAppResourceUri } from "../mcp-tool-result";
 import { projectAgUiToolApprovals } from "./tool-approval";
 import type { AgUiEvent, AgUiInterrupt } from "../types";
@@ -631,12 +632,51 @@ export class RunAggregator {
         break;
       }
 
+      case "ACTIVITY_DELTA": {
+        this.handleActivityDelta(event);
+        break;
+      }
+
       default: {
         this.logger.debug?.("[agui] aggregator ignored event", event);
       }
     }
   }
 
+  private handleActivityDelta(
+    event: Extract<AgUiEvent, { type: "ACTIVITY_DELTA" }>,
+  ): void {
+    const scope = this.scopeOf(event);
+    const existing = this.activityParts.get(
+      this.partKey(scope, `message:${event.messageId}`),
+    );
+    if (!existing) {
+      this.logger.debug?.("[agui] activity delta has no snapshot", event);
+      return;
+    }
+    if (event.patch.length === 0) return;
+    try {
+      const result = jsonpatch.applyPatch(
+        existing.value,
+        event.patch as Operation[],
+        /* validateOperation */ true,
+        /* mutateDocument */ false,
+      );
+      if (!isPlainObject(result.newDocument)) {
+        this.logger.debug?.(
+          "[agui] activity delta produced non-object content",
+          event,
+        );
+        return;
+      }
+      existing.name = `agui-activity/${event.activityType}`;
+      existing.value = result.newDocument;
+      this.emit();
+    } catch (error) {
+      this.logger.error?.("[agui] failed to apply activity delta", error);
+    }
+  }
+
   private handleA2uiActivitySnapshot(
     event: Extract<AgUiEvent, { type: "ACTIVITY_SNAPSHOT" }>,
   ): void {

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

@@ -0,0 +1,48 @@
+import { describe, it, expect, vi } from "vitest";
+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 });
+  return { subscriber, dispatch, logger };
+};
+
+describe("createAgUiSubscriber", () => {
+  it("dispatches ACTIVITY_DELTA as a typed event", () => {
+    const { subscriber, dispatch } = createSubscriber();
+    subscriber.onActivityDeltaEvent?.({
+      event: {
+        type: "ACTIVITY_DELTA",
+        messageId: "act-1",
+        activityType: "progress",
+        patch: [{ op: "replace", path: "/step", value: 2 }],
+      },
+    });
+
+    expect(dispatch).toHaveBeenCalledWith({
+      type: "ACTIVITY_DELTA",
+      messageId: "act-1",
+      activityType: "progress",
+      patch: [{ op: "replace", path: "/step", value: 2 }],
+    });
+  });
+
+  it("logs step boundaries and dispatches nothing", () => {
+    const { subscriber, dispatch, logger } = createSubscriber();
+    const started = { type: "STEP_STARTED", stepName: "plan" };
+    const finished = { type: "STEP_FINISHED", stepName: "plan" };
+    subscriber.onStepStartedEvent?.({ event: started });
+    subscriber.onStepFinishedEvent?.({ event: finished });
+
+    expect(dispatch).not.toHaveBeenCalled();
+    expect(logger.debug).toHaveBeenCalledWith(
+      "[agui] step boundary ignored",
+      started,
+    );
+    expect(logger.debug).toHaveBeenCalledWith(
+      "[agui] step boundary ignored",
+      finished,
+    );
+  });
+});

package/src/runtime/adapter/subscriber.ts

@@ -29,6 +29,7 @@ type Subscriber = {
   onToolCallChunkEvent?: (payload: { event: unknown }) => void;
   onToolCallResultEvent?: (payload: { event: unknown }) => void;
   onActivitySnapshotEvent?: (payload: { event: unknown }) => void;
+  onActivityDeltaEvent?: (payload: { event: unknown }) => void;
   onStateSnapshotEvent?: (payload: { event: unknown }) => void;
   onStateDeltaEvent?: (payload: { event: unknown }) => void;
   onMessagesSnapshotEvent?: (payload: { event: unknown }) => void;
@@ -39,6 +40,8 @@ type Subscriber = {
   onSubagentErrorEvent?: (payload: { event: unknown }) => void;
   onRunFinishedEvent?: (payload: { event: unknown }) => void;
   onRunErrorEvent?: (payload: { event: unknown }) => void;
+  onStepStartedEvent?: (payload: { event: unknown }) => void;
+  onStepFinishedEvent?: (payload: { event: unknown }) => void;
   onRunFinalized?: () => void;
   onRunFailed?: (payload: { error: Error }) => void;
 };
@@ -142,6 +145,8 @@ export const createAgUiSubscriber = (
       dispatchIfValid(event, "TOOL_CALL_RESULT"),
     onActivitySnapshotEvent: ({ event }) =>
       dispatchIfValid(event, "ACTIVITY_SNAPSHOT"),
+    onActivityDeltaEvent: ({ event }) =>
+      dispatchIfValid(event, "ACTIVITY_DELTA"),
     onStateSnapshotEvent: ({ event }) =>
       dispatchIfValid(event, "STATE_SNAPSHOT"),
     onStateDeltaEvent: ({ event }) => dispatchIfValid(event, "STATE_DELTA"),
@@ -174,11 +179,16 @@ export const createAgUiSubscriber = (
       );
       onRunFailed?.(error);
     },
+    onStepStartedEvent: ({ event }) =>
+      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;
       onRunFailed?.(error);
       if (isAbortError(error)) {

package/src/runtime/event-parser.test.ts

@@ -0,0 +1,44 @@
+import { describe, it, expect, vi } from "vitest";
+import { parseAgUiEvent } from "./event-parser";
+
+const delta = {
+  messageId: "act-1",
+  activityType: "progress",
+  patch: [{ op: "replace", path: "/step", value: 2 }],
+};
+
+describe("parseAgUiEvent ACTIVITY_DELTA", () => {
+  it("maps the client shape", () => {
+    expect(
+      parseAgUiEvent({
+        type: "ACTIVITY_DELTA",
+        ...delta,
+        subagentRunId: "sub",
+      }),
+    ).toEqual({ type: "ACTIVITY_DELTA", ...delta, subagentRunId: "sub" });
+  });
+
+  it("maps the bare wire shape without optional fields", () => {
+    expect(
+      parseAgUiEvent({ type: "ACTIVITY_DELTA", ...delta, timestamp: 1 }),
+    ).toEqual({ type: "ACTIVITY_DELTA", ...delta });
+  });
+
+  it("rejects a missing messageId", () => {
+    const logger = { debug: vi.fn(), error: vi.fn() };
+    const { messageId: _, ...rest } = delta;
+    expect(
+      parseAgUiEvent({ type: "ACTIVITY_DELTA", ...rest }, { logger }),
+    ).toBe(null);
+    expect(logger.debug).toHaveBeenCalledWith(
+      "[agui] ACTIVITY_DELTA missing messageId, activityType or patch array",
+      expect.objectContaining({ activityType: "progress" }),
+    );
+  });
+
+  it("rejects a patch that is not an array", () => {
+    expect(
+      parseAgUiEvent({ type: "ACTIVITY_DELTA", ...delta, patch: "nope" }),
+    ).toBe(null);
+  });
+});

package/src/runtime/event-parser.ts

@@ -389,6 +389,22 @@ export const parseAgUiEvent = (
         },
       );
     }
+    case "ACTIVITY_DELTA": {
+      const messageId = getString("messageId");
+      const activityType = getString("activityType");
+      if (!messageId || !activityType || !Array.isArray(payload.patch)) {
+        return reject("missing messageId, activityType or patch array");
+      }
+      return withOptional(
+        {
+          type: "ACTIVITY_DELTA" as const,
+          messageId,
+          activityType,
+          patch: payload.patch as any[],
+        },
+        { subagentRunId: getString("subagentRunId") },
+      );
+    }
     case "RAW":
       return withOptional(
         { type: "RAW" as const, event: payload.event },

package/src/runtime/types.ts

@@ -249,6 +249,13 @@ export type AgUiEvent =
       replace?: boolean;
       subagentRunId?: string;
     }
+  | {
+      type: "ACTIVITY_DELTA";
+      messageId: string;
+      activityType: string;
+      patch: any[];
+      subagentRunId?: string;
+    }
   | { type: "RAW"; event: any; source?: string }
   | { type: "CUSTOM"; name: string; value: any }
   | { type: "STATE_SNAPSHOT"; snapshot: any }

package/src/useAgUiRuntime.fast-refresh.test.tsx

@@ -0,0 +1,276 @@
+// @vitest-environment jsdom
+
+import { act, Activity, version, type ComponentType } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+import type { HttpAgent } from "@ag-ui/client";
+import type {
+  AssistantRuntime,
+  MessageQueueController,
+} from "@assistant-ui/core";
+import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
+import type { UseAgUiRuntimeOptions } from "./runtime/types";
+import { useAgUiRuntime, type AgUiAssistantRuntime } from "./useAgUiRuntime";
+
+const onReact18 = version.startsWith("18.");
+
+const base = vi.hoisted(() => ({ version: 0 }));
+const queue = vi.hoisted(() => ({
+  controller: null as MessageQueueController | null,
+  busy: vi.fn(),
+  idle: vi.fn(),
+}));
+vi.mock("@assistant-ui/core", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@assistant-ui/core")>();
+  return {
+    ...actual,
+    createMessageQueue: (
+      ...args: Parameters<typeof actual.createMessageQueue>
+    ) => {
+      const controller = actual.createMessageQueue(...args);
+      queue.controller = {
+        ...controller,
+        notifyBusy: () => {
+          queue.busy();
+          controller.notifyBusy();
+        },
+        notifyIdle: () => {
+          queue.idle();
+          controller.notifyIdle();
+        },
+      };
+      return queue.controller;
+    },
+  };
+});
+vi.mock("@assistant-ui/core/react", async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import("@assistant-ui/core/react")>();
+  const derived = new WeakMap<AssistantRuntime, AssistantRuntime>();
+  return {
+    ...actual,
+    useExternalStoreRuntime: (
+      store: Parameters<typeof actual.useExternalStoreRuntime>[0],
+    ) => {
+      const runtime = actual.useExternalStoreRuntime(store);
+      if (base.version === 0) return runtime;
+      let next = derived.get(runtime);
+      if (!next) {
+        next = Object.create(runtime) as AssistantRuntime;
+        derived.set(runtime, next);
+      }
+      return next;
+    },
+  };
+});
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+
+afterEach(() => {
+  cleanup();
+  fiberRoots.clear();
+  vi.restoreAllMocks();
+  base.version = 0;
+  queue.controller = null;
+});
+afterAll(() => vi.unstubAllGlobals());
+
+let options: UseAgUiRuntimeOptions;
+let runtime: AgUiAssistantRuntime;
+const createProbe = () => {
+  const Before = () => {
+    runtime = useAgUiRuntime(options);
+    return null;
+  };
+  const After = () => {
+    runtime = useAgUiRuntime(options);
+    return null;
+  };
+  return { Before, After };
+};
+
+const refresh = async (Before: unknown, After: unknown) => {
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const root of fiberRoots) {
+      renderer!.scheduleRefresh(root, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+};
+
+it("keeps the wrapper across refresh and replaces it when the base runtime changes", async () => {
+  const agent = {
+    runAgent: vi.fn(),
+    abortRun: vi.fn(),
+  } as unknown as HttpAgent;
+  options = { agent };
+  const { Before, After } = createProbe();
+  const view = render(<Before />);
+  const wrapper = runtime;
+  await refresh(Before, After);
+  expect(runtime).toBe(wrapper);
+
+  base.version = 1;
+  view.rerender(<After />);
+  await act(async () => {});
+  expect(runtime).not.toBe(wrapper);
+});
+
+it("emits queue edges once across refresh and dispatches the queued message", async () => {
+  let finishFirst!: () => void;
+  const firstRun = new Promise<void>((resolve) => {
+    finishFirst = resolve;
+  });
+  const runAgent = vi.fn(
+    async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {
+      if (runAgent.mock.calls.length === 1) await firstRun;
+      subscriber.onRunFinalized?.();
+    },
+  );
+  options = {
+    agent: { runAgent, abortRun: vi.fn() } as unknown as HttpAgent,
+    unstable_enableMessageQueue: true,
+  };
+  const { Before, After } = createProbe();
+  const view = render(<Before />);
+  expect(queue.idle).toHaveBeenCalledTimes(1);
+
+  await act(async () => {
+    await runtime.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "first" }],
+    });
+  });
+  await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+  expect(queue.busy).toHaveBeenCalledTimes(1);
+
+  await act(async () => {
+    await runtime.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "second" }],
+      parentId: runtime.thread.getState().messages.at(-1)?.id ?? null,
+    });
+  });
+  expect(runAgent).toHaveBeenCalledTimes(1);
+  expect(
+    [
+      ...queue.controller!.adapter.items,
+      ...queue.controller!.adapter.steerItems,
+    ].map((item) => item.prompt),
+  ).toEqual(["second"]);
+
+  await refresh(Before, After);
+  expect(queue.busy).toHaveBeenCalledTimes(1);
+  expect(queue.idle).toHaveBeenCalledTimes(1);
+
+  await act(async () => {
+    finishFirst();
+    await firstRun;
+  });
+  await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+  await waitFor(() =>
+    expect([
+      ...queue.controller!.adapter.items,
+      ...queue.controller!.adapter.steerItems,
+    ]).toHaveLength(0),
+  );
+  await waitFor(() => expect(queue.busy).toHaveBeenCalledTimes(2));
+  await waitFor(() => expect(queue.idle).toHaveBeenCalledTimes(3));
+
+  await refresh(Before, After);
+  expect(queue.busy).toHaveBeenCalledTimes(2);
+  expect(queue.idle).toHaveBeenCalledTimes(3);
+
+  const firstController = queue.controller;
+  options = { ...options, unstable_enableMessageQueue: false };
+  view.rerender(<After />);
+  options = { ...options, unstable_enableMessageQueue: true };
+  view.rerender(<After />);
+  expect(queue.controller).not.toBe(firstController);
+  expect(queue.idle).toHaveBeenCalledTimes(4);
+});
+
+const expectInFlightRunAcrossRefresh = async (
+  exit: "hidden" | "unmount",
+  inActivity: boolean,
+) => {
+  let resolveRun!: () => void;
+  const run = new Promise<void>((resolve) => {
+    resolveRun = resolve;
+  });
+  const agent = {
+    runAgent: vi.fn(() => run),
+    abortRun: vi.fn(),
+  } as unknown as HttpAgent;
+  const attach = vi.spyOn(AgUiThreadRuntimeCore.prototype, "attachRuntime");
+  const detach = vi.spyOn(AgUiThreadRuntimeCore.prototype, "detachRuntime");
+  options = { agent };
+  const { Before, After } = createProbe();
+  const Shell: ComponentType<{ mode: "visible" | "hidden" }> = inActivity
+    ? ({ mode }) => (
+        <Activity mode={mode}>
+          <Before />
+        </Activity>
+      )
+    : () => <Before />;
+  const view = render(<Shell mode="visible" />);
+  const wrapper = runtime;
+  const core = attach.mock.instances[0]!;
+  act(() => {
+    void runtime.thread.append("hello");
+  });
+  await waitFor(() => expect(agent.runAgent).toHaveBeenCalledOnce());
+
+  await refresh(Before, After);
+  expect(agent.abortRun).not.toHaveBeenCalled();
+  expect(detach).not.toHaveBeenCalled();
+  expect(runtime).toBe(wrapper);
+  expect(attach.mock.instances.at(-1)).toBe(core);
+
+  if (exit === "hidden") view.rerender(<Shell mode="hidden" />);
+  else view.unmount();
+  await act(async () => {});
+  expect(detach).toHaveBeenCalledOnce();
+  expect(agent.abortRun).toHaveBeenCalledOnce();
+  resolveRun();
+  await act(async () => {
+    await run;
+  });
+};
+
+// Activity is React 19 only.
+it.skipIf(onReact18).each(["hidden", "unmount"] as const)(
+  "keeps an in-flight run across refresh and detaches when %s",
+  (exit) => expectInFlightRunAcrossRefresh(exit, true),
+);
+
+it("keeps an in-flight run across refresh and detaches when unmounted outside an Activity", () =>
+  expectInFlightRunAcrossRefresh("unmount", false));

package/src/useAgUiRuntime.ts

@@ -14,6 +14,7 @@ import {
   useExternalStoreSharedOptions,
   useRuntimeAdapters,
 } from "@assistant-ui/core/react";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import { createMessageQueue } from "@assistant-ui/core";
 import type {
   MessageQueueController,
@@ -113,6 +114,10 @@ export function useAgUiRuntime(
   // A dispatch whose count never moves was never observed as busy, so no
   // falling edge is coming to release the queue.
   const busyEdgesRef = useRef(0);
+  const lastQueueEdgeRef = useRef<{
+    controller: MessageQueueController;
+    busy: boolean;
+  } | null>(null);
   if (options.unstable_enableMessageQueue && !queueRef.current) {
     queueRef.current = createMessageQueue({
       run: (message) => {
@@ -171,11 +176,18 @@ export function useAgUiRuntime(
   // refusal would lose the message instead of keeping it visible.
   const queueBusy = isRunning || core.getPendingInterrupts() !== null;
   useEffect(() => {
+    if (!queueController) return;
+    if (
+      lastQueueEdgeRef.current?.controller === queueController &&
+      lastQueueEdgeRef.current.busy === queueBusy
+    )
+      return;
+    lastQueueEdgeRef.current = { controller: queueController, busy: queueBusy };
     if (queueBusy) {
       busyEdgesRef.current++;
-      queueController?.notifyBusy();
+      queueController.notifyBusy();
     } else {
-      queueController?.notifyIdle();
+      queueController.notifyIdle();
     }
   }, [queueBusy, queueController]);
 
@@ -309,16 +321,26 @@ export function useAgUiRuntime(
 
   const baseRuntime = useExternalStoreRuntime(store);
 
-  const runtime = useMemo<AgUiAssistantRuntime>(() => {
+  const createRuntime = (): AgUiAssistantRuntime => {
     const wrapper = Object.create(baseRuntime) as AgUiAssistantRuntime;
     wrapper.unstable_getPendingInterrupts = () =>
       core.getPendingInterrupts()?.interrupts ?? [];
     wrapper.unstable_submitInterruptResponses = (responses) =>
       core.submitInterruptResponses(responses);
     return wrapper;
-  }, [baseRuntime, core]);
+  };
+  const [pinnedRuntime, setPinnedRuntime] = useState(() => ({
+    baseRuntime,
+    runtime: createRuntime(),
+  }));
+  let currentRuntime = pinnedRuntime;
+  if (pinnedRuntime.baseRuntime !== baseRuntime) {
+    currentRuntime = { baseRuntime, runtime: createRuntime() };
+    setPinnedRuntime(currentRuntime);
+  }
+  const runtime = currentRuntime.runtime;
 
-  useEffect(() => {
+  useReplaySafeEffect(() => {
     core.attachRuntime(runtime);
     return () => {
       core.detachRuntime();