diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 99f6a5234a1..bba8f0b78ff 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -871,6 +871,35 @@ if MCP_AVAILABLE: return allowed_mcp_servers + async def _get_user_oauth_extra_headers_from_db( + server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Optional[Dict[str, str]]: + """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict.""" + if server.auth_type != MCPAuth.oauth2: + return None + if user_api_key_auth is None: + return None + user_id = getattr(user_api_key_auth, "user_id", None) + server_id = getattr(server, "server_id", None) + if not user_id or not server_id: + return None + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + cred = await get_user_oauth_credential(prisma_client, user_id, server_id) + if cred and cred.get("access_token"): + return {"Authorization": f"Bearer {cred['access_token']}"} + except Exception: + pass + return None + def _prepare_mcp_server_headers( server: MCPServer, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], @@ -1030,6 +1059,12 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) + # If no OAuth2 token came from request headers, fall back to DB lookup + if extra_headers is None and server.auth_type == MCPAuth.oauth2: + extra_headers = await _get_user_oauth_extra_headers_from_db( + server, user_api_key_auth + ) + try: tools = await global_mcp_server_manager._get_tools_from_server( server=server, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 5776ef95acb..03555ed91ff 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,3 +1,4 @@ +import re import traceback from datetime import datetime from typing import ( @@ -43,6 +44,10 @@ ToolParam = Any LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" +# Matches any URL ending in /mcp/ — e.g. http://localhost:4000/mcp/atlassian_test +# Used to auto-route requests targeting the proxy's own MCP endpoint through the internal handler. +_PROXY_MCP_PATH_RE = re.compile(r"^.+/mcp/([^/]+)$") + class LiteLLM_Proxy_MCP_Handler: """ @@ -54,7 +59,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Optional[Iterable[ToolParam]]) -> bool: """ - Returns True if the user passed a MCP tool with server_url="litellm_proxy" + Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. + This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. """ if tools: for tool in tools: @@ -64,6 +70,10 @@ class LiteLLM_Proxy_MCP_Handler: LITELLM_PROXY_MCP_SERVER_URL ): return True + if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match( + server_url + ): + return True return False @staticmethod @@ -87,6 +97,18 @@ class LiteLLM_Proxy_MCP_Handler: LITELLM_PROXY_MCP_SERVER_URL ): mcp_tools_with_litellm_proxy.append(tool) + elif isinstance(server_url, str): + # Also intercept URLs like http://localhost:4000/mcp/atlassian_test + # by rewriting them to the internal litellm_proxy format. + m = _PROXY_MCP_PATH_RE.match(server_url) + if m: + rewritten = dict(tool) + rewritten["server_url"] = ( + f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}" + ) + mcp_tools_with_litellm_proxy.append(rewritten) + else: + other_tools.append(tool) else: other_tools.append(tool) else: