@copilotkit/runtime 1.56.5-canary.1777972218 → 1.57.0

package/src/agent/__tests__/mcp-servers-integration.test.ts

@@ -1,9 +1,8 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
-import { BasicAgent, MCPHeaderResolverError } from "../index";
+import { BasicAgent } from "../index";
 import { EventType } from "@ag-ui/client";
 import { streamText } from "ai";
 import { LLMock, MCPMock } from "@copilotkit/aimock";
-import { logger } from "@copilotkit/shared";
 import {
   mockStreamTextResponse,
   textDelta,
@@ -30,81 +29,6 @@ vi.mock("@ai-sdk/openai", () => ({
 // Do NOT mock @ai-sdk/mcp or @modelcontextprotocol/sdk transports —
 // we want real HTTP connections to the MCPMock server.
 
-/**
- * Spin up an LLMock-mounted MCPMock with a real HTTP listener — the mock
- * has to actually respond so the agent's MCP init + tools/list can
- * complete and the run progresses to streamText.
- */
-async function startMcpServerWithJournal(
-  tools: Array<{ name: string; description?: string }>,
-): Promise<{ mcpUrl: string; llm: LLMock; mcpMock: MCPMock }> {
-  const mock = new MCPMock();
-  for (const t of tools) {
-    mock.addTool({
-      name: t.name,
-      description: t.description ?? `${t.name} tool`,
-      inputSchema: {
-        type: "object",
-        properties: { query: { type: "string" } },
-      },
-    });
-    mock.onToolCall(t.name, () => `result from ${t.name}`);
-  }
-  const server = new LLMock({ port: 0 });
-  server.mount("/mcp", mock);
-  await server.start();
-  return { mcpUrl: `${server.url}/mcp`, llm: server, mcpMock: mock };
-}
-
-/**
- * `server.getRequests()` redacts `Authorization` to `[REDACTED]` (aimock
- * privacy feature) — useless when the test needs to see the actual outgoing
- * auth value. Spy on `globalThis.fetch` instead and read the headers off
- * each call's `RequestInit`. The spy preserves the real fetch so MCPMock
- * still responds. Filter to MCP-bound requests by URL substring to ignore
- * any unrelated traffic that might land on the recorder.
- */
-function spyOnFetch(mcpUrl: string): {
-  records: Array<Record<string, string>>;
-  restore: () => void;
-} {
-  const records: Array<Record<string, string>> = [];
-  const realFetch = globalThis.fetch;
-  const spy = vi
-    .spyOn(globalThis, "fetch")
-    .mockImplementation(async (input, init) => {
-      const url =
-        typeof input === "string"
-          ? input
-          : input instanceof URL
-            ? input.toString()
-            : input.url;
-      if (url.startsWith(mcpUrl)) {
-        const seen: Record<string, string> = {};
-        new Headers(init?.headers ?? {}).forEach((value, key) => {
-          seen[key.toLowerCase()] = value;
-        });
-        records.push(seen);
-      }
-      return realFetch(input, init);
-    });
-  return {
-    records,
-    restore: () => spy.mockRestore(),
-  };
-}
-
-/**
- * x-cpki-user-id is NOT in aimock's redaction list, so journal entries
- * carry the real value. Use this when comparing per-call values.
- */
-function userIdsFrom(server: LLMock): string[] {
-  return server
-    .getRequests()
-    .map((entry) => entry.headers?.["x-cpki-user-id"])
-    .filter((v): v is string => typeof v === "string");
-}
-
 describe("mcpServers — real MCP server integration", () => {
   const originalEnv = process.env;
   let llm: LLMock;
@@ -446,412 +370,4 @@ describe("mcpServers — real MCP server integration", () => {
       await llm2.stop().catch(() => {});
     }
   });
-
-  describe("static headers + per-call getHeaders", () => {
-    it("static `headers` are sent on every outbound MCP request (HTTP)", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            {
-              type: "http",
-              url: result.mcpUrl,
-              headers: { Authorization: "Bearer cpk-test-token" },
-            },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([textDelta("Hello"), finish()]) as any,
-        );
-
-        await collectEvents(agent["run"](baseInput));
-
-        expect(recorder.records.length).toBeGreaterThan(0);
-        for (const headers of recorder.records) {
-          expect(headers["authorization"]).toBe("Bearer cpk-test-token");
-        }
-      } finally {
-        recorder.restore();
-      }
-    });
-
-    it("getHeaders runs per outbound HTTP request, not once per session", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      // Counter-based resolver: returns a different user-id on every
-      // invocation. If the SDK opened the connection once and reused
-      // headers (i.e. cached across calls), all requests would carry
-      // the same user-id.
-      let counter = 0;
-      const resolverInvocations: string[] = [];
-      const agent = new BasicAgent({
-        model: "openai/gpt-4o",
-        mcpServers: [
-          {
-            type: "http",
-            url: result.mcpUrl,
-            getHeaders: () => {
-              const id = `user-${counter++}`;
-              resolverInvocations.push(id);
-              return { "X-Cpki-User-Id": id };
-            },
-          },
-        ],
-      });
-
-      vi.mocked(streamText).mockReturnValue(
-        mockStreamTextResponse([textDelta("ok"), finish()]) as any,
-      );
-      await collectEvents(agent["run"](baseInput));
-
-      // The MCP SDK opens with `initialize` and `tools/list`. Both are
-      // wrapped-fetch invocations, both must hit the resolver.
-      expect(resolverInvocations.length).toBeGreaterThanOrEqual(2);
-
-      // Distinct values across requests on the wire prove no caching
-      // happened. x-cpki-user-id is NOT redacted by aimock so we can read
-      // the actual values from the journal.
-      const userIds = userIdsFrom(result.llm);
-      expect(new Set(userIds).size).toBeGreaterThanOrEqual(2);
-    });
-
-    it("getHeaders receives requestHeaders snapshot + input + mcpServerUrl", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const seenContexts: Array<{
-        requestHeaders: Record<string, string>;
-        threadId: string;
-        mcpServerUrl: string;
-      }> = [];
-
-      const agent = new BasicAgent({
-        model: "openai/gpt-4o",
-        mcpServers: [
-          {
-            type: "http",
-            url: result.mcpUrl,
-            getHeaders: ({ requestHeaders, input, mcpServerUrl }) => {
-              seenContexts.push({
-                requestHeaders: { ...requestHeaders },
-                threadId: input.threadId,
-                mcpServerUrl,
-              });
-              return { "X-Cpki-User-Id": "anyone" };
-            },
-          },
-        ],
-      });
-      // Simulate the runtime's `extractForwardableHeaders` populating headers.
-      agent.headers = { "x-cpki-user-id": "from-bff" };
-
-      vi.mocked(streamText).mockReturnValue(
-        mockStreamTextResponse([finish()]) as any,
-      );
-      await collectEvents(agent["run"](baseInput));
-
-      expect(seenContexts.length).toBeGreaterThan(0);
-      const ctx = seenContexts[0];
-      expect(ctx.requestHeaders["x-cpki-user-id"]).toBe("from-bff");
-      expect(ctx.threadId).toBe("thread1");
-      expect(ctx.mcpServerUrl).toBe(result.mcpUrl);
-    });
-
-    it("getHeaders throwing surfaces RUN_ERROR carrying MCPHeaderResolverError", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const agent = new BasicAgent({
-        model: "openai/gpt-4o",
-        mcpServers: [
-          {
-            type: "http",
-            url: result.mcpUrl,
-            getHeaders: () => {
-              throw new Error("BFF forgot to forward X-Cpki-User-Id");
-            },
-          },
-        ],
-      });
-
-      vi.mocked(streamText).mockReturnValue(
-        mockStreamTextResponse([finish()]) as any,
-      );
-
-      const events: any[] = [];
-      try {
-        await new Promise((resolve, reject) => {
-          agent["run"](baseInput).subscribe({
-            next: (event) => events.push(event),
-            error: (err) => reject(err),
-            complete: () => resolve(events),
-          });
-        });
-      } catch {
-        // Expected — resolver threw, fetch wrapper rethrew, transport failed.
-      }
-
-      const runError = events.find((e) => e.type === EventType.RUN_ERROR);
-      expect(runError).toBeDefined();
-      // The wrapped fetch reports through MCPHeaderResolverError so the
-      // run-error message attributes the failure to the resolver, not the
-      // transport. We assert message content here; the original cause is
-      // preserved on the thrown class instance via ES2022 Error.cause
-      // (visible to subscribers of the Observable error notification, not
-      // on the AG-UI run-error event payload).
-      expect(runError?.message).toContain("MCP header resolver");
-      expect(runError?.message).toContain("BFF forgot to forward");
-      // The thrown error class is exported so user code can branch on it.
-      expect(MCPHeaderResolverError).toBeDefined();
-    });
-
-    it("backwards-compat: existing config with no auth fields still loads tools", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [{ type: "http", url: result.mcpUrl }],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([textDelta("ok"), finish()]) as any,
-        );
-        await collectEvents(agent["run"](baseInput));
-
-        const callArgs = vi.mocked(streamText).mock.calls[0][0];
-        expect(callArgs.tools).toHaveProperty("get_weather");
-        // No Authorization or X-Cpki-User-Id on the wire when no auth fields
-        // are configured.
-        for (const headers of recorder.records) {
-          expect(headers["authorization"]).toBeUndefined();
-          expect(headers["x-cpki-user-id"]).toBeUndefined();
-        }
-      } finally {
-        recorder.restore();
-      }
-    });
-
-    it("getHeaders overrides static `headers` when both set Authorization", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            {
-              type: "http",
-              url: result.mcpUrl,
-              headers: { Authorization: "Bearer cpk-static" },
-              getHeaders: () => ({ Authorization: "Bearer cpk-resolver-wins" }),
-            },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-        await collectEvents(agent["run"](baseInput));
-
-        expect(recorder.records.length).toBeGreaterThan(0);
-        for (const headers of recorder.records) {
-          expect(headers["authorization"]).toBe("Bearer cpk-resolver-wins");
-        }
-      } finally {
-        recorder.restore();
-      }
-    });
-
-    it("requiresUser=true: server is skipped when no user is on the agent", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      const resolver = vi.fn(() => ({ "X-Cpki-User-Id": "never" }));
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            {
-              type: "http",
-              url: result.mcpUrl,
-              requiresUser: true,
-              getHeaders: resolver,
-            },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-
-        const events = await collectEvents(agent["run"](baseInput));
-
-        expect(
-          events.find((e) => e.type === EventType.RUN_ERROR),
-        ).toBeUndefined();
-        expect(recorder.records.length).toBe(0);
-        // Skip happens before the transport opens — resolver never invoked.
-        expect(resolver).not.toHaveBeenCalled();
-      } finally {
-        recorder.restore();
-      }
-    });
-
-    it("requiresUser=true: server is included as normal when a user is set", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            {
-              type: "http",
-              url: result.mcpUrl,
-              requiresUser: true,
-              getHeaders: ({ user }) => ({ "X-Cpki-User-Id": user!.id }),
-            },
-          ],
-        });
-        agent.user = { id: "alice", name: "Alice" };
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-
-        await collectEvents(agent["run"](baseInput));
-
-        expect(recorder.records.length).toBeGreaterThan(0);
-        for (const headers of recorder.records) {
-          expect(headers["x-cpki-user-id"]).toBe("alice");
-        }
-      } finally {
-        recorder.restore();
-      }
-    });
-
-    it("requiresUser only skips the flagged server: peers without the flag still load tools", async () => {
-      const flagged = await startMcpServerWithJournal([{ name: "alpha_tool" }]);
-      const peer = await startMcpServerWithJournal([{ name: "beta_tool" }]);
-      llm = flagged.llm;
-      mcpMock = flagged.mcpMock;
-
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            { type: "http", url: flagged.mcpUrl, requiresUser: true },
-            { type: "http", url: peer.mcpUrl },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-
-        await collectEvents(agent["run"](baseInput));
-
-        const callArgs = vi.mocked(streamText).mock.calls[0][0];
-        expect(callArgs.tools).toHaveProperty("beta_tool");
-        expect(callArgs.tools).not.toHaveProperty("alpha_tool");
-      } finally {
-        await peer.llm.stop().catch(() => {});
-      }
-    });
-
-    it("requiresUser=true: skip is logged at warn level", async () => {
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const warnSpy = vi.spyOn(logger, "warn").mockImplementation(() => {});
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            { type: "http", url: result.mcpUrl, requiresUser: true },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-        await collectEvents(agent["run"](baseInput));
-
-        expect(warnSpy).toHaveBeenCalledWith(
-          expect.objectContaining({ url: result.mcpUrl }),
-          expect.stringContaining("Skipping MCP server"),
-        );
-      } finally {
-        warnSpy.mockRestore();
-      }
-    });
-
-    it("static headers reach the wire on the SSE transport (regression for the silently-dropped-headers bug)", async () => {
-      // MCPMock doesn't speak SSE so the connection ultimately fails, but the
-      // initial GET still goes out via fetch — that's enough to verify the
-      // transport actually attaches `headers` to the outbound request, which
-      // a previous direct-SDK construction silently dropped.
-      const result = await startMcpServerWithJournal([{ name: "get_weather" }]);
-      llm = result.llm;
-      mcpMock = result.mcpMock;
-
-      const recorder = spyOnFetch(result.mcpUrl);
-      try {
-        const agent = new BasicAgent({
-          model: "openai/gpt-4o",
-          mcpServers: [
-            {
-              type: "sse",
-              url: result.mcpUrl,
-              headers: { "X-Test-Auth": "sse-static-token" },
-            },
-          ],
-        });
-
-        vi.mocked(streamText).mockReturnValue(
-          mockStreamTextResponse([finish()]) as any,
-        );
-
-        try {
-          await new Promise((resolve, reject) => {
-            agent["run"](baseInput).subscribe({
-              next: () => {},
-              error: (err) => reject(err),
-              complete: () => resolve(undefined),
-            });
-          });
-        } catch {
-          // Expected — SSE init fails because MCPMock doesn't speak SSE.
-        }
-
-        expect(recorder.records.length).toBeGreaterThan(0);
-        for (const headers of recorder.records) {
-          expect(headers["x-test-auth"]).toBe("sse-static-token");
-        }
-      } finally {
-        recorder.restore();
-      }
-    });
-  });
 });