@assistant-ui/core 0.3.13 → 0.3.15

package/src/react/runtimes/assistant-transport/runManager.ts

@@ -0,0 +1,124 @@
+import { useCallback, useEffect, useRef, useState } from "react";
+import { useLatestRef } from "./useLatestRef";
+
+export type RunManager = Readonly<{
+  isRunning: boolean;
+  schedule: () => void;
+  cancel: () => void;
+}>;
+
+const disposeReason = Symbol("assistant-transport-dispose");
+
+type LifecycleCallbackName = "onCancel" | "onError" | "onFinish";
+
+const reportCallbackError = (name: LifecycleCallbackName, error: unknown) => {
+  console.error(
+    `[assistant-ui] Assistant transport ${name} callback threw an error`,
+    error,
+  );
+};
+
+const invokeCallback = async (
+  name: LifecycleCallbackName,
+  callback: (() => unknown) | undefined,
+): Promise<void> => {
+  try {
+    await callback?.();
+  } catch (error) {
+    try {
+      reportCallbackError(name, error);
+    } catch {}
+  }
+};
+
+export function useRunManager(config: {
+  onRun: (signal: AbortSignal) => Promise<void>;
+  onFinish?: (() => void) | undefined;
+  onCancel?: (() => void) | undefined;
+  onError?: ((error: Error) => void | Promise<void>) | undefined;
+}): RunManager {
+  const [isRunning, setIsRunning] = useState(false);
+  const stateRef = useRef({
+    pending: false,
+    disposed: false,
+    abortController: null as AbortController | null,
+  });
+  const onRunRef = useLatestRef(config.onRun);
+  const onFinishRef = useLatestRef(config.onFinish);
+  const onCancelRef = useLatestRef(config.onCancel);
+  const onErrorRef = useLatestRef(config.onError);
+
+  const startRun = useCallback(() => {
+    setIsRunning(true);
+    stateRef.current.pending = false;
+    const ac = new AbortController();
+    stateRef.current.abortController = ac;
+
+    // The dispose marker rides on the abort reason so a settling run detects
+    // its own disposal even after an effect cycle re-arms the manager.
+    const disposeAborted = () => ac.signal.reason === disposeReason;
+
+    queueMicrotask(async () => {
+      try {
+        if (!disposeAborted()) {
+          await onRunRef.current(ac.signal);
+        }
+      } catch (error) {
+        if (!disposeAborted() && !stateRef.current.disposed) {
+          stateRef.current.pending = false;
+          if (ac.signal.aborted) {
+            void invokeCallback("onCancel", onCancelRef.current);
+          } else {
+            await invokeCallback("onError", () =>
+              onErrorRef.current?.(error as Error),
+            );
+          }
+        }
+      } finally {
+        if (!disposeAborted() && !stateRef.current.disposed) {
+          void invokeCallback("onFinish", onFinishRef.current);
+        }
+        if (!stateRef.current.disposed && stateRef.current.pending) {
+          startRun();
+        } else {
+          setIsRunning(false);
+          if (stateRef.current.abortController === ac) {
+            stateRef.current.abortController = null;
+          }
+        }
+      }
+    });
+  }, [onRunRef, onFinishRef, onErrorRef, onCancelRef]);
+
+  const schedule = useCallback(() => {
+    if (stateRef.current.disposed) return;
+    if (stateRef.current.abortController) {
+      // Coalesce multiple schedules while running into a single follow-up run.
+      stateRef.current.pending = true;
+      return;
+    }
+    startRun();
+  }, [startRun]);
+
+  // Strict-mode effects run setup / cleanup / setup on the same instance;
+  // arming on setup undoes the cleanup's dispose.
+  useEffect(() => {
+    stateRef.current.disposed = false;
+    return () => {
+      stateRef.current.disposed = true;
+      stateRef.current.pending = false;
+      stateRef.current.abortController?.abort(disposeReason);
+    };
+  }, []);
+
+  const cancel = useCallback(() => {
+    stateRef.current.pending = false;
+    stateRef.current.abortController?.abort();
+  }, []);
+
+  return {
+    isRunning,
+    schedule,
+    cancel,
+  };
+}

package/src/react/runtimes/assistant-transport/transport-delivery.test.tsx

@@ -0,0 +1,79 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { FC } from "react";
+import { useAssistantTransportRuntime } from "./useAssistantTransportRuntime";
+import type { AssistantRuntime } from "../../../runtime/api/assistant-runtime";
+import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
+import type {
+  AssistantTransportCommand,
+  AssistantTransportStateConverter,
+} from "./types";
+
+const emptySuccessfulResponse = () =>
+  new Response(
+    new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.close();
+      },
+    }),
+    { status: 200 },
+  );
+
+const setupRuntime = () => {
+  const fetchMock = vi.fn(async () => emptySuccessfulResponse());
+  vi.stubGlobal("fetch", fetchMock);
+
+  const pendingRef: { current: AssistantTransportCommand[] } = { current: [] };
+  const converter: AssistantTransportStateConverter<Record<string, never>> = (
+    _state,
+    { pendingCommands, isSending },
+  ) => {
+    pendingRef.current = pendingCommands;
+    return { messages: [], isRunning: isSending };
+  };
+
+  const runtimeRef: { current: AssistantRuntime | null } = { current: null };
+  const App: FC = () => {
+    const runtime = useAssistantTransportRuntime({
+      initialState: {},
+      api: "http://localhost/api",
+      converter,
+      headers: {},
+    });
+    runtimeRef.current = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        {null}
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  return { App, fetchMock, pendingRef, runtimeRef };
+};
+
+describe("assistant transport delivery contracts", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("clears in-transit commands when a run succeeds without state chunks", async () => {
+    const { App, fetchMock, pendingRef, runtimeRef } = setupRuntime();
+
+    await act(async () => {
+      render(<App />);
+    });
+    await waitFor(() => expect(runtimeRef.current).not.toBeNull());
+
+    await act(async () => {
+      runtimeRef.current!.thread.append("m1");
+    });
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(runtimeRef.current!.thread.getState().isRunning).toBe(false),
+    );
+    await waitFor(() => expect(pendingRef.current).toHaveLength(0));
+  });
+});

package/src/react/runtimes/assistant-transport/transport-parentid.test.tsx

@@ -0,0 +1,97 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { FC } from "react";
+import { useAssistantTransportRuntime } from "./useAssistantTransportRuntime";
+import type { AssistantRuntime } from "../../../runtime/api/assistant-runtime";
+import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
+import type {
+  AssistantTransportCommand,
+  AssistantTransportStateConverter,
+} from "./types";
+
+const emptySuccessfulResponse = () =>
+  new Response(
+    new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.close();
+      },
+    }),
+    { status: 200 },
+  );
+
+const setupRuntime = () => {
+  const requestBodies: Record<string, unknown>[] = [];
+  const fetchMock = vi.fn(
+    async (_url: RequestInfo | URL, init?: RequestInit) => {
+      requestBodies.push(JSON.parse(init!.body as string));
+      return emptySuccessfulResponse();
+    },
+  );
+  vi.stubGlobal("fetch", fetchMock);
+
+  const converter: AssistantTransportStateConverter<Record<string, never>> = (
+    _state,
+    { isSending },
+  ) => ({ messages: [], isRunning: isSending });
+
+  const runtimeRef: { current: AssistantRuntime | null } = { current: null };
+  const App: FC = () => {
+    const runtime = useAssistantTransportRuntime({
+      initialState: {},
+      api: "http://localhost/api",
+      converter,
+      headers: {},
+    });
+    runtimeRef.current = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        {null}
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  return { App, fetchMock, requestBodies, runtimeRef };
+};
+
+describe("assistant transport parentId lifetime", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("consumes parentId per run: the append's run carries it, a later sendCommand run omits it", async () => {
+    const { App, fetchMock, requestBodies, runtimeRef } = setupRuntime();
+
+    await act(async () => {
+      render(<App />);
+    });
+    await waitFor(() => expect(runtimeRef.current).not.toBeNull());
+
+    await act(async () => {
+      runtimeRef.current!.thread.append("m1");
+    });
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(runtimeRef.current!.thread.getState().isRunning).toBe(false),
+    );
+    expect(Object.hasOwn(requestBodies[0]!, "parentId")).toBe(true);
+
+    const extras = runtimeRef.current!.thread.getState().extras as {
+      sendCommand: (command: AssistantTransportCommand) => void;
+    };
+    await act(async () => {
+      extras.sendCommand({
+        type: "add-tool-result",
+        toolCallId: "t1",
+        toolName: "tool",
+        result: {},
+        isError: false,
+      });
+    });
+
+    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
+    expect(Object.hasOwn(requestBodies[1]!, "parentId")).toBe(false);
+  });
+});

package/src/react/runtimes/assistant-transport/transport-scheduling.test.ts

@@ -0,0 +1,414 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { createElement, StrictMode, useRef } from "react";
+import { useCommandQueue } from "./commandQueue";
+import { useRunManager } from "./runManager";
+import type { AssistantTransportCommand } from "./types";
+
+const createMessageCommand = (id: string): AssistantTransportCommand => ({
+  type: "add-message",
+  message: {
+    role: "user",
+    parts: [{ type: "text", text: id }],
+  },
+  parentId: null,
+  sourceId: null,
+});
+
+const createDeferred = () => {
+  let resolve!: () => void;
+  const promise = new Promise<void>((res) => {
+    resolve = res;
+  });
+  return { promise, resolve };
+};
+
+const captureUnhandledRejections = async (
+  callback: () => Promise<void>,
+): Promise<unknown[]> => {
+  const reasons: unknown[] = [];
+  const listener = (reason: unknown) => reasons.push(reason);
+  process.on("unhandledRejection", listener);
+  try {
+    await callback();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    return reasons;
+  } finally {
+    process.off("unhandledRejection", listener);
+  }
+};
+
+const useTransportSchedulingHarness = (
+  opts: {
+    onRun?: (signal: AbortSignal) => Promise<void> | void;
+    onCancel?: (commands: AssistantTransportCommand[]) => void;
+    onError?: (commands: AssistantTransportCommand[]) => void | Promise<void>;
+    onFinish?: () => void;
+  } = {},
+) => {
+  const commandQueueRef = useRef<ReturnType<typeof useCommandQueue> | null>(
+    null,
+  );
+  const runBatchesRef = useRef<AssistantTransportCommand[][]>([]);
+
+  const runManager = useRunManager({
+    onRun: async (signal) => {
+      const batch = commandQueueRef.current!.flush();
+      runBatchesRef.current.push(batch);
+      await opts.onRun?.(signal);
+    },
+    onCancel: () => {
+      const queue = commandQueueRef.current!;
+      const commands = [...queue.state.inTransit, ...queue.state.queued];
+      queue.reset();
+      opts.onCancel?.(commands);
+    },
+    onError: async () => {
+      const queue = commandQueueRef.current!;
+      await opts.onError?.([...queue.state.inTransit]);
+    },
+    onFinish: () => opts.onFinish?.(),
+  });
+
+  const commandQueue = useCommandQueue({
+    onQueue: () => runManager.schedule(),
+  });
+  commandQueueRef.current = commandQueue;
+
+  return {
+    commandQueue,
+    runManager,
+    runBatchesRef,
+  };
+};
+
+describe("assistant transport scheduling contracts", () => {
+  it("runs in single-flight mode and schedules exactly one follow-up run", async () => {
+    const gate = createDeferred();
+    const { result } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: () => gate.promise,
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m1"));
+      result.current.commandQueue.enqueue(createMessageCommand("m2"));
+    });
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+    expect(result.current.runBatchesRef.current[0]).toHaveLength(2);
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m3"));
+    });
+
+    await Promise.resolve();
+    expect(result.current.runBatchesRef.current).toHaveLength(1);
+
+    gate.resolve();
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(2);
+    });
+    expect(result.current.runBatchesRef.current[1]).toHaveLength(1);
+  });
+
+  it("can enqueue without scheduling until a run is started", async () => {
+    const { result } = renderHook(() => useTransportSchedulingHarness());
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("staged"), {
+        schedule: false,
+      });
+    });
+
+    await Promise.resolve();
+    expect(result.current.runBatchesRef.current).toHaveLength(0);
+
+    act(() => {
+      result.current.runManager.schedule();
+    });
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+    expect(result.current.runBatchesRef.current[0]).toHaveLength(1);
+  });
+
+  it("onError receives the live in-transit commands at error time", async () => {
+    const seen: AssistantTransportCommand[][] = [];
+    const { result } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: () => {
+          throw new Error("network error");
+        },
+        onError: (commands) => seen.push(commands),
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m1"));
+    });
+
+    // The flush that moved m1 into transit has not re-rendered yet when the
+    // error fires; the queue state must be read live, not from a render snapshot.
+    await waitFor(() => expect(seen).toHaveLength(1));
+    expect(seen[0]).toEqual([createMessageCommand("m1")]);
+  });
+
+  it("settles the run when onFinish throws", async () => {
+    const error = new Error("telemetry failed");
+    const onFinish = vi.fn<() => void>().mockImplementationOnce(() => {
+      throw error;
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const { result } = renderHook(() =>
+        useTransportSchedulingHarness({ onFinish }),
+      );
+
+      act(() => {
+        result.current.commandQueue.enqueue(createMessageCommand("m1"));
+      });
+
+      await waitFor(() => {
+        expect(result.current.runManager.isRunning).toBe(false);
+      });
+
+      act(() => {
+        result.current.commandQueue.enqueue(createMessageCommand("m2"));
+      });
+
+      await waitFor(() => {
+        expect(result.current.runBatchesRef.current).toHaveLength(2);
+      });
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] Assistant transport onFinish callback threw an error",
+        error,
+      );
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("contains rejected onError callbacks", async () => {
+    const callbackError = new Error("error telemetry failed");
+    const onError = vi.fn().mockRejectedValue(callbackError);
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation((message) => {
+        if (
+          message ===
+          "[assistant-ui] Assistant transport onError callback threw an error"
+        ) {
+          throw new Error("console unavailable");
+        }
+      });
+
+    try {
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        const { result } = renderHook(() =>
+          useTransportSchedulingHarness({
+            onRun: () => Promise.reject(new Error("network failed")),
+            onError,
+          }),
+        );
+
+        act(() => {
+          result.current.commandQueue.enqueue(createMessageCommand("m1"));
+        });
+
+        await waitFor(() => {
+          expect(result.current.runManager.isRunning).toBe(false);
+        });
+      });
+
+      expect(onError).toHaveBeenCalledTimes(1);
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] Assistant transport onError callback threw an error",
+        callbackError,
+      );
+      expect(unhandledRejections).toEqual([]);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("contains throwing onCancel callbacks", async () => {
+    const callbackError = new Error("cancel telemetry failed");
+    const onCancel = vi.fn().mockImplementation(() => {
+      throw callbackError;
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        const { result } = renderHook(() =>
+          useTransportSchedulingHarness({
+            onRun: (signal) =>
+              new Promise<void>((_resolve, reject) => {
+                signal.addEventListener(
+                  "abort",
+                  () => reject(new Error("aborted")),
+                  { once: true },
+                );
+              }),
+            onCancel,
+          }),
+        );
+
+        act(() => {
+          result.current.commandQueue.enqueue(createMessageCommand("m1"));
+        });
+        await waitFor(() => {
+          expect(result.current.runBatchesRef.current).toHaveLength(1);
+        });
+
+        act(() => result.current.runManager.cancel());
+        await waitFor(() => {
+          expect(result.current.runManager.isRunning).toBe(false);
+        });
+      });
+
+      expect(onCancel).toHaveBeenCalledTimes(1);
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] Assistant transport onCancel callback threw an error",
+        callbackError,
+      );
+      expect(unhandledRejections).toEqual([]);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("cancel returns combined in-flight and queued commands", async () => {
+    const onCancel = vi.fn();
+    const { result } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: (signal) =>
+          new Promise<void>((_resolve, reject) => {
+            signal.addEventListener(
+              "abort",
+              () => reject(new Error("aborted")),
+              { once: true },
+            );
+          }),
+        onCancel,
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("in-flight"));
+    });
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("queued"));
+      result.current.runManager.cancel();
+    });
+
+    await waitFor(() => {
+      expect(onCancel).toHaveBeenCalledTimes(1);
+    });
+    expect(onCancel.mock.calls[0]?.[0]).toHaveLength(2);
+  });
+
+  it("unmount aborts the in-flight run without invoking callbacks", async () => {
+    let aborted = false;
+    const onCancel = vi.fn();
+    const onError = vi.fn();
+    const onFinish = vi.fn();
+    const { result, unmount } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: (signal) =>
+          new Promise<void>((_resolve, reject) => {
+            signal.addEventListener(
+              "abort",
+              () => {
+                aborted = true;
+                reject(new Error("aborted"));
+              },
+              { once: true },
+            );
+          }),
+        onCancel,
+        onError,
+        onFinish,
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("in-flight"));
+    });
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+
+    unmount();
+
+    expect(aborted).toBe(true);
+    await act(async () => {});
+    expect(onCancel).not.toHaveBeenCalled();
+    expect(onError).not.toHaveBeenCalled();
+    expect(onFinish).not.toHaveBeenCalled();
+  });
+
+  it("ignores schedules after unmount", async () => {
+    const { result, unmount } = renderHook(() =>
+      useTransportSchedulingHarness(),
+    );
+
+    unmount();
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("late"));
+    });
+
+    await act(async () => {});
+    expect(result.current.runBatchesRef.current).toHaveLength(0);
+  });
+
+  it("survives StrictMode double-mounting", async () => {
+    const gate = createDeferred();
+    const onFinish = vi.fn();
+    const { result } = renderHook(
+      () =>
+        useTransportSchedulingHarness({ onRun: () => gate.promise, onFinish }),
+      { wrapper: ({ children }) => createElement(StrictMode, null, children) },
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m1"));
+    });
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+
+    // enqueued while the first run is in flight — must restart as a follow-up
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m2"));
+    });
+
+    gate.resolve();
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(2);
+    });
+    expect(result.current.runBatchesRef.current[1]).toHaveLength(1);
+    expect(onFinish).toHaveBeenCalled();
+    await waitFor(() => {
+      expect(result.current.runManager.isRunning).toBe(false);
+    });
+  });
+});