@ai-sdk/mcp 2.0.31 → 2.0.33

package/src/tool/mcp-client.ts

@@ -26,14 +26,22 @@ import {
   isCustomMcpTransport,
   type MCPTransport,
   type MCPTransportConfig,
+  type MCPTransportSendOptions,
 } from './mcp-transport';
 import { getMCPAppToolMeta, MCP_APP_MIME_TYPE } from './mcp-apps';
+import {
+  createMCPToolHeaders,
+  getMCPToolHeaderBindings,
+  type MCPToolHeaderBinding,
+} from './mcp-http-headers';
 import {
   CallToolResultSchema,
   CompleteResultSchema,
+  DiscoverResultSchema,
   ElicitationRequestSchema,
   ElicitResultSchema,
   InitializeResultSchema,
+  LATEST_LEGACY_PROTOCOL_VERSION,
   LATEST_PROTOCOL_VERSION,
   ListResourceTemplatesResultSchema,
   ListResourcesResultSchema,
@@ -66,9 +74,12 @@ import {
   type ToolMeta,
   type McpProviderMetadata,
   type InitializeResult,
+  type DiscoverResult,
 } from './types';
 const CLIENT_VERSION = '1.0.0';
 const DEFAULT_MAX_TOOL_CALL_RETRIES = 0;
+const DEFAULT_PROTOCOL_DISCOVERY_TIMEOUT = 1000;
+const MODERN_PROTOCOL_ERROR_CODES = [-32020, -32021, -32022];
 
 const DEFAULT_RETRY_ERROR_CODES = [
   'ConnectionRefused',
@@ -230,6 +241,16 @@ function mcpToModelOutput({
 export interface MCPClientConfig {
   /** Transport configuration for connecting to the MCP server */
   transport: MCPTransportConfig | MCPTransport;
+  /**
+   * Whether transports that support stateless protocol discovery should probe
+   * with `server/discover` before falling back to legacy initialization.
+   *
+   * Disable this for legacy servers that require `initialize` to be the first
+   * request.
+   *
+   * @default true
+   */
+  protocolVersionDiscovery?: boolean;
   /**
    * Options that bound or cancel transport startup and the initialize request.
    */
@@ -388,6 +409,7 @@ export interface MCPClient {
  */
 class DefaultMCPClient implements MCPClient {
   private transport: MCPTransport;
+  private protocolVersionDiscovery: boolean;
   private onUncaughtError?: (error: unknown) => void;
   private maxRetries: number;
   private clientInfo: ClientConfiguration;
@@ -402,11 +424,14 @@ class DefaultMCPClient implements MCPClient {
   private serverCapabilities: ServerCapabilities = {};
   private _serverInfo: Configuration = { name: '', version: '' };
   private _initializeResult: InitializeResult = {
-    protocolVersion: LATEST_PROTOCOL_VERSION,
+    protocolVersion: LATEST_LEGACY_PROTOCOL_VERSION,
     capabilities: {},
     serverInfo: this._serverInfo,
   };
   private _serverInstructions?: string;
+  private protocolEra: 'legacy' | 'modern' = 'legacy';
+  private protocolVersion = LATEST_LEGACY_PROTOCOL_VERSION;
+  private toolHeaderBindings = new Map<string, MCPToolHeaderBinding[]>();
   private isClosed = true;
   private elicitationRequestHandler?: (
     request: ElicitationRequest,
@@ -422,12 +447,14 @@ class DefaultMCPClient implements MCPClient {
     capabilities,
     initialInitializeResult,
     initializationOptions,
+    protocolVersionDiscovery = true,
   }: MCPClientConfig) {
     this.onUncaughtError = onUncaughtError;
     this.maxRetries = prepareMaxRetries(maxRetries);
     this.clientCapabilities = capabilities ?? {};
     this.initialInitializeResult = initialInitializeResult;
     this.initializationOptions = initializationOptions;
+    this.protocolVersionDiscovery = protocolVersionDiscovery;
 
     if (isCustomMcpTransport(transportConfig)) {
       this.transport = transportConfig;
@@ -510,11 +537,25 @@ class DefaultMCPClient implements MCPClient {
         return this;
       }
 
+      if (
+        this.protocolVersionDiscovery &&
+        this.transport.supportsProtocolVersionDiscovery
+      ) {
+        const discovered = await this.tryProtocolDiscovery(signal);
+        if (discovered) {
+          return this;
+        }
+      }
+
+      this.protocolEra = 'legacy';
+      this.protocolVersion = LATEST_LEGACY_PROTOCOL_VERSION;
+      this.setTransportProtocolVersion(this.protocolVersion);
+
       const result = await this.request({
         request: {
           method: 'initialize',
           params: {
-            protocolVersion: LATEST_PROTOCOL_VERSION,
+            protocolVersion: LATEST_LEGACY_PROTOCOL_VERSION,
             capabilities: this.clientCapabilities,
             clientInfo: this.clientInfo,
           },
@@ -565,6 +606,75 @@ class DefaultMCPClient implements MCPClient {
     }
   }
 
+  private async tryProtocolDiscovery(
+    signal: AbortSignal | undefined,
+  ): Promise<boolean> {
+    this.protocolEra = 'modern';
+    this.protocolVersion = LATEST_PROTOCOL_VERSION;
+    this.setTransportProtocolVersion(this.protocolVersion);
+
+    try {
+      const result = await this.request({
+        request: { method: 'server/discover' },
+        resultSchema: DiscoverResultSchema,
+        options: {
+          signal,
+          timeout: DEFAULT_PROTOCOL_DISCOVERY_TIMEOUT,
+        },
+      });
+
+      this.applyDiscoverResult(result);
+      return true;
+    } catch (error) {
+      if (
+        MCPClientError.isInstance(error) &&
+        error.code != null &&
+        MODERN_PROTOCOL_ERROR_CODES.includes(error.code)
+      ) {
+        throw error;
+      }
+
+      return false;
+    }
+  }
+
+  private applyDiscoverResult(result: DiscoverResult): void {
+    if (!result.supportedVersions.includes(this.protocolVersion)) {
+      throw new MCPClientError({
+        message: `Server does not support the requested protocol version: ${this.protocolVersion}`,
+      });
+    }
+
+    const serverInfo = result._meta?.['io.modelcontextprotocol/serverInfo'];
+    if (
+      serverInfo != null &&
+      typeof serverInfo === 'object' &&
+      'name' in serverInfo &&
+      typeof serverInfo.name === 'string' &&
+      'version' in serverInfo &&
+      typeof serverInfo.version === 'string'
+    ) {
+      this._serverInfo = serverInfo as Configuration;
+    }
+
+    this.serverCapabilities = result.capabilities;
+    this._serverInstructions = result.instructions;
+    this._initializeResult = {
+      protocolVersion: this.protocolVersion,
+      capabilities: result.capabilities,
+      serverInfo: this._serverInfo,
+      instructions: result.instructions,
+    };
+  }
+
+  private setTransportProtocolVersion(version: string): void {
+    if (this.transport.setProtocolVersion) {
+      this.transport.setProtocolVersion(version);
+    } else {
+      this.transport.protocolVersion = version;
+    }
+  }
+
   private applyInitializeResult(result: InitializeResult): void {
     if (!SUPPORTED_PROTOCOL_VERSIONS.includes(result.protocolVersion)) {
       throw new MCPClientError({
@@ -573,13 +683,11 @@ class DefaultMCPClient implements MCPClient {
     }
 
     this.serverCapabilities = result.capabilities;
+    this.protocolEra = 'legacy';
+    this.protocolVersion = result.protocolVersion;
     this._serverInfo = result.serverInfo;
     this._initializeResult = result;
-    if (this.transport.setProtocolVersion) {
-      this.transport.setProtocolVersion(result.protocolVersion);
-    } else {
-      this.transport.protocolVersion = result.protocolVersion;
-    }
+    this.setTransportProtocolVersion(result.protocolVersion);
     this._serverInstructions = result.instructions;
   }
 
@@ -591,17 +699,17 @@ class DefaultMCPClient implements MCPClient {
 
   private send(
     message: JSONRPCMessage,
-    signal: AbortSignal | undefined,
+    options?: MCPTransportSendOptions,
   ): Promise<void> {
-    return this.transport.send(
-      message,
-      signal == null ? undefined : { signal },
-    );
+    return options == null
+      ? this.transport.send(message)
+      : this.transport.send(message, options);
   }
 
   private assertCapability(method: string): void {
     switch (method) {
       case 'initialize':
+      case 'server/discover':
         break;
       case 'completion/complete':
         if (!this.serverCapabilities.completions) {
@@ -677,11 +785,29 @@ class DefaultMCPClient implements MCPClient {
       let timeoutId: ReturnType<typeof setTimeout> | undefined;
 
       const messageId = this.requestMessageId++;
+      const preparedRequest =
+        this.protocolEra === 'modern'
+          ? {
+              ...request,
+              params: {
+                ...request.params,
+                _meta: {
+                  ...request.params?._meta,
+                  'io.modelcontextprotocol/protocolVersion':
+                    this.protocolVersion,
+                  'io.modelcontextprotocol/clientCapabilities':
+                    this.clientCapabilities,
+                  'io.modelcontextprotocol/clientInfo': this.clientInfo,
+                },
+              },
+            }
+          : request;
       const jsonrpcRequest: JSONRPCRequest = {
-        ...request,
+        ...preparedRequest,
         jsonrpc: '2.0',
         id: messageId,
       };
+      const headers = this.getToolRequestHeaders(preparedRequest);
 
       const rejectWithAbortError = () => {
         reject(
@@ -729,14 +855,31 @@ class DefaultMCPClient implements MCPClient {
         }
 
         try {
+          if (
+            this.protocolEra === 'modern' &&
+            response.result.resultType == null
+          ) {
+            throw new MCPClientError({
+              message: 'Modern MCP result is missing resultType',
+            });
+          }
+          if (response.result.resultType === 'input_required') {
+            throw new MCPClientError({
+              message:
+                'Server requested additional input, but multi round-trip requests are not supported yet',
+            });
+          }
+
           const result = resultSchema.parse(response.result);
           cleanup();
           resolve(result);
         } catch (error) {
-          const parseError = new MCPClientError({
-            message: 'Failed to parse server response',
-            cause: error,
-          });
+          const parseError = MCPClientError.isInstance(error)
+            ? error
+            : new MCPClientError({
+                message: 'Failed to parse server response',
+                cause: error,
+              });
           rejectAndCleanup(parseError);
         }
       });
@@ -747,10 +890,14 @@ class DefaultMCPClient implements MCPClient {
         timeoutId = setTimeout(onTimeout, timeout);
       }
 
+      const sendOptions: MCPTransportSendOptions = {
+        ...(transportSignal == null ? {} : { signal: transportSignal }),
+        ...(headers == null ? {} : { headers }),
+      };
       const sendPromise =
-        transportSignal == null
-          ? this.transport.send(jsonrpcRequest)
-          : this.send(jsonrpcRequest, transportSignal);
+        Object.keys(sendOptions).length === 0
+          ? this.send(jsonrpcRequest)
+          : this.send(jsonrpcRequest, sendOptions);
 
       sendPromise.catch(error => {
         rejectAndCleanup(error);
@@ -765,11 +912,78 @@ class DefaultMCPClient implements MCPClient {
     params?: PaginatedRequest['params'];
     options?: RequestOptions;
   } = {}): Promise<ListToolsResult> {
-    return this.request({
+    const result = await this.request({
       request: { method: 'tools/list', params },
       resultSchema: ListToolsResultSchema,
       options,
     });
+    return this.prepareToolDefinitions(result, params?.cursor == null);
+  }
+
+  private prepareToolDefinitions(
+    definitions: ListToolsResult,
+    resetHeaderBindings = false,
+  ): ListToolsResult {
+    if (
+      this.protocolEra !== 'modern' ||
+      !this.transport.supportsMcpToolParameterHeaders
+    ) {
+      return definitions;
+    }
+
+    if (resetHeaderBindings) {
+      this.toolHeaderBindings.clear();
+    }
+    const tools = definitions.tools.filter(toolDefinition => {
+      const result = getMCPToolHeaderBindings(toolDefinition.inputSchema);
+      if (!result.success) {
+        this.onError(
+          new MCPClientError({
+            message: `Ignoring MCP tool "${toolDefinition.name}": ${result.error}`,
+          }),
+        );
+        return false;
+      }
+
+      this.toolHeaderBindings.set(toolDefinition.name, result.bindings);
+      return true;
+    });
+
+    return { ...definitions, tools };
+  }
+
+  private getToolRequestHeaders(
+    request: Request,
+  ): Record<string, string> | undefined {
+    if (
+      this.protocolEra !== 'modern' ||
+      request.method !== 'tools/call' ||
+      typeof request.params?.name !== 'string'
+    ) {
+      return undefined;
+    }
+
+    const bindings = this.toolHeaderBindings.get(request.params.name);
+    if (bindings == null || bindings.length === 0) {
+      return undefined;
+    }
+
+    const args = request.params.arguments;
+    if (args == null || typeof args !== 'object' || Array.isArray(args)) {
+      return undefined;
+    }
+
+    try {
+      return createMCPToolHeaders({
+        bindings,
+        args: args as Record<string, unknown>,
+      });
+    } catch (error) {
+      throw new MCPClientError({
+        message: `Failed to create MCP headers for tool "${request.params.name}"`,
+        cause: error,
+      });
+    }
   }
 
   private async callToolWithRetry({
@@ -934,7 +1148,10 @@ class DefaultMCPClient implements MCPClient {
       jsonrpc: '2.0',
     };
     await waitForAbort(
-      this.send(jsonrpcNotification, options?.signal),
+      this.send(
+        jsonrpcNotification,
+        options?.signal == null ? undefined : { signal: options.signal },
+      ),
       options?.signal,
     );
   }
@@ -964,6 +1181,7 @@ class DefaultMCPClient implements MCPClient {
       schemas?: TOOL_SCHEMAS;
     },
   ): McpToolSet<TOOL_SCHEMAS> {
+    definitions = this.prepareToolDefinitions(definitions);
     const tools: Record<string, Tool & { _meta?: ToolMeta }> = {};
 
     for (const {
@@ -1278,6 +1496,17 @@ class DefaultMCPClient implements MCPClient {
   }
 
   private onResponse(response: JSONRPCResponse | JSONRPCError): void {
+    if (response.id == null) {
+      this.onError(
+        new MCPClientError({
+          message: `Protocol error: Received a response without a message ID: ${JSON.stringify(
+            response,
+          )}`,
+        }),
+      );
+      return;
+    }
+
     const messageId = Number(response.id);
     const handler = this.responseHandlers.get(messageId);
 

package/src/tool/mcp-http-headers.ts

@@ -0,0 +1,161 @@
+import { convertUint8ArrayToBase64, isRecord } from '@ai-sdk/provider-utils';
+
+type HeaderValueType = 'boolean' | 'integer' | 'string';
+
+export type MCPToolHeaderBinding = {
+  headerName: string;
+  path: string[];
+  valueType: HeaderValueType;
+};
+
+export type MCPToolHeaderBindingsResult =
+  | { success: true; bindings: MCPToolHeaderBinding[] }
+  | { success: false; error: string };
+
+const HTTP_TOKEN_PATTERN = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/;
+const BASE64_SENTINEL_PATTERN = /^=\?base64\?.*\?=$/;
+
+export function encodeMCPHeaderValue(value: string): string {
+  const isPlainAscii = [...value].every(character => {
+    const code = character.charCodeAt(0);
+    return code === 0x09 || (code >= 0x20 && code <= 0x7e);
+  });
+
+  if (
+    isPlainAscii &&
+    value.trim() === value &&
+    !BASE64_SENTINEL_PATTERN.test(value)
+  ) {
+    return value;
+  }
+
+  return `=?base64?${convertUint8ArrayToBase64(new TextEncoder().encode(value))}?=`;
+}
+
+export function getMCPToolHeaderBindings(
+  inputSchema: unknown,
+): MCPToolHeaderBindingsResult {
+  if (!isRecord(inputSchema)) {
+    return {
+      success: false,
+      error: 'inputSchema must be a JSON Schema object',
+    };
+  }
+
+  const bindings: MCPToolHeaderBinding[] = [];
+  const headerNames = new Set<string>();
+  let error: string | undefined;
+
+  const visit = (
+    value: unknown,
+    path: string[],
+    staticallyReachable: boolean,
+  ): void => {
+    if (error != null || !isRecord(value)) {
+      return;
+    }
+
+    if ('x-mcp-header' in value) {
+      if (!staticallyReachable || path.length === 0) {
+        error = 'x-mcp-header is not on a statically reachable property';
+        return;
+      }
+
+      const headerName = value['x-mcp-header'];
+      if (
+        typeof headerName !== 'string' ||
+        !HTTP_TOKEN_PATTERN.test(headerName)
+      ) {
+        error = 'x-mcp-header must be a non-empty HTTP token';
+        return;
+      }
+
+      const normalizedHeaderName = headerName.toLowerCase();
+      if (headerNames.has(normalizedHeaderName)) {
+        error = `x-mcp-header value "${headerName}" is not unique`;
+        return;
+      }
+
+      const valueType = value.type;
+      if (
+        valueType !== 'boolean' &&
+        valueType !== 'integer' &&
+        valueType !== 'string'
+      ) {
+        error =
+          'x-mcp-header can only annotate boolean, integer, or string properties';
+        return;
+      }
+
+      headerNames.add(normalizedHeaderName);
+      bindings.push({ headerName, path, valueType });
+    }
+
+    for (const [key, child] of Object.entries(value)) {
+      if (key === 'x-mcp-header') {
+        continue;
+      }
+
+      if (key === 'properties' && isRecord(child)) {
+        for (const [propertyName, propertySchema] of Object.entries(child)) {
+          visit(propertySchema, [...path, propertyName], staticallyReachable);
+        }
+      } else {
+        visit(child, path, false);
+      }
+    }
+  };
+
+  visit(inputSchema, [], true);
+
+  return error == null
+    ? { success: true, bindings }
+    : { success: false, error };
+}
+
+function getValueAtPath(
+  value: Record<string, unknown>,
+  path: string[],
+): unknown {
+  let current: unknown = value;
+  for (const segment of path) {
+    if (!isRecord(current)) {
+      return undefined;
+    }
+    current = current[segment];
+  }
+  return current;
+}
+
+export function createMCPToolHeaders({
+  bindings,
+  args,
+}: {
+  bindings: MCPToolHeaderBinding[];
+  args: Record<string, unknown>;
+}): Record<string, string> {
+  const headers: Record<string, string> = {};
+
+  for (const binding of bindings) {
+    const value = getValueAtPath(args, binding.path);
+    if (value == null) {
+      continue;
+    }
+
+    if (
+      (binding.valueType === 'string' && typeof value !== 'string') ||
+      (binding.valueType === 'boolean' && typeof value !== 'boolean') ||
+      (binding.valueType === 'integer' && !Number.isSafeInteger(value))
+    ) {
+      throw new TypeError(
+        `Tool argument "${binding.path.join('.')}" does not match its x-mcp-header type`,
+      );
+    }
+
+    headers[`Mcp-Param-${binding.headerName}`] = encodeMCPHeaderValue(
+      String(value),
+    );
+  }
+
+  return headers;
+}