@assistant-ui/react 0.15.18 → 0.15.20

package/src/mcp-apps/app-frame.test.tsx

@@ -7,6 +7,7 @@ import type {
   SandboxHostProps,
 } from "../sandbox-host/SandboxHost";
 import type { CreateMcpAppBridgeOptions, McpAppBridge } from "./bridge";
+import type * as BridgeModule from "./bridge";
 import { MCP_APP_MIME_TYPE, type McpAppHostContext } from "./types";
 
 const { sandboxHostMock, createMcpAppBridgeMock } = vi.hoisted(() => ({
@@ -28,7 +29,6 @@ import { McpAppFrame } from "./app-frame";
 describe("McpAppFrame", () => {
   afterEach(() => {
     cleanup();
-    vi.clearAllMocks();
   });
 
   it("keeps bridge options scoped to committed renders", async () => {
@@ -100,6 +100,178 @@ describe("McpAppFrame", () => {
     sandboxBridge.dispose();
   });
 
+  it.each([
+    { initial: [], next: ["search"], allowed: true },
+    { initial: ["search"], next: [], allowed: false },
+    { initial: undefined, next: ["other"], allowed: false },
+    { initial: [], next: undefined, allowed: true },
+  ])(
+    "applies replacement tool allowlists: $initial -> $next",
+    async ({ initial, next, allowed }) => {
+      const { createMcpAppBridge } =
+        await vi.importActual<typeof BridgeModule>("./bridge");
+      createMcpAppBridgeMock.mockImplementationOnce(createMcpAppBridge);
+      const captured: { createBridge?: SandboxHostProps["createBridge"] } = {};
+      sandboxHostMock.mockImplementation((props: SandboxHostProps) => {
+        captured.createBridge ??= props.createBridge;
+        return null;
+      });
+      const callTool = vi.fn(() => ({ content: [] }));
+      const view = (allowedTools: readonly string[] | undefined) => (
+        <McpAppFrame
+          app={{ resourceUri: "ui://example/widget" }}
+          resource={{
+            uri: "ui://example/widget",
+            mimeType: MCP_APP_MIME_TYPE,
+            html: "",
+          }}
+          handlers={
+            allowedTools === undefined
+              ? { callTool }
+              : { allowedTools, callTool }
+          }
+        />
+      );
+      const rendered = render(view(initial));
+      if (!captured.createBridge) throw new Error("Frame did not mount");
+      const sendMessage = vi.fn();
+      const bridge = captured.createBridge(
+        {
+          iframe: document.createElement("iframe"),
+          origin: "https://widget.example",
+          sendMessage,
+        },
+        { setHeight: vi.fn() },
+      );
+      try {
+        rendered.rerender(view(next));
+        bridge.onMessage(
+          new MessageEvent("message", {
+            data: {
+              jsonrpc: "2.0",
+              id: 1,
+              method: "tools/call",
+              params: { name: "search" },
+            },
+          }),
+        );
+        await Promise.resolve();
+        expect(sendMessage).toHaveBeenCalledWith(
+          expect.objectContaining(
+            allowed
+              ? { id: 1, result: { content: [] } }
+              : { id: 1, error: expect.objectContaining({ code: -32602 }) },
+          ),
+        );
+        expect(callTool).toHaveBeenCalledTimes(allowed ? 1 : 0);
+      } finally {
+        bridge.dispose();
+      }
+    },
+  );
+
+  it("forwards size changes to the sandbox host and live handler", () => {
+    let createBridge: SandboxHostProps["createBridge"] | null = null;
+    sandboxHostMock.mockImplementation((props: SandboxHostProps) => {
+      createBridge ??= props.createBridge;
+      return null;
+    });
+    const bridge: McpAppBridge = {
+      onMessage: vi.fn(),
+      dispose: vi.fn(),
+      notifyToolInput: vi.fn(),
+      notifyToolResult: vi.fn(),
+      notifyHostContextChanged: vi.fn(),
+    };
+    createMcpAppBridgeMock.mockReturnValue(bridge);
+    const initialOnSizeChange = vi.fn();
+    const replacementOnSizeChange = vi.fn();
+    const view = (onSizeChange: typeof initialOnSizeChange) => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        handlers={{ onSizeChange }}
+      />
+    );
+    const rendered = render(view(initialOnSizeChange));
+
+    const setHeight = vi.fn();
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight },
+    );
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+
+    try {
+      const initialSize = { width: 640, height: 360 };
+      options.handlers?.onSizeChange?.(initialSize);
+      expect(setHeight).toHaveBeenCalledWith(360);
+      expect(initialOnSizeChange).toHaveBeenCalledWith(initialSize);
+
+      rendered.rerender(view(replacementOnSizeChange));
+      const replacementSize = { width: 800, height: 480 };
+      options.handlers?.onSizeChange?.(replacementSize);
+      expect(setHeight).toHaveBeenNthCalledWith(2, 480);
+      expect(initialOnSizeChange).toHaveBeenCalledOnce();
+      expect(replacementOnSizeChange).toHaveBeenCalledOnce();
+      expect(replacementOnSizeChange).toHaveBeenCalledWith(replacementSize);
+
+      const widthOnlySize = { width: 720 };
+      options.handlers?.onSizeChange?.(widthOnlySize);
+      expect(setHeight).toHaveBeenCalledTimes(2);
+      expect(replacementOnSizeChange).toHaveBeenCalledWith(widthOnlySize);
+    } finally {
+      sandboxBridge.dispose();
+    }
+  });
+
+  it("disposes the underlying bridge with the sandbox bridge", () => {
+    let createBridge: SandboxHostProps["createBridge"] | null = null;
+    sandboxHostMock.mockImplementation((props: SandboxHostProps) => {
+      createBridge ??= props.createBridge;
+      return null;
+    });
+    const bridge: McpAppBridge = {
+      onMessage: vi.fn(),
+      dispose: vi.fn(),
+      notifyToolInput: vi.fn(),
+      notifyToolResult: vi.fn(),
+      notifyHostContextChanged: vi.fn(),
+    };
+    createMcpAppBridgeMock.mockReturnValue(bridge);
+    render(
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+      />,
+    );
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+    sandboxBridge.dispose();
+
+    expect(bridge.dispose).toHaveBeenCalledOnce();
+  });
+
   it("only notifies the widget when host context actually changes", () => {
     let createBridge: SandboxHostProps["createBridge"] | null = null;
     sandboxHostMock.mockImplementation((props: SandboxHostProps) => {
@@ -159,6 +331,22 @@ describe("McpAppFrame", () => {
       availableDisplayModes: ["inline", "pip"],
     });
 
+    // A hole is not a value, in either operand order.
+    const sparseModes: ("inline" | "pip")[] = new Array(2);
+    sparseModes[1] = "pip";
+    rendered.rerender(
+      view({ displayMode: "fullscreen", availableDisplayModes: sparseModes }),
+    );
+    expect(bridge.notifyHostContextChanged).toHaveBeenCalledTimes(2);
+
+    rendered.rerender(
+      view({
+        displayMode: "fullscreen",
+        availableDisplayModes: ["inline", "pip"],
+      }),
+    );
+    expect(bridge.notifyHostContextChanged).toHaveBeenCalledTimes(3);
+
     sandboxBridge.dispose();
   });
   it("cancels a queued host context that returns to the delivered value", () => {

package/src/mcp-apps/app-frame.tsx

@@ -42,12 +42,12 @@ const isSameHostContext = (a: unknown, b: unknown, depth = 0): boolean => {
   if (Object.is(a, b)) return true;
   if (depth > 100) return false;
   if (Array.isArray(a) || Array.isArray(b)) {
-    return (
-      Array.isArray(a) &&
-      Array.isArray(b) &&
-      a.length === b.length &&
-      a.every((item, index) => isSameHostContext(item, b[index], depth + 1))
-    );
+    if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length)
+      return false;
+    for (let i = 0; i < a.length; i++) {
+      if (!isSameHostContext(a[i], b[i], depth + 1)) return false;
+    }
+    return true;
   }
   if (!isPlainObject(a) || !isPlainObject(b)) return false;
   const aKeys = Object.keys(a);
@@ -96,25 +96,33 @@ type LiveSnapshot = {
   output: unknown;
 };
 
+type FrameLifecycle = {
+  onInitialized: () => void;
+  onSizeChange: (params: { width?: number; height?: number }) => void;
+};
+
 // Proxy each per-call handler through liveRef so the bridge always dispatches
 // to the latest handler reference (e.g. inline callbacks closing over state).
 // Capability presence is snapshot at mount: a handler added later requires a
 // remount (keyed on resource URI) to expose the capability to the widget.
+// allowedTools is the exception: it is never advertised in the ui/initialize
+// response, so it stays a live getter, which means this object must reach the
+// bridge uncopied and the frame passes its lifecycle work in rather than
+// wrapping the result.
 function buildLiveHandlers(
   initial: McpAppBridgeHandlers | undefined,
   liveRef: { readonly current: LiveSnapshot },
+  lifecycle: FrameLifecycle,
 ): McpAppBridgeHandlers {
   const live = () => liveRef.current.handlers;
   const has = <K extends keyof McpAppBridgeHandlers>(key: K) =>
     initial?.[key] !== undefined;
   const out: McpAppBridgeHandlers = {};
-  if (has("allowedTools")) {
-    Object.defineProperty(out, "allowedTools", {
-      get: () => live()?.allowedTools,
-      enumerable: true,
-      configurable: true,
-    });
-  }
+  Object.defineProperty(out, "allowedTools", {
+    get: () => live()?.allowedTools,
+    enumerable: true,
+    configurable: true,
+  });
   const liveCall = <K extends keyof McpAppBridgeHandlers>(
     key: K,
   ): NonNullable<McpAppBridgeHandlers[K]> =>
@@ -134,8 +142,14 @@ function buildLiveHandlers(
     out.updateModelContext = liveCall("updateModelContext");
   if (has("requestDisplayMode"))
     out.requestDisplayMode = liveCall("requestDisplayMode");
-  out.onSizeChange = (p) => live()?.onSizeChange?.(p);
-  out.onInitialized = () => live()?.onInitialized?.();
+  out.onSizeChange = (p) => {
+    lifecycle.onSizeChange(p);
+    live()?.onSizeChange?.(p);
+  };
+  out.onInitialized = () => {
+    lifecycle.onInitialized();
+    live()?.onInitialized?.();
+  };
   out.onRequestTeardown = (p) => live()?.onRequestTeardown?.(p);
   out.onLog = (p) => live()?.onLog?.(p);
   out.onError = (e) => live()?.onError?.(e);
@@ -215,23 +229,18 @@ export function McpAppFrame({
       }
     };
 
-    const liveHandlers = buildLiveHandlers(current.handlers, liveRef);
-    const liveOnInitialized = liveHandlers.onInitialized;
-    const wrappedHandlers: McpAppBridgeHandlers = {
-      ...liveHandlers,
+    const wrappedHandlers = buildLiveHandlers(current.handlers, liveRef, {
       onInitialized: () => {
         if (initTimeoutId !== null) {
           clearTimeout(initTimeoutId);
           initTimeoutId = null;
         }
         flushPending();
-        liveOnInitialized?.();
       },
       onSizeChange: (p) => {
         if (p.height != null) host.setHeight(p.height);
-        liveHandlers.onSizeChange?.(p);
       },
-    };
+    });
 
     // Safety net: if the widget never sends notifications/initialized (broken
     // or non-spec-compliant), flush the queue anyway so the host doesn't

package/src/mcp-apps/bridge.test.ts

@@ -37,6 +37,104 @@ async function flush() {
 }
 
 describe("createMcpAppBridge", () => {
+  it("ignores requests and notifications after disposal", async () => {
+    const { frame, captured } = makeFrame();
+    const callTool = vi.fn().mockResolvedValue({ ok: true });
+    const onInitialized = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { callTool, onInitialized },
+    });
+
+    bridge.dispose();
+    expect(() => bridge.dispose()).not.toThrow();
+
+    deliver(bridge, {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "tools/call",
+      params: { name: "search" },
+    });
+    deliver(bridge, {
+      jsonrpc: "2.0",
+      method: "notifications/initialized",
+    });
+    await flush();
+
+    expect(callTool).not.toHaveBeenCalled();
+    expect(onInitialized).not.toHaveBeenCalled();
+    expect(captured).toEqual([]);
+  });
+
+  it("does not send a response after disposal", async () => {
+    const { frame, captured } = makeFrame();
+    let resolveCall!: (value: unknown) => void;
+    const callTool = vi.fn(
+      () =>
+        new Promise((resolve) => {
+          resolveCall = resolve;
+        }),
+    );
+    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });
+
+    deliver(bridge, {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "tools/call",
+      params: { name: "search" },
+    });
+    expect(callTool).toHaveBeenCalledOnce();
+
+    bridge.dispose();
+    resolveCall({ ok: true });
+    await flush();
+
+    expect(captured).toEqual([]);
+  });
+
+  it("does not report an async error after disposal", async () => {
+    const { frame, captured } = makeFrame();
+    let rejectCall!: (reason?: unknown) => void;
+    const callTool = vi.fn(
+      () =>
+        new Promise((_resolve, reject) => {
+          rejectCall = reject;
+        }),
+    );
+    const onError = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { callTool, onError },
+    });
+
+    deliver(bridge, {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "tools/call",
+      params: { name: "search" },
+    });
+    expect(callTool).toHaveBeenCalledOnce();
+
+    bridge.dispose();
+    rejectCall(new Error("tool failed"));
+    await flush();
+
+    expect(onError).not.toHaveBeenCalled();
+    expect(captured).toEqual([]);
+  });
+
+  it("does not send host notifications after disposal", () => {
+    const { frame, captured } = makeFrame();
+    const bridge = createMcpAppBridge({ frame });
+
+    bridge.dispose();
+    bridge.notifyToolInput({ query: "hello" });
+    bridge.notifyToolResult({ answer: "world" });
+    bridge.notifyHostContextChanged({ theme: "dark" });
+
+    expect(captured).toEqual([]);
+  });
+
   it("responds to ui/initialize with host info, version, and capabilities", async () => {
     const { frame, captured } = makeFrame();
     const bridge = createMcpAppBridge({

package/src/mcp-apps/bridge.ts

@@ -90,8 +90,10 @@ export function createMcpAppBridge(
     hostInfo = DEFAULT_HOST_INFO,
     hostContext = {},
   } = opts;
+  let disposed = false;
 
   const post = (msg: McpAppJsonRpcMessage) => {
+    if (disposed) return;
     frame.sendMessage(msg);
   };
 
@@ -125,6 +127,7 @@ export function createMcpAppBridge(
   };
 
   const reportError = (error: Error) => {
+    if (disposed) return;
     invokeUserCallback(
       "assistant-ui",
       "MCP App onError",
@@ -435,6 +438,7 @@ export function createMcpAppBridge(
   // The host applies the cross-origin guard before delegating; this only
   // validates the JSON-RPC envelope.
   const onMessage = (event: MessageEvent) => {
+    if (disposed) return;
     if (!isJsonRpcMessage(event.data)) return;
 
     const msg = event.data;
@@ -447,7 +451,9 @@ export function createMcpAppBridge(
 
   return {
     onMessage,
-    dispose: () => {},
+    dispose: () => {
+      disposed = true;
+    },
     notifyToolInput: (input: unknown) => {
       post({
         jsonrpc: "2.0",

package/src/model-context/frame/useAssistantFrameHost.test.tsx

@@ -0,0 +1,42 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { AssistantFrameHost } from "@assistant-ui/core";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { useAssistantFrameHost } from "./useAssistantFrameHost";
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe("useAssistantFrameHost", () => {
+  it("unregisters the host when disposal throws", () => {
+    const disposalError = new Error("tool cancellation failed");
+    const unregistrationError = new Error("unregistration failed");
+    vi.spyOn(AssistantFrameHost.prototype, "dispose").mockImplementation(() => {
+      throw disposalError;
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const unsubscribe = vi.fn(() => {
+      throw unregistrationError;
+    });
+    const register = vi.fn(() => unsubscribe);
+    const iframeRef = {
+      current: {
+        contentWindow: { postMessage: vi.fn() } as unknown as Window,
+      } as HTMLIFrameElement,
+    };
+    const { unmount } = renderHook(() =>
+      useAssistantFrameHost({ iframeRef, register }),
+    );
+
+    expect(() => unmount()).toThrow(disposalError);
+    expect(unsubscribe).toHaveBeenCalledOnce();
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] AssistantFrameHost unregistration failed.",
+      unregistrationError,
+    );
+  });
+});

package/src/model-context/frame/useAssistantFrameHost.ts

@@ -41,8 +41,31 @@ export const useAssistantFrameHost = ({
     const unsubscribe = register(frameHost);
 
     return () => {
-      frameHost.dispose();
-      unsubscribe();
+      let cleanupFailed = false;
+      let cleanupError: unknown;
+
+      try {
+        frameHost.dispose();
+      } catch (error) {
+        cleanupFailed = true;
+        cleanupError = error;
+      }
+
+      try {
+        unsubscribe();
+      } catch (error) {
+        if (cleanupFailed) {
+          console.error(
+            "[assistant-ui] AssistantFrameHost unregistration failed.",
+            error,
+          );
+        } else {
+          cleanupFailed = true;
+          cleanupError = error;
+        }
+      }
+
+      if (cleanupFailed) throw cleanupError;
     };
   }, [iframeRef, targetOrigin, register]);
 };

package/src/primitives/assistantModal/AssistantModalRoot.test.tsx

@@ -0,0 +1,123 @@
+// @vitest-environment jsdom
+
+import {
+  act,
+  cleanup,
+  fireEvent,
+  render,
+  screen,
+} from "@testing-library/react";
+import { useState } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { AssistantRuntimeProvider } from "../../context";
+import { useLocalRuntime } from "../../legacy-runtime/runtime-cores/local/useLocalRuntime";
+import { AssistantModalPrimitiveContent } from "./AssistantModalContent";
+import { AssistantModalPrimitiveRoot } from "./AssistantModalRoot";
+import { AssistantModalPrimitiveTrigger } from "./AssistantModalTrigger";
+
+const adapter = {
+  async *run() {
+    yield { content: [{ type: "text" as const, text: "Hello" }] };
+  },
+};
+
+const Modal = ({
+  controlled = true,
+  onOpenChange,
+  ...props
+}: AssistantModalPrimitiveRoot.Props & { controlled?: boolean }) => {
+  const runtime = useLocalRuntime(adapter);
+  const [open, setOpen] = useState(false);
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <button onClick={() => runtime.thread.append("Hello")}>Start run</button>
+      <AssistantModalPrimitiveRoot
+        {...(controlled ? { open } : {})}
+        onOpenChange={(value) => {
+          setOpen(value);
+          onOpenChange?.(value);
+        }}
+        {...props}
+      >
+        <AssistantModalPrimitiveTrigger>
+          Toggle chat
+        </AssistantModalPrimitiveTrigger>
+        <AssistantModalPrimitiveContent aria-label="Chat">
+          Response
+        </AssistantModalPrimitiveContent>
+      </AssistantModalPrimitiveRoot>
+    </AssistantRuntimeProvider>
+  );
+};
+
+afterEach(cleanup);
+
+describe("AssistantModalPrimitiveRoot run start", () => {
+  it.each([true, false])(
+    "opens on run start and closes through the trigger (controlled: %s)",
+    async (controlled) => {
+      const onOpenChange = vi.fn();
+      render(<Modal controlled={controlled} onOpenChange={onOpenChange} />);
+      expect(screen.queryByRole("dialog")).toBeNull();
+
+      await act(async () => {
+        fireEvent.click(screen.getByRole("button", { name: "Start run" }));
+      });
+      expect(await screen.findByRole("dialog", { name: "Chat" })).toBeDefined();
+      expect(onOpenChange).toHaveBeenCalledExactlyOnceWith(true);
+
+      await act(async () => {
+        fireEvent.click(screen.getByRole("button", { name: "Start run" }));
+      });
+      expect(onOpenChange).toHaveBeenCalledExactlyOnceWith(true);
+
+      fireEvent.click(screen.getByRole("button", { name: "Toggle chat" }));
+      expect(screen.queryByRole("dialog")).toBeNull();
+    },
+  );
+
+  it.each([{ open: true }, { controlled: false, defaultOpen: true }])(
+    "does not request opening when already open (%o)",
+    async (props) => {
+      const onOpenChange = vi.fn();
+      render(<Modal {...props} onOpenChange={onOpenChange} />);
+      expect(screen.getByRole("dialog", { name: "Chat" })).toBeDefined();
+
+      await act(async () => {
+        fireEvent.click(screen.getByRole("button", { name: "Start run" }));
+      });
+
+      expect(onOpenChange).not.toHaveBeenCalled();
+      expect(screen.getByRole("dialog", { name: "Chat" })).toBeDefined();
+    },
+  );
+
+  it("notifies the current owner without overriding its controlled value", async () => {
+    const previous = vi.fn();
+    const current = vi.fn();
+    const view = render(<Modal open={false} onOpenChange={previous} />);
+    view.rerender(<Modal open={false} onOpenChange={current} />);
+
+    await act(async () => {
+      fireEvent.click(screen.getByRole("button", { name: "Start run" }));
+    });
+
+    expect(current).toHaveBeenCalledExactlyOnceWith(true);
+    expect(previous).not.toHaveBeenCalled();
+    expect(screen.queryByRole("dialog")).toBeNull();
+  });
+
+  it("does not request opening when auto-open is disabled", async () => {
+    const onOpenChange = vi.fn();
+    render(
+      <Modal unstable_openOnRunStart={false} onOpenChange={onOpenChange} />,
+    );
+
+    await act(async () => {
+      fireEvent.click(screen.getByRole("button", { name: "Start run" }));
+    });
+
+    expect(onOpenChange).not.toHaveBeenCalled();
+    expect(screen.queryByRole("dialog")).toBeNull();
+  });
+});