branch:
wait-mcp-connections.test.ts
2478 bytesRaw
import { describe, it, expect } from "vitest";
import { MessageType } from "../types";
import { connectChatWS } from "./test-utils";
import type { UIMessage as ChatMessage } from "ai";

/**
 * Tests for the waitForMcpConnections config on AIChatAgent.
 *
 * These verify that:
 * - waitForMcpConnections = true doesn't break message processing
 * - waitForMcpConnections = { timeout: N } doesn't break message processing
 * - waitForMcpConnections = false (default) doesn't break message processing
 *
 * The actual waiting logic is tested at the MCPClientManager level in
 * client-manager.test.ts. These tests verify the config plumbing in AIChatAgent.
 */
describe("waitForMcpConnections config", () => {
  async function sendChatAndWaitForDone(agentPath: string) {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/${agentPath}/${room}`);

    const messages: unknown[] = [];
    let resolvePromise: (value: boolean) => void;
    const donePromise = new Promise<boolean>((res) => {
      resolvePromise = res;
    });

    const timeout = setTimeout(() => resolvePromise(false), 5000);

    ws.addEventListener("message", (e: MessageEvent) => {
      const data = JSON.parse(e.data as string);
      messages.push(data);

      if (data.type === MessageType.CF_AGENT_USE_CHAT_RESPONSE && data.done) {
        clearTimeout(timeout);
        resolvePromise(true);
      }
    });

    const chatMessage: ChatMessage = {
      id: "msg1",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };

    ws.send(
      JSON.stringify({
        type: MessageType.CF_AGENT_USE_CHAT_REQUEST,
        id: "req1",
        init: {
          method: "POST",
          body: JSON.stringify({ messages: [chatMessage] })
        }
      })
    );

    const done = await donePromise;
    ws.close();
    return { done, messages };
  }

  it("should process messages with waitForMcpConnections = true", async () => {
    const { done } = await sendChatAndWaitForDone("wait-mcp-true-agent");
    expect(done).toBe(true);
  });

  it("should process messages with waitForMcpConnections = { timeout: 1000 }", async () => {
    const { done } = await sendChatAndWaitForDone("wait-mcp-timeout-agent");
    expect(done).toBe(true);
  });

  it("should process messages with waitForMcpConnections = false (default)", async () => {
    const { done } = await sendChatAndWaitForDone("wait-mcp-false-agent");
    expect(done).toBe(true);
  });
});