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

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,6 +179,10 @@ 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 });

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,170 @@
+// @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 } 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 }));
+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;
+});
+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);
+});
+
+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,
@@ -309,16 +310,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();