From d0cc1a49f55d5acd2cf19a319db6b8e038a1e4bc Mon Sep 17 00:00:00 2001 From: gym-cmd <186399764+gym-cmd@users.noreply.github.com> Date: Fri, 15 May 2026 13:25:36 +0100 Subject: [PATCH] feat(mcp): support OAuth passthrough discovery --- litellm/experimental_mcp_client/client.py | 14 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 58 +- .../mcp_server/discoverable_endpoints.py | 195 ++++-- .../_experimental/mcp_server/exceptions.py | 27 + .../mcp_server/mcp_server_manager.py | 103 ++- .../_experimental/mcp_server/oauth_utils.py | 103 ++- .../mcp_server/rest_endpoints.py | 41 ++ .../proxy/_experimental/mcp_server/server.py | 154 ++++- litellm/proxy/proxy_server.py | 1 + .../types/mcp_server/mcp_server_manager.py | 24 + .../mcp_server/test_discoverable_endpoints.py | 156 ++++- .../mcp_server/test_discoverable_endpoints.py | 102 ++- .../mcp_server/test_mcp_oauth_passthrough.py | 603 ++++++++++++++++++ 13 files changed, 1473 insertions(+), 108 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/exceptions.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e703a3956b9..8e06ebf34fd 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -419,8 +419,16 @@ class MCPClient: return factory - async def list_tools(self) -> List[MCPTool]: - """List available tools from the server.""" + async def list_tools(self, raise_on_error: bool = False) -> List[MCPTool]: + """List available tools from the server. + + Args: + raise_on_error: When True, re-raise exceptions instead of returning + an empty list. Used by the proxy's pass-through MCP flow so it + can surface upstream HTTP 401 responses as a proper 401 to the + MCP client (triggering the upstream OAuth flow) rather than + masking them as "connected, no tools". + """ verbose_logger.debug( f"MCP client listing tools from {self.server_url or 'stdio'}" ) @@ -456,6 +464,8 @@ class MCPClient: "the MCP server may have crashed, disconnected, or timed out" ) + if raise_on_error: + raise # Return empty list instead of raising to allow graceful degradation return [] diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a05af66118c..49655bbc144 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,3 +1,4 @@ +import re from typing import Dict, List, Optional, Set, Tuple, cast from fastapi import HTTPException @@ -15,6 +16,38 @@ from litellm.proxy._types import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +def _parse_mcp_server_names_from_path(path: str) -> Optional[List[str]]: + """Parse a single MCP server name from /mcp/{name} or /{name}/mcp path patterns. + Returns None for the aggregate /mcp route (no bypass for multi-server paths).""" + m = re.match(r"^/mcp/([^/,?#]+)", path) + if m: + return [m.group(1)] + m = re.match(r"^/([^/,?#]+)/mcp", path) + if m: + return [m.group(1)] + return None + + +def _is_mcp_passthrough_cold_start( + scope: Scope, mcp_servers: Optional[List[str]] +) -> bool: + """True when the request targets a pass-through server with no auth headers — + the cold-start OAuth discovery case per RFC 9728 / MCP Authorization spec. + Lets the route handler's 401 emitter produce the spec-compliant WWW-Authenticate + challenge instead of surfacing a generic admission error.""" + if not mcp_servers: + return False + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + for name in mcp_servers: + server = global_mcp_server_manager.get_mcp_server_by_name(name) + if server is not None and getattr(server, "is_oauth_passthrough", False): + return True + return False + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -36,7 +69,7 @@ class MCPRequestHandler: LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value @staticmethod - async def process_mcp_request( + async def process_mcp_request( # noqa: PLR0915 scope: Scope, ) -> Tuple[ UserAPIKeyAuth, @@ -165,9 +198,26 @@ class MCPRequestHandler: else: raise else: - validated_user_api_key_auth = await user_api_key_auth( - api_key=litellm_api_key, request=request - ) + try: + validated_user_api_key_auth = await user_api_key_auth( + api_key=litellm_api_key, request=request + ) + except (HTTPException, ProxyException): + # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec + # require unauthenticated requests to protected resources to receive + # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers + # for pass-through servers instead of surfacing a generic admission error. + path = scope.get("path", "") + mcp_servers_from_path = _parse_mcp_server_names_from_path(path) + if _is_mcp_passthrough_cold_start( + scope, mcp_servers_from_path or mcp_servers + ): + verbose_logger.debug( + "MCP pass-through cold start: deferring admission to route 401 emitter" + ) + validated_user_api_key_auth = UserAPIKeyAuth() + else: + raise return ( validated_user_api_key_auth, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 1794cd14381..5eaf80641f3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,5 +1,6 @@ import json -from typing import Any, Dict, Optional +import time +from typing import Any, Dict, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request @@ -12,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, - validate_loopback_redirect_uri, + validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -24,6 +25,12 @@ from litellm.proxy.utils import get_server_root_path from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer +# TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. +# Keeps us from hammering the upstream IdP on each discovery request. +# Keyed by (server_id, resource_url) → (expires_at_epoch, payload). +_OAUTH_METADATA_CACHE: Dict[Tuple[str, str], Tuple[float, dict]] = {} +_OAUTH_METADATA_CACHE_TTL_SECONDS = 300 + router = APIRouter( tags=["mcp"], ) @@ -127,15 +134,6 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data -def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str: - """Return a loopback client redirect URI from OAuth state.""" - redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url") - if not redirect_uri or not isinstance(redirect_uri, str): - raise HTTPException(status_code=400, detail="Invalid redirect URI") - validate_loopback_redirect_uri(redirect_uri) - return redirect_uri - - def _append_query_params(url: str, params: Dict[str, str]) -> str: parsed = urlparse(url) query_params = parse_qsl(parsed.query, keep_blank_values=True) @@ -338,12 +336,12 @@ async def authorize_with_server( status_code=400, detail="MCP server authorization url is not set" ) - # Loopback-only redirect_uri. The URI is encrypted into the OAuth - # state and decoded on /callback to redirect the user back; a non- - # loopback URI would be an open-redirect + code-theft primitive - # (VERIA-57 root cause B). MCP clients are native apps — loopback is - # the spec-compliant callback pattern. - validate_loopback_redirect_uri(redirect_uri) + # Loopback OR same-origin redirect_uri. The URI is encrypted into the + # OAuth state and decoded on /callback to redirect the user back; + # restricting to trusted origins blocks the open-redirect + + # code-theft primitive (VERIA-57 root cause B). Loopback supports + # native MCP clients; same-origin supports the proxy's own UI callback. + validate_trusted_redirect_uri(request, redirect_uri) parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) @@ -660,20 +658,22 @@ async def token_endpoint( @router.get("/callback") -async def callback(code: str, state: str): +async def callback(request: Request, code: str, state: str): try: state_data = decode_state_hash(state) + base_url = state_data["base_url"] original_state = state_data["original_state"] - # Re-validate loopback at the sink. /authorize rejects non-loopback + # Re-validate at the sink. /authorize rejects untrusted # redirect_uri before encoding into state, but encrypted states # minted before that check was added have no expiry and remain - # valid indefinitely. Validating here blocks the open-redirect + - # code-theft primitive even for pre-fix states. - redirect_uri = _get_validated_client_redirect_uri(state_data) + # valid indefinitely. Validating here (same-origin OR loopback) + # blocks the open-redirect + code-theft primitive even for pre-fix + # states while allowing the UI's same-origin callback to work. + validate_trusted_redirect_uri(request, base_url) params = {"code": code, "state": original_state} - complete_returned_url = _append_query_params(redirect_uri, params) + complete_returned_url = f"{base_url}?{urlencode(params)}" return RedirectResponse(url=complete_returned_url, status_code=302) except HTTPException: @@ -708,7 +708,91 @@ async def callback(code: str, state: str): """ -def _build_oauth_protected_resource_response( +async def fetch_upstream_oauth_protected_resource( + mcp_server: MCPServer, +) -> Optional[dict]: + """Fetch the upstream MCP server's ``.well-known/oauth-protected-resource`` + metadata for a pass-through server. + + Tries host-only first, then falls back to the RFC 9728 §3.1 path-suffix + form (e.g. ``https://host/.well-known/oauth-protected-resource/mcp``) to + cover upstreams that scope metadata per resource path. + + Responses are cached in-process for ~5 minutes keyed on + ``(server_id, resource_url)`` so we do not hammer the IdP. + + Returns the parsed JSON dict on success, or ``None`` if neither form + responds with a 2xx JSON payload. Raises on network/connection errors so + the caller can emit HTTP 502 rather than fabricate a gateway response. + """ + if not mcp_server.url: + return None + + upstream = urlparse(mcp_server.url) + if not upstream.scheme or not upstream.netloc: + return None + + cache_key = (mcp_server.server_id, mcp_server.url) + cached = _OAUTH_METADATA_CACHE.get(cache_key) + if cached is not None and cached[0] > time.time(): + return cached[1] + + host_base = f"{upstream.scheme}://{upstream.netloc}" + candidates = [f"{host_base}/.well-known/oauth-protected-resource"] + # RFC 9728 §3.1 path fallback + if upstream.path and upstream.path not in ("", "/"): + candidates.append( + f"{host_base}/.well-known/oauth-protected-resource" + f"{upstream.path.rstrip('/')}" + ) + + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + + last_error: Optional[Exception] = None + for candidate in candidates: + try: + response = await async_client.get( + candidate, + headers={"Accept": "application/json"}, + ) + except Exception as exc: # network / connect errors + last_error = exc + continue + if response.status_code == 200: + try: + payload = response.json() + except Exception as exc: + last_error = exc + continue + if isinstance(payload, dict): + _OAUTH_METADATA_CACHE[cache_key] = ( + time.time() + _OAUTH_METADATA_CACHE_TTL_SECONDS, + payload, + ) + return payload + + if last_error is not None and all( + is_network_error(last_error) for _ in candidates + ): + raise last_error + + return None + + +def is_network_error(exc: Exception) -> bool: + """True for transport-layer failures (connection refused, DNS, TLS, timeout) + as opposed to HTTP protocol errors (4xx/5xx with a valid response).""" + name = type(exc).__name__ + return name in { + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "PoolTimeout", + "RemoteProtocolError", + } + + +async def _build_oauth_protected_resource_response( request: Request, mcp_server_name: Optional[str], use_standard_pattern: bool, @@ -716,6 +800,12 @@ def _build_oauth_protected_resource_response( """ Build OAuth protected resource response with the appropriate URL pattern. + For pass-through MCP servers (``MCPServer.is_oauth_passthrough``), the + gateway proxies the upstream's own ``oauth-protected-resource`` metadata + so that standards-compliant MCP clients discover the **upstream** IdP + instead of the gateway. The ``resource`` field is rewritten to the + gateway's own URL so clients present the bearer token back to the gateway. + Args: request: FastAPI Request object mcp_server_name: Name of the MCP server @@ -755,6 +845,30 @@ def _build_oauth_protected_resource_response( else: resource_url = f"{request_base_url}/mcp" + # Pass-through branch: proxy the upstream's own metadata so discovery + # directs the client at the real IdP (Okta, Keycloak, …) instead of us. + if mcp_server is not None and mcp_server.is_oauth_passthrough: + try: + upstream_metadata = await fetch_upstream_oauth_protected_resource( + mcp_server + ) + except Exception as exc: + verbose_logger.warning( + "Failed to fetch upstream oauth-protected-resource metadata " + f"for pass-through MCP server {mcp_server.name!r}: {exc}" + ) + raise HTTPException( + status_code=502, + detail=( + "Failed to fetch upstream oauth-protected-resource " + f"metadata for MCP server {mcp_server.name!r}" + ), + ) + + if upstream_metadata is not None: + response = {**upstream_metadata, "resource": resource_url} + return response + return { "authorization_servers": [ ( @@ -785,7 +899,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam This endpoint is compliant with MCP specification and works with standard MCP clients like mcp-inspector and VSCode Copilot. """ - return _build_oauth_protected_resource_response( + return await _build_oauth_protected_resource_response( request=request, mcp_server_name=mcp_server_name, use_standard_pattern=True, @@ -810,37 +924,18 @@ async def oauth_protected_resource_mcp( This endpoint is kept for backward compatibility. New integrations should use the standard MCP pattern (/mcp/{server_name}) instead. """ - return _build_oauth_protected_resource_response( + return await _build_oauth_protected_resource_response( request=request, mcp_server_name=mcp_server_name, use_standard_pattern=False, ) -""" - https://datatracker.ietf.org/doc/html/rfc8414#section-3.1 - RFC 8414: Path-aware OAuth discovery - If the issuer identifier value contains a path component, any - terminating "/" MUST be removed before inserting "/.well-known/" and - the well-known URI suffix between the host component and the path(include root path) - component. -""" - - -def _build_oauth_authorization_server_response( +async def _build_oauth_authorization_server_response( request: Request, mcp_server_name: Optional[str], ) -> dict: - """ - Build OAuth authorization server metadata response. - - Args: - request: FastAPI Request object - mcp_server_name: Name of the MCP server - - Returns: - OAuth authorization server metadata dict - """ + """Build OAuth authorization server metadata response (gateway-as-AS shape).""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -904,7 +999,7 @@ async def oauth_authorization_server_mcp_standard( Standard pattern: /mcp/{server_name} Discovery path: /.well-known/oauth-authorization-server/mcp/{server_name} """ - return _build_oauth_authorization_server_response( + return await _build_oauth_authorization_server_response( request=request, mcp_server_name=mcp_server_name, ) @@ -923,7 +1018,7 @@ async def oauth_authorization_server_mcp( Supports both legacy pattern (/{server_name}) and root endpoint. """ - return _build_oauth_authorization_server_response( + return await _build_oauth_authorization_server_response( request=request, mcp_server_name=mcp_server_name, ) @@ -994,7 +1089,7 @@ async def oauth_authorization_server_legacy(request: Request, mcp_server_name: s """ OAuth authorization server discovery for legacy /{server_name}/mcp pattern. """ - return _build_oauth_authorization_server_response( + return await _build_oauth_authorization_server_response( request=request, mcp_server_name=mcp_server_name, ) diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py new file mode 100644 index 00000000000..675f63113d6 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -0,0 +1,27 @@ +"""Exceptions raised by the LiteLLM MCP proxy.""" + +from typing import Optional + + +class MCPUpstreamAuthError(Exception): + """Raised when an upstream MCP server returns an authentication failure + (typically HTTP 401) and the gateway should surface it transparently to + the client instead of swallowing it. + + Only relevant for pass-through MCP servers (see + ``MCPServer.is_oauth_passthrough``). The gateway converts this exception + into an HTTP 401 response on single-server routes, preserving any + ``WWW-Authenticate`` challenge emitted by the upstream so standards- + compliant MCP clients can trigger the upstream OAuth flow. + """ + + def __init__( + self, + status_code: int, + www_authenticate: Optional[str], + server_name: str, + ) -> None: + self.status_code = status_code + self.www_authenticate = www_authenticate + self.server_name = server_name + super().__init__(f"Upstream MCP server {server_name!r} returned {status_code}") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 50dd7efda13..399912f8574 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -47,6 +47,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, @@ -117,6 +118,67 @@ _AZURE_ENTRA_HOSTS = { } +def _extract_upstream_auth_failure( + exc: BaseException, +) -> Optional[Tuple[int, Optional[str]]]: + """Walk the exception tree looking for an HTTP 401/403 response from the + upstream MCP server. + + The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and + may chain through ``__cause__`` / ``__context__``. We inspect all of those + layers for an ``httpx.Response``-bearing exception (typically + ``httpx.HTTPStatusError``) and extract the status code and any upstream + ``WWW-Authenticate`` header. + + Returns ``(status_code, www_authenticate)`` on match, else ``None``. + """ + seen: Set[int] = set() + stack: List[BaseException] = [exc] + while stack: + current = stack.pop() + if id(current) in seen: + continue + seen.add(id(current)) + + response = getattr(current, "response", None) + if response is not None: + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int) and status_code in (401, 403): + www_authenticate: Optional[str] = None + headers = getattr(response, "headers", None) + if headers is not None: + try: + www_authenticate = headers.get("www-authenticate") + except Exception: + www_authenticate = None + return status_code, www_authenticate + + if isinstance(current, HTTPStatusError): + status_code = getattr(current.response, "status_code", None) + if isinstance(status_code, int) and status_code in (401, 403): + www_authenticate = None + try: + www_authenticate = current.response.headers.get("www-authenticate") + except Exception: + pass + return status_code, www_authenticate + + # anyio / PEP 654 ExceptionGroup + sub_exceptions = getattr(current, "exceptions", None) + if sub_exceptions: + stack.extend(sub_exceptions) + + if current.__cause__ is not None: + stack.append(current.__cause__) + if ( + current.__context__ is not None + and current.__context__ is not current.__cause__ + ): + stack.append(current.__context__) + + return None + + def _warn_on_server_name_fields( *, server_id: str, @@ -1330,7 +1392,9 @@ class MCPServerManager: ] return tools else: - tools = await self._fetch_tools_with_timeout(client, server.name) + tools = await self._fetch_tools_with_timeout( + client, server.name, server=server + ) self._remember_upstream_initialize_instructions(server, client) prefixed_or_original_tools = self._create_prefixed_tools( @@ -1339,6 +1403,11 @@ class MCPServerManager: return prefixed_or_original_tools + except MCPUpstreamAuthError: + # Pass-through 401 must surface to single-server routes so the + # client triggers the upstream OAuth flow. The multi-server + # aggregator catches this explicitly to keep absorbing. + raise except Exception as e: verbose_logger.warning( f"Failed to get tools from server {server.name}: {str(e)}" @@ -1940,7 +2009,10 @@ class MCPServerManager: return None async def _fetch_tools_with_timeout( - self, client: MCPClient, server_name: str + self, + client: MCPClient, + server_name: str, + server: Optional[MCPServer] = None, ) -> List[MCPTool]: """ Fetch tools from MCP client with timeout and error handling. @@ -1948,16 +2020,28 @@ class MCPServerManager: Uses anyio.fail_after() instead of asyncio.wait_for() to avoid conflicts with the MCP SDK's anyio TaskGroup. See GitHub issue #20715 for details. + For pass-through MCP servers (``MCPServer.is_oauth_passthrough``) an + upstream HTTP 401 is converted into :class:`MCPUpstreamAuthError` + instead of being swallowed to an empty tool list. That lets the + single-server HTTP routes surface a proper 401 + ``WWW-Authenticate`` + challenge so standards-compliant MCP clients trigger the upstream + OAuth flow. Non-pass-through servers keep today's swallow-and-log + behaviour so the multi-server ``/mcp`` aggregator doesn't get + tainted by a single bad server. + Args: client: MCP client instance server_name: Name of the server for logging + server: Optional MCPServer; when pass-through, auth errors are + re-raised as :class:`MCPUpstreamAuthError`. Returns: List of tools from the server """ + is_passthrough = bool(server is not None and server.is_oauth_passthrough) try: with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): - tools = await client.list_tools() + tools = await client.list_tools(raise_on_error=is_passthrough) verbose_logger.debug(f"Tools from {server_name}: {tools}") return tools except TimeoutError: @@ -1974,6 +2058,19 @@ class MCPServerManager: ) return [] except Exception as e: + if is_passthrough: + auth_info = _extract_upstream_auth_failure(e) + if auth_info is not None: + status_code, www_authenticate = auth_info + verbose_logger.info( + f"Upstream auth failure from pass-through MCP server " + f"{server_name}: HTTP {status_code}" + ) + raise MCPUpstreamAuthError( + status_code=status_code, + www_authenticate=www_authenticate, + server_name=server_name, + ) from e verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}") return [] diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b13cf83058c..6ac3c8cf2a7 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -2,14 +2,60 @@ (BYOK + discoverable / pass-through OAuth proxy).""" from ipaddress import ip_address -from urllib.parse import urlparse +from urllib.parse import urlparse, urlunparse -from fastapi import HTTPException +from fastapi import HTTPException, Request + +from litellm._logging import verbose_logger +from litellm.proxy.auth.ip_address_utils import IPAddressUtils # RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses # must not be cached — both success and error bodies may reveal secrets. TOKEN_NO_CACHE_HEADERS = {"Cache-Control": "no-store", "Pragma": "no-cache"} +def get_request_base_url(request: Request) -> str: + """ + Get the base URL for the request, considering X-Forwarded-* headers. + + X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured + when the request comes from a configured trusted proxy + (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). + Otherwise the request's literal ``base_url`` is returned, so an + untrusted caller cannot poison OAuth-discovery / redirect_uri values + by injecting headers. + + Args: + request: FastAPI Request object + + Returns: + The reconstructed base URL (e.g., "https://proxy.example.com") + """ + base_url = str(request.base_url).rstrip("/") + parsed = urlparse(base_url) + + if not IPAddressUtils.is_request_from_trusted_proxy(request): + return base_url + + x_forwarded_proto = request.headers.get("X-Forwarded-Proto") + x_forwarded_host = request.headers.get("X-Forwarded-Host") + x_forwarded_port = request.headers.get("X-Forwarded-Port") + + scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme + + if x_forwarded_host: + if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): + netloc = x_forwarded_host + elif x_forwarded_port: + netloc = f"{x_forwarded_host}:{x_forwarded_port}" + else: + netloc = x_forwarded_host + else: + netloc = parsed.netloc + if x_forwarded_port and ":" not in netloc: + netloc = f"{netloc}:{x_forwarded_port}" + + return urlunparse((scheme, netloc, parsed.path, "", "", "")) + def validate_loopback_redirect_uri(redirect_uri: str) -> None: """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 @@ -46,3 +92,56 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: # don't let it bubble up as a 500. pass raise HTTPException(status_code=400, detail="invalid_request") + + +def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: + """Accept same-origin (proxy's own origin) OR loopback ``redirect_uri``. + + Same-origin is required for the LiteLLM UI's OAuth flow: the UI + redirects to ``/ui/mcp/oauth/callback`` which is not loopback + but is on the proxy's own trusted HTTPS origin. An attacker cannot + host content on the proxy's own origin without already owning the + proxy, so the open-redirect / code-theft primitive that motivated + :func:`validate_loopback_redirect_uri` does not apply here. + + Loopback continues to be accepted for native MCP clients (per + OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3). + + Use this in the discoverable OAuth proxy endpoints that serve both + native clients and the proxy's own UI. BYOK endpoints that only + support native clients should keep + :func:`validate_loopback_redirect_uri`. + """ + try: + parsed = urlparse(redirect_uri) + except ValueError: + raise HTTPException(status_code=400, detail="invalid_request") + if parsed.scheme not in ("http", "https"): + raise HTTPException(status_code=400, detail="invalid_request") + if parsed.fragment: + raise HTTPException(status_code=400, detail="invalid_request") + + try: + proxy_base = urlparse(get_request_base_url(request)) + if ( + parsed.netloc + and parsed.scheme == proxy_base.scheme + and parsed.netloc.lower() == proxy_base.netloc.lower() + ): + return + except Exception as exc: + verbose_logger.warning( + "validate_trusted_redirect_uri: could not determine proxy origin, " + "falling back to loopback-only check. error=%s", + exc, + ) + + host = (parsed.hostname or "").lower() + if host == "localhost": + return + try: + if ip_address(host).is_loopback: + return + except ValueError: + pass + raise HTTPException(status_code=400, detail="invalid_request") diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3047fb73325..bcee8192183 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -5,6 +5,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Set, from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.ui_session_utils import ( build_effective_auth_contexts, ) @@ -441,6 +442,18 @@ if MCP_AVAILABLE: user_api_key_dict, extra_headers=user_oauth_extra_headers, ) + except MCPUpstreamAuthError as e: + # Pass-through server returned 401 — surface it to the client so + # standards-compliant MCP clients trigger the upstream OAuth flow. + raise HTTPException( + status_code=e.status_code, + detail="Unauthorized", + headers=( + {"www-authenticate": e.www_authenticate} + if e.www_authenticate + else None + ), + ) except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") return { @@ -530,6 +543,18 @@ if MCP_AVAILABLE: user_api_key_dict, extra_headers=user_oauth_extra_headers, ) + except MCPUpstreamAuthError as e: + # Pass-through server returned 401 — surface it to the client so + # standards-compliant MCP clients trigger the upstream OAuth flow. + raise HTTPException( + status_code=e.status_code, + detail="Unauthorized", + headers=( + {"www-authenticate": e.www_authenticate} + if e.www_authenticate + else None + ), + ) except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") return { @@ -696,6 +721,22 @@ if MCP_AVAILABLE: ), } + except HTTPException as http_exc: + # Preserve 401s emitted by the single-server pass-through path so + # clients receive the upstream WWW-Authenticate challenge and can + # start the upstream OAuth flow. 403 etc. keep flowing through the + # legacy "error dict" response shape so the existing contract + # stays intact. + if http_exc.status_code == 401: + raise + verbose_logger.exception( + "HTTPException in list_tool_rest_api: %s", str(http_exc) + ) + return { + "tools": [], + "error": "unexpected_error", + "message": (f"An unexpected error occurred: {http_exc.detail}"), + } except Exception as e: verbose_logger.exception( "Unexpected error in list_tool_rest_api: %s", str(e) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 276a6e8a3bb..85fa920be4f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -926,6 +926,42 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _client_has_passthrough_authorization( + server: MCPServer, + oauth2_headers: Optional[Dict[str, str]], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + ) -> bool: + """True if the incoming request already carries an ``Authorization`` + header the gateway will forward to this pass-through server. + + The client may supply the bearer as either the top-level + ``Authorization`` header (surfaced via ``oauth2_headers``) or a + per-server ``x-mcp-auth-`` style header (surfaced via + ``mcp_server_auth_headers``). Either form skips the pre-emptive 401. + """ + if oauth2_headers: + for k in oauth2_headers.keys(): + if k.lower() == "authorization": + return True + if mcp_server_auth_headers: + for key in (server.alias, server.server_name, server.name): + if not key: + continue + server_headers = None + for k, v in mcp_server_auth_headers.items(): + if k.lower() == key.lower(): + server_headers = v + break + if server_headers is None: + continue + if isinstance(server_headers, str) and server_headers.strip(): + return True + if isinstance(server_headers, dict): + for hk in server_headers.keys(): + if hk.lower() == "authorization": + return True + return False + async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth], @@ -2754,6 +2790,84 @@ if MCP_AVAILABLE: ) return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + async def _raise_preemptive_401_for_unauthenticated_servers( + scope: Scope, + mcp_servers: Optional[List[str]], + oauth2_headers: Optional[Dict[str, str]], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + user_api_key_auth: Optional[UserAPIKeyAuth], + client_ip: Optional[str], + ) -> None: + """Fail fast with HTTP 401 for MCP servers that need user auth but + didn't receive it on this request. Covers both gateway-managed OAuth2 + (points clients at the gateway AS metadata) and pass-through OAuth + (points clients at the upstream resource-metadata via our well-known).""" + for server_name in mcp_servers or []: + server = global_mcp_server_manager.get_mcp_server_by_name( + server_name, client_ip=client_ip + ) + if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: + # For per-user OAuth servers, only skip the pre-emptive 401 when + # a stored token actually exists for this user+server pair. + # If no stored token exists, fail fast with 401 so clients can + # kick off PKCE/interactive OAuth flow immediately. + if server.needs_user_oauth_token: + stored_oauth_headers = await _get_user_oauth_extra_headers_from_db( + server=server, + user_api_key_auth=user_api_key_auth, + ) + if stored_oauth_headers: + continue + + request = StarletteRequest(scope) + base_url = get_request_base_url(request) + _path = scope.get("_original_path") or scope.get("path", "") or "" + + # Pick the well-known AS-metadata form that matches the inbound route + # so strict RFC 9728 §3.2 clients can resolve it correctly. + if _path.startswith(f"/mcp/{server_name}"): + _as_url = f"{base_url}/.well-known/oauth-authorization-server/mcp/{server_name}" + else: + _as_url = f"{base_url}/.well-known/oauth-authorization-server/{server_name}" + authorization_uri = f"Bearer authorization_uri={_as_url}" + + raise HTTPException( + status_code=401, + detail="Unauthorized", + headers={"www-authenticate": authorization_uri}, + ) + + # Pass-through OAuth: when the admin has opted a server into + # forwarding the client's bearer token (is_oauth_passthrough) and + # the client hasn't supplied one, fail fast with 401 and point + # them at the gateway's oauth-protected-resource well-known URL. + # That endpoint proxies the upstream's metadata so the client + # kicks off OAuth against the real upstream IdP, not the gateway. + if ( + server + and server.is_oauth_passthrough + and not _client_has_passthrough_authorization( + server, oauth2_headers, mcp_server_auth_headers + ) + ): + request = StarletteRequest(scope) + base_url = get_request_base_url(request) + _path = scope.get("_original_path") or scope.get("path", "") or "" + + # Pick the well-known resource-metadata form that matches the inbound + # route pattern so the metadata's `resource` field round-trips to what + # the client actually hit (RFC 9728 §3.2). + if _path.startswith(f"/{server_name}/mcp"): + resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" + else: + resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + www_authenticate = f'Bearer resource_metadata="{resource_metadata_url}"' + raise HTTPException( + status_code=401, + detail="Unauthorized", + headers={"www-authenticate": www_authenticate}, + ) + async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send ) -> None: @@ -2779,38 +2893,14 @@ if MCP_AVAILABLE: f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response - for server_name in mcp_servers or []: - server = global_mcp_server_manager.get_mcp_server_by_name( - server_name, client_ip=_client_ip - ) - if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: - # For per-user OAuth servers, only skip the pre-emptive 401 when - # a stored token actually exists for this user+server pair. - # If no stored token exists, fail fast with 401 so clients can - # kick off PKCE/interactive OAuth flow immediately. - if server.needs_user_oauth_token: - stored_oauth_headers = ( - await _get_user_oauth_extra_headers_from_db( - server=server, - user_api_key_auth=user_api_key_auth, - ) - ) - if stored_oauth_headers: - continue - - request = StarletteRequest(scope) - base_url = get_request_base_url(request) - - authorization_uri = ( - f"Bearer authorization_uri=" - f"{base_url}/.well-known/oauth-authorization-server/{server_name}" - ) - - raise HTTPException( - status_code=401, - detail="Unauthorized", - headers={"www-authenticate": authorization_uri}, - ) + await _raise_preemptive_401_for_unauthenticated_servers( + scope=scope, + mcp_servers=mcp_servers, + oauth2_headers=oauth2_headers, + mcp_server_auth_headers=mcp_server_auth_headers, + user_api_key_auth=user_api_key_auth, + client_ip=_client_ip, + ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. scope["headers"] = [ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3ed22e273b..002c1725d1a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15276,6 +15276,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): # Create a new scope with the correct path format that the MCP handler expects # Transform /{mcp_server_name}/mcp to /mcp/{mcp_server_name} scope = dict(request.scope) + scope["_original_path"] = scope.get("path", "") # preserve for 401 URL form selection scope["path"] = f"/mcp/{mcp_server_name}" # Import the MCP handler diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 8f8673b0a7d..46a788c0404 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -127,3 +127,27 @@ class MCPServer(BaseModel): return any(h.lower() in auth_header_names for h in self.extra_headers) return False + + @property + def is_oauth_passthrough(self) -> bool: + """True iff the gateway should transparently forward upstream OAuth + (discovery + 401s) rather than participating as an authorization + server itself. + + A server is pass-through for OAuth purposes when both conditions hold: + 1. ``auth_type`` is ``None`` or ``MCPAuth.none`` (the gateway does + not manage OAuth for this server). + 2. ``extra_headers`` includes ``Authorization`` — the admin has + opted this server into forwarding the client's bearer token + straight to the upstream MCP server. + + This is intentionally narrower than ``requires_per_user_auth``, + which also covers PATs (``x-api-key``, ``api-key``, ``apikey``). + Those are static credentials, not OAuth bearer tokens, so they + must not trigger upstream OAuth discovery or 401 propagation. + """ + if self.auth_type not in (None, MCPAuth.none): + return False + if not self.extra_headers: + return False + return any(h.lower() == "authorization" for h in self.extra_headers) diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 8785e450a4b..c0e4b4702fc 100644 --- a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1,8 +1,26 @@ """Tests for MCP OAuth discoverable endpoints""" import pytest +from fastapi import HTTPException from unittest.mock import AsyncMock, MagicMock, patch +TRUSTED_PROXY_IP = "10.0.0.5" +TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] + + +def set_request_from_trusted_proxy(mock_request): + mock_request.client = MagicMock() + mock_request.client.host = TRUSTED_PROXY_IP + + +@pytest.fixture +def trusted_proxy_origin_headers(): + with patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_mcp_trusted_proxy_ranges", + return_value=TRUSTED_PROXY_RANGES, + ): + yield + @pytest.mark.asyncio async def test_authorize_endpoint_includes_response_type(): @@ -56,7 +74,7 @@ async def test_authorize_endpoint_includes_response_type(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -244,10 +262,15 @@ async def test_register_client_without_mcp_server_name_returns_dummy(): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + global_mcp_server_manager.registry.clear() + mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} @@ -410,7 +433,9 @@ async def test_register_client_remote_registration_success(): @pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_proto(): +async def test_authorize_endpoint_respects_x_forwarded_proto( + trusted_proxy_origin_headers, +): """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -449,6 +474,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -461,7 +487,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -476,7 +502,9 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): @pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_proto(): +async def test_token_endpoint_respects_x_forwarded_proto( + trusted_proxy_origin_headers, +): """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -515,6 +543,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -666,7 +695,9 @@ async def test_oauth_protected_resource_legacy_pattern(): @pytest.mark.asyncio -async def test_oauth_protected_resource_respects_x_forwarded_proto(): +async def test_oauth_protected_resource_respects_x_forwarded_proto( + trusted_proxy_origin_headers, +): """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -704,6 +735,7 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_protected_resource_mcp( @@ -719,7 +751,9 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): @pytest.mark.asyncio -async def test_oauth_authorization_server_respects_x_forwarded_proto(): +async def test_oauth_authorization_server_respects_x_forwarded_proto( + trusted_proxy_origin_headers, +): """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -757,6 +791,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_authorization_server_mcp( @@ -773,20 +808,28 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): @pytest.mark.asyncio -async def test_register_client_respects_x_forwarded_proto(): +async def test_register_client_respects_x_forwarded_proto( + trusted_proxy_origin_headers, +): """Test that register_client uses X-Forwarded-Proto for redirect_uris""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + global_mcp_server_manager.registry.clear() + # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://proxy.litellm.example/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", @@ -803,7 +846,9 @@ async def test_register_client_respects_x_forwarded_proto(): @pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_host(): +async def test_authorize_endpoint_respects_x_forwarded_host( + trusted_proxy_origin_headers, +): """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -847,6 +892,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } + set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -859,7 +905,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -875,7 +921,9 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): @pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_host(): +async def test_token_endpoint_respects_x_forwarded_host( + trusted_proxy_origin_headers, +): """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -917,6 +965,7 @@ async def test_token_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } + set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -1075,7 +1124,12 @@ async def test_token_endpoint_respects_x_forwarded_host(): ], ) def test_get_request_base_url_comprehensive( - base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url + base_url, + x_forwarded_proto, + x_forwarded_host, + x_forwarded_port, + expected_url, + trusted_proxy_origin_headers, ): """Comprehensive test for get_request_base_url with various header combinations""" try: @@ -1089,6 +1143,7 @@ def test_get_request_base_url_comprehensive( # Create mock request mock_request = MagicMock(spec=Request) mock_request.base_url = base_url + set_request_from_trusted_proxy(mock_request) # Build headers dict headers = {} @@ -1116,3 +1171,82 @@ def test_get_request_base_url_comprehensive( f"X-Forwarded-Host={x_forwarded_host}, " f"X-Forwarded-Port={x_forwarded_port}" ) + + +def test_get_request_base_url_ignores_forwarded_headers_from_untrusted_client(): + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://gateway.example.com/mcp" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "attacker.example.com", + "X-Forwarded-Port": "443", + } + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.10" + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_mcp_trusted_proxy_ranges", + return_value=TRUSTED_PROXY_RANGES, + ): + assert get_request_base_url(mock_request) == "https://gateway.example.com/mcp" + + +def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): + try: + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://gateway.example.com/" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "attacker.example.com", + } + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.10" + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_mcp_trusted_proxy_ranges", + return_value=TRUSTED_PROXY_RANGES, + ), pytest.raises(HTTPException): + validate_trusted_redirect_uri( + mock_request, + "https://attacker.example.com/callback", + ) + + +def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy( + trusted_proxy_origin_headers, +): + try: + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + set_request_from_trusted_proxy(mock_request) + + validate_trusted_redirect_uri( + mock_request, + "https://proxy.example.com/callback", + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 85d5d6ba466..e4521ffb964 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -5,6 +5,23 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +TRUSTED_PROXY_IP = "10.0.0.5" +TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] + + +def set_request_from_trusted_proxy(mock_request): + mock_request.client = MagicMock() + mock_request.client.host = TRUSTED_PROXY_IP + + +@pytest.fixture +def trusted_proxy_origin_headers(): + with patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ): + yield + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @@ -561,6 +578,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -629,6 +647,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -707,6 +726,7 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_protected_resource_mcp( @@ -762,6 +782,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_authorization_server_mcp( @@ -800,6 +821,7 @@ async def test_register_client_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://proxy.litellm.example/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy + set_request_from_trusted_proxy(mock_request) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", @@ -862,6 +884,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } + set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -934,6 +957,7 @@ async def test_token_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } + set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -1091,7 +1115,11 @@ async def test_token_endpoint_respects_x_forwarded_host(): ], ) def test_get_request_base_url_comprehensive( - base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url + base_url, + x_forwarded_proto, + x_forwarded_host, + x_forwarded_port, + expected_url, ): """Comprehensive test for get_request_base_url with various header combinations. @@ -1110,6 +1138,7 @@ def test_get_request_base_url_comprehensive( mock_request = MagicMock(spec=Request) mock_request.base_url = base_url + set_request_from_trusted_proxy(mock_request) headers = {} if x_forwarded_proto: @@ -1264,6 +1293,71 @@ def test_xff_misconfig_warning_emitted_once(caplog): ), f"expected exactly one warning, got {len(matching)}: {[r.getMessage() for r in matching]}" +def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://gateway.example.com/" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "attacker.example.com", + } + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.10" + + with patch( + "litellm.proxy.proxy_server.general_settings", + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, + }, + create=True, + ), pytest.raises(HTTPException): + validate_trusted_redirect_uri( + mock_request, + "https://attacker.example.com/callback", + ) + + +def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy(): + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + set_request_from_trusted_proxy(mock_request) + + with patch( + "litellm.proxy.proxy_server.general_settings", + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, + }, + create=True, + ): + validate_trusted_redirect_uri( + mock_request, + "https://proxy.example.com/callback", + ) + + # ------------------------------------------------------------------- # Tests for scopes_supported when mcp_server.scopes is None # ------------------------------------------------------------------- @@ -1313,7 +1407,7 @@ async def test_oauth_protected_resource_returns_empty_scopes_when_none(): mock_request.headers = {} try: - response = _build_oauth_protected_resource_response( + response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name="atlassian_mcp", use_standard_pattern=False, @@ -1367,7 +1461,7 @@ async def test_oauth_authorization_server_returns_empty_scopes_when_none(): mock_request.headers = {} try: - response = _build_oauth_authorization_server_response( + response = await _build_oauth_authorization_server_response( request=mock_request, mcp_server_name="atlassian_mcp", ) @@ -1756,7 +1850,7 @@ async def test_discovery_root_includes_server_name_prefix(): try: # Call with mcp_server_name=None (root discovery) - response = _build_oauth_authorization_server_response( + response = await _build_oauth_authorization_server_response( request=mock_request, mcp_server_name=None, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py new file mode 100644 index 00000000000..32a5fb1f6bc --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py @@ -0,0 +1,603 @@ +"""Unit tests for the MCP OAuth pass-through patch (EAI-506 / Idea G). + +Covers: +- `MCPServer.is_oauth_passthrough` property semantics. +- `/.well-known/oauth-protected-resource/...` pass-through branch (proxies + upstream metadata, normalizes the `resource` field, caches, and surfaces + network errors as HTTP 502). +- `MCPServerManager._fetch_tools_with_timeout` converting upstream 401s into + `MCPUpstreamAuthError` for pass-through servers while keeping the silent + empty-list fallback for gateway-managed / aggregator paths. +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from fastapi import HTTPException, Request + +sys.path.insert(0, "../../../../../") + + +from litellm.proxy._experimental.mcp_server import discoverable_endpoints +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _OAUTH_METADATA_CACHE, + _build_oauth_protected_resource_response, +) +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + _extract_upstream_auth_failure, +) +from litellm.proxy._types import MCPTransport +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +@pytest.fixture(autouse=True) +def _mock_mcp_client_ip(): + """Bypass IP-based access control in tests.""" + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints" + ".IPAddressUtils.get_mcp_client_ip", + return_value=None, + ): + yield + + +@pytest.fixture(autouse=True) +def _clear_metadata_cache(): + """Prevent cross-test cache bleed for the oauth-protected-resource TTL cache.""" + _OAUTH_METADATA_CACHE.clear() + yield + _OAUTH_METADATA_CACHE.clear() + + +def _make_request(base_url: str = "https://gateway.example.com/") -> Request: + request = MagicMock(spec=Request) + request.base_url = base_url + request.headers = {} + return request + + +# -------------------------------------------------------------------------- +# is_oauth_passthrough property +# -------------------------------------------------------------------------- + + +def test_is_oauth_passthrough_true_when_none_auth_and_authorization_header(): + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + assert server.is_oauth_passthrough is True + + +def test_is_oauth_passthrough_true_when_auth_type_none_and_mixed_case_header(): + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + auth_type=None, + extra_headers=["authorization", "x-request-id"], + ) + assert server.is_oauth_passthrough is True + + +def test_is_oauth_passthrough_false_for_oauth2_server(): + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + extra_headers=["Authorization"], + ) + assert server.is_oauth_passthrough is False + + +def test_is_oauth_passthrough_false_without_authorization_header(): + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["x-api-key"], + ) + assert server.is_oauth_passthrough is False + + +def test_is_oauth_passthrough_false_without_extra_headers(): + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + assert server.is_oauth_passthrough is False + + +# -------------------------------------------------------------------------- +# _build_oauth_protected_resource_response: pass-through branch +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_passthrough_proxies_upstream_metadata(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + passthrough_server = MCPServer( + server_id="passthrough-1", + name="jet_knowledge_qa", + server_name="jet_knowledge_qa", + alias="jet_knowledge_qa", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + global_mcp_server_manager.registry[passthrough_server.server_id] = ( + passthrough_server + ) + + upstream_payload = { + "resource": "https://upstream.example.com/mcp", + "authorization_servers": ["https://okta.example.com/oauth2/default"], + "scopes_supported": ["openid", "profile"], + "bearer_methods_supported": ["header"], + } + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = upstream_payload + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=mock_client + ): + result = await _build_oauth_protected_resource_response( + request=_make_request(), + mcp_server_name="jet_knowledge_qa", + use_standard_pattern=True, + ) + + assert result["authorization_servers"] == [ + "https://okta.example.com/oauth2/default" + ] + # resource is normalized to the gateway URL so bearers are sent back to us + assert result["resource"].endswith("/mcp/jet_knowledge_qa") + assert result["scopes_supported"] == ["openid", "profile"] + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_passthrough_cache_hit(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + passthrough_server = MCPServer( + server_id="passthrough-2", + name="jet_knowledge_qa", + server_name="jet_knowledge_qa", + alias="jet_knowledge_qa", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + global_mcp_server_manager.registry[passthrough_server.server_id] = ( + passthrough_server + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "authorization_servers": ["https://okta.example.com"], + } + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=mock_client + ): + await _build_oauth_protected_resource_response( + request=_make_request(), + mcp_server_name="jet_knowledge_qa", + use_standard_pattern=True, + ) + await _build_oauth_protected_resource_response( + request=_make_request(), + mcp_server_name="jet_knowledge_qa", + use_standard_pattern=True, + ) + + assert mock_client.get.await_count == 1 + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_passthrough_network_error_returns_502(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + passthrough_server = MCPServer( + server_id="passthrough-3", + name="jet_knowledge_qa", + server_name="jet_knowledge_qa", + alias="jet_knowledge_qa", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + global_mcp_server_manager.registry[passthrough_server.server_id] = ( + passthrough_server + ) + + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + + with patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=mock_client + ): + with pytest.raises(HTTPException) as exc_info: + await _build_oauth_protected_resource_response( + request=_make_request(), + mcp_server_name="jet_knowledge_qa", + use_standard_pattern=True, + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_oauth_protected_resource_gateway_managed_unchanged(): + """Regression guard: OAuth2 servers still advertise the gateway as AS.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="oauth2-1", + name="keycloak_whoami", + server_name="keycloak_whoami", + alias="keycloak_whoami", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://keycloak/auth", + token_url="https://keycloak/token", + scopes=["read"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # If the code mistakenly fetched upstream metadata for a gateway-managed + # server, this spy would catch it. + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=mock_client + ): + result = await _build_oauth_protected_resource_response( + request=_make_request(), + mcp_server_name="keycloak_whoami", + use_standard_pattern=True, + ) + + mock_client.get.assert_not_awaited() + assert result["authorization_servers"] == [ + "https://gateway.example.com/keycloak_whoami" + ] + assert result["scopes_supported"] == ["read"] + + +# -------------------------------------------------------------------------- +# _extract_upstream_auth_failure helper +# -------------------------------------------------------------------------- + + +def test_extract_upstream_auth_failure_finds_401_in_http_status_error(): + response = httpx.Response( + status_code=401, + headers={"www-authenticate": 'Bearer resource_metadata="https://x"'}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + exc = httpx.HTTPStatusError("401", request=response.request, response=response) + + result = _extract_upstream_auth_failure(exc) + assert result == (401, 'Bearer resource_metadata="https://x"') + + +def test_extract_upstream_auth_failure_walks_exception_group(): + response = httpx.Response( + status_code=401, + headers={"www-authenticate": "Bearer"}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + inner = httpx.HTTPStatusError("401", request=response.request, response=response) + + try: + raise ExceptionGroup("wrapped", [inner]) # noqa: F821 (PEP 654, py3.11+) + except Exception as group: + result = _extract_upstream_auth_failure(group) + + assert result == (401, "Bearer") + + +def test_extract_upstream_auth_failure_returns_none_for_non_auth(): + assert _extract_upstream_auth_failure(RuntimeError("boom")) is None + + +# -------------------------------------------------------------------------- +# _fetch_tools_with_timeout pass-through behaviour +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_fetch_tools_from_passthrough_raises_on_upstream_401(): + manager = MCPServerManager() + passthrough_server = MCPServer( + server_id="p1", + name="jet_knowledge_qa", + url="https://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + + response = httpx.Response( + status_code=401, + headers={"www-authenticate": 'Bearer resource_metadata="https://upstream"'}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + upstream_error = httpx.HTTPStatusError( + "401", request=response.request, response=response + ) + + mock_client = MagicMock() + mock_client.list_tools = AsyncMock(side_effect=upstream_error) + + with pytest.raises(MCPUpstreamAuthError) as exc_info: + await manager._fetch_tools_with_timeout( + mock_client, passthrough_server.name, server=passthrough_server + ) + + assert exc_info.value.status_code == 401 + assert exc_info.value.www_authenticate == ( + 'Bearer resource_metadata="https://upstream"' + ) + assert exc_info.value.server_name == "jet_knowledge_qa" + mock_client.list_tools.assert_awaited_with(raise_on_error=True) + + +@pytest.mark.asyncio +async def test_fetch_tools_from_passthrough_returns_tools_on_success(): + manager = MCPServerManager() + passthrough_server = MCPServer( + server_id="p1", + name="jet_knowledge_qa", + url="https://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + + # list_tools returns a pre-baked tools list directly (MCPClient contract). + tool = MagicMock() + tool.name = "list_knowledge_bases" + mock_client = MagicMock() + mock_client.list_tools = AsyncMock(return_value=[tool]) + + tools = await manager._fetch_tools_with_timeout( + mock_client, passthrough_server.name, server=passthrough_server + ) + assert tools == [tool] + + +@pytest.mark.asyncio +async def test_fetch_tools_from_gateway_managed_swallows_errors(): + """Regression guard: non-pass-through servers keep returning [] on errors + so the multi-server aggregator isn't tainted by a single bad server.""" + manager = MCPServerManager() + oauth2_server = MCPServer( + server_id="o1", + name="keycloak_whoami", + url="https://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + + response = httpx.Response( + status_code=401, + headers={}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + upstream_error = httpx.HTTPStatusError( + "401", request=response.request, response=response + ) + mock_client = MagicMock() + mock_client.list_tools = AsyncMock(side_effect=upstream_error) + + tools = await manager._fetch_tools_with_timeout( + mock_client, oauth2_server.name, server=oauth2_server + ) + assert tools == [] + mock_client.list_tools.assert_awaited_with(raise_on_error=False) + + +# -------------------------------------------------------------------------- +# §2.1 — Admission cold-start: 401 + matching resource_metadata URL +# -------------------------------------------------------------------------- + + +def _make_scope(path: str, headers: list = None) -> dict: + """Build a minimal ASGI HTTP scope for testing.""" + raw_headers = [(k.encode(), v.encode()) for k, v in (headers or [])] + return { + "type": "http", + "method": "POST", + "path": path, + "headers": raw_headers, + "query_string": b"", + "server": ("localhost", 4000), + "scheme": "http", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route,expected_metadata_path", + [ + ( + "/mcp/jet_knowledge_qa", + "/.well-known/oauth-protected-resource/mcp/jet_knowledge_qa", + ), + ( + "/jet_knowledge_qa/mcp", + "/.well-known/oauth-protected-resource/jet_knowledge_qa/mcp", + ), + ], +) +async def test_passthrough_cold_start_emits_401_with_matching_resource_metadata( + route, expected_metadata_path +): + """No auth headers on a pass-through server route → 401 with resource_metadata URL + that matches the inbound path so RFC 9728 §3.2 strict clients accept it.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _is_mcp_passthrough_cold_start, + _parse_mcp_server_names_from_path, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + passthrough_server = MCPServer( + server_id="pt-cold-start", + name="jet_knowledge_qa", + server_name="jet_knowledge_qa", + alias="jet_knowledge_qa", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + global_mcp_server_manager.registry[passthrough_server.server_id] = ( + passthrough_server + ) + + # For /mcp/{name}: scope path stays as-is. + # For /{name}/mcp: dynamic_mcp_route rewrites to /mcp/{name} and sets _original_path. + if route.startswith("/mcp/"): + scope = _make_scope(route) + else: + scope = _make_scope(f"/mcp/jet_knowledge_qa") + scope["_original_path"] = route + + # Verify cold-start detection fires for this path + effective_path = scope.get("_original_path") or scope.get("path", "") + servers = _parse_mcp_server_names_from_path( + scope.get("path", "") # always /mcp/{name} by the time admission runs + ) + assert _is_mcp_passthrough_cold_start(scope, servers) is True + + # Verify resource_metadata_url form selection + server_name = "jet_knowledge_qa" + base_url = "http://localhost:4000" + path = scope.get("_original_path") or scope.get("path", "") or "" + if path.startswith(f"/{server_name}/mcp"): + resource_metadata_url = ( + f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" + ) + else: + resource_metadata_url = ( + f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + ) + + assert resource_metadata_url == f"{base_url}{expected_metadata_path}", ( + f"resource_metadata_url {resource_metadata_url!r} does not match " + f"expected {base_url + expected_metadata_path!r}" + ) + + +# -------------------------------------------------------------------------- +# §2.1 — Regression: non-pass-through servers bypass is NOT applied +# -------------------------------------------------------------------------- + + +def test_is_mcp_passthrough_cold_start_false_for_oauth2_server(): + """Gateway-managed OAuth2 servers must not trigger the cold-start bypass.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _is_mcp_passthrough_cold_start, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="oauth2-cold", + name="keycloak_whoami", + server_name="keycloak_whoami", + alias="keycloak_whoami", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://keycloak/auth", + token_url="https://keycloak/token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + scope = _make_scope("/mcp/keycloak_whoami") + result = _is_mcp_passthrough_cold_start(scope, ["keycloak_whoami"]) + assert result is False + + +def test_is_mcp_passthrough_cold_start_false_for_empty_servers(): + """Aggregate /mcp route (no server list) must not trigger bypass.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _is_mcp_passthrough_cold_start, + ) + + scope = _make_scope("/mcp") + assert _is_mcp_passthrough_cold_start(scope, None) is False + assert _is_mcp_passthrough_cold_start(scope, []) is False + + +# -------------------------------------------------------------------------- +# §2.1 — _parse_mcp_server_names_from_path +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "path,expected", + [ + ("/mcp/jet_knowledge_qa", ["jet_knowledge_qa"]), + ("/mcp/jet_knowledge_qa/tools/list", ["jet_knowledge_qa"]), + ("/jet_knowledge_qa/mcp", ["jet_knowledge_qa"]), + ("/jet_knowledge_qa/mcp/tools/list", ["jet_knowledge_qa"]), + ("/mcp", None), + ("/mcp/", None), + ("/other/path", None), + ], +) +def test_parse_mcp_server_names_from_path(path, expected): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _parse_mcp_server_names_from_path, + ) + + assert _parse_mcp_server_names_from_path(path) == expected