diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 831e588e5ba..c1c90233bee 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -520,13 +520,28 @@ async def _list_tools_operation(session: ClientSession): # Return empty list instead of raising to allow graceful degradation return [] + @staticmethod + def error_tool_result(exc: Exception) -> MCPCallToolResult: + """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" + return MCPCallToolResult( + content=[TextContent(type="text", text=f"{type(exc).__name__}: {str(exc)}")], + isError=True, + ) + async def call_tool( self, call_tool_request_params: MCPCallToolRequestParams, host_progress_callback: Optional[Callable] = None, + raise_on_error: bool = False, ) -> MCPCallToolResult: """ Call an MCP Tool. + + Args: + raise_on_error: When True, re-raise the underlying exception instead of returning an + ``isError=True`` result. The token-exchange (OBO) tool-call path uses this to detect + an upstream 401 so it can re-mint the exchanged token and retry once; every other + caller keeps the default and gets graceful ``isError`` degradation. """ verbose_logger.info(f"MCP client calling tool '{call_tool_request_params.name}'") @@ -579,11 +594,10 @@ async def _call_tool_operation(session: ClientSession): "MCP client detected broken connection/stream - " "the MCP server may have crashed, disconnected, or timed out." ) + if raise_on_error: + raise # Return a default error result instead of raising - return MCPCallToolResult( - content=[TextContent(type="text", text=f"{error_type}: {str(e)}")], # Empty content for error case - isError=True, - ) + return self.error_tool_result(e) async def list_prompts(self) -> List[Prompt]: """List available prompts from the server.""" diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index c9d1d09b4f0..200b148922f 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1184,6 +1184,10 @@ async def _build_oauth_protected_resource_response( detail=(f"Upstream oauth-protected-resource metadata unavailable for MCP server {mcp_server.name!r}"), ) + obo_response = _obo_protected_resource_response(mcp_server, resource_url) + if obo_response is not None: + return obo_response + return { "authorization_servers": [ (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") @@ -1193,6 +1197,51 @@ async def _build_oauth_protected_resource_response( } +def _obo_protected_resource_response(mcp_server: Optional[MCPServer], resource_url: str) -> Optional[dict]: + """The OBO (token_exchange) PRM, or None when this server is not OBO / no issuer is configured. + + The client SSOs with the IdP to obtain a subject token, which LiteLLM then exchanges, so discovery + points at the JWT-auth issuer(s) LiteLLM trusts (the same IdP that issues and validates the + subject), not the gateway. None falls the caller back to the gateway default so discovery still + returns metadata; it just can't name the IdP. + """ + if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: + return None + issuers = _jwt_auth_issuers() + if not issuers: + return None + return { + "authorization_servers": issuers, + "resource": resource_url, + "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), + } + + +def _jwt_auth_issuers() -> list: + """The OAuth issuer identifier(s) LiteLLM's JWT auth trusts, for the OBO PRM authorization_servers. + + In token_exchange the IdP that issues the subject JWT is the same one LiteLLM validates it + against, so OBO discovery points clients at the JWT-auth issuer to obtain a subject token. + Sourced from ``JWT_ISSUER`` and any configured ``litellm_jwtauth.issuers``. + """ + import os # noqa: PLC0415 + + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 + + issuers: list = [] + env_issuer = os.getenv("JWT_ISSUER") + if env_issuer: + issuers.append(env_issuer) + + jwtauth = general_settings.get("litellm_jwtauth") if isinstance(general_settings, dict) else None + raw_issuers = jwtauth.get("issuers") if isinstance(jwtauth, dict) else getattr(jwtauth, "issuers", None) + for cfg in raw_issuers or []: + issuer = cfg.get("issuer") if isinstance(cfg, dict) else getattr(cfg, "issuer", None) + if issuer and issuer not in issuers: + issuers.append(issuer) + return issuers + + # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) @router.get( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b5832533b17..800cabbca53 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -17,6 +17,7 @@ from urllib.parse import urlparse import anyio +import httpx from fastapi import HTTPException from httpx import HTTPStatusError from mcp import ReadResourceResult, Resource @@ -63,6 +64,7 @@ ) from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_public, + raise_token_exchange_challenge, raise_user_oauth_challenge, to_server_spec, to_subject, @@ -70,8 +72,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( + build_token_exchanger, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, + ServerSpec, + TokenExchangeConfig, ) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, @@ -104,7 +111,7 @@ from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, get_server_root_path from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPStdioConfig @@ -206,6 +213,10 @@ def _should_strip_caller_authorization( ``Authorization`` is the upstream OAuth token and must be forwarded, so we keep it. """ + if mcp_server.auth_type == MCPAuth.oauth2_token_exchange: + # OBO: the inbound Authorization is the subject token. It is exchanged at the IdP and only the + # exchanged token is sent upstream, so the raw caller bearer must never be forwarded. + return True if mcp_server.has_client_credentials: return True if mcp_server.auth_type == MCPAuth.oauth2 and to_server_spec(mcp_server) is not None: @@ -528,9 +539,22 @@ def _resolve_oauth2_flow( return "client_credentials" return None + @staticmethod + def _obo_needs_endpoint_discovery( + auth_type: Optional[MCPAuthType], + token_exchange_endpoint: Optional[str], + token_url: Optional[str], + ) -> bool: + """An ``oauth2_token_exchange`` server with no configured token endpoint can have it + discovered (RFC 9728 -> RFC 8414) the same way the ``oauth2`` flow already does; an explicitly + configured ``token_exchange_endpoint``/``token_url`` wins and skips the discovery round-trip. + """ + return auth_type == MCPAuth.oauth2_token_exchange and not (token_exchange_endpoint or token_url) + def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None): self._cred_provider = cred_provider or UpstreamCredentialProvider( - oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id) + oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id), + token_exchanger=build_token_exchanger(), ) self.registry: Dict[str, MCPServer] = {} self.config_mcp_servers: Dict[str, MCPServer] = {} @@ -711,9 +735,17 @@ async def load_servers_from_config( ) auth_type = server_config.get("auth_type", None) - if server_url and auth_type is not None and auth_type == MCPAuth.oauth2: + if server_url and ( + auth_type == MCPAuth.oauth2 + or self._obo_needs_endpoint_discovery( + auth_type, + server_config.get("token_exchange_endpoint"), + server_config.get("token_url"), + ) + ): mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, + allow_origin_fallback=auth_type == MCPAuth.oauth2, ) else: mcp_oauth_metadata = None @@ -1090,9 +1122,19 @@ async def build_mcp_server_from_table( auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url - needs_discovery = bool(server_url) and auth_type == MCPAuth.oauth2 and not mcp_server.authorization_url + needs_discovery = bool(server_url) and ( + (auth_type == MCPAuth.oauth2 and not mcp_server.authorization_url) + or self._obo_needs_endpoint_discovery( + auth_type, + credentials_dict.get("token_exchange_endpoint") if credentials_dict else None, + mcp_server.token_url, + ) + ) mcp_oauth_metadata = ( - await self._descovery_metadata(server_url=server_url) # type: ignore[arg-type] + await self._descovery_metadata( + server_url=server_url, # type: ignore[arg-type] + allow_origin_fallback=auth_type == MCPAuth.oauth2, + ) if needs_discovery else None ) @@ -1166,8 +1208,49 @@ async def build_mcp_server_from_table( timeout=getattr(mcp_server, "timeout", None), ) _warn_internal_delegate_pkce_if_applicable(new_server, source="database") + await self._persist_discovered_obo_token_url( + server_id=mcp_server.server_id, + auth_type=auth_type, + existing_token_url=mcp_server.token_url, + discovered_token_url=new_server.token_url, + ) return new_server + async def _persist_discovered_obo_token_url( + self, + *, + server_id: str, + auth_type: Optional[MCPAuthType], + existing_token_url: Optional[str], + discovered_token_url: Optional[str], + ) -> None: + """Write a freshly discovered OBO token endpoint back onto the DB row. + + ``build_mcp_server_from_table`` resolves ``token_url`` via RFC 9728 -> RFC 8414 for an + ``oauth2_token_exchange`` server that has none configured, but that resolved value otherwise + lives only on the returned in-memory object; the row keeps ``token_url=None`` so every rebuild + re-runs discovery, and a transient upstream outage during a rebuild leaves the server with no + endpoint until discovery next succeeds. Persisting it makes ``_obo_needs_endpoint_discovery`` + return False on the next build. Fires at most once per server (skipped once the row has a + value), and is best-effort: a write failure just means discovery runs again next time. + """ + if auth_type != MCPAuth.oauth2_token_exchange: + return + if existing_token_url or not discovered_token_url: + return + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 + + if prisma_client is None: + return + try: + await MCPServerRepository(prisma_client).table.update( + where={"server_id": server_id}, + data={"token_url": discovered_token_url}, + ) + verbose_logger.debug("Persisted discovered OBO token_url for MCP server %s", server_id) + except Exception as exc: # noqa: BLE001 - best-effort; a failed write re-discovers next build + verbose_logger.warning("Failed to persist discovered OBO token_url for MCP server %s: %s", server_id, exc) + async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): """Register OpenAPI tools if the server has a spec_path configured.""" if server.spec_path: @@ -1593,6 +1676,21 @@ def _extract_bearer_token( return auth_value return None + def _obo_subject_token( + self, + server: MCPServer, + raw_headers: Optional[dict[str, str]], + ) -> Optional[str]: + """The caller's bearer as the token_exchange (OBO) subject token, for that mode only. + + Prompts/resources discovery and reads on a token_exchange server must exchange the caller's + token like the tools paths do, not connect with no credential. Other modes never read the + inbound bearer, so return None to avoid forwarding it. + """ + if server.auth_type != MCPAuth.oauth2_token_exchange: + return None + return self._extract_bearer_token(None, raw_headers) + def _build_stdio_env( self, server: MCPServer, @@ -1783,6 +1881,84 @@ async def _load_user_env_vars( _write_user_env_vars_cache(user_id, server.server_id, values) return values + async def _resolve_v2_auth( + self, + *, + server: MCPServer, + spec: ServerSpec, + provider: UpstreamCredentialProvider, + subject_token: Optional[str], + user_api_key_auth: Optional[UserAPIKeyAuth], + extra_headers: Optional[dict[str, str]], + ) -> tuple[Optional[httpx.Auth], Optional[dict[str, str]]]: + """Resolve a v2-owned server's upstream credential into ``(resolved_auth, extra_headers)``. + + On a missing/rejected per-user credential this raises the mode's discovery challenge + (authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any + other ``CredError`` onto its public HTTP status; it never returns an error as a value. + """ + match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): + case Ok(auth): + # NoOpAuth has no header_name and so never conflicts. + header_name = getattr(auth, "header_name", None) + conflicts = bool( + header_name and extra_headers and any(key.lower() == header_name.lower() for key in extra_headers) + ) + if not conflicts: + return auth, extra_headers + if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig)): + # The resolver owns the per-user credential here (token_exchange's exchanged + # token, authorization_code's stored token). It is authoritative: a guardrail such + # as MCPJWTSigner, static_headers, or any other injected Authorization must NOT + # shadow it (otherwise the upstream gets e.g. the signer's JWT instead of the + # exchanged token and rejects it). Drop the conflicting header so the resolved + # token reaches upstream. + return auth, _without_authorization(extra_headers) + # Other modes: an Authorization already supplied via extra_headers (a forwarded caller + # header or static_headers) is intentional and wins; v1 applies those last. + return None, extra_headers + case Error(err): + if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): + # authorization_code's missing per-user token -> the per-server browser-OAuth + # challenge, built here where the full MCPServer is in hand. + raise_user_oauth_challenge(server, root_path=get_server_root_path()) + if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): + # token_exchange (OBO): a missing/rejected subject token -> the RFC 9728 challenge + # pointing at the IdP the client must SSO with to obtain one, rather than an opaque + # 401. No gateway-side browser flow. + raise_token_exchange_challenge(server, root_path=get_server_root_path()) + raise_public(err) + + async def preflight_token_exchange( + self, + server: MCPServer, + oauth2_headers: Optional[dict[str, str]], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> None: + """Run the OBO exchange for a caller-supplied subject at the transport edge. + + Single-server routes call this before the MCP session opens, where an HTTP status and + ``WWW-Authenticate`` still reach the client. A rejected subject raises the RFC 9728 + challenge and any other ``CredError`` maps onto its public HTTP status, so an exchange + failure surfaces as a failure instead of the session continuing into an empty tool list. + A successful exchange is cached by the exchanger, so the session's list/call reuses it. + """ + if server.auth_type != MCPAuth.oauth2_token_exchange: + return + subject_token = self._extract_bearer_token(oauth2_headers, None) + if not subject_token: + return + spec = to_server_spec(server) + if spec is None or not isinstance(spec.config, TokenExchangeConfig): + return + match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): + case Ok(_): + return + case Error(err): + if err.tag == "unauthorized": + raise_token_exchange_challenge(server, root_path=get_server_root_path()) + raise_public(err) + async def _create_mcp_client( self, server: MCPServer, @@ -1817,11 +1993,17 @@ async def _create_mcp_client( spec = None if transport == MCPTransport.stdio else to_server_spec(server) provider = cred_provider or self._cred_provider # A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path - # so it wins - except for authorization_code, whose per-user token the v2 resolver owns. A - # caller must not be able to substitute another user's stored credential, so we keep the v2 - # spec and ignore the override there; the REST tools preview supplies its not-yet-persisted - # token through the resolver (cred_provider), never this path. - if spec is not None and mcp_auth_header and not isinstance(spec.config, AuthorizationCodeConfig): + # so it wins - except for the per-user modes the v2 resolver owns (authorization_code's + # stored token and token_exchange's RFC 8693 minted token). A caller must not be able to + # substitute another user's stored credential, nor silently disable the OBO exchange and + # forward an arbitrary bearer upstream, so we keep the v2 spec and ignore the override for + # both; the REST tools preview supplies its not-yet-persisted token through the resolver + # (cred_provider), never this path. + if ( + spec is not None + and mcp_auth_header + and not isinstance(spec.config, (AuthorizationCodeConfig, TokenExchangeConfig)) + ): spec = None auth_value = ( await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) if spec is None else None @@ -1887,26 +2069,14 @@ async def _create_mcp_client( server_url = server.url or "" if spec is not None: - match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): - case Ok(auth): - resolved_auth = auth - # Do not override an Authorization already supplied via extra_headers - # (a guardrail hook such as the JWT signer, static_headers, or a - # forwarded caller header): v1 applies those last, so they win. NoOpAuth - # has no header_name and so never skips. - header_name = getattr(resolved_auth, "header_name", None) - if ( - header_name - and extra_headers - and any(key.lower() == header_name.lower() for key in extra_headers) - ): - resolved_auth = None - case Error(err): - if err.tag == "unauthorized": - # The arm signals a missing per-user token semantically; raise the - # per-server OAuth challenge here, where the full MCPServer is in hand. - raise_user_oauth_challenge(server) - raise_public(err) + resolved_auth, extra_headers = await self._resolve_v2_auth( + server=server, + spec=spec, + provider=provider, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + extra_headers=extra_headers, + ) return MCPClient( server_url=server_url, transport_type=transport, @@ -1951,6 +2121,7 @@ async def _get_tools_from_server( add_prefix: bool = True, raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, + oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -2024,11 +2195,21 @@ async def _get_tools_from_server( stdio_env = self._build_stdio_env(server, raw_headers) + # token_exchange (OBO) discovery needs the caller's token too: list it with the user's own + # token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes + # never read the inbound bearer, so leave subject_token None to avoid forwarding it. + subject_token = ( + self._extract_bearer_token(oauth2_headers, raw_headers) + if server.auth_type == MCPAuth.oauth2_token_exchange + else None + ) + client = await self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, stdio_env=stdio_env, + subject_token=subject_token, user_api_key_auth=user_api_key_auth, ) @@ -2068,12 +2249,16 @@ async def _get_tools_from_server( # aggregator catches this explicitly to keep absorbing. raise except HTTPException as e: - headers = e.headers or {} - www_authenticate = headers.get("WWW-Authenticate") or headers.get("www-authenticate") - if e.status_code == 401 and www_authenticate is not None: + # A v2 resolver auth challenge (token_exchange's RFC 9728 401, authorization_code's + # browser-OAuth 401, or a 403) is raised at client-build time, inside this try. Route it + # through the same MCPUpstreamAuthError channel as pass-through so single-server routes + # surface the challenge (the client re-authenticates) while the aggregator keeps absorbing. + # Non-auth HTTP errors stay absorbed so one misconfigured server can't blank the listing. + if e.status_code in (401, 403): + headers = e.headers or {} raise MCPUpstreamAuthError( - status_code=401, - www_authenticate=www_authenticate, + status_code=e.status_code, + www_authenticate=headers.get("WWW-Authenticate") or headers.get("www-authenticate"), server_name=server.name, ) from e verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}") @@ -2113,12 +2298,14 @@ async def get_prompts_from_server( extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) + subject_token = self._obo_subject_token(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, + subject_token=subject_token, ) prompts = await client.list_prompts() @@ -2153,12 +2340,14 @@ async def get_resources_from_server( extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) + subject_token = self._obo_subject_token(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, + subject_token=subject_token, ) resources = await client.list_resources() @@ -2193,12 +2382,14 @@ async def get_resource_templates_from_server( extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) + subject_token = self._obo_subject_token(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, + subject_token=subject_token, ) resource_templates = await client.list_resource_templates() @@ -2232,12 +2423,14 @@ async def read_resource_from_server( extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) + subject_token = self._obo_subject_token(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, + subject_token=subject_token, ) return await client.read_resource(url) @@ -2262,12 +2455,14 @@ async def get_prompt_from_server( extra_headers.update(server.static_headers) stdio_env = self._build_stdio_env(server, raw_headers) + subject_token = self._obo_subject_token(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, + subject_token=subject_token, ) get_prompt_request_params = GetPromptRequestParams( @@ -2320,8 +2515,17 @@ async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> Any: async def _descovery_metadata( self, server_url: str, + *, + allow_origin_fallback: bool = True, ) -> Optional[MCPOAuthMetadata]: - """Discover OAuth metadata by following RFC 9728 (protected resource metadata discovery).""" + """Discover OAuth metadata by following RFC 9728 (protected resource metadata discovery). + + ``allow_origin_fallback`` controls the last-resort guess that treats the resource server's own + origin as its authorization server when nothing is advertised. The browser ``oauth2`` flow keeps + it (a human sees the redirect), but token_exchange (OBO) sets it False so the gateway never + exchanges a subject token against an endpoint it inferred rather than one explicitly configured + or authoritatively advertised via RFC 9728 / RFC 8414. + """ try: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -2371,7 +2575,7 @@ async def _descovery_metadata( ) = await self._attempt_well_known_discovery(server_url) metadata = None - if not authorization_servers: + if allow_origin_fallback and not authorization_servers: try: parsed_url = urlparse(server_url) if parsed_url.scheme and parsed_url.netloc: @@ -2533,6 +2737,14 @@ async def _fetch_single_authorization_server_metadata( continue scopes = self._extract_scopes(data.get("scopes_supported")) + verbose_logger.debug( + "Authorization server metadata from %s: issuer=%s grant_types_supported=%s " + "token_endpoint_auth_methods_supported=%s", + url, + data.get("issuer"), + data.get("grant_types_supported"), + data.get("token_endpoint_auth_methods_supported"), + ) metadata = MCPOAuthMetadata( scopes=scopes, authorization_url=data.get("authorization_endpoint"), @@ -3142,6 +3354,46 @@ def _create_during_hook_task( ) ) + async def _obo_call_tool_with_retry( + self, + *, + client: MCPClient, + call_tool_params: MCPCallToolRequestParams, + host_progress_callback: Optional[Callable], + mcp_server: MCPServer, + server_auth_header: str | dict[str, str] | None, + extra_headers: Optional[dict[str, str]], + stdio_env: Optional[dict[str, str]], + subject_token: Optional[str], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> CallToolResult: + """Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry. + + The exchanged token is baked into the client at build time, so the retry invalidates the + cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth + failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange + that now fails surfaces its own 401 challenge from ``_create_mcp_client``. + """ + try: + return await client.call_tool( + call_tool_params, host_progress_callback=host_progress_callback, raise_on_error=True + ) + except Exception as exc: + if _extract_upstream_auth_failure(exc) is None: + return MCPClient.error_tool_result(exc) + spec = to_server_spec(mcp_server) + if spec is not None: + await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec) + retry_client = await self._create_mcp_client( + server=mcp_server, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + stdio_env=stdio_env, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + ) + return await retry_client.call_tool(call_tool_params, host_progress_callback=host_progress_callback) + async def _call_regular_mcp_tool( self, mcp_server: MCPServer, @@ -3300,10 +3552,29 @@ async def _call_regular_mcp_tool( arguments=arguments, ) - async def _call_tool_via_client(client, params): - return await client.call_tool(params, host_progress_callback=host_progress_callback) + if mcp_server.auth_type == MCPAuth.oauth2_token_exchange and subject_token: + # OBO: the exchanged token may have been revoked/rotated upstream since it was cached, so + # an upstream 401 gets one re-mint + retry. Gated to this mode; all others keep the plain + # single call below. + tool_call_coro = self._obo_call_tool_with_retry( + client=client, + call_tool_params=call_tool_params, + host_progress_callback=host_progress_callback, + mcp_server=mcp_server, + server_auth_header=server_auth_header, + extra_headers=extra_headers, + stdio_env=stdio_env, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + ) + else: + + async def _call_tool_via_client(client, params): + return await client.call_tool(params, host_progress_callback=host_progress_callback) + + tool_call_coro = _call_tool_via_client(client, call_tool_params) - tasks.append(asyncio.create_task(_call_tool_via_client(client, call_tool_params))) + tasks.append(asyncio.create_task(tool_call_coro)) _timeout = mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT try: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 815fc2ba29d..3ba8baab427 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -26,6 +26,7 @@ ServerSpec, SharedKey, Subject, + TokenExchangeConfig, ) from litellm.types.mcp import MCPAuth @@ -61,8 +62,9 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes, - all shared-key), and ``oauth2`` per-user tokens (``authorization_code``); client_credentials - (M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1. + all shared-key), ``oauth2`` per-user tokens (``authorization_code``), and + ``oauth2_token_exchange`` (RFC 8693 OBO); client_credentials (M2M), delegated/passthrough + oauth2, and SigV4 return None and stay on v1. """ if server.is_byok: return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) @@ -92,11 +94,41 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: ) # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 return None - case MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4: - return None # token exchange and SigV4 are not migrated yet -> defer to v1 + case MCPAuth.oauth2_token_exchange: + return _token_exchange_spec(server, resource) + case MCPAuth.aws_sigv4: + return None # SigV4 is not migrated yet -> defer to v1 assert_never(auth_type) +def _token_exchange_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: + """Build a token_exchange (RFC 8693 OBO) spec, or defer (None) when it is not OBO-configured. + + An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the + ``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at + the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the + gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is + nothing to own, so the server stays on v1 (parity-safe). ``audience`` is forwarded only when the + operator set it; a missing one is omitted, not derived. + """ + endpoint = server.token_exchange_endpoint or server.token_url + if not server.client_id or not server.client_secret: + return None + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=TokenExchangeConfig( + subject_token_type=server.subject_token_type or "urn:ietf:params:oauth:token-type:access_token", + token_exchange_endpoint=endpoint, + audience=server.audience, + client_id=server.client_id, + client_secret=SecretStr(server.client_secret), + token_endpoint_auth_method=server.token_endpoint_auth_method, + scopes=tuple(server.scopes or ()), + ), + ) + + def _shared_key_spec( server: MCPServer, resource: str, @@ -148,23 +180,52 @@ def raise_public(error: CredError) -> NoReturn: assert_never(error.tag) -def raise_user_oauth_challenge(server: MCPServer) -> NoReturn: - """Raise the 401 an ``authorization_code`` server returns at egress when the user has no token. +def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str: + """The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges. - Points at the server's RFC 9728 Protected Resource Metadata (``resource_metadata``), which names - the upstream authorization server the client must complete OAuth with. The URL is per-server and - relative, so it resolves against the caller's own host (correct even behind a reverse proxy) - without needing request context. The listing-phase 401 still emits the RFC 8414 ``authorization_uri`` - form pending the format unification; both target the same server, so the difference is cosmetic. + ``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell) + so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is + relative, so it resolves against the caller's own host (correct even behind a reverse proxy). """ - from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 - - root = get_server_root_path() - prefix = "" if root == "/" else root + prefix = "" if root_path == "/" else root_path name = server.alias or server.server_name or server.name or server.server_id - resource_metadata = f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + + +def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn: + """Raise the 401 an ``authorization_code`` server returns at egress when the user has no token. + + Points at the server's RFC 9728 Protected Resource Metadata, which names the upstream + authorization server the client must complete OAuth with. The listing-phase 401 still emits the + RFC 8414 ``authorization_uri`` form pending the format unification; both target the same server, + so the difference is cosmetic. + """ + resource_metadata = oauth_protected_resource_path(root_path, server) raise HTTPException( status_code=401, detail="Unauthorized", headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata}"'}, ) + + +def raise_token_exchange_challenge(server: MCPServer, *, root_path: str) -> NoReturn: + """Raise the RFC 9728 / RFC 6750 challenge an OBO (``token_exchange``) server returns when the + caller's subject token is missing or the IdP rejected it. + + Points at the server's Protected Resource Metadata, whose ``authorization_servers`` names the IdP + the client must SSO with to obtain a subject token; ``error="invalid_token"`` tells a + spec-compliant MCP client to discover that AS and retry with a fresh bearer. Mirrors + ``raise_user_oauth_challenge`` but for the exchange flow: there is no gateway-side browser OAuth — + the client re-authenticates directly with the IdP, and LiteLLM then exchanges the resulting token. + """ + resource_metadata = oauth_protected_resource_path(root_path, server) + www_authenticate = ( + f'Bearer resource_metadata="{resource_metadata}", ' + 'error="invalid_token", ' + 'error_description="Missing or invalid subject token; authenticate with the IdP and retry"' + ) + raise HTTPException( + status_code=401, + detail="Unauthorized", + headers={"WWW-Authenticate": www_authenticate}, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index f9a9fa00b23..c82ce1037d6 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -8,7 +8,8 @@ at runtime instead of returning `None`. `none` and `api_key` (shared-key source) are live, as is `authorization_code`, which reads the -user's token from the injected `OAuthTokenStore`. The remaining arms are `not_implemented` stubs +user's token from the injected `OAuthTokenStore`, and `token_exchange`, which swaps the caller's +inbound token through the injected `TokenExchanger`. The remaining arms are `not_implemented` stubs that each land in a follow-up PR with their seam. Pure v2: no imports from v1. """ @@ -31,6 +32,9 @@ Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + TokenExchanger, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -55,16 +59,36 @@ async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: return None +class _NullTokenExchanger: + """Fail-closed default: with no exchanger wired, token_exchange cannot produce a credential.""" + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> Result[OAuthToken, CredError]: + return Error(CredError.of_misconfigured("token exchange collaborator not wired")) + + async def invalidate( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> None: + return None + + class UpstreamCredentialProvider: """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode. Collaborators (the per-mode credential stores and token fetchers) are injected as each arm is built; the live `none` and `api_key`-shared arms read from the config and need none, while - `authorization_code` reads the user's token from the injected `OAuthTokenStore`. + `authorization_code` reads the user's token from the injected `OAuthTokenStore` and + `token_exchange` swaps the caller's token through the injected `TokenExchanger`. """ - def __init__(self, oauth_token_store: OAuthTokenStore | None = None) -> None: + def __init__( + self, + oauth_token_store: OAuthTokenStore | None = None, + token_exchanger: TokenExchanger | None = None, + ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() + self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -76,8 +100,8 @@ async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Res return _not_implemented(AuthSpecKind.passthrough) case ClientCredentialsConfig(): return _not_implemented(AuthSpecKind.client_credentials) - case TokenExchangeConfig(): - return _not_implemented(AuthSpecKind.token_exchange) + case TokenExchangeConfig() as config: + return await self._token_exchange(subject, server, config) case AuthorizationCodeConfig(): return await self._authorization_code(subject, server) case AwsSigV4Config(): @@ -110,6 +134,43 @@ async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Res return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server.")) return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + async def _token_exchange( + self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig + ) -> Result[StaticHeaderAuth, CredError]: + """RFC 8693 OBO: exchange the caller's inbound token for an upstream-bound bearer. + + No inbound token means there is nothing to exchange, so the arm fails closed with a 401 rather + than falling through to a weaker source (§1.5); the exchanger handles the IdP round-trip and + caching and returns the upstream token or a typed error. + """ + inbound = subject.inbound_token + if inbound is None: + return Error( + CredError.of_unauthorized( + "Token exchange requires a caller token to exchange (OBO).", + www_authenticate='Bearer error="invalid_request"', + ) + ) + match await self._token_exchanger.exchange( + inbound.get_secret_value(), server, config, tenant_id=subject.tenant_id + ): + case Ok(token): + return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + case Error(err): + return Error(err) + + async def invalidate_credentials(self, subject: Subject, server: ServerSpec) -> None: + """Drop any cached credential the resolver owns for this `(subject, server)`. + + Used after an upstream rejects the injected credential, so the next resolve re-mints rather + than serving the same rejected token until TTL. Only `token_exchange` holds a re-mintable + cached credential here; other modes are a no-op. + """ + if isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None: + await self._token_exchanger.invalidate( + subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id + ) + async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None: """The user's authorization_code token, or None when absent or the store is unreachable. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py new file mode 100644 index 00000000000..7161dbb32d4 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py @@ -0,0 +1,105 @@ +"""Composition root for the v2-native token_exchange (OBO) exchanger. + +Wires the pure ``Rfc8693TokenExchanger`` to its runtime edges: the real httpx POST against the IdP and +the configured cache sizing/TTL constants. ``build_token_exchanger`` is built once at egress +construction and reused, so the in-process exchanged-token cache survives across requests. Unlike the +per-user store, nothing here reads a runtime global at build time (the httpx client is acquired per +call), so it needs no lazy wrapper. +""" + +from __future__ import annotations + +import httpx + +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, + SubjectTokenRejected, + TokenExchangeClientError, +) + +# RFC 6749 5.2 error codes that mean the gateway's own request/credentials are wrong (not the +# caller's subject token), so they surface as a 500 the caller can't fix by re-authenticating. +_GATEWAY_FAULT_OAUTH_ERRORS = frozenset( + {"invalid_client", "unauthorized_client", "unsupported_grant_type", "invalid_target", "invalid_scope"} +) + + +def _oauth_error_code(response: httpx.Response) -> str | None: + """Read the RFC 6749 5.2 ``error`` code from a token-endpoint error body, or None if absent. + + The ``error_description`` is deliberately not read: it can carry IdP internals and must never + reach the caller. Only the standard machine code drives classification. + """ + try: + body: object = response.json() + except Exception: # noqa: BLE001 + return None + if isinstance(body, dict): + code = body.get("error") + if isinstance(code, str): + return code + return None + + +async def _post_exchange_endpoint( + url: str, form: dict[str, str], client_auth_headers: dict[str, str] +) -> dict[str, object] | None: + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + # litellm's httpx handler and httpx.Response are only partially typed; the IdP returns a JSON + # object and the exchanger validates each field, so the untyped boundary is contained here. + # A 4xx is the IdP rejecting the subject (non-retryable -> 401 via SubjectTokenRejected); any + # other failure is a miss (-> None -> upstream_unavailable -> 503), matching v1's fail-closed. + headers = {"Accept": "application/json", **client_auth_headers} + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore + response = await client.post(url, headers=headers, data=form) # pyright: ignore + response.raise_for_status() # pyright: ignore + parsed: object = response.json() # pyright: ignore + except httpx.HTTPStatusError as status_err: + status_code = status_err.response.status_code + if 400 <= status_code < 500: + oauth_error = _oauth_error_code(status_err.response) + if oauth_error in _GATEWAY_FAULT_OAUTH_ERRORS: + verbose_logger.warning( + "MCP token exchange rejected as %s (HTTP %d); check the gateway client credentials, " + "audience, and scope for this server", + oauth_error, + status_code, + ) + raise TokenExchangeClientError(oauth_error) from status_err + raise SubjectTokenRejected(f"IdP rejected the subject token (HTTP {status_code})") from status_err + verbose_logger.warning("MCP token exchange request failed: %s", status_err) + return None + except Exception as exc: # noqa: BLE001 + verbose_logger.warning("MCP token exchange request failed: %s", exc) + return None + if not isinstance(parsed, dict): + # A valid-but-non-object JSON body (list/string/number) would crash the field parsing; map it + # to a miss so it surfaces as a typed upstream_unavailable, not a 500. + verbose_logger.warning("MCP token exchange returned non-object JSON (%s)", type(parsed).__name__) + return None + return parsed # pyright: ignore + + +def build_token_exchanger() -> Rfc8693TokenExchanger: + return Rfc8693TokenExchanger( + _post_exchange_endpoint, + cache=InMemoryTokenCacheBackend(max_size=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE), + default_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + min_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + expiry_buffer_seconds=MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py new file mode 100644 index 00000000000..9f1787da612 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py @@ -0,0 +1,298 @@ +"""v2-native RFC 8693 token exchange (OBO): swap the caller's token for an upstream one. + +The pure core of the ``token_exchange`` mode. Given the caller's ``subject_token`` and the server's +``TokenExchangeConfig``, ``Rfc8693TokenExchanger.exchange`` POSTs the RFC 8693 token-exchange grant to +the configured endpoint and returns the upstream-bound ``access_token`` as a typed ``OAuthToken``, or a +typed ``CredError`` - never a raise (the HTTP edge is the injected ``ExchangeHttpPost``, whose adapter +contains the I/O). The exchanged token is cached and single-flighted per ``(subject_token, server)`` so +a repeated caller token skips the IdP round-trip and concurrent calls collapse to one exchange, reusing +the shared in-process cache + coordinator foundation. A rotated caller token hashes to a new key and +re-exchanges. Pure v2 apart from the shared RFC 6749 client-auth helper, which carries no v1 state. + +A missing/expired exchange is an error, never a fall-through to a weaker source (§1.5): the caller +presenting no token is the resolver arm's 401, and an IdP that does not return a usable token is an +``upstream_unavailable`` here. +""" + +from __future__ import annotations + +import hashlib +import time +from collections.abc import Awaitable, Callable +from typing import Protocol + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, + InProcessRefreshCoordinator, + OAuthToken, + RefreshCoordinator, + TokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( + build_token_endpoint_client_auth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + CredError, + ServerSpec, + TokenExchangeConfig, +) + +# A token with no declared expiry is cached for this long; one with an expiry is cached until then +# minus the skew buffer, floored at the minimum. Values mirror v1's MCP_OAUTH2_* constants; the +# composition root injects the configured ones. +_DEFAULT_TTL_SECONDS = 3600.0 +_MIN_TTL_SECONDS = 10.0 +_EXPIRY_BUFFER_SECONDS = 60.0 + +_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" + +# RFC 8693 3 token-type URNs that are not usable as an upstream Bearer access token. token_type +# already rejects the common non-access case (N_A); this catches a malformed STS that mints one of +# these but still labels it Bearer. An access_token / jwt / absent / unknown type is accepted (lenient). +_NON_ACCESS_ISSUED_TOKEN_TYPES = frozenset( + { + "urn:ietf:params:oauth:token-type:refresh_token", + "urn:ietf:params:oauth:token-type:id_token", + "urn:ietf:params:oauth:token-type:saml1", + "urn:ietf:params:oauth:token-type:saml2", + } +) + +# The IdP returns an opaque JSON object; the post adapter hands it over untyped and the exchanger +# validates each field, so no Any leaks past this seam (None == any transport/HTTP failure). The +# second dict is the form body; the third is the client-auth headers (HTTP Basic for +# client_secret_basic, empty for client_secret_post). +ExchangeHttpPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable["dict[str, object] | None"]] + + +class SubjectTokenRejected(Exception): + """The IdP refused to exchange the subject token (an RFC 8693 4xx, e.g. ``invalid_grant``). + + Distinct from a transport / IdP-availability failure, which the post adapter maps to ``None`` -> + ``upstream_unavailable`` -> 503 (retryable). A rejected subject is the caller's problem, not the + gateway's, so the arm surfaces it as a non-retryable 401 (the OBO challenge) instead. + """ + + +class TokenExchangeClientError(Exception): + """The IdP rejected the exchange for a reason that is the gateway's fault, not the caller's. + + RFC 6749 5.2 codes such as ``invalid_client`` (the gateway's own STS credentials are wrong), + ``unauthorized_client`` / ``unsupported_grant_type`` (the gateway is not permitted to exchange), + ``invalid_target`` / ``invalid_scope`` (the gateway's audience/scope config for this server is + wrong). The caller cannot fix these by re-authenticating, so the arm surfaces them as a 500 + (``misconfigured``), not the 401 OBO challenge. The IdP ``error_description`` is never carried. + """ + + +class TokenExchanger(Protocol): + """Exchanges a caller token for an upstream-bound one, per the server's token_exchange config.""" + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> Result[OAuthToken, CredError]: ... + + async def invalidate( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> None: ... + + +def _cache_key(subject_token: str, tenant_id: str, config: TokenExchangeConfig) -> str: + """Bind the cache entry to the caller token, the tenant, AND the exchange config that minted it. + + A rotated caller token, a different tenant, endpoint, audience, scope, client_id, secret, auth + method, or subject_token_type all change the key, so two tenants behind the same opaque token + never share an entry and a config change forces a fresh exchange instead of serving a token + minted for the old config until TTL. Everything is hashed, so no secret is held in the key. + """ + secret = config.client_secret.get_secret_value() if config.client_secret else "" + material = "\x00".join( + ( + subject_token, + tenant_id, + config.token_exchange_endpoint or "", + config.audience or "", + config.subject_token_type, + config.client_id or "", + secret, + config.token_endpoint_auth_method or "", + " ".join(config.scopes), + ) + ) + return hashlib.sha256(material.encode()).hexdigest() + + +def _parse_expires_in(raw: object) -> int | None: + if isinstance(raw, bool): + return None + if isinstance(raw, (int, float)): + return int(raw) + if isinstance(raw, str): + try: + return int(float(raw)) + except ValueError: + return None + return None + + +def _build_exchange_form( + *, + subject_token: str, + subject_token_type: str, + audience: str | None, + scopes: tuple[str, ...], +) -> dict[str, str]: + return { + "grant_type": _GRANT_TYPE, + "subject_token": subject_token, + "subject_token_type": subject_token_type, + **({"audience": audience} if audience else {}), + **({"scope": " ".join(scopes)} if scopes else {}), + } + + +class Rfc8693TokenExchanger: + """``TokenExchanger`` that runs the RFC 8693 grant once per caller token, then caches the result. + + The HTTP post is injected (``None`` on any IdP failure, mirroring v1: a failed exchange is a miss, + not a 500). The cache and single-flight coordinator default to the in-process foundation; a + deployment with no shared state needs nothing more (v1's exchanged-token cache is per-process too). + The clock is injected so TTL/expiry is deterministic in tests. + """ + + def __init__( + self, + http_post: ExchangeHttpPost, + *, + cache: TokenCacheBackend | None = None, + coordinator: RefreshCoordinator | None = None, + clock: Callable[[], float] = time.time, + default_ttl_seconds: float = _DEFAULT_TTL_SECONDS, + min_ttl_seconds: float = _MIN_TTL_SECONDS, + expiry_buffer_seconds: float = _EXPIRY_BUFFER_SECONDS, + ) -> None: + self._http_post = http_post + self._cache: TokenCacheBackend = cache or InMemoryTokenCacheBackend(clock=clock) + self._coordinator: RefreshCoordinator = coordinator or InProcessRefreshCoordinator() + self._clock = clock + self._default_ttl_seconds = default_ttl_seconds + self._min_ttl_seconds = min_ttl_seconds + self._expiry_buffer_seconds = expiry_buffer_seconds + + async def exchange( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> Result[OAuthToken, CredError]: + endpoint = config.token_exchange_endpoint + client_id = config.client_id + client_secret = config.client_secret + if not endpoint: + # No endpoint configured and none discoverable: fail closed (412) rather than guess an IdP + # or fall back to a weaker source. The caller's token is never sent anywhere. + return Error( + CredError.of_precondition_required("token exchange endpoint is not configured for this server") + ) + if not client_id or client_secret is None: + return Error(CredError.of_misconfigured("token_exchange requires client_id and client_secret")) + + cache_key = _cache_key(subject_token, tenant_id, config) + server_id = server.server_id + cached = await self._cache.get(cache_key, server_id) + if cached is not None: + verbose_logger.debug("MCP token exchange cache hit for server %s", server_id) + return Ok(cached) + + client_auth = build_token_endpoint_client_auth( + auth_method=config.token_endpoint_auth_method, + client_id=client_id, + client_secret=client_secret.get_secret_value(), + ) + form = { + **_build_exchange_form( + subject_token=subject_token, + subject_token_type=config.subject_token_type, + audience=config.audience, + scopes=config.scopes, + ), + **client_auth.body, + } + + async def run_exchange() -> OAuthToken | None: + fresh = await self._cache.get(cache_key, server_id) + if fresh is not None: + return fresh + verbose_logger.debug( + "Exchanging token for MCP server %s at %s (audience=%s)", server_id, endpoint, config.audience + ) + body = await self._http_post(endpoint, form, client_auth.headers) + if body is None: + return None + token = self._token_from_body(body) + if token is None: + return None + await self._cache.set(cache_key, server_id, token, self._ttl_seconds(token)) + verbose_logger.info("Token exchange succeeded for MCP server %s", server_id) + return token + + async def reread() -> OAuthToken | None: + return await self._cache.get(cache_key, server_id) + + try: + token = await self._coordinator.run(cache_key, server_id, refresh=run_exchange, reread=reread) + except SubjectTokenRejected as rejected: + # The IdP rejected the subject token (4xx). This is non-retryable: the caller must + # re-authenticate with the IdP, so it surfaces as a 401 (the OBO challenge), not a 503. + return Error(CredError.of_unauthorized(str(rejected) or "subject token rejected by the IdP")) + except TokenExchangeClientError: + # RFC 6749 5.2 gateway-fault code (invalid_client / invalid_target / ...): the caller can't + # fix it by re-authenticating, so surface a 500 rather than the OBO 401 challenge. The + # specific code is logged at the edge; the user-facing summary stays generic. + return Error( + CredError.of_misconfigured( + "token exchange configuration error: the gateway's credentials, audience, or scope " + "for this server were not accepted by the IdP" + ) + ) + if token is None: + return Error(CredError.of_upstream_unavailable("token exchange did not return a usable access token")) + return Ok(token) + + async def invalidate( + self, subject_token: str, server: ServerSpec, config: TokenExchangeConfig, *, tenant_id: str = "" + ) -> None: + """Drop the cached exchanged token so the next call re-exchanges (e.g. after an upstream 401).""" + await self._cache.delete(_cache_key(subject_token, tenant_id, config), server.server_id) + + def _token_from_body(self, body: dict[str, object]) -> OAuthToken | None: + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return None + # token_type is forwarded downstream as Bearer, so a present-but-non-Bearer type (e.g. N_A) + # must fail closed rather than be minted as a bogus Bearer; an absent type defaults to Bearer. + token_type = body.get("token_type") + if isinstance(token_type, str) and token_type.strip().lower() != "bearer": + verbose_logger.warning( + "MCP token exchange returned unusable token_type %r; refusing to forward it as Bearer", token_type + ) + return None + # issued_token_type says what representation was minted; reject a clearly-non-access type + # (refresh/id/saml) even if token_type claimed Bearer. access_token / jwt / absent / unknown pass. + issued_token_type = body.get("issued_token_type") + if isinstance(issued_token_type, str) and issued_token_type in _NON_ACCESS_ISSUED_TOKEN_TYPES: + return None + expires_in = _parse_expires_in(body.get("expires_in")) + expires_at = self._clock() + expires_in if expires_in is not None else None + return OAuthToken(access_token=access_token, expires_at=expires_at) + + def _ttl_seconds(self, token: OAuthToken) -> float: + if token.expires_at is None: + return self._default_ttl_seconds + lifetime = max(0.0, token.expires_at - self._clock()) + # Floor at min_ttl, but never cache past the token's own expiry: a token whose remaining + # lifetime is below the buffer (or even below min_ttl) must not be served stale upstream. + return min(max(lifetime - self._expiry_buffer_seconds, self._min_ttl_seconds), lifetime) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 671de63eabe..d475f931a25 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -183,17 +183,23 @@ class ClientCredentialsConfig(BaseModel): class TokenExchangeConfig(BaseModel): """RFC 8693 OBO; swap the caller's live subject_token for a token bound to the upstream's - audience (`server.resource`, RFC 8707). The gateway authenticates to the exchange endpoint - as an OAuth client (`client_id`/`client_secret`); the inbound token is sent only to that - endpoint, never to the upstream. + audience. The gateway authenticates to the exchange endpoint as an OAuth client + (`client_id`/`client_secret`); the inbound token is sent only to that endpoint, never to the + upstream. + + `audience` is the RFC 8693 target; it is optional and sent only when the operator configured + one, since both `audience` and `resource` are optional in the spec and the authorization server + applies its own default when neither is sent (fabricating one risks `invalid_target`). """ model_config = ConfigDict(frozen=True) kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token" token_exchange_endpoint: str | None = None + audience: str | None = None client_id: str | None = None client_secret: SecretStr | None = None + token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] | None = None scopes: tuple[str, ...] = () diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 4b55510a629..67a726d8d18 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1625,6 +1625,7 @@ async def _fetch_and_filter_server_tools( add_prefix=True, # Always add server prefix raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, + oauth2_headers=oauth2_headers, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) @@ -2473,12 +2474,18 @@ async def execute_mcp_tool( # Forward named client headers to OpenAPI tool upstream requests. # MCPServer.extra_headers lists header names to copy from raw_headers. - # OAuth2 M2M: never take Authorization from the caller (matches - # _prepare_mcp_server_headers for managed MCP). + # The strip decision is centralized in _should_strip_caller_authorization so this + # OpenAPI/local path agrees with the managed paths: M2M and the resolver-owned modes + # (token_exchange's raw subject token, authorization_code's stored token) must never + # have the caller's Authorization forwarded verbatim upstream. forwarded_headers: Optional[Dict[str, str]] = None if mcp_server and mcp_server.extra_headers and raw_headers: normalized_raw = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} - skip_caller_authorization = bool(mcp_server.has_client_credentials) + skip_caller_authorization = _should_strip_caller_authorization( + mcp_server=mcp_server, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) for header_name in mcp_server.extra_headers: if not isinstance(header_name, str): continue @@ -3240,6 +3247,36 @@ async def _raise_preemptive_401_for_unauthenticated_servers( headers={"www-authenticate": authorization_uri}, ) + # token_exchange (OBO): the caller supplied no subject token. Challenge at connect + # (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata + # so the client discovers the IdP, SSOs, and retries with a subject token, which LiteLLM + # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the + # header lost, so the discovery flow needs this pre-emptive challenge. + if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers: + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + raise_token_exchange_challenge, + ) + from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 + + raise_token_exchange_challenge(server, root_path=get_server_root_path()) + + # token_exchange (OBO) with a subject present: run the exchange here at the transport + # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its + # public status) instead of the session opening and list_tools masking the failure as + # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps + # absorbing per-server auth failures so one bad server cannot 401 the whole connect. + if ( + server + and server.auth_type == MCPAuth.oauth2_token_exchange + and oauth2_headers + and len(mcp_servers or []) == 1 + ): + await global_mcp_server_manager.preflight_token_exchange( + server=server, + oauth2_headers=oauth2_headers, + user_api_key_auth=user_api_key_auth, + ) + # 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 diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index bb13a7ce8cc..321a17fbb03 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -955,6 +955,7 @@ async def mock_get_tools_side_effect( add_prefix=False, raw_headers=None, user_api_key_auth=None, + oauth2_headers=None, ): if server.server_id == "server1_id": return [mock_tool_1] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 383dc255607..cdd88c59923 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -7,12 +7,12 @@ import base64 from types import SimpleNamespace -from unittest.mock import patch import pytest from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + oauth_protected_resource_path, raise_public, raise_user_oauth_challenge, to_server_spec, @@ -24,6 +24,7 @@ CredError, NoneConfig, SharedKey, + TokenExchangeConfig, ) from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -65,9 +66,7 @@ def test_authorization_schemes_map_with_their_prefix(auth_type, prefix): def test_basic_scheme_base64_encodes_the_token(): - spec = to_server_spec( - _server(auth_type=MCPAuth.basic, authentication_token="user:pass") - ) + spec = to_server_spec(_server(auth_type=MCPAuth.basic, authentication_token="user:pass")) assert spec is not None and isinstance(spec.config, ApiKeyConfig) assert spec.config.value_prefix == "Basic" expected = base64.b64encode(b"user:pass").decode() @@ -89,17 +88,16 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): [ _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured + _server(auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials"), # M2M -> v1 + _server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1 + _server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1 _server( - auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials" - ), # M2M -> v1 - _server( - auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True - ), # delegated upstream OAuth -> v1 - _server(auth_type=MCPAuth.oauth2_token_exchange), + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp/token", + client_id="cid", + ), # missing client_secret -> incomplete -> v1 _server(auth_type=MCPAuth.aws_sigv4), - _server( - auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"] - ), + _server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]), ], ) def test_unmigrated_modes_defer_to_v1(server): @@ -107,6 +105,91 @@ def test_unmigrated_modes_defer_to_v1(server): assert to_server_spec(server) is None +def test_token_exchange_maps_full_config(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + url="https://up.example.com/mcp", + token_exchange_endpoint="https://idp.example.com/token", + audience="https://up.example.com", + client_id="cid", + client_secret="csec", + subject_token_type="urn:ietf:params:oauth:token-type:jwt", + scopes=["a", "b"], + ) + ) + assert spec is not None + config = spec.config + assert isinstance(config, TokenExchangeConfig) + assert config.token_exchange_endpoint == "https://idp.example.com/token" + assert config.audience == "https://up.example.com" + assert config.client_id == "cid" + assert config.client_secret is not None + assert config.client_secret.get_secret_value() == "csec" + assert config.subject_token_type == "urn:ietf:params:oauth:token-type:jwt" + assert config.scopes == ("a", "b") + + +def test_token_exchange_falls_back_to_token_url_when_no_exchange_endpoint(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.token_exchange_endpoint == "https://idp.example.com/token" + + +def test_token_exchange_with_creds_but_no_endpoint_is_owned_for_fail_closed(): + # An OBO server with client credentials but no endpoint is still owned by v2 (spec, not None) so + # it fails closed at the exchanger (412) rather than silently deferring to v1 and connecting + # unauthenticated. The endpoint stays None for the exchanger to reject. + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.token_exchange_endpoint is None + + +def test_token_exchange_omits_audience_when_unset(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.audience is None + + +def test_token_exchange_empty_subject_token_type_normalizes_to_default(): + # Parity with v1: a falsy subject_token_type must not be sent verbatim to the IdP. + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + subject_token_type="", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:access_token" + + @pytest.mark.parametrize( "server", [ @@ -115,9 +198,7 @@ def test_unmigrated_modes_defer_to_v1(server): # static token must not route a BYOK server to a v2 shared-key spec with the wrong value. _server(auth_type=MCPAuth.bearer_token, is_byok=True, authentication_token="x"), _server(auth_type=MCPAuth.basic, is_byok=True, authentication_token="x"), - _server( - auth_type=MCPAuth.authorization, is_byok=True, authentication_token="x" - ), + _server(auth_type=MCPAuth.authorization, is_byok=True, authentication_token="x"), _server(auth_type=MCPAuth.token, is_byok=True, authentication_token="x"), _server(auth_type=None, is_byok=True), ], @@ -161,9 +242,7 @@ def test_raise_public_maps_each_error_to_its_status(error, status): def test_raise_public_emits_unauthorized_challenge(): body = {"error": "byok_auth_required", "server_id": "s1"} - error = CredError.of_unauthorized( - "needs key", www_authenticate='Bearer resource_metadata="/x"', body=body - ) + error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body) with pytest.raises(HTTPException) as exc_info: raise_public(error) exc = exc_info.value @@ -182,41 +261,72 @@ def test_raise_public_plain_unauthorized_has_no_challenge(): assert exc.headers is None -_ROOT_PATH = "litellm.proxy.utils.get_server_root_path" +@pytest.mark.parametrize( + "root_path, expected_prefix", + [ + ("/", ""), # "/" means no prefix + ("", ""), # empty means no prefix + ("/api/v1", "/api/v1"), # a real root path is prepended verbatim + ], +) +def test_oauth_protected_resource_path_honors_root_path(root_path, expected_prefix): + path = oauth_protected_resource_path(root_path, _server(alias="my-srv")) + assert path == f"/.well-known/oauth-protected-resource{expected_prefix}/mcp/my-srv" + + +@pytest.mark.parametrize( + "kwargs, expected_name", + [ + ({"alias": "a", "server_name": "sn"}, "a"), # alias wins + ({"server_name": "sn"}, "sn"), # then server_name + ({}, "n"), # then the name field (server_id is the last fallback) + ], +) +def test_oauth_protected_resource_path_name_fallback(kwargs, expected_name): + assert oauth_protected_resource_path("/", _server(**kwargs)).endswith(f"/mcp/{expected_name}") def test_raise_user_oauth_challenge_points_at_per_server_prm(): - with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info: - raise_user_oauth_challenge(_server(alias="my-srv")) + with pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/") exc = exc_info.value assert exc.status_code == 401 assert ( - exc.headers["WWW-Authenticate"] - == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"' + exc.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"' ) def test_raise_user_oauth_challenge_includes_server_root_path(): - with ( - patch(_ROOT_PATH, return_value="/api/v1"), - pytest.raises(HTTPException) as exc_info, - ): - raise_user_oauth_challenge(_server(alias="my-srv")) + with pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/api/v1") assert ( exc_info.value.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/my-srv"' ) -@pytest.mark.parametrize( - "kwargs, expected_name", - [ - ({"alias": "a", "server_name": "sn"}, "a"), # alias wins - ({"server_name": "sn"}, "sn"), # then server_name - ({}, "n"), # then the name field (server_id is the last fallback) - ], -) -def test_raise_user_oauth_challenge_name_fallback(kwargs, expected_name): - with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info: - raise_user_oauth_challenge(_server(**kwargs)) - assert f'/mcp/{expected_name}"' in exc_info.value.headers["WWW-Authenticate"] +def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_token_exchange_challenge, + ) + + with pytest.raises(HTTPException) as exc_info: + raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/") + exc = exc_info.value + www = exc.headers["WWW-Authenticate"] + assert exc.status_code == 401 + # RFC 9728 resource_metadata so the client can discover the IdP, plus RFC 6750 invalid_token. + assert 'resource_metadata="/.well-known/oauth-protected-resource/mcp/obo-srv"' in www + assert 'error="invalid_token"' in www + assert "error_description=" in www + + +def test_raise_token_exchange_challenge_includes_server_root_path(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_token_exchange_challenge, + ) + + with pytest.raises(HTTPException) as exc_info: + raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/api/v1") + www = exc_info.value.headers["WWW-Authenticate"] + assert 'resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/obo-srv"' in www diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 73e9a52b937..f8e45b38b49 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -1,9 +1,9 @@ """Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -`none`, `api_key` (shared-key source), and `authorization_code` are implemented; every other arm, -plus the `api_key` BYOK source, returns a typed `not_implemented` error until its mode lands. -Parametrizing the stubs over one config each also guards reachability: a dropped `case` would hit -`assert_never` and raise instead of returning the stub. +`none`, `api_key` (shared-key source), `authorization_code`, and `token_exchange` are implemented; +every other arm, plus the `api_key` BYOK source, returns a typed `not_implemented` error until its +mode lands. Parametrizing the stubs over one config each also guards reachability: a dropped `case` +would hit `assert_never` and raise instead of returning the stub. """ import httpx @@ -16,11 +16,13 @@ AwsSigV4Config, Byok, ClientCredentialsConfig, + CredError, Error, NoneConfig, NoOpAuth, Ok, PassthroughConfig, + Result, ServerSpec, SharedKey, StaticHeaderAuth, @@ -37,9 +39,7 @@ def _spec(config): - return ServerSpec( - server_id="s", resource="https://upstream.example.com", config=config - ) + return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) def _emitted(auth: httpx.Auth) -> httpx.Headers: @@ -52,9 +52,7 @@ def _emitted(auth: httpx.Auth) -> httpx.Headers: @pytest.mark.asyncio async def test_none_mode_yields_a_no_op_auth(): - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(NoneConfig()) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(NoneConfig())) assert isinstance(result, Ok) assert isinstance(result.ok, NoOpAuth) @@ -66,9 +64,7 @@ async def test_api_key_shared_emits_the_configured_header(): value_prefix="", key_source=SharedKey(value=SecretStr("secret-key")), ) - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Ok) assert isinstance(result.ok, StaticHeaderAuth) assert _emitted(result.ok)["X-API-Key"] == "secret-key" @@ -81,9 +77,7 @@ async def test_api_key_shared_honors_authorization_scheme(): value_prefix="Bearer", key_source=SharedKey(value=SecretStr("tok")), ) - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer tok" @@ -101,9 +95,7 @@ async def fetch(self, user_id: str, server_id: str): @pytest.mark.asyncio async def test_authorization_code_emits_bearer_for_a_stored_token(): store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")}) - result = await UpstreamCredentialProvider( - oauth_token_store=store - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=store).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Ok) @@ -112,9 +104,7 @@ async def test_authorization_code_emits_bearer_for_a_stored_token(): @pytest.mark.asyncio async def test_authorization_code_without_token_is_semantically_unauthorized(): - result = await UpstreamCredentialProvider( - oauth_token_store=_FakeTokenStore({}) - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Error) @@ -131,9 +121,7 @@ class _Unavailable: async def fetch(self, user_id: str, server_id: str): raise TokenStoreUnavailable("down") - result = await UpstreamCredentialProvider( - oauth_token_store=_Unavailable() - ).resolve_credentials( + result = await UpstreamCredentialProvider(oauth_token_store=_Unavailable()).resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) assert isinstance(result, Error) @@ -156,22 +144,15 @@ async def test_authorization_code_isolates_by_subject(): alice = await provider.resolve_credentials( Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) ) - bob = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig()) - ) - assert ( - isinstance(alice, Ok) - and _emitted(alice.ok)["Authorization"] == "Bearer at-alice" - ) + bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig())) + assert isinstance(alice, Ok) and _emitted(alice.ok)["Authorization"] == "Bearer at-alice" assert isinstance(bob, Error) and bob.error.tag == "unauthorized" @pytest.mark.asyncio async def test_has_user_token_reflects_the_stored_token(): present = UpstreamCredentialProvider( - oauth_token_store=_FakeTokenStore( - {("alice", "s"): OAuthToken(access_token="at")} - ) + oauth_token_store=_FakeTokenStore({("alice", "s"): OAuthToken(access_token="at")}) ) absent = UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})) spec = _spec(AuthorizationCodeConfig()) @@ -185,17 +166,97 @@ async def test_has_user_token_false_for_a_non_per_user_mode(): # A none-mode server has no per-user token to check. provider = UpstreamCredentialProvider() spec = _spec(NoneConfig()) - assert ( - await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) - is False + assert await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) is False + + +class _FakeExchanger: + def __init__(self, result: Result[OAuthToken, CredError]) -> None: + self._result = result + self.calls: list[tuple[str, str, str]] = [] + self.invalidations: list[tuple[str, str, str]] = [] + + async def exchange(self, subject_token, server, config, *, tenant_id=""): + self.calls.append((subject_token, tenant_id, server.server_id)) + return self._result + + async def invalidate(self, subject_token, server, config, *, tenant_id=""): + self.invalidations.append((subject_token, tenant_id, server.server_id)) + + +_OBO = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), +) + + +@pytest.mark.asyncio +async def test_token_exchange_emits_the_exchanged_bearer(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + subject = Subject(tenant_id="acme", subject_id="alice", inbound_token=SecretStr("caller-jwt")) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(_OBO)) + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer exchanged-at" + # The arm hands the unwrapped caller token AND the tenant to the exchanger, never the upstream. + assert exchanger.calls == [("caller-jwt", "acme", "s")] + + +@pytest.mark.asyncio +async def test_invalidate_credentials_drops_the_exchanged_token_for_the_subject_and_tenant(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + provider = UpstreamCredentialProvider(token_exchanger=exchanger) + subject = Subject(tenant_id="acme", subject_id="alice", inbound_token=SecretStr("caller-jwt")) + await provider.invalidate_credentials(subject, _spec(_OBO)) + assert exchanger.invalidations == [("caller-jwt", "acme", "s")] + + +@pytest.mark.asyncio +async def test_invalidate_credentials_is_a_noop_without_a_caller_token(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="never"))) + provider = UpstreamCredentialProvider(token_exchanger=exchanger) + await provider.invalidate_credentials(Subject(tenant_id="acme", subject_id="alice"), _spec(_OBO)) + assert exchanger.invalidations == [] + + +@pytest.mark.asyncio +async def test_token_exchange_without_caller_token_is_unauthorized(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="never"))) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_OBO) + ) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + assert result.error.unauthorized.www_authenticate == 'Bearer error="invalid_request"' + # No caller token means nothing to exchange: the IdP is never hit. + assert exchanger.calls == [] + + +@pytest.mark.asyncio +async def test_token_exchange_propagates_the_exchanger_error(): + err = CredError.of_upstream_unavailable("idp down") + result = await UpstreamCredentialProvider(token_exchanger=_FakeExchanger(Error(err))).resolve_credentials( + Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr("jwt")), + _spec(_OBO), ) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_token_exchange_without_an_exchanger_fails_closed(): + # The fail-closed default (no exchanger wired) must not produce a credential. + result = await UpstreamCredentialProvider().resolve_credentials( + Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr("jwt")), + _spec(_OBO), + ) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), ("passthrough", PassthroughConfig()), ("client_credentials", ClientCredentialsConfig()), - ("token_exchange", TokenExchangeConfig()), ("aws_sigv4", AwsSigV4Config(region="us-east-1")), ] @@ -203,8 +264,6 @@ async def test_has_user_token_false_for_a_non_per_user_mode(): @pytest.mark.asyncio @pytest.mark.parametrize("label, config", _STUBBED) async def test_unbuilt_arms_fail_closed_with_not_implemented(label, config): - result = await UpstreamCredentialProvider().resolve_credentials( - _SUBJECT, _spec(config) - ) + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Error) assert result.error.tag == "not_implemented" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py new file mode 100644 index 00000000000..7cea99f0b5f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py @@ -0,0 +1,118 @@ +"""Tests for the token_exchange composition root: the built exchanger and the HTTP edge contract. + +`build_token_exchanger` wires the pure exchanger to its runtime edges; `_post_exchange_endpoint` is +the I/O edge that maps any transport/HTTP failure to None and parses a JSON body on success. +""" + +from unittest.mock import patch + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( + _post_exchange_endpoint, + build_token_exchanger, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, + SubjectTokenRejected, + TokenExchangeClientError, +) + +_HTTP_CLIENT = "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + + +def _client_raising_4xx(body: object): + """An httpx client whose POST returns a 4xx whose ``raise_for_status`` raises an HTTPStatusError + carrying ``body`` as its JSON, so the RFC 6749 error-code classification can be driven.""" + import httpx + + request = httpx.Request("POST", "https://idp/token") + response = httpx.Response(400, json=body, request=request) + + class _Resp: + def raise_for_status(self) -> None: + raise httpx.HTTPStatusError("bad request", request=request, response=response) + + class _Client: + async def post(self, url, headers, data): + return _Resp() + + return _Client() + + +def test_build_token_exchanger_returns_an_exchanger(): + assert isinstance(build_token_exchanger(), Rfc8693TokenExchanger) + + +def test_build_gives_each_caller_an_independent_cache(): + # Separate builds must not share a cache, so one egress instance cannot serve another's tokens. + assert build_token_exchanger() is not build_token_exchanger() + + +@pytest.mark.asyncio +async def test_post_returns_none_on_transport_error(): + with patch(_HTTP_CLIENT, side_effect=RuntimeError("boom")): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result is None + + +@pytest.mark.asyncio +async def test_post_parses_json_body_on_success(): + class _Resp: + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, object]: + return {"access_token": "x", "expires_in": 60} + + class _Client: + async def post(self, url, headers, data): + return _Resp() + + with patch(_HTTP_CLIENT, return_value=_Client()): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result == {"access_token": "x", "expires_in": 60} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code", ["invalid_client", "unauthorized_client", "unsupported_grant_type", "invalid_target", "invalid_scope"] +) +async def test_post_maps_gateway_fault_4xx_to_client_error(code): + # RFC 6749 5.2 gateway-fault codes must raise TokenExchangeClientError (-> 500), not the caller 401. + with patch(_HTTP_CLIENT, return_value=_client_raising_4xx({"error": code})): + with pytest.raises(TokenExchangeClientError): + await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [{"error": "invalid_grant"}, {"error": "invalid_request"}, {}, {"error": 123}, "not-json-object"], + ids=["invalid_grant", "invalid_request", "no_error", "non_str_error", "non_dict"], +) +async def test_post_maps_subject_fault_4xx_to_subject_rejected(body): + # A subject-fault code (or an unparseable/absent error) is the caller's problem -> SubjectTokenRejected (401). + with patch(_HTTP_CLIENT, return_value=_client_raising_4xx(body)): + with pytest.raises(SubjectTokenRejected): + await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [["a", "b"], "a-string", 42], ids=["list", "str", "int"]) +async def test_post_returns_none_on_non_object_json(payload): + # A valid-but-non-object JSON body must become a miss, not crash field parsing downstream. + class _Resp: + def raise_for_status(self) -> None: + return None + + def json(self) -> object: + return payload + + class _Client: + async def post(self, url, headers, data): + return _Resp() + + with patch(_HTTP_CLIENT, return_value=_Client()): + result = await _post_exchange_endpoint("https://idp/token", {"grant_type": "x"}, {}) + assert result is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py new file mode 100644 index 00000000000..13a0652474b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py @@ -0,0 +1,476 @@ +"""Tests for the pure RFC 8693 token exchanger: the OBO swap, caching, and single-flight. + +Drives `Rfc8693TokenExchanger` through an injected fake HTTP post and clock, so the exchange, the +form it sends, the per-caller-token cache, and the failure mapping are pinned without a live IdP. +""" + +import asyncio + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + Error, + Ok, + ServerSpec, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + Rfc8693TokenExchanger, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + TokenExchangeConfig, +) + +_GRANT = "urn:ietf:params:oauth:grant-type:token-exchange" + +_CONFIG = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + audience="https://up.example.com", + client_id="cid", + client_secret=SecretStr("csec"), + scopes=("s1", "s2"), +) +_SERVER = ServerSpec(server_id="srv", resource="https://up.example.com", config=_CONFIG) + + +class _Clock: + def __init__(self, now: float = 1000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +class _RecordingPost: + def __init__(self, body: dict[str, object] | None) -> None: + self._body = body + self.calls: list[tuple[str, dict[str, str]]] = [] + self.headers: list[dict[str, str]] = [] + + async def __call__(self, url: str, form: dict[str, str], headers: dict[str, str]) -> dict[str, object] | None: + self.calls.append((url, dict(form))) + self.headers.append(dict(headers)) + return self._body + + +def _spec(config: TokenExchangeConfig) -> ServerSpec: + return ServerSpec(server_id="srv", resource="https://up.example.com", config=config) + + +@pytest.mark.asyncio +async def test_exchange_emits_token_and_sends_rfc8693_form(): + post = _RecordingPost({"access_token": "exchanged", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("caller-jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) + assert result.ok.access_token == "exchanged" + url, form = post.calls[0] + assert url == "https://idp.example.com/token" + assert form == { + "grant_type": _GRANT, + "subject_token": "caller-jwt", + "subject_token_type": "urn:ietf:params:oauth:token-type:access_token", + "client_id": "cid", + "client_secret": "csec", + "audience": "https://up.example.com", + "scope": "s1 s2", + } + # client_secret_post is the default: creds in the body, no client-auth header + assert post.headers[0] == {} + + +@pytest.mark.asyncio +async def test_client_secret_basic_sends_authorization_header_and_omits_body_creds(): + import base64 + + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), + token_endpoint_auth_method="client_secret_basic", + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + assert isinstance(result, Ok) + _, form = post.calls[0] + assert "client_id" not in form and "client_secret" not in form + assert post.headers[0]["Authorization"] == "Basic " + base64.b64encode(b"cid:csec").decode() + + +@pytest.mark.asyncio +async def test_client_secret_post_keeps_creds_in_body_with_no_auth_header(): + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret=SecretStr("csec"), + token_endpoint_auth_method="client_secret_post", + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + _, form = post.calls[0] + assert form["client_id"] == "cid" and form["client_secret"] == "csec" + assert "Authorization" not in post.headers[0] + + +@pytest.mark.asyncio +async def test_exchange_maps_idp_rejection_to_unauthorized(): + """An IdP 4xx (surfaced as SubjectTokenRejected by the post adapter) is non-retryable: it maps + to ``unauthorized`` (the 401 OBO challenge), not the retryable ``upstream_unavailable`` (503).""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + SubjectTokenRejected, + ) + + async def _rejecting_post(url, form, headers): + raise SubjectTokenRejected("IdP rejected the token exchange (HTTP 400)") + + result = await Rfc8693TokenExchanger(_rejecting_post, clock=_Clock()).exchange("bad-jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_exchange_maps_gateway_fault_to_misconfigured(): + """A gateway-fault RFC 6749 code (invalid_client / invalid_target / ..., surfaced as + TokenExchangeClientError) is the gateway's problem, not the caller's, so it maps to misconfigured + (500) rather than the retryable 503 or the 401 OBO challenge the caller can't act on.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( + TokenExchangeClientError, + ) + + async def _client_error_post(url, form, headers): + raise TokenExchangeClientError("invalid_client") + + result = await Rfc8693TokenExchanger(_client_error_post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + +@pytest.mark.asyncio +async def test_exchange_maps_transport_failure_to_upstream_unavailable(): + """A post returning None (5xx / network / timeout / malformed body) stays retryable: 503.""" + result = await Rfc8693TokenExchanger(_RecordingPost(None), clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_exchange_caches_per_caller_token(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + second = await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert isinstance(second, Ok) and second.ok.access_token == "x" + assert len(post.calls) == 1 + + +@pytest.mark.asyncio +async def test_rotated_caller_token_re_exchanges(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt-1", _SERVER, _CONFIG) + await exchanger.exchange("jwt-2", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_same_token_different_tenant_does_not_share_cache(): + # Two tenants presenting the same opaque token (e.g. a shared/service token) must not collide on + # one cache entry: tenant_id is part of the key, so each tenant gets its own exchange. + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="globex") + assert len(post.calls) == 2 + # Same tenant + token still hits the cache. + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_invalidate_forces_re_exchange(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + await exchanger.invalidate("jwt", _SERVER, _CONFIG, tenant_id="acme") + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_invalidate_targets_only_the_matching_tenant(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="globex") + await exchanger.invalidate("jwt", _SERVER, _CONFIG, tenant_id="acme") + # globex's entry survives; only acme re-exchanges. + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="globex") + await exchanger.exchange("jwt", _SERVER, _CONFIG, tenant_id="acme") + assert len(post.calls) == 3 + + +@pytest.mark.asyncio +async def test_rotated_config_re_exchanges_before_ttl(): + # Same caller token + server, but the operator rotated the audience/scope: the cached token was + # minted for the old config, so it must re-exchange (not serve the stale token) before TTL. + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + rotated = _CONFIG.model_copy(update={"audience": "https://new.example.com", "scopes": ("s3",)}) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + await exchanger.exchange("jwt", _SERVER, rotated) + assert len(post.calls) == 2 + _, second_form = post.calls[1] + assert second_form["audience"] == "https://new.example.com" + assert second_form["scope"] == "s3" + + +@pytest.mark.asyncio +async def test_rotated_token_endpoint_auth_method_re_exchanges_before_ttl(): + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + rotated = _CONFIG.model_copy(update={"token_endpoint_auth_method": "client_secret_basic"}) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + await exchanger.exchange("jwt", _SERVER, rotated) + assert len(post.calls) == 2 + _, second_form = post.calls[1] + assert "client_id" not in second_form + assert "client_secret" not in second_form + assert "Authorization" in post.headers[1] + + +@pytest.mark.asyncio +async def test_concurrent_callers_single_flight_one_exchange(): + release = asyncio.Event() + + class _Blocking: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, url, form, headers): + self.calls += 1 + await release.wait() + return {"access_token": "x", "expires_in": 3600} + + post = _Blocking() + exchanger = Rfc8693TokenExchanger(post, clock=_Clock()) + first = asyncio.create_task(exchanger.exchange("jwt", _SERVER, _CONFIG)) + second = asyncio.create_task(exchanger.exchange("jwt", _SERVER, _CONFIG)) + await asyncio.sleep(0.02) + release.set() + r1, r2 = await asyncio.gather(first, second) + assert post.calls == 1 + assert isinstance(r1, Ok) and isinstance(r2, Ok) + + +@pytest.mark.asyncio +async def test_idp_failure_is_upstream_unavailable(): + result = await Rfc8693TokenExchanger(_RecordingPost(None), clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_missing_access_token_is_upstream_unavailable(): + post = _RecordingPost({"token_type": "Bearer"}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_type", ["N_A", "n_a", "DPoP", "mac"]) +async def test_non_bearer_token_type_is_refused(token_type): + # RFC 8693 2.2.1: the resolver forwards the exchanged token as `Bearer`. A non-Bearer token_type + # (e.g. N_A = not a standalone access token) must fail closed, not be minted as a bogus Bearer. + post = _RecordingPost({"access_token": "x", "token_type": token_type, "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_non_bearer_token_type_is_logged(): + from unittest.mock import patch + + post = _RecordingPost({"access_token": "x", "token_type": "N_A", "expires_in": 3600}) + with patch( + "litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger.verbose_logger" + ) as mock_logger: + await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert mock_logger.warning.called + assert "N_A" in repr(mock_logger.warning.call_args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_type", ["Bearer", "bearer", "BEARER"]) +async def test_bearer_token_type_is_accepted_case_insensitively(token_type): + post = _RecordingPost({"access_token": "x", "token_type": token_type, "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +async def test_absent_token_type_defaults_to_bearer(): + # Many IdPs omit token_type; absence must not fail the exchange (RFC 6750 default is Bearer). + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "issued_token_type", + [ + "urn:ietf:params:oauth:token-type:refresh_token", + "urn:ietf:params:oauth:token-type:id_token", + "urn:ietf:params:oauth:token-type:saml2", + ], +) +async def test_non_access_issued_token_type_is_refused_even_when_bearer(issued_token_type): + # A malformed STS could mint a refresh/id/saml token but label it Bearer; issued_token_type must + # still fail it closed rather than forward a non-access token as an upstream access credential. + post = _RecordingPost( + {"access_token": "x", "token_type": "Bearer", "issued_token_type": issued_token_type, "expires_in": 3600} + ) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "issued_token_type", + ["urn:ietf:params:oauth:token-type:access_token", "urn:ietf:params:oauth:token-type:jwt", "custom-unknown", None], +) +async def test_access_or_unknown_issued_token_type_is_accepted(issued_token_type): + # access_token / jwt are usable; an absent or unrecognized type is accepted (lenient), so real + # IdPs that omit issued_token_type or use a custom URN keep working. + body: dict[str, object] = {"access_token": "x", "token_type": "Bearer", "expires_in": 3600} + if issued_token_type is not None: + body["issued_token_type"] = issued_token_type + result = await Rfc8693TokenExchanger(_RecordingPost(body), clock=_Clock()).exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "config", + [ + TokenExchangeConfig(token_exchange_endpoint="https://idp/token", client_secret=SecretStr("s")), + TokenExchangeConfig(token_exchange_endpoint="https://idp/token", client_id="c"), + ], +) +async def test_incomplete_config_is_misconfigured_without_hitting_idp(config): + post = _RecordingPost({"access_token": "x"}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + assert post.calls == [] + + +@pytest.mark.asyncio +async def test_missing_endpoint_is_precondition_required_without_hitting_idp(): + # No endpoint configured (and none discoverable): fail closed with a 412-mapped precondition + # rather than guessing an IdP or falling back. The subject token is never POSTed anywhere. + config = TokenExchangeConfig(client_id="c", client_secret=SecretStr("s")) + post = _RecordingPost({"access_token": "x"}) + result = await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert post.calls == [] + + +@pytest.mark.asyncio +async def test_cached_token_expires_after_its_ttl(): + clock = _Clock(1000.0) + # expires_in=120, buffer=60 -> ttl 60 -> cached until t=1060. + post = _RecordingPost({"access_token": "x", "expires_in": 120}) + exchanger = Rfc8693TokenExchanger(post, clock=clock) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1059.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 1 + clock.now = 1061.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_audience_and_scope_omitted_when_unset(): + config = TokenExchangeConfig( + token_exchange_endpoint="https://idp/token", + client_id="cid", + client_secret=SecretStr("csec"), + ) + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + await Rfc8693TokenExchanger(post, clock=_Clock()).exchange("jwt", _spec(config), config) + _, form = post.calls[0] + assert "audience" not in form + assert "scope" not in form + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires_in", ["120", 120.0, "120.0"], ids=["str", "float", "str_float"]) +async def test_numeric_expires_in_is_honored(expires_in): + # A JSON int/float/numeric-string expires_in must drive the TTL, not fall back to the default. + clock = _Clock(1000.0) + post = _RecordingPost({"access_token": "x", "expires_in": expires_in}) + exchanger = Rfc8693TokenExchanger(post, clock=clock) # ttl = max(120-60, 10) = 60 -> until 1060 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1061.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +async def test_short_lived_token_is_not_cached_past_its_expiry(): + # expires_in (5s) below the buffer/min floor must NOT be served stale: cache only until expiry. + clock = _Clock(1000.0) + post = _RecordingPost({"access_token": "x", "expires_in": 5}) + exchanger = Rfc8693TokenExchanger(post, clock=clock) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1004.0 # within the 5s lifetime -> still cached + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 1 + clock.now = 1006.0 # past expiry -> re-exchange, not a stale bearer + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"access_token": "x"}, + {"access_token": "x", "expires_in": "not-a-number"}, + {"access_token": "x", "expires_in": True}, + ], + ids=["missing", "unparseable", "bool"], +) +async def test_unusable_expires_in_falls_back_to_default_ttl(body): + clock = _Clock(1000.0) + post = _RecordingPost(body) + exchanger = Rfc8693TokenExchanger(post, clock=clock, default_ttl_seconds=300.0) + await exchanger.exchange("jwt", _SERVER, _CONFIG) + clock.now = 1299.0 + await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert len(post.calls) == 1 + + +@pytest.mark.asyncio +async def test_distributed_coordinator_refresh_and_reread_use_the_cache(): + # Mimics the cross-replica coordinator contract: the winner's refresh populates the cache, a + # re-entrant refresh sees the fresh entry, and a loser reads it back via reread, all without a + # second IdP call. + class _ReplayCoordinator: + async def run(self, user_id, server_id, refresh, reread): + first = await refresh() + second = await refresh() + via_reread = await reread() + assert first is not None and second is not None and via_reread is not None + return via_reread + + post = _RecordingPost({"access_token": "x", "expires_in": 3600}) + exchanger = Rfc8693TokenExchanger(post, coordinator=_ReplayCoordinator(), clock=_Clock()) + result = await exchanger.exchange("jwt", _SERVER, _CONFIG) + assert isinstance(result, Ok) and result.ok.access_token == "x" + assert len(post.calls) == 1 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 3e08d6f35e6..d8167f17ef9 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 @@ -96,9 +96,7 @@ async def test_authorize_endpoint_includes_response_type(): mock_request.headers = {} # Mock the encryption functions to avoid needing a signing key - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint @@ -160,9 +158,7 @@ async def test_authorize_endpoint_preserves_existing_query_params(): mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" response = await authorize( @@ -176,9 +172,7 @@ async def test_authorize_endpoint_preserves_existing_query_params(): location = response.headers["location"] # Must NOT have double '?' — existing params must be merged correctly - assert ( - location.count("?") == 1 - ), f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" + assert location.count("?") == 1, f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" assert "tenant=system" in location assert "client_id=test_client_id" in location assert "response_type=code" in location @@ -228,9 +222,7 @@ async def test_authorize_endpoint_forwards_pkce_parameters(): mock_request.headers = {} # Mock the encryption function - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" # Call authorize endpoint with PKCE parameters @@ -338,10 +330,7 @@ async def test_token_endpoint_forwards_code_verifier(): # Check the data parameter includes code_verifier assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" assert call_args[1]["data"]["code"] == "4/test_authorization_code" - assert ( - call_args[1]["data"]["client_id"] - == "669428968603-test.apps.googleusercontent.com" - ) + assert call_args[1]["data"]["client_id"] == "669428968603-test.apps.googleusercontent.com" assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" assert call_args[1]["data"]["grant_type"] == "authorization_code" @@ -428,9 +417,7 @@ async def test_register_client_returns_existing_server_credentials(): "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): - result = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) + result = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) finally: global_mcp_server_manager.registry.clear() @@ -505,9 +492,7 @@ async def test_register_client_remote_registration_success(): return_value=mock_async_client, ), ): - response = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) + response = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) finally: global_mcp_server_manager.registry.clear() @@ -524,14 +509,9 @@ async def test_register_client_remote_registration_success(): "Content-Type": "application/json", "Accept": "application/json", } - assert call_args.kwargs["json"]["redirect_uris"] == [ - "https://proxy.litellm.example/callback" - ] + assert call_args.kwargs["json"]["redirect_uris"] == ["https://proxy.litellm.example/callback"] assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] - assert ( - call_args.kwargs["json"]["token_endpoint_auth_method"] - == request_payload["token_endpoint_auth_method"] - ) + assert call_args.kwargs["json"]["token_endpoint_auth_method"] == request_payload["token_endpoint_auth_method"] @pytest.mark.asyncio @@ -1120,9 +1100,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint @@ -1217,10 +1195,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): # Verify that the redirect_uri sent to the provider uses HTTPS call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://litellm-proxy.example.com/callback" - ) + assert call_args[1]["data"]["redirect_uri"] == "https://litellm-proxy.example.com/callback" @pytest.mark.asyncio @@ -1272,9 +1247,7 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): ) # Verify response uses HTTPS URLs - assert response["authorization_servers"][0].startswith( - "https://litellm.example.com/" - ) + assert response["authorization_servers"][0].startswith("https://litellm.example.com/") assert response["scopes_supported"] == oauth2_server.scopes @@ -1421,9 +1394,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): } # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint @@ -1440,8 +1411,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): # The redirect_uri parameter should use the external URL assert ( - "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" - in location + "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" in location or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location ) @@ -1522,10 +1492,7 @@ async def test_token_endpoint_respects_x_forwarded_host(): # Verify that the redirect_uri sent to the provider uses the external URL call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://proxy.example.com/github/mcp/callback" - ) + assert call_args[1]["data"]["redirect_uri"] == "https://proxy.example.com/github/mcp/callback" @pytest.mark.parametrize( @@ -1733,9 +1700,7 @@ def mock_get(header_name, default=None): ), ], ) -def test_get_request_base_url_xff_trust_gate( - general_settings, direct_ip, expect_xff_honoured -): +def test_get_request_base_url_xff_trust_gate(general_settings, direct_ip, expect_xff_honoured): """Verify the X-Forwarded-* trust gate. With XFF poisoning attempted, the helper must return either the literal @@ -1813,12 +1778,10 @@ def test_xff_misconfig_warning_emitted_once(caplog): for _ in range(3): get_request_base_url(mock_request) - matching = [ - rec for rec in caplog.records if "mcp_trusted_proxy_ranges" in rec.getMessage() - ] - assert ( - len(matching) == 1 - ), f"expected exactly one warning, got {len(matching)}: {[r.getMessage() for r in matching]}" + matching = [rec for rec in caplog.records if "mcp_trusted_proxy_ranges" in rec.getMessage()] + assert len(matching) == 1, ( + f"expected exactly one warning, got {len(matching)}: {[r.getMessage() for r in matching]}" + ) def test_get_request_base_url_honors_proxy_base_url_env(monkeypatch): @@ -1849,9 +1812,7 @@ def test_get_request_base_url_honors_proxy_base_url_env(monkeypatch): assert get_request_base_url(mock_request) == "https://litellm.example.com" -def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection( - caplog, monkeypatch -): +def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection(caplog, monkeypatch): try: from fastapi import HTTPException, Request @@ -1898,8 +1859,7 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection( matching = [r for r in caplog.records if "rejecting redirect_uri" in r.getMessage()] assert len(matching) == 1, ( - "expected exactly one diagnostic warning, got " - f"{[r.getMessage() for r in caplog.records]}" + f"expected exactly one diagnostic warning, got {[r.getMessage() for r in caplog.records]}" ) msg = matching[0].getMessage() assert "https://litellm.example.com/ui/mcp/oauth/callback" in msg @@ -1918,9 +1878,7 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection( "not a url at all", ], ) -def test_get_request_base_url_rejects_malformed_proxy_base_url( - bad_value, monkeypatch, caplog -): +def test_get_request_base_url_rejects_malformed_proxy_base_url(bad_value, monkeypatch, caplog): try: from fastapi import Request @@ -1950,26 +1908,16 @@ def test_get_request_base_url_rejects_malformed_proxy_base_url( result = get_request_base_url(mock_request) assert result == "http://litellm-internal:4000", ( - f"malformed PROXY_BASE_URL={bad_value!r} should be ignored, " f"got {result!r}" + f"malformed PROXY_BASE_URL={bad_value!r} should be ignored, got {result!r}" ) - matching = [ - r - for r in caplog.records - if "PROXY_BASE_URL" in r.getMessage() and "ignored" in r.getMessage() - ] + matching = [r for r in caplog.records if "PROXY_BASE_URL" in r.getMessage() and "ignored" in r.getMessage()] assert len(matching) == 1, ( - "expected one diagnostic for malformed PROXY_BASE_URL, got " - f"{[r.getMessage() for r in caplog.records]}" - ) - assert ( - repr(bad_value) in matching[0].getMessage() - or bad_value in matching[0].getMessage() + f"expected one diagnostic for malformed PROXY_BASE_URL, got {[r.getMessage() for r in caplog.records]}" ) + assert repr(bad_value) in matching[0].getMessage() or bad_value in matching[0].getMessage() -def test_get_request_base_url_malformed_proxy_base_url_warning_is_one_shot( - monkeypatch, caplog -): +def test_get_request_base_url_malformed_proxy_base_url_warning_is_one_shot(monkeypatch, caplog): try: from fastapi import Request @@ -1999,14 +1947,8 @@ def test_get_request_base_url_malformed_proxy_base_url_warning_is_one_shot( for _ in range(5): get_request_base_url(mock_request) - matching = [ - r - for r in caplog.records - if "PROXY_BASE_URL" in r.getMessage() and "ignored" in r.getMessage() - ] - assert ( - len(matching) == 1 - ), f"expected exactly one warning across 5 calls, got {len(matching)}" + matching = [r for r in caplog.records if "PROXY_BASE_URL" in r.getMessage() and "ignored" in r.getMessage()] + assert len(matching) == 1, f"expected exactly one warning across 5 calls, got {len(matching)}" # ------------------------------------------------------------------- @@ -2219,12 +2161,8 @@ async def test_authorize_root_fails_with_multiple_oauth2_servers(): pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() - server1 = _create_oauth2_server( - server_id="server1", name="server1", server_name="server1", alias="server1" - ) - server2 = _create_oauth2_server( - server_id="server2", name="server2", server_name="server2", alias="server2" - ) + server1 = _create_oauth2_server(server_id="server1", name="server1", server_name="server1", alias="server1") + server2 = _create_oauth2_server(server_id="server2", name="server2", server_name="server2", alias="server2") global_mcp_server_manager.registry[server1.server_id] = server1 global_mcp_server_manager.registry[server2.server_id] = server2 @@ -2582,9 +2520,7 @@ async def test_oauth_callback_redirects_with_state(): "client_redirect_uri": "http://localhost:3000/ui/mcp/oauth/callback", } - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = mock_state_data # Call callback endpoint with code and state @@ -2596,10 +2532,7 @@ async def test_oauth_callback_redirects_with_state(): # Should redirect to the client callback URL with code and original state assert response.status_code == 302 - assert ( - "http://localhost:3000/ui/mcp/oauth/callback" - in response.headers["location"] - ) + assert "http://localhost:3000/ui/mcp/oauth/callback" in response.headers["location"] assert "code=test_authorization_code_12345" in response.headers["location"] assert "state=test-uuid-state-123" in response.headers["location"] @@ -2619,17 +2552,13 @@ async def test_oauth_callback_preserves_client_redirect_uri_query(): except ImportError: pytest.skip("MCP discoverable endpoints not available") - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = { "base_url": "http://localhost:3000/ui/mcp/oauth/callback", "original_state": "test-uuid-state-123", "code_challenge": "test_challenge", "code_challenge_method": "S256", - "client_redirect_uri": ( - "http://localhost:3000/ui/mcp/oauth/callback?session=abc" - ), + "client_redirect_uri": ("http://localhost:3000/ui/mcp/oauth/callback?session=abc"), } response = await callback( @@ -2657,9 +2586,7 @@ async def test_oauth_callback_handles_invalid_state(): pytest.skip("MCP discoverable endpoints not available") # Mock state decoding to raise an exception - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.side_effect = Exception("Failed to decrypt state") # Call callback endpoint with invalid state @@ -2682,9 +2609,7 @@ async def test_oauth_callback_accepts_same_origin_ui_redirect(): callback, ) - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = { "base_url": "https://proxy.example.com/ui/mcp/oauth/callback", "original_state": "state-123", @@ -2700,10 +2625,7 @@ async def test_oauth_callback_accepts_same_origin_ui_redirect(): ) assert response.status_code == 302 - assert ( - "https://proxy.example.com/ui/mcp/oauth/callback" - in response.headers["location"] - ) + assert "https://proxy.example.com/ui/mcp/oauth/callback" in response.headers["location"] assert "code=auth-code-123" in response.headers["location"] assert "state=state-123" in response.headers["location"] @@ -2740,9 +2662,7 @@ async def test_oauth_authorize_includes_scopes_from_server_config(): mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "encrypted_state" # Call authorize without explicit scope parameter @@ -2762,8 +2682,7 @@ async def test_oauth_authorize_includes_scopes_from_server_config(): assert response.status_code in (307, 302) redirect_url = response.headers["location"] assert ( - "scope=api+read_user+ai_workflows" in redirect_url - or "scope=api%20read_user%20ai_workflows" in redirect_url + "scope=api+read_user+ai_workflows" in redirect_url or "scope=api%20read_user%20ai_workflows" in redirect_url ) @@ -2798,9 +2717,7 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config(): mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: mock_encrypt.return_value = "encrypted_state" # Call authorize WITH explicit scope parameter @@ -2820,8 +2737,7 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config(): assert response.status_code in (307, 302) redirect_url = response.headers["location"] assert ( - "scope=custom_scope1+custom_scope2" in redirect_url - or "scope=custom_scope1%20custom_scope2" in redirect_url + "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url ) assert "default_scope" not in redirect_url @@ -3078,9 +2994,7 @@ async def test_callback_revalidates_loopback_on_decoded_base_url(): callback, ) - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = { "base_url": "https://attacker.example.com/cb", "original_state": "s", @@ -3104,9 +3018,7 @@ async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri(): callback, ) - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = { "base_url": "http://localhost:3000/cb", "original_state": "s", @@ -3130,9 +3042,7 @@ async def test_callback_rejects_state_missing_redirect_uri(): callback, ) - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: mock_decode.return_value = { "original_state": "s", "code_challenge": None, @@ -3261,9 +3171,7 @@ async def test_token_exchange_omits_expires_in_when_upstream_omits_it(): rotation) returns no ``expires_in``. The exchange must mirror that and omit ``expires_in`` rather than fabricate a 1-hour TTL, so the stored credential is treated as non-expiring instead of dying after an hour.""" - body = await _exchange_with_upstream_token_response( - {"access_token": "tok", "token_type": "Bearer"} - ) + body = await _exchange_with_upstream_token_response({"access_token": "tok", "token_type": "Bearer"}) assert "expires_in" not in body @@ -3277,6 +3185,120 @@ async def test_token_exchange_passes_through_upstream_expires_in(): assert body["expires_in"] == 43200 +# ------------------------------------------------------------------- +# OBO (token_exchange) Protected Resource Metadata: discovery must name the +# JWT-auth issuer the client SSOs with, not the gateway. +# ------------------------------------------------------------------- + +_OBO_RESOURCE = "https://litellm.example.com/mcp/obo_mcp" +_PATCH_ISSUERS = "litellm.proxy._experimental.mcp_server.discoverable_endpoints._jwt_auth_issuers" + + +def _obo_server(scopes=None): + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id="obo_mcp", + name="obo_mcp", + server_name="obo_mcp", + alias="obo_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + scopes=scopes, + ) + + +def test_obo_protected_resource_response_names_jwt_issuers(): + """An OBO server's PRM points authorization_servers at the configured JWT issuers (the IdP that + mints and validates the subject token), with the gateway resource echoed back.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _obo_protected_resource_response, + ) + + with patch(_PATCH_ISSUERS, return_value=["https://idp.example.com"]): + response = _obo_protected_resource_response(_obo_server(scopes=["read"]), _OBO_RESOURCE) + assert response == { + "authorization_servers": ["https://idp.example.com"], + "resource": _OBO_RESOURCE, + "scopes_supported": ["read"], + } + + +def test_obo_protected_resource_response_scopes_default_empty(): + """A scopeless OBO server reports scopes_supported as [] rather than None.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _obo_protected_resource_response, + ) + + with patch(_PATCH_ISSUERS, return_value=["https://idp.example.com"]): + response = _obo_protected_resource_response(_obo_server(scopes=None), _OBO_RESOURCE) + assert response["scopes_supported"] == [] + + +def test_obo_protected_resource_response_falls_back_when_no_issuer(): + """With no JWT issuer configured, the OBO branch returns None so the caller falls back to the + gateway-default PRM (discovery still works, it just can't name the IdP).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _obo_protected_resource_response, + ) + + with patch(_PATCH_ISSUERS, return_value=[]): + assert _obo_protected_resource_response(_obo_server(), _OBO_RESOURCE) is None + + +def test_obo_protected_resource_response_ignores_non_obo_server(): + """Non-OBO servers are not handled by this branch (returns None -> gateway default).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _obo_protected_resource_response, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + oauth2_server = MCPServer( + server_id="oauth2_mcp", + name="oauth2_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + assert _obo_protected_resource_response(oauth2_server, _OBO_RESOURCE) is None + + +@pytest.mark.asyncio +async def test_build_oauth_protected_resource_response_obo_end_to_end(): + """End to end through the response builder: an OBO server's PRM advertises the JWT issuer as + authorization_servers, proving the extracted branch is wired into the public discovery path.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry["obo_mcp"] = _obo_server(scopes=["read"]) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + try: + with patch(_PATCH_ISSUERS, return_value=["https://idp.example.com"]): + response = await _build_oauth_protected_resource_response( + request=mock_request, + mcp_server_name="obo_mcp", + use_standard_pattern=True, + ) + assert response["authorization_servers"] == ["https://idp.example.com"] + assert response["resource"] == "https://litellm.example.com/mcp/obo_mcp" + finally: + global_mcp_server_manager.registry.clear() + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e790dea0506..d006e00436c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,12 +5,14 @@ import os import sys from datetime import datetime -from typing import Any, Dict +from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../../../") @@ -47,18 +49,14 @@ def _reload_mcp_manager_module(): utils_module = sys.modules["litellm.proxy._experimental.mcp_server.utils"] - manager_module = sys.modules[ - "litellm.proxy._experimental.mcp_server.mcp_server_manager" - ] + manager_module = sys.modules["litellm.proxy._experimental.mcp_server.mcp_server_manager"] importlib.reload(utils_module) reloaded = importlib.reload(manager_module) # After reload, server.py still holds a stale reference to the old # global_mcp_server_manager. Update it so tests that exercise server.py # functions (e.g. _get_tools_from_mcp_servers) use the fresh instance. server_module = sys.modules.get("litellm.proxy._experimental.mcp_server.server") - if server_module is not None and hasattr( - server_module, "global_mcp_server_manager" - ): + if server_module is not None and hasattr(server_module, "global_mcp_server_manager"): server_module.global_mcp_server_manager = reloaded.global_mcp_server_manager return reloaded @@ -168,9 +166,7 @@ async def resolve_credentials(self, subject, server): auth_type=MCPAuth.oauth2, # oauth2 + no client creds + not delegate -> authorization_code ) - client = await manager._create_mcp_client( - server, mcp_auth_header="Bearer caller-supplied-token" - ) + client = await manager._create_mcp_client(server, mcp_auth_header="Bearer caller-supplied-token") # the v2 resolver ran (the caller override did NOT defer to v1); the stored token wins assert calls == [("", "authz-srv")] @@ -441,9 +437,7 @@ async def test_list_tools_with_server_specific_auth_headers(self): # Mock get_allowed_mcp_servers to return our test servers manager.get_allowed_mcp_servers = AsyncMock(return_value=["github", "zapier"]) - manager.get_mcp_server_by_id = MagicMock( - side_effect=lambda x: server1 if x == "github" else server2 - ) + manager.get_mcp_server_by_id = MagicMock(side_effect=lambda x: server1 if x == "github" else server2) # Mock _get_tools_from_server to return different results async def mock_get_tools_from_server( @@ -470,9 +464,7 @@ async def mock_get_tools_from_server( "zapier": "zapier-api-key", } - result = await manager.list_tools( - mcp_server_auth_headers=mcp_server_auth_headers - ) + result = await manager.list_tools(mcp_server_auth_headers=mcp_server_auth_headers) # Verify that both servers were called with their specific auth headers assert len(result) == 3 # 2 from github + 1 from zapier @@ -541,9 +533,7 @@ async def mock_get_tools_from_server( mcp_auth_header=None, **kwargs, ): - assert ( - mcp_auth_header == "server-specific-token" - ) # Should use server-specific header + assert mcp_auth_header == "server-specific-token" # Should use server-specific header tool = MagicMock() tool.name = "github_tool_1" return [tool] @@ -574,9 +564,7 @@ async def test_call_regular_mcp_tool_case_insensitive_extra_headers(self): ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock( - return_value=CallToolResult(content=[], isError=False) - ) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -603,6 +591,332 @@ async def capture_create_mcp_client( assert captured_extra_headers == {"Authorization": "Bearer token"} assert isinstance(result, CallToolResult) + async def _capture_list_subject_token(self, server, oauth2_headers, raw_headers=None): + """Run _get_tools_from_server and return the subject_token it threaded to _create_mcp_client.""" + manager = MCPServerManager() + captured = {} + + async def capture_create_mcp_client( + server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs + ): # pragma: no cover - helper + captured["subject_token"] = subject_token + return AsyncMock() + + manager._create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + await manager._get_tools_from_server(server=server, oauth2_headers=oauth2_headers, raw_headers=raw_headers) + return captured["subject_token"] + + @pytest.mark.asyncio + async def test_list_threads_subject_token_for_token_exchange(self): + """tools/list discovery must hand the caller's bearer to the resolver for OBO servers.""" + server = MCPServer( + server_id="te-list", + name="te-list-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + subject_token = await self._capture_list_subject_token( + server, oauth2_headers={"Authorization": "Bearer subj-jwt"} + ) + assert subject_token == "subj-jwt" + + @pytest.mark.asyncio + async def test_list_does_not_thread_subject_token_for_non_token_exchange(self): + """A non-OBO server must not get the caller's bearer threaded (no leak across modes).""" + server = MCPServer( + server_id="none-list", + name="none-list-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + subject_token = await self._capture_list_subject_token( + server, oauth2_headers={"Authorization": "Bearer subj-jwt"} + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_list_subject_token_none_without_oauth2_headers(self): + """Background/registry refresh (no oauth2 headers) lists with no subject token, as before.""" + server = MCPServer( + server_id="te-list-bg", + name="te-list-bg-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + subject_token = await self._capture_list_subject_token(server, oauth2_headers=None) + assert subject_token is None + + @pytest.mark.asyncio + async def test_list_surfaces_resolver_401_as_upstream_auth_error(self): + """A v2 resolver auth challenge (HTTPException 401) raised while building the client must + surface as MCPUpstreamAuthError with its WWW-Authenticate preserved, so single-server routes + challenge the client instead of the old behavior of masking it to an empty tool list.""" + server = MCPServer( + server_id="te-401", + name="te-401-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + manager = MCPServerManager() + challenge = ( + 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/te-401-server", error="invalid_token"' + ) + manager._create_mcp_client = AsyncMock( + side_effect=HTTPException(status_code=401, detail="Unauthorized", headers={"WWW-Authenticate": challenge}) + ) + with pytest.raises(MCPUpstreamAuthError) as exc_info: + await manager._get_tools_from_server(server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"}) + assert exc_info.value.status_code == 401 + assert exc_info.value.www_authenticate == challenge + + @pytest.mark.asyncio + async def test_list_absorbs_non_auth_httpexception(self): + """A non-auth HTTP error (e.g. 412 no endpoint, 503 IdP down) must stay absorbed to [] so one + misconfigured/unavailable server does not blank the whole aggregate listing.""" + server = MCPServer( + server_id="te-412", + name="te-412-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + manager = MCPServerManager() + manager._create_mcp_client = AsyncMock( + side_effect=HTTPException(status_code=412, detail="token exchange endpoint is not configured") + ) + result = await manager._get_tools_from_server( + server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"} + ) + assert result == [] + + def _token_exchange_server(self, server_id: str) -> "MCPServer": + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + + async def _capture_subject_token(self, call) -> Optional[str]: + """Run a manager method (via ``call(manager)``) and return the subject_token it threaded + into ``_create_mcp_client``.""" + manager = MCPServerManager() + captured: Dict[str, Any] = {} + + async def capture_create_mcp_client( + server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs + ): # pragma: no cover - helper + captured["subject_token"] = subject_token + return AsyncMock() + + manager._create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) + await call(manager) + return captured.get("subject_token") + + @pytest.mark.asyncio + async def test_prompts_thread_subject_token_for_token_exchange(self): + """prompts/list on an OBO server must exchange the caller's bearer, not connect with none.""" + server = self._token_exchange_server("te-prompts") + st = await self._capture_subject_token( + lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + ) + assert st == "subj-jwt" + + @pytest.mark.asyncio + async def test_resources_thread_subject_token_for_token_exchange(self): + """resources/list on an OBO server must exchange the caller's bearer.""" + server = self._token_exchange_server("te-resources") + st = await self._capture_subject_token( + lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + ) + assert st == "subj-jwt" + + @pytest.mark.asyncio + async def test_read_resource_threads_subject_token_for_token_exchange(self): + """resources/read on an OBO server must exchange the caller's bearer.""" + server = self._token_exchange_server("te-read") + st = await self._capture_subject_token( + lambda m: m.read_resource_from_server( + server=server, + url="https://up.example.com/r", + raw_headers={"authorization": "Bearer subj-jwt"}, + ) + ) + assert st == "subj-jwt" + + @pytest.mark.asyncio + async def test_prompts_no_subject_token_for_non_token_exchange(self): + """A non-OBO server must not get the caller's bearer threaded (no cross-mode leak).""" + server = MCPServer( + server_id="none-prompts", + name="none-prompts-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + st = await self._capture_subject_token( + lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + ) + assert st is None + + @pytest.mark.asyncio + async def test_caller_header_cannot_bypass_v2_for_token_exchange(self): + """A caller-supplied per-server header (x-mcp-*) must NOT disable the OBO exchange: + _create_mcp_client keeps the v2 spec and runs the resolver (which exchanges the subject), + rather than deferring to v1 and forwarding the caller's header verbatim upstream.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + exchanged_subjects = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + exchanged_subjects.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-bypass") + + client = await manager._create_mcp_client( + server, + mcp_auth_header="Bearer x-mcp-caller-header", + subject_token="subj-jwt", + ) + + # The resolver ran and exchanged the subject despite the per-server header; not bypassed. + assert exchanged_subjects == ["subj-jwt"] + assert client is not None + + @pytest.mark.asyncio + async def test_injected_authorization_does_not_shadow_obo_minted_token(self): + """A guardrail/static Authorization (e.g. MCPJWTSigner) must NOT shadow the exchanged OBO + token. The resolver-owned credential is authoritative: the conflicting header is dropped and + the minted token is what reaches the upstream, not the injected JWT.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-shadow") + + client = await manager._create_mcp_client( + server, + extra_headers={"Authorization": "Bearer signer-jwt"}, # simulate the JWT signer + subject_token="subj-jwt", + ) + + # minted token wins (resolved_auth kept), signer's header dropped from extra_headers + assert client._resolved_auth is not None + assert "authorization" not in {k.lower() for k in (client.extra_headers or {})} + + @pytest.mark.asyncio + async def test_preflight_token_exchange_challenges_on_rejected_subject(self): + """A subject the IdP rejects must raise the RFC 9728 401 challenge from the preflight, so a + single-server route fails observably at the transport edge instead of the old behavior of + the session opening and list_tools masking the failed exchange as an empty tool list.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_unauthorized("subject token rejected by the IdP")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-preflight-401") + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer rejected-subject"}, + user_api_key_auth=None, + ) + assert exc_info.value.status_code == 401 + headers = exc_info.value.headers or {} + www_authenticate = headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "" + assert "resource_metadata" in www_authenticate + + @pytest.mark.asyncio + async def test_preflight_token_exchange_maps_gateway_fault_to_public_status(self): + """A gateway-fault CredError (e.g. invalid_client) must surface its public status (500) + from the preflight, not the OBO 401 challenge and not an empty-success session.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_misconfigured("token exchange configuration error")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-preflight-500") + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer subj"}, + user_api_key_auth=None, + ) + assert exc_info.value.status_code == 500 + + @pytest.mark.asyncio + async def test_preflight_token_exchange_noop_on_success_and_without_subject(self): + """A successful exchange returns without raising, and a request with no bearer never + reaches the resolver (the no-subject case is the existing preemptive challenge's job).""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + resolved = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-preflight-ok") + + assert ( + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer good-subject"}, + user_api_key_auth=None, + ) + is None + ) + assert resolved == ["good-subject"] + + await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) + assert resolved == ["good-subject"] + @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key( self, @@ -624,9 +938,7 @@ async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admis ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock( - return_value=CallToolResult(content=[], isError=False) - ) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -682,17 +994,10 @@ async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization( ) # Migrated authorization_code => the centralized strip decision says drop the # caller's Authorization (the v2 resolver injects the stored token). - assert ( - _should_strip_caller_authorization( - mcp_server=server, raw_headers=None, user_api_key_auth=None - ) - is True - ) + assert _should_strip_caller_authorization(mcp_server=server, raw_headers=None, user_api_key_auth=None) is True mock_client = AsyncMock() - mock_client.call_tool = AsyncMock( - return_value=CallToolResult(content=[], isError=False) - ) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = "unset" async def capture_create_mcp_client( @@ -735,9 +1040,7 @@ def test_without_authorization_drops_only_the_credential(self): # Only Authorization present -> nothing left -> None (case-insensitive) assert _without_authorization({"authorization": "Bearer x"}) is None # Authorization dropped, other headers kept - assert _without_authorization( - {"Authorization": "Bearer x", "X-Trace-Id": "t"} - ) == {"X-Trace-Id": "t"} + assert _without_authorization({"Authorization": "Bearer x", "X-Trace-Id": "t"}) == {"X-Trace-Id": "t"} @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_forwards_authorization_with_admission_header( @@ -760,9 +1063,7 @@ async def test_call_regular_mcp_tool_passthrough_forwards_authorization_with_adm ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock( - return_value=CallToolResult(content=[], isError=False) - ) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -795,9 +1096,7 @@ async def capture_create_mcp_client( user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), ) - assert captured_extra_headers == { - "Authorization": "Bearer upstream-oauth-bearer" - } + assert captured_extra_headers == {"Authorization": "Bearer upstream-oauth-bearer"} @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_forwards_authorization_for_anonymous_admission( @@ -821,9 +1120,7 @@ async def test_call_regular_mcp_tool_passthrough_forwards_authorization_for_anon ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock( - return_value=CallToolResult(content=[], isError=False) - ) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -853,9 +1150,7 @@ async def capture_create_mcp_client( user_api_key_auth=UserAPIKeyAuth(api_key=None), ) - assert captured_extra_headers == { - "Authorization": "Bearer upstream-oauth-bearer" - } + assert captured_extra_headers == {"Authorization": "Bearer upstream-oauth-bearer"} @pytest.mark.asyncio async def test_get_prompts_from_server_success(self): @@ -944,9 +1239,7 @@ async def test_get_resources_from_server_success(self): mock_client = AsyncMock() mock_resources = [Resource(name="file", uri="https://example.com/file")] mock_client.list_resources = AsyncMock(return_value=mock_resources) - prefixed_resources = [ - Resource(name="alias-server-file", uri="https://example.com/file") - ] + prefixed_resources = [Resource(name="alias-server-file", uri="https://example.com/file")] with ( patch.object( @@ -1030,6 +1323,7 @@ async def test_get_resource_templates_from_server_success(self): mcp_auth_header="auth", extra_headers=None, stdio_env=None, + subject_token=None, ) mock_client.list_resource_templates.assert_awaited_once() mock_prefix.assert_called_once_with(mock_templates, server, add_prefix=False) @@ -1077,9 +1371,7 @@ async def test_read_resource_from_server_success(self): mock_create_client.assert_called_once() called_kwargs = mock_create_client.call_args.kwargs assert called_kwargs["extra_headers"] == {"X-Test": "1", "X-Static": "1"} - mock_client.read_resource.assert_awaited_once_with( - "https://example.com/resource" - ) + mock_client.read_resource.assert_awaited_once_with("https://example.com/resource") assert result is read_result @pytest.mark.asyncio @@ -1184,9 +1476,7 @@ def build_response(url: str, **kwargs): request = httpx.Request("GET", url) response_obj = httpx.Response(status_code=404, request=request) mock_response.raise_for_status = MagicMock( - side_effect=httpx.HTTPStatusError( - "not found", request=request, response=response_obj - ) + side_effect=httpx.HTTPStatusError("not found", request=request, response=response_obj) ) return mock_response @@ -1200,19 +1490,11 @@ def build_response(url: str, **kwargs): # The Azure issuer is cross-origin against the server_url — use # the issuer itself as server_url so the test exercises the # well-known fetch logic without needing real DNS. - result = await manager._fetch_single_authorization_server_metadata( - issuer, issuer - ) + result = await manager._fetch_single_authorization_server_metadata(issuer, issuer) assert result is not None - assert ( - result.authorization_url - == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" - ) - assert ( - result.token_url - == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" - ) + assert result.authorization_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" assert result.scopes == ["api://some-scope/.default"] @pytest.mark.asyncio @@ -1226,9 +1508,7 @@ async def test_fetch_single_authorization_server_metadata_derives_azure_metadata response_obj = httpx.Response(status_code=404, request=request) mock_response = MagicMock() mock_response.raise_for_status = MagicMock( - side_effect=httpx.HTTPStatusError( - "not found", request=request, response=response_obj - ) + side_effect=httpx.HTTPStatusError("not found", request=request, response=response_obj) ) mock_client = MagicMock() @@ -1238,19 +1518,11 @@ async def test_fetch_single_authorization_server_metadata_derives_azure_metadata "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", return_value=mock_client, ): - result = await manager._fetch_single_authorization_server_metadata( - issuer, issuer - ) + result = await manager._fetch_single_authorization_server_metadata(issuer, issuer) assert result is not None - assert ( - result.authorization_url - == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" - ) - assert ( - result.token_url - == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" - ) + assert result.authorization_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + assert result.token_url == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" @pytest.mark.asyncio async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self): @@ -1265,9 +1537,7 @@ async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self ) def raise_http_error(): - raise httpx.HTTPStatusError( - "unauthorized", request=request, response=response_obj - ) + raise httpx.HTTPStatusError("unauthorized", request=request, response=response_obj) response_obj.raise_for_status = MagicMock(side_effect=raise_http_error) @@ -1319,8 +1589,10 @@ async def test_load_servers_from_config_overrides_discovery_metadata(self): registration_url="https://discovered.example.com/register", ) - async def fake_discovery(server_url: str): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): assert server_url == "https://example.com/mcp" + # oauth2 (browser flow) keeps the origin fallback; only OBO disables it. + assert allow_origin_fallback is True return discovered_metadata manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] @@ -1384,9 +1656,7 @@ async def mock_get_tools_from_server( mcp_auth_header=None, **kwargs, ): - assert ( - mcp_auth_header == "server-specific-token" - ) # Should use server-specific header via server_name + assert mcp_auth_header == "server-specific-token" # Should use server-specific header via server_name tool = MagicMock() tool.name = "github_tool_1" return [tool] @@ -1453,9 +1723,7 @@ async def test_health_check_server_unhealthy(self): # Mock failed client.run_with_session mock_client = AsyncMock() - mock_client.run_with_session = AsyncMock( - side_effect=Exception("Connection timeout") - ) + mock_client.run_with_session = AsyncMock(side_effect=Exception("Connection timeout")) manager._create_mcp_client = AsyncMock(return_value=mock_client) # Perform health check @@ -1577,9 +1845,7 @@ async def test_health_check_server_with_static_headers(self): # Capture the extra_headers passed to _create_mcp_client captured_extra_headers = None - async def capture_create_mcp_client( - server, mcp_auth_header, extra_headers, stdio_env - ): + async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env): nonlocal captured_extra_headers captured_extra_headers = extra_headers return mock_client @@ -1921,9 +2187,7 @@ async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) @@ -1967,13 +2231,8 @@ async def test_pre_call_tool_check_allowed_tools_list_blocks_tool(self): ) assert exc_info.value.status_code == 403 - assert ( - "Tool blocked_tool is not allowed for server test-server" - in exc_info.value.detail["error"] - ) - assert ( - "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] - ) + assert "Tool blocked_tool is not allowed for server test-server" in exc_info.value.detail["error"] + assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] @pytest.mark.asyncio async def test_pre_call_tool_check_disallowed_tools_list_allows_tool(self): @@ -1997,9 +2256,7 @@ async def test_pre_call_tool_check_disallowed_tools_list_allows_tool(self): proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) @@ -2043,13 +2300,8 @@ async def test_pre_call_tool_check_disallowed_tools_list_blocks_tool(self): ) assert exc_info.value.status_code == 403 - assert ( - "Tool banned_tool is not allowed for server test-server" - in exc_info.value.detail["error"] - ) - assert ( - "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] - ) + assert "Tool banned_tool is not allowed for server test-server" in exc_info.value.detail["error"] + assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] @pytest.mark.asyncio async def test_pre_call_tool_check_no_restrictions_allows_any_tool(self): @@ -2073,9 +2325,7 @@ async def test_pre_call_tool_check_no_restrictions_allows_any_tool(self): proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) @@ -2112,9 +2362,7 @@ async def test_pre_call_tool_check_allowed_tools_takes_precedence(self): proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) @@ -2140,10 +2388,7 @@ async def test_pre_call_tool_check_allowed_tools_takes_precedence(self): ) assert exc_info.value.status_code == 403 - assert ( - "Tool tool3 is not allowed for server test-server" - in exc_info.value.detail["error"] - ) + assert "Tool tool3 is not allowed for server test-server" in exc_info.value.detail["error"] async def test_get_tools_from_server_add_prefix(self): """Verify _get_tools_from_server respects add_prefix True/False.""" @@ -2174,9 +2419,7 @@ async def test_get_tools_from_server_add_prefix(self): assert tools_prefixed[0].name == "zapier-send_email" # Case 2: add_prefix=False (single-server) -> expect unprefixed - tools_unprefixed = await manager._get_tools_from_server( - server, add_prefix=False - ) + tools_unprefixed = await manager._get_tools_from_server(server, add_prefix=False) assert len(tools_unprefixed) == 1 assert tools_unprefixed[0].name == "send_email" @@ -2254,9 +2497,7 @@ def test_resolve_mcp_server_for_tool_call_via_alias(self): manager.registry = {"srv-uuid-123": server} manager.tool_name_to_mcp_server_name_mapping["create_zap"] = "zapier" - resolved = manager._resolve_mcp_server_for_tool_call( - "zapier-alias", "create_zap" - ) + resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") assert resolved is server def test_resolve_mcp_server_for_tool_call_unknown_tool_with_empty_mapping(self): @@ -2537,13 +2778,9 @@ def test_create_prefixed_tools_updates_mapping_for_both_forms(self): # Mapping should include both original and prefixed names -> resolves calls either way assert manager.tool_name_to_mcp_server_name_mapping["create_issue"] == "jira" - assert ( - manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira" - ) + assert manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira" assert manager.tool_name_to_mcp_server_name_mapping["close_issue"] == "jira" - assert ( - manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira" - ) + assert manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira" def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self): """After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server.""" @@ -2573,9 +2810,7 @@ def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self): assert resolved_server_unpref.server_id == server.server_id # Prefixed resolution - resolved_server_pref = manager._get_mcp_server_from_tool_name( - "zapier-create_zap" - ) + resolved_server_pref = manager._get_mcp_server_from_tool_name("zapier-create_zap") assert resolved_server_pref is not None assert resolved_server_pref.server_id == server.server_id @@ -2620,9 +2855,7 @@ async def test_rest_endpoint_filters_by_allowed_tools(self): new=AsyncMock(return_value=[tool1, tool2, tool3]), ): # Call the REST endpoint helper - filtered_response = await _get_tools_for_single_server( - server, server_auth_header=None - ) + filtered_response = await _get_tools_for_single_server(server, server_auth_header=None) # Verify only allowed tools are in the response assert len(filtered_response) == 2 @@ -2672,9 +2905,7 @@ async def test_rest_endpoint_shows_all_when_allowed_tools_is_none(self): new=AsyncMock(return_value=[tool1, tool2, tool3]), ): # Call the REST endpoint helper - all_tools_response = await _get_tools_for_single_server( - server, server_auth_header=None - ) + all_tools_response = await _get_tools_for_single_server(server, server_auth_header=None) # Verify all tools are returned (no filtering) assert len(all_tools_response) == 3 @@ -2719,9 +2950,7 @@ async def test_rest_endpoint_shows_all_when_allowed_tools_is_empty_list(self): new=AsyncMock(return_value=[tool1, tool2]), ): # Call the REST endpoint helper - all_tools_response = await _get_tools_for_single_server( - server, server_auth_header=None - ) + all_tools_response = await _get_tools_for_single_server(server, server_auth_header=None) # Verify all tools are returned (no filtering) assert len(all_tools_response) == 2 @@ -2791,9 +3020,7 @@ async def test_key_tool_permission_allows_permitted_tool(self): ) proxy_logging = MagicMock() - proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging.pre_call_hook = AsyncMock(return_value=None) @@ -2836,9 +3063,7 @@ async def test_key_tool_permission_blocks_unpermitted_tool(self): ) proxy_logging = MagicMock() - proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging.pre_call_hook = AsyncMock(return_value=None) @@ -2983,9 +3208,7 @@ async def test_allowed_tools_with_mixed_prefixed_and_unprefixed_names(self): proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) @@ -3021,13 +3244,8 @@ async def test_allowed_tools_with_mixed_prefixed_and_unprefixed_names(self): ) assert exc_info.value.status_code == 403 - assert ( - "Tool deletepet is not allowed for server my_api_mcp" - in exc_info.value.detail["error"] - ) - assert ( - "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] - ) + assert "Tool deletepet is not allowed for server my_api_mcp" in exc_info.value.detail["error"] + assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] @pytest.mark.asyncio async def test_call_tool_without_broken_pipe_error(self): @@ -3052,9 +3270,7 @@ async def test_call_tool_without_broken_pipe_error(self): # Register the server and map a tool to it manager.registry = {"test-server": server} manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server" - manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = ( - "test-server" - ) + manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server" # Create mock client that tracks call_tool usage mock_client = AsyncMock() @@ -3078,9 +3294,7 @@ async def mock_call_tool(params, host_progress_callback=None): # Mock proxy logging proxy_logging_obj = MagicMock() - proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( - return_value={} - ) + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -3145,9 +3359,7 @@ async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): # Verify MCPRequestHandler.get_allowed_mcp_servers was called with user_api_key_auth mock_get_allowed.assert_called_once() call_args = mock_get_allowed.call_args - assert ( - call_args[0][0] is user_api_key_auth - ) # First positional arg should be user_api_key_auth + assert call_args[0][0] is user_api_key_auth # First positional arg should be user_api_key_auth assert call_args[0][0].user_id == "user-123" assert call_args[0][0].object_permission_id == "perm_123" assert call_args[0][0].object_permission is not None @@ -3183,9 +3395,7 @@ async def test_no_mcp_servers_sentinel_blocks_allow_all_keys(self): ) with ( - patch.object( - manager, "get_allow_all_keys_server_ids", return_value=["global-server"] - ), + patch.object(manager, "get_allow_all_keys_server_ids", return_value=["global-server"]), patch.object( MCPRequestHandler, "get_allowed_mcp_servers", @@ -3339,6 +3549,180 @@ async def test_build_mcp_server_from_table_reads_token_endpoint_auth_method(self default_server = await manager.build_mcp_server_from_table(default_record, credentials_are_encrypted=False) assert default_server.token_endpoint_auth_method is None + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_discovers_obo_token_url_when_unset(self): + """DB path: an OBO server with no token_exchange_endpoint in credentials and no token_url + column runs discovery, and the resolved endpoint lands on the returned MCPServer.""" + manager = MCPServerManager() + calls: list[bool] = [] + + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + calls.append(allow_origin_fallback) + return MCPOAuthMetadata( + scopes=None, + authorization_url=None, + token_url="https://discovered.example.com/token", + registration_url=None, + ) + + manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] + + record = LiteLLM_MCPServerTable( + server_id="obo-discover-db-1", + server_name="obo_discover_db", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + credentials={"client_id": "cid", "client_secret": "csec"}, + ) + + # prisma_client None -> the write-back no-ops; this test isolates the discovery behavior. + with patch("litellm.proxy.proxy_server.prisma_client", None): + server = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + + assert calls == [False] # discovery ran once, origin fallback disabled for OBO + assert server.token_url == "https://discovered.example.com/token" + + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_skips_discovery_when_obo_endpoint_configured(self): + """DB path: a configured token_exchange_endpoint in the credentials JSON wins and skips + discovery entirely, even though the token_url column is empty (the DB-specific lookup uses + credentials_dict["token_exchange_endpoint"], not the column).""" + manager = MCPServerManager() + calls: list[str] = [] + + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + calls.append(server_url) + raise AssertionError("discovery must not run when token_exchange_endpoint is configured") + + manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] + + record = LiteLLM_MCPServerTable( + server_id="obo-configured-db-1", + server_name="obo_configured_db", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + credentials={ + "client_id": "cid", + "client_secret": "csec", + "token_exchange_endpoint": "https://idp.example.com/token", + }, + ) + + server = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + + assert calls == [] # discovery never ran + assert server.token_exchange_endpoint == "https://idp.example.com/token" + + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_persists_discovered_obo_token_url(self): + """A DB-backed OBO server with no configured endpoint discovers token_url and must write it + back to the row, so the next rebuild skips discovery instead of re-running it every time.""" + manager = MCPServerManager() + + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + assert server_url == "https://example.com/mcp" + assert allow_origin_fallback is False # OBO never guesses the origin + return MCPOAuthMetadata( + scopes=None, + authorization_url=None, + token_url="https://discovered.example.com/token", + registration_url=None, + ) + + manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] + + record = LiteLLM_MCPServerTable( + server_id="obo-persist-1", + server_name="obo_persist", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + credentials={"client_id": "cid", "client_secret": "csec", "audience": "aud"}, + ) + + update_mock = AsyncMock() + repo_instance = MagicMock() + repo_instance.table.update = update_mock + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repo_instance, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + server = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + + assert server.token_url == "https://discovered.example.com/token" + update_mock.assert_awaited_once() + assert update_mock.call_args.kwargs["where"] == {"server_id": "obo-persist-1"} + assert update_mock.call_args.kwargs["data"] == {"token_url": "https://discovered.example.com/token"} + + @pytest.mark.asyncio + async def test_persist_discovered_obo_token_url_skips_when_not_needed(self): + """The write-back fires only for an OBO server that discovered a new endpoint: a row that + already has token_url, a non-OBO auth_type, or a discovery that found nothing all no-op.""" + manager = MCPServerManager() + update_mock = AsyncMock() + repo_instance = MagicMock() + repo_instance.table.update = update_mock + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repo_instance, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + # already populated -> no write + await manager._persist_discovered_obo_token_url( + server_id="s", + auth_type=MCPAuth.oauth2_token_exchange, + existing_token_url="https://already.example.com/token", + discovered_token_url="https://new.example.com/token", + ) + # not an OBO server -> no write + await manager._persist_discovered_obo_token_url( + server_id="s", + auth_type=MCPAuth.oauth2, + existing_token_url=None, + discovered_token_url="https://new.example.com/token", + ) + # discovery found nothing -> no write + await manager._persist_discovered_obo_token_url( + server_id="s", + auth_type=MCPAuth.oauth2_token_exchange, + existing_token_url=None, + discovered_token_url=None, + ) + + update_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_persist_discovered_obo_token_url_is_best_effort(self): + """A write-back failure must not propagate; discovery just re-runs on the next build.""" + manager = MCPServerManager() + update_mock = AsyncMock(side_effect=Exception("db unavailable")) + repo_instance = MagicMock() + repo_instance.table.update = update_mock + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repo_instance, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + await manager._persist_discovered_obo_token_url( + server_id="s", + auth_type=MCPAuth.oauth2_token_exchange, + existing_token_url=None, + discovered_token_url="https://new.example.com/token", + ) + + update_mock.assert_awaited_once() + def test_build_mcp_server_table_preserves_timestamps(self): """_build_mcp_server_table must use the MCPServer's stored timestamps, not datetime.now().""" manager = MCPServerManager() @@ -3404,9 +3788,7 @@ def test_deserialize_json_list_normalizes_pydantic_models(self): ``_deserialize_json_list`` must hand back plain dicts so ``MCPServer`` (typed ``List[Dict[str, Any]]``) validates.""" env_vars = [ - MCPEnvVar( - name="GITHUB_TOKEN", scope=MCPEnvVarScope.user, description="PAT" - ), + MCPEnvVar(name="GITHUB_TOKEN", scope=MCPEnvVarScope.user, description="PAT"), MCPEnvVar(name="REGION", value="us-east-1", scope=MCPEnvVarScope.global_), ] result = _deserialize_json_list(env_vars) @@ -3752,10 +4134,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: def test_get_returns_none_when_empty(self): """Empty cache returns None for any key.""" manager = MCPServerManager() - assert ( - manager._upstream_initialize_instructions_by_server_id.get("nonexistent") - is None - ) + assert manager._upstream_initialize_instructions_by_server_id.get("nonexistent") is None def test_remember_stores_stripped_value(self): """_remember_upstream_initialize_instructions stores a stripped string.""" @@ -3763,9 +4142,7 @@ def test_remember_stores_stripped_value(self): fake_server = MagicMock(server_id="srv") fake_client = MagicMock(_last_initialize_instructions=" hello \n") manager._remember_upstream_initialize_instructions(fake_server, fake_client) - assert ( - manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello" - ) + assert manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello" def test_remember_ignores_empty_string(self): """Whitespace-only instructions are not stored.""" @@ -3843,9 +4220,7 @@ def test_empty_list_returns_empty(self): def test_expands_server_name(self): manager = MCPServerManager() - manager.config_mcp_servers["id-usw1"] = self._make_server( - "id-usw1", server_name="a" - ) + manager.config_mcp_servers["id-usw1"] = self._make_server("id-usw1", server_name="a") assert manager.expand_permission_list(["a"]) == ["id-usw1"] @@ -3869,9 +4244,7 @@ def test_passes_through_unknown_entry(self): def test_name_collision_expands_to_all_matches(self): """Two servers sharing a server_name both resolve — the documented behavior.""" manager = MCPServerManager() - manager.config_mcp_servers["id-config"] = self._make_server( - "id-config", server_name="shared" - ) + manager.config_mcp_servers["id-config"] = self._make_server("id-config", server_name="shared") manager.registry["id-db"] = self._make_server("id-db", server_name="shared") assert sorted(manager.expand_permission_list(["shared"])) == [ @@ -3881,9 +4254,7 @@ def test_name_collision_expands_to_all_matches(self): def test_searches_config_and_registry_union(self): manager = MCPServerManager() - manager.config_mcp_servers["cfg-id"] = self._make_server( - "cfg-id", server_name="a" - ) + manager.config_mcp_servers["cfg-id"] = self._make_server("cfg-id", server_name="a") manager.registry["reg-id"] = self._make_server("reg-id", server_name="b") assert manager.expand_permission_list(["a"]) == ["cfg-id"] @@ -3895,23 +4266,15 @@ def test_id_match_takes_precedence_over_name_match(self): servers whose server_name happens to equal that id. """ manager = MCPServerManager() - manager.config_mcp_servers["id-1"] = self._make_server( - "id-1", server_name="other_name" - ) - manager.config_mcp_servers["id-2"] = self._make_server( - "id-2", server_name="id-1" - ) + manager.config_mcp_servers["id-1"] = self._make_server("id-1", server_name="other_name") + manager.config_mcp_servers["id-2"] = self._make_server("id-2", server_name="id-1") assert manager.expand_permission_list(["id-1"]) == ["id-1"] def test_mixed_ids_and_names_in_same_list(self): manager = MCPServerManager() - manager.config_mcp_servers["uuid-1"] = self._make_server( - "uuid-1", server_name="a" - ) - manager.config_mcp_servers["uuid-2"] = self._make_server( - "uuid-2", server_name="b" - ) + manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="a") + manager.config_mcp_servers["uuid-2"] = self._make_server("uuid-2", server_name="b") # ["uuid-1", "b"] -> uuid-1 passes through, "b" resolves to uuid-2 assert sorted(manager.expand_permission_list(["uuid-1", "b"])) == [ @@ -3922,9 +4285,7 @@ def test_mixed_ids_and_names_in_same_list(self): def test_deduplicates_overlapping_id_and_name_entries(self): """If a list references the same server by both id and name, return it once.""" manager = MCPServerManager() - manager.config_mcp_servers["uuid-1"] = self._make_server( - "uuid-1", server_name="a" - ) + manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="a") assert manager.expand_permission_list(["uuid-1", "a"]) == ["uuid-1"] @@ -3934,14 +4295,10 @@ def test_simulates_cross_region_portability(self): the cross-region portability the customer is asking for. """ usw1 = MCPServerManager() - usw1.config_mcp_servers["hash-usw1"] = self._make_server( - "hash-usw1", server_name="a" - ) + usw1.config_mcp_servers["hash-usw1"] = self._make_server("hash-usw1", server_name="a") usc1 = MCPServerManager() - usc1.config_mcp_servers["hash-usc1"] = self._make_server( - "hash-usc1", server_name="a" - ) + usc1.config_mcp_servers["hash-usc1"] = self._make_server("hash-usc1", server_name="a") assert usw1.expand_permission_list(["a"]) == ["hash-usw1"] assert usc1.expand_permission_list(["a"]) == ["hash-usc1"] @@ -3970,18 +4327,14 @@ def test_rewrites_name_key_to_server_id(self): concrete server_id, otherwise `.get(server_id)` misses and the tool restriction is silently dropped (caller treats None as allow-all).""" manager = MCPServerManager() - manager.config_mcp_servers["uuid-a"] = self._make_server( - "uuid-a", server_name="my-alias" - ) + manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="my-alias") result = manager.expand_tool_permissions({"my-alias": ["read_file"]}) assert result == {"uuid-a": ["read_file"]} def test_passes_through_existing_server_id_key(self): manager = MCPServerManager() - manager.config_mcp_servers["uuid-a"] = self._make_server( - "uuid-a", server_name="alpha" - ) + manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="alpha") result = manager.expand_tool_permissions({"uuid-a": ["read_file"]}) assert result == {"uuid-a": ["read_file"]} @@ -4000,9 +4353,7 @@ def test_name_collision_unions_tool_lists(self): """Two servers sharing a server_name both match; their tool lists get the restriction (matches the list-expansion collision semantics).""" manager = MCPServerManager() - manager.config_mcp_servers["uuid-1"] = self._make_server( - "uuid-1", server_name="shared" - ) + manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="shared") manager.registry["uuid-2"] = self._make_server("uuid-2", server_name="shared") result = manager.expand_tool_permissions({"shared": ["read_file"]}) @@ -4015,13 +4366,9 @@ def test_id_and_name_keys_pointing_at_same_server_union_tools(self): both refer to the same server, the tool lists are unioned rather than one overwriting the other.""" manager = MCPServerManager() - manager.config_mcp_servers["uuid-a"] = self._make_server( - "uuid-a", server_name="alias-a" - ) + manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="alias-a") - result = manager.expand_tool_permissions( - {"uuid-a": ["read_file"], "alias-a": ["write_file"]} - ) + result = manager.expand_tool_permissions({"uuid-a": ["read_file"], "alias-a": ["write_file"]}) assert sorted(result["uuid-a"]) == ["read_file", "write_file"] @@ -4048,9 +4395,7 @@ def fake_getaddrinfo(host, port, *args, **kwargs): if host not in mapping: raise _socket.gaierror(f"unknown host {host}") family = _socket.AF_INET - return [ - (family, _socket.SOCK_STREAM, 0, "", (ip, port)) for ip in mapping[host] - ] + return [(family, _socket.SOCK_STREAM, 0, "", (ip, port)) for ip in mapping[host]] monkeypatch.setattr( "litellm.litellm_core_utils.url_utils.socket.getaddrinfo", @@ -4088,9 +4433,7 @@ def test_same_host_different_port_uses_safe_fetch_path(self): ], ) @pytest.mark.asyncio - async def test_cross_origin_blocked_when_resolves_to_unsafe_ip( - self, monkeypatch, ip - ): + async def test_cross_origin_blocked_when_resolves_to_unsafe_ip(self, monkeypatch, ip): self._patch_resolves(monkeypatch, {"attacker.example.com": [ip]}) manager = MCPServerManager() @@ -4111,9 +4454,7 @@ async def test_cross_origin_blocked_when_resolves_to_unsafe_ip( @pytest.mark.asyncio async def test_cross_origin_allowed_when_resolves_to_public_ip(self, monkeypatch): - self._patch_resolves( - monkeypatch, {"login.microsoftonline.com": ["20.190.151.7"]} - ) + self._patch_resolves(monkeypatch, {"login.microsoftonline.com": ["20.190.151.7"]}) manager = MCPServerManager() mock_response = MagicMock() @@ -4140,10 +4481,7 @@ async def test_cross_origin_allowed_when_resolves_to_public_ip(self, monkeypatch assert scopes == ["mcp.read"] mock_client.get.assert_awaited_once() assert mock_client.get.await_args.kwargs["follow_redirects"] is False - assert ( - mock_client.get.await_args.kwargs["headers"]["Host"] - == "login.microsoftonline.com" - ) + assert mock_client.get.await_args.kwargs["headers"]["Host"] == "login.microsoftonline.com" @pytest.mark.asyncio async def test_cross_origin_blocked_when_unresolvable(self, monkeypatch): @@ -4194,9 +4532,7 @@ async def test_non_http_scheme_is_not_safe(self): async def test_dual_resolution_blocked_if_any_ip_unsafe(self, monkeypatch): # If the attacker controls a DNS record returning multiple A records, # one of which is private, async_safe_get rejects before any network call. - self._patch_resolves( - monkeypatch, {"dual-stack.example.com": ["8.8.8.8", "127.0.0.1"]} - ) + self._patch_resolves(monkeypatch, {"dual-stack.example.com": ["8.8.8.8", "127.0.0.1"]}) manager = MCPServerManager() mock_client = MagicMock() @@ -4421,9 +4757,7 @@ def _make_server(self, server_id: str, approval_status): ("approved", True), ], ) - async def test_add_server_respects_approval_status( - self, approval_status, expect_in_registry - ): + async def test_add_server_respects_approval_status(self, approval_status, expect_in_registry): manager = MCPServerManager() server_id = f"sid-{approval_status}" await manager.add_server(self._make_server(server_id, approval_status)) @@ -4434,19 +4768,13 @@ async def test_update_server_evicts_when_transitioned_away_from_active(self): # The stale registry entry must be evicted so subsequent tool calls # and health probes can't reach it. manager = MCPServerManager() - await manager.add_server( - self._make_server("evict-me", MCPApprovalStatus.active) - ) + await manager.add_server(self._make_server("evict-me", MCPApprovalStatus.active)) assert "evict-me" in manager.registry - await manager.update_server( - self._make_server("evict-me", MCPApprovalStatus.rejected) - ) + await manager.update_server(self._make_server("evict-me", MCPApprovalStatus.rejected)) assert "evict-me" not in manager.registry - async def test_update_server_eviction_clears_openapi_routing_artifacts( - self, tmp_path - ): + async def test_update_server_eviction_clears_openapi_routing_artifacts(self, tmp_path): """Rejecting a server must remove its OpenAPI tools and name mappings.""" from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, @@ -4457,9 +4785,7 @@ async def test_update_server_eviction_clears_openapi_routing_artifacts( ) manager = MCPServerManager() - await manager.add_server( - self._make_server("evict-openapi", MCPApprovalStatus.active) - ) + await manager.add_server(self._make_server("evict-openapi", MCPApprovalStatus.active)) assert "evict-openapi" in manager.registry server = manager.registry["evict-openapi"] @@ -4479,9 +4805,7 @@ async def _noop_handler(**kwargs): manager.tool_name_to_mcp_server_name_mapping["demo_tool"] = prefix manager.tool_name_to_mcp_server_name_mapping[prefixed] = prefix - await manager.update_server( - self._make_server("evict-openapi", MCPApprovalStatus.rejected) - ) + await manager.update_server(self._make_server("evict-openapi", MCPApprovalStatus.rejected)) assert "evict-openapi" not in manager.registry assert prefixed not in global_mcp_tool_registry.tools @@ -4494,9 +4818,7 @@ async def test_update_server_noop_for_unregistered_pending(self): # so a future refactor can't accidentally route the pending row to # build_mcp_server_from_table. manager = MCPServerManager() - await manager.update_server( - self._make_server("never-seen", MCPApprovalStatus.pending_review) - ) + await manager.update_server(self._make_server("never-seen", MCPApprovalStatus.pending_review)) assert "never-seen" not in manager.registry @@ -4710,9 +5032,7 @@ def test_returns_empty_when_whitelist_is_none(self): @patch("litellm.public_mcp_servers", []) def test_returns_empty_when_whitelist_is_empty(self): """Explicit empty whitelist → hub returns nothing.""" - manager = self._make_manager( - [self._make_server("a", available_on_public_internet=True)] - ) + manager = self._make_manager([self._make_server("a", available_on_public_internet=True)]) assert manager.get_public_mcp_servers() == [] @patch("litellm.public_mcp_servers", ["a"]) @@ -4750,9 +5070,7 @@ def test_does_not_leak_servers_via_internal_flag(self): @patch("litellm.public_mcp_servers", ["does-not-exist"]) def test_stale_whitelist_id_returns_empty(self): """Whitelist references an unknown server_id → no spurious results.""" - manager = self._make_manager( - [self._make_server("a", available_on_public_internet=True)] - ) + manager = self._make_manager([self._make_server("a", available_on_public_internet=True)]) assert manager.get_public_mcp_servers() == [] @@ -4831,9 +5149,7 @@ async def test_none_mode_resolves_to_noop_auth(self): NoOpAuth, ) - client = await MCPServerManager()._create_mcp_client( - self._http_server(auth_type=None) - ) + client = await MCPServerManager()._create_mcp_client(self._http_server(auth_type=None)) assert isinstance(client._resolved_auth, NoOpAuth) assert client._mcp_auth_value is None @@ -4847,9 +5163,7 @@ async def test_none_mode_resolves_to_noop_auth(self): (MCPAuth.authorization, "raw-123", "Authorization", "raw-123"), ], ) - async def test_static_family_emits_expected_header( - self, auth_type, token, expected_name, expected_value - ): + async def test_static_family_emits_expected_header(self, auth_type, token, expected_name, expected_value): from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( StaticHeaderAuth, ) @@ -4877,9 +5191,7 @@ async def test_basic_mode_base64_encodes(self): encoded = base64.b64encode(b"user:pass").decode() assert isinstance(client._resolved_auth, StaticHeaderAuth) assert client._resolved_auth.header_name == "Authorization" - assert ( - client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}" - ) + assert client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}" async def test_m2m_client_credentials_defers_to_v1(self): # M2M (oauth2 client_credentials) is not migrated: to_server_spec returns @@ -4944,9 +5256,7 @@ async def test_per_request_override_defers_to_v1(self): # A per-request override (mcp_auth_header) must win over the shared static token, # exactly as v1 did, so a migrated static server defers to v1 when one is present. client = await MCPServerManager()._create_mcp_client( - self._http_server( - auth_type=MCPAuth.bearer_token, authentication_token="shared-tok" - ), + self._http_server(auth_type=MCPAuth.bearer_token, authentication_token="shared-tok"), mcp_auth_header="caller-override", ) @@ -4958,9 +5268,7 @@ async def test_conflicting_extra_header_skips_resolved_auth_on_v2(self): # signer, static_headers, or a forwarded caller header) must win. The server stays on # the v2 path but skips resolved_auth, so nothing overwrites the inbound header. client = await MCPServerManager()._create_mcp_client( - self._http_server( - auth_type=MCPAuth.bearer_token, authentication_token="shared-tok" - ), + self._http_server(auth_type=MCPAuth.bearer_token, authentication_token="shared-tok"), extra_headers={"Authorization": "Bearer hook-jwt"}, ) @@ -4991,9 +5299,7 @@ def _upstream_status_error(status_code: int, challenge: str) -> httpx.HTTPStatus headers={"WWW-Authenticate": challenge}, request=request, ) - return httpx.HTTPStatusError( - "upstream rejected token", request=request, response=response - ) + return httpx.HTTPStatusError("upstream rejected token", request=request, response=response) class TestMCPToolsListAuthSurfacing: @@ -5053,9 +5359,7 @@ async def test_get_tools_from_server_surfaces_unusable_user_token(self): ) manager = MCPServerManager() - server = MCPServer( - server_id="oauth-srv", name="oauth-srv", transport=MCPTransport.http - ) + server = MCPServer(server_id="oauth-srv", name="oauth-srv", transport=MCPTransport.http) challenge = 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/oauth-srv"' manager._create_mcp_client = AsyncMock( side_effect=HTTPException( @@ -5074,13 +5378,13 @@ async def test_get_tools_from_server_surfaces_unusable_user_token(self): @pytest.mark.asyncio async def test_get_tools_from_server_absorbs_non_challenge_http_error(self): + """A non-auth HTTPException (500) stays absorbed so one misconfigured server cannot blank + the listing; 401/403 are the challenge-class statuses routed to MCPUpstreamAuthError.""" manager = MCPServerManager() - server = MCPServer( - server_id="stdio-srv", name="stdio-srv", transport=MCPTransport.http - ) + server = MCPServer(server_id="stdio-srv", name="stdio-srv", transport=MCPTransport.http) manager._create_mcp_client = AsyncMock( side_effect=HTTPException( - status_code=403, + status_code=500, detail="MCP stdio command 'foo' is not in the allowlist", ) ) @@ -5118,5 +5422,256 @@ async def fake_get_tools(server, **kwargs): assert [t.name for t in result] == ["good-do_thing"] +def test_should_strip_caller_authorization_for_token_exchange(): + """OBO: the inbound bearer is the subject token (exchanged), never forwarded upstream raw.""" + server = MCPServer( + server_id="te-strip", + name="te-strip-server", + url="https://up.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + assert _should_strip_caller_authorization(mcp_server=server, raw_headers=None, user_api_key_auth=None) is True + + +class _UpstreamAuthError(Exception): + """Mimics a wrapped upstream 401 the way _extract_upstream_auth_failure detects it.""" + + def __init__(self, status_code: int = 401) -> None: + super().__init__(f"HTTP {status_code}") + self.response = httpx.Response(status_code) + + +class _RetryFakeClient: + """A fake MCPClient whose call_tool fails on the first attempt and (optionally) succeeds after.""" + + def __init__(self, *, raises=None, result=None) -> None: + from litellm.experimental_mcp_client.client import MCPClient + + self._raises = raises + self._result = result + self._MCPClient = MCPClient + self.attempts = 0 + + async def call_tool(self, params, host_progress_callback=None, raise_on_error=False): + self.attempts += 1 + if self._raises is not None: + if raise_on_error: + raise self._raises + return self._MCPClient.error_tool_result(self._raises) + return self._result + + +def _obo_server() -> MCPServer: + return MCPServer( + server_id="obo-srv", + name="obo", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + + +class TestOBOCallToolRetry: + """The token_exchange (OBO) tool-call path re-mints the exchanged token once on an upstream 401.""" + + def _manager(self): + manager = MCPServerManager() + manager._cred_provider = MagicMock() + manager._cred_provider.invalidate_credentials = AsyncMock() + return manager + + @pytest.mark.asyncio + async def test_upstream_401_invalidates_and_retries_once(self): + manager = self._manager() + success = CallToolResult(content=[], isError=False) + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + retry = _RetryFakeClient(result=success) + manager._create_mcp_client = AsyncMock(return_value=retry) + + result = await manager._obo_call_tool_with_retry( + client=first, + call_tool_params=MagicMock(), + host_progress_callback=None, + mcp_server=_obo_server(), + server_auth_header=None, + extra_headers=None, + stdio_env=None, + subject_token="caller-jwt", + user_api_key_auth=None, + ) + + assert result is success + manager._cred_provider.invalidate_credentials.assert_awaited_once() + manager._create_mcp_client.assert_awaited_once() + assert first.attempts == 1 and retry.attempts == 1 + + @pytest.mark.asyncio + async def test_non_auth_error_does_not_retry(self): + manager = self._manager() + first = _RetryFakeClient(raises=ValueError("tool blew up")) + manager._create_mcp_client = AsyncMock() + + result = await manager._obo_call_tool_with_retry( + client=first, + call_tool_params=MagicMock(), + host_progress_callback=None, + mcp_server=_obo_server(), + server_auth_header=None, + extra_headers=None, + stdio_env=None, + subject_token="caller-jwt", + user_api_key_auth=None, + ) + + assert result.isError is True + manager._cred_provider.invalidate_credentials.assert_not_awaited() + manager._create_mcp_client.assert_not_awaited() + assert first.attempts == 1 + + @pytest.mark.asyncio + async def test_second_401_degrades_without_looping(self): + manager = self._manager() + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + # The retry client still fails; with raise_on_error defaulting False it returns isError. + retry = _RetryFakeClient(raises=_UpstreamAuthError(401)) + manager._create_mcp_client = AsyncMock(return_value=retry) + + result = await manager._obo_call_tool_with_retry( + client=first, + call_tool_params=MagicMock(), + host_progress_callback=None, + mcp_server=_obo_server(), + server_auth_header=None, + extra_headers=None, + stdio_env=None, + subject_token="caller-jwt", + user_api_key_auth=None, + ) + + assert result.isError is True + manager._create_mcp_client.assert_awaited_once() + assert first.attempts == 1 and retry.attempts == 1 + + +class TestOBOEndpointDiscovery: + """An oauth2_token_exchange server with no configured token endpoint discovers it (RFC 9728 -> + RFC 8414) like the oauth2 flow does; an explicitly configured endpoint skips discovery.""" + + @pytest.mark.parametrize( + "auth_type, endpoint, token_url, expected", + [ + (MCPAuth.oauth2_token_exchange, None, None, True), # OBO, nothing configured -> discover + (MCPAuth.oauth2_token_exchange, "https://idp/token", None, False), # endpoint set -> skip + (MCPAuth.oauth2_token_exchange, None, "https://idp/token", False), # token_url set -> skip + (MCPAuth.oauth2, None, None, False), # not OBO + (MCPAuth.none, None, None, False), + (None, None, None, False), + ], + ) + def test_decision(self, auth_type, endpoint, token_url, expected): + assert MCPServerManager._obo_needs_endpoint_discovery(auth_type, endpoint, token_url) is expected + + @pytest.mark.asyncio + async def test_config_obo_without_endpoint_discovers_token_endpoint(self): + manager = MCPServerManager() + discovered = MCPOAuthMetadata( + scopes=None, + authorization_url=None, + token_url="https://discovered.example.com/token", + registration_url=None, + ) + seen = [] + + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + seen.append((server_url, allow_origin_fallback)) + return discovered + + manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] + + await manager.load_servers_from_config( + { + "obo": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_token_exchange, + "client_id": "cid", + "client_secret": "csec", + } + } + ) + + server = next(iter(manager.config_mcp_servers.values())) + # OBO discovery runs, and never guesses the resource origin as the IdP (origin fallback off). + assert seen == [("https://example.com/mcp", False)] + # The discovered token endpoint lands on token_url, which _token_exchange_spec reads. + assert server.token_url == "https://discovered.example.com/token" + + @pytest.mark.asyncio + async def test_config_obo_with_configured_endpoint_skips_discovery(self): + manager = MCPServerManager() + + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + raise AssertionError("discovery must not run when the endpoint is configured") + + manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] + + await manager.load_servers_from_config( + { + "obo": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_token_exchange, + "token_exchange_endpoint": "https://configured.example.com/token", + "client_id": "cid", + "client_secret": "csec", + } + } + ) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.token_exchange_endpoint == "https://configured.example.com/token" + + @pytest.mark.asyncio + async def test_descovery_metadata_does_not_guess_origin_when_disallowed(self): + # The authoritative RFC 9728 -> RFC 8414 chain stays, but with allow_origin_fallback=False the + # resource origin is never assumed to be the IdP, so no token endpoint is invented. + manager = MCPServerManager() + server_url = "https://example.com/public/mcp" + request = httpx.Request("GET", server_url) + response_obj = httpx.Response( + status_code=401, request=request, headers={"WWW-Authenticate": 'Bearer scope="read"'} + ) + response_obj.raise_for_status = MagicMock( + side_effect=lambda: (_ for _ in ()).throw( + httpx.HTTPStatusError("unauthorized", request=request, response=response_obj) + ) + ) + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=response_obj) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ), + patch.object(manager, "_fetch_oauth_metadata_from_resource", AsyncMock(return_value=([], None))), + patch.object(manager, "_attempt_well_known_discovery", AsyncMock(return_value=([], None))), + patch.object(manager, "_fetch_authorization_server_metadata", AsyncMock()) as mock_fetch_auth, + ): + result = await manager._descovery_metadata(server_url, allow_origin_fallback=False) + + # No advertised AS -> with the guess disabled, the AS-metadata fetch is never attempted, so no + # token endpoint is discovered (only the scopes parsed from the challenge survive). + mock_fetch_auth.assert_not_awaited() + assert result is None or result.token_url is None + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index e4e1890a45a..76614a35b53 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -294,15 +294,9 @@ async def _stateless_capture(s, r, se): # Verify the mcp-session-id header was stripped header_names = [k for k, v in captured_scope.get("headers", [])] - assert ( - b"mcp-session-id" not in header_names - ), "Stale mcp-session-id header should have been stripped from the scope" - assert ( - stateless_handle_request.called - ), "Stale non-initialize requests should route stateless" - assert ( - not stateful_handle_request.called - ), "Stale non-initialize requests should not route stateful" + assert b"mcp-session-id" not in header_names, "Stale mcp-session-id header should have been stripped from the scope" + assert stateless_handle_request.called, "Stale non-initialize requests should route stateless" + assert not stateful_handle_request.called, "Stale non-initialize requests should not route stateful" @pytest.mark.asyncio @@ -366,9 +360,7 @@ async def test_delete_stale_mcp_session_returns_success(): await handle_streamable_http_mcp(scope, receive, send) # Verify session manager was NOT called (request was handled early) - assert ( - not mock_handle_request.called - ), "Session manager should not be called for DELETE on non-existent session" + assert not mock_handle_request.called, "Session manager should not be called for DELETE on non-existent session" # Verify a success response was sent assert send.called, "A response should have been sent" @@ -523,9 +515,7 @@ async def mock_handle_request(s, r, se): # Verify the mcp-session-id header was preserved header_names = [k for k, v in captured_scope.get("headers", [])] - assert ( - b"mcp-session-id" in header_names - ), "Valid mcp-session-id header should have been preserved" + assert b"mcp-session-id" in header_names, "Valid mcp-session-id header should have been preserved" @pytest.mark.asyncio @@ -967,6 +957,94 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns challenge = exc_info.value.headers["www-authenticate"] assert "resource_metadata=" in challenge assert "authorization_uri=" not in challenge - assert ( - "/.well-known/oauth-protected-resource/delegated_oauth_server/mcp" in challenge + assert "/.well-known/oauth-protected-resource/delegated_oauth_server/mcp" in challenge + + +@pytest.mark.asyncio +async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns_preemptive_resource_metadata_401(): + """An ``oauth2_token_exchange`` (OBO) server with no caller subject token must fail fast at + connect with a 401 carrying the RFC 9728 ``resource_metadata`` + RFC 6750 ``invalid_token`` + challenge, so the client discovers the IdP and retries with a subject token. A tool-call-time + 401 would be wrapped into a JSON-RPC error and the WWW-Authenticate lost, so this preemptive + challenge is what drives the discovery flow.""" + from fastapi import HTTPException + + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/obo_server", + "_original_path": "/mcp/obo_server", + "scheme": "https", + "query_string": b"", + "root_path": "", + "server": ("litellm.example.com", 443), + "headers": [ + (b"content-type", b"application/json"), + (b"host", b"litellm.example.com"), + ], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', + "more_body": False, + } ) + send = AsyncMock() + user_auth = MagicMock() + user_auth.user_id = None + obo_server = MagicMock() + obo_server.auth_type = MCPAuth.oauth2_token_exchange + obo_server.alias = None + obo_server.server_name = "obo_server" + obo_server.name = "obo_server" + obo_server.server_id = "obo-server" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["obo_server"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new_callable=AsyncMock, + return_value=False, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + return_value=obo_server, + ), + patch.object( + session_manager_stateful, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request, + ): + with pytest.raises(HTTPException) as exc_info: + await handle_streamable_http_mcp(scope, receive, send) + + assert mock_handle_request.await_count == 0 + assert exc_info.value.status_code == 401 + headers = {k.lower(): v for k, v in (exc_info.value.headers or {}).items()} + challenge = headers["www-authenticate"] + # Structural invariants only: the exact root-path prefix is exercised in the adapter's + # oauth_protected_resource_path unit test, so this handler test stays hermetic w.r.t. + # SERVER_ROOT_PATH (which other tests in the shard may have left set in the environment). + assert "resource_metadata=" in challenge + assert "/.well-known/oauth-protected-resource" in challenge + assert challenge.split('resource_metadata="', 1)[1].split('"', 1)[0].endswith("/mcp/obo_server") + assert 'error="invalid_token"' in challenge