@copilotkit/react-core 1.55.1 → 1.55.2-next.0

package/src/v2/components/chat/__tests__/MCPAppsUiMessage.e2e.test.tsx

@@ -0,0 +1,458 @@
+/**
+ * Tests for MCP Apps ui/message handler behavior.
+ *
+ * Verifies the followUp logic that controls whether the agent is invoked
+ * after an MCP app sends a ui/message request via JSON-RPC:
+ *
+ *   shouldFollowUp = params.followUp ?? role === "user"
+ *
+ * - User-role messages invoke runAgent (default followUp = true)
+ * - Assistant-role messages do NOT invoke runAgent (default followUp = false)
+ * - followUp: false skips runAgent regardless of role
+ * - followUp: true forces runAgent regardless of role
+ * - addMessage is always called when textContent is present
+ */
+import { fireEvent, screen, waitFor, act } from "@testing-library/react";
+import { vi } from "vitest";
+import {
+  activitySnapshotEvent,
+  renderWithCopilotKit,
+  runFinishedEvent,
+  runStartedEvent,
+  testId,
+} from "../../../__tests__/utils/test-helpers";
+import { MCPAppsActivityType } from "../../../components/MCPAppsActivityRenderer";
+import {
+  AbstractAgent,
+  RunAgentInput,
+  RunAgentResult,
+  BaseEvent,
+  EventType,
+} from "@ag-ui/client";
+import { Observable, Subject } from "rxjs";
+
+/**
+ * MockMCPProxyAgent with spying support for ui/message tests.
+ */
+class MockMCPProxyAgent extends AbstractAgent {
+  private subject = new Subject<BaseEvent>();
+  public runAgentCalls: Array<{ input: Partial<RunAgentInput> }> = [];
+  public addMessageCalls: Array<{
+    id: string;
+    role: string;
+    content: string;
+  }> = [];
+
+  private runAgentResponses: Map<string, unknown> = new Map();
+
+  setRunAgentResponse(method: string, response: unknown) {
+    this.runAgentResponses.set(method, response);
+  }
+
+  emit(event: BaseEvent) {
+    if (event.type === EventType.RUN_STARTED) {
+      this.isRunning = true;
+    } else if (
+      event.type === EventType.RUN_FINISHED ||
+      event.type === EventType.RUN_ERROR
+    ) {
+      this.isRunning = false;
+    }
+    act(() => {
+      this.subject.next(event);
+    });
+  }
+
+  complete() {
+    this.isRunning = false;
+    act(() => {
+      this.subject.complete();
+    });
+  }
+
+  clone(): MockMCPProxyAgent {
+    const cloned = new MockMCPProxyAgent();
+    cloned.agentId = this.agentId;
+    type Internal = {
+      subject: Subject<BaseEvent>;
+      runAgentCalls: Array<{ input: Partial<RunAgentInput> }>;
+      addMessageCalls: Array<{ id: string; role: string; content: string }>;
+      runAgentResponses: Map<string, unknown>;
+    };
+    (cloned as unknown as Internal).subject = (
+      this as unknown as Internal
+    ).subject;
+    (cloned as unknown as Internal).runAgentCalls = (
+      this as unknown as Internal
+    ).runAgentCalls;
+    (cloned as unknown as Internal).addMessageCalls = (
+      this as unknown as Internal
+    ).addMessageCalls;
+    (cloned as unknown as Internal).runAgentResponses = (
+      this as unknown as Internal
+    ).runAgentResponses;
+
+    const registry = this;
+    Object.defineProperty(cloned, "isRunning", {
+      get() {
+        return registry.isRunning;
+      },
+      set(v: boolean) {
+        registry.isRunning = v;
+      },
+      configurable: true,
+      enumerable: true,
+    });
+
+    const proto = MockMCPProxyAgent.prototype;
+    cloned.runAgent = async function (
+      input?: Partial<RunAgentInput>,
+    ): Promise<RunAgentResult> {
+      const proxiedRequest = input?.forwardedProps?.__proxiedMCPRequest;
+      if (proxiedRequest) {
+        return registry.runAgent(input);
+      }
+      return proto.runAgent.call(cloned, input);
+    };
+
+    // Track addMessage calls on the clone (the component uses the clone)
+    const origAddMessage = cloned.addMessage.bind(cloned);
+    cloned.addMessage = function (msg: Parameters<typeof origAddMessage>[0]) {
+      registry.addMessageCalls.push(msg as any);
+      return origAddMessage(msg);
+    };
+
+    // Proxy run() calls so spies on the registry's run() see clone invocations
+    cloned.run = function (input: RunAgentInput): Observable<BaseEvent> {
+      return registry.run(input);
+    };
+
+    return cloned;
+  }
+
+  async detachActiveRun(): Promise<void> {}
+
+  run(_input: RunAgentInput): Observable<BaseEvent> {
+    return this.subject.asObservable();
+  }
+
+  async runAgent(input?: Partial<RunAgentInput>): Promise<RunAgentResult> {
+    const proxiedRequest = input?.forwardedProps?.__proxiedMCPRequest as
+      | {
+          serverHash?: string;
+          serverId?: string;
+          method: string;
+          params?: Record<string, unknown>;
+        }
+      | undefined;
+
+    if (proxiedRequest) {
+      if (input) {
+        this.runAgentCalls.push({ input });
+      }
+      const method = proxiedRequest.method;
+      const response = this.runAgentResponses.get(method);
+      if (response !== undefined) {
+        return { result: response, newMessages: [] };
+      }
+      if (method === "resources/read") {
+        return {
+          result: {
+            contents: [
+              {
+                uri: proxiedRequest.params?.uri,
+                mimeType: "text/html",
+                text: "<html><body>Test content</body></html>",
+              },
+            ],
+          },
+          newMessages: [],
+        };
+      }
+      return { result: {}, newMessages: [] };
+    }
+
+    return super.runAgent(input);
+  }
+}
+
+function mcpAppsActivityContent(overrides: {
+  resourceUri?: string;
+  serverHash?: string;
+}) {
+  return {
+    resourceUri: overrides.resourceUri ?? "ui://test-server/test-resource",
+    serverHash: overrides.serverHash ?? "abc123hash",
+    toolInput: {},
+    result: {
+      content: [{ type: "text", text: "Tool output" }],
+      isError: false,
+    },
+  };
+}
+
+/**
+ * Set up the agent, render, emit MCP activity, wait for iframe creation,
+ * then simulate sandbox-proxy-ready so the message handler gets installed.
+ */
+async function setupMCPActivity(
+  agent: MockMCPProxyAgent,
+  agentId: string,
+  userMessage: string,
+): Promise<HTMLIFrameElement> {
+  agent.setRunAgentResponse("resources/read", {
+    contents: [
+      {
+        uri: "ui://test/app",
+        mimeType: "text/html",
+        text: "<html><body>App</body></html>",
+      },
+    ],
+  });
+
+  // Use a unique threadId per test to avoid module-level mcpAppsRequestQueue
+  // state leaking between tests (the queue keys by threadId).
+  const threadId = testId("thread");
+
+  renderWithCopilotKit({
+    agents: { [agentId]: agent },
+    agentId,
+    threadId,
+  });
+
+  const input = await screen.findByRole("textbox");
+  fireEvent.change(input, { target: { value: userMessage } });
+  fireEvent.keyDown(input, { key: "Enter", code: "Enter" });
+
+  await waitFor(() => {
+    expect(screen.getByText(userMessage)).toBeDefined();
+  });
+
+  agent.emit(runStartedEvent());
+  agent.emit(
+    activitySnapshotEvent({
+      messageId: testId("mcp-activity"),
+      activityType: MCPAppsActivityType,
+      content: mcpAppsActivityContent({
+        resourceUri: "ui://test/app",
+        serverHash: "test-hash",
+      }),
+    }),
+  );
+  agent.emit(runFinishedEvent());
+
+  // Wait for iframe to be created
+  let iframe: HTMLIFrameElement | null = null;
+  await waitFor(
+    () => {
+      iframe = document.querySelector("iframe[srcdoc]");
+      expect(iframe).not.toBeNull();
+    },
+    { timeout: 3000 },
+  );
+
+  // Simulate sandbox-proxy-ready notification from the iframe.
+  // The message handler checks event.source === iframe.contentWindow.
+  // In jsdom, iframe.contentWindow exists for srcdoc iframes.
+  const readyEvent = new MessageEvent("message", {
+    data: {
+      jsonrpc: "2.0",
+      method: "ui/notifications/sandbox-proxy-ready",
+    },
+    source: iframe!.contentWindow,
+    origin: "",
+  });
+
+  await act(async () => {
+    window.dispatchEvent(readyEvent);
+    // Give async setup() time to install the messageHandler
+    await new Promise((resolve) => setTimeout(resolve, 100));
+  });
+
+  return iframe!;
+}
+
+/**
+ * Send a ui/message JSON-RPC request as if coming from the iframe.
+ */
+async function sendUiMessage(
+  iframe: HTMLIFrameElement,
+  params: {
+    role?: string;
+    content?: Array<{ type: string; text?: string }>;
+    followUp?: boolean;
+  },
+) {
+  const msg = new MessageEvent("message", {
+    data: {
+      jsonrpc: "2.0",
+      id: testId("req"),
+      method: "ui/message",
+      params,
+    },
+    source: iframe.contentWindow,
+    origin: "",
+  });
+
+  await act(async () => {
+    window.dispatchEvent(msg);
+    await new Promise((resolve) => setTimeout(resolve, 200));
+  });
+}
+
+describe("MCP Apps ui/message followUp behavior", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("user-role message: addMessage IS called and runAgent IS invoked", async () => {
+    const agent = new MockMCPProxyAgent();
+    agent.agentId = "ui-msg-agent-user";
+
+    const iframe = await setupMCPActivity(
+      agent,
+      "ui-msg-agent-user",
+      "User role test",
+    );
+
+    const runSpy = vi.spyOn(agent, "run");
+
+    await sendUiMessage(iframe, {
+      role: "user",
+      content: [{ type: "text", text: "Hello from MCP app" }],
+    });
+
+    // addMessage should have been called
+    const userMsgCalls = agent.addMessageCalls.filter(
+      (c) => c.content === "Hello from MCP app" && c.role === "user",
+    );
+    expect(userMsgCalls.length).toBeGreaterThanOrEqual(1);
+
+    // runAgent should have been invoked (user role defaults to followUp: true)
+    expect(runSpy.mock.calls.length).toBeGreaterThan(0);
+  });
+
+  it("assistant-role message: addMessage IS called but runAgent is NOT invoked", async () => {
+    const agent = new MockMCPProxyAgent();
+    agent.agentId = "ui-msg-agent-assist";
+
+    const iframe = await setupMCPActivity(
+      agent,
+      "ui-msg-agent-assist",
+      "Assist role test",
+    );
+
+    const runSpy = vi.spyOn(agent, "run");
+
+    await sendUiMessage(iframe, {
+      role: "assistant",
+      content: [{ type: "text", text: "Response from MCP" }],
+    });
+
+    // addMessage should have been called
+    const assistCalls = agent.addMessageCalls.filter(
+      (c) => c.content === "Response from MCP" && c.role === "assistant",
+    );
+    expect(assistCalls.length).toBeGreaterThanOrEqual(1);
+
+    // run() should NOT have been called
+    expect(runSpy.mock.calls.length).toBe(0);
+  });
+
+  it("followUp: false on user-role message: addMessage IS called but runAgent is NOT invoked", async () => {
+    const agent = new MockMCPProxyAgent();
+    agent.agentId = "ui-msg-agent-nofollowup";
+
+    const iframe = await setupMCPActivity(
+      agent,
+      "ui-msg-agent-nofollowup",
+      "No followUp test",
+    );
+
+    const runSpy = vi.spyOn(agent, "run");
+
+    await sendUiMessage(iframe, {
+      role: "user",
+      content: [{ type: "text", text: "Display only message" }],
+      followUp: false,
+    });
+
+    // addMessage should have been called
+    const calls = agent.addMessageCalls.filter(
+      (c) => c.content === "Display only message",
+    );
+    expect(calls.length).toBeGreaterThanOrEqual(1);
+
+    // run() should NOT have been called
+    expect(runSpy.mock.calls.length).toBe(0);
+  });
+
+  it("followUp: true on assistant-role message: addMessage IS called AND runAgent IS invoked", async () => {
+    const agent = new MockMCPProxyAgent();
+    agent.agentId = "ui-msg-agent-force";
+
+    const iframe = await setupMCPActivity(
+      agent,
+      "ui-msg-agent-force",
+      "Force followUp test",
+    );
+
+    const runSpy = vi.spyOn(agent, "run");
+
+    await sendUiMessage(iframe, {
+      role: "assistant",
+      content: [{ type: "text", text: "Assistant with followUp" }],
+      followUp: true,
+    });
+
+    // addMessage should have been called
+    const calls = agent.addMessageCalls.filter(
+      (c) => c.content === "Assistant with followUp",
+    );
+    expect(calls.length).toBeGreaterThanOrEqual(1);
+
+    // run() should have been called
+    expect(runSpy.mock.calls.length).toBeGreaterThan(0);
+  });
+
+  it("message with text content always adds to agent messages regardless of followUp", async () => {
+    const agent = new MockMCPProxyAgent();
+    agent.agentId = "ui-msg-agent-all";
+
+    const iframe = await setupMCPActivity(
+      agent,
+      "ui-msg-agent-all",
+      "All messages test",
+    );
+
+    await sendUiMessage(iframe, {
+      role: "user",
+      content: [{ type: "text", text: "User msg" }],
+    });
+
+    await sendUiMessage(iframe, {
+      role: "assistant",
+      content: [{ type: "text", text: "Assistant msg" }],
+    });
+
+    await sendUiMessage(iframe, {
+      role: "user",
+      content: [{ type: "text", text: "No followUp msg" }],
+      followUp: false,
+    });
+
+    const userCalls = agent.addMessageCalls.filter(
+      (c) => c.content === "User msg",
+    );
+    const assistCalls = agent.addMessageCalls.filter(
+      (c) => c.content === "Assistant msg",
+    );
+    const noFollowCalls = agent.addMessageCalls.filter(
+      (c) => c.content === "No followUp msg",
+    );
+
+    expect(userCalls.length).toBeGreaterThanOrEqual(1);
+    expect(assistCalls.length).toBeGreaterThanOrEqual(1);
+    expect(noFollowCalls.length).toBeGreaterThanOrEqual(1);
+  });
+});