diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index b1c11f73dc2..5952ce9df78 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -433,6 +433,10 @@ if MCP_AVAILABLE: if server_id not in allowed_server_ids: _server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + if _server is None: + # server_id may be a name string — look up by name for accurate + # IP-filter error reporting (get_mcp_server_by_id only matches UUIDs). + _server = global_mcp_server_manager.get_mcp_server_by_name(server_id) if ( _server is not None and _rest_client_ip is not None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bba8f0b78ff..c23362cf8f1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -8,7 +8,7 @@ import contextlib import time import traceback import uuid -from datetime import datetime +from datetime import datetime, timezone from typing import ( Any, AsyncIterator, @@ -871,11 +871,30 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _is_oauth_cred_expired(cred: Dict[str, Any]) -> bool: + """Return True if the OAuth2 credential's access_token has expired.""" + expires_at = cred.get("expires_at") + if not expires_at: + return False + try: + exp_dt = datetime.fromisoformat(expires_at) + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + return datetime.now(timezone.utc) > exp_dt + except (ValueError, TypeError): + return False + async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth], + prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: - """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict.""" + """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict. + + Args: + prefetched_creds: Optional dict keyed by server_id with credential payloads. + When provided, avoids a per-server DB round-trip. + """ if server.auth_type != MCPAuth.oauth2: return None if user_api_key_auth is None: @@ -884,21 +903,62 @@ if MCP_AVAILABLE: server_id = getattr(server, "server_id", None) if not user_id or not server_id: return None + try: + if prefetched_creds is not None: + cred = prefetched_creds.get(server_id) + else: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + ) + from litellm.proxy.utils import ( # noqa: PLC0415 + get_prisma_client_or_throw, + ) + + 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"): + if _is_oauth_cred_expired(cred): + verbose_logger.debug( + f"_get_user_oauth_extra_headers_from_db: token expired for " + f"user={user_id} server={server_id}" + ) + return None + return {"Authorization": f"Bearer {cred['access_token']}"} + except Exception as e: + verbose_logger.warning( + f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " + f"user={user_id} server={server_id}: {e}" + ) + return None + + async def _prefetch_oauth_creds_for_user( + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Dict[str, Dict[str, Any]]: + """Fetch all OAuth2 credentials for the user in one DB query. + + Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. + """ + user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + if not user_id: + return {} try: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 - get_user_oauth_credential, + list_user_oauth_credentials, ) 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 + creds = await list_user_oauth_credentials(prisma_client, user_id) + return {c["server_id"]: c for c in creds if "server_id" in c} + except Exception as e: + verbose_logger.warning( + f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}" + ) + return {} def _prepare_mcp_server_headers( server: MCPServer, @@ -1044,6 +1104,10 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) + # Pre-fetch all OAuth credentials for this user once to avoid N+1 DB queries + # inside the asyncio.gather loop below. + _prefetched_oauth_creds = await _prefetch_oauth_creds_for_user(user_api_key_auth) + async def _fetch_and_filter_server_tools( server: MCPServer, ) -> List[MCPTool]: @@ -1059,10 +1123,10 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) - # If no OAuth2 token came from request headers, fall back to DB lookup + # If no OAuth2 token came from request headers, fall back to pre-fetched creds 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 + server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds ) try: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 03555ed91ff..45157a47151 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -44,9 +44,11 @@ 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/([^/]+)$") +# Matches full proxy URLs of the form http(s):///mcp/ where /mcp/ is +# directly under the root path (no sub-path prefix). This ensures external MCP servers +# whose paths happen to contain "/mcp/" (e.g. https://mcp.atlassian.com/v1/mcp/...) are +# NOT rewritten as internal proxy routes. +_PROXY_MCP_PATH_RE = re.compile(r"^https?://[^/]+/mcp/([^/]+)$") class LiteLLM_Proxy_MCP_Handler: