From c1fcbf6219a8110027cc0a72e93c27f8476f0c10 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 5 Mar 2026 18:29:32 -0800 Subject: [PATCH] fix(mcp): fix OAuth2 BYOK token injection for OpenAPI tools Three bugs prevented stored OAuth2 tokens from being injected as Bearer headers when OpenAPI MCP tools were called: 1. _get_tools_from_server called _create_mcp_client before checking spec_path, so for BYOK OAuth2 servers resolve_mcp_auth tried a client_credentials token exchange (GitHub rejects this with a non-JSON body), causing a JSON parse error that made tools/list return []. Fix: check spec_path first and read from tool registry directly, skipping MCP client creation entirely for OpenAPI servers. 2. REST tool call passed user_api_key_auth=data.get("user_api_key_auth") which is None unless set in request metadata, so _get_byok_credential couldn't look up the stored token by user_id. Fix: fall back to user_api_key_dict from the route dependency. 3. OpenAPI tools are registered in the local registry under the prefixed name (e.g. github_user-get_authenticated_user) but callers may use the bare name (get_authenticated_user). get_tool(bare_name) returned None and execution fell through to the managed MCP client path. Fix: when local_tool is None and we know the server has a spec_path, retry get_tool with the prefixed name. Also fixes load_servers_from_config to propagate is_byok, byok_description, and byok_api_key_help_url from YAML config into the MCPServer object. --- .../mcp_server/mcp_server_manager.py | 23 +++++++++++-------- .../mcp_server/rest_endpoints.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 8 +++++++ 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7c17da36bb7..4dd7a87dca1 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -315,6 +315,10 @@ class MCPServerManager: command=server_config.get("command", None) or "", args=server_config.get("args", None) or [], env=server_config.get("env", None) or {}, + # byok fields + is_byok=bool(server_config.get("is_byok", False)), + byok_description=server_config.get("byok_description", []) or [], + byok_api_key_help_url=server_config.get("byok_api_key_help_url", None), # oauth specific fields client_id=server_config.get("client_id", None), client_secret=server_config.get("client_secret", None), @@ -1008,22 +1012,21 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) - stdio_env = self._build_stdio_env(server, raw_headers) - - client = await self._create_mcp_client( - server=server, - mcp_auth_header=mcp_auth_header, - extra_headers=extra_headers, - stdio_env=stdio_env, - ) - - ## HANDLE OPENAPI TOOLS + ## HANDLE OPENAPI TOOLS — skip MCP client creation entirely if server.spec_path: _tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name) tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type( _tools ) else: + stdio_env = self._build_stdio_env(server, raw_headers) + + client = await self._create_mcp_client( + server=server, + mcp_auth_header=mcp_auth_header, + extra_headers=extra_headers, + stdio_env=stdio_env, + ) tools = await self._fetch_tools_with_timeout(client, server.name) prefixed_or_original_tools = self._create_prefixed_tools( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 16f8f835430..60277b91d5d 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -511,7 +511,7 @@ if MCP_AVAILABLE: arguments=tool_arguments, allowed_mcp_servers=allowed_mcp_servers, start_time=datetime.now(), - user_api_key_auth=data.get("user_api_key_auth"), + user_api_key_auth=data.get("user_api_key_auth") or user_api_key_dict, mcp_auth_header=data.get("mcp_auth_header"), mcp_server_auth_headers=data.get("mcp_server_auth_headers"), oauth2_headers=data.get("oauth2_headers"), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5c063839304..f728582d5c6 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1762,6 +1762,14 @@ if MCP_AVAILABLE: # These tools are registered with their prefixed names ######################################################### local_tool = global_mcp_tool_registry.get_tool(name) + # OpenAPI tools are always registered with their prefixed name. If the + # caller used the bare (unprefixed) name and we already resolved the + # server, construct the prefixed name and try again. + if local_tool is None and mcp_server is not None and mcp_server.spec_path: + prefixed_name = add_server_prefix_to_name(name, get_server_prefix(mcp_server)) + local_tool = global_mcp_tool_registry.get_tool(prefixed_name) + if local_tool: + name = prefixed_name if local_tool: verbose_logger.debug(f"Executing local registry tool: {name}") # For BYOK servers the credential must be injected via a ContextVar