@assistant-ui/react 0.15.17 → 0.15.19

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

@@ -7,7 +7,8 @@ import type {
   SandboxHostProps,
 } from "../sandbox-host/SandboxHost";
 import type { CreateMcpAppBridgeOptions, McpAppBridge } from "./bridge";
-import { MCP_APP_MIME_TYPE } from "./types";
+import type * as BridgeModule from "./bridge";
+import { MCP_APP_MIME_TYPE, type McpAppHostContext } from "./types";
 
 const { sandboxHostMock, createMcpAppBridgeMock } = vi.hoisted(() => ({
   sandboxHostMock: vi.fn(),
@@ -99,4 +100,439 @@ 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) => {
+      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 view = (hostContext: McpAppHostContext) => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        hostContext={hostContext}
+      />
+    );
+    const rendered = render(
+      view({ displayMode: "inline", availableDisplayModes: ["inline", "pip"] }),
+    );
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+    options.handlers?.onInitialized?.();
+
+    rendered.rerender(
+      view({ displayMode: "inline", availableDisplayModes: ["inline", "pip"] }),
+    );
+    expect(bridge.notifyHostContextChanged).not.toHaveBeenCalled();
+
+    rendered.rerender(
+      view({
+        displayMode: "fullscreen",
+        availableDisplayModes: ["inline", "pip"],
+      }),
+    );
+    expect(bridge.notifyHostContextChanged).toHaveBeenCalledTimes(1);
+    expect(bridge.notifyHostContextChanged).toHaveBeenCalledWith({
+      displayMode: "fullscreen",
+      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", () => {
+    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 view = (hostContext: McpAppHostContext) => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        hostContext={hostContext}
+      />
+    );
+    const rendered = render(view({ displayMode: "inline" }));
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+
+    rendered.rerender(view({ displayMode: "fullscreen" }));
+    rendered.rerender(view({ displayMode: "inline" }));
+
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+    options.handlers?.onInitialized?.();
+
+    expect(bridge.notifyHostContextChanged).not.toHaveBeenCalled();
+
+    sandboxBridge.dispose();
+  });
+
+  it("treats an equal host context as unchanged when it is not strict JSON", () => {
+    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 view = () => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        hostContext={{ displayMode: "inline", workspace: undefined }}
+      />
+    );
+    const rendered = render(view());
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+    options.handlers?.onInitialized?.();
+
+    rendered.rerender(view());
+    expect(bridge.notifyHostContextChanged).not.toHaveBeenCalled();
+
+    sandboxBridge.dispose();
+  });
+  it("notifies when a context value changes behind an unchanged key set", () => {
+    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 view = (updatedAt: Date) => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        hostContext={{ updatedAt }}
+      />
+    );
+    const rendered = render(view(new Date(1)));
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+    options.handlers?.onInitialized?.();
+
+    rendered.rerender(view(new Date(2)));
+    expect(bridge.notifyHostContextChanged).toHaveBeenCalledTimes(1);
+
+    sandboxBridge.dispose();
+  });
+
+  it("does not repeat the empty context the bridge already delivered", () => {
+    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 view = (hostContext?: McpAppHostContext) => (
+      <McpAppFrame
+        app={{ resourceUri: "ui://example/widget" }}
+        resource={{
+          uri: "ui://example/widget",
+          mimeType: MCP_APP_MIME_TYPE,
+          html: "",
+        }}
+        {...(hostContext === undefined ? {} : { hostContext })}
+      />
+    );
+    const rendered = render(view());
+
+    const sandboxBridge = createBridge!(
+      {
+        iframe: document.createElement("iframe"),
+        origin: "https://widget.example",
+        sendMessage: vi.fn(),
+      },
+      { setHeight: vi.fn() },
+    );
+    const options = createMcpAppBridgeMock.mock
+      .calls[0]![0] as CreateMcpAppBridgeOptions;
+    options.handlers?.onInitialized?.();
+
+    rendered.rerender(view({}));
+    expect(bridge.notifyHostContextChanged).not.toHaveBeenCalled();
+
+    sandboxBridge.dispose();
+  });
 });

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

@@ -18,11 +18,48 @@ import type {
   McpAppFrameProps,
   McpAppHostContext,
 } from "./types";
+import { isRecord } from "@assistant-ui/core/internal";
 
 const DEFAULT_PRODUCT = "assistant-ui-mcp-app";
 const INIT_TIMEOUT_MS = 5000;
 const DEFAULT_MAX_HEIGHT = 800;
 
+// Only a plain object can be compared by its own keys. Structured clone carries
+// a Date, Map, or class instance to the widget intact, and those expose no
+// enumerable keys, so walking them would report two different values as equal.
+const isPlainObject = (value: unknown): value is Record<string, unknown> => {
+  if (!isRecord(value)) return false;
+  const prototype = Object.getPrototypeOf(value);
+  return prototype === null || prototype === Object.prototype;
+};
+
+// Hosts commonly rebuild an equal context object on every render, and the
+// widget cannot tell a repeated notification apart from a real change. The
+// open index signature on McpAppHostContext admits values JSON rejects (an
+// undefined property value is enough), so leaves the walk cannot compare fall
+// back to identity rather than failing the whole comparison.
+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)) {
+    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);
+  return (
+    aKeys.length === Object.keys(b).length &&
+    aKeys.every(
+      (key) =>
+        Object.hasOwn(b, key) && isSameHostContext(a[key], b[key], depth + 1),
+    )
+  );
+};
+
 function useBridgeNotify<T>(
   value: T | undefined,
   bridgeRef: MutableRefObject<McpAppBridge | null>,
@@ -30,18 +67,24 @@ function useBridgeNotify<T>(
   pendingRef: MutableRefObject<T | undefined>,
   lastSentRef: MutableRefObject<T | undefined>,
   notify: (bridge: McpAppBridge, v: T) => void,
+  isEqual: (sent: T | undefined, next: T) => boolean = Object.is,
 ) {
   useEffect(() => {
     if (!bridgeRef.current) return;
-    if (value === undefined) return;
-    if (lastSentRef.current === value) return;
+    // Nothing to deliver, either because there is no value or because the
+    // widget already holds it; a queued value here is superseded and would
+    // otherwise flush after initialization as a stale update.
+    if (value === undefined || isEqual(lastSentRef.current, value)) {
+      pendingRef.current = undefined;
+      return;
+    }
     if (!widgetReadyRef.current) {
       pendingRef.current = value;
       return;
     }
     notify(bridgeRef.current, value);
     lastSentRef.current = value;
-    // oxlint-disable-next-line react/exhaustive-deps -- refs are stable; notify is assumed stable; re-run only when value changes
+    // oxlint-disable-next-line react/exhaustive-deps -- refs are stable; notify and isEqual are assumed stable; re-run only when value changes
   }, [value]);
 }
 
@@ -53,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]> =>
@@ -91,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);
@@ -172,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
@@ -208,8 +260,10 @@ export function McpAppFrame({
 
     if (current.input !== undefined) pendingInputRef.current = current.input;
     if (current.output !== undefined) pendingOutputRef.current = current.output;
-    // hostContext is delivered inside the ui/initialize response; subsequent
-    // changes flow through useBridgeNotify's pending path.
+    // hostContext is delivered inside the ui/initialize response, where the
+    // bridge defaults it to {}; recording that same normalized value keeps the
+    // first later change from repeating what the widget already holds.
+    lastSentHostContextRef.current = current.hostContext ?? {};
 
     return {
       onMessage: bridge.onMessage,
@@ -254,6 +308,7 @@ export function McpAppFrame({
     pendingHostContextRef,
     lastSentHostContextRef,
     (b, v) => b.notifyHostContextChanged(v),
+    isSameHostContext,
   );
 
   return (