diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 3e11ebfde77..9c8890efd2b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -307,12 +307,43 @@ def _gateway_dcr_challenge_target( targets: Final = _parse_mcp_server_names_from_path(route, mcp_servers) if targets is None: return None - server: Final = global_mcp_server_manager.get_mcp_server_by_name(targets[0], client_ip=client_ip) + server: Final = global_mcp_server_manager.get_mcp_server_by_identifier(targets[0], client_ip=client_ip) if server is None or not server.advertises_gateway_authorization_server: return None return targets[0] +def gateway_dcr_challenge_scope( + route: str, + mcp_servers: Sequence[str] | None, + client_ip: str | None, +) -> str | None: + if MCPRequestHandler.extract_target_server_names_from_path(route) or not mcp_servers: + return None + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import gateway_server_scope + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + resolved: Final = tuple( + (entry, global_mcp_server_manager.get_mcp_server_by_identifier(entry, client_ip=client_ip)) + for entry in dict.fromkeys(mcp_servers) + ) + if any(server is None or not server.advertises_gateway_authorization_server for _, server in resolved): + return None + servers: Final = tuple((entry, server) for entry, server in resolved if server is not None) + first_by_id: Final = {server.server_id: (entry, server) for entry, server in reversed(servers)} + unique_servers: Final = tuple( + first_by_id[server_id] for server_id in dict.fromkeys(server.server_id for _, server in servers) + ) + tokens: Final = tuple( + gateway_server_scope(entry) or gateway_server_scope(server.server_id) for entry, server in unique_servers + ) + if any(token is None for token in tokens): + return None + return " ".join(token for token in tokens if token is not None) + + def _is_gateway_dcr_challenge_scope( route: str, mcp_servers: list[str] | None, @@ -327,9 +358,10 @@ def _is_gateway_dcr_challenge_scope( Fires only for a genuine 401 with no client-supplied MCP auth headers (those mean the caller is not a cold-start DCR client), on the scopes the gateway's keyless flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request - (the resource the client configured is still ``/mcp``), or a per-server path whose - single target advertises gateway-owned sign-in. Every other named target keeps - its existing behavior, failing closed to the original admission error.""" + (the resource the client configured is still ``/mcp``; a single eligible header + target is also carried in the OAuth scope), or a per-server path whose single target + advertises gateway-owned sign-in. Every other named target keeps its existing + behavior, failing closed to the original admission error.""" if not _is_litellm_auth_admission_error(exc): return False if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): @@ -355,20 +387,23 @@ def _gateway_dcr_challenge( present a bearer that failed admission (expired or revoked), telling spec-compliant clients to re-authorize rather than retry; a request with no credentials at all gets the bare challenge per RFC 6750 section 3.1.""" - target: Final = _gateway_dcr_challenge_target(route, mcp_servers, IPAddressUtils.get_mcp_client_ip(request)) + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) + target: Final = _gateway_dcr_challenge_target(route, mcp_servers, client_ip) + server_scope: Final = gateway_dcr_challenge_scope(route, mcp_servers, client_ip) if target is None else None resource_metadata_url: Final = ( get_passthrough_resource_metadata_url(request.scope, target) if target is not None else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" ) error_attr: Final = 'error="invalid_token", ' if invalid_token else "" + scope_attr: Final = f', scope="{server_scope}"' if server_scope is not None else "" return HTTPException( status_code=401, detail={ "error": "authentication_required", "message": "Authenticate with the gateway to use the MCP endpoint.", }, - headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"'}, + headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"{scope_attr}'}, ) @@ -1123,7 +1158,7 @@ class MCPRequestHandler: case SessionBearerAdmitted(): try: admitted: Final = await MCPRequestHandler.reload_admitted_user(result.principal.user_id) - admitted.mcp_session_resource_server_id = result.principal.resource_server_id + admitted.mcp_session_resource_server_ids = result.principal.resource_server_ids await MCPRequestHandler._enforce_admitted_live_policy( admitted=admitted, request=request, route=route ) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 8876fd2dd87..cca722f45c2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -3113,10 +3113,28 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": f"{request_base_url}/mcp", - "scopes_supported": [], + "scopes_supported": _aggregate_header_scopes(request), } +def _aggregate_header_scopes(request: Request) -> tuple[str, ...]: + """Gateway server-scope tokens for the ``x-mcp-servers`` header a client sends with its + discovery request, so a client that skips the 401 challenge (Claude Code re-authenticating + a server it already knows needs auth) still carries the selection to ``/authorize`` through + ``scopes_supported``. Empty when the header is absent or any entry does not resolve.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( # noqa: PLC0415 # proxy import cycle + MCPRequestHandler, + gateway_dcr_challenge_scope, + ) + + header: Final = request.headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME) + if not header: + return () + names: Final = tuple(dict.fromkeys(entry.strip() for entry in header.split(",") if entry.strip())) + scope: Final = gateway_dcr_challenge_scope("/mcp", names, IPAddressUtils.get_mcp_client_ip(request)) + return tuple(scope.split()) if scope else () + + def _build_aggregate_authorization_server_response( request: Request, token_exchange_available: bool ) -> dict[str, object]: diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index a895dfed162..dd8c82a99fe 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -42,7 +42,8 @@ import hmac import html import secrets from base64 import urlsafe_b64encode -from collections.abc import Awaitable, Callable, Iterable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar @@ -50,7 +51,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, model_serializer, model_validator +from pydantic_core.core_schema import SerializerFunctionWrapHandler from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._internal_context import with_service_target @@ -106,6 +108,8 @@ GATEWAY_DCR_CLIENT_ID_PREFIX: Final = "llm_dcrc_" endpoints can route an aggregate-flow request without decrypting, and existing per-server flows (whose client_ids are upstream-issued) are never captured by the aggregate arm.""" +GATEWAY_SERVER_SCOPE_PREFIX: Final = "litellm:mcp_server:" + GATEWAY_AUTH_CODE_PREFIX: Final = "llm_gcode_" """Marker prefix on the gateway-sealed authorization code, distinct from the bridge ``llm_bcode_`` so neither flow can consume the other's codes.""" @@ -301,13 +305,55 @@ class GatewayDcrClient(LiteLLMBaseModel): iat: int -class _ConnectFlow(LiteLLMBaseModel): +_SEALED_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _legacy_server_scope_payload(payload: object) -> object: + if not isinstance(payload, Mapping): + return payload + values: Final = _SEALED_PAYLOAD_ADAPTER.validate_python(payload) + if "resource_server_id" not in values: + return values + if "resource_server_ids" in values: + raise ValueError("sealed payload cannot contain both resource_server_id and resource_server_ids") + resource_server_id: Final = values["resource_server_id"] + return {key: value for key, value in values.items() if key != "resource_server_id"} | ( + {"resource_server_ids": (resource_server_id,)} if resource_server_id is not None else {} + ) + + +def _server_scope_wire_payload( + payload: Mapping[str, object], resource_server_ids: tuple[str, ...] | None +) -> Mapping[str, object]: + if resource_server_ids is not None and len(resource_server_ids) != 1: + return payload + unscoped: Final = {key: value for key, value in payload.items() if key != "resource_server_ids"} + if resource_server_ids is None: + return unscoped + return unscoped | {"resource_server_id": resource_server_ids[0]} + + +class _ServerScopedSealedModel(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + resource_server_ids: tuple[str, ...] | None = Field(default=None, min_length=1) + + @model_validator(mode="before") + @classmethod + def _accept_legacy_server_scope(cls, payload: object) -> object: + return _legacy_server_scope_payload(payload) + + @model_serializer(mode="wrap") + def _serialize_server_scope(self, handler: SerializerFunctionWrapHandler) -> Mapping[str, object]: + payload: Final = _SEALED_PAYLOAD_ADAPTER.validate_python(handler(self)) + return _server_scope_wire_payload(payload, self.resource_server_ids) + + +class _ConnectFlow(_ServerScopedSealedModel): """One in-flight authorize: the SSO user it belongs to and the client parameters needed to mint the code at the finish step. Sealed into the per-flow cookie. ``jti`` makes the flow single-use at complete; ``extra="forbid"`` rejects cross-type confusion.""" - model_config = ConfigDict(frozen=True, extra="forbid") user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) redirect_uri: str = Field(min_length=1) @@ -315,19 +361,17 @@ class _ConnectFlow(LiteLLMBaseModel): code_challenge: str = Field(min_length=1) jti: str = Field(min_length=1) exp: int - resource_server_id: str | None = None audience: SessionAudience | None = None requested_scopes: tuple[str, ...] = () upstream_scopes: tuple[UpstreamScope, ...] = () -class _GatewayAuthCode(LiteLLMBaseModel): +class _GatewayAuthCode(_ServerScopedSealedModel): """The gateway-sealed authorization code: the user consent it represents and the bindings the token endpoint must verify (client, redirect URI, PKCE challenge), plus a ``jti`` for the single-use guard. ``extra="forbid"`` rejects cross-type confusion.""" - model_config = ConfigDict(frozen=True, extra="forbid") user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) redirect_uri: str = Field(min_length=1) @@ -335,7 +379,6 @@ class _GatewayAuthCode(LiteLLMBaseModel): jti: str = Field(min_length=1) iat: int exp: int - resource_server_id: str | None = None audience: SessionAudience | None = None team_id: str | None = None @@ -516,6 +559,34 @@ def relative_request_url(request: Request) -> str: return f"{path}?{request.url.query}" if request.url.query else path +def gateway_server_scope(server_name: str) -> str | None: + if not server_name or not all( + ord(character) == 0x21 or 0x23 <= ord(character) <= 0x5B or 0x5D <= ord(character) <= 0x7E + for character in server_name + ): + return None + return f"{GATEWAY_SERVER_SCOPE_PREFIX}{server_name}" + + +def _is_aggregate_resource(request: Request, resource: str) -> bool: + canonical: Final = canonical_resource_uri(resource) + if canonical is None: + return False + base: Final = canonicalize_url_identity(get_request_base_url(request)) + return canonical in (base, f"{base}/mcp") + + +def _gateway_flow_server(name: str) -> MCPServer | None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # proxy import cycle + global_mcp_server_manager, + ) + + server: Final = global_mcp_server_manager.get_mcp_server_by_identifier(name) + if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server): + return None + return server + + def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None: """Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it names, or ``None`` for every other shape: absent, the aggregate resource, a foreign @@ -536,22 +607,40 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC if canonical is None: return None base: Final = canonicalize_url_identity(get_request_base_url(request)) - if canonical == f"{base}/mcp" or not canonical.startswith(f"{base}/"): + if _is_aggregate_resource(request, resource) or not canonical.startswith(f"{base}/"): return None from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( # noqa: PLC0415 # proxy import cycle MCPRequestHandler, ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # proxy import cycle - global_mcp_server_manager, - ) names: Final = MCPRequestHandler.extract_target_server_names_from_path(canonical[len(base) :]) if len(names) != 1: return None - server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0]) - if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server): + return _gateway_flow_server(names[0]) + + +@dataclass(frozen=True, slots=True) +class InvalidGatewayScope: + pass + + +def resolve_scoped_servers_from_scope( + scope: str | None, +) -> tuple[MCPServer, ...] | InvalidGatewayScope | None: + if scope is None: return None - return server + prefixed_tokens: Final = tuple(token for token in scope.split() if token.startswith(GATEWAY_SERVER_SCOPE_PREFIX)) + if not prefixed_tokens: + return None + server_names: Final = tuple(dict.fromkeys(token[len(GATEWAY_SERVER_SCOPE_PREFIX) :] for token in prefixed_tokens)) + if any(not name for name in server_names): + return InvalidGatewayScope() + resolved: Final = tuple(_gateway_flow_server(name) for name in server_names) + if any(server is None for server in resolved): + return InvalidGatewayScope() + servers: Final = tuple(server for server in resolved if server is not None) + first_by_id: Final = {server.server_id: server for server in reversed(servers)} + return tuple(first_by_id[server_id] for server_id in dict.fromkeys(server.server_id for server in servers)) def aggregate_authorize( @@ -570,32 +659,42 @@ def aggregate_authorize( LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a per-flow cookie. - A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the - flow to that one server: the scope is sealed into the flow, carried into the code, and - bound into the session token. The connect URL carries only the flow handle; the page - learns the client origin, the scoped server, and whether its vendor OAuth is done from - :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry - steers which server the page authorizes or names on the confirmation. + A per-server RFC 8707 ``resource`` or gateway-server scope tokens on an aggregate + request scope the flow to those servers. The scope is sealed into the flow, carried + into the code, and bound into the session token. - Validation failures respond directly with 400 and never redirect: per RFC 6749 - section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and - once the client is at fault there is no trusted place to send the browser. + Validation failures respond directly with 400, except for an unknown gateway scope, + which is redirected to the already-validated ``redirect_uri`` as ``invalid_scope`` per + RFC 6749 section 4.1.2.1. """ rejected: Final = _rejected_authorize_request( client_id, redirect_uri, state, code_challenge, code_challenge_method, response_type ) if rejected is not None: return rejected + resource_scoped_server: Final = resolve_scoped_resource_server(request, resource) + scope_resolution: Final = None if resource_scoped_server is not None else resolve_scoped_servers_from_scope(scope) base_url: Final = get_request_base_url(request) if session_user_id is None: return _login_redirect(base_url, request) - scoped_server: Final = resolve_scoped_resource_server(request, resource) + if isinstance(scope_resolution, InvalidGatewayScope): + invalid_scope_location: Final = _append_query_params( + redirect_uri, + ( + ("error", "invalid_scope"), + ("error_description", "unknown or unsupported MCP server scope"), + *_state_param(state), + ), + ) + return RedirectResponse(invalid_scope_location, status_code=303) + scoped_servers: Final = (resource_scoped_server,) if resource_scoped_server is not None else scope_resolution or () scope_tokens: Final = tuple(token for token in (scope or "").split() if token.startswith(UPSTREAM_SCOPE_PREFIX)) decoded_scopes: Final = tuple(decode_upstream_scope(token) for token in scope_tokens) if any(item is None for item in decoded_scopes): return _oauth_error(400, "invalid_scope", "invalid upstream permission scope") upstream_scopes: Final = tuple(item for item in decoded_scopes if item is not None) - if scoped_server is not None and any(item.server_id != scoped_server.server_id for item in upstream_scopes): + scoped_server_ids: Final = frozenset(server.server_id for server in scoped_servers) + if scoped_server_ids and any(item.server_id not in scoped_server_ids for item in upstream_scopes): return _oauth_error(400, "invalid_scope", "upstream permission does not match the requested resource") handle: Final = secrets.token_urlsafe(24) flow: Final = _new_connect_flow( @@ -604,9 +703,11 @@ def aggregate_authorize( redirect_uri=redirect_uri, state=state, code_challenge=code_challenge or "", - resource_server_id=scoped_server.server_id if scoped_server is not None else None, + resource_server_ids=tuple(server.server_id for server in scoped_servers) or None, audience=None, - requested_scopes=tuple((scope or "").split()) if scoped_server is not None and not upstream_scopes else (), + requested_scopes=tuple((scope or "").split()) + if resource_scoped_server is not None and not upstream_scopes + else (), upstream_scopes=upstream_scopes, ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) @@ -653,7 +754,7 @@ async def native_client_authorize( redirect_uri=redirect_uri, state=state, code_challenge=code_challenge or "", - resource_server_id=None, + resource_server_ids=None, audience=PROXY_API_AUDIENCE, ) page: Final = render_native_client_consent_page( @@ -765,7 +866,7 @@ def _new_connect_flow( redirect_uri: str, state: str, code_challenge: str, - resource_server_id: str | None, + resource_server_ids: tuple[str, ...] | None, audience: SessionAudience | None, requested_scopes: tuple[str, ...] = (), upstream_scopes: tuple[UpstreamScope, ...] = (), @@ -779,7 +880,7 @@ def _new_connect_flow( code_challenge=code_challenge, jti=secrets.token_urlsafe(24), exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, - resource_server_id=resource_server_id, + resource_server_ids=resource_server_ids, audience=audience, requested_scopes=requested_scopes, upstream_scopes=upstream_scopes, @@ -862,17 +963,82 @@ async def _flow_target( return state, server +class ConnectFlowServer(TypedDict): + server_id: ReadOnly[str] + server_name: ReadOnly[str] + connected: ReadOnly[bool] + + class ConnectFlowDescription(TypedDict): """What the connect page is allowed to know about one in-flight flow.""" - state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]] + state: ReadOnly[Literal["unscoped", "interactive", "m2m", "multi", "stale"]] client_origin: ReadOnly[str] server_id: ReadOnly[str | None] server_name: ReadOnly[str | None] connected: ReadOnly[bool | None] + servers: ReadOnly[Sequence[ConnectFlowServer] | None] requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]] +async def _describe_connect_flow_server( + user_id: str, + server: MCPServer, + interactive: bool, + lookup_vendor_credential: LookupVendorCredential, +) -> ConnectFlowServer | Response: + credential: Final = await lookup_vendor_credential(user_id, server.server_id) if interactive else "present" + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + return { + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + + +def _unscoped_or_stale_description(flow: _ConnectFlow, state: Literal["unscoped", "stale"]) -> ConnectFlowDescription: + return { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None, + "server_name": None, + "connected": None, + "servers": None, + } + + +async def _describe_multi_server_flow( + flow: _ConnectFlow, + server_ids: tuple[str, ...], + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + targets: Final = tuple( + [await _flow_target(flow.user_id, server_id, lookup_server_reachability) for server_id in server_ids] + ) + if any(server is None for _, server in targets): + return _unscoped_or_stale_description(flow, "stale") + described: Final = tuple( + [ + await _describe_connect_flow_server(flow.user_id, server, state == "interactive", lookup_vendor_credential) + for state, server in targets + if server is not None + ] + ) + unavailable: Final = next((item for item in described if isinstance(item, Response)), None) + if unavailable is not None: + return unavailable + return { + "state": "multi", + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None, + "server_name": None, + "connected": None, + "servers": [item for item in described if not isinstance(item, Response)], + } + + async def _describe_opened_flow( flow: _ConnectFlow, lookup_vendor_credential: LookupVendorCredential, @@ -880,7 +1046,10 @@ async def _describe_opened_flow( ) -> ConnectFlowDescription | Response: if flow.upstream_scopes: return await _describe_upstream_scope_flow(flow, lookup_vendor_credential, lookup_server_reachability) - state, server = await _flow_target(flow.user_id, flow.resource_server_id, lookup_server_reachability) + server_ids: Final = flow.resource_server_ids or () + if len(server_ids) > 1: + return await _describe_multi_server_flow(flow, server_ids, lookup_vendor_credential, lookup_server_reachability) + state, server = await _flow_target(flow.user_id, next(iter(server_ids), None), lookup_server_reachability) if state == "interactive" and server is not None: return await _describe_interactive_target(flow, server, flow.requested_scopes, lookup_vendor_credential) described: Final[ConnectFlowDescription] = { @@ -889,6 +1058,7 @@ async def _describe_opened_flow( "server_id": None if server is None else server.server_id, "server_name": None if server is None else (server.server_name or server.alias or server.name), "connected": state == "m2m" or None, + "servers": None, } return described @@ -912,6 +1082,7 @@ async def _describe_interactive_target( "server_id": server.server_id, "server_name": server.server_name or server.alias or server.name, "connected": credential == "present", + "servers": None, } return {**described, "requested_scopes": scopes} if scopes else described @@ -946,15 +1117,12 @@ async def _describe_upstream_scope_flow( for description in descriptions: if isinstance(description, Response) or not description["connected"]: return description - if flow.resource_server_id is not None: - return descriptions[0] - return { - "state": "unscoped", - "client_origin": _origin_only(flow.redirect_uri), - "server_id": None, - "server_name": None, - "connected": None, - } + sealed: Final = flow.resource_server_ids + if sealed is None: + return _unscoped_or_stale_description(flow, "unscoped") + if len(sealed) > 1: + return await _describe_multi_server_flow(flow, sealed, lookup_vendor_credential, lookup_server_reachability) + return descriptions[0] async def describe_connect_flow( @@ -1006,6 +1174,8 @@ async def complete_connect_flow( return described if described["state"] == "stale": return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + if described["state"] == "multi" and not all(server["connected"] for server in described["servers"] or ()): + return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") if described["connected"] is False: return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") flow_refusal: Final = _claim_refusal( @@ -1026,12 +1196,12 @@ async def complete_connect_flow( return response -def _state_param(flow: _ConnectFlow) -> tuple[tuple[str, str], ...]: - return (("state", flow.state),) if flow.state else () +def _state_param(state: str) -> tuple[tuple[str, str], ...]: + return (("state", state),) if state else () def _denied_flow_response(flow: _ConnectFlow) -> Response: - params: Final = (("error", "access_denied"), *_state_param(flow)) + params: Final = (("error", "access_denied"), *_state_param(flow.state)) return RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303) @@ -1048,12 +1218,12 @@ def _approved_flow_response(flow: _ConnectFlow, delivery: str | None, team_id: s jti=secrets.token_urlsafe(24), iat=int(now.timestamp()), exp=int(now.timestamp()) + code_ttl, - resource_server_id=flow.resource_server_id, + resource_server_ids=flow.resource_server_ids, audience=flow.audience, team_id=(team_id or None) if flow.audience == PROXY_API_AUDIENCE else None, ), ) - callback_url: Final = _append_query_params(flow.redirect_uri, (("code", code), *_state_param(flow))) + callback_url: Final = _append_query_params(flow.redirect_uri, (("code", code), *_state_param(flow.state))) if manual_delivery: return _manual_delivery_response(callback_url) return RedirectResponse(callback_url, status_code=303) @@ -1313,17 +1483,19 @@ def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response: def _resource_conflicts_with_scope( - request: Request, resource: str | None, sealed_resource_server_id: str | None + request: Request, resource: str | None, sealed_resource_server_ids: tuple[str, ...] | None ) -> bool: """True when a scoped grant is being redeemed for a DIFFERENT resource than the one sealed into it (RFC 8707 section 2.2: reject with ``invalid_target``). An absent - ``resource`` never conflicts (the sealed scope still binds the minted session), and an - unscoped grant ignores the parameter entirely, exactly as the endpoint always has, so - no pre-existing client breaks.""" - if sealed_resource_server_id is None or resource is None: + ``resource`` never conflicts (the sealed scope still binds the minted session), the + aggregate resource does not conflict with a single-server grant, and an unscoped grant + ignores the parameter entirely.""" + if sealed_resource_server_ids is None or resource is None: + return False + if _is_aggregate_resource(request, resource): return False resolved: Final = resolve_scoped_resource_server(request, resource) - return resolved is None or resolved.server_id != sealed_resource_server_id + return resolved is None or sealed_resource_server_ids != (resolved.server_id,) async def aggregate_token( @@ -1545,7 +1717,7 @@ async def _authorization_code_grant( return _oauth_error(400, "invalid_grant", "the authorization code has expired") if client_id != parsed.client_id or redirect_uri != parsed.redirect_uri: return _oauth_error(400, "invalid_grant", "the authorization code was issued to a different client") - if _resource_conflicts_with_scope(request, resource, parsed.resource_server_id): + if _resource_conflicts_with_scope(request, resource, parsed.resource_server_ids): return _oauth_error(400, "invalid_target", "resource does not match the scope this code was issued for") if not _pkce_verifier_matches(code_verifier, parsed.code_challenge): return _oauth_error(400, "invalid_grant", "PKCE verification failed") @@ -1555,7 +1727,7 @@ async def _authorization_code_grant( SessionPrincipal( user_id=parsed.user_id, client_id=client_id, - resource_server_id=parsed.resource_server_id, + resource_server_ids=parsed.resource_server_ids, audience=parsed.audience, team_id=parsed.team_id, ), @@ -1587,7 +1759,7 @@ async def _refresh_token_grant( opened: Final = _open_presented_refresh_token(refresh_token, client_id, signing) if isinstance(opened, Response): return opened - if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_id): + if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_ids): return _oauth_error(400, "invalid_target", "resource does not match the scope this token was issued for") # Refresh-token rotation (OAuth 2.0 Security BCP section 4.13): the presented refresh token is # single-use, so a captured or replayed refresh token cannot mint a second pair after the @@ -1713,7 +1885,18 @@ def _active_introspection_response(opened: OpenedSessionToken) -> Response: for key, value in ( ("token_type", "Bearer" if opened.kind == "session" else None), ("team_id", principal.team_id), - ("resource_server_id", principal.resource_server_id), + ( + "resource_server_id", + principal.resource_server_ids[0] + if principal.resource_server_ids is not None and len(principal.resource_server_ids) == 1 + else None, + ), + ( + "resource_server_ids", + list(principal.resource_server_ids) + if principal.resource_server_ids is not None and len(principal.resource_server_ids) > 1 + else None, + ), ("audience", principal.audience), ) if value is not None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4e4741b2d17..fe34e18aecd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -196,6 +196,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( normalize_server_name, openapi_tool_name, parse_admin_env_vars, + server_answers_to, strip_known_server_prefix, validate_mcp_server_name, ) @@ -3616,8 +3617,8 @@ class MCPServerManager: return open_ids @staticmethod - def _admitted_session_resource_scope(user_api_key_auth: UserAPIKeyAuth | None) -> str | None: - """The single server an admitted session subject's bearer was scoped to at authorize + def _admitted_session_resource_scope(user_api_key_auth: UserAPIKeyAuth | None) -> frozenset[str] | None: + """The servers an admitted session subject's bearer was scoped to at authorize time (RFC 8707 resource), or None for every other principal shape and for unscoped sessions. Read at every return path of :meth:`get_allowed_mcp_servers`, including the exception fallback, and applied AFTER every union (grants, operator-open, @@ -3625,7 +3626,8 @@ class MCPServerManager: fault therefore never widens a scoped bearer to the allow-all set.""" if user_api_key_auth is None or not is_mcp_admitted_user_subject(user_api_key_auth): return None - return user_api_key_auth.mcp_session_resource_server_id + resource_server_ids: Final = user_api_key_auth.mcp_session_resource_server_ids + return None if resource_server_ids is None else frozenset(resource_server_ids) async def get_allowed_mcp_servers( self, @@ -3740,7 +3742,7 @@ class MCPServerManager: if len(combined_servers) == 0: verbose_logger.debug("No allowed MCP Servers found for user api key auth.") scope = MCPServerManager._admitted_session_resource_scope(user_api_key_auth) - return [server_id for server_id in combined_servers if scope is None or server_id == scope] + return [server_id for server_id in combined_servers if scope is None or server_id in scope] except Exception: # noqa: BLE001 verbose_logger.exception( "Failed to get allowed MCP servers; team-level object_permission " @@ -3750,7 +3752,7 @@ class MCPServerManager: return [ server_id for server_id in dict.fromkeys(allow_all_server_ids + submitted_server_ids) - if not explicit_grants_only and (scope is None or server_id == scope) + if not explicit_grants_only and (scope is None or server_id in scope) ] @with_service_target(MCP_SERVERS_TARGET) @@ -7329,6 +7331,19 @@ class MCPServerManager: return server return None + def get_mcp_server_by_identifier(self, identifier: str, client_ip: str | None = None) -> MCPServer | None: + server: Final = self.get_mcp_server_by_name(identifier, client_ip) + if server is not None: + return server + return next( + ( + server + for server in self.get_filtered_registry(client_ip).values() + if server_answers_to(server, identifier) + ), + None, + ) + def get_filtered_registry(self, client_ip: str | None = None) -> Mapping[str, MCPServer]: """ Get registry filtered by client IP access control. diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 5047948350c..43435a6fcb2 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -146,6 +146,9 @@ from litellm.proxy._experimental.mcp_server.utils import ( split_server_prefix_from_name, strip_known_server_prefix, ) +from litellm.proxy._experimental.mcp_server.utils import ( + server_answers_to as _server_answers_to, +) from litellm.proxy._types import ( UserAPIKeyAuth, ) @@ -578,9 +581,36 @@ def _http_detail_message(detail: object) -> str: return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) -def _server_answers_to(server: MCPServer, name: str) -> bool: - requested: Final = name.lower() - return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) +_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def raise_if_unscoped_aggregate_request( + mcp_servers: Sequence[str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + toolset_id: str | None, +) -> None: + from litellm.proxy import proxy_server # noqa: PLC0415 # proxy import cycle + + raw_settings: Final[object] = getattr(proxy_server, "general_settings", None) + if not raw_settings: + return + settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(raw_settings) + if not settings.get("mcp_require_explicit_server_scope", False): + return + if mcp_servers or toolset_id is not None: + return + if user_api_key_auth is not None and user_api_key_auth.mcp_session_resource_server_ids is not None: + return + raise HTTPException( + status_code=400, + detail={ + "error": "mcp_server_scope_required", + "message": ( + "This gateway requires an explicit MCP server scope: connect to /mcp/ " + "or send the x-mcp-servers header." + ), + }, + ) async def raise_denied_scoped_mcp_access( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 319ce81d619..fb475cf0e3b 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -37,6 +37,7 @@ from __future__ import annotations import secrets from collections import Counter +from collections.abc import Sequence from datetime import datetime, timedelta from functools import lru_cache from typing import Final, Literal, TypeAlias @@ -110,20 +111,27 @@ class SessionPrincipal(LiteLLMBaseModel): gateway-sealed) DCR client identifier the token was issued to; the token endpoint requires it to match on the refresh grant. - ``resource_server_id`` is the single MCP server this session was authorized for when - the client requested a per-server RFC 8707 resource at authorize time, or ``None`` for - the aggregate scope. It is a RESTRICTION carried for admission to intersect against - the live grant resolution, never a grant by itself; the refresh grant re-mints from - this principal so the restriction survives rotation. + ``resource_server_ids`` are the MCP servers this session was authorized for when the + client requested an RFC 8707 resource or gateway scope at authorize time, or ``None`` + for the aggregate scope. They are a RESTRICTION carried for admission to intersect + against the live grant resolution, never a grant by itself; the refresh grant re-mints + from this principal so the restriction survives rotation. """ model_config = ConfigDict(frozen=True) user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) - resource_server_id: str | None = None + resource_server_ids: tuple[str, ...] | None = Field(default=None, min_length=1) audience: SessionAudience | None = None team_id: str | None = None + @field_validator("resource_server_ids") + @classmethod + def _resource_server_ids_are_nonempty(cls, value: tuple[str, ...] | None) -> tuple[str, ...] | None: + if value is not None and not value: + raise ValueError("resource_server_ids must not be empty") + return value + class SessionKeys(LiteLLMBaseModel): """Injected key material: the HS256 signing key. @@ -307,6 +315,7 @@ class _SessionClaims(LiteLLMBaseModel): user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) resource_server_id: str | None = None + resource_server_ids: Sequence[str] | None = Field(default=None, min_length=2) audience: SessionAudience | None = None team_id: str | None = None family: str | None = Field(default=None, min_length=1) @@ -414,7 +423,16 @@ def _mint( kind=kind, user_id=principal.user_id, client_id=principal.client_id, - resource_server_id=principal.resource_server_id, + resource_server_id=( + principal.resource_server_ids[0] + if principal.resource_server_ids is not None and len(principal.resource_server_ids) == 1 + else None + ), + resource_server_ids=( + list(principal.resource_server_ids) + if principal.resource_server_ids is not None and len(principal.resource_server_ids) >= 2 + else None + ), audience=principal.audience, team_id=principal.team_id, family=family, @@ -465,13 +483,21 @@ def _open( return claims if claims.kind != expected_kind: return SessionMalformed() + if claims.resource_server_id is not None and claims.resource_server_ids is not None: + return SessionMalformed() if now.timestamp() >= claims.exp: return SessionExpired() return OpenedSessionToken( principal=SessionPrincipal( user_id=claims.user_id, client_id=claims.client_id, - resource_server_id=claims.resource_server_id, + resource_server_ids=( + (claims.resource_server_id,) + if claims.resource_server_id is not None + else tuple(claims.resource_server_ids) + if claims.resource_server_ids is not None + else None + ), audience=claims.audience, team_id=claims.team_id, ), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 76156e9878e..b3426089375 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1272,6 +1272,14 @@ if MCP_AVAILABLE: verbose_logger.warning("Rejected MCP request from a disallowed client application: %s", rejection.details) raise HTTPException(status_code=403, detail=rejection.response_body) + def admit_mcp_request( + scope: Scope, user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: Sequence[str] | None + ) -> None: + reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) + operations.raise_if_unscoped_aggregate_request( + mcp_servers=mcp_servers, user_api_key_auth=user_api_key_auth, toolset_id=mcp_active_toolset_id.get() + ) + async def extract_mcp_auth_context(scope, path): """ Extracts mcp_servers from the path and processes the MCP request for auth context. @@ -1776,8 +1784,8 @@ if MCP_AVAILABLE: detail=f"API key does not have access to toolset '{toolset_id}'.", ) if is_mcp_admitted_user_subject(acting): - resource_server_id: Final = acting.mcp_session_resource_server_id - if resource_server_id is not None and resource_server_id not in ( + resource_server_ids: Final = acting.mcp_session_resource_server_ids + if resource_server_ids is not None and frozenset(resource_server_ids).isdisjoint( await operations.global_mcp_server_manager.resolve_toolset_tool_permissions( toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy ) @@ -2413,7 +2421,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) + admit_mcp_request(scope, user_api_key_auth, mcp_servers) body: Final = await peek.body() if peek is not None else b"" is_initialize: Final = _is_initialize_request(body) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 @@ -2780,7 +2788,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) + admit_mcp_request(scope, user_api_key_auth, mcp_servers) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index b2e3eae8e59..1684b9c5c30 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -379,6 +379,11 @@ def iter_known_server_prefixes(server: _McpServerLike) -> Iterator[str]: yield from _emit(server_id) +def server_answers_to(server: _McpServerLike, name: str) -> bool: + requested: Final = name.lower() + return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) + + def iter_known_tool_name_spellings(tool_name: str, server: MCPServer) -> Iterator[str]: """Yield every name that denotes the bare ``tool_name`` on ``server``: the bare name, then its wire spelling under each prefix ``iter_known_server_prefixes`` accepts. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 682a5a56ee9..c3bb76e0f28 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2897,6 +2897,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.", ) + mcp_require_explicit_server_scope: bool = Field( + default=False, + description="Require aggregate MCP requests to include x-mcp-servers unless a toolset or gateway OAuth session already scopes the server. Off by default.", + ) mcp_client_id_header: str | None = Field( None, description="Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs.", diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 18dd4e53701..c2f7c487b7b 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -168,7 +168,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # path via post-construction assignment and stripped from validated input like the markers # above; a forged value could at most narrow, but the stripping keeps the field's provenance # single-owner so its meaning stays trustworthy. - mcp_session_resource_server_id: str | None = Field(default=None, exclude=True) + mcp_session_resource_server_ids: tuple[str, ...] | None = Field(default=None, exclude=True) mcp_toolset_id: str | None = Field(default=None, exclude=True) authenticated_by_custom_auth: bool = Field(default=False, exclude=True) via_virtual_key: bool = Field( @@ -229,7 +229,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("requires_fresh_policy", None) values.pop("mcp_explicit_grants_only", None) values.pop("mcp_source_team_rpm_limits", None) - values.pop("mcp_session_resource_server_id", None) + values.pop("mcp_session_resource_server_ids", None) values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) values.pop("authenticated_by_custom_auth", None) diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0f4600d2a37..c96e759b1a3 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7789,6 +7789,73 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" + @pytest.mark.parametrize( + ("header", "expected_scope"), + [ + ("alpha,beta", 'scope="litellm:mcp_server:alpha litellm:mcp_server:beta"'), + ("ALPHA", 'scope="litellm:mcp_server:ALPHA"'), + ("alpha-id", 'scope="litellm:mcp_server:alpha-id"'), + ("alpha,unknown", None), + ("access-group", None), + ("token-exchange", None), + ], + ) + async def test_aggregate_header_challenge_scopes_resolved_gateway_servers( + self, header: str, expected_scope: str | None + ) -> None: + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + alpha: Final = MCPServer( + server_id="alpha-id", + name="alpha", + server_name="alpha", + alias="alpha", + url="https://upstream.example/mcp", + transport="http", + auth_type=MCPAuth("oauth2"), + ) + beta: Final = MCPServer( + server_id="beta-id", + name="beta", + server_name="beta", + alias="beta", + url="https://upstream.example/mcp", + transport="http", + auth_type=MCPAuth("oauth2"), + ) + token_exchange: Final = MCPServer( + server_id="token-exchange-id", + name="token-exchange", + server_name="token-exchange", + alias="token-exchange", + url="https://upstream.example/mcp", + transport="http", + auth_type=MCPAuth("oauth2_token_exchange"), + ) + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as manager, + ): + manager.get_mcp_server_by_identifier.side_effect = lambda identifier, client_ip=None: { + "alpha": alpha, + "alpha-id": alpha, + "beta": beta, + "token-exchange": token_exchange, + }.get(identifier.lower()) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-servers", header.encode()),)) + ) + + expected: Final = ( + f"Bearer {self._EXPECTED_RESOURCE_METADATA}" + if expected_scope is None + else f"Bearer {self._EXPECTED_RESOURCE_METADATA}, {expected_scope}" + ) + assert (exc_info.value.headers or {})["WWW-Authenticate"] == expected + assert exc_info.value.status_code == 401 + @pytest.mark.parametrize( "auth_type", (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token", "oauth2"), @@ -7898,6 +7965,7 @@ class TestAggregateGatewayDcrChallenge: ) as mock_mgr, ): mock_mgr.get_mcp_server_by_name.return_value = resolved + mock_mgr.get_mcp_server_by_identifier.return_value = resolved with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( self._scope(path=path, extra_headers=((b"authorization", b"Bearer not-a-key"),)) @@ -7942,12 +8010,12 @@ class TestAggregateGatewayDcrChallenge: with patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" ) as mock_mgr: - mock_mgr.get_mcp_server_by_name.return_value = resolved + mock_mgr.get_mcp_server_by_identifier.return_value = resolved assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None assert _gateway_dcr_challenge_target("/mcp", None, None) is None with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr: - mock_mgr.get_mcp_server_by_name.return_value = _server(MCPAuth.oauth2) + mock_mgr.get_mcp_server_by_identifier.return_value = _server(MCPAuth.oauth2) assert _gateway_dcr_challenge_target("/mcp/srv", ["other"], None) is None async def test_no_challenge_for_path_named_server(self): @@ -10658,7 +10726,7 @@ class TestScopedSessionAdmission: _MASTER_KEY = "sk-scoped-session-admission-master-key" - def _bearer(self, resource_server_id): + def _bearer(self, resource_server_ids): from datetime import datetime, timezone from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( @@ -10671,11 +10739,11 @@ class TestScopedSessionAdmission: keys = session_keys_from_master_key(self._MASTER_KEY) principal = SessionPrincipal( - user_id="scoped-user", client_id="llm_dcrc_abc", resource_server_id=resource_server_id + user_id="scoped-user", client_id="llm_dcrc_abc", resource_server_ids=resource_server_ids ) return mint_session_token(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() - @pytest.mark.parametrize("scope", ["github-server-id", None]) + @pytest.mark.parametrize("scope", [("github-server-id",), None]) async def test_admission_carries_sealed_resource_scope(self, scope): token = self._bearer(scope) scope_dict = { @@ -10708,11 +10776,11 @@ class TestScopedSessionAdmission: ): auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope_dict) assert auth_result.mcp_admitted_user_subject is True - assert auth_result.mcp_session_resource_server_id == scope + assert auth_result.mcp_session_resource_server_ids == scope def test_scope_field_cannot_be_forged_through_construction(self): - forged = UserAPIKeyAuth(user_id="u1", mcp_session_resource_server_id="any-server") - assert forged.mcp_session_resource_server_id is None + forged = UserAPIKeyAuth(user_id="u1", mcp_session_resource_server_ids=("any-server",)) + assert forged.mcp_session_resource_server_ids is None @pytest.mark.asyncio @@ -10789,13 +10857,13 @@ async def test_catalog_refresh_preserves_non_database_admission_and_resource_sco caller = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS if kind == "master" else "custom-subject") caller.via_virtual_key = kind == "master" caller.authenticated_by_custom_auth = kind == "custom" - caller.mcp_session_resource_server_id = "only-this-server" + caller.mcp_session_resource_server_ids = ("only-this-server",) caller.mcp_toolset_id = "only-this-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) assert refreshed is not caller assert refreshed.api_key == caller.api_key assert refreshed.authenticated_by_custom_auth == caller.authenticated_by_custom_auth - assert refreshed.mcp_session_resource_server_id == "only-this-server" + assert refreshed.mcp_session_resource_server_ids == ("only-this-server",) assert refreshed.mcp_toolset_id == "only-this-toolset" assert refreshed.requires_fresh_policy is True assert caller.requires_fresh_policy is False @@ -10818,11 +10886,11 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) caller = UserAPIKeyAuth(user_id="catalog-user", org_id="previous-org", user_role="proxy_admin") caller.mcp_admitted_user_subject = True - caller.mcp_session_resource_server_id = "scoped-server" + caller.mcp_session_resource_server_ids = ("scoped-server",) refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) assert refreshed.org_id == "current-org" assert refreshed.user_role == "internal_user" - assert refreshed.mcp_session_resource_server_id == "scoped-server" + assert refreshed.mcp_session_resource_server_ids == ("scoped-server",) assert refreshed.mcp_admitted_user_subject is True assert caller.org_id == "previous-org" @@ -10851,7 +10919,7 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session access_group_ids=["original-group"], ) caller.via_virtual_key = True - caller.mcp_session_resource_server_id = "session-server" + caller.mcp_session_resource_server_ids = ("session-server",) caller.mcp_toolset_id = "session-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) @@ -10863,7 +10931,7 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session "new-project", "new-owner", ) - assert refreshed.mcp_session_resource_server_id == "session-server" + assert refreshed.mcp_session_resource_server_ids == ("session-server",) assert refreshed.mcp_toolset_id == "session-toolset" assert refreshed.via_virtual_key and refreshed.requires_fresh_policy assert caller.team_id == "old-team" and not caller.requires_fresh_policy diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index bcde8a99566..641aca2446e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -219,6 +219,8 @@ def test_empty_principal_fields_rejected_at_construction(): SessionPrincipal(user_id="", client_id="c") with pytest.raises(ValidationError): SessionPrincipal(user_id="u", client_id="") + with pytest.raises(ValidationError): + SessionPrincipal(user_id="u", client_id="c", resource_server_ids=()) def test_short_signing_key_rejected_at_construction(): @@ -248,6 +250,46 @@ def _decoded_claims(token: str, prefix: str) -> dict: ) +def test_single_server_scope_keeps_the_scalar_wire_claim(): + principal = SessionPrincipal( + user_id="user-123", + client_id="llm_client_abc", + resource_server_ids=("alpha-id",), + ) + minted = mint_session_token(principal, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + + claims = _decoded_claims(minted.token.get_secret_value(), SESSION_TOKEN_PREFIX) + opened = open_session_token(minted.token.get_secret_value(), KEYS, NOW) + assert claims["resource_server_id"] == "alpha-id" + assert "resource_server_ids" not in claims + assert isinstance(opened, OpenedSessionToken) + assert opened.principal.resource_server_ids == ("alpha-id",) + + +def test_multi_server_scope_round_trips_through_the_list_claim(): + principal = SessionPrincipal( + user_id="user-123", + client_id="llm_client_abc", + resource_server_ids=("alpha-id", "beta-id"), + ) + minted = mint_session_token(principal, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + + claims = _decoded_claims(minted.token.get_secret_value(), SESSION_TOKEN_PREFIX) + opened = open_session_token(minted.token.get_secret_value(), KEYS, NOW) + assert claims["resource_server_ids"] == ["alpha-id", "beta-id"] + assert "resource_server_id" not in claims + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == principal + + +def test_token_with_both_server_scope_claims_is_malformed(): + token = _sign_claims(_valid_claims(resource_server_id="alpha-id", resource_server_ids=["alpha-id", "beta-id"])) + + assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed) + + def test_mcp_principal_wire_claims_carry_no_audience_or_team_keys(): access_claims = _decoded_claims(_mint_access(), SESSION_TOKEN_PREFIX) refresh_claims = _decoded_claims(_mint_refresh(), SESSION_REFRESH_PREFIX) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_contracts.py b/tests/unit/proxy/_experimental/mcp_server/test_contracts.py index db877bd4395..adb17590d74 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_contracts.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_contracts.py @@ -22,7 +22,7 @@ def test_guardrail_bearer_stays_out_of_egress_and_context_representations() -> N def test_operation_context_isolates_nested_headers_and_caller_permissions(): caller = UserAPIKeyAuth(user_id="alpha", models=["allowed"]) caller.mcp_admitted_user_subject = True - caller.mcp_session_resource_server_id = "alpha-server" + caller.mcp_session_resource_server_ids = ("alpha-server",) caller.mcp_toolset_id = "toolset-alpha" caller.mcp_source_team_rpm_limits = {"team": {"alpha-server": 2}} headers = {"x-caller": "alpha"} @@ -37,7 +37,7 @@ def test_operation_context_isolates_nested_headers_and_caller_permissions(): assert captured is not None assert captured.models == ["allowed"] assert captured.mcp_admitted_user_subject is True - assert captured.mcp_session_resource_server_id == "alpha-server" + assert captured.mcp_session_resource_server_ids == ("alpha-server",) assert captured.mcp_toolset_id == "toolset-alpha" assert captured.mcp_source_team_rpm_limits == {"team": {"alpha-server": 2}} captured.models.append("also-forbidden") diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index ab45ecb4fbc..aa0166864ce 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -105,6 +105,7 @@ def isolate_global_mcp_registry(monkeypatch): from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.proxy._experimental.mcp_server.catalog import CatalogSnapshots + monkeypatch.setattr(global_mcp_server_manager, "catalog", CatalogSnapshots(global_mcp_server_manager)) snapshot = dict(global_mcp_server_manager.registry) yield @@ -1493,11 +1494,15 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b replace_stale_client=True, ) - expected_binding = { - "dcr_issuer": None, - "dcr_server_url": None, - "dcr_redirect_uris": ["https://proxy.litellm.example/callback"], - } if auth_type == MCPAuth.oauth2 else {} + expected_binding = ( + { + "dcr_issuer": None, + "dcr_server_url": None, + "dcr_redirect_uris": ["https://proxy.litellm.example/callback"], + } + if auth_type == MCPAuth.oauth2 + else {} + ) assert json.loads(response.body.decode("utf-8")) == {**mock_response.json.return_value, **expected_binding} return mock_update.await_count > 0 @@ -1874,7 +1879,9 @@ async def test_register_client_returns_reused_client_when_concurrent_persist_win persisted_server = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) persisted_server.credentials = {"client_id": "persisted-client"} empty_server = persisted_server.model_copy(update={"credentials": None}) - mock_get_mcp_server = AsyncMock(side_effect=[empty_server, empty_server, persisted_server] if late_winner else [empty_server, persisted_server]) + mock_get_mcp_server = AsyncMock( + side_effect=[empty_server, empty_server, persisted_server] if late_winner else [empty_server, persisted_server] + ) mock_update_mcp_server = AsyncMock(return_value=None) mock_update_server = AsyncMock() @@ -3911,9 +3918,15 @@ async def test_authorize_resolves_server_by_id_when_name_lookup_fails(): request.headers = {} with ( - patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam - patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam - patch.object(discoverable_endpoints, "encrypt_value_helper", return_value="encrypted-state"), # test-quality-ok: flow seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_name", return_value=None + ) as by_name, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_id", return_value=server + ) as by_id, # test-quality-ok: resolver seam + patch.object( + discoverable_endpoints, "encrypt_value_helper", return_value="encrypted-state" + ), # test-quality-ok: flow seam ): response = await discoverable_endpoints.authorize( request=request, @@ -3947,9 +3960,15 @@ async def test_token_endpoint_resolves_server_by_id_when_name_lookup_fails(): client.post = AsyncMock(return_value=response) with ( - patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam - patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam - patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_name", return_value=None + ) as by_name, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_id", return_value=server + ) as by_id, # test-quality-ok: resolver seam + patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=client + ), # test-quality-ok: HTTP seam ): result = await discoverable_endpoints.token_endpoint( request=request, @@ -3986,10 +4005,18 @@ async def test_register_client_resolves_server_by_id_when_name_lookup_fails(): client.post = AsyncMock(return_value=response) with ( - patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam - patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam - patch.object(discoverable_endpoints, "read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: request seam - patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_name", return_value=None + ) as by_name, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_id", return_value=server + ) as by_id, # test-quality-ok: resolver seam + patch.object( + discoverable_endpoints, "read_request_body", new=AsyncMock(return_value={}) + ), # test-quality-ok: request seam + patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=client + ), # test-quality-ok: HTTP seam ): result = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_id) @@ -4011,8 +4038,12 @@ async def test_protected_resource_metadata_resolves_server_by_id_when_name_looku request.headers = {} with ( - patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam - patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_name", return_value=None + ) as by_name, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_id", return_value=server + ) as by_id, # test-quality-ok: resolver seam ): result = await discoverable_endpoints._build_oauth_protected_resource_response( request=request, @@ -4038,8 +4069,12 @@ def test_authorization_server_metadata_resolves_server_by_id_when_name_lookup_fa request.headers = {} with ( - patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam - patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_name", return_value=None + ) as by_name, # test-quality-ok: resolver seam + patch.object( + global_mcp_server_manager, "get_mcp_server_by_id", return_value=server + ) as by_id, # test-quality-ok: resolver seam ): result = discoverable_endpoints._build_oauth_authorization_server_response( request=request, @@ -4882,7 +4917,9 @@ async def test_callback_rejects_supplied_issuer_without_expected_identity(monkey @pytest.mark.parametrize("issuer", [None, "https://idp.example.com"]) async def test_callback_accepts_missing_unadvertised_issuer(monkeypatch, issuer): response, _ = await _authorize_then_callback( - _issuer_anchored_oauth_server(issuer), iss=None, monkeypatch=monkeypatch, + _issuer_anchored_oauth_server(issuer), + iss=None, + monkeypatch=monkeypatch, ) assert response.status_code == 302 @@ -8275,9 +8312,11 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ proxy_globals.user_api_key_cache = UserApiKeyCache() proxy_globals.prisma_client = object() - with patch( # test-quality-ok: get_user_object is the DB seam that wraps the fault; same patch as the outage sibling - "litellm.proxy.auth.auth_checks.get_user_object", - new=AsyncMock(side_effect=_wrapped_user_lookup_error(MismatchedVersionsError(expected="1", got="2"))), + with ( + patch( # test-quality-ok: get_user_object is the DB seam that wraps the fault; same patch as the outage sibling + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(MismatchedVersionsError(expected="1", got="2"))), + ) ): assert await reload_active_user_by_id("sso-user-7") == "faulted" @@ -8470,7 +8509,6 @@ async def test_token_endpoint_client_secret_basic_without_secret_returns_400(): assert exc_info.value.status_code == 400 - def _access_group_none_server(server_name="access_group_server"): """A non-oauth2, access-group gated MCP server: no client_id, no OAuth.""" from litellm.proxy._types import MCPTransport @@ -8622,9 +8660,7 @@ async def test_oauth_protected_resource_for_gateway_owned_auth(auth_type, use_st ) server = _access_group_none_server().model_copy(update={"auth_type": auth_type}) - request = Request( - {"type": "http", "scheme": "https", "path": "/", "headers": [(b"host", b"litellm.example.com")]} - ) + request = Request({"type": "http", "scheme": "https", "path": "/", "headers": [(b"host", b"litellm.example.com")]}) global_mcp_server_manager.registry.clear() global_mcp_server_manager.registry[server.server_id] = server try: @@ -9548,7 +9584,12 @@ async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch, lega monkeypatch.setattr(enc, "get_salt_key", lambda: "salt-hydrate-key") stored_blob = safe_dumps( encrypt_credentials( - credentials={**({} if legacy else {"dcr_issuer": "https://idp.example", "dcr_server_url": "https://resource.example/mcp"}), + credentials={ + **( + {} + if legacy + else {"dcr_issuer": "https://idp.example", "dcr_server_url": "https://resource.example/mcp"} + ), "client_id": "stored-client", "client_secret": "stored-secret", "token_endpoint_auth_method": "client_secret_basic", @@ -9605,7 +9646,13 @@ async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): monkeypatch.setattr(enc, "get_salt_key", lambda: "salt-reuse-key") blob = safe_dumps( encrypt_credentials( - credentials={"dcr_issuer": "https://idp.example", "dcr_server_url": "https://resource.example/mcp","client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + credentials={ + "dcr_issuer": "https://idp.example", + "dcr_server_url": "https://resource.example/mcp", + "client_id": "stored-client", + "client_secret": "sec", + "redirect_uris": ["https://x/callback"], + }, encryption_key="salt-reuse-key", ) ) @@ -11526,9 +11573,7 @@ class TestPerRequestRootPathDiscovery: get_request_root_path, ) - server = _create_oauth2_server( - server_id="srv_a", name="server_a", server_name="server_a", alias="server_a" - ) + server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a") _isolated_mcp_registry[server.server_id] = server app = FastAPI() @@ -11613,6 +11658,61 @@ def _native_client_app(monkeypatch): return client, session_cookie, minted +@pytest.mark.parametrize("authorize_path", ["/authorize/mcp-session", "/authorize"]) +def test_gateway_authorize_routes_forward_server_scope(monkeypatch: pytest.MonkeyPatch, authorize_path: str) -> None: + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import CONNECT_FLOW_COOKIE_PREFIX + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + monkeypatch.setattr( + global_mcp_server_manager, + "registry", + global_mcp_server_manager.registry.copy(), + ) + client, session_cookie, _ = _native_client_app(monkeypatch) + server: Final = _create_oauth2_server( + server_id="scope-server-id", + name="alpha", + server_name="alpha", + alias="alpha", + ) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + + redirect_uri: Final = "https://client.example/callback" + registration: Final = client.post("/register", json={"redirect_uris": [redirect_uri]}) + assert registration.status_code == 201, registration.text + client_id: Final = registration.json()["client_id"] + response: Final = client.get( + authorize_path, + params={ + "client_id": client_id, + "redirect_uri": redirect_uri, + "state": "scope-test-state", + "code_challenge": _s256("scope-test-verifier"), + "code_challenge_method": "S256", + "response_type": "code", + "resource": "http://testserver/mcp", + "scope": "litellm:mcp_server:alpha", + }, + cookies={"token": session_cookie}, + follow_redirects=False, + ) + assert response.status_code == 303, response.text + + flow: Final = parse_qs(urlparse(response.headers["location"]).query)["connect_flow"][0] + flow_cookies: Final = SimpleCookie() + flow_cookies.load(response.headers["set-cookie"]) + sealed_flow: Final = flow_cookies[f"{CONNECT_FLOW_COOKIE_PREFIX}{flow}"].value + decoded_flow: Final = decrypt_value_helper(sealed_flow, "gateway_connect_flow", return_original_value=False) + assert isinstance(decoded_flow, str) + flow_payload: Final = json.loads(decoded_flow) + assert flow_payload["resource_server_id"] == "scope-server-id" + assert "resource_server_ids" not in flow_payload + + def _consent_flow_handle(page: str) -> str: import re @@ -11856,7 +11956,9 @@ def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serv litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field), ) monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler) - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled, "supported_db_objects": []}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled, "supported_db_objects": []} + ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] @@ -11976,41 +12078,66 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal from urllib.parse import parse_qs, urlparse from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _oauth_state_cookie_name, authorize_with_server, callback, open_bridge_authorization_code, + _oauth_state_cookie_name, + authorize_with_server, + callback, + open_bridge_authorization_code, ) from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt") server = MCPServer( - server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, - client_id="client", authorization_url="https://idp.example.com/authorize", + server_id="srv", + name="srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="client", + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", oauth_identity_binding=MCPOAuthIdentityBinding( - mode="enforce", issuer="https://idp.example.com", audiences=["client"], + mode="enforce", + issuer="https://idp.example.com", + audiences=["client"], ), ) - request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443), - "path": "/authorize", "query_string": b"", "headers": [(b"authorization", b"Bearer sk-alice")]}) + request = Request( + { + "type": "http", + "scheme": "https", + "server": ("proxy.example.com", 443), + "path": "/authorize", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-alice")], + } + ) with ( patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.authorize_oauth_credential_request", - new=AsyncMock(return_value="alice")), + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.authorize_oauth_credential_request", + new=AsyncMock(return_value="alice"), + ), patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._user_can_reach_mcp_server", - new=AsyncMock(return_value=True)), + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._user_can_reach_mcp_server", + new=AsyncMock(return_value=True), + ), ): authorized = await authorize_with_server( - request, server, "client", "http://127.0.0.1:6274/callback", state="client-state", - code_challenge="pkce-challenge", code_challenge_method="S256", + request, + server, + "client", + "http://127.0.0.1:6274/callback", + state="client-state", + code_challenge="pkce-challenge", + code_challenge_method="S256", ) query = parse_qs(urlparse(authorized.headers["location"]).query) assert len(query["nonce"][0]) >= 32 cookies = SimpleCookie() cookies.load(authorized.headers["set-cookie"]) name = _oauth_state_cookie_name(query["state"][0]) - callback_request = Request({**request.scope, "path": "/callback", - "headers": [(b"cookie", f"{name}={cookies[name].value}".encode())]}) + callback_request = Request( + {**request.scope, "path": "/callback", "headers": [(b"cookie", f"{name}={cookies[name].value}".encode())]} + ) completed = await callback(callback_request, code="upstream-code", state=query["state"][0]) returned = parse_qs(urlparse(completed.headers["location"]).query) sealed = open_bridge_authorization_code(returned["code"][0]) @@ -12030,33 +12157,53 @@ async def test_enforced_login_warms_verified_token_readable_without_database_loo from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _store_per_user_token_server_side from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof - from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend + from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import ( + DualCacheTokenCacheBackend, + ) from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import CachedOAuthTokenStore - from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import LazyPerUserOAuthTokenStore + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( + LazyPerUserOAuthTokenStore, + ) from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import V2PerUserTokenStore from litellm.types.mcp_server.mcp_server_manager import MCPServer server = MCPServer( - server_id="srv", name="srv", transport="http", auth_type="oauth2", - oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"], - "caller_field": "user_id", "principal_claim": "sub"}, + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={ + "mode": "enforce", + "issuer": "https://idp.example", + "audiences": ["client"], + "caller_field": "user_id", + "principal_claim": "sub", + }, ) proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") cache = DualCache() monkeypatch.setenv("LITELLM_SALT_KEY", "test-cache-warm-encryption-salt") monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) - monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", SimpleNamespace(invalidate_user_oauth_token_cache=AsyncMock())) + monkeypatch.setattr( + mcp_server_manager, "global_mcp_server_manager", SimpleNamespace(invalidate_user_oauth_token_cache=AsyncMock()) + ) monkeypatch.setattr(db, "store_user_oauth_credential", AsyncMock()) await _store_per_user_token_server_side( - server=server, user_id="alice", token_response={"access_token": "alice-token", "refresh_token": "private", "expires_in": 3600}, + server=server, + user_id="alice", + token_response={"access_token": "alice-token", "refresh_token": "private", "expires_in": 3600}, identity_binding_proof=proof, ) read = AsyncMock(return_value=None) cached = CachedOAuthTokenStore( - V2PerUserTokenStore(read), default_ttl_seconds=300, backend=DualCacheTokenCacheBackend(cache, mcp_per_user_token_cache._codec()), + V2PerUserTokenStore(read), + default_ttl_seconds=300, + backend=DualCacheTokenCacheBackend(cache, mcp_per_user_token_cache._codec()), + ) + store = LazyPerUserOAuthTokenStore( + lambda _: server, store_builder=lambda _: (cached, True), redis_available=lambda: True ) - store = LazyPerUserOAuthTokenStore(lambda _: server, store_builder=lambda _: (cached, True), redis_available=lambda: True) token = await store.fetch("alice", "srv") assert token is not None and token.access_token == "alice-token" assert token.identity_binding_proof == proof @@ -12081,7 +12228,9 @@ async def test_dcr_refusal_is_actionable_without_upstream_body( ) server: Final = _bridge_server( - auth_type=auth_type, dcr_bridge=dcr_bridge, server_id=f"refused-{auth_type}-{dcr_bridge}-{flow}-{upstream_status}", + auth_type=auth_type, + dcr_bridge=dcr_bridge, + server_id=f"refused-{auth_type}-{dcr_bridge}-{flow}-{upstream_status}", client_id=None, ) import respx @@ -12095,8 +12244,12 @@ async def test_dcr_refusal_is_actionable_without_upstream_body( mint_ephemeral_dcr_client(_bridge_mock_request(), server) if flow == "mint" else register_client_with_server( - request=_bridge_mock_request(), mcp_server=server, client_name="Test client", - grant_types=None, response_types=None, token_endpoint_auth_method=None, + request=_bridge_mock_request(), + mcp_server=server, + client_name="Test client", + grant_types=None, + response_types=None, + token_endpoint_auth_method=None, client_redirect_uris=["http://localhost:9999/callback"], ) ) @@ -12405,7 +12558,8 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - extract_user_id_from_request, authorize_oauth_credential_request, + extract_user_id_from_request, + authorize_oauth_credential_request, ) allowed_servers: Final = AsyncMock(return_value=["server-a"]) @@ -12436,7 +12590,8 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( request: Final = _token_request({"Authorization": f"Bearer {bearer}"}) result: Final = ( await authorize_oauth_credential_request(request, "server-a") - if credential_write else await extract_user_id_from_request(request) + if credential_write + else await extract_user_id_from_request(request) ) assert result is None allowed_servers.assert_not_awaited() @@ -12680,7 +12835,9 @@ async def test_oauth_refresh_revalidates_the_same_active_user_rule( @pytest.mark.asyncio @pytest.mark.parametrize("mapped", [False, True]) -@pytest.mark.parametrize("state", ["allowed", "route_denied", "server_denied", "blocked", "expired", "lookup_error", "cancelled"]) +@pytest.mark.parametrize( + "state", ["allowed", "route_denied", "server_denied", "blocked", "expired", "lookup_error", "cancelled"] +) async def test_oauth_credential_write_keeps_virtual_key_permissions( jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], monkeypatch: pytest.MonkeyPatch, @@ -12718,7 +12875,13 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( manager: Final = MagicMock() manager.get_allowed_mcp_servers = AsyncMock( return_value=[] if state == "server_denied" else ["server-a"], - side_effect=(asyncio.CancelledError() if state == "cancelled" else RuntimeError("permission lookup unavailable") if state == "lookup_error" else None), + side_effect=( + asyncio.CancelledError() + if state == "cancelled" + else RuntimeError("permission lookup unavailable") + if state == "lookup_error" + else None + ), ) monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) bearer: Final = _oauth_identity_jwt(signing_key) if mapped else key @@ -12728,7 +12891,9 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( await authorize_oauth_credential_request(request, "server-a") manager.get_allowed_mcp_servers.assert_awaited_once() return - assert await authorize_oauth_credential_request(request, "server-a") == ("jwt-owner" if state == "allowed" else None) + assert await authorize_oauth_credential_request(request, "server-a") == ( + "jwt-owner" if state == "allowed" else None + ) if state in ("allowed", "server_denied", "lookup_error"): manager.get_allowed_mcp_servers.assert_awaited_once() writer: Final = manager.get_allowed_mcp_servers.call_args.args[0] @@ -12785,7 +12950,8 @@ async def test_oauth_writer_preserves_claimed_team_instead_of_expanding_user_ros @pytest.mark.asyncio async def test_oauth_write_denial_does_not_erase_identity_binding( - jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], monkeypatch: pytest.MonkeyPatch, + jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager from litellm.proxy._types import MCPTransport @@ -12797,21 +12963,36 @@ async def test_oauth_write_denial_does_not_erase_identity_binding( manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) server: Final = MCPServer( - server_id="bound-server", name="bound-server", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client", + server_id="bound-server", + name="bound-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="client", token_url="https://upstream.example.test/token", oauth_identity_binding=MCPOAuthIdentityBinding( - mode="enforce", issuer="https://upstream.example.test", audiences=["client"], + mode="enforce", + issuer="https://upstream.example.test", + audiences=["client"], ), ) request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}) code: Final = discoverable_endpoints.seal_bridge_authorization_code( - "upstream-code", "another-owner", server.server_id, "bound-nonce", + "upstream-code", + "another-owner", + server.server_id, + "bound-nonce", ) with pytest.raises(HTTPException) as denied: await discoverable_endpoints.exchange_token_with_server( - request=request, mcp_server=server, grant_type="authorization_code", code=code, - redirect_uri="http://localhost/callback", client_id="client", client_secret=None, code_verifier="verifier", + request=request, + mcp_server=server, + grant_type="authorization_code", + code=code, + redirect_uri="http://localhost/callback", + client_id="client", + client_secret=None, + code_verifier="verifier", ) assert denied.value.status_code == 403 assert denied.value.detail == {"error": "oauth_principal_mismatch"} @@ -12839,15 +13020,23 @@ async def test_signed_oauth_callback_honors_credential_write_policy( from litellm.types.mcp_server.mcp_server_manager import MCPServer server: Final = MCPServer( - server_id="signed-server", name="signed-server", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id=None if cimd else "client", + server_id="signed-server", + name="signed-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id=None if cimd else "client", client_id_metadata_document_supported=cimd, token_url="https://upstream.example.test/token", ) - monkeypatch.setattr(proxy_server, "general_settings", { - "enable_jwt_auth": True, - "admin_only_routes": [f"/v1/mcp/server/{server.server_id}/oauth-user-credential"] if admin_only else [], - }) + monkeypatch.setattr( + proxy_server, + "general_settings", + { + "enable_jwt_auth": True, + "admin_only_routes": [f"/v1/mcp/server/{server.server_id}/oauth-user-credential"] if admin_only else [], + }, + ) monkeypatch.setenv("LITELLM_SALT_KEY", "signed-oauth-test-salt") monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com") manager: Final = MagicMock() @@ -12871,10 +13060,14 @@ async def test_signed_oauth_callback_honors_credential_write_policy( upstream.client = transport clients.set_cache("async_httpx_client" + httpxSpecialProvider.Oauth2Check, upstream) response: Final = await discoverable_endpoints.exchange_token_with_server( - request=_token_request({}, path="/signed-server/token"), mcp_server=server, + request=_token_request({}, path="/signed-server/token"), + mcp_server=server, grant_type="authorization_code", code=discoverable_endpoints.seal_bridge_authorization_code("upstream-code", "jwt-owner", server.server_id), - redirect_uri="http://localhost/callback", client_id="client", client_secret=None, code_verifier=None, + redirect_uri="http://localhost/callback", + client_id="client", + client_secret=None, + code_verifier=None, ) assert response.status_code == 200 assert json.loads(response.body)["access_token"] == "upstream-token" @@ -12883,14 +13076,15 @@ async def test_signed_oauth_callback_honors_credential_write_policy( else: table.upsert.assert_awaited_once() assert table.upsert.call_args.kwargs["where"]["user_id_server_id"] == { - "user_id": "jwt-owner", "server_id": server.server_id, + "user_id": "jwt-owner", + "server_id": server.server_id, } from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper - payload: Final = json.loads(decrypt_value_helper( - table.upsert.call_args.kwargs["data"]["create"]["credential_b64"], "credential_b64" - )) + payload: Final = json.loads( + decrypt_value_helper(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"], "credential_b64") + ) if cimd: assert payload["cimd_client_id"] == "https://gateway.example.com/oauth/client-metadata.json" else: @@ -12899,11 +13093,33 @@ async def test_signed_oauth_callback_honors_credential_write_policy( @pytest.mark.asyncio @pytest.mark.parametrize("allowed", [False, True]) -@pytest.mark.parametrize("credential", [ - "jwt", "key", "expired_jwt", "wrong_audience", "bad_signature", "malformed_jwt", "missing_issuer", - "foreign_explicit", "blank_explicit", "unknown_key", "blocked_key", "expired_key", "opaque_record", - "opaque_outage", "opaque_oidc", "opaque_custom", "foreign_unscoped", "foreign_configured", "encrypted", "invalid_encrypted", "envelope", "master", -]) +@pytest.mark.parametrize( + "credential", + [ + "jwt", + "key", + "expired_jwt", + "wrong_audience", + "bad_signature", + "malformed_jwt", + "missing_issuer", + "foreign_explicit", + "blank_explicit", + "unknown_key", + "blocked_key", + "expired_key", + "opaque_record", + "opaque_outage", + "opaque_oidc", + "opaque_custom", + "foreign_unscoped", + "foreign_configured", + "encrypted", + "invalid_encrypted", + "envelope", + "master", + ], +) async def test_identity_bound_authorize_preserves_presented_jwt_permissions( jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], monkeypatch: pytest.MonkeyPatch, @@ -12931,27 +13147,44 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions( if credential == "foreign_unscoped": monkeypatch.delenv("JWT_ISSUER") if credential == "foreign_configured": - handler.litellm_jwtauth.issuers = [JWTIssuerConfig( - issuer="https://unrelated.example.test", jwks_url="https://idp.example.test/jwks", - audience="litellm-proxy", user_id_jwt_field="identity.user_id", - )] + handler.litellm_jwtauth.issuers = [ + JWTIssuerConfig( + issuer="https://unrelated.example.test", + jwks_url="https://idp.example.test/jwks", + audience="litellm-proxy", + user_id_jwt_field="identity.user_id", + ) + ] proxy_server.prisma_client.get_data = AsyncMock( - return_value=None, side_effect=RuntimeError("database unavailable") if credential == "opaque_outage" else None, + return_value=None, + side_effect=RuntimeError("database unavailable") if credential == "opaque_outage" else None, ) handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner")) key: Final = "opaque-record" if credential == "opaque_record" else "sk-browser-gateway-key" if credential in ("key", "blocked_key", "expired_key", "opaque_record"): - handler.user_api_key_cache.set_cache(hash_token(key), UserAPIKeyAuth( - token=hash_token(key), user_id="jwt-owner", blocked=credential in ("blocked_key", "opaque_record"), - expires=datetime.now(timezone.utc) - timedelta(seconds=60) if credential == "expired_key" else None, - )) + handler.user_api_key_cache.set_cache( + hash_token(key), + UserAPIKeyAuth( + token=hash_token(key), + user_id="jwt-owner", + blocked=credential in ("blocked_key", "opaque_record"), + expires=datetime.now(timezone.utc) - timedelta(seconds=60) if credential == "expired_key" else None, + ), + ) monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt") server: Final = MCPServer( - server_id="bound-server", name="bound-server", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client", - authorization_url="https://upstream.example.test/authorize", token_url="https://upstream.example.test/token", + server_id="bound-server", + name="bound-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="client", + authorization_url="https://upstream.example.test/authorize", + token_url="https://upstream.example.test/token", oauth_identity_binding=MCPOAuthIdentityBinding( - mode="enforce", issuer="https://upstream.example.test", audiences=["client"], + mode="enforce", + issuer="https://upstream.example.test", + audiences=["client"], ), ) manager: Final = MagicMock() @@ -12961,36 +13194,55 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions( ) monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) bearer: Final = ( - key if credential in ("key", "blocked_key", "expired_key", "opaque_record", "unknown_key") - else "opaque-bearer" if credential in ("opaque_outage", "opaque_oidc", "opaque_custom") - else "not.a.jwt" if credential == "malformed_jwt" - else "llm_env_invalid" if credential == "envelope" - else "v2:gcm:invalid" if credential == "invalid_encrypted" - else master if credential == "master" + key + if credential in ("key", "blocked_key", "expired_key", "opaque_record", "unknown_key") + else "opaque-bearer" + if credential in ("opaque_outage", "opaque_oidc", "opaque_custom") + else "not.a.jwt" + if credential == "malformed_jwt" + else "llm_env_invalid" + if credential == "envelope" + else "v2:gcm:invalid" + if credential == "invalid_encrypted" + else master + if credential == "master" else ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( LiteLLM_UserTable(user_id="jwt-owner", user_role="internal_user"), - ) if credential == "encrypted" + ) + if credential == "encrypted" else jwt.encode({"iss": "https://idp.example.test"}, "wrong-signing-key-at-least-32-bytes", algorithm="HS256") if credential == "bad_signature" - else jwt.encode({"sub": "jwt-owner"}, signing_key, algorithm="RS256") if credential == "missing_issuer" + else jwt.encode({"sub": "jwt-owner"}, signing_key, algorithm="RS256") + if credential == "missing_issuer" else _oauth_identity_jwt( signing_key, expires_in=-60 if credential == "expired_jwt" else 300, audience="another-service" if credential == "wrong_audience" else "litellm-proxy", - issuer="https://unrelated.example.test" if credential.startswith("foreign_") or credential == "blank_explicit" else "https://idp.example.test", + issuer="https://unrelated.example.test" + if credential.startswith("foreign_") or credential == "blank_explicit" + else "https://idp.example.test", ) ) cookie: Final = jwt.encode( - {"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + 300}, master, algorithm="HS256", + {"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + 300}, + master, + algorithm="HS256", ) response: Final = await discoverable_endpoints.authorize_with_server( - request=_token_request({ - "Authorization": f"Bearer {bearer}", "Cookie": f"token={cookie}", - **({"x-litellm-api-key": bearer} if credential == "foreign_explicit" else {}), - **({"x-litellm-api-key": ""} if credential == "blank_explicit" else {}), - }), - mcp_server=server, client_id="client", redirect_uri="http://127.0.0.1:6274/callback", - state="client-state", code_challenge="pkce-challenge", code_challenge_method="S256", + request=_token_request( + { + "Authorization": f"Bearer {bearer}", + "Cookie": f"token={cookie}", + **({"x-litellm-api-key": bearer} if credential == "foreign_explicit" else {}), + **({"x-litellm-api-key": ""} if credential == "blank_explicit" else {}), + } + ), + mcp_server=server, + client_id="client", + redirect_uri="http://127.0.0.1:6274/callback", + state="client-state", + code_challenge="pkce-challenge", + code_challenge_method="S256", ) redirect: Final = urlparse(response.headers["location"]) query: Final = parse_qs(redirect.query) @@ -13035,31 +13287,52 @@ async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session( handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner")) proxy_server.prisma_client.get_data = AsyncMock(return_value=None) server: Final = MCPServer( - server_id="bound-server", name="bound-server", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client", - authorization_url="https://upstream.example.test/authorize", token_url="https://upstream.example.test/token", + server_id="bound-server", + name="bound-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="client", + authorization_url="https://upstream.example.test/authorize", + token_url="https://upstream.example.test/token", oauth_identity_binding=MCPOAuthIdentityBinding( - mode="enforce", issuer="https://upstream.example.test", audiences=["client"], + mode="enforce", + issuer="https://upstream.example.test", + audiences=["client"], ), ) manager: Final = MagicMock() - manager.get_allowed_mcp_servers = AsyncMock(return_value=[] if cookie_state == "server_denied" else [server.server_id]) + manager.get_allowed_mcp_servers = AsyncMock( + return_value=[] if cookie_state == "server_denied" else [server.server_id] + ) monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) bearer: Final = ( _oauth_identity_jwt(signing_key, issuer="https://unrelated.example.test") - if credential == "foreign_jwt" else "unrelated-upstream-bearer" + if credential == "foreign_jwt" + else "unrelated-upstream-bearer" ) cookie: Final = jwt.encode( - {"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + (-60 if cookie_state == "expired" else 300)}, - master, algorithm="HS256", + { + "user_id": "cookie-owner", + "login_method": "sso", + "exp": int(time.time()) + (-60 if cookie_state == "expired" else 300), + }, + master, + algorithm="HS256", ) response: Final = await discoverable_endpoints.authorize_with_server( - request=_token_request({ - **({"Authorization": f"Bearer {bearer}"} if credential != "none" else {}), - **({"Cookie": f"token={cookie}"} if cookie_state != "missing" else {}), - }), - mcp_server=server, client_id="client", redirect_uri="http://127.0.0.1:6274/callback", - state="client-state", code_challenge="pkce-challenge", code_challenge_method="S256", + request=_token_request( + { + **({"Authorization": f"Bearer {bearer}"} if credential != "none" else {}), + **({"Cookie": f"token={cookie}"} if cookie_state != "missing" else {}), + } + ), + mcp_server=server, + client_id="client", + redirect_uri="http://127.0.0.1:6274/callback", + state="client-state", + code_challenge="pkce-challenge", + code_challenge_method="S256", ) redirect: Final = urlparse(response.headers["location"]) query: Final = parse_qs(redirect.query) @@ -13111,7 +13384,12 @@ async def test_authorize_observes_committed_peer_server_changes(monkeypatch, cha updated_at=stamp + timedelta(seconds=1), ) read_rows = AsyncMock(return_value=[] if change == "delete" else [row]) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows), litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)))) + prisma = SimpleNamespace( + db=SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=read_rows), + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + ) + ) prisma.writer_db = prisma.db monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr( @@ -13288,28 +13566,35 @@ def test_invalidating_an_idle_server_leaves_no_generation_behind(): @pytest.mark.asyncio async def test_callback_rejects_missing_advertised_issuer(monkeypatch): - server = _issuer_anchored_oauth_server().model_copy( - update={"authorization_response_iss_parameter_supported": True} - ) + server = _issuer_anchored_oauth_server().model_copy(update={"authorization_response_iss_parameter_supported": True}) response, _ = await _authorize_then_callback(server, iss=None, monkeypatch=monkeypatch) assert response.status_code == 400 assert "location" not in response.headers assert b"upstream-auth-code" not in response.body -@pytest.mark.parametrize("issuer,url", [ - ("https://other.example", "https://resource.example/mcp"), - (None, "https://other.example/mcp"), -]) +@pytest.mark.parametrize( + "issuer,url", + [ + ("https://other.example", "https://resource.example/mcp"), + (None, "https://other.example/mcp"), + ], +) def test_persisted_client_is_not_applied_to_another_upstream(issuer, url): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _apply_persisted_dcr_credentials, _PersistedDcrCredentials, + _apply_persisted_dcr_credentials, + _PersistedDcrCredentials, ) + server = _issuer_anchored_oauth_server(issuer).model_copy(update={"url": url, "client_id": None}) - stored = _PersistedDcrCredentials.model_validate({ - "client_id": "old-client", "client_secret": "old-secret", - "dcr_issuer": "https://idp.example.com", "dcr_server_url": "https://resource.example/mcp", - }) + stored = _PersistedDcrCredentials.model_validate( + { + "client_id": "old-client", + "client_secret": "old-secret", + "dcr_issuer": "https://idp.example.com", + "dcr_server_url": "https://resource.example/mcp", + } + ) assert _apply_persisted_dcr_credentials(server, stored) is False assert server.client_id is None assert server.client_secret is None @@ -13322,10 +13607,17 @@ async def test_registration_does_not_write_client_after_upstream_edit(monkeypatc from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _persist_dcr_client_registration - server = _issuer_anchored_oauth_server().model_copy(update={"url": "https://resource.example/mcp", "client_id": None}) + server = _issuer_anchored_oauth_server().model_copy( + update={"url": "https://resource.example/mcp", "client_id": None} + ) current = models.LiteLLM_MCPServerTable.model_construct( - server_id=server.server_id, url="https://other.example/mcp", issuer="https://other.example", - auth_type="oauth2", transport="http", credentials=None, updated_at=datetime.now(timezone.utc), + server_id=server.server_id, + url="https://other.example/mcp", + issuer="https://other.example", + auth_type="oauth2", + transport="http", + credentials=None, + updated_at=datetime.now(timezone.utc), ) prisma = MagicMock() table = AsyncMock() @@ -13333,29 +13625,46 @@ async def test_registration_does_not_write_client_after_upstream_edit(monkeypatc table.update.return_value = current prisma.db.litellm_mcpservertable = table monkeypatch.setattr(proxy_server, "prisma_client", prisma) - result = await _persist_dcr_client_registration(server, {"client_id": "old-issuer-client"}, "https://proxy.example/callback") + result = await _persist_dcr_client_registration( + server, {"client_id": "old-issuer-client"}, "https://proxy.example/callback" + ) assert result == "failed" table.update.assert_not_called() table.update_many.assert_not_called() assert server.client_id is None -@pytest.mark.parametrize("response_issuer", ["https://IDP.example.com", "https://idp.example.com/", "https://idp.example.com:443"]) +@pytest.mark.parametrize( + "response_issuer", ["https://IDP.example.com", "https://idp.example.com/", "https://idp.example.com:443"] +) def test_issuer_validation_uses_exact_identifier(response_issuer): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _authorization_response_issuer_is_trusted - assert _authorization_response_issuer_is_trusted(response_issuer, {"expected_issuer": "https://idp.example.com"}) is False + + assert ( + _authorization_response_issuer_is_trusted(response_issuer, {"expected_issuer": "https://idp.example.com"}) + is False + ) -@pytest.mark.parametrize("issuer,url", [ - ("https://idp.example.com", "https://resource.example/mcp"), - ("https://idp.example.com", "https://other-resource.example/mcp"), -]) +@pytest.mark.parametrize( + "issuer,url", + [ + ("https://idp.example.com", "https://resource.example/mcp"), + ("https://idp.example.com", "https://other-resource.example/mcp"), + ], +) def test_persisted_client_remains_usable_for_its_issuer(issuer, url, monkeypatch): - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _apply_persisted_dcr_credentials, _PersistedDcrCredentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _apply_persisted_dcr_credentials, + _PersistedDcrCredentials, + ) + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") server = _issuer_anchored_oauth_server(issuer).model_copy(update={"url": url, "client_id": None}) credentials = _PersistedDcrCredentials( - client_id="registered-client", client_secret="registered-secret", dcr_issuer=issuer, + client_id="registered-client", + client_secret="registered-secret", + dcr_issuer=issuer, dcr_server_url="https://resource.example/mcp", ) assert _apply_persisted_dcr_credentials(server, credentials) is True @@ -13367,7 +13676,10 @@ def test_persisted_client_remains_usable_for_its_issuer(issuer, url, monkeypatch @pytest.mark.asyncio async def test_callback_does_not_forward_error_from_another_issuer(monkeypatch): response, _ = await _authorize_then_callback( - _issuer_anchored_oauth_server(), iss="https://other.example", monkeypatch=monkeypatch, error="access_denied", + _issuer_anchored_oauth_server(), + iss="https://other.example", + monkeypatch=monkeypatch, + error="access_denied", ) assert response.status_code == 400 assert "location" not in response.headers @@ -13382,22 +13694,36 @@ async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, op from litellm.types.mcp_server.mcp_server_manager import MCPServer server = MCPServer( - server_id="bound-client", name="bound-client", transport="http", auth_type="oauth2", - url="https://new.example/mcp", issuer="https://new.example", client_id="old-client", - dcr_issuer="https://old.example", dcr_server_url="https://old.example/mcp", + server_id="bound-client", + name="bound-client", + transport="http", + auth_type="oauth2", + url="https://new.example/mcp", + issuer="https://new.example", + client_id="old-client", + dcr_issuer="https://old.example", + dcr_server_url="https://old.example/mcp", ) monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) http_client = MagicMock() monkeypatch.setattr(endpoints, "get_async_httpx_client", http_client) request = ( endpoints.authorize_with_server( - mcp_server=server, request=MagicMock(spec=Request), redirect_uri="http://localhost/callback", client_id="old-client", + mcp_server=server, + request=MagicMock(spec=Request), + redirect_uri="http://localhost/callback", + client_id="old-client", ) if operation == "authorize" else endpoints.exchange_token_with_server( - mcp_server=server, request=MagicMock(spec=Request), grant_type="authorization_code", - code="test-code", redirect_uri="http://localhost/callback", client_id="old-client", - client_secret=None, code_verifier="verifier", + mcp_server=server, + request=MagicMock(spec=Request), + grant_type="authorization_code", + code="test-code", + redirect_uri="http://localhost/callback", + client_id="old-client", + client_secret=None, + code_verifier="verifier", ) ) with pytest.raises(HTTPException) as error: @@ -13412,10 +13738,14 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): from fastapi import Request from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints - server = _dcr_redirect_test_server("old-client").model_copy(update={ - "url": "https://new.example/mcp", "issuer": "https://new.example", - "dcr_issuer": "https://old.example", "client_secret": "old-secret", - }) + server = _dcr_redirect_test_server("old-client").model_copy( + update={ + "url": "https://new.example/mcp", + "issuer": "https://new.example", + "dcr_issuer": "https://old.example", + "client_secret": "old-secret", + } + ) monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) monkeypatch.setattr(endpoints, "_reuse_persisted_dcr_client_if_available", AsyncMock(return_value=False)) response = MagicMock() @@ -13449,7 +13779,9 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints monkeypatch.setattr(proxy_server, "prisma_client", None) - server = _dcr_redirect_test_server(None).model_copy(update={"issuer": "https://idp.example", "url": "https://resource.example/mcp"}) + server = _dcr_redirect_test_server(None).model_copy( + update={"issuer": "https://idp.example", "url": "https://resource.example/mcp"} + ) monkeypatch.setattr(endpoints, "_server_with_oauth_endpoints", AsyncMock(return_value=server)) upstream = MagicMock() upstream.json.return_value = {"client_id": "temporary-client", "client_secret": "temporary-secret"} @@ -13478,14 +13810,22 @@ async def test_registration_does_not_overwrite_credentials_after_failed_identity server = _dcr_redirect_test_server(None) monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) - monkeypatch.setattr(db, "get_mcp_server", AsyncMock(return_value=None, side_effect=None if config_store else RuntimeError("unavailable"))) - monkeypatch.setattr(db, "get_mcp_server_oauth_client_credentials", AsyncMock(side_effect=RuntimeError("unavailable"))) + monkeypatch.setattr( + db, + "get_mcp_server", + AsyncMock(return_value=None, side_effect=None if config_store else RuntimeError("unavailable")), + ) + monkeypatch.setattr( + db, "get_mcp_server_oauth_client_credentials", AsyncMock(side_effect=RuntimeError("unavailable")) + ) monkeypatch.setattr(global_mcp_server_manager, "is_config_declared_server", lambda _: config_store) update = AsyncMock() upsert = AsyncMock() monkeypatch.setattr(db, "update_mcp_server", update) monkeypatch.setattr(db, "upsert_mcp_server_oauth_client_credentials", upsert) - result = await endpoints._persist_dcr_client_registration(server, {"client_id": "new-client"}, "https://gateway.example/callback") + result = await endpoints._persist_dcr_client_registration( + server, {"client_id": "new-client"}, "https://gateway.example/callback" + ) assert result == "failed" assert server.client_id is None update.assert_not_awaited() @@ -13501,17 +13841,28 @@ async def test_registration_losing_conditional_write_reuses_only_a_matching_winn server = _dcr_redirect_test_server(None).model_copy(update={"url": "https://upstream.example/mcp"}) row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True)) - winner = row.model_copy(update={"credentials": { - "client_id": "winner-client", "redirect_uris": ["https://gateway.example/callback"], - "dcr_server_url": server.url, - }}) if winner_available else row + winner = ( + row.model_copy( + update={ + "credentials": { + "client_id": "winner-client", + "redirect_uris": ["https://gateway.example/callback"], + "dcr_server_url": server.url, + } + } + ) + if winner_available + else row + ) monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) read = AsyncMock(side_effect=[row, winner]) monkeypatch.setattr(db, "get_mcp_server", read) monkeypatch.setattr(endpoints, "_refresh_persisted_dcr_server", AsyncMock()) update = AsyncMock(return_value=None) monkeypatch.setattr(db, "update_mcp_server", update) - result = await endpoints._persist_dcr_client_registration(server, {"client_id": "losing-client"}, "https://gateway.example/callback") + result = await endpoints._persist_dcr_client_registration( + server, {"client_id": "losing-client"}, "https://gateway.example/callback" + ) assert result == ("reused" if winner_available else "failed") assert update.await_args.kwargs["expected_updated_at"] == row.updated_at assert server.client_id == ("winner-client" if winner_available else None) @@ -13901,27 +14252,46 @@ async def test_optional_cimd_discovery_preserves_the_callers_configured_endpoint discovery: Final = respx_mock.route().respond(503) manager: Final = manager_module.MCPServerManager() monkeypatch.setattr(manager_module, "global_mcp_server_manager", manager) - await manager.load_servers_from_config({"manual": { - "url": "https://mcp.example.com/mcp", "transport": "http", "auth_type": "oauth2", - "oauth2_flow": "authorization_code", - **({"token_url": "https://idp.example.com/token"} if flow == "refresh" - else {"authorization_url": "https://idp.example.com/authorize"}), - }}) + await manager.load_servers_from_config( + { + "manual": { + "url": "https://mcp.example.com/mcp", + "transport": "http", + "auth_type": "oauth2", + "oauth2_flow": "authorization_code", + **( + {"token_url": "https://idp.example.com/token"} + if flow == "refresh" + else {"authorization_url": "https://idp.example.com/authorize"} + ), + } + } + ) server: Final = next(iter(manager.config_mcp_servers.values())) async with manager.catalog.operation(): if flow == "refresh": response: Final = await endpoints.exchange_token_with_server( - request=_cimd_request(), mcp_server=server, grant_type="refresh_token", - refresh_token="existing-refresh", client_id="existing-client", - code=None, redirect_uri=None, client_secret=None, code_verifier=None, + request=_cimd_request(), + mcp_server=server, + grant_type="refresh_token", + refresh_token="existing-refresh", + client_id="existing-client", + code=None, + redirect_uri=None, + client_secret=None, + code_verifier=None, ) assert response.status_code == 200 assert json.loads(response.body)["access_token"] == "upstream-access" assert parse_qs(token.calls[0].request.content.decode())["refresh_token"] == ["existing-refresh"] else: redirect: Final = await endpoints.authorize_with_server( - _cimd_request(), server, "existing-client", "https://gateway.example.com/ui/", - code_challenge="a" * 43, code_challenge_method="S256", + _cimd_request(), + server, + "existing-client", + "https://gateway.example.com/ui/", + code_challenge="a" * 43, + code_challenge_method="S256", ) assert redirect.status_code == 307 assert redirect.headers["location"].startswith("https://idp.example.com/authorize?") @@ -13953,7 +14323,9 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change( table.find_unique = AsyncMock(return_value=None) table.upsert = AsyncMock() await endpoints._store_per_user_token_server_side( - server, "jwt-owner", {"access_token": "expired", "refresh_token": "saved-refresh", "expires_in": -1}, + server, + "jwt-owner", + {"access_token": "expired", "refresh_token": "saved-refresh", "expires_in": -1}, cimd_client_id=identity, ) @@ -13971,7 +14343,14 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change( ) response: Final = await endpoints.exchange_token_with_server( _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(key, scope='litellm_proxy_admin')}"}), - server, "refresh_token", None, None, identity, None, None, refresh_token="saved-refresh", + server, + "refresh_token", + None, + None, + identity, + None, + None, + refresh_token="saved-refresh", ) assert response.status_code == 200 assert parse_qs(upstream.calls[0].request.content.decode())["client_id"] == [identity] @@ -13982,12 +14361,23 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change( @pytest.mark.asyncio -@pytest.mark.parametrize("state", [ - "matching", "unicode_grant", "foreign_refresh", "missing_refresh", "missing_grant", "database_missing", - "database_outage", "static_client", "anonymous", -]) +@pytest.mark.parametrize( + "state", + [ + "matching", + "unicode_grant", + "foreign_refresh", + "missing_refresh", + "missing_grant", + "database_missing", + "database_outage", + "static_client", + "anonymous", + ], +) async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( - monkeypatch: pytest.MonkeyPatch, state: str, + monkeypatch: pytest.MonkeyPatch, + state: str, ) -> None: from types import SimpleNamespace @@ -14005,14 +14395,18 @@ async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( table.upsert = AsyncMock() monkeypatch.setattr(proxy_server, "prisma_client", database) await db.store_user_oauth_credential( - database, "alice", server.server_id, "expired", - refresh_token=None if state == "missing_refresh" else grant, cimd_client_id=identity, + database, + "alice", + server.server_id, + "expired", + refresh_token=None if state == "missing_refresh" else grant, + cimd_client_id=identity, ) table.find_unique.reset_mock() table.find_unique.return_value = ( - None if state == "missing_grant" else SimpleNamespace( - credential_b64=table.upsert.call_args.kwargs["data"]["create"]["credential_b64"] - ) + None + if state == "missing_grant" + else SimpleNamespace(credential_b64=table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) ) if state == "database_missing": monkeypatch.setattr(proxy_server, "prisma_client", None) @@ -14021,7 +14415,8 @@ async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( if state == "static_client": server.client_id = "configured-client" resolved: Final = await endpoints._saved_refresh_credential( - server, None if state == "anonymous" else "alice", + server, + None if state == "anonymous" else "alice", "foreign-refresh" if state == "foreign_refresh" else grant, ) if state in ("matching", "unicode_grant", "static_client"): @@ -14032,12 +14427,18 @@ async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( if state in ("anonymous", "database_missing"): table.find_unique.assert_not_awaited() else: - table.find_unique.assert_awaited_once_with(where={"user_id_server_id": {"user_id": "alice", "server_id": server.server_id}}) + table.find_unique.assert_awaited_once_with( + where={"user_id_server_id": {"user_id": "alice", "server_id": server.server_id}} + ) @pytest.mark.asyncio -@pytest.mark.parametrize("prior_scopes, expected", [(["tools.read"], "tools.read tools.write"), (None, "tools.write"), ([], "tools.write")]) -async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_identity, monkeypatch, prior_scopes, expected): +@pytest.mark.parametrize( + "prior_scopes, expected", [(["tools.read"], "tools.read tools.write"), (None, "tools.write"), ([], "tools.write")] +) +async def test_managed_authorize_retains_the_callers_previous_grant( + jwt_oauth_identity, monkeypatch, prior_scopes, expected +): from types import SimpleNamespace from urllib.parse import parse_qs, urlsplit @@ -14049,13 +14450,22 @@ async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_id monkeypatch.setenv("LITELLM_SALT_KEY", "scope-grant-test-salt") server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client", "scopes": ["tools.read"]}) global_mcp_server_manager.registry[server.server_id] = server - credential = db.encrypt_value_helper(json.dumps({"type": "oauth2", "access_token": "old-access", "scopes": prior_scopes})) + credential = db.encrypt_value_helper( + json.dumps({"type": "oauth2", "access_token": "old-access", "scopes": prior_scopes}) + ) table = proxy_server.prisma_client.db.litellm_mcpusercredentials table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) - request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + request = _token_request( + {"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"} + ) response = await endpoints.authorize_with_server( - request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.write", - code_challenge="challenge", code_challenge_method="S256", + request, + server, + "scope-client", + "http://127.0.0.1:9/cb", + scope="tools.write", + code_challenge="challenge", + code_challenge_method="S256", ) assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] assert table.find_unique.call_args.kwargs["where"] == { @@ -14064,9 +14474,15 @@ async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_id @pytest.mark.asyncio -@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])]) +@pytest.mark.parametrize( + "returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])] +) async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted( - jwt_oauth_identity, monkeypatch, respx_mock, returned_scope, expected, + jwt_oauth_identity, + monkeypatch, + respx_mock, + returned_scope, + expected, ): from http.cookies import SimpleCookie from urllib.parse import parse_qs, urlsplit @@ -14083,10 +14499,17 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted table = proxy_server.prisma_client.db.litellm_mcpusercredentials table.find_unique = AsyncMock(return_value=None) table.upsert = AsyncMock() - request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + request = _token_request( + {"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"} + ) started = await endpoints.authorize_with_server( - request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read tools.write", - code_challenge="challenge", code_challenge_method="S256", + request, + server, + "scope-client", + "http://127.0.0.1:9/cb", + scope="tools.read tools.write", + code_challenge="challenge", + code_challenge_method="S256", ) state = parse_qs(urlsplit(started.headers["location"]).query)["state"][0] cookies = SimpleCookie() @@ -14094,12 +14517,24 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted callback_request = _token_request({"Cookie": "; ".join(f"{name}={value.value}" for name, value in cookies.items())}) returned = await endpoints.callback(callback_request, code="provider-code", state=state, error=None, iss=None) code = parse_qs(urlsplit(returned.headers["location"]).query)["code"][0] - upstream = respx_mock.post(server.token_url).respond(200, json={ - "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", - **({"scope": returned_scope} if returned_scope is not None else {}), - }) + upstream = respx_mock.post(server.token_url).respond( + 200, + json={ + "access_token": "new-access", + "token_type": "Bearer", + "refresh_token": "new-refresh", + **({"scope": returned_scope} if returned_scope is not None else {}), + }, + ) exchanged = await endpoints.exchange_token_with_server( - request, server, "authorization_code", client_id="scope-client", code=code, code_verifier="verifier", redirect_uri="http://127.0.0.1:9/cb", client_secret=None, + request, + server, + "authorization_code", + client_id="scope-client", + code=code, + code_verifier="verifier", + redirect_uri="http://127.0.0.1:9/cb", + client_secret=None, ) assert parse_qs(upstream.calls[0].request.content.decode())["code"] == ["provider-code"] saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) @@ -14109,7 +14544,9 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted @pytest.mark.asyncio @pytest.mark.parametrize("advertised, expected", [(True, "tools.read offline_access"), (False, "tools.read")]) -async def test_managed_authorize_requests_only_selected_scope_and_supported_refresh_access(monkeypatch, advertised, expected): +async def test_managed_authorize_requests_only_selected_scope_and_supported_refresh_access( + monkeypatch, advertised, expected +): from urllib.parse import parse_qs, urlsplit from mcp.shared.auth import OAuthMetadata @@ -14118,14 +14555,21 @@ async def test_managed_authorize_requests_only_selected_scope_and_supported_refr monkeypatch.setenv("LITELLM_SALT_KEY", "metadata-scope-test-salt") metadata = OAuthMetadata( - issuer="https://idp.example.com", authorization_endpoint="https://idp.example.com/authorize", - token_endpoint="https://idp.example.com/token", response_types_supported=["code"], + issuer="https://idp.example.com", + authorization_endpoint="https://idp.example.com/authorize", + token_endpoint="https://idp.example.com/token", + response_types_supported=["code"], scopes_supported=["tools.read", "admin"] + (["offline_access"] if advertised else []), ) server = _cimd_oauth_server().model_copy(update={"authorization_server_metadata": metadata.model_dump(mode="json")}) response = await endpoints.authorize_with_server( - _cimd_request(), server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read", - code_challenge="challenge", code_challenge_method="S256", + _cimd_request(), + server, + "scope-client", + "http://127.0.0.1:9/cb", + scope="tools.read", + code_challenge="challenge", + code_challenge_method="S256", ) assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] @@ -14139,29 +14583,52 @@ async def test_managed_scope_code_cannot_cross_oauth_boundaries(jwt_oauth_identi monkeypatch.setenv("LITELLM_SALT_KEY", "scope-boundary-test-salt") server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) payload = endpoints._ManagedAuthorizationCode( - upstream_code="provider-code", mcp_server_id="other-server" if mismatch == "server" else server.server_id, + upstream_code="provider-code", + mcp_server_id="other-server" if mismatch == "server" else server.server_id, issuer="https://other-issuer.test" if mismatch == "issuer" else server.issuer, - requested_scope="tools.read tools.write", scope_user_id="other-user" if mismatch == "user" else "jwt-owner", + requested_scope="tools.read tools.write", + scope_user_id="other-user" if mismatch == "user" else "jwt-owner", + ) + code = endpoints._MANAGED_AUTH_CODE_PREFIX + ( + "invalid" + if mismatch == "tampered" + else endpoints.encrypt_value_helper("{}" if mismatch == "invalid_payload" else payload.model_dump_json()) + ) + request = _token_request( + {"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"} ) - code = endpoints._MANAGED_AUTH_CODE_PREFIX + ("invalid" if mismatch == "tampered" else endpoints.encrypt_value_helper("{}" if mismatch == "invalid_payload" else payload.model_dump_json())) - request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) if mismatch == "mode": server.auth_type = MCPAuth.true_passthrough with pytest.raises(HTTPException) as raised: await endpoints.exchange_token_with_server( - request, server, "authorization_code", code=code, redirect_uri="http://127.0.0.1:9/cb", - client_id="scope-client", client_secret=None, code_verifier="verifier", + request, + server, + "authorization_code", + code=code, + redirect_uri="http://127.0.0.1:9/cb", + client_id="scope-client", + client_secret=None, + code_verifier="verifier", ) assert raised.value.status_code == 400 assert respx_mock.calls.call_count == 0 @pytest.mark.asyncio -@pytest.mark.parametrize("scope, requested_scope, expected", [ - (None, None, ["tools.read", "tools.write"]), ("tools.read", None, ["tools.read"]), ("", None, []), - (None, "tools.write", ["tools.write"]), ("tools.read", "tools.write", ["tools.read"]), ("", "tools.write", []), -]) -async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch, respx_mock, scope, requested_scope, expected): +@pytest.mark.parametrize( + "scope, requested_scope, expected", + [ + (None, None, ["tools.read", "tools.write"]), + ("tools.read", None, ["tools.read"]), + ("", None, []), + (None, "tools.write", ["tools.write"]), + ("tools.read", "tools.write", ["tools.read"]), + ("", "tools.write", []), + ], +) +async def test_managed_refresh_retains_only_the_actual_grant( + jwt_oauth_identity, monkeypatch, respx_mock, scope, requested_scope, expected +): from types import SimpleNamespace from litellm.proxy import proxy_server @@ -14173,25 +14640,48 @@ async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) global_mcp_server_manager.registry[server.server_id] = server - credential = db.encrypt_value_helper(json.dumps({ - "type": "oauth2", "access_token": "old-access", "refresh_token": "owner-refresh", "scopes": ["tools.read", "tools.write"], - })) + credential = db.encrypt_value_helper( + json.dumps( + { + "type": "oauth2", + "access_token": "old-access", + "refresh_token": "owner-refresh", + "scopes": ["tools.read", "tools.write"], + } + ) + ) table = proxy_server.prisma_client.db.litellm_mcpusercredentials table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) table.upsert = AsyncMock() - upstream = respx_mock.post(server.token_url).respond(200, json={ - "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", - **({"scope": scope} if scope is not None else {}), - }) - request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + upstream = respx_mock.post(server.token_url).respond( + 200, + json={ + "access_token": "new-access", + "token_type": "Bearer", + "refresh_token": "new-refresh", + **({"scope": scope} if scope is not None else {}), + }, + ) + request = _token_request( + {"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"} + ) await endpoints.exchange_token_with_server( - request, server, "refresh_token", code=None, redirect_uri=None, client_id="scope-client", client_secret=None, - code_verifier=None, refresh_token="owner-refresh", scope=requested_scope, + request, + server, + "refresh_token", + code=None, + redirect_uri=None, + client_id="scope-client", + client_secret=None, + code_verifier=None, + refresh_token="owner-refresh", + scope=requested_scope, ) saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) assert saved is not None and saved.get("scopes") == expected assert upstream.call_count == 1 from urllib.parse import parse_qs + sent = parse_qs(upstream.calls.last.request.content.decode()) assert sent.get("scope") == ([requested_scope] if requested_scope else None) @@ -14202,21 +14692,27 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - server = _cimd_oauth_server().model_copy(update={ - "scopes": ["tools.read", "offline_access"], "per_server_oauth_discovery": per_server_relay, - }) + server = _cimd_oauth_server().model_copy( + update={ + "scopes": ["tools.read", "offline_access"], + "per_server_oauth_discovery": per_server_relay, + } + ) global_mcp_server_manager.registry[server.server_id] = server result = await endpoints._build_oauth_protected_resource_response(_cimd_request(), server.server_id, True) assert result["scopes_supported"] == ["tools.read"] @pytest.mark.asyncio -@pytest.mark.parametrize("granted, required, expected", [ - (["tools.read"], ("tools.write",), "absent"), - (["tools.read", "tools.write"], ("tools.write",), "present"), - (None, ("tools.write",), "absent"), - (["tools.read"], (), "present"), -]) +@pytest.mark.parametrize( + "granted, required, expected", + [ + (["tools.read"], ("tools.write",), "absent"), + (["tools.read", "tools.write"], ("tools.write",), "present"), + (None, ("tools.write",), "absent"), + (["tools.read"], (), "present"), + ], +) async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints @@ -14227,3 +14723,56 @@ async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, monkeypatch.setattr(db, "get_user_oauth_credential", read) assert await endpoints._vendor_credential_state("owner", "server", required) == expected read.assert_awaited_once_with(prisma, "owner", "server") + + +@pytest.mark.parametrize( + ("header", "expected"), + [ + ("alpha", ("litellm:mcp_server:alpha",)), + ("alpha, beta", ("litellm:mcp_server:alpha", "litellm:mcp_server:beta")), + ("alpha,nope", ()), + ("", ()), + (None, ()), + ], +) +def test_aggregate_protected_resource_advertises_header_server_scopes( + header: str | None, expected: tuple[str, ...] +) -> None: + """A client that re-authenticates from cached discovery never sees the 401 challenge, so the + aggregate PRM carries the x-mcp-servers selection as scopes_supported instead.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_protected_resource_response, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + servers: Final = { + name: MCPServer( + server_id=f"{name}-id", + name=name, + server_name=name, + url=f"http://{name}", + transport="http", + auth_type=MCPAuth.none, + ) + for name in ("alpha", "beta") + } + manager: Final = MagicMock() + manager.get_mcp_server_by_identifier.side_effect = lambda entry, client_ip=None: servers.get(entry) + headers: Final = [(b"host", b"localhost:4000")] + ( + [(b"x-mcp-servers", header.encode())] if header is not None else [] + ) + request: Final = Request( + { + "type": "http", + "method": "GET", + "path": "/.well-known/oauth-protected-resource/mcp", + "headers": headers, + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "root_path": "", + } + ) + with patch.object(mcp_server_manager, "global_mcp_server_manager", manager): + assert _build_aggregate_protected_resource_response(request)["scopes_supported"] == expected diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index c79b2efde5e..1417e77c8e2 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -11,6 +11,7 @@ from urllib.parse import parse_qs, urlparse import pytest from starlette.requests import Request +from starlette.responses import Response from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( @@ -27,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( MintedProxyCredential, SubjectIdentity, SubjectTokenRefusal, + _ConnectFlow, _GatewayAuthCode, _open_sealed, _seal, @@ -57,6 +59,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_refresh_token, mint_session_token, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" @@ -264,10 +267,17 @@ async def test_tampered_client_id_does_not_open(): def _authorize( - client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code" + client_id, + session_user_id, + redirect_uri=REDIRECT_URI, + challenge=CODE_CHALLENGE, + method="S256", + response_type="code", + scope: str | None = None, ): + query: Final = f"client_id={client_id}" if scope is None else f"client_id={client_id}&scope={scope}" return aggregate_authorize( - request=_request(query=f"client_id={client_id}"), + request=_request(query=query), client_id=client_id, redirect_uri=redirect_uri, state="client-state-123", @@ -275,6 +285,7 @@ def _authorize( code_challenge_method=method, response_type=response_type, session_user_id=session_user_id, + scope=scope, ) @@ -295,11 +306,12 @@ async def test_authorize_validation_failures_never_redirect_to_client(): @pytest.mark.asyncio async def test_authorize_without_session_redirects_to_login_with_return_to(): client_id = (await _register([REDIRECT_URI]))["client_id"] - response = _authorize(client_id, session_user_id=None) + response = _authorize(client_id, session_user_id=None, scope="litellm:mcp_server:missing") assert response.status_code == 303 location = response.headers["location"] assert location.startswith("https://llm.example.com/sso/key/generate?return_to=") assert "return_to=%2Fauthorize" in location + assert "scope=litellm:mcp_server:missing" in parse_qs(urlparse(location).query)["return_to"][0] @pytest.mark.asyncio @@ -991,7 +1003,12 @@ SCOPED_RESOURCE = "https://llm.example.com/mcp/github" _MANAGER_PATCH = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" -def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs): +def _scoped_authorize( + client_id: str, + resource: str | None, + session_user_id: str | None = "u1", + scope: str | None = None, +) -> Response: return aggregate_authorize( request=_request(query=f"client_id={client_id}"), client_id=client_id, @@ -1002,7 +1019,7 @@ def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs): response_type="code", session_user_id=session_user_id, resource=resource, - **kwargs, + scope=scope, ) @@ -1094,6 +1111,105 @@ def _sealed_wire_json(sealed, prefix, debug_key): return json.loads(raw) +def test_legacy_server_scope_payloads_open_as_single_server_ids(): + flow_payload: Final = { + "user_id": "u1", + "client_id": "client-1", + "redirect_uri": REDIRECT_URI, + "state": "state", + "code_challenge": CODE_CHALLENGE, + "jti": "flow-jti", + "exp": 2, + "resource_server_id": "server-1", + } + encrypted_flow_payload: Final = encrypt_value_helper(json.dumps(flow_payload)) + assert isinstance(encrypted_flow_payload, str) + opened_flow: Final = _open_sealed(encrypted_flow_payload, "", _ConnectFlow, "gateway_connect_flow") + assert opened_flow is not None + assert opened_flow.resource_server_ids == ("server-1",) + rejected_flow_payload: Final = encrypt_value_helper( + json.dumps(flow_payload | {"resource_server_ids": ["server-1"]}) + ) + assert isinstance(rejected_flow_payload, str) + assert _open_sealed(rejected_flow_payload, "", _ConnectFlow, "gateway_connect_flow") is None + + code_payload: Final = { + "user_id": "u1", + "client_id": "client-1", + "redirect_uri": REDIRECT_URI, + "code_challenge": CODE_CHALLENGE, + "jti": "code-jti", + "iat": 1, + "exp": 2, + "resource_server_id": "server-1", + } + encrypted_code_payload: Final = encrypt_value_helper(json.dumps(code_payload)) + assert isinstance(encrypted_code_payload, str) + opened_code: Final = _open_sealed( + f"{GATEWAY_AUTH_CODE_PREFIX}{encrypted_code_payload}", + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode, + _AUTH_CODE_DEBUG_KEY, + ) + assert opened_code is not None + assert opened_code.resource_server_ids == ("server-1",) + rejected_code_payload: Final = encrypt_value_helper( + json.dumps(code_payload | {"resource_server_ids": ["server-1"]}) + ) + assert isinstance(rejected_code_payload, str) + assert ( + _open_sealed( + f"{GATEWAY_AUTH_CODE_PREFIX}{rejected_code_payload}", + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode, + _AUTH_CODE_DEBUG_KEY, + ) + is None + ) + + +def test_single_server_flow_and_code_payloads_keep_the_legacy_wire_key(): + flow_payload: Final = _sealed_wire_json( + _seal( + "", + _ConnectFlow( + user_id="u1", + client_id="client-1", + redirect_uri=REDIRECT_URI, + state="state", + code_challenge=CODE_CHALLENGE, + jti="flow-jti", + exp=2, + resource_server_ids=("server-1",), + ), + ), + "", + "gateway_connect_flow", + ) + assert flow_payload["resource_server_id"] == "server-1" + assert "resource_server_ids" not in flow_payload + + code_payload: Final = _sealed_wire_json( + _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id="u1", + client_id="client-1", + redirect_uri=REDIRECT_URI, + code_challenge=CODE_CHALLENGE, + jti="code-jti", + iat=1, + exp=2, + resource_server_ids=("server-1",), + ), + ), + GATEWAY_AUTH_CODE_PREFIX, + _AUTH_CODE_DEBUG_KEY, + ) + assert code_payload["resource_server_id"] == "server-1" + assert "resource_server_ids" not in code_payload + + @pytest.mark.asyncio async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2 @@ -1106,15 +1222,15 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): client_id = (await _register([REDIRECT_URI]))["client_id"] github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = github + manager.get_mcp_server_by_identifier.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) assert response.status_code == 303 location = urlparse(response.headers["location"]) assert location.path == "/ui/connect" assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) - assert ( - _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id" + assert _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == ( + "github-id" ) described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent")) assert json.loads(described.body) == { @@ -1123,6 +1239,7 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): "server_id": "github-id", "server_name": "github", "connected": False, + "servers": None, } cache = DualCache() premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache) @@ -1133,17 +1250,233 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): assert completed.status_code == 303 assert present.calls == [("u1", "github-id")] code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] - assert ( - _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] - == "github-id" + assert _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == ( + "github-id" ) token_response = await _redeem(code, client_id) assert token_response.status_code == 200 principal = _opened_principal(json.loads(token_response.body)) - assert principal.resource_server_id == "github-id" + assert principal.resource_server_ids == ("github-id",) assert principal.user_id == "u1" +@pytest.mark.asyncio +@pytest.mark.parametrize("resource", ["https://llm.example.com/mcp", "https://llm.example.com"]) +async def test_aggregate_scope_seals_server_through_redemption_and_refresh(resource: str) -> None: + from unittest.mock import patch + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + github: Final = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.return_value = github + response: Final = _scoped_authorize( + client_id, + resource, + scope="litellm:mcp_server:github", + ) + assert response.status_code == 303 + _, cookies = _flow_cookie_from(response) + sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert sealed_flow["resource_server_id"] == "github-id" + + code: Final = await _finish_connect_page(response, scoped_server=github) + cache: Final = DualCache() + token_response: Final = await _redeem( + code, + client_id, + cache=cache, + resource=resource, + ) + assert token_response.status_code == 200 + payload: Final = json.loads(token_response.body) + assert _opened_principal(payload).resource_server_ids == ("github-id",) + + rotated: Final = await _redeem( + None, + client_id, + cache=cache, + grant_type="refresh_token", + refresh_token=payload["refresh_token"], + resource=resource, + ) + assert rotated.status_code == 200 + assert _opened_principal(json.loads(rotated.body)).resource_server_ids == ("github-id",) + + +@pytest.mark.asyncio +async def test_multi_server_scope_seals_set_through_consent_redemption_and_refresh() -> None: + from unittest.mock import patch + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + alpha: Final = _scoped_mcp_server(name="alpha") + beta: Final = _scoped_mcp_server(name="beta", oauth2_flow="client_credentials") + servers_by_name: Final = {"alpha": alpha, "beta": beta} + servers_by_id: Final = {alpha.server_id: alpha, beta.server_id: beta} + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.side_effect = servers_by_name.get + response: Final = _scoped_authorize( + client_id, + "https://llm.example.com/mcp", + scope="litellm:mcp_server:alpha litellm:mcp_server:beta", + ) + assert response.status_code == 303 + handle, cookies = _flow_cookie_from(response) + sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert sealed_flow["resource_server_ids"] == ["alpha-id", "beta-id"] + + vendor: Final = _VendorCredential("absent") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = servers_by_id.get + described: Final = await describe_connect_flow( + request=_request("/authorize/flow", cookies=cookies), + flow_handle=handle, + session_user_id="u1", + lookup_vendor_credential=vendor, + lookup_server_reachability=_ServerReachability(), + ) + assert json.loads(described.body) == { + "state": "multi", + "client_origin": "https://claude.ai", + "server_id": None, + "server_name": None, + "connected": None, + "servers": [ + {"server_id": "alpha-id", "server_name": "alpha", "connected": False}, + {"server_id": "beta-id", "server_name": "beta", "connected": True}, + ], + } + assert vendor.calls == [("u1", "alpha-id")] + + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = servers_by_id.get + blocked: Final = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=DualCache(), + lookup_vendor_credential=_VendorCredential("absent"), + lookup_server_reachability=_ServerReachability(), + ) + assert blocked.status_code == 400 + assert "authorize the requested MCP server" in json.loads(blocked.body)["error_description"] + + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = servers_by_id.get + completed: Final = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=DualCache(), + lookup_vendor_credential=_VendorCredential("present"), + lookup_server_reachability=_ServerReachability(), + ) + assert completed.status_code == 303 + code: Final = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + sealed_code: Final = _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") + assert sealed_code["resource_server_ids"] == ["alpha-id", "beta-id"] + + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.return_value = alpha + conflict: Final = await _redeem(code, client_id, resource="https://llm.example.com/mcp/alpha") + assert json.loads(conflict.body)["error"] == "invalid_target" + + cache: Final = DualCache() + token_response: Final = await _redeem( + code, + client_id, + cache=cache, + resource="https://llm.example.com/mcp", + ) + assert token_response.status_code == 200 + payload: Final = json.loads(token_response.body) + assert _opened_principal(payload).resource_server_ids == ("alpha-id", "beta-id") + + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.return_value = alpha + refresh_conflict: Final = await _redeem( + None, + client_id, + cache=cache, + grant_type="refresh_token", + refresh_token=payload["refresh_token"], + resource="https://llm.example.com/mcp/alpha", + ) + assert json.loads(refresh_conflict.body)["error"] == "invalid_target" + + refreshed: Final = await _redeem( + None, + client_id, + cache=cache, + grant_type="refresh_token", + refresh_token=payload["refresh_token"], + resource="https://llm.example.com/mcp", + ) + assert refreshed.status_code == 200 + assert _opened_principal(json.loads(refreshed.body)).resource_server_ids == ("alpha-id", "beta-id") + + +@pytest.mark.asyncio +async def test_per_server_resource_takes_precedence_over_conflicting_scope() -> None: + from unittest.mock import patch + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + github: Final = _scoped_mcp_server() + linear: Final = _scoped_mcp_server(name="linear") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.side_effect = (github, linear) + response: Final = _scoped_authorize( + client_id, + SCOPED_RESOURCE, + scope="litellm:mcp_server:linear", + ) + _, cookies = _flow_cookie_from(response) + sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert sealed_flow["resource_server_id"] == "github-id" + + +@pytest.mark.asyncio +async def test_unknown_server_scope_redirects_with_invalid_scope_and_state() -> None: + from unittest.mock import patch + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.return_value = None + response: Final = _scoped_authorize( + client_id, + "https://llm.example.com/mcp", + session_user_id="u1", + scope="openid litellm:mcp_server:missing", + ) + + assert response.status_code == 303 + location: Final = urlparse(response.headers["location"]) + assert location.scheme == "https" + assert location.netloc == "claude.ai" + assert location.path == "/api/mcp/auth_callback" + assert parse_qs(location.query) == { + "error": ["invalid_scope"], + "error_description": ["unknown or unsupported MCP server scope"], + "state": ["client-state-123"], + } + + +@pytest.mark.parametrize( + ("server_name", "expected"), + [ + ("alpha", "litellm:mcp_server:alpha"), + ("server~name", "litellm:mcp_server:server~name"), + ("", None), + ("bad name", None), + ('bad"name', None), + ("naïve", None), + ], +) +def test_gateway_server_scope_accepts_only_rfc6749_scope_tokens(server_name: str, expected: str | None) -> None: + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import gateway_server_scope + + assert gateway_server_scope(server_name) == expected + + @pytest.mark.asyncio @pytest.mark.parametrize( "resource, resolves", @@ -1168,14 +1501,14 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server() + manager.get_mcp_server_by_identifier.return_value = None if resolves is None else _scoped_mcp_server() response = _scoped_authorize(client_id, resource) assert response.status_code == 303 location = urlparse(response.headers["location"]) assert location.path == "/ui/connect" assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) - assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert "resource_server_ids" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") vendor = _VendorCredential("absent") described = await _describe_page(response, vendor=vendor) assert json.loads(described.body)["state"] == "unscoped" @@ -1183,13 +1516,14 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, completed = await _complete_page(response, vendor=vendor) assert vendor.calls == [] code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] - assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") + assert "resource_server_ids" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") token_response = await _redeem(code, client_id) payload = json.loads(token_response.body) - assert _opened_principal(payload).resource_server_id is None + assert _opened_principal(payload).resource_server_ids is None jwt_payload_segment = payload["access_token"].removeprefix("llm_session_").split(".")[1] claims = json.loads(base64.urlsafe_b64decode(jwt_payload_segment + "=" * (-len(jwt_payload_segment) % 4))) assert "resource_server_id" not in claims + assert "resource_server_ids" not in claims @pytest.mark.asyncio @@ -1201,7 +1535,7 @@ async def test_scoped_authorize_delegate_server_stays_unscoped(): client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) + manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) response = _scoped_authorize(client_id, SCOPED_RESOURCE) location = urlparse(response.headers["location"]) assert location.path == "/ui/connect" @@ -1209,7 +1543,7 @@ async def test_scoped_authorize_delegate_server_stays_unscoped(): assert json.loads((await _describe_page(response)).body)["server_id"] is None code = await _finish_connect_page(response) token_response = await _redeem(code, client_id) - assert _opened_principal(json.loads(token_response.body)).resource_server_id is None + assert _opened_principal(json.loads(token_response.body)).resource_server_ids is None @pytest.mark.asyncio @@ -1221,7 +1555,7 @@ async def test_m2m_scoped_flow_mints_without_a_user_credential(): client_id = (await _register([REDIRECT_URI]))["client_id"] m2m = _scoped_mcp_server(oauth2_flow="client_credentials") with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = m2m + manager.get_mcp_server_by_identifier.return_value = m2m response = _scoped_authorize(client_id, SCOPED_RESOURCE) vendor = _VendorCredential("unavailable") described = await _describe_page(response, scoped_server=m2m, vendor=vendor) @@ -1241,7 +1575,7 @@ async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): client_id = (await _register([REDIRECT_URI]))["client_id"] server = _scoped_mcp_server(oauth2_flow=oauth2_flow) with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = server + manager.get_mcp_server_by_identifier.return_value = server response = _scoped_authorize(client_id, SCOPED_RESOURCE) reachable = _ServerReachability(False) vendor = _VendorCredential("present") @@ -1252,6 +1586,7 @@ async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): "server_id": None, "server_name": None, "connected": None, + "servers": None, } cache = DualCache() refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache) @@ -1270,7 +1605,7 @@ async def test_stale_scoped_flow_remains_distinct_from_unscoped(): client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server() response = _scoped_authorize(client_id, SCOPED_RESOURCE) described = await _describe_page(response, scoped_server=None) assert json.loads(described.body)["state"] == "stale" @@ -1290,7 +1625,7 @@ async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credentia client_id = (await _register([REDIRECT_URI]))["client_id"] github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = github + manager.get_mcp_server_by_identifier.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) cache = DualCache() skipped_reachability = _ServerReachability(False) @@ -1350,22 +1685,22 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): github = _scoped_mcp_server() linear = _scoped_mcp_server(name="linear") with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = github + manager.get_mcp_server_by_identifier.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) code = await _finish_connect_page(response, scoped_server=github) with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = linear + manager.get_mcp_server_by_identifier.return_value = linear mismatched = await _redeem(code, client_id, resource="https://llm.example.com/mcp/linear") assert json.loads(mismatched.body)["error"] == "invalid_target" cache = DualCache() token_response = await _redeem(code, client_id, cache=cache) payload = json.loads(token_response.body) - assert _opened_principal(payload).resource_server_id == "github-id" + assert _opened_principal(payload).resource_server_ids == ("github-id",) with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = linear + manager.get_mcp_server_by_identifier.return_value = linear refresh_mismatch = await _redeem( None, client_id, @@ -1380,7 +1715,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"] ) assert rotated.status_code == 200 - assert _opened_principal(json.loads(rotated.body)).resource_server_id == "github-id" + assert _opened_principal(json.loads(rotated.body)).resource_server_ids == ("github-id",) @pytest.mark.asyncio @@ -1406,7 +1741,7 @@ async def test_resolve_scoped_resource_server_matrix(): (None, github, None), ]: with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = resolved_server + manager.get_mcp_server_by_identifier.return_value = resolved_server result = resolve_scoped_resource_server(request, resource) assert (result.server_id if result is not None else None) == expected, resource @@ -1422,10 +1757,10 @@ async def test_resource_resolution_is_identity_not_ip_filtered_access(): from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import resolve_scoped_resource_server with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server() result = resolve_scoped_resource_server(_request(), SCOPED_RESOURCE) assert result is not None - manager.get_mcp_server_by_name.assert_called_once_with("github") + manager.get_mcp_server_by_identifier.assert_called_once_with("github") LOOPBACK_REDIRECT_URI = "http://127.0.0.1:51234/callback" @@ -1560,7 +1895,7 @@ async def test_native_authorize_renders_consent_page_and_sets_flow_cookie(): assert flow["client_id"] == client_id assert flow["redirect_uri"] == LOOPBACK_REDIRECT_URI assert flow["user_id"] == "u1" - assert "resource_server_id" not in flow + assert "resource_server_ids" not in flow assert handle not in body.replace(f'value="{handle}"', "") @@ -2214,6 +2549,34 @@ async def test_introspect_active_access_token_reports_rfc7662_claims(): assert body["jti"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("resource_server_ids", "claim", "expected", "other_claim"), + [ + (("alpha-id",), "resource_server_id", "alpha-id", "resource_server_ids"), + (("alpha-id", "beta-id"), "resource_server_ids", ["alpha-id", "beta-id"], "resource_server_id"), + ], +) +async def test_introspect_reports_server_scope_as_scalar_or_list( + resource_server_ids: tuple[str, ...], + claim: str, + expected: str | list[str], + other_claim: str, +): + keys, now, _ = _introspection_fixtures() + principal = SessionPrincipal( + user_id="u1", + client_id="llm_dcrc_client", + resource_server_ids=resource_server_ids, + ) + minted = mint_session_token(principal, keys, now) + status, body = await _introspect(minted.token.get_secret_value()) + + assert status == 200 + assert body[claim] == expected + assert other_claim not in body + + @pytest.mark.asyncio async def test_introspect_invalid_tokens_answer_active_false(): keys, now, principal = _introspection_fixtures() @@ -2380,7 +2743,7 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a server = _scoped_mcp_server(auth_type=auth_type) vendor = _VendorCredential("absent") with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = server + manager.get_mcp_server_by_identifier.return_value = server response = _scoped_authorize(client_id, resource) described = await _describe_page(response, scoped_server=server, vendor=vendor) assert json.loads(described.body) == { @@ -2389,6 +2752,7 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a "server_id": "github-id", "server_name": "github", "connected": True, + "servers": None, } unreachable = await _complete_page(response, scoped_server=server, reachable=_ServerReachability(False)) assert unreachable.status_code == 400 @@ -2398,16 +2762,16 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a assert vendor.calls == [] code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = server + manager.get_mcp_server_by_identifier.return_value = server redeemed = await _redeem(code, client_id, cache=cache, resource=resource) assert redeemed.status_code == 200 payload = json.loads(redeemed.body) - assert _opened_principal(payload).resource_server_id == "github-id" + assert _opened_principal(payload).resource_server_ids == ("github-id",) renewed = await _redeem( None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"] ) assert renewed.status_code == 200 - assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" + assert _opened_principal(json.loads(renewed.body)).resource_server_ids == ("github-id",) JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" @@ -2594,7 +2958,7 @@ async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it(): client_id = (await _register([REDIRECT_URI]))["client_id"] server = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = server + manager.get_mcp_server_by_identifier.return_value = server response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope="tools.write") vendor = AsyncMock(return_value="absent") described = await _describe_page(response, scoped_server=server, vendor=vendor) @@ -2631,7 +2995,7 @@ async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_g code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] tokens = await _redeem(code, client_id) assert tokens.status_code == 200 - assert _opened_principal(json.loads(tokens.body)).resource_server_id is None + assert _opened_principal(json.loads(tokens.body)).resource_server_ids is None @pytest.mark.asyncio @@ -2710,7 +3074,7 @@ async def test_named_resource_rejects_another_upstreams_permission_before_consen client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server() response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope=encode_upstream_scope("other-id", "write")) assert response.status_code == 400 assert json.loads(response.body)["error"] == "invalid_scope" @@ -2825,7 +3189,7 @@ async def test_sdk_completes_aggregate_step_up_discovery_consent_and_pkce(route) tokens = await _redeem(cache=cache, **params) assert tokens.status_code == 200 payload = json.loads(tokens.body) - assert _opened_principal(payload).resource_server_id is None + assert _opened_principal(payload).resource_server_ids is None retry = await flow.asend(httpx2.Response(200, json=payload, request=request)) assert retry.url == initial.url assert retry.headers["authorization"] == f"Bearer {payload['access_token']}" @@ -2854,7 +3218,7 @@ async def test_named_step_up_finish_describes_and_retains_its_single_server_ceil client_id = (await _register([REDIRECT_URI]))["client_id"] server = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = server + manager.get_mcp_server_by_identifier.return_value = server response = _scoped_authorize( client_id, SCOPED_RESOURCE, scope=encode_upstream_scope(server.server_id, "tools.write") ) @@ -2868,4 +3232,67 @@ async def test_named_step_up_finish_describes_and_retains_its_single_server_ceil code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] tokens = await _redeem(code, client_id) assert tokens.status_code == 200 - assert _opened_principal(json.loads(tokens.body)).resource_server_id == server.server_id + assert _opened_principal(json.loads(tokens.body)).resource_server_ids == (server.server_id,) + + +@pytest.mark.asyncio +async def test_upstream_step_up_inside_multi_server_scope_describes_every_sealed_server() -> None: + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + alpha: Final = _scoped_mcp_server(name="alpha") + beta: Final = _scoped_mcp_server(name="beta") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.side_effect = {"alpha": alpha, "beta": beta}.get + response: Final = _scoped_authorize( + client_id, + "https://llm.example.com/mcp", + scope=" ".join( + ["litellm:mcp_server:alpha", "litellm:mcp_server:beta", encode_upstream_scope(alpha.server_id, "w")] + ), + ) + assert response.status_code == 303 + handle, cookies = _flow_cookie_from(response) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = {alpha.server_id: alpha, beta.server_id: beta}.get + described: Final = await describe_connect_flow( + request=_request("/authorize/flow", cookies=cookies), + flow_handle=handle, + session_user_id="u1", + lookup_vendor_credential=AsyncMock(return_value="present"), + lookup_server_reachability=_ServerReachability(), + ) + body: Final = json.loads(described.body) + assert body["state"] == "multi" + assert [item["server_id"] for item in body["servers"]] == ["alpha-id", "beta-id"] + + +@pytest.mark.asyncio +async def test_unscoped_flow_cookie_keeps_the_legacy_wire_shape() -> None: + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + response: Final = _scoped_authorize(client_id, "https://llm.example.com/mcp") + assert response.status_code == 303 + _, cookies = _flow_cookie_from(response) + sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert "resource_server_ids" not in sealed_flow + assert "resource_server_id" not in sealed_flow + + +@pytest.mark.asyncio +async def test_gateway_scope_still_narrows_when_the_resource_does_not_resolve() -> None: + from unittest.mock import patch + + client_id: Final = (await _register([REDIRECT_URI]))["client_id"] + alpha: Final = _scoped_mcp_server(name="alpha") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_identifier.side_effect = {"alpha": alpha}.get + manager.get_mcp_server_by_name.return_value = None + response: Final = _scoped_authorize( + client_id, "https://llm.example.com/mcp/deleted-server", scope="litellm:mcp_server:alpha" + ) + assert response.status_code == 303 + _, cookies = _flow_cookie_from(response) + sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") + assert sealed_flow["resource_server_id"] == "alpha-id" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ef0c3de0d99..536c1eb0b44 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -379,6 +379,22 @@ class TestMCPServerManager: with patch.object(manager, "_get_general_settings", return_value={}): assert manager.get_mcp_server_by_id(server.server_id, client_ip="8.8.8.8") is None + def test_get_mcp_server_by_identifier_matches_alias_case_and_server_id(self): + manager = MCPServerManager() + server = MCPServer( + server_id="alpha-id", + name="alpha", + server_name="alpha", + alias="alpha-alias", + transport=MCPTransport.http, + ) + manager.registry[server.server_id] = server + + assert manager.get_mcp_server_by_identifier("ALPHA") is server + assert manager.get_mcp_server_by_identifier("alpha-alias") is server + assert manager.get_mcp_server_by_identifier("alpha-id") is server + assert manager.get_mcp_server_by_identifier("unknown") is None + async def test_create_mcp_client_stdio(self, monkeypatch): """Test creating MCP client for stdio transport""" monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") @@ -13765,7 +13781,7 @@ class TestSessionResourceScopeIntersect: auth = UserAPIKeyAuth(user_id="scoped-user") auth.mcp_admitted_user_subject = True - auth.mcp_session_resource_server_id = scope + auth.mcp_session_resource_server_ids = None if scope is None else (scope,) return auth def test_scope_reader_is_none_for_keys_and_unscoped_subjects(self): @@ -13779,7 +13795,14 @@ class TestSessionResourceScopeIntersect: def test_scope_reader_returns_sealed_scope_for_admitted_subjects(self): from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager - assert MCPServerManager._admitted_session_resource_scope(self._admitted_auth("b")) == "b" + assert MCPServerManager._admitted_session_resource_scope(self._admitted_auth("b")) == frozenset({"b"}) + + def test_scope_reader_returns_all_sealed_server_ids(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + auth = self._admitted_auth("first") + auth.mcp_session_resource_server_ids = ("first", "second") + assert MCPServerManager._admitted_session_resource_scope(auth) == frozenset({"first", "second"}) @pytest.mark.asyncio async def test_admin_registry_seed_still_bounded_by_session_resource_scope(self): @@ -13811,9 +13834,41 @@ class TestSessionResourceScopeIntersect: ), ): assert await manager.get_allowed_mcp_servers(auth) == ["granted-id"] - auth.mcp_session_resource_server_id = None + auth.mcp_session_resource_server_ids = None assert set(await manager.get_allowed_mcp_servers(auth)) == {"granted-id", "other-id"} + @pytest.mark.asyncio + async def test_admin_registry_seed_keeps_both_servers_in_a_multi_server_session(self): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy._types import LitellmUserRoles + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager = MCPServerManager() + for server_id in ("alpha-id", "beta-id", "gamma-id"): + manager.registry[server_id] = MCPServer( + server_id=server_id, + name=server_id, + server_name=server_id, + url="https://example.com/mcp", + transport=MCPTransport.http, + ) + auth = self._admitted_auth("alpha-id") + auth.mcp_session_resource_server_ids = ("alpha-id", "beta-id") + auth.user_role = LitellmUserRoles.PROXY_ADMIN + with ( + patch.object(MCPServerManager, "get_allow_all_keys_server_ids", return_value=[]), + patch.object( + MCPServerManager, + "_get_active_submitted_mcp_server_ids_for_user", + new_callable=AsyncMock, + return_value=[], + ), + ): + assert set(await manager.get_allowed_mcp_servers(auth)) == {"alpha-id", "beta-id"} + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_scopes_past_operator_open_union(self): """The intersect applies AFTER the operator-open (allow_all_keys) union, so a scoped diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index 097b0924b8d..b9eebca3baa 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -182,13 +182,12 @@ class TestApplyToolsetScope: admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None) admitted.mcp_admitted_user_subject = True - admitted.mcp_session_resource_server_id = "server-own" + admitted.mcp_session_resource_server_ids = ("server-own", "server-own-other") admitted.requires_fresh_policy = True granted = AsyncMock(return_value=frozenset({"toolset-123"})) resolve = AsyncMock(return_value={"server-team": ["tool1"]}) with patch( - "litellm.proxy._experimental.mcp_server.server." - "global_mcp_server_manager.resolve_toolset_tool_permissions", + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.resolve_toolset_tool_permissions", new=resolve, ): with pytest.raises(HTTPException) as exc_info: @@ -203,18 +202,17 @@ class TestApplyToolsetScope: admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None) admitted.mcp_admitted_user_subject = True - admitted.mcp_session_resource_server_id = "server-team" + admitted.mcp_session_resource_server_ids = ("server-team",) granted = AsyncMock(return_value=frozenset({"toolset-123"})) resolve = AsyncMock(return_value={"server-team": ["tool1"], "server-other": ["tool2"]}) with patch( - "litellm.proxy._experimental.mcp_server.server." - "global_mcp_server_manager.resolve_toolset_tool_permissions", + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.resolve_toolset_tool_permissions", new=resolve, ): result = await apply_toolset_scope(admitted, "toolset-123", granted=granted) assert result.mcp_toolset_id == "toolset-123" - assert result.mcp_session_resource_server_id == "server-team" + assert result.mcp_session_resource_server_ids == ("server-team",) resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=False) @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 5d8db1e6e2b..47578344f77 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1,6 +1,6 @@ import asyncio from collections.abc import Callable, Sequence -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, patch import pytest @@ -44,6 +44,50 @@ async def _served_catalog_tool() -> str: return "ok" +@pytest.mark.parametrize( + ("settings", "mcp_servers", "toolset_id", "auth_state", "raises"), + [ + ({"mcp_require_explicit_server_scope": True}, None, None, None, True), + ({"mcp_require_explicit_server_scope": True}, None, None, "unscoped", True), + ({"mcp_require_explicit_server_scope": True}, ["alpha"], None, None, False), + ({"mcp_require_explicit_server_scope": True}, [], None, None, True), + ({"mcp_require_explicit_server_scope": True}, None, "toolset-id", None, False), + ({"mcp_require_explicit_server_scope": True}, None, None, "sealed", False), + ({"mcp_require_explicit_server_scope": False}, None, None, None, False), + ({}, None, None, None, False), + ], +) +def test_explicit_server_scope_setting_only_rejects_unscoped_aggregate_requests( + settings: dict[str, bool], + mcp_servers: list[str] | None, + toolset_id: str | None, + auth_state: Literal["unscoped", "sealed"] | None, + raises: bool, +) -> None: + from fastapi import HTTPException + + auth: Final[UserAPIKeyAuth | None] = ( + None if auth_state is None else UserAPIKeyAuth(api_key=None, user_id="scope-test") + ) + if auth_state == "sealed" and auth is not None: + auth.mcp_session_resource_server_ids = ("alpha-id",) + + with patch("litellm.proxy.proxy_server.general_settings", settings): + if raises: + with pytest.raises(HTTPException) as exc_info: + operations.raise_if_unscoped_aggregate_request(mcp_servers, auth, toolset_id) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == { + "error": "mcp_server_scope_required", + "message": ( + "This gateway requires an explicit MCP server scope: connect to /mcp/ " + "or send the x-mcp-servers header." + ), + } + else: + operations.raise_if_unscoped_aggregate_request(mcp_servers, auth, toolset_id) + + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["direct", "virtual", "proxy"]) @pytest.mark.parametrize("verdict", ["allow", "block", "obo_error", "evaluate_error", "missing"]) diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index d9c38d60a29..68bfe4b34e5 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -1,4 +1,5 @@ import json +from typing import Final import pytest from pydantic import ValidationError @@ -32,32 +33,39 @@ SERVER_ONLY_MARKERS = ( "billing_agent_policy", "mcp_admitted_user_subject", "mcp_source_team_rpm_limits", - "mcp_session_resource_server_id", + "mcp_session_resource_server_ids", "via_virtual_key", ) +def _server_only_marker_value(marker: str, value: str) -> str | tuple[str, ...]: + return (value,) if marker == "mcp_session_resource_server_ids" else value + + @pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS) def test_a_caller_cannot_forge_a_server_only_marker_through_the_constructor(marker): - auth = UserAPIKeyAuth(**{marker: "forged-by-caller"}) + marker_value: Final = _server_only_marker_value(marker, "forged-by-caller") + auth = UserAPIKeyAuth(**{marker: marker_value}) - assert getattr(auth, marker) != "forged-by-caller" + assert getattr(auth, marker) != marker_value @pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS) def test_a_caller_cannot_forge_a_server_only_marker_through_model_validate(marker): - auth = UserAPIKeyAuth.model_validate({marker: "forged-by-caller"}) + marker_value: Final = _server_only_marker_value(marker, "forged-by-caller") + auth = UserAPIKeyAuth.model_validate({marker: marker_value}) - assert getattr(auth, marker) != "forged-by-caller" + assert getattr(auth, marker) != marker_value @pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS) def test_the_server_sets_a_marker_by_assignment_after_construction(marker): auth = UserAPIKeyAuth() + marker_value: Final = _server_only_marker_value(marker, "set-by-the-server") - setattr(auth, marker, "set-by-the-server") + setattr(auth, marker, marker_value) - assert getattr(auth, marker) == "set-by-the-server" + assert getattr(auth, marker) == marker_value def test_a_virtual_key_is_hashed_out_of_the_auth_object(): diff --git a/tests/unit/proxy/test_proxy_types.py b/tests/unit/proxy/test_proxy_types.py index 4d0689f4260..574435c0cb2 100644 --- a/tests/unit/proxy/test_proxy_types.py +++ b/tests/unit/proxy/test_proxy_types.py @@ -266,7 +266,7 @@ def test_an_llm_backed_injection_check_needs_the_call_it_would_make(): [ ("mcp_admitted_user_subject", "someone-else", False), ("mcp_source_team_rpm_limits", {"team-1": 10_000}, None), - ("mcp_session_resource_server_id", "server-1", None), + ("mcp_session_resource_server_ids", ("server-1",), None), ("via_virtual_key", "sk-someone-elses-key", False), ], ) diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx index 5caf15d1fce..d3a24d7b9a5 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -3,16 +3,19 @@ import { render, screen } from "@testing-library/react"; import type { ConnectFlowStatus } from "@/components/networking"; import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner"; +const { startOAuthFlow } = vi.hoisted(() => ({ startOAuthFlow: vi.fn() })); + vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "https://gateway.example.com", })); vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ - useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }), + useUserMcpOAuthFlow: () => ({ startOAuthFlow, status: "idle" }), })); afterEach(() => { vi.restoreAllMocks(); + startOAuthFlow.mockClear(); }); const unscoped = (client_origin: string): ConnectFlowStatus => ({ @@ -21,6 +24,7 @@ const unscoped = (client_origin: string): ConnectFlowStatus => ({ server_id: null, server_name: null, connected: null, + servers: null, }); const renderBanner = (clientOrigin: string) => @@ -47,6 +51,53 @@ describe("ConnectFlowBanner", () => { expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); }); + it("lists multi-server consent and requires every server before finishing", () => { + const flow: ConnectFlowStatus = { + state: "multi", + client_origin: "https://claude.ai", + server_id: null, + server_name: null, + connected: null, + servers: [ + { server_id: "alpha-id", server_name: "alpha", connected: false }, + { server_id: "beta-id", server_name: "beta", connected: false }, + ], + }; + const { rerender } = render( + , + ); + + expect(screen.getByText("Allow https://claude.ai to use alpha, beta")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Connect alpha" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Connect beta" })).toBeInTheDocument(); + expect(startOAuthFlow).not.toHaveBeenCalled(); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); + + rerender( + ({ ...server, connected: true })) ?? null, + }} + accessToken="tok" + onConnected={vi.fn()} + failed={false} + />, + ); + + expect( + screen.getByText("Click Finish connecting to give https://claude.ai access to alpha, beta as you."), + ).toBeInTheDocument(); + }); + it("offers manual delivery only for a loopback client, posted only when checked", () => { const loopback = renderBanner("http://localhost:3118"); const checkbox = loopback.container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement; diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index e32ae028097..4ebfce698eb 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -27,6 +27,7 @@ export function isLoopbackOrigin(origin: string | null): boolean { const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly [string, string] => { const clientLabel = flow?.client_origin ?? "the application"; const serverLabel = flow?.server_name ?? "the requested MCP server"; + const serverNames = flow?.servers?.map((server) => server.server_name).join(", ") ?? ""; if (failed || flow === undefined || flow.state === "stale") { return [ "The connection cannot continue", @@ -39,6 +40,16 @@ const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly `Authorize the servers you want to use below, then click Finish connecting to return to ${clientLabel}.`, ]; } + if (flow.state === "multi") { + const allServersConnected = + flow.servers !== null && flow.servers.length > 0 && flow.servers.every((server) => server.connected); + return [ + `Allow ${clientLabel} to use ${serverNames}`, + allServersConnected + ? `Click Finish connecting to give ${clientLabel} access to ${serverNames} as you.` + : `Authorize each requested server below to continue, or cancel to send ${clientLabel} away.`, + ]; + } if (flow.state === "interactive" && !flow.connected) { return [ `Allow ${clientLabel} to use ${serverLabel}`, @@ -51,16 +62,33 @@ const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly ]; }; +const canFinishConnectFlow = (state: ConnectFlowStatus["state"], flow: ConnectFlowStatus | undefined): boolean => { + if (state === "unscoped") return true; + if (state === "stale" || flow === undefined) return false; + if (state !== "multi") return flow.connected === true; + return flow.servers !== null && flow.servers.length > 0 && flow.servers.every((server) => server.connected); +}; + +const vendorServersFor = ( + state: ConnectFlowStatus["state"], + flow: ConnectFlowStatus | undefined, +): { server_id: string; server_name: string | null }[] => { + if (state === "multi") { + return (flow?.servers ?? []) + .filter((server) => !server.connected) + .map(({ server_id, server_name }) => ({ server_id, server_name })); + } + if (state !== "interactive" || flow?.connected !== false || flow.server_id === null) return []; + return [{ server_id: flow.server_id, server_name: flow.server_name }]; +}; + const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onConnected, failed }) => { const action = `${getProxyBaseUrl()}/authorize/complete`; const state = failed || flow === undefined ? "stale" : flow.state; - const canFinish = state === "unscoped" || (state !== "stale" && flow?.connected === true); + const canFinish = canFinishConnectFlow(state, flow); const canCancel = state !== "unscoped"; const loopbackClient = isLoopbackOrigin(flow?.client_origin ?? null); - const vendorServer = - state === "interactive" && flow?.connected === false && flow.server_id !== null - ? { server_id: flow.server_id, server_name: flow.server_name } - : null; + const vendorServers = vendorServersFor(state, flow); const copy = copyFor(flow, failed); return ( @@ -74,17 +102,20 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onC
- {vendorServer !== null && ( + {vendorServers.map((server) => ( - )} + ))}
{canFinish && ( diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx index cf59dcfc368..a7248a3a3cb 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx @@ -32,12 +32,19 @@ vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ }, })); -const flow = (state: "unscoped" | "interactive" | "m2m" | "stale", connected: boolean | null = null) => ({ +const flow = (state: "unscoped" | "interactive" | "m2m" | "multi" | "stale", connected: boolean | null = null) => ({ state, client_origin: "https://claude.ai", server_id: state === "interactive" || state === "m2m" ? "s-design" : null, server_name: state === "interactive" || state === "m2m" ? "design_tool" : null, connected, + servers: + state === "multi" + ? [ + { server_id: "s-alpha", server_name: "alpha", connected: false }, + { server_id: "s-beta", server_name: "beta", connected: false }, + ] + : null, }); const renderSurface = () => diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index d8386afeb93..5cc590ffa79 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -32,6 +32,7 @@ interface OAuth2ConnectButtonProps { variant?: "badge" | "button"; autoStartKey?: string | null; scopes?: string[]; + buttonLabel?: string; } export const OAuth2ConnectButton: React.FC = ({ @@ -41,6 +42,7 @@ export const OAuth2ConnectButton: React.FC = ({ variant = "badge", autoStartKey = null, scopes, + buttonLabel = "Connect", }) => { const name = server.server_name ?? server.alias ?? server.server_id; const flowOptions = { @@ -64,7 +66,7 @@ export const OAuth2ConnectButton: React.FC = ({ return ( ); } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index ac1f7d84bcb..69623d05232 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4840,12 +4840,13 @@ export const fetchDiscoverableMCPServers = async (accessToken: string) => { }; export interface ConnectFlowStatus { - state: "unscoped" | "interactive" | "m2m" | "stale"; + state: "unscoped" | "interactive" | "m2m" | "multi" | "stale"; client_origin: string; server_id: string | null; server_name: string | null; connected: boolean | null; requested_scopes?: string[]; + servers: { server_id: string; server_name: string; connected: boolean }[] | null; } /** diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d94fc976fe..7abd84c4ad7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30051,6 +30051,12 @@ export interface components { * @description When true, a gateway-managed OAuth2 MCP server whose authorization server advertises Client ID Metadata Document support identifies itself with the gateway's public metadata document URL even when that authorization server also offers dynamic client registration. Requires a public HTTPS PROXY_BASE_URL the authorization server can fetch. Default false: dynamic client registration is used whenever the authorization server offers it, and the metadata document only when it does not. */ mcp_prefer_client_id_metadata_document?: boolean | null; + /** + * Mcp Require Explicit Server Scope + * @description Require aggregate MCP requests to include x-mcp-servers unless a toolset or gateway OAuth session already scopes the server. Off by default. + * @default false + */ + mcp_require_explicit_server_scope: boolean; /** * Mcp Required Fields * @description List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).