@copilotkit/runtime 1.56.4-canary.1777538870 → 1.56.5

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

@@ -1,5 +1,5 @@
 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";
@@ -29,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;
@@ -445,283 +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("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();
-      }
-    });
-  });
 });

package/src/agent/converters/tanstack.ts

@@ -266,12 +266,30 @@ export async function* convertTanStackStream(
     }
   }
 
+  // TanStack's chat() engine runs a multi-turn agent loop: after the model
+  // returns tool calls, the engine tries to execute them and re-prompt. This
+  // produces a second round of TOOL_CALL_START / TOOL_CALL_END events that
+  // duplicate the ones from the first streaming pass. The CopilotKit runtime
+  // handles tool execution externally (via the frontend SDK), so we must stop
+  // converting events once the TanStack adapter signals the first turn is
+  // complete with RUN_FINISHED.
+  let runFinished = false;
+
   for await (const chunk of stream) {
     if (abortSignal.aborted) break;
 
     const raw = chunk as Record<string, unknown>;
     const type = raw.type as string;
 
+    // Stop converting after the first RUN_FINISHED — any subsequent events
+    // come from TanStack's internal tool-execution loop and would produce
+    // duplicate TOOL_CALL_END events that violate the ag-ui verify middleware.
+    if (type === "RUN_FINISHED") {
+      runFinished = true;
+      continue;
+    }
+    if (runFinished) continue;
+
     if (type === "TEXT_MESSAGE_CONTENT" && raw.delta != null) {
       yield* closeReasoningIfOpen();
       const textEvent: TextMessageChunkEvent = {

package/src/agent/index.ts

@@ -1,6 +1,8 @@
-import type {
+import {
+  AbstractAgent,
   BaseEvent,
   RunAgentInput,
+  EventType,
   Message,
   ReasoningEndEvent,
   ReasoningMessageContentEvent,
@@ -18,9 +20,9 @@ import type {
   StateSnapshotEvent,
   StateDeltaEvent,
 } from "@ag-ui/client";
-import { AbstractAgent, EventType } from "@ag-ui/client";
 import type { AgentCapabilities } from "@ag-ui/core";
-import type {
+import {
+  streamText,
   LanguageModel,
   ModelMessage,
   AssistantModelMessage,
@@ -32,12 +34,12 @@ import type {
   TextPart,
   ImagePart,
   FilePart,
+  tool as createVercelAISDKTool,
   ToolChoice,
   ToolSet,
+  stepCountIs,
 } from "ai";
-import { streamText, tool as createVercelAISDKTool, stepCountIs } from "ai";
-import { createMCPClient } from "@ai-sdk/mcp";
-import type { MCPClient } from "@ai-sdk/mcp";
+import { experimental_createMCPClient as createMCPClient } from "@ai-sdk/mcp";
 import { Observable } from "rxjs";
 import { createOpenAI } from "@ai-sdk/openai";
 import { createAnthropic } from "@ai-sdk/anthropic";
@@ -50,13 +52,11 @@ import { schemaToJsonSchema } from "@copilotkit/shared";
 import { jsonSchema as aiJsonSchema } from "ai";
 import { convertAISDKStream } from "./converters/aisdk";
 import { convertTanStackStream } from "./converters/tanstack";
-import type { StreamableHTTPClientTransportOptions } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import {
-  CopilotKitMCPTransport,
-  MCPHeaderResolverError,
-  type MCPRequestContext,
-  type MCPRuntimeUser,
-} from "./mcp-transport";
+  StreamableHTTPClientTransport,
+  StreamableHTTPClientTransportOptions,
+} from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import { randomUUID } from "@copilotkit/shared";
 
 /**
@@ -113,112 +113,58 @@ export type BuiltInAgentModel =
  */
 export type ModelSpecifier = string | LanguageModel;
 
-// MCPRequestContext and MCPHeaderResolverError now live in mcp-transport.ts.
-// Re-export so existing imports of these symbols from agent/index continue to
-// work.
-export { MCPHeaderResolverError, type MCPRequestContext };
-
 /**
- * MCP Client configuration for HTTP transport.
+ * MCP Client configuration for HTTP transport
  */
 export interface MCPClientConfigHTTP {
-  /** Type of MCP client */
-  type: "http";
-  /** URL of the MCP server */
-  url: string;
   /**
-   * Optional transport options for the underlying
-   * `StreamableHTTPClientTransport`. Pre-existing escape hatch for advanced
-   * use cases (custom `fetch`, `requestInit`, OAuth `authProvider`, etc.).
+   * Type of MCP client
    */
-  options?: StreamableHTTPClientTransportOptions;
+  type: "http";
   /**
-   * Static HTTP headers, merged into every outbound request to this server.
-   * For per-call values, use {@link MCPClientConfigHTTP.getHeaders} instead.
+   * URL of the MCP server
    */
-  headers?: Record<string, string>;
+  url: string;
   /**
-   * Per-call header resolver. Invoked on **every** outbound HTTP request to
-   * this server (initialize, tools/list, tools/call, SSE reconnects). The
-   * returned headers are merged on top of `headers` and any
-   * `options.requestInit.headers`, so a resolver can override either.
-   *
-   * Throwing from the resolver causes the agent run to emit `RUN_ERROR`
-   * carrying a {@link MCPHeaderResolverError}.
+   * Optional transport options for HTTP client
    */
-  getHeaders?: (
-    ctx: MCPRequestContext,
-  ) => Record<string, string> | Promise<Record<string, string>>;
+  options?: StreamableHTTPClientTransportOptions;
 }
 
 /**
- * MCP Client configuration for SSE transport.
+ * MCP Client configuration for SSE transport
  */
 export interface MCPClientConfigSSE {
-  /** Type of MCP client */
+  /**
+   * Type of MCP client
+   */
   type: "sse";
-  /** URL of the MCP server */
+  /**
+   * URL of the MCP server
+   */
   url: string;
-  /** Optional HTTP headers (e.g., for authentication). */
+  /**
+   * Optional HTTP headers (e.g., for authentication)
+   */
   headers?: Record<string, string>;
 }
 
 /**
- * MCP Client configuration.
+ * MCP Client configuration
  */
 export type MCPClientConfig = MCPClientConfigHTTP | MCPClientConfigSSE;
 
 /**
- * A user-managed MCP client that provides tools to the agent. Structural
- * alias for `@ai-sdk/mcp`'s `MCPClient.tools` slice — pass any value built by
- * `createMCPClient()` directly, or supply a custom `{ tools(): ... }` object
- * for tests/caching layers.
+ * A user-managed MCP client that provides tools to the agent.
+ * The user is responsible for creating, configuring, and closing the client.
+ * Compatible with the return type of @ai-sdk/mcp's createMCPClient().
  *
- * Unlike `mcpServers`, the agent does NOT create or close these clients. The
- * user controls the full lifecycle.
+ * Unlike mcpServers, the agent does NOT create or close these clients.
+ * This allows persistent connections, custom auth, and tool caching.
  */
-export type MCPClientProvider = Pick<MCPClient, "tools">;
-
-/**
- * Open an MCP client for the given server config.
- *
- * - HTTP always goes through {@link CopilotKitMCPTransport} (preserves the
- *   pre-existing `options` escape hatch and adds per-call `getHeaders`
- *   resolution).
- * - SSE goes through `@ai-sdk/mcp`'s `createMCPClient`, whose built-in
- *   `SseMCPTransport` correctly applies static `headers` on every outbound
- *   request.
- */
-async function openMcpClient(
-  config: MCPClientConfig,
-  context: {
-    requestHeaders: Record<string, string>;
-    input: RunAgentInput;
-    user?: MCPRuntimeUser;
-  },
-): Promise<MCPClient> {
-  if (config.type === "http") {
-    const transport = new CopilotKitMCPTransport({
-      url: config.url,
-      headers: config.headers,
-      getHeaders: config.getHeaders,
-      options: config.options,
-      requestHeaders: context.requestHeaders,
-      input: context.input,
-      user: context.user,
-    });
-    return createMCPClient({ transport });
-  }
-
-  // SSE: hand to Vercel's transport. Static `headers` are applied via the
-  // SseMCPTransport's common-header pipeline.
-  return createMCPClient({
-    transport: {
-      type: "sse",
-      url: config.url,
-      headers: config.headers,
-    },
-  });
+export interface MCPClientProvider {
+  /** Return tools to be merged into the agent's tool set. */
+  tools(): Promise<ToolSet>;
 }
 
 /**
@@ -905,21 +851,6 @@ function isFactoryConfig(
 
 export class BuiltInAgent extends AbstractAgent {
   private abortController?: AbortController;
-  /**
-   * Headers populated per-request by the runtime's
-   * `extractForwardableHeaders` (the incoming request's `Authorization` +
-   * every `x-*` header, lower-cased). Available to MCP header resolvers via
-   * {@link MCPRequestContext.requestHeaders}; kept here as a plain field so
-   * the runtime's `configureAgentForRequest` feature-detect activates.
-   */
-  public headers: Record<string, string> = {};
-  /**
-   * End-user identity for the current request, populated by the runtime by
-   * invoking `identifyUser(request)` (Intelligence mode). Surfaced to MCP
-   * header resolvers via {@link MCPRequestContext.user}; remains undefined
-   * for runs that aren't going through a runtime with `identifyUser` set.
-   */
-  public user?: { id: string; name: string };
 
   constructor(private config: BuiltInAgentConfiguration) {
     super();
@@ -1173,7 +1104,7 @@ export class BuiltInAgent extends AbstractAgent {
       }
 
       // Set up MCP clients if configured and process the stream
-      const mcpClients: MCPClient[] = [];
+      const mcpClients: Array<{ close: () => Promise<void> }> = [];
 
       (async () => {
         let terminalEventEmitted = false;
@@ -1258,27 +1189,33 @@ export class BuiltInAgent extends AbstractAgent {
 
           // Initialize MCP clients and get their tools
           if (this.config.mcpServers && this.config.mcpServers.length > 0) {
-            // Snapshot the agent's per-run state (forwarded headers + user)
-            // once at run-start. Resolvers see this immutable snapshot for the
-            // lifetime of the run, including any reconnections fired after the
-            // initial run completes.
-            const requestHeaders: Record<string, string> = { ...this.headers };
-            const user = this.user ? { ...this.user } : undefined;
-
             for (const serverConfig of this.config.mcpServers) {
-              const mcpClient = await openMcpClient(serverConfig, {
-                requestHeaders,
-                input,
-                user,
-              });
-              mcpClients.push(mcpClient);
-
-              // Get tools from this MCP server and merge with existing tools
-              const mcpTools = await mcpClient.tools();
-              streamTextParams.tools = {
-                ...streamTextParams.tools,
-                ...mcpTools,
-              } as ToolSet;
+              let transport;
+
+              if (serverConfig.type === "http") {
+                const url = new URL(serverConfig.url);
+                transport = new StreamableHTTPClientTransport(
+                  url,
+                  serverConfig.options,
+                );
+              } else if (serverConfig.type === "sse") {
+                transport = new SSEClientTransport(
+                  new URL(serverConfig.url),
+                  serverConfig.headers,
+                );
+              }
+
+              if (transport) {
+                const mcpClient = await createMCPClient({ transport });
+                mcpClients.push(mcpClient);
+
+                // Get tools from this MCP server and merge with existing tools
+                const mcpTools = await mcpClient.tools();
+                streamTextParams.tools = {
+                  ...streamTextParams.tools,
+                  ...mcpTools,
+                } as ToolSet;
+              }
             }
           }