From 64868f2fc8ad2f32cad550cf1735442b3c660ba4 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 23:59:18 -0700 Subject: [PATCH 01/16] feat(mcp): aggregate gateway DCR discovery front door behind mcp_gateway_dcr --- .../mcp_server/auth/user_api_key_auth_mcp.py | 156 ++++++++++++----- .../mcp_server/discoverable_endpoints.py | 98 +++++++++++ .../_experimental/mcp_server/oauth_utils.py | 23 +++ .../auth/test_user_api_key_auth_mcp.py | 133 ++++++++++++++ .../mcp_server/test_discoverable_endpoints.py | 165 ++++++++++++++++++ 5 files changed, 532 insertions(+), 43 deletions(-) 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 421f1dcfbeab..76589ff07349 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 @@ -10,6 +10,10 @@ import litellm from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + get_request_base_url, + is_mcp_gateway_dcr_enabled, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, @@ -120,6 +124,96 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _is_aggregate_gateway_dcr_challenge_scope( + route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, +) -> bool: + """True when an unauthenticated request to the aggregate ``/mcp`` endpoint + should receive the RFC 9728 401 challenge that advertises the gateway as + the authorization server (``mcp_gateway_dcr`` front door). + + Fires only for a genuine 401 on the aggregate scope: any named target + (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and + client-supplied MCP auth headers mean the caller is not a cold-start DCR + client. Fails closed to the original admission error otherwise.""" + if not is_mcp_gateway_dcr_enabled(): + return False + if not _is_litellm_auth_admission_error(exc): + return False + if mcp_servers: + return False + if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): + return False + return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + + +def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException: + """The RFC 9728 challenge for the aggregate endpoint: points the client at + the gateway's own protected-resource metadata so a DCR client discovers + the gateway as its authorization server and starts the sign-in flow. + + ``invalid_token`` adds the RFC 6750 error code for a request that DID + 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.""" + error_attr = 'error="invalid_token", ' if invalid_token else "" + resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + 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}"'}, + ) + + +def _admission_failure_fallback( + request: Request, + request_route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, + bearer_presented: bool, +) -> UserAPIKeyAuth: + """Map a failed LiteLLM admission to its anonymous fallback or challenge. + + Two fallbacks exist, both gated on a genuine 401 with no client-supplied + MCP auth headers. The pass-through cold start (RFC 9728 / MCP + Authorization spec discovery return) admits anonymously so the route's + 401 emitter can produce the per-server challenge. The aggregate + gateway-DCR scope converts the failure into the gateway's own + resource_metadata challenge, with the RFC 6750 ``invalid_token`` error + code when the caller DID present a bearer (an expired gateway session + must re-authorize, not retry a dead token). Anything else re-raises the + original admission error unchanged.""" + mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) + if ( + mcp_servers_from_path is not None + and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers) + and _is_litellm_auth_admission_error(exc) + and _is_mcp_passthrough_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + ): + verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") + return UserAPIKeyAuth() + if _is_aggregate_gateway_dcr_challenge_scope( + route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + ): + raise _aggregate_gateway_dcr_challenge(request, invalid_token=bearer_presented) from exc + raise exc + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -271,56 +365,32 @@ async def mock_body(): elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and - # propagates. The sole anonymous fallback is the auth_type=none - # pass-through cold-start (RFC 9728 discovery return), gated on a 401 - # so a recognized-but-forbidden key still fails closed. - client_ip = IPAddressUtils.get_mcp_client_ip(request) + # propagates unless a fallback in _admission_failure_fallback applies. try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as e: - # ProxyException.code is normalized to str (possibly "None"), so - # compare both int and str forms rather than coercing. - status = e.status_code if isinstance(e, HTTPException) else e.code - is_unauthenticated = status in (401, "401") - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - if ( - is_unauthenticated - and mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug( - "MCP pass-through return: forwarding Authorization as upstream OAuth token for delegated auth" - ) - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=e, + bearer_presented=True, + ) else: try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: - # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec - # require unauthenticated requests to protected resources to receive - # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers - # for pass-through servers instead of surfacing a generic admission error. - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - client_ip = IPAddressUtils.get_mcp_client_ip(request) - if ( - mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_litellm_auth_admission_error(exc) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + bearer_presented=False, + ) return ( validated_user_api_key_auth, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 54aff86aab2a..0ad9a5e3ea23 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, + is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1640,6 +1641,12 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery + # describes the gateway itself as the authorization server for the + # aggregate /mcp resource instead of narrowing to one server. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_protected_resource_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1770,6 +1777,92 @@ def _jwt_auth_issuers() -> list: return issuers +def _build_aggregate_protected_resource_response(request: Request) -> dict: + """RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is + the authorization server. No per-server names or scopes leak here; access + is resolved after sign-in from the authenticated user's grants. + + The advertised authorization server is ``{base}/mcp`` (not the bare + origin) so RFC 8414 path-insertion resolves its metadata at + ``/.well-known/oauth-authorization-server/mcp``, a route this module + owns. The bare-origin well-known is registered first by the BYOK OAuth + feature and describes the BYOK flow, so it must not be the aggregate + discovery entry point (same pattern as the per-server documents, which + advertise ``{base}/{server_name}``).""" + request_base_url = get_request_base_url(request) + return { + "authorization_servers": [f"{request_base_url}/mcp"], + "resource": f"{request_base_url}/mcp", + "scopes_supported": [], + } + + +def _build_aggregate_authorization_server_response(request: Request) -> dict: + """RFC 8414 metadata for the gateway as the aggregate authorization server. + + The issuer is ``{base}/mcp`` and must stay equal to the value the + aggregate protected-resource document advertises: spec clients verify the + issuer in the metadata matches the one that derived the well-known URL. + Advertises the root /authorize, /token, and /register endpoints and + ``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR + clients (Claude Desktop, MCP Inspector) register as public clients; PKCE + S256 is mandatory in the gateway's authorize flow.""" + request_base_url = get_request_base_url(request) + return { + "issuer": f"{request_base_url}/mcp", + "authorization_endpoint": f"{request_base_url}/authorize", + "token_endpoint": f"{request_base_url}/token", + "registration_endpoint": f"{request_base_url}/register", + "response_types_supported": ["code"], + "scopes_supported": [], + "grant_types_supported": ["authorization_code", "refresh_token"], + "code_challenge_methods_supported": ["S256"], + "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], + } + + +def _raise_404_unless_gateway_dcr_enabled() -> None: + """The aggregate well-known routes exist only under the gateway-level DCR + front door; flag-off they 404 exactly like the previously-absent routes so + discovery behavior is byte-identical for existing deployments.""" + if is_mcp_gateway_dcr_enabled(): + return + raise HTTPException(status_code=404, detail="Not Found") + + +# RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client +# pointed at {base}/mcp inserts the well-known segment before the resource +# path, so this exact route must exist for aggregate discovery to work at all. +# Declared before the parameterized well-known routes below: Starlette matches +# in registration order, and /.well-known/oauth-authorization-server/{name} +# would otherwise capture the "/mcp" suffix as a server name. +@router.get( + f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_protected_resource_aggregate(request: Request): + """ + OAuth protected resource discovery for the aggregate /mcp endpoint + (gateway-level DCR front door; 404 when the flag is off). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_protected_resource_response(request) + + +@router.get( + f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_authorization_server_aggregate(request: Request): + """ + OAuth authorization server discovery for the aggregate /mcp endpoint, the + RFC 8414 path-inserted form for a client that treats {base}/mcp as its + authorization base URL (gateway-level DCR front door; 404 when the flag + is off, indistinguishable from an unknown server name on the + parameterized route below). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_authorization_server_response(request) + + # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) @router.get( @@ -1829,6 +1922,11 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery keeps + # advertising the gateway's own /authorize, /token, and /register. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_authorization_server_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858e..74b56cf424bf 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,6 +70,29 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" +MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" + + +def is_mcp_gateway_dcr_enabled() -> bool: + """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into + the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root + OAuth discovery advertises the gateway itself as the authorization server + (instead of resolving the single configured oauth2 server), and the + anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` + challenge so DCR clients (Claude Desktop, MCP Inspector) can start the + sign-in flow. Off by default; flag-off behavior is unchanged.""" + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load + + if not isinstance(general_settings, dict): + return False + raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() == "true" + return False + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 6f132aaae9cd..7625eec2f5ce 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5950,3 +5950,136 @@ async def test_admit_helper_raises_500_when_no_db_connection(self): route="/mcp/bridge_delegate_server", ) assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +class TestAggregateGatewayDcrChallenge: + """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must + carry the RFC 9728 resource_metadata challenge pointing at the gateway's + own protected-resource metadata, and must NOT fire for named-server + targets, explicit litellm keys, non-401 failures, or with the flag off.""" + + _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" + _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" + _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' + + def _scope(self, path="/mcp", extra_headers=()): + return { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), *extra_headers], + } + + def _auth_401(self): + async def _raise(api_key, request): + raise ProxyException( + message="Authentication Error: Invalid API key", + type="auth_error", + param="api_key", + code=401, + ) + + return _raise + + async def test_challenge_on_anonymous_aggregate_mcp(self): + """Anonymous request to the aggregate /mcp with the flag on: 401 plus + the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" + + async def test_challenge_invalid_token_on_failed_bearer(self): + """A bearer that fails LiteLLM admission at aggregate scope (an expired + gateway session, a revoked key) re-challenges with error=invalid_token + so a spec client re-authorizes instead of retrying the dead token.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"authorization", b"Bearer expired-session-token"),)) + ) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + + async def test_no_challenge_when_flag_off(self): + """Flag off: the original admission error propagates untouched, both + with and without a bearer.""" + for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=False), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) + assert str(exc_info.value.code) == "401" + + async def test_no_challenge_for_explicit_litellm_key(self): + """An explicit x-litellm-api-key declares a litellm-key client; a typo + there must surface the real auth error, never a DCR challenge that + would send SDKs into a sign-in flow.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-litellm-api-key", b"sk-typo"),)) + ) + + async def test_no_challenge_for_named_servers_header(self): + """x-mcp-servers names explicit targets; the per-server challenge paths + own those, so the aggregate challenge must not fire.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-servers", b"github"),)) + ) + + async def test_no_challenge_for_path_named_server(self): + """/mcp/{server} targets one server; the aggregate challenge must not + fire even when that server does not resolve.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) + + async def test_no_challenge_for_client_supplied_mcp_auth(self): + """Per-server x-mcp-{alias}-authorization headers mean the caller is + not a cold-start DCR client; keep the original error.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-github-authorization", b"Bearer upstream"),)) + ) + + async def test_no_challenge_for_non_401_failure(self): + """Only genuine 401s convert to a challenge; a 500 stays a 500.""" + + async def _raise_500(api_key, request): + raise ProxyException(message="boom", type="server_error", param=None, code=500) + + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert str(exc_info.value.code) == "500" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f5ac229d119f..30a814c8ea41 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7130,3 +7130,168 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert response.status_code == 502 body = json.loads(response.body) assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} + + +def _patch_gateway_dcr_flag(enabled: bool): + return patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", + return_value=enabled, + ) + + +@pytest.mark.asyncio +async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): + """Flag on: root discovery must keep describing the gateway as the + authorization server for the aggregate /mcp resource even when exactly one + OAuth2 server exists (flag off, resolution narrows to that server; that + behavior is pinned by test_discovery_root_includes_server_name_prefix).""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_authorization_server_response, + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with _patch_gateway_dcr_flag(True): + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name=None, + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, + mcp_server_name=None, + use_standard_pattern=True, + ) + + assert authorization_response["issuer"] == "https://llm.example.com/mcp" + assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" + assert authorization_response["token_endpoint"] == "https://llm.example.com/token" + assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" + assert "none" in authorization_response["token_endpoint_auth_methods_supported"] + assert authorization_response["code_challenge_methods_supported"] == ["S256"] + assert authorization_response["scopes_supported"] == [] + + assert resource_response["resource"] == "https://llm.example.com/mcp" + assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert resource_response["scopes_supported"] == [] + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_gateway_dcr_named_discovery_unaffected_by_flag(): + """Flag on must not change named-server discovery: a named oauth2 server + still resolves to its own per-server document.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_authorization_server_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with _patch_gateway_dcr_flag(True): + response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name="test_oauth", + ) + assert "/test_oauth/authorize" in response["authorization_endpoint"] + assert response["scopes_supported"] == ["read", "write"] + finally: + global_mcp_server_manager.registry.clear() + + +def test_aggregate_wellknown_routes_404_when_flag_off(): + """Flag off, the aggregate well-known routes answer 404 exactly like the + previously-absent routes: discovery behavior is byte-identical for + existing deployments.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with _patch_gateway_dcr_flag(False): + assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 + assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 + + +def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): + """Flag on, both path-appended aggregate routes serve the gateway + documents. Exercises real routing, so this also pins registration order: + /.well-known/oauth-authorization-server/{name} would otherwise capture + the /mcp suffix as a server name and 404.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with _patch_gateway_dcr_flag(True): + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + + +def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): + """The flag reader accepts YAML booleans and env-interpolated strings, and + fails closed on anything else.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + is_mcp_gateway_dcr_enabled, + ) + from litellm.proxy.proxy_server import general_settings + + for raw, expected in ( + (True, True), + (False, False), + ("true", True), + ("True", True), + ("false", False), + ("yes", False), + (1, False), + (None, False), + ): + with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): + assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" + + with patch.dict(general_settings, {}, clear=True): + assert is_mcp_gateway_dcr_enabled() is False From 943d7aa0a91019d4e6b05ef781e32fb55644382e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 17:27:22 -0700 Subject: [PATCH 02/16] refactor(mcp): make the aggregate DCR front door always-on, remove the mcp_gateway_dcr flag The flag guarded no breaking change: the aggregate discovery lives at new /mcp-suffixed routes, the challenge only fires at aggregate scope, and the authorize/token/register/admission arms self-gate on the llm_dcrc_/llm_session_ prefixes. Bare-origin and per-server discovery are left exactly as they were, and a server literally named mcp keeps its own discovery via disambiguation, so turning it on for everyone changes nothing about existing flows. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 55 +++-- .../_experimental/mcp_server/oauth_utils.py | 23 --- .../auth/test_user_api_key_auth_mcp.py | 24 +-- .../mcp_server/test_discoverable_endpoints.py | 190 ++++++------------ 5 files changed, 93 insertions(+), 204 deletions(-) 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 76589ff07349..bf88c0be5da4 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 @@ -12,7 +12,6 @@ from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, - is_mcp_gateway_dcr_enabled, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -133,14 +132,12 @@ def _is_aggregate_gateway_dcr_challenge_scope( ) -> bool: """True when an unauthenticated request to the aggregate ``/mcp`` endpoint should receive the RFC 9728 401 challenge that advertises the gateway as - the authorization server (``mcp_gateway_dcr`` front door). + the authorization server. Fires only for a genuine 401 on the aggregate scope: any named target (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and client-supplied MCP auth headers mean the caller is not a cold-start DCR client. Fails closed to the original admission error otherwise.""" - if not is_mcp_gateway_dcr_enabled(): - return False if not _is_litellm_auth_admission_error(exc): return False if mcp_servers: diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0ad9a5e3ea23..39b84bdbd4eb 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,7 +42,6 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, - is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1641,12 +1640,6 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery - # describes the gateway itself as the authorization server for the - # aggregate /mcp resource instead of narrowing to one server. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_protected_resource_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1821,13 +1814,20 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _raise_404_unless_gateway_dcr_enabled() -> None: - """The aggregate well-known routes exist only under the gateway-level DCR - front door; flag-off they 404 exactly like the previously-absent routes so - discovery behavior is byte-identical for existing deployments.""" - if is_mcp_gateway_dcr_enabled(): - return - raise HTTPException(status_code=404, detail="Not Found") +def _mcp_named_server_exists(request: Request) -> bool: + """True when a server literally named ``mcp`` is configured and visible to this caller. + + Its per-server authorization-server document is served at + ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the + aggregate path. When such a server exists the real server wins the route, so that + deployment keeps its per-server discovery regardless of whether the aggregate front door + is on.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load + global_mcp_server_manager, + ) + + client_ip = IPAddressUtils.get_mcp_client_ip(request) + return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client @@ -1841,10 +1841,12 @@ def _raise_404_unless_gateway_dcr_enabled() -> None: ) async def oauth_protected_resource_aggregate(request: Request): """ - OAuth protected resource discovery for the aggregate /mcp endpoint - (gateway-level DCR front door; 404 when the flag is off). + OAuth protected resource discovery for the aggregate /mcp endpoint. + + The single-segment ``/mcp`` path does not collide with any per-server PRM pattern + (those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously + describes the aggregate resource. """ - _raise_404_unless_gateway_dcr_enabled() return _build_aggregate_protected_resource_response(request) @@ -1853,13 +1855,15 @@ async def oauth_protected_resource_aggregate(request: Request): ) async def oauth_authorization_server_aggregate(request: Request): """ - OAuth authorization server discovery for the aggregate /mcp endpoint, the - RFC 8414 path-inserted form for a client that treats {base}/mcp as its - authorization base URL (gateway-level DCR front door; 404 when the flag - is off, indistinguishable from an unknown server name on the - parameterized route below). + OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 + path-inserted form for a client that treats {base}/mcp as its authorization base URL. + + This single-segment path collides with the parameterized ``/{mcp_server_name}`` route + below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; + only when no such server exists is the aggregate document served. """ - _raise_404_unless_gateway_dcr_enabled() + if _mcp_named_server_exists(request): + return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) @@ -1922,11 +1926,6 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery keeps - # advertising the gateway's own /authorize, /token, and /register. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_authorization_server_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 74b56cf424bf..6edb22dd858e 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,29 +70,6 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" -MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" - - -def is_mcp_gateway_dcr_enabled() -> bool: - """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into - the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root - OAuth discovery advertises the gateway itself as the authorization server - (instead of resolving the single configured oauth2 server), and the - anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` - challenge so DCR clients (Claude Desktop, MCP Inspector) can start the - sign-in flow. Off by default; flag-off behavior is unchanged.""" - from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load - - if not isinstance(general_settings, dict): - return False - raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) - if isinstance(raw, bool): - return raw - if isinstance(raw, str): - return raw.strip().lower() == "true" - return False - - def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 7625eec2f5ce..2a28ff612717 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5957,9 +5957,8 @@ class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must carry the RFC 9728 resource_metadata challenge pointing at the gateway's own protected-resource metadata, and must NOT fire for named-server - targets, explicit litellm keys, non-401 failures, or with the flag off.""" + targets, explicit litellm keys, or non-401 failures.""" - _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' @@ -5983,11 +5982,10 @@ async def _raise(api_key, request): return _raise async def test_challenge_on_anonymous_aggregate_mcp(self): - """Anonymous request to the aggregate /mcp with the flag on: 401 plus + """Anonymous request to the aggregate /mcp: 401 plus the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) @@ -6001,7 +5999,6 @@ async def test_challenge_invalid_token_on_failed_bearer(self): so a spec client re-authorizes instead of retrying the dead token.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request( @@ -6011,25 +6008,12 @@ async def test_challenge_invalid_token_on_failed_bearer(self): www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' - async def test_no_challenge_when_flag_off(self): - """Flag off: the original admission error propagates untouched, both - with and without a bearer.""" - for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): - with ( - patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=False), - ): - with pytest.raises(ProxyException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) - assert str(exc_info.value.code) == "401" - async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that would send SDKs into a sign-in flow.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6041,7 +6025,6 @@ async def test_no_challenge_for_named_servers_header(self): own those, so the aggregate challenge must not fire.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6053,7 +6036,6 @@ async def test_no_challenge_for_path_named_server(self): fire even when that server does not resolve.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) @@ -6063,7 +6045,6 @@ async def test_no_challenge_for_client_supplied_mcp_auth(self): not a cold-start DCR client; keep the original error.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6078,7 +6059,6 @@ async def _raise_500(api_key, request): with ( patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 30a814c8ea41..af4b2caca722 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7132,72 +7132,79 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} -def _patch_gateway_dcr_flag(enabled: bool): - return patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", - return_value=enabled, +def test_aggregate_wellknown_routes_serve_gateway_metadata(): + """Both path-appended aggregate routes serve the gateway documents. Exercises real + routing, so this also pins registration order: the parameterized + /.well-known/oauth-authorization-server/{name} route would otherwise capture the /mcp + suffix as a server name.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, ) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) -@pytest.mark.asyncio -async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): - """Flag on: root discovery must keep describing the gateway as the - authorization server for the aggregate /mcp resource even when exactly one - OAuth2 server exists (flag off, resolution narrows to that server; that - behavior is pinned by test_discovery_root_includes_server_name_prefix).""" - from fastapi import Request + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_authorization_server_response, - _build_oauth_protected_resource_response, - ) + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + assert "none" in asm.json()["token_endpoint_auth_methods_supported"] + + +def test_as_aggregate_route_prefers_a_real_server_named_mcp(): + """A server literally named ``mcp`` wins the single-segment + /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized + /{server_name} route) and keeps its per-server discovery; the aggregate document is + served only when no such server exists.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) global_mcp_server_manager.registry.clear() - oauth2_server = _create_oauth2_server() - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://llm.example.com/" - mock_request.headers = {} + server_named_mcp = _create_oauth2_server(server_id="mcp_srv", name="mcp", server_name="mcp", alias="mcp") + global_mcp_server_manager.registry[server_named_mcp.server_id] = server_named_mcp + app = FastAPI() + app.include_router(router) + client = TestClient(app) try: - with _patch_gateway_dcr_flag(True): - authorization_response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name=None, - ) - resource_response = await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name=None, - use_standard_pattern=True, - ) - - assert authorization_response["issuer"] == "https://llm.example.com/mcp" - assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" - assert authorization_response["token_endpoint"] == "https://llm.example.com/token" - assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" - assert "none" in authorization_response["token_endpoint_auth_methods_supported"] - assert authorization_response["code_challenge_methods_supported"] == ["S256"] - assert authorization_response["scopes_supported"] == [] - - assert resource_response["resource"] == "https://llm.example.com/mcp" - assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] - assert resource_response["scopes_supported"] == [] + asm = client.get("/.well-known/oauth-authorization-server/mcp") + assert asm.status_code == 200 + # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), + # not the aggregate one (whose issuer would be {base}/mcp) + assert asm.json()["issuer"] == "http://testserver" + assert "/mcp/authorize" in asm.json()["authorization_endpoint"] finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio -async def test_gateway_dcr_named_discovery_unaffected_by_flag(): - """Flag on must not change named-server discovery: a named oauth2 server - still resolves to its own per-server document.""" +async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): + """The always-on aggregate front door must not change bare-origin discovery: with one + oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server, + protected-resource} still resolves THAT server, so an existing single-server deployment's + discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes.""" from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_authorization_server_response, + _build_oauth_protected_resource_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -7212,86 +7219,15 @@ async def test_gateway_dcr_named_discovery_unaffected_by_flag(): mock_request.headers = {} try: - with _patch_gateway_dcr_flag(True): - response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name="test_oauth", - ) - assert "/test_oauth/authorize" in response["authorization_endpoint"] - assert response["scopes_supported"] == ["read", "write"] + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, mcp_server_name=None + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name=None, use_standard_pattern=True + ) + # per-server, not aggregate: the single server's name is in the endpoints + assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] + assert authorization_response["issuer"] == "https://llm.example.com" + assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] finally: global_mcp_server_manager.registry.clear() - - -def test_aggregate_wellknown_routes_404_when_flag_off(): - """Flag off, the aggregate well-known routes answer 404 exactly like the - previously-absent routes: discovery behavior is byte-identical for - existing deployments.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(False): - assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 - assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 - - -def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): - """Flag on, both path-appended aggregate routes serve the gateway - documents. Exercises real routing, so this also pins registration order: - /.well-known/oauth-authorization-server/{name} would otherwise capture - the /mcp suffix as a server name and 404.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - global_mcp_server_manager.registry.clear() - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(True): - prm = client.get("/.well-known/oauth-protected-resource/mcp") - asm = client.get("/.well-known/oauth-authorization-server/mcp") - - assert prm.status_code == 200 - assert prm.json()["resource"] == "http://testserver/mcp" - assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] - - assert asm.status_code == 200 - assert asm.json()["issuer"] == "http://testserver/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" - - -def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): - """The flag reader accepts YAML booleans and env-interpolated strings, and - fails closed on anything else.""" - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - is_mcp_gateway_dcr_enabled, - ) - from litellm.proxy.proxy_server import general_settings - - for raw, expected in ( - (True, True), - (False, False), - ("true", True), - ("True", True), - ("false", False), - ("yes", False), - (1, False), - (None, False), - ): - with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): - assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" - - with patch.dict(general_settings, {}, clear=True): - assert is_mcp_gateway_dcr_enabled() is False From e7e5265f658b345085bfed9a8186fdae5200aa75 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 22:37:35 -0700 Subject: [PATCH 03/16] fix(mcp): reserve mcp for the aggregate AS and root-path the discovery challenges Two RFC 9728 / 8414 discovery fixes on the aggregate front door, both raised by Bugbot on this PR The aggregate authorization-server document at /.well-known/oauth-authorization-server/mcp used to defer to a per-server row literally named "mcp", serving issuer {base} while the aggregate protected-resource document advertises {base}/mcp as its authorization server. A spec client following that chain fails the RFC 8414 issuer check and cannot sign in. The single segment /mcp is now reserved for the aggregate so the issuer stays {base}/mcp and matches the protected-resource document; a server named "mcp" keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp The 401 challenges built the resource_metadata URL as {base}/.well-known/oauth-protected-resource/mcp with no SERVER_ROOT_PATH segment, but the routes are registered with the path-inserted root segment, so a proxy mounted under a sub-path pointed DCR clients at a URL that 404s. Both the aggregate challenge and the pre-existing per-server pass-through challenge now derive the path from one well_known_root_suffix helper that the route registrations also use, so the advertised URL cannot drift from the served route --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 54 +++++-------------- .../_experimental/mcp_server/oauth_utils.py | 12 +++++ .../proxy/_experimental/mcp_server/server.py | 6 ++- .../auth/test_user_api_key_auth_mcp.py | 17 ++++++ .../mcp_server/test_discoverable_endpoints.py | 45 ++++++++++++---- 6 files changed, 87 insertions(+), 52 deletions(-) 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 bf88c0be5da4..4f324e6b4230 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 @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -157,7 +158,9 @@ def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> H 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.""" error_attr = 'error="invalid_token", ' if invalid_token else "" - resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + resource_metadata_url = ( + f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" + ) return HTTPException( status_code=401, detail={ diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 39b84bdbd4eb..ed6c791064d4 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,6 +43,7 @@ TOKEN_NO_CACHE_HEADERS, get_request_base_url, validate_trusted_redirect_uri, + well_known_root_suffix, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -50,7 +51,6 @@ encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.utils import get_server_root_path from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1814,31 +1814,13 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _mcp_named_server_exists(request: Request) -> bool: - """True when a server literally named ``mcp`` is configured and visible to this caller. - - Its per-server authorization-server document is served at - ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the - aggregate path. When such a server exists the real server wins the route, so that - deployment keeps its per-server discovery regardless of whether the aggregate front door - is on.""" - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load - global_mcp_server_manager, - ) - - client_ip = IPAddressUtils.get_mcp_client_ip(request) - return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None - - # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client # pointed at {base}/mcp inserts the well-known segment before the resource # path, so this exact route must exist for aggregate discovery to work at all. # Declared before the parameterized well-known routes below: Starlette matches # in registration order, and /.well-known/oauth-authorization-server/{name} # would otherwise capture the "/mcp" suffix as a server name. -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp") async def oauth_protected_resource_aggregate(request: Request): """ OAuth protected resource discovery for the aggregate /mcp endpoint. @@ -1850,28 +1832,26 @@ async def oauth_protected_resource_aggregate(request: Request): return _build_aggregate_protected_resource_response(request) -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp") async def oauth_authorization_server_aggregate(request: Request): """ OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 path-inserted form for a client that treats {base}/mcp as its authorization base URL. - This single-segment path collides with the parameterized ``/{mcp_server_name}`` route - below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; - only when no such server exists is the aggregate document served. + The single-segment /mcp is reserved for the aggregate so the discovery chain stays + consistent: the aggregate protected-resource document advertises {base}/mcp as its + authorization server, so the document served here must have issuer {base}/mcp. A server + literally named ``mcp`` therefore does not take this route; it keeps its standard + two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the + per-server row win here instead would serve an issuer of {base} against a resource that + advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - if _mcp_named_server_exists(request): - return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_name: str): """ OAuth protected resource discovery endpoint using standard MCP URL pattern. @@ -1891,9 +1871,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") @router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp(request: Request, mcp_server_name: Optional[str] = None): """ @@ -1963,9 +1941,7 @@ def _build_oauth_authorization_server_response( # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -1980,9 +1956,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n # LiteLLM legacy pattern and root endpoint -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_mcp(request: Request, mcp_server_name: Optional[str] = None): """ diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858e..ccee3fc8ac00 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -132,6 +132,18 @@ def get_request_base_url(request: Request) -> str: return urlunparse((scheme, _strip_default_port(scheme, netloc), parsed.path, "", "", "")) +def well_known_root_suffix() -> str: + """The ``SERVER_ROOT_PATH`` segment inserted into a ``.well-known`` path (RFC 8414 / 9728 + path insertion), empty for a root-mounted proxy or an explicit ``/``. + + The discovery route registrations and the 401 challenges that advertise those routes both + derive their path from this one function, so the ``resource_metadata`` URL a client is told + to fetch cannot drift from the route that actually serves it. + """ + root = os.getenv("SERVER_ROOT_PATH", "") + return "" if root == "/" else root + + def validate_loopback_redirect_uri(redirect_uri: str) -> None: """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3 native-app pattern). MCP clients are native apps that listen on diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 68a61b851758..9513388014e9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -48,6 +48,7 @@ ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, @@ -3536,9 +3537,10 @@ def _get_passthrough_resource_metadata_url(scope: Scope, server_name: str) -> st base_url = get_request_base_url(request) _path = scope.get("_original_path") or scope.get("path", "") or "" + suffix = well_known_root_suffix() if _path.startswith(f"/{server_name}/mcp"): - return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" - return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/mcp/{server_name}" def _get_passthrough_www_authenticate( scope: Scope, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 2a28ff612717..e38101267023 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6008,6 +6008,23 @@ async def test_challenge_invalid_token_on_failed_bearer(self): www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + async def test_challenge_inserts_server_root_path(self): + """With SERVER_ROOT_PATH set the resource_metadata URL must carry the same path-inserted + root segment the aggregate PRM route is registered with (both derive it from + well_known_root_suffix), so a DCR client behind a sub-path is pointed at a route that + exists instead of a 404. Regression: the challenge used to hard-code /mcp and omit the + root path the route inserts.""" + import os + + with ( + patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}), + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/litellm/mcp"' in www_authenticate + async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index af4b2caca722..7e9ff4692b51 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7163,11 +7163,13 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata(): assert "none" in asm.json()["token_endpoint_auth_methods_supported"] -def test_as_aggregate_route_prefers_a_real_server_named_mcp(): - """A server literally named ``mcp`` wins the single-segment - /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized - /{server_name} route) and keeps its per-server discovery; the aggregate document is - served only when no such server exists.""" +def test_as_aggregate_route_reserves_mcp_for_the_aggregate(): + """The single-segment /.well-known/oauth-authorization-server/mcp is reserved for the + aggregate even when a server is literally named ``mcp``. The aggregate protected-resource + document advertises {base}/mcp as its authorization server, so the document served here + must carry issuer {base}/mcp for the RFC 8414 issuer check to pass. Letting the per-server + row win (issuer {base}) breaks that chain, so the aggregate wins and the mcp-named server + keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp.""" from fastapi import FastAPI from fastapi.testclient import TestClient @@ -7186,14 +7188,39 @@ def test_as_aggregate_route_prefers_a_real_server_named_mcp(): try: asm = client.get("/.well-known/oauth-authorization-server/mcp") assert asm.status_code == 200 - # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), - # not the aggregate one (whose issuer would be {base}/mcp) - assert asm.json()["issuer"] == "http://testserver" - assert "/mcp/authorize" in asm.json()["authorization_endpoint"] + # the aggregate document, whose issuer matches what the aggregate PRM advertises + assert asm.json()["issuer"] == "http://testserver/mcp" + + prm = client.get("/.well-known/oauth-protected-resource/mcp") + assert prm.status_code == 200 + assert prm.json()["authorization_servers"] == [asm.json()["issuer"]] + + # the mcp-named server keeps its own document on the standard two-segment route + per_server = client.get("/.well-known/oauth-authorization-server/mcp/mcp") + assert per_server.status_code == 200 + assert "/mcp/authorize" in per_server.json()["authorization_endpoint"] finally: global_mcp_server_manager.registry.clear() +def test_well_known_root_suffix_reflects_server_root_path(): + """The single path segment both the discovery routes and the 401 challenges insert for RFC + 8414/9728 path insertion: empty for a root-mounted proxy or an explicit ``/``, the configured + path otherwise. Sharing this one function is what keeps the advertised resource_metadata URL + equal to the route that serves it.""" + import os + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import well_known_root_suffix + + with patch.dict(os.environ, {"SERVER_ROOT_PATH": ""}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/"}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}): + assert well_known_root_suffix() == "/litellm" + + @pytest.mark.asyncio async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): """The always-on aggregate front door must not change bare-origin discovery: with one From 90195afa06c1c8454850d0a9b0e693175be8f7ce Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 00:14:10 -0700 Subject: [PATCH 04/16] test(mcp): isolate MCP discovery tests from a leaked SERVER_ROOT_PATH tests/test_litellm/proxy/test_custom_proxy.py sets SERVER_ROOT_PATH at import time (its app mounts under a custom path) and never restores it, so in a shared shard the value leaks into the process. The discovery routes and the 401 challenges now read SERVER_ROOT_PATH to path-insert it where they previously ignored it, so a leaked value rewrites every resource_metadata URL and the exact-URL assertions in the delegate, pass-through, and aggregate challenge tests fail depending on shard order An autouse fixture clears SERVER_ROOT_PATH for the MCP discovery tests so they deterministically exercise the default root-mounted deployment; the tests that assert a sub-path deployment set the value explicitly within their own body. No assertion changed; the leak was invisible before only because the code ignored the variable --- .../_experimental/mcp_server/conftest.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/conftest.py diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py new file mode 100644 index 000000000000..b477bf3f4062 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -0,0 +1,22 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True) +def _hermetic_server_root_path(): + """Isolate MCP discovery tests from a leaked ``SERVER_ROOT_PATH``. + + ``tests/test_litellm/proxy/test_custom_proxy.py`` sets ``SERVER_ROOT_PATH`` at import time + (its app mounts under a custom path) and never restores it, so in a shared shard the value + leaks into this process. The discovery routes and the 401 challenges read it, so a leaked + value would silently rewrite every ``resource_metadata`` URL and make these tests depend on + shard ordering. Clearing it here pins the default (root-mounted) deployment; a test that + exercises a sub-path deployment sets the value explicitly within its own body. + """ + saved = os.environ.pop("SERVER_ROOT_PATH", None) + try: + yield + finally: + if saved is not None: + os.environ["SERVER_ROOT_PATH"] = saved From 5b2877f7420e0f18e8312ee65a7ab689435d0d52 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 00:18:29 -0700 Subject: [PATCH 05/16] feat(mcp): identity-only session tokens for the gateway DCR front door --- .../session_credentials.py | 190 ++++++++++ .../outbound_credentials/session_token.py | 356 ++++++++++++++++++ .../test_session_credentials.py | 135 +++++++ .../test_session_token.py | 206 ++++++++++ 4 files changed, 887 insertions(+) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py new file mode 100644 index 000000000000..08d5cc8b1f13 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -0,0 +1,190 @@ +"""Producer and consumer helpers for the gateway-level DCR session token. + +The aggregate ``/mcp`` front door (``mcp_gateway_dcr``) issues the identity-only session +tokens defined in :mod:`.session_token`. The gateway token endpoint mints them (producer) +after SSO sign-in, and at the MCP admission edge the gateway derives the session signing +key from the proxy ``master_key``, opens the bearer, and admits the request under the +recovered litellm user (consumer), reloading the live user record and policy before +anything runs. This module is the pure surface for both sides; the token-endpoint and +admission wiring live in their respective call sites. + +The signing key is derived with the same memory-hard scrypt construction as +:func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain +label, so session tokens and bridge envelopes never share key material: a token of one +family is unverifiable in the other by key separation, on top of the distinct issuers, +prefixes, and claim shapes. +""" + +import hashlib +from datetime import datetime +from functools import lru_cache +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + OpenedSessionToken, + SessionExpired, + SessionKeys, + SessionPrincipal, + is_session_refresh_token, + is_session_token, + open_session_refresh_token, + open_session_token, +) + +_SESSION_SIGNING_KEY_DOMAIN = b"litellm-mcp-gateway:session-signing:" + +# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured +# session token is not a cheap offline oracle for the master key. +_SCRYPT_N = 2**15 +_SCRYPT_R = 8 +_SCRYPT_P = 1 +_SCRYPT_MAXMEM = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 +_DERIVED_KEY_BYTES = 32 + + +@lru_cache(maxsize=8) +def session_keys_from_master_key(master_key: str) -> SessionKeys: + """Derive the session signing key from the proxy master key. + + A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a + 256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on + the key without persisting any. The domain label differs from both envelope labels in + :mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses + into the other. The result is cached (the master key is fixed for a process); rotating + ``master_key`` invalidates every outstanding session, which is the intended behavior + for a signing-key change. + """ + signing = hashlib.scrypt( + master_key.encode(), + salt=_SESSION_SIGNING_KEY_DOMAIN, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=_DERIVED_KEY_BYTES, + ).hex() + return SessionKeys(signing_key=SecretStr(signing)) + + +class NotSessionBearer(BaseModel): + """The bearer is not session-shaped; admission continues on its normal path.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_session_bearer"] = "not_session_bearer" + + +class SessionBearerAdmitted(BaseModel): + """A valid session access token: the principal to admit under after a live reload.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["admitted"] = "admitted" + principal: SessionPrincipal + + +class SessionBearerInvalid(BaseModel): + """The bearer is session-shaped but must not admit (expired, tampered, wrong key, or a + refresh token presented at the tool-call edge); admission fails closed with the + ``invalid_token`` challenge rather than falling through to another arm. ``expired`` + distinguishes a routine expiry (debug-log worthy) from a tampered or foreign token.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + expired: bool = False + + +SessionBearerResult: TypeAlias = NotSessionBearer | SessionBearerAdmitted | SessionBearerInvalid + + +def _strip_bearer(value: str) -> str: + parts = value.split(None, 1) + if len(parts) == 2 and parts[0].lower() == "bearer": + return parts[1] + return value + + +def is_session_bearer_shaped(authorization_value: str) -> bool: + """Cheap, keyless test that an ``Authorization`` value carries a session token of either + kind (optional ``Bearer`` scheme stripped). The admission edge engages the session arm + for an access token (to admit) and for a refresh token (to reject it explicitly, since + a refresh credential is never usable at the tool-call edge); anything else falls + through to normal admission.""" + candidate = _strip_bearer(authorization_value) + return is_session_token(candidate) or is_session_refresh_token(candidate) + + +def resolve_session_bearer( + authorization_value: str, + keys: SessionKeys, + now: datetime, +) -> SessionBearerResult: + """Classify an ``Authorization`` value presented at the aggregate MCP edge. + + Strips an optional ``Bearer`` scheme, then returns ``NotSessionBearer`` for a + non-session bearer (normal admission continues), ``SessionBearerAdmitted`` with the + recovered principal for a valid access token, and ``SessionBearerInvalid`` for a + session-shaped bearer that must not admit. Never raises: total over hostile input via + :func:`~.session_token.open_session_token`. + + A refresh token is ``SessionBearerInvalid`` here: it is a valid gateway credential but + only ever presented back to the token endpoint, so admission must fail it closed rather + than let it fall through to another arm. + """ + candidate = _strip_bearer(authorization_value) + if is_session_refresh_token(candidate): + return SessionBearerInvalid() + if not is_session_token(candidate): + return NotSessionBearer() + opened = open_session_token(candidate, keys, now) + if isinstance(opened, OpenedSessionToken): + return SessionBearerAdmitted(principal=opened.principal) + return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) + + +class SessionRefreshOpened(BaseModel): + """A valid session refresh token presented to the token endpoint: the principal to + re-validate and renew under.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["opened"] = "opened" + principal: SessionPrincipal + + +class SessionRefreshInvalid(BaseModel): + """The presented refresh grant is not a valid session refresh token for this client + (not refresh-shaped, will not open, or bound to a different ``client_id``); the token + endpoint fails the refresh closed.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + + +SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid + + +def open_session_refresh_bearer( + refresh_value: str, + keys: SessionKeys, + now: datetime, + expected_client_id: str, +) -> SessionRefreshResult: + """Open a session refresh token presented on a ``refresh_token`` grant. + + The token-endpoint mirror of :func:`resolve_session_bearer`: strips an optional + ``Bearer`` scheme, then returns ``SessionRefreshOpened`` with the recovered principal, + or ``SessionRefreshInvalid`` for anything that is not a valid session refresh token + issued to ``expected_client_id``. Never raises. The client binding (RFC 6749 section 6) + stops a refresh token stolen from one DCR client from being renewed through another; + ``client_id`` is not a secret (the caller presents it), so a plain equality check is + sufficient and, unlike ``hmac.compare_digest`` on ``str``, does not raise on non-ASCII. + """ + candidate = _strip_bearer(refresh_value) + if not is_session_refresh_token(candidate): + return SessionRefreshInvalid() + opened = open_session_refresh_token(candidate, keys, now) + if not isinstance(opened, OpenedSessionToken): + return SessionRefreshInvalid() + if opened.principal.client_id != expected_client_id: + return SessionRefreshInvalid() + return SessionRefreshOpened(principal=opened.principal) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py new file mode 100644 index 000000000000..78b1f7e49166 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -0,0 +1,356 @@ +"""Identity-only session tokens for the gateway-level (aggregate ``/mcp``) DCR front door. + +A DCR client that signs in through LiteLLM SSO holds ONE bearer that carries ONLY a +litellm identity; unlike the :mod:`.envelope` bridge bearer it seals no upstream +credential, because the custody model vaults every upstream token server-side in +``LiteLLM_MCPUserCredentials`` and egress resolves them by user at call time. The token +is therefore a stable REFERENCE, not an authorization: admission reloads the live user +record and policy on every request, so deactivating the user (or their team) kills +outstanding sessions immediately without a revocation store. + +Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, +the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` +plus ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token +to the DCR client it was issued to (RFC 6749 section 6) and is carried on the access +token for parity and audit. There is no encrypted payload: nothing in a session token +is secret beyond the signature, and reprs never print the signed value because minted +tokens are ``SecretStr``. + +This module is pure and unwired: it imports nothing from endpoint or edge code, reads +no proxy globals, and takes all key material and the clock as explicit parameters. +Failures are values: :func:`open_session_token` and :func:`open_session_refresh_token` +are total over hostile, attacker-controlled input and return a +``SessionTokenOpenError`` variant rather than raising. PyJWT's ``iat``/``nbf``/``exp`` +validators are disabled for the same reasons documented in :mod:`.envelope` (they +raise on hostile claim types and compare against the wall clock instead of the +injected ``now``); the strict pydantic claims model is the sole, total type gate. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import jwt +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError + +SESSION_TOKEN_PREFIX = "llm_session_" +"""Marker prefix on every serialized session ACCESS token so the admission edge can cheaply +tell a gateway session from a litellm key, JWT, or bridge envelope before doing any +cryptography. Distinct from the ``llm_env_``/``llm_refresh_`` envelope prefixes.""" + +SESSION_REFRESH_PREFIX = "llm_srefresh_" +"""Marker prefix on every serialized session REFRESH token. A distinct prefix keeps the two +credentials routable without crypto and, together with the signed ``kind`` claim, stops one +from being presented where the other is expected: the refresh token is only ever presented +back to the token endpoint, never at the MCP edge.""" + +SESSION_ISSUER = "litellm-mcp-gateway" +"""``iss`` claim stamped into every session token and required back on open. Distinct from +the envelope issuer so a token of one family can never validate in the other even under a +hypothetical shared signing key.""" + +SESSION_TTL_SECONDS = 3600 +"""Session ACCESS token lifetime (1h), matching the access-envelope and BYOK session bearer +windows: a client-held credential never outlives a bounded window, and each refresh +re-validates the live user before re-minting.""" + +SESSION_REFRESH_TTL_SECONDS = 1209600 +"""Session REFRESH token lifetime (14 days), matching the refresh-envelope bound. Each +renewal re-validates the sealed user against the live record (deactivation gates it) and +rotates the refresh token, so the practical bound is idle time, not a fixed session.""" + +MAX_SESSION_TOKEN_BYTES = 4096 +"""Size cap on the serialized token (prefix + JWT, in bytes) and on any candidate accepted +by the openers. Session claims are small; the only variable-length field is ``client_id`` +(a sealed DCR client record), and 4096 leaves ample headroom under common 8-16KB header +limits while bounding hostile input before JWT parsing.""" + +_SESSION_JWT_ALGORITHM = "HS256" + +SessionTokenKind = Literal["session", "session_refresh"] +"""Which credential a session token is. Stamped into the signed claims and required to match +on open, so a signature-valid token of one kind cannot be replayed as the other even if its +wire prefix is swapped (the prefix is not part of the signed payload; this claim is).""" + + +class SessionPrincipal(BaseModel): + """The litellm user a session token identifies and the DCR client it was issued to. + + ``user_id`` is the SSO-established litellm user subject, never a credential: admission + reloads the live user record by it, so current role, team, and revocation state are + enforced at use time rather than frozen at mint time. ``client_id`` is the (stateless, + gateway-sealed) DCR client identifier the token was issued to; the token endpoint + requires it to match on the refresh grant. + """ + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + + +class SessionKeys(BaseModel): + """Injected key material: the HS256 signing key. + + ``signing_key`` must be at least 32 bytes: HS256's HMAC-SHA256 has a 256-bit security + level, RFC 7518 requires a key of at least that size, and a shorter key makes PyJWT + emit ``InsecureKeyLengthWarning``. + """ + + model_config = ConfigDict(frozen=True) + signing_key: SecretStr = Field(min_length=32) + + +class MintedSessionToken(BaseModel): + """A minted session token: the client-held bearer value and when it expires.""" + + model_config = ConfigDict(frozen=True) + token: SecretStr + expires_at: datetime + + +class OpenedSessionToken(BaseModel): + """A validated session token of either kind: the principal it was minted for.""" + + model_config = ConfigDict(frozen=True) + principal: SessionPrincipal + + +class SessionTokenTooLarge(BaseModel): + """The serialized token exceeded ``MAX_SESSION_TOKEN_BYTES``; carries sizes only. Only + reachable through an oversized ``client_id``, which registration should have bounded.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_token_too_large"] = "session_token_too_large" + size_bytes: int + max_bytes: int + + +SessionTokenMintError: TypeAlias = SessionTokenTooLarge + + +class NotASessionToken(BaseModel): + """The candidate does not carry the expected session prefix.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_a_session_token"] = "not_a_session_token" + + +class SessionBadSignature(BaseModel): + """The JWT signature does not verify under the provided signing key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_bad_signature"] = "session_bad_signature" + + +class SessionExpired(BaseModel): + """The token's ``exp`` is not in the future relative to the provided ``now``.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_expired"] = "session_expired" + + +class SessionMalformed(BaseModel): + """The token is not a well-formed session token: undecodable JWT, wrong issuer, wrong + ``kind``, or missing/mistyped/extra claims.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_malformed"] = "session_malformed" + + +SessionTokenOpenError: TypeAlias = NotASessionToken | SessionBadSignature | SessionExpired | SessionMalformed + + +class _SessionClaims(BaseModel): + """Decoded-claims boundary that pins the exact shape the mints emit. + + ``user_id``/``client_id`` mirror the ``min_length`` constraints of + :class:`SessionPrincipal` so any claim set that validates here also constructs a + principal, keeping the openers raise-free: a correctly signed JWT with an empty + identity claim fails here and maps to ``SessionMalformed``. ``strict`` rejects coerced + types (``exp: "123"``) and ``extra="forbid"`` rejects any claim the gateway never + mints; PyJWT's own registered-claim validators are disabled at decode (see module + docstring), so this model is the sole, total type gate for every claim. + """ + + model_config = ConfigDict(frozen=True, strict=True, extra="forbid") + iss: str + iat: int + exp: int + kind: SessionTokenKind + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + + +def is_session_token(candidate: str) -> bool: + """Cheap prefix check for a session ACCESS token so the admission edge can route gateway + sessions vs keys, JWTs, and envelopes without crypto.""" + return candidate.startswith(SESSION_TOKEN_PREFIX) + + +def is_session_refresh_token(candidate: str) -> bool: + """Cheap prefix check for a session REFRESH token so the token endpoint can route a + refresh grant without crypto.""" + return candidate.startswith(SESSION_REFRESH_PREFIX) + + +def mint_session_token( + principal: SessionPrincipal, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenMintError: + """Mint the short-lived session ACCESS token for ``principal``. + + ``exp`` is ``SESSION_TTL_SECONDS`` from ``now``. Returns ``SessionTokenTooLarge`` when + the serialized token exceeds ``MAX_SESSION_TOKEN_BYTES``. + """ + return _mint( + kind="session", + prefix=SESSION_TOKEN_PREFIX, + principal=principal, + expires_at=now + timedelta(seconds=SESSION_TTL_SECONDS), + keys=keys, + now=now, + ) + + +def mint_session_refresh_token( + principal: SessionPrincipal, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenMintError: + """Mint the long-lived session REFRESH token for ``principal``. + + ``exp`` is ``SESSION_REFRESH_TTL_SECONDS`` from ``now``. Minting a distinct + ``kind="session_refresh"`` claim is what keeps a refresh token from ever opening as an + access credential at the MCP edge. + """ + return _mint( + kind="session_refresh", + prefix=SESSION_REFRESH_PREFIX, + principal=principal, + expires_at=now + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS), + keys=keys, + now=now, + ) + + +def open_session_token( + candidate: str, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Validate a session ACCESS ``candidate`` and recover the principal. + + Never raises for bad input: every invalid, expired, tampered, or wrong-kind candidate + maps to a distinct ``SessionTokenOpenError`` variant. + """ + return _open(candidate, prefix=SESSION_TOKEN_PREFIX, expected_kind="session", keys=keys, now=now) + + +def open_session_refresh_token( + candidate: str, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Validate a session REFRESH ``candidate`` and recover the principal. + + Total over hostile input exactly like :func:`open_session_token`. The + ``kind="session_refresh"`` claim is required, so an access token re-prefixed as a + refresh one is rejected as ``SessionMalformed``. + """ + return _open(candidate, prefix=SESSION_REFRESH_PREFIX, expected_kind="session_refresh", keys=keys, now=now) + + +def _mint( + kind: SessionTokenKind, + prefix: str, + principal: SessionPrincipal, + expires_at: datetime, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenTooLarge: + """Sign the claims for either token kind and enforce the size cap. Shared by both mints + so the JWT shape, issuer, and size guard cannot drift between access and refresh.""" + claims = _SessionClaims( + iss=SESSION_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + kind=kind, + user_id=principal.user_id, + client_id=principal.client_id, + ) + token = prefix + jwt.encode( + claims.model_dump(), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM + ) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_SESSION_TOKEN_BYTES: + return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES) + return MintedSessionToken(token=SecretStr(token), expires_at=expires_at) + + +def _open( + candidate: str, + prefix: str, + expected_kind: SessionTokenKind, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an + attacker-controlled candidate, shared by both openers so the security gate is identical + for access and refresh. Returns the opened token or a distinct error; never raises.""" + if not candidate.startswith(prefix): + return NotASessionToken() + # UTF-8 byte length is never below character length, so a character count already over + # the cap rejects an oversize candidate in O(1) without encoding it; the exact byte + # check then runs only on candidates already bounded to the cap in characters. + if len(candidate) > MAX_SESSION_TOKEN_BYTES: + return SessionMalformed() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES: + return SessionMalformed() + claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + if not isinstance(claims, _SessionClaims): + return claims + if claims.kind != expected_kind: + return SessionMalformed() + if now.timestamp() >= claims.exp: + return SessionExpired() + return OpenedSessionToken(principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id)) + + +def _decode_claims( + compact: str, + signing_key: SecretStr, +) -> _SessionClaims | SessionBadSignature | SessionMalformed: + """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + + ``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller. + PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected + ``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature + mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces + as ``UnicodeEncodeError`` (a ``ValueError``), a non-string registered claim as a + ``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid + token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate. + """ + try: + payload = jwt.decode( + compact, + signing_key.get_secret_value(), + algorithms=[_SESSION_JWT_ALGORITHM], + issuer=SESSION_ISSUER, + options={ + "verify_exp": False, + "verify_iat": False, + "verify_nbf": False, + "require": ["iss", "iat", "exp"], + }, + ) + except jwt.InvalidSignatureError: + return SessionBadSignature() + except (jwt.InvalidTokenError, ValueError, TypeError): + return SessionMalformed() + try: + return _SessionClaims.model_validate(payload) + except ValidationError: + return SessionMalformed() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py new file mode 100644 index 000000000000..8fa7c15d2d36 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -0,0 +1,135 @@ +"""Tests for the session-token KDF and the edge/token-endpoint resolvers.""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + NotSessionBearer, + SessionBearerAdmitted, + SessionBearerInvalid, + SessionRefreshInvalid, + SessionRefreshOpened, + is_session_bearer_shaped, + open_session_refresh_bearer, + resolve_session_bearer, + session_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_TTL_SECONDS, + MintedSessionToken, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, +) + +NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) +MASTER_KEY = "sk-master-key-for-tests" +KEYS = session_keys_from_master_key(MASTER_KEY) +PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc") + + +def _access_token() -> str: + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _refresh_token() -> str: + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def test_kdf_is_deterministic_and_key_length_is_256_bit(): + again = session_keys_from_master_key(MASTER_KEY) + assert again.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value() + assert len(bytes.fromhex(KEYS.signing_key.get_secret_value())) == 32 + + +def test_kdf_domain_separated_from_envelope_keys(): + envelope_keys = envelope_keys_from_master_key(MASTER_KEY) + session_signing = KEYS.signing_key.get_secret_value() + assert session_signing != envelope_keys.signing_key.get_secret_value() + assert session_signing != envelope_keys.encryption_key.get_secret_value() + + +def test_kdf_differs_across_master_keys(): + other = session_keys_from_master_key("sk-a-different-master-key") + assert other.signing_key.get_secret_value() != KEYS.signing_key.get_secret_value() + + +@pytest.mark.parametrize( + "value,expected", + [ + ("Bearer sk-1234", False), + ("sk-1234", False), + ("Bearer llm_env_abc", False), + ("Bearer llm_refresh_abc", False), + ("llm_session_abc", True), + ("Bearer llm_session_abc", True), + ("bearer llm_srefresh_abc", True), + ], +) +def test_is_session_bearer_shaped(value, expected): + assert is_session_bearer_shaped(value) is expected + + +def test_resolve_admits_valid_access_token_with_and_without_scheme(): + token = _access_token() + for value in (token, f"Bearer {token}", f"bearer {token}"): + result = resolve_session_bearer(value, KEYS, NOW) + assert isinstance(result, SessionBearerAdmitted) + assert result.principal == PRINCIPAL + + +def test_resolve_passes_non_session_bearers_through(): + for value in ("Bearer sk-1234", "Bearer llm_env_whatever", "Bearer eyJhbGciOi"): + assert isinstance(resolve_session_bearer(value, KEYS, NOW), NotSessionBearer) + + +def test_resolve_fails_expired_token_closed_and_flags_expiry(): + token = _access_token() + later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result = resolve_session_bearer(f"Bearer {token}", KEYS, later) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is True + + +def test_resolve_fails_tampered_token_closed_without_expiry_flag(): + token = _access_token() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + result = resolve_session_bearer(f"Bearer {tampered}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_resolve_rejects_refresh_token_at_the_edge(): + result = resolve_session_bearer(f"Bearer {_refresh_token()}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_resolve_wrong_master_key_fails_closed(): + other_keys = session_keys_from_master_key("sk-rotated-master-key") + result = resolve_session_bearer(f"Bearer {_access_token()}", other_keys, NOW) + assert isinstance(result, SessionBearerInvalid) + + +def test_refresh_grant_opens_for_the_issued_client(): + result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_abc") + assert isinstance(result, SessionRefreshOpened) + assert result.principal == PRINCIPAL + + +def test_refresh_grant_rejects_a_different_client(): + result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_other") + assert isinstance(result, SessionRefreshInvalid) + + +def test_refresh_grant_rejects_access_token_presented_as_refresh(): + result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc") + assert isinstance(result, SessionRefreshInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py new file mode 100644 index 000000000000..551270f8d4b6 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -0,0 +1,206 @@ +"""Tests for the identity-only gateway session token (mint/open, hostile-input totality).""" + +from datetime import datetime, timedelta, timezone + +import jwt +import pytest +from pydantic import SecretStr, ValidationError + +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + MAX_SESSION_TOKEN_BYTES, + SESSION_ISSUER, + SESSION_REFRESH_PREFIX, + SESSION_REFRESH_TTL_SECONDS, + SESSION_TOKEN_PREFIX, + SESSION_TTL_SECONDS, + MintedSessionToken, + NotASessionToken, + OpenedSessionToken, + SessionBadSignature, + SessionExpired, + SessionKeys, + SessionMalformed, + SessionPrincipal, + SessionTokenTooLarge, + is_session_refresh_token, + is_session_token, + mint_session_refresh_token, + mint_session_token, + open_session_refresh_token, + open_session_token, +) + +NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) +KEYS = SessionKeys(signing_key=SecretStr("k" * 32)) +OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32)) +PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc") + + +def _mint_access() -> str: + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _mint_refresh() -> str: + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _sign_claims(payload: dict, prefix: str = SESSION_TOKEN_PREFIX, keys: SessionKeys = KEYS) -> str: + return prefix + jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm="HS256") + + +def _valid_claims(**overrides) -> dict: + base = { + "iss": SESSION_ISSUER, + "iat": int(NOW.timestamp()), + "exp": int((NOW + timedelta(seconds=600)).timestamp()), + "kind": "session", + "user_id": "user-123", + "client_id": "llm_client_abc", + } + return {**base, **overrides} + + +def test_access_round_trip_recovers_principal_and_caps_ttl(): + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.expires_at == NOW + timedelta(seconds=SESSION_TTL_SECONDS) + token = minted.token.get_secret_value() + assert is_session_token(token) + assert not is_session_refresh_token(token) + opened = open_session_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_refresh_round_trip_recovers_principal_and_caps_ttl(): + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.expires_at == NOW + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS) + token = minted.token.get_secret_value() + assert is_session_refresh_token(token) + opened = open_session_refresh_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_access_token_reprefixed_as_refresh_is_rejected_by_signed_kind(): + body = _mint_access().removeprefix(SESSION_TOKEN_PREFIX) + swapped = SESSION_REFRESH_PREFIX + body + assert isinstance(open_session_refresh_token(swapped, KEYS, NOW), SessionMalformed) + + +def test_refresh_token_reprefixed_as_access_is_rejected_by_signed_kind(): + body = _mint_refresh().removeprefix(SESSION_REFRESH_PREFIX) + swapped = SESSION_TOKEN_PREFIX + body + assert isinstance(open_session_token(swapped, KEYS, NOW), SessionMalformed) + + +def test_refresh_token_is_not_an_access_token_at_the_edge(): + assert isinstance(open_session_token(_mint_refresh(), KEYS, NOW), NotASessionToken) + + +def test_expired_access_token_is_expired_not_malformed(): + token = _mint_access() + at_expiry = NOW + timedelta(seconds=SESSION_TTL_SECONDS) + assert isinstance(open_session_token(token, KEYS, at_expiry), SessionExpired) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(token, KEYS, after), SessionExpired) + + +def test_still_valid_one_second_before_expiry(): + token = _mint_access() + just_before = NOW + timedelta(seconds=SESSION_TTL_SECONDS - 1) + assert isinstance(open_session_token(token, KEYS, just_before), OpenedSessionToken) + + +def test_tampered_signature_is_bad_signature(): + token = _mint_access() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, KEYS, NOW), SessionBadSignature) + + +def test_key_rotation_invalidates_outstanding_tokens(): + token = _mint_access() + assert isinstance(open_session_token(token, OTHER_KEYS, NOW), SessionBadSignature) + + +@pytest.mark.parametrize( + "candidate,expected", + [ + ("sk-1234", NotASessionToken), + ("llm_env_something", NotASessionToken), + ("", NotASessionToken), + (SESSION_TOKEN_PREFIX, SessionMalformed), + (SESSION_TOKEN_PREFIX + "not-a-jwt", SessionMalformed), + (SESSION_TOKEN_PREFIX + "\ud800garbage", SessionMalformed), + (SESSION_TOKEN_PREFIX + "a" * (MAX_SESSION_TOKEN_BYTES + 1), SessionMalformed), + ], +) +def test_hostile_candidates_never_raise(candidate, expected): + assert isinstance(open_session_token(candidate, KEYS, NOW), expected) + + +def test_multibyte_candidate_over_byte_cap_but_under_char_cap_is_rejected(): + filler = "€" * (MAX_SESSION_TOKEN_BYTES // 3) + candidate = SESSION_TOKEN_PREFIX + filler + assert len(candidate) <= MAX_SESSION_TOKEN_BYTES + assert isinstance(open_session_token(candidate, KEYS, NOW), SessionMalformed) + + +def test_alg_none_token_is_rejected(): + unsigned = jwt.api_jws.encode(b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none") + assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, KEYS, NOW), SessionMalformed) + + +@pytest.mark.parametrize( + "claims", + [ + _valid_claims(iss="wrong-issuer"), + _valid_claims(exp=str(int((NOW + timedelta(seconds=600)).timestamp()))), + _valid_claims(iat="evil"), + _valid_claims(kind="access"), + _valid_claims(user_id=""), + _valid_claims(nbf=0), + {k: v for k, v in _valid_claims().items() if k != "client_id"}, + {k: v for k, v in _valid_claims().items() if k != "exp"}, + ], +) +def test_signed_but_malformed_claims_are_rejected_without_raising(claims): + token = _sign_claims(claims) + assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed) + + +def test_signed_claims_with_exact_shape_open(): + token = _sign_claims(_valid_claims()) + opened = open_session_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal.user_id == "user-123" + + +def test_oversized_client_id_fails_mint_with_typed_error_not_truncation(): + principal = SessionPrincipal(user_id="user-123", client_id="c" * (MAX_SESSION_TOKEN_BYTES + 100)) + minted = mint_session_token(principal, KEYS, NOW) + assert isinstance(minted, SessionTokenTooLarge) + assert minted.max_bytes == MAX_SESSION_TOKEN_BYTES + + +def test_empty_principal_fields_rejected_at_construction(): + with pytest.raises(ValidationError): + SessionPrincipal(user_id="", client_id="c") + with pytest.raises(ValidationError): + SessionPrincipal(user_id="u", client_id="") + + +def test_short_signing_key_rejected_at_construction(): + with pytest.raises(ValidationError): + SessionKeys(signing_key=SecretStr("short")) + + +def test_minted_token_repr_never_leaks_value(): + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.token.get_secret_value() not in repr(minted) From 6e7238633ef8493ac82b496479f736cd1d7de063 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 00:31:06 -0700 Subject: [PATCH 06/16] feat(mcp): add jti claim for per-mint session token uniqueness --- .../mcp_server/outbound_credentials/session_token.py | 7 ++++++- .../mcp_server/outbound_credentials/test_session_token.py | 8 ++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) 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 78b1f7e49166..9325428f0492 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -10,7 +10,9 @@ Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` -plus ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token +plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never +collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and +``client_id``; ``client_id`` binds the refresh token to the DCR client it was issued to (RFC 6749 section 6) and is carried on the access token for parity and audit. There is no encrypted payload: nothing in a session token is secret beyond the signature, and reprs never print the signed value because minted @@ -28,6 +30,7 @@ from __future__ import annotations +import secrets from datetime import datetime, timedelta from typing import Literal, TypeAlias @@ -177,6 +180,7 @@ class _SessionClaims(BaseModel): iss: str iat: int exp: int + jti: str = Field(min_length=1) kind: SessionTokenKind user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) @@ -276,6 +280,7 @@ def _mint( iss=SESSION_ISSUER, iat=int(now.timestamp()), exp=int(expires_at.timestamp()), + jti=secrets.token_urlsafe(16), kind=kind, user_id=principal.user_id, client_id=principal.client_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 551270f8d4b6..a43592ebe183 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -57,6 +57,7 @@ def _valid_claims(**overrides) -> dict: "iss": SESSION_ISSUER, "iat": int(NOW.timestamp()), "exp": int((NOW + timedelta(seconds=600)).timestamp()), + "jti": "jti-fixed", "kind": "session", "user_id": "user-123", "client_id": "llm_client_abc", @@ -200,6 +201,13 @@ def test_short_signing_key_rejected_at_construction(): SessionKeys(signing_key=SecretStr("short")) +def test_two_mints_of_the_same_principal_are_distinct_tokens(): + first = mint_session_token(PRINCIPAL, KEYS, NOW) + second = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(first, MintedSessionToken) and isinstance(second, MintedSessionToken) + assert first.token.get_secret_value() != second.token.get_secret_value() + + def test_minted_token_repr_never_leaks_value(): minted = mint_session_token(PRINCIPAL, KEYS, NOW) assert isinstance(minted, MintedSessionToken) From 22a7174e3c88dc5658277eb64328d00431bc95b2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:21:14 -0700 Subject: [PATCH 07/16] fix(ui): stamp exp on the UI session cookie so bounded-lifetime readers accept it --- litellm/proxy/auth/login_utils.py | 26 +++++++++ litellm/proxy/management_endpoints/ui_sso.py | 8 +-- litellm/proxy/proxy_server.py | 30 ++-------- .../proxy/auth/test_login_utils.py | 55 +++++++++++++++++++ 4 files changed, 90 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 11f12e597b97..f35d94c986e1 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -7,12 +7,15 @@ import os import secrets +from datetime import datetime, timedelta, timezone from typing import Literal, Optional, cast +import jwt from fastapi import HTTPException import litellm from litellm.constants import LITELLM_PROXY_ADMIN_NAME, LITELLM_UI_SESSION_DURATION +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_UserTable, LitellmUserRoles, @@ -313,6 +316,29 @@ async def authenticate_user( ) +def _ui_session_exp_timestamp() -> int: + """The ``exp`` claim (unix seconds) for a UI session cookie, ``LITELLM_UI_SESSION_DURATION`` + from now. The virtual key sealed inside the cookie already expires after this same + duration; stamping the JWT itself gives the cookie the bounded lifetime the dashboard's + client-side expiry check and the server-side session-cookie readers both assume, instead + of a token that stays signature-valid until the master key rotates.""" + ttl_seconds = duration_in_seconds(LITELLM_UI_SESSION_DURATION) + return int((datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).timestamp()) + + +def encode_ui_session_jwt(returned_ui_token_object: ReturnedUITokenObject, master_key: str) -> str: + """Encode a UI session cookie JWT with a bounded ``exp``. + + The single choke point every UI login path (SSO and username/password /login, /v2, + /v3) uses to mint the ``token`` cookie, so the cookie's lifetime is set in exactly one + place and cannot drift between paths. Without the ``exp`` the cookie is valid until the + master key rotates, and the session-cookie readers that require a bounded lifetime + (the MCP interactive sign-in) reject it. + """ + claims = {**cast(dict, returned_ui_token_object), "exp": _ui_session_exp_timestamp()} + return jwt.encode(claims, master_key, algorithm="HS256") + + def create_ui_token_object( login_result: LoginResult, general_settings: dict, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0475566192e2..fe6682e4221b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3216,11 +3216,9 @@ async def get_redirect_response_from_openid( server_root_path=get_server_root_path(), ) - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - master_key or "", - algorithm="HS256", - ) + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "") # Control-plane cross-origin: store JWT behind a single-use opaque # code (60s TTL) so the token never appears in browser history / logs. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0eb42b266e1..0dbb3cf94b22 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13339,7 +13339,7 @@ async def fallback_login(request: Request): @router.post("/login", include_in_schema=False) # hidden since this is a helper for UI sso login async def login(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url form = await request.form() @@ -13362,13 +13362,7 @@ async def login(request: Request): ) # Generate JWT token - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) # Build redirect URL litellm_dashboard_ui = get_custom_url(str(request.base_url)) @@ -13387,7 +13381,7 @@ async def login(request: Request): @router.post("/v2/login", include_in_schema=False) # hidden helper for UI logins via API async def login_v2(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13408,13 +13402,7 @@ async def login_v2(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -13458,7 +13446,7 @@ async def login_v2(request: Request): ) # control-plane login — always returns token in body for cross-origin use async def login_v3(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13487,13 +13475,7 @@ async def login_v3(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 288e2533b72d..e301e878bf4b 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -559,3 +559,58 @@ def mock_find_first(**kwargs): assert isinstance(result, LoginResult) assert result.user_id == "test-user-123" assert result.user_email == user_email + + +class TestEncodeUiSessionJwt: + """The UI session cookie must carry a bounded exp so it does not stay + signature-valid until the master key rotates, and so the session-cookie readers + that require a bounded lifetime (the MCP interactive sign-in) accept it.""" + + def _decode(self, token: str) -> dict: + import jwt + + return jwt.decode(token, "sk-master-for-tests", algorithms=["HS256"]) + + def test_encoded_cookie_carries_bounded_exp(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "u1", "key": "sk-abc", "login_method": "username_password"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + claims = self._decode(token) + assert claims["user_id"] == "u1" + assert claims["login_method"] == "username_password" + remaining = claims["exp"] - int(time.time()) + assert 23 * 3600 < remaining <= 24 * 3600 + + def test_duration_is_honored_from_env(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "1h"): + token = encode_ui_session_jwt({"user_id": "u1"}, "sk-master-for-tests") + remaining = self._decode(token)["exp"] - int(time.time()) + assert 0 < remaining <= 3600 + + def test_cookie_is_accepted_by_the_exp_requiring_session_reader(self): + """The regression this change exists for: before it, the UI cookie carried no + exp and _user_id_from_session_cookie (require=["exp"]) rejected every real login, + so the MCP interactive sign-in could never capture identity. A cookie minted by + this helper must now be accepted.""" + from unittest.mock import MagicMock + + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + _user_id_from_session_cookie, + ) + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "cornell-user", "key": "sk-abc", "login_method": "sso"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + request = MagicMock() + request.cookies = {"token": token} + with patch("litellm.proxy.proxy_server.master_key", "sk-master-for-tests"): + assert _user_id_from_session_cookie(request) == "cornell-user" From f546cfcb5f47264d24295eaf3f316cd1b3fa2f47 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:33:45 -0700 Subject: [PATCH 08/16] feat(mcp): aggregate DCR register, authorize, complete, and token flow for the gateway front door --- .../mcp_server/discoverable_endpoints.py | 75 ++- .../mcp_server/gateway_dcr_flow.py | 519 ++++++++++++++++++ litellm/proxy/management_endpoints/ui_sso.py | 21 +- .../mcp_server/test_discoverable_endpoints.py | 61 ++ .../mcp_server/test_gateway_dcr_flow.py | 379 +++++++++++++ .../proxy/management_endpoints/test_ui_sso.py | 21 + 6 files changed, 1070 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ed6c791064d4..ab9339593c47 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -39,6 +39,14 @@ dcr_fault_detail, render_token_fault, ) +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + aggregate_authorize, + aggregate_token, + complete_connect_flow, + is_gateway_dcr_client_id, + register_aggregate_client, + relative_request_url, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, @@ -217,14 +225,25 @@ def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None return None +def _session_cookie_user_id(request: Request) -> str | None: + """The signed-in litellm user for a browser request, or ``None``. Thin wrapper so the + aggregate DCR flow's verbs receive the identity as a plain value instead of parsing + cookies themselves.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # circular import at module load + _user_id_from_session_cookie, + ) + + return _user_id_from_session_cookie(request) + + def _redirect_to_litellm_login(request: Request) -> RedirectResponse: """Send an unauthenticated browser through litellm login before the interactive bridge authorize can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code, - so a session is required; without one there is nothing to bind. After login the user re-initiates - the connection, which then finds the session cookie (the seamless return-to round-trip, which is - origin-validated against the control-plane URL, is a follow-up).""" + so a session is required; without one there is nothing to bind. A same-origin relative + ``return_to`` (honored by the SSO callback) brings the browser straight back to this authorize + request after login instead of stranding it on the dashboard.""" base_url = get_request_base_url(request) - return RedirectResponse(f"{base_url}/sso/key/generate") + return RedirectResponse(f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}") # LIT-4197: some upstream authorization servers reject an over-long ``state`` @@ -1253,6 +1272,18 @@ async def authorize( global_mcp_server_manager, ) + if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id) and is_mcp_gateway_dcr_enabled(): + return aggregate_authorize( + request=request, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + response_type=response_type, + session_user_id=_session_cookie_user_id(request), + ) + lookup_name: Optional[str] = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( @@ -1316,6 +1347,25 @@ async def token_endpoint( global_mcp_server_manager, ) + if mcp_server_name is None and is_gateway_dcr_client_id(client_id) and is_mcp_gateway_dcr_enabled(): + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load + master_key, + user_api_key_cache, + ) + + return await aggregate_token( + request=request, + grant_type=grant_type, + code=code, + redirect_uri=redirect_uri, + client_id=client_id, + code_verifier=code_verifier, + refresh_token=refresh_token, + master_key=master_key, + reload_user=_reload_active_user_by_id, + cache=user_api_key_cache, + ) + lookup_name = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) @@ -1337,6 +1387,21 @@ async def token_endpoint( ) +@router.post("/authorize/complete") +async def authorize_complete(request: Request, flow: str = Form(...)): + """Finish an aggregate connect flow (``mcp_gateway_dcr``): mint the gateway + authorization code for the signed-in user and redirect back to the DCR client. POST + plus the per-flow HttpOnly cookie set at /authorize; 404 when the flag is off so the + route is byte-invisible to existing deployments.""" + if not is_mcp_gateway_dcr_enabled(): + raise HTTPException(status_code=404, detail="Not Found") + return complete_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + ) + + # Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request # redirects back to the configured redirect URI with ``error`` / # ``error_description`` / ``error_uri`` query params and no ``code``. The MCP @@ -2061,6 +2126,8 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non } client_ip = IPAddressUtils.get_mcp_client_ip(request) if not mcp_server_name: + if is_mcp_gateway_dcr_enabled(): + return await register_aggregate_client(request=request, request_body=data) resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py new file mode 100644 index 000000000000..da1c30a547d4 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -0,0 +1,519 @@ +"""The gateway-level DCR flow for the aggregate ``/mcp`` endpoint (``mcp_gateway_dcr``). + +An OAuth-only DCR client (Claude Desktop, Claude Code, MCP Inspector) pointed at the +aggregate ``/mcp`` endpoint discovers the gateway as its authorization server (PR 1 of +this track) and then walks the flow implemented here: + +1. ``POST /register``: stateless dynamic client registration. The ``client_id`` IS the + registration: the client's redirect URIs are sealed into it with the repo's + authenticated symmetric helper, so nothing is persisted and a forged or tampered + client_id simply fails to open. Clients are always public (``token_endpoint_auth_method + "none"``); PKCE S256 is what protects the code. +2. ``GET /authorize``: validates the client and redirect URI, requires S256 PKCE, and + interposes LiteLLM sign-in. Without a session cookie the browser is sent through + ``/sso/key/generate`` with a same-origin ``return_to`` so it lands back here after + login. With a session, the flow parameters and the SSO user are sealed into a per-flow + HttpOnly cookie (the same pattern as the upstream OAuth state relay) and the browser is + sent to the connect page, where the user authorizes individual servers (vaulting those + tokens server-side) before finishing. +3. ``POST /authorize/complete``: the deliberate finish step. A POST (not GET) bound to the + SameSite=Lax flow cookie, so a cross-site link cannot silently mint a code with the + victim's session, and the signed-in user must match the user sealed into the flow. + Mints a short-lived, single-use, gateway-sealed authorization code and redirects to the + client's registered redirect URI. +4. ``POST /token``: exchanges the code (PKCE-verified, client- and redirect-bound, + single-use) for the identity-only session tokens of + :mod:`.outbound_credentials.session_token`, re-validating that the litellm user is + still active first; the ``refresh_token`` grant rotates the pair the same way. + +Nothing here stores state server-side except the single-use code guard (a TTL cache +entry). Every sealed value is authenticated encryption over the proxy salt/master key +family, opened totally (bad input maps to an OAuth error, never a raise), and every +identity is a stable reference re-validated live at mint, refresh, and (in the admission +PR) tool-call time. Upstream server credentials never appear anywhere in this flow; they +are vaulted per user by the existing ``/v1/mcp`` authorize endpoints and resolved at +egress by user id. +""" + +from __future__ import annotations + +import hashlib +import hmac +import secrets +from base64 import urlsafe_b64encode +from datetime import datetime, timezone +from typing import Awaitable, Callable, Literal, TypeVar +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse + +from fastapi import Request +from fastapi.responses import JSONResponse, RedirectResponse, Response +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from litellm._logging import verbose_logger +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + TOKEN_NO_CACHE_HEADERS, + get_request_base_url, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + SessionRefreshOpened, + open_session_refresh_bearer, + session_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + MintedSessionToken, + SessionKeys, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) + +GATEWAY_DCR_CLIENT_ID_PREFIX = "llm_dcrc_" +"""Marker prefix on every gateway-issued DCR client_id so the root authorize/token +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_AUTH_CODE_PREFIX = "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.""" + +CONNECT_FLOW_COOKIE_PREFIX = "mcp_connect_flow_" +"""Per-flow HttpOnly cookie holding the sealed connect flow, keyed by a short random +handle carried in the connect-page URL (the same handle-plus-cookie pattern as the +``mcp_oauth_state_`` upstream relay, for the same reasons: replica-safe with no +server-side session store, and the sealed value never appears in a URL).""" + +CONNECT_FLOW_TTL_SECONDS = 600 +GATEWAY_AUTH_CODE_TTL_SECONDS = 120 +_USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:" + +MAX_REDIRECT_URIS = 3 +MAX_REDIRECT_URI_LENGTH = 256 +MAX_CLIENT_ID_LENGTH = 2048 +"""Registration bounds. They exist to bound the sealed client_id, which rides inside +every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably +under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP +Inspector register one or two redirect URIs.""" + +_CLIENT_RECORD_DEBUG_KEY = "gateway_dcr_client" +_CONNECT_FLOW_DEBUG_KEY = "gateway_connect_flow" +_AUTH_CODE_DEBUG_KEY = "gateway_authorization_code" + +ReloadUserFailure = Literal["unresolvable", "unavailable", "no_active_key"] +ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] +"""Injected live-user revalidation (the token endpoint's mirror of admission): +``None`` means the user is active; ``unavailable`` is a retryable DB outage; anything +else fails the grant closed.""" + + +class GatewayDcrClient(BaseModel): + """The registration record sealed into a gateway DCR ``client_id``.""" + + model_config = ConfigDict(frozen=True) + redirect_uris: tuple[str, ...] = Field(min_length=1, max_length=MAX_REDIRECT_URIS) + iat: int + + +class _ConnectFlow(BaseModel): + """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.""" + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + redirect_uri: str = Field(min_length=1) + state: str + code_challenge: str = Field(min_length=1) + exp: int + + +class _GatewayAuthCode(BaseModel): + """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.""" + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + redirect_uri: str = Field(min_length=1) + code_challenge: str = Field(min_length=1) + jti: str = Field(min_length=1) + iat: int + exp: int + + +def is_gateway_dcr_client_id(client_id: str | None) -> bool: + """Cheap prefix routing test so the root endpoints only enter the aggregate arm for + clients this flow registered; every other client_id keeps today's behavior.""" + return bool(client_id) and str(client_id).startswith(GATEWAY_DCR_CLIENT_ID_PREFIX) + + +def _oauth_error(status_code: int, error: str, description: str) -> JSONResponse: + """RFC 6749 section 5.2 / RFC 7591 section 3.2.2 error body. Descriptions carry no + token, code, or URL material so they are safe to relay to any client.""" + return JSONResponse( + status_code=status_code, + content={"error": error, "error_description": description}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + +def _seal(prefix: str, payload: BaseModel) -> str: + return prefix + encrypt_value_helper(payload.model_dump_json()) + + +_SealedModelT = TypeVar("_SealedModelT", bound=BaseModel) + + +def _open_sealed(value: str, prefix: str, model: type[_SealedModelT], debug_key: str) -> _SealedModelT | None: + """Open a sealed value totally: anything that is not prefix-shaped, does not decrypt, + or does not validate returns ``None`` for the caller to map onto an OAuth error.""" + if not value.startswith(prefix): + return None + decrypted = decrypt_value_helper(value[len(prefix) :], debug_key, return_original_value=False) + if not isinstance(decrypted, str): + return None + try: + return model.model_validate_json(decrypted) + except ValidationError: + return None + + +def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None: + return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY) + + +def _redirect_uri_acceptable(uri: str) -> bool: + """https for real clients, plus http strictly on a loopback host for local dev + clients (RFC 8252 section 7.3). No fragments (RFC 6749 section 3.1.2).""" + if len(uri) > MAX_REDIRECT_URI_LENGTH: + return False + parsed = urlparse(uri) + if parsed.fragment or not parsed.netloc: + return False + if parsed.scheme == "https": + return True + return parsed.scheme == "http" and (parsed.hostname or "").lower() in ("localhost", "127.0.0.1", "::1") + + +async def register_aggregate_client(request: Request, request_body: dict) -> Response: + """RFC 7591 dynamic registration against the gateway itself, statelessly. + + Only ``redirect_uris`` is authoritative; every client is registered as a public + ``token_endpoint_auth_method "none"`` client regardless of what it asked for (RFC + 7591 lets the server override metadata), because the gateway never issues client + secrets: possession of a secret would add nothing over the mandatory S256 PKCE, and a + stateless registration has nowhere to keep one. Nothing is persisted, so open + registration cannot be used to fill storage. + """ + raw_uris = request_body.get("redirect_uris") + if not isinstance(raw_uris, list) or not raw_uris or len(raw_uris) > MAX_REDIRECT_URIS: + return _oauth_error( + 400, + "invalid_redirect_uri", + f"redirect_uris must be a list of 1 to {MAX_REDIRECT_URIS} URIs", + ) + if not all(isinstance(uri, str) and _redirect_uri_acceptable(uri) for uri in raw_uris): + return _oauth_error( + 400, + "invalid_redirect_uri", + "each redirect URI must be https (or http on a loopback host), " + f"fragment-free, and at most {MAX_REDIRECT_URI_LENGTH} characters", + ) + now = datetime.now(timezone.utc) + client_id = _seal( + GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient(redirect_uris=tuple(raw_uris), iat=int(now.timestamp())) + ) + if len(client_id) > MAX_CLIENT_ID_LENGTH: + return _oauth_error(400, "invalid_client_metadata", "registered metadata is too large") + return JSONResponse( + status_code=201, + content={ + "client_id": client_id, + "client_id_issued_at": int(now.timestamp()), + "redirect_uris": list(raw_uris), + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + }, + ) + + +def _flow_cookie_name(handle: str) -> str: + return f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" + + +def _cookie_path_and_secure(request: Request) -> tuple[str, bool]: + parsed = urlparse(get_request_base_url(request)) + return parsed.path or "/", parsed.scheme == "https" + + +def _append_query_params(url: str, params: dict[str, str]) -> str: + parsed = urlparse(url) + query = parse_qsl(parsed.query, keep_blank_values=True) + list(params.items()) + return urlunparse(parsed._replace(query=urlencode(query))) + + +def relative_request_url(request: Request) -> str: + """The request's own path and query as a same-origin ``return_to`` target for the + login round-trip; relative by construction, so it can never leave the gateway.""" + path = request.url.path + return f"{path}?{request.url.query}" if request.url.query else path + + +def aggregate_authorize( + request: Request, + client_id: str, + redirect_uri: str, + state: str, + code_challenge: str | None, + code_challenge_method: str | None, + response_type: str | None, + session_user_id: str | None, +) -> Response: + """The aggregate authorize verb: validate the client, require S256 PKCE, interpose + LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a + per-flow cookie. + + 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. + """ + client = open_gateway_dcr_client(client_id) + if client is None: + return _oauth_error(400, "invalid_client", "unknown or malformed client_id") + if redirect_uri not in client.redirect_uris: + return _oauth_error(400, "invalid_request", "redirect_uri is not registered for this client") + if response_type != "code": + return _oauth_error(400, "unsupported_response_type", "response_type must be 'code'") + if not code_challenge or code_challenge_method != "S256": + return _oauth_error( + 400, + "invalid_request", + "PKCE is required: send code_challenge with code_challenge_method=S256", + ) + base_url = get_request_base_url(request) + if session_user_id is None: + login_url = f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}" + return RedirectResponse(login_url, status_code=303) + now = datetime.now(timezone.utc) + handle = secrets.token_urlsafe(24) + flow = _ConnectFlow( + user_id=session_user_id, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, + ) + connect_url = _append_query_params( + f"{base_url}/ui/chat/integrations", + {"connect_flow": handle, "connect_client": _origin_only(redirect_uri)}, + ) + response = RedirectResponse(connect_url, status_code=303) + path, secure = _cookie_path_and_secure(request) + response.set_cookie( + key=_flow_cookie_name(handle), + value=_seal("", flow), + max_age=CONNECT_FLOW_TTL_SECONDS, + path=path, + secure=secure, + httponly=True, + samesite="lax", + ) + return response + + +def _origin_only(url: str) -> str: + """Scheme+host for display on the connect page; never the full redirect URI, whose + path or query could carry values that do not belong in a page URL or logs.""" + parsed = urlparse(url) + return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" + + +def complete_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, +) -> Response: + """The deliberate finish step of the connect flow: mint the gateway authorization + code and send the browser back to the client. + + Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly + per-flow cookie plus an exact match between the signed-in user and the user sealed + into the flow: a link crafted by another party dies here with ``access_denied`` + instead of minting a code for the victim's identity. + """ + sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow = _open_sealed(sealed_flow, "", _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + now = datetime.now(timezone.utc) + if now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id=flow.user_id, + client_id=flow.client_id, + redirect_uri=flow.redirect_uri, + code_challenge=flow.code_challenge, + jti=secrets.token_urlsafe(24), + iat=int(now.timestamp()), + exp=int(now.timestamp()) + GATEWAY_AUTH_CODE_TTL_SECONDS, + ), + ) + params = {"code": code, **({"state": flow.state} if flow.state else {})} + response = RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303) + path, secure = _cookie_path_and_secure(request) + response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") + return response + + +def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool: + digest = hashlib.sha256(code_verifier.encode("ascii", "replace")).digest() + computed = urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return hmac.compare_digest(computed, code_challenge) + + +class _SingleUseGuard: + """Best-effort single-use marking for gateway authorization codes over the injected + proxy cache (in-memory always, Redis when the deployment wires it, in which case the + guard holds across replicas). The code's 120s TTL is the hard bound either way; the + guard exists so a same-process or shared-cache replay fails ``invalid_grant``.""" + + def __init__(self, cache: DualCache) -> None: + self._cache = cache + + async def already_used(self, jti: str) -> bool: + return await self._cache.async_get_cache(f"{_USED_CODE_CACHE_PREFIX}{jti}") is not None + + async def mark_used(self, jti: str) -> None: + await self._cache.async_set_cache( + f"{_USED_CODE_CACHE_PREFIX}{jti}", "1", ttl=GATEWAY_AUTH_CODE_TTL_SECONDS + 60 + ) + + +def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response: + access = mint_session_token(principal, keys, now) + refresh = mint_session_refresh_token(principal, keys, now) + if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken): + return _oauth_error(500, "server_error", "failed to mint the session credential") + return JSONResponse( + status_code=200, + content={ + "access_token": access.token.get_secret_value(), + "token_type": "Bearer", + "expires_in": int((access.expires_at - now).total_seconds()), + "refresh_token": refresh.token.get_secret_value(), + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + +def _reload_failure_response(failure: ReloadUserFailure) -> Response: + if failure == "unavailable": + return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + if failure == "unresolvable": + return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") + return _oauth_error(400, "invalid_grant", "the user for this grant is no longer active") + + +async def aggregate_token( + request: Request, + grant_type: str, + code: str | None, + redirect_uri: str | None, + client_id: str, + code_verifier: str | None, + refresh_token: str | None, + master_key: str | None, + reload_user: ReloadUser, + cache: DualCache, +) -> Response: + """The aggregate token verb: authorization_code and refresh_token grants for the + identity-only session pair. Every path re-validates the litellm user live before + minting, so a deactivated user cannot obtain or renew a session.""" + if master_key is None: + verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") + return _oauth_error(500, "server_error", "the gateway has no master key configured") + keys = session_keys_from_master_key(master_key) + now = datetime.now(timezone.utc) + if grant_type == "authorization_code": + return await _authorization_code_grant( + code=code, + redirect_uri=redirect_uri, + client_id=client_id, + code_verifier=code_verifier, + keys=keys, + now=now, + reload_user=reload_user, + guard=_SingleUseGuard(cache), + ) + if grant_type == "refresh_token": + return await _refresh_token_grant( + refresh_token=refresh_token, + client_id=client_id, + keys=keys, + now=now, + reload_user=reload_user, + ) + return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token") + + +async def _authorization_code_grant( + code: str | None, + redirect_uri: str | None, + client_id: str, + code_verifier: str | None, + keys: SessionKeys, + now: datetime, + reload_user: ReloadUser, + guard: _SingleUseGuard, +) -> Response: + if not code or not redirect_uri or not code_verifier: + return _oauth_error(400, "invalid_request", "code, redirect_uri, and code_verifier are required") + parsed = _open_sealed(code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY) + if parsed is None: + return _oauth_error(400, "invalid_grant", "the authorization code is invalid") + if now.timestamp() >= parsed.exp: + 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 not _pkce_verifier_matches(code_verifier, parsed.code_challenge): + return _oauth_error(400, "invalid_grant", "PKCE verification failed") + if await guard.already_used(parsed.jti): + return _oauth_error(400, "invalid_grant", "the authorization code was already used") + await guard.mark_used(parsed.jti) + failure = await reload_user(parsed.user_id) + if failure is not None: + return _reload_failure_response(failure) + return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now) + + +async def _refresh_token_grant( + refresh_token: str | None, + client_id: str, + keys: SessionKeys, + now: datetime, + reload_user: ReloadUser, +) -> Response: + if not refresh_token: + return _oauth_error(400, "invalid_request", "refresh_token is required") + opened = open_session_refresh_bearer(refresh_token, keys, now, expected_client_id=client_id) + if not isinstance(opened, SessionRefreshOpened): + return _oauth_error(400, "invalid_grant", "the refresh token is invalid for this client") + failure = await reload_user(opened.principal.user_id) + if failure is not None: + return _reload_failure_response(failure) + return _session_token_pair(opened.principal, keys, now) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index fe6682e4221b..41e70d5f1eb3 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -950,7 +950,7 @@ async def google_login( request=request, ) if return_to is not None and sso_redirect is not None: - if SSOAuthenticationHandler._validate_return_to(return_to): + if _is_same_origin_return_path(return_to) or SSOAuthenticationHandler._validate_return_to(return_to): sso_redirect.set_cookie( key="litellm_cp_return_to", value=return_to, @@ -2419,6 +2419,15 @@ async def sso_readiness(): ) +def _is_same_origin_return_path(return_to: str) -> bool: + """True for a strictly relative return path (starts with ``/``, not + protocol-relative ``//``, no backslash tricks browsers normalize to slashes), which + stays on the gateway's own origin by construction and is therefore safe to honor + without a configured ``control_plane_url``. Used by the MCP gateway DCR authorize + round-trip so a browser sent through login lands back on the authorize request.""" + return return_to.startswith("/") and not return_to.startswith("//") and "\\" not in return_to + + class SSOAuthenticationHandler: """ Handler for SSO Authentication across all SSO providers @@ -3055,7 +3064,6 @@ async def get_redirect_response_from_openid( jwt_handler: Optional[JWTHandler] = None, return_to: Optional[str] = None, ) -> RedirectResponse: - import jwt from litellm.proxy.proxy_server import ( general_settings, @@ -3220,6 +3228,15 @@ async def get_redirect_response_from_openid( jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "") + # Same-origin relative return (the MCP gateway DCR authorize round-trip): + # set the session cookie exactly like the dashboard path, then send the + # browser back to where it came from instead of the dashboard. + if return_to is not None and _is_same_origin_return_path(return_to): + redirect_response = RedirectResponse(url=return_to, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token) + redirect_response.delete_cookie("litellm_cp_return_to") + return redirect_response + # Control-plane cross-origin: store JWT behind a single-use opaque # code (60s TTL) so the token never appears in browser history / logs. # The control plane redeems it via POST /v3/login/exchange. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 7e9ff4692b51..1dd342b7b28c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7258,3 +7258,64 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] finally: global_mcp_server_manager.registry.clear() + + +def test_gateway_dcr_flow_routing_engages_only_for_llm_dcrc_clients(monkeypatch): + """The aggregate DCR arms engage for llm_dcrc_ client_ids (register always mints one, + authorize/token route into the aggregate flow); a non-gateway client_id keeps the + per-server behavior, and /authorize/complete exists but 400s without a valid flow.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit3637") + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637", raising=False) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + registered = client.post("/register", json={"redirect_uris": ["https://claude.ai/cb"]}) + assert registered.status_code == 201 + assert registered.json()["client_id"].startswith("llm_dcrc_") + assert registered.json()["token_endpoint_auth_method"] == "none" + + authorize_params = { + "client_id": "llm_dcrc_bogus", + "redirect_uri": "https://claude.ai/cb", + "response_type": "code", + "code_challenge": "c" * 43, + "code_challenge_method": "S256", + } + bogus_client = client.get("/authorize", params=authorize_params) + assert bogus_client.status_code == 400 + assert bogus_client.json()["error"] == "invalid_client" + + no_cookie = client.post("/authorize/complete", data={"flow": "h"}) + assert no_cookie.status_code == 400 + assert no_cookie.json()["error"] == "invalid_request" + + token_response = client.post( + "/token", + data={ + "grant_type": "authorization_code", + "client_id": "llm_dcrc_bogus", + "code": "x", + "redirect_uri": "https://claude.ai/cb", + "code_verifier": "v" * 43, + }, + ) + assert token_response.status_code == 400 + assert token_response.json()["error"] == "invalid_grant" + + # a non-gateway (upstream-issued) client_id is not routed into the aggregate arm; it + # falls to the per-server exchange, which 404s for an unknown server + upstream_shaped = client.post( + "/token", + data={"grant_type": "authorization_code", "client_id": "regular-upstream-client", "code": "x"}, + ) + assert upstream_shaped.status_code == 404 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py new file mode 100644 index 000000000000..6db337c6f4f6 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -0,0 +1,379 @@ +"""Tests for the aggregate gateway DCR flow (register, authorize, complete, token).""" + +import hashlib +import json +from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone +from http.cookies import SimpleCookie +from urllib.parse import parse_qs, urlparse + +import pytest +from starlette.requests import Request + +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + CONNECT_FLOW_COOKIE_PREFIX, + GATEWAY_AUTH_CODE_PREFIX, + GATEWAY_AUTH_CODE_TTL_SECONDS, + GATEWAY_DCR_CLIENT_ID_PREFIX, + _GatewayAuthCode, + _seal, + aggregate_authorize, + aggregate_token, + complete_connect_flow, + is_gateway_dcr_client_id, + open_gateway_dcr_client, + register_aggregate_client, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + resolve_session_bearer, + session_keys_from_master_key, + SessionBearerAdmitted, +) + +MASTER_KEY = "sk-gateway-dcr-flow-tests" +REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" +CODE_VERIFIER = "verifier-" + "v" * 43 +CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii") + + +@pytest.fixture(autouse=True) +def _salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", MASTER_KEY) + + +def _request(path="/authorize", query="", cookies=None, method="GET"): + cookie_header = [] + if cookies: + cookie = SimpleCookie() + for name, value in cookies.items(): + cookie[name] = value + cookie_header = [(b"cookie", cookie.output(header="", sep="; ").strip().encode())] + return Request( + { + "type": "http", + "method": method, + "scheme": "https", + "path": path, + "query_string": query.encode(), + "headers": [(b"host", b"llm.example.com"), *cookie_header], + } + ) + + +async def _register(redirect_uris) -> dict: + response = await register_aggregate_client(request=_request("/register"), request_body={"redirect_uris": redirect_uris}) + return json.loads(response.body) + + +async def _reload_user_active(user_id: str): + return None + + +@pytest.mark.asyncio +async def test_register_mints_stateless_public_client(): + body = await _register([REDIRECT_URI]) + assert body["token_endpoint_auth_method"] == "none" + assert "client_secret" not in body + assert body["redirect_uris"] == [REDIRECT_URI] + assert is_gateway_dcr_client_id(body["client_id"]) + record = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == (REDIRECT_URI,) + + +@pytest.mark.asyncio +async def test_register_allows_loopback_http_for_dev_clients(): + body = await _register(["http://localhost:6274/oauth/callback"]) + assert is_gateway_dcr_client_id(body["client_id"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "redirect_uris", + [ + [], + "not-a-list", + ["http://evil.example.com/callback"], + ["https://claude.ai/cb#fragment"], + ["ftp://claude.ai/cb"], + ["https://a.example.com/" + "p" * 300], + ["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"], + [12345], + ], +) +async def test_register_rejects_bad_redirect_uris(redirect_uris): + response = await register_aggregate_client(request=_request("/register"), request_body={"redirect_uris": redirect_uris}) + assert response.status_code == 400 + assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") + + +@pytest.mark.asyncio +async def test_tampered_client_id_does_not_open(): + body = await _register([REDIRECT_URI]) + tampered = body["client_id"][:-4] + "AAAA" + assert open_gateway_dcr_client(tampered) is None + assert open_gateway_dcr_client("llm_dcrc_garbage") is None + assert open_gateway_dcr_client("other_prefix") is None + + +def _authorize(client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code"): + return aggregate_authorize( + request=_request(query=f"client_id={client_id}"), + client_id=client_id, + redirect_uri=redirect_uri, + state="client-state-123", + code_challenge=challenge, + code_challenge_method=method, + response_type=response_type, + session_user_id=session_user_id, + ) + + +@pytest.mark.asyncio +async def test_authorize_validation_failures_never_redirect_to_client(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + for response, expected_error in ( + (_authorize("llm_dcrc_bogus", "u1"), "invalid_client"), + (_authorize(client_id, "u1", redirect_uri="https://attacker.example.com/cb"), "invalid_request"), + (_authorize(client_id, "u1", response_type="token"), "unsupported_response_type"), + (_authorize(client_id, "u1", challenge=None), "invalid_request"), + (_authorize(client_id, "u1", method="plain"), "invalid_request"), + ): + assert response.status_code == 400 + assert json.loads(response.body)["error"] == expected_error + + +@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) + 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 + + +@pytest.mark.asyncio +async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_cookie(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + assert response.status_code == 303 + location = urlparse(response.headers["location"]) + assert location.path == "/ui/chat/integrations" + params = parse_qs(location.query) + handle = params["connect_flow"][0] + assert params["connect_client"] == ["https://claude.ai"] + set_cookie = response.headers["set-cookie"] + assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie + assert "HttpOnly" in set_cookie + return handle, set_cookie + + +def _flow_cookie_from(response) -> tuple: + location = urlparse(response.headers["location"]) + handle = parse_qs(location.query)["connect_flow"][0] + cookie = SimpleCookie() + cookie.load(response.headers["set-cookie"]) + name = f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" + return handle, {name: cookie[name].value} + + +@pytest.mark.asyncio +async def test_full_walk_register_authorize_complete_token_and_replay(): + """The whole front door on one deterministic walk: register -> authorize -> + complete -> token, then the security edges on the same artifacts (user mismatch, + PKCE mismatch, single-use replay, refresh rotation, cross-client refresh).""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + authorize_response = _authorize(client_id, session_user_id="u1") + handle, cookies = _flow_cookie_from(authorize_response) + + denied = complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="attacker", + ) + assert denied.status_code == 403 + + anonymous = complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id=None, + ) + assert anonymous.status_code == 401 + + completed = complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + ) + assert completed.status_code == 303 + redirect = urlparse(completed.headers["location"]) + assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI + params = parse_qs(redirect.query) + assert params["state"] == ["client-state-123"] + code = params["code"][0] + assert code.startswith(GATEWAY_AUTH_CODE_PREFIX) + + cache = DualCache() + + async def _token(**overrides): + arguments = { + "request": _request("/token", method="POST"), + "grant_type": "authorization_code", + "code": code, + "redirect_uri": REDIRECT_URI, + "client_id": client_id, + "code_verifier": CODE_VERIFIER, + "refresh_token": None, + "master_key": MASTER_KEY, + "reload_user": _reload_user_active, + "cache": cache, + } + return await aggregate_token(**{**arguments, **overrides}) + + wrong_verifier = await _token(code_verifier="wrong-" + "w" * 43) + assert json.loads(wrong_verifier.body)["error"] == "invalid_grant" + + wrong_client = await _token(client_id=(await _register([REDIRECT_URI]))["client_id"]) + assert json.loads(wrong_client.body)["error"] == "invalid_grant" + + token_response = await _token() + assert token_response.status_code == 200 + payload = json.loads(token_response.body) + assert payload["token_type"] == "Bearer" + assert 0 < payload["expires_in"] <= 3600 + + keys = session_keys_from_master_key(MASTER_KEY) + admitted = resolve_session_bearer(f"Bearer {payload['access_token']}", keys, datetime.now(timezone.utc)) + assert isinstance(admitted, SessionBearerAdmitted) + assert admitted.principal.user_id == "u1" + assert admitted.principal.client_id == client_id + + replay = await _token() + assert json.loads(replay.body)["error"] == "invalid_grant" + + refreshed = await _token(grant_type="refresh_token", code=None, refresh_token=payload["refresh_token"]) + assert refreshed.status_code == 200 + rotated = json.loads(refreshed.body) + assert rotated["refresh_token"] != payload["refresh_token"] + + cross_client = await _token( + grant_type="refresh_token", + code=None, + refresh_token=payload["refresh_token"], + client_id=(await _register([REDIRECT_URI]))["client_id"], + ) + assert json.loads(cross_client.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_complete_rejects_missing_tampered_and_expired_flows(): + missing = complete_connect_flow( + request=_request("/authorize/complete", method="POST"), flow_handle="nope", session_user_id="u1" + ) + assert missing.status_code == 400 + + tampered = complete_connect_flow( + request=_request("/authorize/complete", cookies={f"{CONNECT_FLOW_COOKIE_PREFIX}h1": "garbage"}, method="POST"), + flow_handle="h1", + session_user_id="u1", + ) + assert tampered.status_code == 400 + + +@pytest.mark.asyncio +async def test_token_rejects_expired_code_and_missing_configuration(): + expired_code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id="u1", + client_id="llm_dcrc_x", + redirect_uri=REDIRECT_URI, + code_challenge=CODE_CHALLENGE, + jti="jti-1", + iat=int((datetime.now(timezone.utc) - timedelta(seconds=500)).timestamp()), + exp=int((datetime.now(timezone.utc) - timedelta(seconds=500 - GATEWAY_AUTH_CODE_TTL_SECONDS)).timestamp()), + ), + ) + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=expired_code, + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert json.loads(response.body)["error"] == "invalid_grant" + + no_master_key = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code="llm_gcode_x", + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=None, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert no_master_key.status_code == 500 + assert json.loads(no_master_key.body)["error"] == "server_error" + + unsupported = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="password", + code=None, + redirect_uri=None, + client_id="llm_dcrc_x", + code_verifier=None, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert json.loads(unsupported.body)["error"] == "unsupported_grant_type" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure,expected_status,expected_error", + [ + ("no_active_key", 400, "invalid_grant"), + ("unavailable", 503, "temporarily_unavailable"), + ("unresolvable", 500, "server_error"), + ], +) +async def test_token_gates_on_live_user_revalidation(failure, expected_status, expected_error): + client_id = (await _register([REDIRECT_URI]))["client_id"] + authorize_response = _authorize(client_id, session_user_id="deactivated-user") + handle, cookies = _flow_cookie_from(authorize_response) + completed = complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="deactivated-user", + ) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + + async def _reload_user_failing(user_id: str): + return failure + + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=code, + redirect_uri=REDIRECT_URI, + client_id=client_id, + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_failing, + cache=DualCache(), + ) + assert response.status_code == expected_status + assert json.loads(response.body)["error"] == expected_error diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 92d1b870d757..2d7067726dff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7391,3 +7391,24 @@ async def test_auth_callback_without_oauth_error_proceeds_to_normal_flow(): assert exc_info.value.status_code == 500 assert "DB not connected" in str(exc_info.value.detail) + + +class TestSameOriginReturnPath: + """The same-origin relative return_to arm added for the MCP gateway DCR authorize + round-trip: only strictly relative paths qualify, so login can never redirect the + browser off the gateway origin.""" + + def test_accepts_relative_paths(self): + from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path + + assert _is_same_origin_return_path("/authorize?client_id=llm_dcrc_x&state=s") is True + assert _is_same_origin_return_path("/some_server/authorize") is True + + def test_rejects_absolute_protocol_relative_and_backslash_paths(self): + from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path + + assert _is_same_origin_return_path("https://evil.example.com/authorize") is False + assert _is_same_origin_return_path("//evil.example.com/authorize") is False + assert _is_same_origin_return_path("/\\evil.example.com") is False + assert _is_same_origin_return_path("javascript:alert(1)") is False + assert _is_same_origin_return_path("") is False From 05c55d016bc236f4b27d812bd1271512839e9d1b Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 09:52:35 -0700 Subject: [PATCH 09/16] refactor(mcp): harden gateway DCR flow per adversarial review - atomic single-use guard (async_increment_cache) + reload-before-claim so a transient DB blip does not burn a valid code - PKCE verify over bytes so a non-ASCII code_challenge fails invalid_grant instead of raising a 500; validate code_verifier length (RFC 7636) - flag-off byte-identical for a server literally named mcp (AS well-known delegates to the named-server document) - connect flow is single-use (atomic jti claim) so a double-submit cannot mint two codes - extra=forbid on the sealed models; bound state length; drop unused request param and coarse dict on register - _reload_failure_response exhaustive match+assert_never; dedupe ReloadUserFailure with _KeyResolutionFailure - reject control/whitespace chars in the same-origin return_to --- .../mcp_server/discoverable_endpoints.py | 28 ++-- .../mcp_server/gateway_dcr_flow.py | 127 ++++++++++++----- litellm/proxy/management_endpoints/ui_sso.py | 19 ++- .../mcp_server/test_discoverable_endpoints.py | 80 ++++++----- .../mcp_server/test_gateway_dcr_flow.py | 130 ++++++++++++++++-- 5 files changed, 287 insertions(+), 97 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ab9339593c47..6ebbb66ffedb 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -29,6 +29,7 @@ _finish_bridge_mint, _prepare_bridge_mint, _prepare_bridge_refresh, + _reload_active_user_by_id, ) from litellm.proxy._experimental.mcp_server.faults import ( CallerRejected, @@ -1272,7 +1273,7 @@ async def authorize( global_mcp_server_manager, ) - if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id) and is_mcp_gateway_dcr_enabled(): + if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id): return aggregate_authorize( request=request, client_id=client_id, @@ -1347,7 +1348,7 @@ async def token_endpoint( global_mcp_server_manager, ) - if mcp_server_name is None and is_gateway_dcr_client_id(client_id) and is_mcp_gateway_dcr_enabled(): + if mcp_server_name is None and is_gateway_dcr_client_id(client_id): from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load master_key, user_api_key_cache, @@ -1389,16 +1390,16 @@ async def token_endpoint( @router.post("/authorize/complete") async def authorize_complete(request: Request, flow: str = Form(...)): - """Finish an aggregate connect flow (``mcp_gateway_dcr``): mint the gateway - authorization code for the signed-in user and redirect back to the DCR client. POST - plus the per-flow HttpOnly cookie set at /authorize; 404 when the flag is off so the - route is byte-invisible to existing deployments.""" - if not is_mcp_gateway_dcr_enabled(): - raise HTTPException(status_code=404, detail="Not Found") - return complete_connect_flow( + """Finish an aggregate connect flow: mint the gateway authorization code for the + signed-in user and redirect back to the DCR client. POST plus the per-flow HttpOnly + cookie set at /authorize; an anonymous or bad-flow request just 400s.""" + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # circular import at module load + + return await complete_connect_flow( request=request, flow_handle=flow, session_user_id=_session_cookie_user_id(request), + cache=user_api_key_cache, ) @@ -2126,8 +2127,13 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non } client_ip = IPAddressUtils.get_mcp_client_ip(request) if not mcp_server_name: - if is_mcp_gateway_dcr_enabled(): - return await register_aggregate_client(request=request, request_body=data) + # A real DCR request carries redirect_uris (RFC 7591): route it to the aggregate DCR + # endpoint the aggregate authorization-server metadata advertises. A single-server + # deployment registers at /{server}/register instead (its bare-origin discovery + # advertises that), so this does not affect it. A request without redirect_uris is not + # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. + if data.get("redirect_uris"): + return await register_aggregate_client(request_body=data) resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index da1c30a547d4..017826ac9b54 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -41,6 +41,7 @@ import hmac import secrets from base64 import urlsafe_b64encode +from collections.abc import Mapping from datetime import datetime, timezone from typing import Awaitable, Callable, Literal, TypeVar from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -48,6 +49,7 @@ from fastapi import Request from fastapi.responses import JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, ValidationError +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -89,7 +91,9 @@ CONNECT_FLOW_TTL_SECONDS = 600 GATEWAY_AUTH_CODE_TTL_SECONDS = 120 +_CLAIM_TTL_BUFFER_SECONDS = 60 _USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:" +_USED_FLOW_CACHE_PREFIX = "mcp_gateway_dcr_flow_used:" MAX_REDIRECT_URIS = 3 MAX_REDIRECT_URI_LENGTH = 256 @@ -99,6 +103,22 @@ under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP Inspector register one or two redirect URIs.""" +MAX_STATE_LENGTH = 1024 +"""Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code +redirect. An unbounded ``state`` can push the sealed cookie past the browser's ~4KB cap +(silently dropped, breaking the flow); spec clients send a short opaque value.""" + +MIN_CODE_VERIFIER_LENGTH = 43 +MAX_CODE_VERIFIER_LENGTH = 128 +"""RFC 7636 section 4.1 bounds for the PKCE ``code_verifier``. Enforced so an out-of-range +verifier gets a clean ``invalid_request`` instead of an opaque PKCE-mismatch.""" + +_UNPREFIXED = "" +"""Prefix for a sealed value that carries no wire marker because it is never routed by +prefix (the connect flow lives only in its own per-handle cookie, opened by that one +handle). Named so the empty-string argument to ``_seal`` / ``_open_sealed`` reads as +deliberate rather than a typo.""" + _CLIENT_RECORD_DEBUG_KEY = "gateway_dcr_client" _CONNECT_FLOW_DEBUG_KEY = "gateway_connect_flow" _AUTH_CODE_DEBUG_KEY = "gateway_authorization_code" @@ -111,32 +131,41 @@ class GatewayDcrClient(BaseModel): - """The registration record sealed into a gateway DCR ``client_id``.""" + """The registration record sealed into a gateway DCR ``client_id``. - model_config = ConfigDict(frozen=True) + ``extra="forbid"`` so a sealed value of another type (an auth code, a connect flow) + that happened to decrypt under the shared key can never validate as a client record: + cross-type confusion is rejected at the model boundary, not left to differing required + fields.""" + + model_config = ConfigDict(frozen=True, extra="forbid") redirect_uris: tuple[str, ...] = Field(min_length=1, max_length=MAX_REDIRECT_URIS) iat: int class _ConnectFlow(BaseModel): """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.""" + 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) + 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) state: str code_challenge: str = Field(min_length=1) + jti: str = Field(min_length=1) exp: int class _GatewayAuthCode(BaseModel): """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.""" + plus a ``jti`` for the single-use guard. ``extra="forbid"`` rejects cross-type + confusion.""" - model_config = ConfigDict(frozen=True) + 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) @@ -149,7 +178,7 @@ class _GatewayAuthCode(BaseModel): def is_gateway_dcr_client_id(client_id: str | None) -> bool: """Cheap prefix routing test so the root endpoints only enter the aggregate arm for clients this flow registered; every other client_id keeps today's behavior.""" - return bool(client_id) and str(client_id).startswith(GATEWAY_DCR_CLIENT_ID_PREFIX) + return client_id is not None and client_id.startswith(GATEWAY_DCR_CLIENT_ID_PREFIX) def _oauth_error(status_code: int, error: str, description: str) -> JSONResponse: @@ -200,7 +229,7 @@ def _redirect_uri_acceptable(uri: str) -> bool: return parsed.scheme == "http" and (parsed.hostname or "").lower() in ("localhost", "127.0.0.1", "::1") -async def register_aggregate_client(request: Request, request_body: dict) -> Response: +async def register_aggregate_client(request_body: Mapping[str, object]) -> Response: """RFC 7591 dynamic registration against the gateway itself, statelessly. Only ``redirect_uris`` is authoritative; every client is registered as a public @@ -296,6 +325,8 @@ def aggregate_authorize( "invalid_request", "PKCE is required: send code_challenge with code_challenge_method=S256", ) + if len(state) > MAX_STATE_LENGTH: + return _oauth_error(400, "invalid_request", f"state must be at most {MAX_STATE_LENGTH} characters") base_url = get_request_base_url(request) if session_user_id is None: login_url = f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}" @@ -308,6 +339,7 @@ def aggregate_authorize( redirect_uri=redirect_uri, state=state, code_challenge=code_challenge, + jti=secrets.token_urlsafe(24), exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, ) connect_url = _append_query_params( @@ -318,7 +350,7 @@ def aggregate_authorize( path, secure = _cookie_path_and_secure(request) response.set_cookie( key=_flow_cookie_name(handle), - value=_seal("", flow), + value=_seal(_UNPREFIXED, flow), max_age=CONNECT_FLOW_TTL_SECONDS, path=path, secure=secure, @@ -335,10 +367,11 @@ def _origin_only(url: str) -> str: return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" -def complete_connect_flow( +async def complete_connect_flow( request: Request, flow_handle: str, session_user_id: str | None, + cache: DualCache, ) -> Response: """The deliberate finish step of the connect flow: mint the gateway authorization code and send the browser back to the client. @@ -346,12 +379,13 @@ def complete_connect_flow( Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly per-flow cookie plus an exact match between the signed-in user and the user sealed into the flow: a link crafted by another party dies here with ``access_denied`` - instead of minting a code for the victim's identity. + instead of minting a code for the victim's identity. The flow is single-use (an atomic + claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. """ sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle)) if sealed_flow is None: return _oauth_error(400, "invalid_request", "unknown or expired connect flow") - flow = _open_sealed(sealed_flow, "", _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + flow = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) if flow is None: return _oauth_error(400, "invalid_request", "unknown or expired connect flow") now = datetime.now(timezone.utc) @@ -361,6 +395,10 @@ def complete_connect_flow( return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") if session_user_id != flow.user_id: return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + if not await _SingleUseGuard(cache).claim( + f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + ): + return _oauth_error(400, "invalid_request", "this connect flow was already completed; restart the connection") code = _seal( GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode( @@ -381,27 +419,37 @@ def complete_connect_flow( def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool: + """RFC 7636 S256 verification, total over hostile input. The comparison is over bytes + so a non-ASCII ``code_challenge`` (which reaches here unvalidated from the client's + authorize request) simply fails to match instead of raising ``TypeError`` the way + ``hmac.compare_digest`` does on two ``str`` with non-ASCII content. The verifier is + ASCII per spec; a compliant client's challenge is base64url and matches.""" digest = hashlib.sha256(code_verifier.encode("ascii", "replace")).digest() - computed = urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") - return hmac.compare_digest(computed, code_challenge) + computed = urlsafe_b64encode(digest).rstrip(b"=") + return hmac.compare_digest(computed, code_challenge.encode("utf-8")) class _SingleUseGuard: - """Best-effort single-use marking for gateway authorization codes over the injected - proxy cache (in-memory always, Redis when the deployment wires it, in which case the - guard holds across replicas). The code's 120s TTL is the hard bound either way; the - guard exists so a same-process or shared-cache replay fails ``invalid_grant``.""" + """Atomic single-use claim for a one-time id (an auth-code or connect-flow ``jti``) over + the injected proxy cache. + + Uses an atomic increment rather than a get-then-set: two concurrent redemptions of the + same id cannot both observe "unused", because exactly one increment returns 1. With + Redis wired this holds across replicas (``INCR`` is atomic); single-replica it holds in + the in-memory cache. The id's own TTL is the outer bound. A claim is the gate, not a + marker to check separately, so it fails closed: if the cache cannot record the claim + (no backend at all) the id is refused rather than admitted. For the auth code, PKCE + binding is the primary defense against interception; this makes the RFC 6749 4.1.2 + single-use property reliable on top of it.""" def __init__(self, cache: DualCache) -> None: self._cache = cache - async def already_used(self, jti: str) -> bool: - return await self._cache.async_get_cache(f"{_USED_CODE_CACHE_PREFIX}{jti}") is not None - - async def mark_used(self, jti: str) -> None: - await self._cache.async_set_cache( - f"{_USED_CODE_CACHE_PREFIX}{jti}", "1", ttl=GATEWAY_AUTH_CODE_TTL_SECONDS + 60 - ) + async def claim(self, key: str, ttl_seconds: int) -> bool: + """Atomically claim ``key``. ``True`` iff this caller is the first (increment to 1); + ``False`` on a replay (>1) or when the claim could not be recorded (fail closed).""" + count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds) + return count == 1 def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response: @@ -422,11 +470,17 @@ def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: dat def _reload_failure_response(failure: ReloadUserFailure) -> Response: - if failure == "unavailable": - return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") - if failure == "unresolvable": - return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") - return _oauth_error(400, "invalid_grant", "the user for this grant is no longer active") + """Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new + ``ReloadUserFailure`` member is a type error here rather than silently 400ing.""" + match failure: + case "unavailable": + return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + case "unresolvable": + return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") + case "no_active_key": + return _oauth_error(400, "invalid_grant", "the user for this grant is no longer active") + case _: + assert_never(failure) async def aggregate_token( @@ -483,6 +537,8 @@ async def _authorization_code_grant( ) -> Response: if not code or not redirect_uri or not code_verifier: return _oauth_error(400, "invalid_request", "code, redirect_uri, and code_verifier are required") + if not MIN_CODE_VERIFIER_LENGTH <= len(code_verifier) <= MAX_CODE_VERIFIER_LENGTH: + return _oauth_error(400, "invalid_request", "code_verifier must be 43 to 128 characters (RFC 7636)") parsed = _open_sealed(code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY) if parsed is None: return _oauth_error(400, "invalid_grant", "the authorization code is invalid") @@ -492,12 +548,17 @@ async def _authorization_code_grant( return _oauth_error(400, "invalid_grant", "the authorization code was issued to a different client") if not _pkce_verifier_matches(code_verifier, parsed.code_challenge): return _oauth_error(400, "invalid_grant", "PKCE verification failed") - if await guard.already_used(parsed.jti): - return _oauth_error(400, "invalid_grant", "the authorization code was already used") - await guard.mark_used(parsed.jti) + # Revalidate the user BEFORE claiming the code, so a transient DB outage (a retryable + # 503) does not consume a still-valid code and force the client to restart sign-in. failure = await reload_user(parsed.user_id) if failure is not None: return _reload_failure_response(failure) + # Atomic single-use claim is the gate: on a concurrent double-redeem exactly one caller + # wins, and a claim that cannot be recorded fails closed. + if not await guard.claim( + f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}", GATEWAY_AUTH_CODE_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + ): + return _oauth_error(400, "invalid_grant", "the authorization code was already used") return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 41e70d5f1eb3..ab383ee1f0fc 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2420,12 +2420,19 @@ async def sso_readiness(): def _is_same_origin_return_path(return_to: str) -> bool: - """True for a strictly relative return path (starts with ``/``, not - protocol-relative ``//``, no backslash tricks browsers normalize to slashes), which - stays on the gateway's own origin by construction and is therefore safe to honor - without a configured ``control_plane_url``. Used by the MCP gateway DCR authorize - round-trip so a browser sent through login lands back on the authorize request.""" - return return_to.startswith("/") and not return_to.startswith("//") and "\\" not in return_to + """True for a strictly relative return path that stays on the gateway's own origin by + construction, and is therefore safe to honor without a configured ``control_plane_url``. + Used by the MCP gateway DCR authorize round-trip so a browser sent through login lands + back on the authorize request. + + Requires a single leading ``/`` (not protocol-relative ``//``), no backslash (browsers + fold ``\\`` to ``/``, so ``/\\evil.com`` would escape the origin), and no control or + whitespace characters. Rejecting control chars keeps a ``\\r\\n``/tab-bearing value out + of the redirect ``Location`` and the ``litellm_cp_return_to`` cookie entirely, rather + than relying on downstream header encoding to neutralize it.""" + if not return_to.startswith("/") or return_to.startswith("//") or "\\" in return_to: + return False + return not any(ord(ch) < 0x20 or ch in (" ", "\x7f") for ch in return_to) class SSOAuthenticationHandler: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 1dd342b7b28c..7bab11300099 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2899,19 +2899,20 @@ async def test_token_root_does_not_resolve_private_server_for_external_client(): @pytest.mark.asyncio -async def test_register_root_resolves_single_oauth2_server(): - """When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" - try: - from fastapi import Request +async def test_register_root_does_aggregate_dcr_not_single_server_resolution(): + """Root /register is the aggregate DCR endpoint: it mints a stateless llm_dcrc_ client + from the request's redirect_uris and does NOT resolve a single configured oauth2 server + (a single-server deployment registers at /{server}/register instead).""" + import json - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() @@ -2922,33 +2923,37 @@ async def test_register_root_resolves_single_oauth2_server(): mock_request.headers = {} try: - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}), + ), + patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"), ): - result = await register_client(request=mock_request, mcp_server_name=None) + response = await register_client(request=mock_request, mcp_server_name=None) - # Should resolve to the single server and return its name as client_id - assert result["client_id"] == "test_oauth" - assert "redirect_uris" in result + body = json.loads(response.body) + assert body["client_id"].startswith("llm_dcrc_") + assert body["client_id"] != "test_oauth" + assert body["token_endpoint_auth_method"] == "none" finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio -async def test_register_root_does_not_resolve_private_server_for_external_client(): - """Root /register must not reveal or use a hidden MCP server.""" - try: - from fastapi import Request +async def test_register_root_does_not_leak_a_private_server(): + """Root /register never resolves or reveals a configured server, so a private one cannot + leak to an external caller: it always mints the aggregate DCR client instead.""" + import json - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server(available_on_public_internet=False) @@ -2962,17 +2967,19 @@ async def test_register_root_does_not_resolve_private_server_for_external_client with ( patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), + new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}), ), patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", return_value="198.51.100.10", ), + patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"), ): - result = await register_client(request=mock_request, mcp_server_name=None) + response = await register_client(request=mock_request, mcp_server_name=None) - assert result["client_id"] == "dummy_client" - assert result["redirect_uris"] == ["https://llm.example.com/callback"] + body = json.loads(response.body) + assert body["client_id"].startswith("llm_dcrc_") + assert "test_oauth" not in body["client_id"] finally: global_mcp_server_manager.registry.clear() @@ -7260,6 +7267,7 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): global_mcp_server_manager.registry.clear() + def test_gateway_dcr_flow_routing_engages_only_for_llm_dcrc_clients(monkeypatch): """The aggregate DCR arms engage for llm_dcrc_ client_ids (register always mints one, authorize/token route into the aggregate flow); a non-gateway client_id keeps the @@ -7312,8 +7320,6 @@ def test_gateway_dcr_flow_routing_engages_only_for_llm_dcrc_clients(monkeypatch) assert token_response.status_code == 400 assert token_response.json()["error"] == "invalid_grant" - # a non-gateway (upstream-issued) client_id is not routed into the aggregate arm; it - # falls to the per-server exchange, which 404s for an unknown server upstream_shaped = client.post( "/token", data={"grant_type": "authorization_code", "client_id": "regular-upstream-client", "code": "x"}, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 6db337c6f4f6..6bcf68fb05d5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -62,7 +62,7 @@ def _request(path="/authorize", query="", cookies=None, method="GET"): async def _register(redirect_uris) -> dict: - response = await register_aggregate_client(request=_request("/register"), request_body={"redirect_uris": redirect_uris}) + response = await register_aggregate_client(request_body={"redirect_uris": redirect_uris}) return json.loads(response.body) @@ -103,7 +103,7 @@ async def test_register_allows_loopback_http_for_dev_clients(): ], ) async def test_register_rejects_bad_redirect_uris(redirect_uris): - response = await register_aggregate_client(request=_request("/register"), request_body={"redirect_uris": redirect_uris}) + response = await register_aggregate_client(request_body={"redirect_uris": redirect_uris}) assert response.status_code == 400 assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") @@ -117,7 +117,9 @@ async def test_tampered_client_id_does_not_open(): assert open_gateway_dcr_client("other_prefix") is None -def _authorize(client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code"): +def _authorize( + client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code" +): return aggregate_authorize( request=_request(query=f"client_id={client_id}"), client_id=client_id, @@ -188,24 +190,27 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): authorize_response = _authorize(client_id, session_user_id="u1") handle, cookies = _flow_cookie_from(authorize_response) - denied = complete_connect_flow( + denied = await complete_connect_flow( request=_request("/authorize/complete", cookies=cookies, method="POST"), flow_handle=handle, session_user_id="attacker", + cache=DualCache(), ) assert denied.status_code == 403 - anonymous = complete_connect_flow( + anonymous = await complete_connect_flow( request=_request("/authorize/complete", cookies=cookies, method="POST"), flow_handle=handle, session_user_id=None, + cache=DualCache(), ) assert anonymous.status_code == 401 - completed = complete_connect_flow( + completed = await complete_connect_flow( request=_request("/authorize/complete", cookies=cookies, method="POST"), flow_handle=handle, session_user_id="u1", + cache=DualCache(), ) assert completed.status_code == 303 redirect = urlparse(completed.headers["location"]) @@ -269,15 +274,19 @@ async def _token(**overrides): @pytest.mark.asyncio async def test_complete_rejects_missing_tampered_and_expired_flows(): - missing = complete_connect_flow( - request=_request("/authorize/complete", method="POST"), flow_handle="nope", session_user_id="u1" + missing = await complete_connect_flow( + request=_request("/authorize/complete", method="POST"), + flow_handle="nope", + session_user_id="u1", + cache=DualCache(), ) assert missing.status_code == 400 - tampered = complete_connect_flow( + tampered = await complete_connect_flow( request=_request("/authorize/complete", cookies={f"{CONNECT_FLOW_COOKIE_PREFIX}h1": "garbage"}, method="POST"), flow_handle="h1", session_user_id="u1", + cache=DualCache(), ) assert tampered.status_code == 400 @@ -353,10 +362,11 @@ async def test_token_gates_on_live_user_revalidation(failure, expected_status, e client_id = (await _register([REDIRECT_URI]))["client_id"] authorize_response = _authorize(client_id, session_user_id="deactivated-user") handle, cookies = _flow_cookie_from(authorize_response) - completed = complete_connect_flow( + completed = await complete_connect_flow( request=_request("/authorize/complete", cookies=cookies, method="POST"), flow_handle=handle, session_user_id="deactivated-user", + cache=DualCache(), ) code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] @@ -377,3 +387,103 @@ async def _reload_user_failing(user_id: str): ) assert response.status_code == expected_status assert json.loads(response.body)["error"] == expected_error + + +@pytest.mark.asyncio +async def test_flow_is_single_use_shared_cache_rejects_second_complete(): + """A double-submit of the finish step mints only ONE code: the second complete over the + same cache fails invalid_request (atomic flow claim), so one sign-in cannot yield two codes.""" + cache = DualCache() + client_id = (await _register([REDIRECT_URI]))["client_id"] + handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1")) + + first = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache, + ) + assert first.status_code == 303 + second = await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache, + ) + assert second.status_code == 400 + assert json.loads(second.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_token_rejects_out_of_range_code_verifier(): + """RFC 7636: a code_verifier outside 43-128 chars is invalid_request, not a confusing + invalid_grant PKCE-mismatch.""" + for bad in ["short", "x" * 200]: + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code="llm_gcode_whatever", + redirect_uri=REDIRECT_URI, + client_id="llm_dcrc_x", + code_verifier=bad, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_authorize_rejects_over_long_state(): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = aggregate_authorize( + request=_request(query=f"client_id={client_id}"), + client_id=client_id, + redirect_uri=REDIRECT_URI, + state="s" * 2000, + code_challenge=CODE_CHALLENGE, + code_challenge_method="S256", + response_type="code", + session_user_id="u1", + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + + +@pytest.mark.asyncio +async def test_non_ascii_code_challenge_fails_grant_not_500(): + """A non-ASCII code_challenge (unvalidated from the client) must yield a clean + invalid_grant, never a TypeError-driven 500 (bytes comparison, not str).""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + # Seal a code carrying a non-ASCII challenge directly (authorize requires S256 shape, + # but the challenge charset is not validated there, so this state is reachable). + from datetime import datetime, timezone + + code = _seal( + GATEWAY_AUTH_CODE_PREFIX, + _GatewayAuthCode( + user_id="u1", + client_id=client_id, + redirect_uri=REDIRECT_URI, + code_challenge="challenge-with-€-non-ascii", + jti="jti-x", + iat=int(datetime.now(timezone.utc).timestamp()), + exp=int(datetime.now(timezone.utc).timestamp()) + 120, + ), + ) + response = await aggregate_token( + request=_request("/token", method="POST"), + grant_type="authorization_code", + code=code, + redirect_uri=REDIRECT_URI, + client_id=client_id, + code_verifier=CODE_VERIFIER, + refresh_token=None, + master_key=MASTER_KEY, + reload_user=_reload_user_active, + cache=DualCache(), + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" From 4e9ffb889d9a42c14bd77aa2ca8c107483ff7bc3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:44:59 -0700 Subject: [PATCH 10/16] feat(mcp): admit gateway DCR session bearers at the aggregate /mcp scope --- .../mcp_server/auth/user_api_key_auth_mcp.py | 87 +++++++++- .../auth/test_user_api_key_auth_mcp.py | 162 ++++++++++++++++++ 2 files changed, 246 insertions(+), 3 deletions(-) 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 4f324e6b4230..774c7181aa77 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 @@ -25,6 +25,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( EnvelopeIdentity, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + is_session_bearer_shaped, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -124,6 +127,17 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _is_aggregate_mcp_scope(route: str, mcp_servers: list[str] | None) -> bool: + """True when a request targets the aggregate ``/mcp`` endpoint rather than any named + server. Named targets arrive either through ``x-mcp-servers`` (``mcp_servers``) or a + path segment (``/mcp/{server}`` / ``/{server}/mcp``); the aggregate scope has neither. + The gateway-DCR session arm and challenge fire only here, so a per-server flow is never + affected.""" + if mcp_servers: + return False + return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + + def _is_aggregate_gateway_dcr_challenge_scope( route: str, mcp_servers: list[str] | None, @@ -141,11 +155,9 @@ def _is_aggregate_gateway_dcr_challenge_scope( client. Fails closed to the original admission error otherwise.""" if not _is_litellm_auth_admission_error(exc): return False - if mcp_servers: - return False if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): return False - return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + return _is_aggregate_mcp_scope(route, mcp_servers) def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException: @@ -362,6 +374,22 @@ async def mock_body(): request=request, route=request_route, ) + elif ( + is_mcp_gateway_dcr_enabled() + and _is_aggregate_mcp_scope(request_route, mcp_servers) + and oauth2_headers + and is_session_bearer_shaped(oauth2_headers["Authorization"]) + ): + # A gateway DCR session bearer at the aggregate /mcp scope: open the + # identity-only session token and admit under the live litellm user it + # references. A session-shaped bearer that does not open fails closed with + # the aggregate invalid_token challenge; a non-session bearer never reaches + # here (is_session_bearer_shaped is false) and falls through to the oauth2 arm. + validated_user_api_key_auth = await MCPRequestHandler._admit_gateway_session( + authorization_value=oauth2_headers["Authorization"], + request=request, + route=request_route, + ) elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and @@ -626,6 +654,59 @@ async def _admit_dcr_bridge_delegate( case _: assert_never(result) + @staticmethod + async def _admit_gateway_session( + authorization_value: str, + request: Request, + route: str, + ) -> UserAPIKeyAuth: + """Open a gateway DCR session bearer and admit the live litellm user it references. + + The custody sibling of :meth:`_admit_dcr_bridge_delegate`: the session token seals + no upstream credential (those are vaulted per user and resolved at egress), so this + admits identity only and injects no per-server header. The token's signature proves + the user signed in when it was minted, but authorization is resolved fresh here, the + sealed ``user_id`` reloads the current user record through the SAME + :meth:`_reload_admitted_user` the bridge user-subject path uses, and the admitted + identity runs through the centralized policy gate, so the user's present team, org, + budget, and SCIM state gate the request rather than a snapshot frozen at mint time. + + Fails closed with the aggregate ``invalid_token`` challenge on an expired, tampered, + or foreign token, on a refresh token presented at the tool edge, and when the + referenced user is missing, deactivated, or rejected by the policy gate. The + pre-DB gates (size, IP, route allowlist) run first, mirroring the bridge arm and the + standard pipeline, so a caller blocked by IP or route is turned away before any + crypto or DB read.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + NotSessionBearer, + SessionBearerAdmitted, + SessionBearerInvalid, + resolve_session_bearer, + session_keys_from_master_key, + ) + from litellm.proxy.proxy_server import master_key + + if not master_key: + raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + + await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) + + keys = session_keys_from_master_key(master_key) + result = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc)) + match result: + case SessionBearerAdmitted(): + admitted = await MCPRequestHandler._reload_admitted_user(result.principal.user_id) + await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) + return admitted + case SessionBearerInvalid(): + raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) + case NotSessionBearer(): + # is_session_bearer_shaped gated entry, so a non-session bearer here means a + # session-shaped-but-empty value; fail closed with the same challenge. + raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) + case _: + assert_never(result) + @staticmethod async def _run_pre_db_read_auth_checks(request: Request, route: str) -> None: """Run the proxy-wide gates ``user_api_key_auth`` applies before any key lookup: the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e38101267023..e31e07827222 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6080,3 +6080,165 @@ async def _raise_500(api_key, request): with pytest.raises(ProxyException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) assert str(exc_info.value.code) == "500" + + +@pytest.mark.asyncio +class TestGatewaySessionAdmission: + """The aggregate /mcp session-bearer admission arm (mcp_gateway_dcr). A valid session + token admits under the LIVE litellm user it references; an invalid/expired/refresh/foreign + token fails closed with the aggregate invalid_token challenge; the arm fires ONLY at the + aggregate scope with the flag on, never for named servers or per-server flows.""" + + _MASTER_KEY = "sk-gateway-session-admission-master-key" + _FLAG = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" + + def _session_bearer(self, user_id="sso-user-42", client_id="llm_dcrc_abc"): + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SessionPrincipal, + mint_session_token, + mint_session_refresh_token, + ) + + keys = session_keys_from_master_key(self._MASTER_KEY) + principal = SessionPrincipal(user_id=user_id, client_id=client_id) + return mint_session_token, mint_session_refresh_token, principal, keys + + def _access_token(self, **kw): + from datetime import datetime, timezone + + mint, _refresh, principal, keys = self._session_bearer(**kw) + return mint(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + + def _scope(self, bearer, path="/mcp", extra_headers=()): + return { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), (b"authorization", f"Bearer {bearer}".encode()), *extra_headers], + } + + @staticmethod + @contextlib.contextmanager + def _patch_user_reload(*, user_id, active=True): + get_user_object = AsyncMock( + return_value=MagicMock( + user_id=user_id, + metadata={"scim_active": active} if not active else {"scim_active": True}, + user_role=None, + object_permission=None, + object_permission_id=None, + ) + ) + with ( + patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_user_object + + async def test_valid_session_admits_under_live_user_at_aggregate_scope(self): + token = self._access_token(user_id="sso-user-42") + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + self._patch_user_reload(user_id="sso-user-42") as get_user_object, + ): + auth_result, _h, _servers, mcp_server_auth_headers, _o, _r = await MCPRequestHandler.process_mcp_request( + self._scope(token) + ) + assert get_user_object.await_args.kwargs["user_id"] == "sso-user-42" + assert auth_result.user_id == "sso-user-42" + mock_auth.assert_not_called() + # Identity-only admission injects no per-server upstream credential (unlike the + # bridge envelope arm); the headers dict is whatever the request carried, here empty. + assert not mcp_server_auth_headers + + async def test_expired_session_fails_closed_with_invalid_token_challenge(self): + from datetime import datetime, timezone + + mint, _refresh, principal, keys = self._session_bearer() + token = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert exc_info.value.status_code == 401 + assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"] + + async def test_tampered_session_fails_closed(self): + token = self._access_token() + tampered = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb") + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(tampered)) + assert exc_info.value.status_code == 401 + + async def test_refresh_token_is_not_admitted_at_the_tool_edge(self): + from datetime import datetime, timezone + + _mint, refresh, principal, keys = self._session_bearer() + refresh_token = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(refresh_token)) + assert exc_info.value.status_code == 401 + + async def test_foreign_key_session_fails_closed(self): + token = self._access_token() + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", "sk-a-totally-different-master-key"), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert exc_info.value.status_code == 401 + + async def test_arm_does_not_fire_when_flag_off(self): + """Flag off: a session-shaped bearer is treated as an ordinary bearer and hits the + oauth2 arm, which validates it as a litellm credential and fails it there (not the + session arm). Proven by user_api_key_auth being called, unlike the flag-on path.""" + token = self._access_token() + with ( + patch(self._FLAG, return_value=False), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401), + ) as mock_auth, + ): + with pytest.raises((HTTPException, ProxyException)): + await MCPRequestHandler.process_mcp_request(self._scope(token)) + mock_auth.assert_called_once() + + async def test_arm_does_not_fire_for_named_server(self): + """A session-shaped bearer aimed at a named server (path scope) does not enter the + aggregate arm; it is treated as an ordinary bearer on that server.""" + token = self._access_token() + with ( + patch(self._FLAG, return_value=True), + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401), + ) as mock_auth, + ): + with pytest.raises((HTTPException, ProxyException)): + await MCPRequestHandler.process_mcp_request(self._scope(token, path="/mcp/github")) + mock_auth.assert_called_once() From 1b28128b22132ac0b4038c3ebaec6da2414957fe Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 09:57:17 -0700 Subject: [PATCH 11/16] refactor(mcp): bind org_id on session admission and clarify unreachable arm - _reload_admitted_user binds the user's org_id so the org-level MCP ceiling stays in force for a gateway session (a ceiling can only narrow; multi-org users are capped conservatively to their primary org) instead of being silently skipped - trim the NotSessionBearer arm comment to state it is simply unreachable --- .../mcp_server/auth/user_api_key_auth_mcp.py | 23 +++++---- .../auth/test_user_api_key_auth_mcp.py | 49 ++++++++----------- 2 files changed, 34 insertions(+), 38 deletions(-) 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 774c7181aa77..0ef7416b6778 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 @@ -375,8 +375,7 @@ async def mock_body(): route=request_route, ) elif ( - is_mcp_gateway_dcr_enabled() - and _is_aggregate_mcp_scope(request_route, mcp_servers) + _is_aggregate_mcp_scope(request_route, mcp_servers) and oauth2_headers and is_session_bearer_shaped(oauth2_headers["Authorization"]) ): @@ -701,8 +700,8 @@ async def _admit_gateway_session( case SessionBearerInvalid(): raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) case NotSessionBearer(): - # is_session_bearer_shaped gated entry, so a non-session bearer here means a - # session-shaped-but-empty value; fail closed with the same challenge. + # Unreachable: the arm is entered only for an is_session_bearer_shaped + # value. Kept for match exhaustiveness and fails closed regardless. raise _aggregate_gateway_dcr_challenge(request, invalid_token=True) case _: assert_never(result) @@ -752,12 +751,15 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: The DCR client authenticates via SSO at the bridged authorize, which yields a user subject rather than a virtual key, so the envelope admits under the user's own - identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the - returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then - computes which servers the user may reach, so the user's litellm MCP grants and access groups - gate the request exactly as a key's do. Only the user's OWN object permission is bound: a - ``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so - team-inherited MCP grants for a user are a follow-up (they need a many-teams union + identity: the reloaded ``user_id``, the user's own MCP object permission, and the user's + ``org_id`` ride on the returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` + the key path uses then computes which servers the user may reach, so the user's litellm MCP + grants and access groups gate the request exactly as a key's do. Binding ``org_id`` keeps the + org-level MCP ceiling in force for this admission rather than silently skipping it; a user's + primary organization is used, so a user who spans organizations is capped conservatively (the + ceiling can only narrow the result, never broaden it). Only the user's OWN object permission is + bound: a ``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, + so team-inherited MCP grants for a user are a follow-up (they need a many-teams union ``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed. @@ -805,6 +807,7 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: return UserAPIKeyAuth( user_id=user_object.user_id, user_role=user_object.user_role, + org_id=user_object.organization_id, object_permission=object_permission, object_permission_id=user_object.object_permission_id, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e31e07827222..33e78027872a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5121,6 +5121,7 @@ async def test_user_subject_envelope_admits_under_the_reloaded_user(self): self._patch_user_reload( return_value=MagicMock( user_id="sso-user-7", + organization_id=None, metadata={"scim_active": True}, user_role=None, object_permission=None, @@ -5163,6 +5164,7 @@ async def test_user_subject_envelope_carries_the_users_mcp_object_permission(sel self._patch_user_reload( return_value=MagicMock( user_id="sso-user-7", + organization_id=None, metadata={"scim_active": True}, user_role=None, object_permission=object_permission, @@ -5240,7 +5242,7 @@ async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self with ( patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), - self._patch_user_reload(return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False})), + self._patch_user_reload(return_value=MagicMock(user_id="offboarded-user", organization_id=None, metadata={"scim_active": False})), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: @@ -6087,10 +6089,9 @@ class TestGatewaySessionAdmission: """The aggregate /mcp session-bearer admission arm (mcp_gateway_dcr). A valid session token admits under the LIVE litellm user it references; an invalid/expired/refresh/foreign token fails closed with the aggregate invalid_token challenge; the arm fires ONLY at the - aggregate scope with the flag on, never for named servers or per-server flows.""" + aggregate scope, never for named servers or per-server flows.""" _MASTER_KEY = "sk-gateway-session-admission-master-key" - _FLAG = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" def _session_bearer(self, user_id="sso-user-42", client_id="llm_dcrc_abc"): from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( @@ -6122,10 +6123,11 @@ def _scope(self, bearer, path="/mcp", extra_headers=()): @staticmethod @contextlib.contextmanager - def _patch_user_reload(*, user_id, active=True): + def _patch_user_reload(*, user_id, active=True, organization_id=None): get_user_object = AsyncMock( return_value=MagicMock( user_id=user_id, + organization_id=organization_id, metadata={"scim_active": active} if not active else {"scim_active": True}, user_role=None, object_permission=None, @@ -6139,10 +6141,24 @@ def _patch_user_reload(*, user_id, active=True): ): yield get_user_object + async def test_session_admission_binds_org_id_so_the_org_ceiling_applies(self): + """The admitted auth carries the user's org_id, so get_allowed_mcp_servers keeps the + org-level MCP ceiling in force for a gateway session instead of skipping it.""" + token = self._access_token(user_id="org-user") + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + self._patch_user_reload(user_id="org-user", organization_id="org-123"), + ): + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(self._scope(token)) + assert auth_result.org_id == "org-123" + async def test_valid_session_admits_under_live_user_at_aggregate_scope(self): token = self._access_token(user_id="sso-user-42") with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", @@ -6166,7 +6182,6 @@ async def test_expired_session_fails_closed_with_invalid_token_challenge(self): mint, _refresh, principal, keys = self._session_bearer() token = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), ): with pytest.raises(HTTPException) as exc_info: @@ -6178,7 +6193,6 @@ async def test_tampered_session_fails_closed(self): token = self._access_token() tampered = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb") with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), ): with pytest.raises(HTTPException) as exc_info: @@ -6191,7 +6205,6 @@ async def test_refresh_token_is_not_admitted_at_the_tool_edge(self): _mint, refresh, principal, keys = self._session_bearer() refresh_token = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value() with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), ): with pytest.raises(HTTPException) as exc_info: @@ -6201,37 +6214,17 @@ async def test_refresh_token_is_not_admitted_at_the_tool_edge(self): async def test_foreign_key_session_fails_closed(self): token = self._access_token() with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", "sk-a-totally-different-master-key"), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope(token)) assert exc_info.value.status_code == 401 - async def test_arm_does_not_fire_when_flag_off(self): - """Flag off: a session-shaped bearer is treated as an ordinary bearer and hits the - oauth2 arm, which validates it as a litellm credential and fails it there (not the - session arm). Proven by user_api_key_auth being called, unlike the flag-on path.""" - token = self._access_token() - with ( - patch(self._FLAG, return_value=False), - patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), - patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - new_callable=AsyncMock, - side_effect=ProxyException(message="bad key", type="auth_error", param="api_key", code=401), - ) as mock_auth, - ): - with pytest.raises((HTTPException, ProxyException)): - await MCPRequestHandler.process_mcp_request(self._scope(token)) - mock_auth.assert_called_once() - async def test_arm_does_not_fire_for_named_server(self): """A session-shaped bearer aimed at a named server (path scope) does not enter the aggregate arm; it is treated as an ordinary bearer on that server.""" token = self._access_token() with ( - patch(self._FLAG, return_value=True), patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", From b82fb75292432d2b37e45c589d486331d212ebbe Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:51:06 -0700 Subject: [PATCH 12/16] feat(mcp): union team-inherited MCP grants across all of a user's teams for keyless admission --- .../mcp_server/auth/user_api_key_auth_mcp.py | 97 ++++++++++++++++--- .../auth/test_user_api_key_auth_mcp.py | 94 ++++++++++++++++++ 2 files changed, 180 insertions(+), 11 deletions(-) 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 0ef7416b6778..4965b2da1ed8 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,3 +1,4 @@ +import asyncio import re from datetime import datetime, timezone from typing import Dict, List, Optional, Set, Tuple, cast @@ -1593,10 +1594,86 @@ async def _get_allowed_mcp_servers_for_key( @staticmethod async def _get_allowed_mcp_servers_for_team( - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: - """ - Get allowed MCP servers for a team. + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> list[str]: + """ + Get allowed MCP servers a caller inherits from team membership. + + For a key-based caller the ``team_id`` on the auth is the one team, and the result + is that team's grants (byte-identical to before this method learned about multiple + teams). For a user-subject caller admitted WITHOUT a key (the gateway DCR session + bearer and the bridge user-envelope, which carry a ``user_id`` and no ``api_key`` or + ``team_id``), a ``UserAPIKeyAuth`` can only pin one team while the user may belong to + many, so the inherited grant is the UNION across every team the user belongs to. + Without this a signed-in user would see only servers granted to them directly and + none granted through their teams, which is how servers are meant to be shared + (assign teams, not individuals). Key-based auth never enters the union branch, so its + access is unchanged. + """ + team_ids = await MCPRequestHandler._team_ids_for_mcp_grant(user_api_key_auth) + if not team_ids: + return [] + per_team = await asyncio.gather( + *( + MCPRequestHandler._allowed_mcp_servers_for_single_team(team_id, user_api_key_auth) + for team_id in team_ids + ) + ) + return list({server for servers in per_team for server in servers}) + + @staticmethod + async def _team_ids_for_mcp_grant(user_api_key_auth: UserAPIKeyAuth | None) -> list[str]: + """The team ids whose MCP grants a caller inherits. + + A key-based caller (``api_key`` set) or any caller with an explicit ``team_id`` uses + that single team, so key auth is unchanged. Only a keyless user-subject caller (no + ``api_key``, no ``team_id``, a ``user_id``) fans out to the user's full team list, + resolved once from the live user record. The ``UI_TEAM_ID`` sentinel resolves to no + teams exactly as before.""" + if user_api_key_auth is None: + return [] + if user_api_key_auth.team_id: + return [] if user_api_key_auth.team_id == UI_TEAM_ID else [user_api_key_auth.team_id] + if user_api_key_auth.api_key is not None or not user_api_key_auth.user_id: + return [] + return await MCPRequestHandler._resolve_user_team_ids(user_api_key_auth.user_id, user_api_key_auth) + + @staticmethod + async def _resolve_user_team_ids(user_id: str, user_api_key_auth: UserAPIKeyAuth) -> list[str]: + """The distinct team ids a user belongs to, from the live user record. Returns [] on + no DB, a missing user, or any resolution failure so a lookup blip narrows access + rather than raising; the caller's direct grants still apply.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return [] + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises + verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {str(e)}") + return [] + if user_object is None or not user_object.teams: + return [] + return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID)) + + @staticmethod + async def _allowed_mcp_servers_for_single_team( + team_id: str, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> list[str]: + """Allowed MCP servers granted by ONE team. Unions two sources: - Legacy team.object_permission (mcp_servers, mcp_access_groups, @@ -1620,17 +1697,15 @@ async def _get_allowed_mcp_servers_for_team( user_api_key_cache, ) - if user_api_key_auth is None or not user_api_key_auth.team_id or prisma_client is None: + if not team_id or team_id == UI_TEAM_ID or prisma_client is None: return [] - if user_api_key_auth.team_id == UI_TEAM_ID: - return [] - - team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( - team_id=user_api_key_auth.team_id, + parent_otel_span = user_api_key_auth.parent_otel_span if user_api_key_auth is not None else None + team_obj: LiteLLM_TeamTable | None = await get_team_object( + team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, + parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) if team_obj is None: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 33e78027872a..25f7feb099bf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6235,3 +6235,97 @@ async def test_arm_does_not_fire_for_named_server(self): with pytest.raises((HTTPException, ProxyException)): await MCPRequestHandler.process_mcp_request(self._scope(token, path="/mcp/github")) mock_auth.assert_called_once() + + +@pytest.mark.asyncio +class TestUserSubjectTeamUnion: + """_get_allowed_mcp_servers_for_team unions across ALL a user's teams for a keyless + user-subject caller (the gateway DCR session bearer and bridge user-envelope), while a + key-based caller keeps its single-team behavior byte-identically.""" + + def _team(self, team_id, mcp_servers): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + + return LiteLLM_TeamTable( + team_id=team_id, + access_group_ids=[], + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers + ), + ) + + @contextlib.contextmanager + def _patch(self, *, teams_by_id, user_teams=None): + async def _get_team_object(team_id, **kw): + return teams_by_id.get(team_id) + + async def _get_user_object(user_id, **kw): + return MagicMock(user_id=user_id, teams=user_teams or []) + + with ( + patch("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object), + patch("litellm.proxy.auth.auth_checks.get_user_object", _get_user_object), + patch("litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", AsyncMock(return_value=[])), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ): + yield + + async def test_keyless_user_unions_servers_across_all_their_teams(self): + teams = {"team-a": self._team("team-a", ["srv1", "srv2"]), "team-b": self._team("team-b", ["srv2", "srv3"])} + auth = UserAPIKeyAuth(user_id="sso-user", api_key=None) + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv1", "srv2", "srv3"} + + async def test_key_based_caller_uses_single_team_only(self): + """A key-based caller (api_key set) with a team_id sees ONLY that team, even though the + same user belongs to other teams: key auth must be byte-identical to before.""" + teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2", "srv3"])} + auth = UserAPIKeyAuth(user_id="sso-user", api_key="sk-hash", team_id="team-a") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv1"} + + async def test_keyless_user_with_explicit_team_id_uses_that_team_only(self): + """A keyless caller that already pins a team_id (not the user-subject fan-out shape) + resolves only that team; the union is strictly for the no-team-id user-subject case.""" + teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])} + auth = UserAPIKeyAuth(user_id="sso-user", api_key=None, team_id="team-a") + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv1"} + + async def test_keyless_user_with_no_teams_gets_nothing_from_teams(self): + auth = UserAPIKeyAuth(user_id="lonely-user", api_key=None) + with self._patch(teams_by_id={}, user_teams=[]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert result == [] + + async def test_ui_session_team_id_still_resolves_to_nothing(self): + from litellm.proxy._types import UI_TEAM_ID + + auth = UserAPIKeyAuth(user_id="dash-user", api_key="sk-hash", team_id=UI_TEAM_ID) + with self._patch(teams_by_id={}, user_teams=["team-a"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert result == [] + + async def test_team_ids_helper_gates_on_shape(self): + from litellm.proxy._types import UI_TEAM_ID + + # key-based with team -> that team + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key="sk", team_id="t1", user_id="u") + ) == ["t1"] + # keyless user-subject, no team -> resolved from user record + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key=None, user_id="u") + ) == ["t2", "t3"] + # keyless, no user_id -> nothing + assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == [] + # UI sentinel -> nothing + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key="sk", team_id=UI_TEAM_ID, user_id="u") + ) == [] From 5df5fa50faaeda34ff031b3f9bf9b995575ce114 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 10:01:13 -0700 Subject: [PATCH 13/16] fix(mcp): gate the team-grant union on an admission marker, not api_key absence A JWT-authenticated caller is also keyless with a user_id and, absent a team claim, no team_id, so gating the multi-team union on api_key-is-None silently broadened JWT MCP access to the union of every team the user belongs to. _reload_admitted_user now stamps MCP_ADMITTED_USER_SUBJECT_METADATA and the union fires only for that positive marker, so the gateway session and bridge user paths union while JWT and other keyless auth keep their prior behavior. Regression-tested. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 53 +++++++++++++------ .../auth/test_user_api_key_auth_mcp.py | 34 +++++++++--- 2 files changed, 65 insertions(+), 22 deletions(-) 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 4965b2da1ed8..94650b64f011 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 @@ -128,6 +128,21 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +MCP_ADMITTED_USER_SUBJECT_METADATA = "mcp_admitted_user_subject" +"""Marker key stamped into ``UserAPIKeyAuth.metadata`` by ``_reload_admitted_user`` for a +subject admitted keyless through the gateway session / bridge user path. It is what lets +``_team_ids_for_mcp_grant`` union the user's teams for exactly those admissions without also +broadening JWT auth, which produces a structurally identical keyless auth.""" + + +def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth) -> bool: + """True when this auth is a keyless subject admitted by the gateway session / bridge user + path (stamped at admission), as opposed to a JWT or other keyless auth that merely lacks a + ``team_id``.""" + metadata = user_api_key_auth.metadata + return isinstance(metadata, dict) and metadata.get(MCP_ADMITTED_USER_SUBJECT_METADATA) is True + + def _is_aggregate_mcp_scope(route: str, mcp_servers: list[str] | None) -> bool: """True when a request targets the aggregate ``/mcp`` endpoint rather than any named server. Named targets arrive either through ``x-mcp-servers`` (``mcp_servers``) or a @@ -811,6 +826,7 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: org_id=user_object.organization_id, object_permission=object_permission, object_permission_id=user_object.object_permission_id, + metadata={MCP_ADMITTED_USER_SUBJECT_METADATA: True}, ) @staticmethod @@ -1599,16 +1615,17 @@ async def _get_allowed_mcp_servers_for_team( """ Get allowed MCP servers a caller inherits from team membership. - For a key-based caller the ``team_id`` on the auth is the one team, and the result - is that team's grants (byte-identical to before this method learned about multiple - teams). For a user-subject caller admitted WITHOUT a key (the gateway DCR session - bearer and the bridge user-envelope, which carry a ``user_id`` and no ``api_key`` or - ``team_id``), a ``UserAPIKeyAuth`` can only pin one team while the user may belong to - many, so the inherited grant is the UNION across every team the user belongs to. - Without this a signed-in user would see only servers granted to them directly and - none granted through their teams, which is how servers are meant to be shared - (assign teams, not individuals). Key-based auth never enters the union branch, so its - access is unchanged. + For a caller with a ``team_id`` (every key-based caller) the result is that one team's + grants, byte-identical to before this method learned about multiple teams. For a + subject admitted keyless through the gateway DCR session or bridge user path, a + ``UserAPIKeyAuth`` can only pin one team while the user may belong to many, so the + inherited grant is the UNION across every team the user belongs to. Without this a + signed-in user would see only servers granted to them directly and none granted + through their teams, which is how servers are meant to be shared (assign teams, not + individuals). The union is gated on the admission marker + ``_team_ids_for_mcp_grant`` checks, NOT on ``api_key is None``, so JWT auth (also + keyless, also possibly team-less) keeps its prior behavior and is not silently + broadened. """ team_ids = await MCPRequestHandler._team_ids_for_mcp_grant(user_api_key_auth) if not team_ids: @@ -1625,16 +1642,20 @@ async def _get_allowed_mcp_servers_for_team( async def _team_ids_for_mcp_grant(user_api_key_auth: UserAPIKeyAuth | None) -> list[str]: """The team ids whose MCP grants a caller inherits. - A key-based caller (``api_key`` set) or any caller with an explicit ``team_id`` uses - that single team, so key auth is unchanged. Only a keyless user-subject caller (no - ``api_key``, no ``team_id``, a ``user_id``) fans out to the user's full team list, - resolved once from the live user record. The ``UI_TEAM_ID`` sentinel resolves to no - teams exactly as before.""" + A caller with an explicit ``team_id`` (every key-based caller, and any auth that pins + a team) uses that single team, so key auth is byte-identical. The fan-out to the + user's full team list happens ONLY for a subject admitted keyless through the gateway + session or bridge user path, which ``_reload_admitted_user`` stamps with + ``MCP_ADMITTED_USER_SUBJECT_METADATA``. Gating on that positive marker rather than on + ``api_key is None`` is deliberate: JWT auth also produces a keyless ``user_id`` auth + with no ``team_id``, and it must keep its prior behavior (no team-inherited grants) + rather than silently gaining the union across every team the user belongs to. The + ``UI_TEAM_ID`` sentinel resolves to no teams exactly as before.""" if user_api_key_auth is None: return [] if user_api_key_auth.team_id: return [] if user_api_key_auth.team_id == UI_TEAM_ID else [user_api_key_auth.team_id] - if user_api_key_auth.api_key is not None or not user_api_key_auth.user_id: + if not user_api_key_auth.user_id or not _is_mcp_admitted_user_subject(user_api_key_auth): return [] return await MCPRequestHandler._resolve_user_team_ids(user_api_key_auth.user_id, user_api_key_auth) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 25f7feb099bf..47656afa3fc6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6272,9 +6272,17 @@ async def _get_user_object(user_id, **kw): ): yield + @staticmethod + def _admitted_subject(user_id): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCP_ADMITTED_USER_SUBJECT_METADATA, + ) + + return UserAPIKeyAuth(user_id=user_id, api_key=None, metadata={MCP_ADMITTED_USER_SUBJECT_METADATA: True}) + async def test_keyless_user_unions_servers_across_all_their_teams(self): teams = {"team-a": self._team("team-a", ["srv1", "srv2"]), "team-b": self._team("team-b", ["srv2", "srv3"])} - auth = UserAPIKeyAuth(user_id="sso-user", api_key=None) + auth = self._admitted_subject("sso-user") with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) assert set(result) == {"srv1", "srv2", "srv3"} @@ -6298,7 +6306,7 @@ async def test_keyless_user_with_explicit_team_id_uses_that_team_only(self): assert set(result) == {"srv1"} async def test_keyless_user_with_no_teams_gets_nothing_from_teams(self): - auth = UserAPIKeyAuth(user_id="lonely-user", api_key=None) + auth = self._admitted_subject("lonely-user") with self._patch(teams_by_id={}, user_teams=[]): result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) assert result == [] @@ -6318,14 +6326,28 @@ async def test_team_ids_helper_gates_on_shape(self): assert await MCPRequestHandler._team_ids_for_mcp_grant( UserAPIKeyAuth(api_key="sk", team_id="t1", user_id="u") ) == ["t1"] - # keyless user-subject, no team -> resolved from user record + # keyless subject admitted by the gateway/bridge path (marked), no team -> resolved from record with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): - assert await MCPRequestHandler._team_ids_for_mcp_grant( - UserAPIKeyAuth(api_key=None, user_id="u") - ) == ["t2", "t3"] + assert await MCPRequestHandler._team_ids_for_mcp_grant(self._admitted_subject("u")) == ["t2", "t3"] # keyless, no user_id -> nothing assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == [] + # keyless with a user_id but NOT admission-marked (JWT auth) -> nothing (unchanged behavior) + with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): + assert await MCPRequestHandler._team_ids_for_mcp_grant( + UserAPIKeyAuth(api_key=None, user_id="jwt-user") + ) == [] # UI sentinel -> nothing assert await MCPRequestHandler._team_ids_for_mcp_grant( UserAPIKeyAuth(api_key="sk", team_id=UI_TEAM_ID, user_id="u") ) == [] + + async def test_jwt_keyless_user_without_team_claim_does_not_union(self): + """Regression for the review finding: a JWT-authenticated caller is also keyless with a + user_id and (with no team claim) no team_id, but it is NOT admission-marked, so it must + keep its prior behavior of inheriting no team grants rather than silently gaining the + union across every team the user belongs to.""" + teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])} + jwt_auth = UserAPIKeyAuth(user_id="jwt-user", api_key=None) # no admission marker + with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(jwt_auth) + assert result == [] From 262e93ec276e3e4bfcbfc6b1b6b4e0b098fefc43 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:55:03 -0700 Subject: [PATCH 14/16] feat(ui): connect-flow interlude banner on the MCP apps grid for gateway DCR sign-in --- .../src/app/chat/integrations/page.tsx | 9 +++ .../chat/ConnectFlowBanner.test.tsx | 33 +++++++++++ .../src/components/chat/ConnectFlowBanner.tsx | 55 +++++++++++++++++++ 3 files changed, 97 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx index f85dd591199c..452c7918bf63 100644 --- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx @@ -4,6 +4,7 @@ import { Suspense, useEffect } from "react"; import { useRouter, useSearchParams } from "next/navigation"; import { useChatShell } from "@/contexts/ChatShellContext"; import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; +import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; // useSearchParams() requires a Suspense boundary for static export. function IntegrationsPageContent() { @@ -11,6 +12,13 @@ function IntegrationsPageContent() { const router = useRouter(); const searchParams = useSearchParams(); const oauthReturn = searchParams.get("mcpOauthReturn"); + // Set by the gateway DCR authorize when a DCR client sends the user here to + // authorize servers before finishing sign-in (see gateway_dcr_flow.py). The + // handle keys the sealed per-flow cookie; connect_client is the client origin + // for display only. connect_flow is NOT cleaned from the URL: the finish form + // needs it, and the sealed cookie (not the URL) is the security boundary. + const connectFlow = searchParams.get("connect_flow"); + const connectClient = searchParams.get("connect_client"); // Clean up the OAuth return param after it's been consumed — real routing means // we no longer need it to pick a tab, but it should not linger in the address bar. @@ -24,6 +32,7 @@ function IntegrationsPageContent() { return (
+ {connectFlow && }
); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx new file mode 100644 index 000000000000..188ba00becbf --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -0,0 +1,33 @@ +import { describe, expect, it, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import ConnectFlowBanner from "./ConnectFlowBanner"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: () => "https://gateway.example.com", +})); + +describe("ConnectFlowBanner", () => { + it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => { + const { container } = render(); + + const form = container.querySelector("form")!; + expect(form.getAttribute("method")).toBe("POST"); + expect(form.getAttribute("action")).toBe("https://gateway.example.com/authorize/complete"); + + const hidden = form.querySelector('input[name="flow"]') as HTMLInputElement; + expect(hidden.value).toBe("flow-handle-123"); + // No token, code, or secret is ever placed in the form; the sealed cookie carries them. + expect(form.innerHTML).not.toContain("token"); + }); + + it("shows the client origin so the user knows what they are connecting to", () => { + render(); + expect(screen.getAllByText(/claude\.ai/).length).toBeGreaterThan(0); + expect(screen.getByRole("button", { name: /finish connecting/i })).toBeInTheDocument(); + }); + + it("falls back to a generic label when the client origin is unknown", () => { + render(); + expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx new file mode 100644 index 000000000000..f4728b540c83 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -0,0 +1,55 @@ +"use client"; + +import React from "react"; +import { CheckCircle } from "lucide-react"; +import { getProxyBaseUrl } from "@/components/networking"; + +interface Props { + flowHandle: string; + clientOrigin: string | null; +} + +/** + * The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user + * through the gateway sign-in and lands them on the apps grid to authorize servers. The + * grid below authorizes individual servers into the per-user vault; this banner is the + * deliberate finish step. + * + * "Finish connecting" is a native form POST to the proxy's /authorize/complete, not a + * fetch: the endpoint 303-redirects the browser back to the DCR client's own redirect URI + * with the gateway authorization code, and only a full-page navigation carries the + * HttpOnly per-flow cookie and follows that cross-origin redirect. The flow handle is the + * only field; the sealed flow cookie set at /authorize holds everything else. + */ +const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => { + const action = `${getProxyBaseUrl()}/authorize/complete`; + const clientLabel = clientOrigin ?? "the application"; + + return ( +
+
+
+ +
+

Connect your MCP servers to {clientLabel}

+

+ Authorize the servers you want to use below. When you are ready, finish connecting and you will be + returned to {clientLabel}. +

+
+
+
+ + +
+
+
+ ); +}; + +export default ConnectFlowBanner; From 5ab7a855c3fbe917d0acefef751631d7353fc851 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 18:51:57 -0700 Subject: [PATCH 15/16] feat(ui): simplify DCR connect grid, auto-finish on tab close, fix connect-status flash In the gateway DCR connect flow the apps grid now reads as "authorize your servers" rather than a chat feature; the connectMode prop drops the Beta badge, the "use in chat" subtitle, and the tool-count chrome Closing the connect tab now best-effort finishes the flow via navigator.sendBeacon to /authorize/complete, so the gateway authorization code still reaches the client's loopback without an explicit click; the explicit "Finish connecting" button stays as the reliable path. The beacon is skipped while a per-server authorize is navigating away and after the button was pressed, so it never double-delivers or fires mid-authorize Authorized servers previously flashed "Connect" for a second before flipping to "Connected" because the per-user credential checks ran only after the whole tool-count fetch finished. They now fire in parallel with the tool-count load, and each card shows a skeleton in the button slot until its status resolves, so the state never flips under the user --- ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../src/app/chat/integrations/page.tsx | 7 +- .../chat/ConnectFlowBanner.test.tsx | 47 ++++++++- .../src/components/chat/ConnectFlowBanner.tsx | 42 ++++++-- .../src/components/chat/MCPAppsPanel.tsx | 98 ++++++++++++------- .../src/hooks/mcpOAuthUtils.ts | 8 ++ .../src/hooks/useUserMcpOAuthFlow.tsx | 3 +- 7 files changed, 155 insertions(+), 52 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9de171397db6..14a3076c27eb 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1712,7 +1712,7 @@ }, "src/components/chat/MCPAppsPanel.tsx": { "no-nested-ternary": { - "count": 7 + "count": 6 } }, "src/components/chat/MCPConnectPicker.tsx": { diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx index 452c7918bf63..30ce62d8081f 100644 --- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx @@ -33,7 +33,12 @@ function IntegrationsPageContent() { return (
{connectFlow && } - +
); } diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx index 188ba00becbf..9baf52966a61 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -1,11 +1,17 @@ -import { describe, expect, it, vi } from "vitest"; -import { render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; import ConnectFlowBanner from "./ConnectFlowBanner"; +import { PERSERVER_CONNECTING_KEY } from "@/hooks/mcpOAuthUtils"; vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "https://gateway.example.com", })); +afterEach(() => { + vi.restoreAllMocks(); + sessionStorage.clear(); +}); + describe("ConnectFlowBanner", () => { it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => { const { container } = render(); @@ -30,4 +36,41 @@ describe("ConnectFlowBanner", () => { render(); expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0); }); + + it("best-effort auto-finishes on pagehide (closing the tab)", () => { + const beaconMock = vi.fn(() => true); + vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock }); + render(); + + window.dispatchEvent(new Event("pagehide")); + + expect(beaconMock).toHaveBeenCalledTimes(1); + const [url, body] = beaconMock.mock.calls[0] as unknown as [string, URLSearchParams]; + expect(url).toBe("https://gateway.example.com/authorize/complete"); + expect(body.toString()).toContain("flow=flow-xyz"); + }); + + it("does NOT auto-finish while a per-server connect is navigating away", () => { + const beaconMock = vi.fn(() => true); + vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock }); + render(); + + // the per-server connect flow sets this right before it navigates to the upstream IdP + sessionStorage.setItem(PERSERVER_CONNECTING_KEY, "1"); + window.dispatchEvent(new Event("pagehide")); + + expect(beaconMock).not.toHaveBeenCalled(); + }); + + it("does NOT double-fire the auto-finish after the button was pressed", () => { + const beaconMock = vi.fn(() => true); + vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock }); + const { container } = render(); + + // jsdom does not submit forms; fire the form's submit so onSubmit marks it finished + fireEvent.submit(container.querySelector("form")!); + window.dispatchEvent(new Event("pagehide")); + + expect(beaconMock).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index f4728b540c83..a46b4e59959b 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -1,8 +1,9 @@ "use client"; -import React from "react"; +import React, { useEffect, useRef } from "react"; import { CheckCircle } from "lucide-react"; import { getProxyBaseUrl } from "@/components/networking"; +import { PERSERVER_CONNECTING_KEY } from "@/hooks/mcpOAuthUtils"; interface Props { flowHandle: string; @@ -13,17 +14,38 @@ interface Props { * The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user * through the gateway sign-in and lands them on the apps grid to authorize servers. The * grid below authorizes individual servers into the per-user vault; this banner is the - * deliberate finish step. + * finish step that returns the user to the client. * - * "Finish connecting" is a native form POST to the proxy's /authorize/complete, not a - * fetch: the endpoint 303-redirects the browser back to the DCR client's own redirect URI - * with the gateway authorization code, and only a full-page navigation carries the - * HttpOnly per-flow cookie and follows that cross-origin redirect. The flow handle is the - * only field; the sealed flow cookie set at /authorize holds everything else. + * Finishing happens two ways, both hitting the proxy's /authorize/complete, which mints the + * gateway authorization code and 303-redirects to the DCR client's own redirect URI: + * - The explicit "Finish connecting" button is a native form POST, so the full-page + * navigation carries the HttpOnly per-flow cookie and follows the cross-origin redirect + * to the client's loopback. This is the reliable path. + * - Closing (or navigating away from) the tab fires a best-effort navigator.sendBeacon to the + * same endpoint. The browser follows the 303 to the client's loopback, so in most browsers + * the code still reaches the client without an explicit click. This is a convenience, not a + * consent gate: consent already happened at sign-in, so returning the user is safe. It is + * skipped while a per-server connect is navigating away (that is not leaving the flow), + * and after the button was pressed (which already delivers the code). */ const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => { const action = `${getProxyBaseUrl()}/authorize/complete`; const clientLabel = clientOrigin ?? "the application"; + const finishedRef = useRef(false); + + useEffect(() => { + sessionStorage.removeItem(PERSERVER_CONNECTING_KEY); + + const autoFinishOnLeave = () => { + if (finishedRef.current) return; + if (sessionStorage.getItem(PERSERVER_CONNECTING_KEY) === "1") return; + if (typeof navigator.sendBeacon === "function") { + navigator.sendBeacon(action, new URLSearchParams({ flow: flowHandle })); + } + }; + window.addEventListener("pagehide", autoFinishOnLeave); + return () => window.removeEventListener("pagehide", autoFinishOnLeave); + }, [action, flowHandle]); return (
@@ -33,12 +55,12 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => {

Connect your MCP servers to {clientLabel}

- Authorize the servers you want to use below. When you are ready, finish connecting and you will be - returned to {clientLabel}. + Authorize the servers you want to use below, then finish connecting to return to {clientLabel}. Closing + this tab finishes for you.

-
+ (finishedRef.current = true)}>