@assistant-ui/react 0.14.14 → 0.14.16

package/src/sandbox-host/SandboxHost.test.tsx

@@ -0,0 +1,231 @@
+// @vitest-environment jsdom
+import { act } from "react";
+import { createRoot, type Root } from "react-dom/client";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const { renderHtmlMock } = vi.hoisted(() => ({ renderHtmlMock: vi.fn() }));
+
+vi.mock("safe-content-frame", () => ({
+  SafeContentFrame: class {
+    renderHtml = renderHtmlMock;
+  },
+}));
+
+import {
+  SandboxHost,
+  isSandboxFrameMessage,
+  type SandboxBridge,
+  type SandboxHostApi,
+} from "./SandboxHost";
+
+const validData = { jsonrpc: "2.0", method: "x" };
+
+function makeFrame() {
+  const iframe = document.createElement("iframe");
+  document.body.appendChild(iframe);
+  return { iframe, origin: "https://app.example" };
+}
+
+function fakeRendered() {
+  const iframe = document.createElement("iframe");
+  document.body.appendChild(iframe);
+  return {
+    iframe,
+    origin: "https://fake.scf.test",
+    sendMessage: vi.fn(),
+    dispose: vi.fn(),
+    fullyLoadedPromiseWithTimeout: vi.fn(),
+  };
+}
+
+const flush = () =>
+  act(async () => {
+    await new Promise((r) => setTimeout(r, 0));
+  });
+
+describe("isSandboxFrameMessage", () => {
+  it("accepts a message from the frame's contentWindow at its origin", () => {
+    const frame = makeFrame();
+    const event = new MessageEvent("message", {
+      data: validData,
+      origin: frame.origin,
+      source: frame.iframe.contentWindow,
+    });
+    expect(isSandboxFrameMessage(event, frame)).toBe(true);
+  });
+
+  it("rejects a message from a different origin", () => {
+    const frame = makeFrame();
+    const event = new MessageEvent("message", {
+      data: validData,
+      origin: "https://attacker.example",
+      source: frame.iframe.contentWindow,
+    });
+    expect(isSandboxFrameMessage(event, frame)).toBe(false);
+  });
+
+  it("rejects a message from a different source window", () => {
+    const frame = makeFrame();
+    const event = new MessageEvent("message", {
+      data: validData,
+      origin: frame.origin,
+      source: window,
+    });
+    expect(isSandboxFrameMessage(event, frame)).toBe(false);
+  });
+});
+
+describe("SandboxHost", () => {
+  let container: HTMLDivElement;
+  let root: Root;
+
+  beforeEach(() => {
+    container = document.createElement("div");
+    document.body.appendChild(container);
+    root = createRoot(container);
+    renderHtmlMock.mockReset();
+  });
+
+  afterEach(() => {
+    act(() => {
+      try {
+        root.unmount();
+      } catch {
+        // already unmounted by the test
+      }
+    });
+    container.remove();
+  });
+
+  it("delivers only frame-validated messages to the bridge", async () => {
+    const rendered = fakeRendered();
+    renderHtmlMock.mockResolvedValue(rendered);
+    const onMessage = vi.fn();
+    const bridge: SandboxBridge = { onMessage, dispose: vi.fn() };
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "" }}
+          contentKey="k"
+          createBridge={() => bridge}
+        />,
+      );
+    });
+    await flush();
+
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: validData,
+        origin: rendered.origin,
+        source: rendered.iframe.contentWindow,
+      }),
+    );
+    expect(onMessage).toHaveBeenCalledTimes(1);
+
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: validData,
+        origin: "https://attacker.example",
+        source: rendered.iframe.contentWindow,
+      }),
+    );
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: validData,
+        origin: rendered.origin,
+        source: window,
+      }),
+    );
+    expect(onMessage).toHaveBeenCalledTimes(1);
+  });
+
+  it("clamps the bridge-reported height to maxHeight and ignores invalid values", async () => {
+    const rendered = fakeRendered();
+    renderHtmlMock.mockResolvedValue(rendered);
+    let host!: SandboxHostApi;
+    const bridge: SandboxBridge = { onMessage: vi.fn(), dispose: vi.fn() };
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "" }}
+          contentKey="k"
+          maxHeight={800}
+          createBridge={(_frame, h) => {
+            host = h;
+            return bridge;
+          }}
+        />,
+      );
+    });
+    await flush();
+
+    const div = container.firstElementChild as HTMLDivElement;
+    await act(async () => host.setHeight(200));
+    expect(div.style.height).toBe("200px");
+    await act(async () => host.setHeight(5000));
+    expect(div.style.height).toBe("800px");
+    await act(async () => host.setHeight(0));
+    expect(div.style.height).toBe("800px");
+  });
+
+  it("disposes the bridge before the frame and detaches the listener on unmount", async () => {
+    const rendered = fakeRendered();
+    renderHtmlMock.mockResolvedValue(rendered);
+    const order: string[] = [];
+    const onMessage = vi.fn();
+    const bridge: SandboxBridge = {
+      onMessage,
+      dispose: vi.fn(() => order.push("bridge")),
+    };
+    rendered.dispose = vi.fn(() => order.push("frame"));
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "" }}
+          contentKey="k"
+          createBridge={() => bridge}
+        />,
+      );
+    });
+    await flush();
+
+    await act(async () => {
+      root.unmount();
+    });
+
+    expect(order).toEqual(["bridge", "frame"]);
+
+    onMessage.mockClear();
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: validData,
+        origin: rendered.origin,
+        source: rendered.iframe.contentWindow,
+      }),
+    );
+    expect(onMessage).not.toHaveBeenCalled();
+  });
+
+  it("calls onError when rendering rejects", async () => {
+    renderHtmlMock.mockRejectedValue(new Error("boom"));
+    const onError = vi.fn();
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "" }}
+          contentKey="k"
+          createBridge={() => ({ onMessage: vi.fn(), dispose: vi.fn() })}
+          onError={onError}
+        />,
+      );
+    });
+    await flush();
+
+    expect(onError).toHaveBeenCalledTimes(1);
+    expect(onError.mock.calls[0]![0].message).toBe("boom");
+  });
+});

package/src/sandbox-host/SandboxHost.tsx

@@ -0,0 +1,185 @@
+"use client";
+
+import { type CSSProperties, useEffect, useRef, useState } from "react";
+import {
+  type RenderedFrame,
+  SafeContentFrame,
+  type SandboxOption,
+} from "safe-content-frame";
+
+const DEFAULT_PRODUCT = "assistant-ui-sandbox";
+const DEFAULT_MAX_HEIGHT = 800;
+
+export type SandboxHostConfig = {
+  sandbox?: SandboxOption[];
+  useShadowDom?: boolean;
+  enableBrowserCaching?: boolean;
+  salt?: string;
+  product?: string;
+  className?: string;
+  style?: CSSProperties;
+  unsafeDocumentWrite?: boolean;
+};
+
+export type SandboxHostFrame = Pick<
+  RenderedFrame,
+  "iframe" | "origin" | "sendMessage"
+>;
+
+export type SandboxHostApi = {
+  setHeight: (height: number) => void;
+};
+
+export type SandboxBridge = {
+  onMessage: (event: MessageEvent) => void;
+  dispose: () => void;
+};
+
+export type SandboxContent = { html: string };
+
+export type SandboxHostProps = {
+  content: SandboxContent;
+  contentKey: string;
+  sandbox?: SandboxHostConfig | undefined;
+  maxHeight?: number | undefined;
+  createBridge: (
+    frame: SandboxHostFrame,
+    host: SandboxHostApi,
+  ) => SandboxBridge;
+  onError?: ((error: Error) => void) | undefined;
+  containerProps?: Record<string, string | undefined> | undefined;
+};
+
+export function isSandboxFrameMessage(
+  event: MessageEvent,
+  frame: { iframe: HTMLIFrameElement; origin: string },
+): boolean {
+  return (
+    event.source === frame.iframe.contentWindow && event.origin === frame.origin
+  );
+}
+
+type LiveSnapshot = {
+  content: SandboxContent;
+  sandbox: SandboxHostConfig | undefined;
+  createBridge: SandboxHostProps["createBridge"];
+  onError: SandboxHostProps["onError"];
+};
+
+export function SandboxHost({
+  content,
+  contentKey,
+  sandbox,
+  maxHeight = DEFAULT_MAX_HEIGHT,
+  createBridge,
+  onError,
+  containerProps,
+}: SandboxHostProps) {
+  const containerRef = useRef<HTMLDivElement>(null);
+  const [contentHeight, setContentHeight] = useState<number | undefined>(
+    undefined,
+  );
+
+  const liveRef = useRef<LiveSnapshot>(null!);
+  liveRef.current = { content, sandbox, createBridge, onError };
+
+  useEffect(() => {
+    const container = containerRef.current;
+    if (!container) return;
+
+    let cancelled = false;
+    let frame: RenderedFrame | null = null;
+    let bridge: SandboxBridge | null = null;
+    let onMessage: ((event: MessageEvent) => void) | null = null;
+
+    const { content: liveContent, sandbox: sb } = liveRef.current;
+
+    const scf = new SafeContentFrame(sb?.product ?? DEFAULT_PRODUCT, {
+      ...(sb?.sandbox !== undefined && { sandbox: sb.sandbox }),
+      ...(sb?.useShadowDom !== undefined && { useShadowDom: sb.useShadowDom }),
+      ...(sb?.enableBrowserCaching !== undefined && {
+        enableBrowserCaching: sb.enableBrowserCaching,
+      }),
+      ...(sb?.salt !== undefined && { salt: sb.salt }),
+    });
+
+    const renderOpts =
+      sb?.unsafeDocumentWrite !== undefined
+        ? { unsafeDocumentWrite: sb.unsafeDocumentWrite }
+        : undefined;
+
+    scf
+      .renderHtml(liveContent.html, container, renderOpts)
+      .then((rendered) => {
+        if (cancelled) {
+          rendered.dispose();
+          return;
+        }
+        frame = rendered;
+
+        const hostApi: SandboxHostApi = {
+          setHeight: (height) => {
+            if (
+              typeof height === "number" &&
+              Number.isFinite(height) &&
+              height > 0
+            ) {
+              setContentHeight(height);
+            }
+          },
+        };
+
+        bridge = liveRef.current.createBridge(
+          {
+            iframe: rendered.iframe,
+            origin: rendered.origin,
+            sendMessage: rendered.sendMessage,
+          },
+          hostApi,
+        );
+
+        // Single owner of the window listener; the cross-origin guard runs
+        // here so the bridge only sees frame-validated messages.
+        onMessage = (event) => {
+          if (!isSandboxFrameMessage(event, rendered)) return;
+          bridge?.onMessage(event);
+        };
+        window.addEventListener("message", onMessage);
+      })
+      .catch((err) => {
+        liveRef.current.onError?.(
+          err instanceof Error ? err : new Error(String(err)),
+        );
+      });
+
+    return () => {
+      cancelled = true;
+      if (onMessage) {
+        window.removeEventListener("message", onMessage);
+        onMessage = null;
+      }
+      bridge?.dispose();
+      bridge = null;
+      frame?.dispose();
+      frame = null;
+      setContentHeight(undefined);
+    };
+    // oxlint-disable-next-line react/exhaustive-deps -- re-init only on contentKey change; live values flow through liveRef
+  }, [contentKey]);
+
+  const resolvedHeight =
+    contentHeight != null ? Math.min(contentHeight, maxHeight) : undefined;
+  const mergedStyle =
+    resolvedHeight != null
+      ? { ...sandbox?.style, height: resolvedHeight }
+      : sandbox?.style;
+
+  return (
+    <div
+      {...containerProps}
+      ref={containerRef}
+      className={sandbox?.className}
+      style={mergedStyle}
+    />
+  );
+}

package/src/tests/local-runtime-queue.test.tsx

@@ -0,0 +1,305 @@
+// @vitest-environment jsdom
+
+import { render, act } from "@testing-library/react";
+import type { FC } from "react";
+import { describe, it, expect } from "vitest";
+import { useAui } from "@assistant-ui/store";
+import { AssistantRuntimeProvider } from "../context";
+import { useLocalRuntime } from "../legacy-runtime/runtime-cores/local/useLocalRuntime";
+import type { ChatModelAdapter } from "../legacy-runtime/runtime-cores/local/ChatModelAdapter";
+
+const flush = () => new Promise<void>((resolve) => setTimeout(resolve, 0));
+
+const createCountingAdapter = () => {
+  const releases: Array<() => void> = [];
+  let runCount = 0;
+  const adapter: ChatModelAdapter = {
+    async *run({ abortSignal }) {
+      runCount++;
+      await new Promise<void>((resolve) => {
+        releases.push(resolve);
+        abortSignal.addEventListener("abort", () => resolve(), { once: true });
+      });
+      yield { content: [{ type: "text", text: "done" }] };
+    },
+  };
+  return { adapter, releases, getRunCount: () => runCount };
+};
+
+const userTexts = (aui: ReturnType<typeof useAui>) =>
+  aui
+    .thread()
+    .getState()
+    .messages.filter((m) => m.role === "user")
+    .map((m) =>
+      m.content.map((p) => (p.type === "text" ? p.text : "")).join(""),
+    );
+
+const renderWithRuntime = (adapter: ChatModelAdapter, enableQueue: boolean) => {
+  const captured: { aui?: ReturnType<typeof useAui> } = {};
+  const Capture: FC = () => {
+    captured.aui = useAui();
+    return null;
+  };
+  const App: FC = () => {
+    const runtime = useLocalRuntime(adapter, {
+      unstable_enableMessageQueue: enableQueue,
+    });
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <Capture />
+      </AssistantRuntimeProvider>
+    );
+  };
+  render(<App />);
+  return captured.aui!;
+};
+
+const send = async (aui: ReturnType<typeof useAui>, text: string) => {
+  await act(async () => {
+    aui.thread().composer().setText(text);
+    aui.thread().composer().send();
+    await flush();
+  });
+};
+
+describe("local runtime message queue", () => {
+  it("buffers a send while running and flushes it when the run ends", async () => {
+    const { adapter, releases } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    expect(aui.thread().getState().isRunning).toBe(true);
+    expect(aui.thread().getState().capabilities.queue).toBe(true);
+
+    await send(aui, "second");
+    expect(
+      aui
+        .thread()
+        .composer()
+        .getState()
+        .queue.map((q) => q.prompt),
+    ).toEqual(["second"]);
+    expect(userTexts(aui)).toEqual(["first"]);
+
+    await act(async () => {
+      releases[0]!();
+      await flush();
+      await flush();
+    });
+    expect(aui.thread().composer().getState().queue).toEqual([]);
+    expect(userTexts(aui)).toContain("second");
+  });
+
+  it("drains two queued items in separate runs, not all at once", async () => {
+    const { adapter, releases, getRunCount } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    expect(getRunCount()).toBe(1);
+
+    await send(aui, "a");
+    await send(aui, "b");
+    expect(aui.thread().composer().getState().queue).toHaveLength(2);
+
+    await act(async () => {
+      releases[0]!();
+      await flush();
+      await flush();
+    });
+    expect(getRunCount()).toBe(2);
+    expect(
+      aui
+        .thread()
+        .composer()
+        .getState()
+        .queue.map((q) => q.prompt),
+    ).toEqual(["b"]);
+  });
+
+  it("queueItem(index).remove() drops a queued message", async () => {
+    const { adapter } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    await send(aui, "a");
+    await send(aui, "b");
+    expect(aui.thread().composer().getState().queue).toHaveLength(2);
+
+    await act(async () => {
+      aui.thread().composer().queueItem({ index: 0 }).remove();
+      await flush();
+    });
+    const queue = aui.thread().composer().getState().queue;
+    expect(queue).toHaveLength(1);
+    expect(queue[0]!.prompt).toBe("b");
+  });
+
+  it("clears the queue when the run is cancelled, without flushing", async () => {
+    const { adapter, getRunCount } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    await send(aui, "a");
+    await send(aui, "b");
+    expect(aui.thread().composer().getState().queue).toHaveLength(2);
+
+    await act(async () => {
+      aui.thread().cancelRun();
+      await flush();
+      await flush();
+    });
+
+    expect(aui.thread().composer().getState().queue).toEqual([]);
+    // cancelling must not start the next queued message
+    expect(getRunCount()).toBe(1);
+  });
+
+  it("applies an edit instead of queuing it, dropping pending items", async () => {
+    const { adapter, releases } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    await act(async () => {
+      releases[0]!();
+      await flush();
+    });
+
+    // start a second run and queue a message behind it
+    await send(aui, "second");
+    await send(aui, "queued");
+    expect(aui.thread().composer().getState().queue).toHaveLength(1);
+
+    // edit the first message while the run is in progress
+    await act(async () => {
+      const message = aui.thread().message({ index: 0 });
+      message.composer().beginEdit();
+      message.composer().setText("edited");
+      message.composer().send();
+      await flush();
+    });
+
+    // the edit is applied (branches the thread) and the stale queue is cleared
+    expect(aui.thread().composer().getState().queue).toEqual([]);
+  });
+
+  it("buffers a send during a regenerate instead of interrupting it", async () => {
+    const { adapter, releases, getRunCount } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    await act(async () => {
+      releases[0]!();
+      await flush();
+    });
+    expect(getRunCount()).toBe(1);
+
+    // regenerate the assistant message: a run started outside the queue
+    await act(async () => {
+      aui.thread().message({ index: 1 }).reload();
+      await flush();
+    });
+    expect(getRunCount()).toBe(2);
+    expect(aui.thread().getState().isRunning).toBe(true);
+
+    // sending now must buffer, not interrupt the regenerate
+    await act(async () => {
+      aui.thread().composer().setText("Y");
+      aui.thread().composer().send();
+      await flush();
+    });
+    expect(getRunCount()).toBe(2);
+    expect(
+      aui
+        .thread()
+        .composer()
+        .getState()
+        .queue.map((q) => q.prompt),
+    ).toEqual(["Y"]);
+  });
+
+  it("advances exactly once after a failed run, without deadlocking", async () => {
+    const releases: Array<() => void> = [];
+    let runCount = 0;
+    const adapter: ChatModelAdapter = {
+      async *run({ abortSignal }) {
+        runCount++;
+        if (runCount === 2) throw new Error("model boom");
+        await new Promise<void>((resolve) => {
+          releases.push(resolve);
+          abortSignal.addEventListener("abort", () => resolve(), {
+            once: true,
+          });
+        });
+        yield { content: [{ type: "text", text: "done" }] };
+      },
+    };
+    const aui = renderWithRuntime(adapter, true);
+
+    await send(aui, "first");
+    await send(aui, "a");
+    await send(aui, "b");
+    expect(aui.thread().composer().getState().queue).toHaveLength(2);
+
+    // run 1 settles -> "a" drains (run 2 throws) -> "b" drains (run 3)
+    await act(async () => {
+      releases[0]!();
+      await flush();
+      await flush();
+    });
+    expect(runCount).toBe(3);
+    expect(aui.thread().composer().getState().queue).toEqual([]);
+
+    // "b" is running; a new send must buffer behind it, not interrupt it
+    await act(async () => {
+      aui.thread().composer().setText("c");
+      aui.thread().composer().send();
+      await flush();
+    });
+    expect(runCount).toBe(3);
+    expect(
+      aui
+        .thread()
+        .composer()
+        .getState()
+        .queue.map((q) => q.prompt),
+    ).toEqual(["c"]);
+  });
+
+  it("does not expose the queue capability when the flag is off", async () => {
+    const { adapter } = createCountingAdapter();
+    const aui = renderWithRuntime(adapter, false);
+    expect(aui.thread().getState().capabilities.queue).toBe(false);
+  });
+
+  it("tears the queue down when the flag is toggled off at runtime", async () => {
+    const { adapter } = createCountingAdapter();
+    const captured: { aui?: ReturnType<typeof useAui> } = {};
+    const Capture: FC = () => {
+      captured.aui = useAui();
+      return null;
+    };
+    const App: FC<{ enabled: boolean }> = ({ enabled }) => {
+      const runtime = useLocalRuntime(adapter, {
+        unstable_enableMessageQueue: enabled,
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Capture />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    const { rerender } = render(<App enabled={true} />);
+    await act(async () => {
+      await flush();
+    });
+    expect(captured.aui!.thread().getState().capabilities.queue).toBe(true);
+
+    await act(async () => {
+      rerender(<App enabled={false} />);
+      await flush();
+    });
+    expect(captured.aui!.thread().getState().capabilities.queue).toBe(false);
+  });
+});