@assistant-ui/react 0.14.0 → 0.14.2

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

@@ -0,0 +1,391 @@
+// @vitest-environment jsdom
+import { describe, expect, it, vi } from "vitest";
+import { createMcpAppBridge, type McpAppBridgeFrame } from "./bridge";
+import type {
+  McpAppJsonRpcMessage,
+  McpAppJsonRpcRequest,
+  McpAppJsonRpcResponse,
+} from "./types";
+import { MCP_APP_PROTOCOL_VERSION } from "./types";
+
+type Captured = McpAppJsonRpcMessage;
+
+function makeFrame() {
+  const captured: Captured[] = [];
+  const iframe = document.createElement("iframe");
+  document.body.appendChild(iframe);
+  const frame: McpAppBridgeFrame = {
+    iframe,
+    origin: "https://app.example",
+    sendMessage: (data) => {
+      captured.push(data as Captured);
+    },
+  };
+  return { frame, captured };
+}
+
+function dispatch(frame: McpAppBridgeFrame, message: McpAppJsonRpcMessage) {
+  const event = new MessageEvent("message", {
+    data: message,
+    origin: frame.origin,
+    source: frame.iframe.contentWindow,
+  });
+  window.dispatchEvent(event);
+}
+
+async function flush() {
+  await new Promise((r) => setTimeout(r, 0));
+}
+
+describe("createMcpAppBridge", () => {
+  it("responds to ui/initialize with host info, version, and capabilities", async () => {
+    const { frame, captured } = makeFrame();
+    const bridge = createMcpAppBridge({
+      frame,
+      hostInfo: { name: "test-host", version: "9.9.9" },
+      hostContext: { theme: "dark" },
+      handlers: {
+        callTool: vi.fn(),
+        sendMessage: vi.fn(),
+      },
+    });
+
+    const req: McpAppJsonRpcRequest = {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "ui/initialize",
+    };
+    dispatch(frame, req);
+    await flush();
+
+    expect(captured).toHaveLength(1);
+    const res = captured[0] as McpAppJsonRpcResponse;
+    expect(res.id).toBe(1);
+    const result = res.result as Record<string, any>;
+    expect(result["protocolVersion"]).toBe(MCP_APP_PROTOCOL_VERSION);
+    expect(result["host"]).toEqual({ name: "test-host", version: "9.9.9" });
+    expect(result["hostContext"]).toEqual({ theme: "dark" });
+    expect(result["capabilities"]["tools"]).toBeDefined();
+    expect(result["capabilities"]["ui"]["sendMessage"]).toBe(true);
+    expect(result["capabilities"]["ui"]["openLink"]).toBe(false);
+
+    bridge.dispose();
+  });
+
+  it("routes tools/call to handler", async () => {
+    const { frame, captured } = makeFrame();
+    const callTool = vi.fn().mockResolvedValue({ ok: true });
+    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 7,
+      method: "tools/call",
+      params: { name: "search", arguments: { q: "hi" } },
+    });
+    await flush();
+
+    expect(callTool).toHaveBeenCalledWith({
+      name: "search",
+      arguments: { q: "hi" },
+    });
+    expect(captured[0]).toEqual({
+      jsonrpc: "2.0",
+      id: 7,
+      result: { ok: true },
+    });
+    bridge.dispose();
+  });
+
+  it("rejects tools/call for disallowed tool with -32602", async () => {
+    const { frame, captured } = makeFrame();
+    const callTool = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { callTool, allowedTools: ["search"] },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 2,
+      method: "tools/call",
+      params: { name: "delete_everything" },
+    });
+    await flush();
+
+    expect(callTool).not.toHaveBeenCalled();
+    const res = captured[0] as McpAppJsonRpcResponse;
+    expect(res.error?.code).toBe(-32602);
+    bridge.dispose();
+  });
+
+  it("returns -32601 when no callTool handler", async () => {
+    const { frame, captured } = makeFrame();
+    const bridge = createMcpAppBridge({ frame });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 3,
+      method: "tools/call",
+      params: { name: "x" },
+    });
+    await flush();
+
+    const res = captured[0] as McpAppJsonRpcResponse;
+    expect(res.error?.code).toBe(-32601);
+    bridge.dispose();
+  });
+
+  it("rejects tools/call with non-object arguments via -32602", async () => {
+    const { frame, captured } = makeFrame();
+    const callTool = vi.fn();
+    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 11,
+      method: "tools/call",
+      params: { name: "x", arguments: "not-an-object" },
+    });
+    await flush();
+
+    expect(callTool).not.toHaveBeenCalled();
+    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
+    bridge.dispose();
+  });
+
+  it("rejects requestDisplayMode with unknown mode via -32602", async () => {
+    const { frame, captured } = makeFrame();
+    const requestDisplayMode = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { requestDisplayMode },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 13,
+      method: "requestDisplayMode",
+      params: { mode: "sidebar" },
+    });
+    await flush();
+
+    expect(requestDisplayMode).not.toHaveBeenCalled();
+    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
+    bridge.dispose();
+  });
+
+  it("rejects openLink for non-http(s) schemes via -32602", async () => {
+    const { frame, captured } = makeFrame();
+    const openLink = vi.fn();
+    const bridge = createMcpAppBridge({ frame, handlers: { openLink } });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 12,
+      method: "openLink",
+      params: { url: "javascript:alert(1)" },
+    });
+    await flush();
+
+    expect(openLink).not.toHaveBeenCalled();
+    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
+    bridge.dispose();
+  });
+
+  it("invokes onSizeChange / onInitialized for notifications", () => {
+    const { frame } = makeFrame();
+    const onSizeChange = vi.fn();
+    const onInitialized = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { onSizeChange, onInitialized },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      method: "notifications/size_changed",
+      params: { width: 320, height: 240 },
+    });
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      method: "notifications/initialized",
+    });
+
+    expect(onSizeChange).toHaveBeenCalledWith({ width: 320, height: 240 });
+    expect(onInitialized).toHaveBeenCalled();
+    bridge.dispose();
+  });
+
+  it("notifyToolInput / notifyToolResult / notifyHostContextChanged post correct notifications", () => {
+    const { frame, captured } = makeFrame();
+    const bridge = createMcpAppBridge({ frame });
+
+    bridge.notifyToolInput({ a: 1 });
+    bridge.notifyToolResult({ ok: 1 });
+    bridge.notifyHostContextChanged({ theme: "light" });
+
+    expect(captured).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "notifications/tools/call/input",
+        params: { input: { a: 1 } },
+      },
+      {
+        jsonrpc: "2.0",
+        method: "notifications/tools/call/result",
+        params: { result: { ok: 1 } },
+      },
+      {
+        jsonrpc: "2.0",
+        method: "notifications/host_context/changed",
+        params: { theme: "light" },
+      },
+    ]);
+    bridge.dispose();
+  });
+
+  it("routes resources/read and resources/list to handlers", async () => {
+    const { frame, captured } = makeFrame();
+    const readResource = vi.fn().mockResolvedValue({ contents: [] });
+    const listResources = vi.fn().mockResolvedValue({ resources: [] });
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { readResource, listResources },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 20,
+      method: "resources/read",
+      params: { uri: "ui://app/x" },
+    });
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 21,
+      method: "resources/list",
+    });
+    await flush();
+
+    expect(readResource).toHaveBeenCalledWith({ uri: "ui://app/x" });
+    expect(listResources).toHaveBeenCalled();
+    expect(captured.map((c) => (c as McpAppJsonRpcResponse).id)).toEqual([
+      20, 21,
+    ]);
+    bridge.dispose();
+  });
+
+  it("returns -32601 for resources/read and resources/list when no handler", async () => {
+    const { frame, captured } = makeFrame();
+    const bridge = createMcpAppBridge({ frame });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 22,
+      method: "resources/read",
+      params: { uri: "ui://x" },
+    });
+    dispatch(frame, { jsonrpc: "2.0", id: 23, method: "resources/list" });
+    await flush();
+
+    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32601);
+    expect((captured[1] as McpAppJsonRpcResponse).error?.code).toBe(-32601);
+    bridge.dispose();
+  });
+
+  it("routes sendMessage and updateModelContext to handlers", async () => {
+    const { frame, captured } = makeFrame();
+    const sendMessage = vi.fn().mockResolvedValue({ ok: true });
+    const updateModelContext = vi.fn().mockResolvedValue({ ok: true });
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { sendMessage, updateModelContext },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 30,
+      method: "sendMessage",
+      params: { text: "hi" },
+    });
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      id: 31,
+      method: "updateModelContext",
+      params: { foo: "bar" },
+    });
+    await flush();
+
+    expect(sendMessage).toHaveBeenCalledWith({ text: "hi" });
+    expect(updateModelContext).toHaveBeenCalledWith({ foo: "bar" });
+    expect(captured.map((c) => (c as McpAppJsonRpcResponse).id)).toEqual([
+      30, 31,
+    ]);
+    bridge.dispose();
+  });
+
+  it("invokes onLog / onError / onRequestTeardown for notifications", () => {
+    const { frame } = makeFrame();
+    const onLog = vi.fn();
+    const onError = vi.fn();
+    const onRequestTeardown = vi.fn();
+    const bridge = createMcpAppBridge({
+      frame,
+      handlers: { onLog, onError, onRequestTeardown },
+    });
+
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      method: "notifications/log",
+      params: { level: "info", message: "hello" },
+    });
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      method: "notifications/error",
+      params: { message: "kaboom" },
+    });
+    dispatch(frame, {
+      jsonrpc: "2.0",
+      method: "notifications/request_teardown",
+      params: { reason: "done" },
+    });
+
+    expect(onLog).toHaveBeenCalledWith({ level: "info", message: "hello" });
+    expect(onError).toHaveBeenCalled();
+    expect(onRequestTeardown).toHaveBeenCalledWith({ reason: "done" });
+    bridge.dispose();
+  });
+
+  it("ignores messages from wrong origin or wrong source", async () => {
+    const { frame, captured } = makeFrame();
+    const callTool = vi.fn();
+    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });
+
+    const msg: McpAppJsonRpcMessage = {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "tools/call",
+      params: { name: "search" },
+    };
+
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: msg,
+        origin: "https://attacker.example",
+        source: frame.iframe.contentWindow,
+      }),
+    );
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        data: msg,
+        origin: frame.origin,
+        source: window,
+      }),
+    );
+    await flush();
+
+    expect(callTool).not.toHaveBeenCalled();
+    expect(captured).toHaveLength(0);
+    bridge.dispose();
+  });
+});