From 691598aa011ddc6dcd9d4a6d6577a4eb5a993973 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:28:51 -0700 Subject: [PATCH 1/8] fix(mcp): preserve scopes during managed OAuth reauthorization --- .../mcp_server/discoverable_endpoints.py | 234 +++++++++++++++--- .../_experimental/mcp_server/exceptions.py | 12 +- .../mcp_server/faults/list_outcomes.py | 27 +- .../mcp_server/gateway_dcr_flow.py | 21 +- .../mcp_server/mcp_server_manager.py | 51 +++- litellm/proxy/_lazy_openapi_snapshot.json | 16 ++ .../types/mcp_server/mcp_server_manager.py | 3 + tests/integration/mcp/test_mcp_oauth_flows.py | 81 +++++- .../mcp_server/faults/test_list_outcomes.py | 29 +++ .../mcp_server/test_discoverable_endpoints.py | 201 ++++++++++++++- .../mcp_server/test_gateway_dcr_flow.py | 25 +- .../test_mcp_oauth_passthrough_tools.py | 12 + .../mcp_server/test_mcp_server_manager.py | 37 +++ .../ConnectFlowBanner.integration.test.tsx | 50 ++++ .../src/components/chat/ConnectFlowBanner.tsx | 1 + .../src/components/chat/MCPAppsPanel.tsx | 3 + .../src/components/networking.tsx | 1 + .../src/hooks/useUserMcpOAuthFlow.test.tsx | 68 +++++ .../src/hooks/useUserMcpOAuthFlow.tsx | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 1 + 20 files changed, 816 insertions(+), 59 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0253719f410..b1cd855ceb1 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,6 +12,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response +from mcp.client.auth.utils import get_client_metadata_scopes, union_scopes +from mcp.shared.auth import OAuthMetadata from pydantic import ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_logger @@ -110,6 +112,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload from litellm.proxy._types import LiteLLM_MCPServerTable # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. @@ -219,6 +222,8 @@ def encode_state_with_base_url( expected_issuer: str | None = None, authorization_response_iss_parameter_supported: bool = False, oauth_nonce: str | None = None, + managed_scope: str | None = None, + scope_user_id: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -251,6 +256,7 @@ def encode_state_with_base_url( """ state_data: Final = { "oauth_nonce": oauth_nonce, + **({"managed_scope": managed_scope, "scope_user_id": scope_user_id} if managed_scope is not None else {}), "base_url": base_url, "original_state": original_state, "code_challenge": code_challenge, @@ -290,6 +296,18 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +_MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_" + + +class _ManagedAuthorizationCode(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + upstream_code: str = Field(min_length=1) + mcp_server_id: str = Field(min_length=1) + issuer: str | None = None + requested_scope: str + scope_user_id: str | None = None + + _BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_" @@ -298,6 +316,8 @@ class _BridgeAuthorizationCode(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) oauth_nonce: str | None = None + requested_scope: str | None = None + expected_issuer: str | None = None upstream_code: str = Field(min_length=1) litellm_user_id: str = Field(min_length=1) mcp_server_id: str = Field(min_length=1) @@ -314,6 +334,8 @@ def seal_bridge_authorization_code( litellm_user_id: str, mcp_server_id: str, oauth_nonce: str | None = None, + requested_scope: str | None = None, + expected_issuer: str | None = None, ) -> str: """Seal the upstream authorization code and the SSO-captured litellm user into a gateway authorization code. The DCR client only echoes this opaque value back at the token endpoint; the @@ -328,6 +350,8 @@ def seal_bridge_authorization_code( "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id, "oauth_nonce": oauth_nonce, + "requested_scope": requested_scope, + "expected_issuer": expected_issuer, }, sort_keys=True, ) @@ -643,6 +667,7 @@ async def _store_per_user_token_server_side( token_response: dict[str, Any], identity_binding_proof: str | None = None, cimd_client_id: str | None = None, + requested_scope: str | None = None, ) -> None: """Persist the OAuth token server-side and warm the Redis cache. @@ -667,7 +692,7 @@ async def _store_per_user_token_server_side( expires_in = None refresh_token: Final[str | None] = token_response.get("refresh_token") or None - raw_scope: Final = token_response.get("scope") + raw_scope: Final = token_response.get("scope", requested_scope) scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None try: @@ -954,6 +979,46 @@ async def _resolve_oauth_authorization_user( return user_id +async def _managed_authorization_scope( + request: Request, server: MCPServer, scope: str | None, authorized_user_id: str | None +) -> tuple[str | None, str | None]: + from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential + from litellm.proxy.proxy_server import prisma_client + + gateway_credential: Final = authorized_user_id is None and await oauth_authorization_uses_gateway_credential( + request + ) + session_user: Final = ( + _session_cookie_user_id(request) if authorized_user_id is None and not gateway_credential else None + ) + scope_user_id: Final = authorized_user_id or ( + await authorize_oauth_credential_request(request, server.server_id) + if gateway_credential + else session_user + if session_user and await _user_can_reach_mcp_server(session_user, server.server_id) + else None + ) + prior_grant: Final = ( + await get_user_oauth_credential(prisma_client, scope_user_id, server.server_id) + if scope_user_id and prisma_client is not None + else None + ) + merged_scope: Final = union_scopes(" ".join(prior_grant.get("scopes") or ()) if prior_grant else None, scope) + selected_scope: Final = ( + get_client_metadata_scopes( + merged_scope, + None, + OAuthMetadata.model_validate(server.authorization_server_metadata) + if server.authorization_server_metadata is not None + else None, + ["authorization_code", "refresh_token"], + ) + if merged_scope is not None + else None + ) + return selected_scope, scope_user_id + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -1029,6 +1094,17 @@ async def authorize_with_server( return subject litellm_user_id = subject + managed_per_user: Final = resolved_server.is_gateway_managed_oauth2 and resolved_server.needs_user_oauth_token + initial_scope: Final = union_scopes( + "openid" if enforce_binding else None, + scope or (" ".join(resolved_server.scopes) if resolved_server.scopes else None), + ) + selected_scope, scope_user_id = ( + await _managed_authorization_scope(request, resolved_server, initial_scope, litellm_user_id) + if managed_per_user + else (initial_scope, None) + ) + oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None encoded_state: Final = encode_state_with_base_url( base_url=base_url, @@ -1038,7 +1114,11 @@ async def authorize_with_server( code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, litellm_user_id=litellm_user_id, - mcp_server_id=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None, + mcp_server_id=resolved_server.server_id + if (litellm_user_id or ephemeral_dcr_client or managed_per_user) + else None, + managed_scope=selected_scope if managed_per_user else None, + scope_user_id=scope_user_id, dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None, dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None, dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method @@ -1057,13 +1137,8 @@ async def authorize_with_server( } if oauth_nonce: params["nonce"] = oauth_nonce - if scope: - params["scope"] = scope - elif resolved_server.scopes: - params["scope"] = " ".join(resolved_server.scopes) - - if enforce_binding and "openid" not in params.get("scope", "").split(): - params["scope"] = f"openid {params.get('scope', '')}".strip() + if selected_scope: + params["scope"] = selected_scope if code_challenge: params["code_challenge"] = code_challenge @@ -1090,18 +1165,11 @@ def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: return "gateway_stored" if mcp_server.client_id or get_cimd_client_id(mcp_server) else "caller_supplied" -async def _saved_cimd_refresh_client_id( +async def _saved_refresh_credential( server: MCPServer, user_id: str | None, refresh_token: str | None -) -> str | None: - """Reuse only the client identity bound to this caller's presented refresh grant.""" - if ( - not user_id - or not refresh_token - or not server.needs_user_oauth_token - or server.auth_type != MCPAuth.oauth2 - or server.client_id - or server.client_secret - ): +) -> "OAuthCredentialPayload | None": + """Reuse only metadata bound to this caller's presented refresh grant.""" + if not user_id or not refresh_token or not server.needs_user_oauth_token or server.auth_type != MCPAuth.oauth2: return None from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential @@ -1117,7 +1185,33 @@ async def _saved_cimd_refresh_client_id( stored_refresh: Final = credential.get("refresh_token") if not stored_refresh or not secrets.compare_digest(stored_refresh.encode(), refresh_token.encode()): return None - return credential.get("cimd_client_id") + return credential + + +def _open_managed_authorization_code( + code: str | None, server: MCPServer, grant_type: str, user_id: str | None +) -> _ManagedAuthorizationCode | None: + if not isinstance(code, str) or not code.startswith(_MANAGED_AUTH_CODE_PREFIX): + return None + decoded_code: Final = decrypt_value_helper( + code[len(_MANAGED_AUTH_CODE_PREFIX) :], "managed_authorization_code", return_original_value=False + ) + if not isinstance(decoded_code, str): + raise HTTPException(status_code=400, detail="Invalid authorization code") + try: + managed_code: Final = _ManagedAuthorizationCode.model_validate_json(decoded_code) + except ValidationError as exc: + raise HTTPException(status_code=400, detail="Invalid authorization code") from exc + if ( + grant_type != "authorization_code" + or not server.is_gateway_managed_oauth2 + or not server.needs_user_oauth_token + or managed_code.mcp_server_id != server.server_id + or managed_code.issuer != server.issuer + or (managed_code.scope_user_id is not None and managed_code.scope_user_id != user_id) + ): + raise HTTPException(status_code=400, detail="Authorization code does not match this OAuth flow") + return managed_code async def exchange_token_with_server( @@ -1159,10 +1253,18 @@ async def exchange_token_with_server( if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None else None ) - cimd_client_id: Final = ( - await _saved_cimd_refresh_client_id(resolved_server, request_user_id, refresh_token) + managed_code: Final = _open_managed_authorization_code(code, resolved_server, grant_type, request_user_id) + authorization_code: Final = managed_code.upstream_code if managed_code else code + + refresh_credential: Final = ( + await _saved_refresh_credential(resolved_server, request_user_id, refresh_token) if grant_type == "refresh_token" else None + ) + cimd_client_id: Final = ( + refresh_credential.get("cimd_client_id") + if refresh_credential and not resolved_server.client_id and not resolved_server.client_secret + else None ) or get_cimd_client_id(resolved_server) # The id, secret, and token-endpoint auth method must come from the same source. When the @@ -1234,7 +1336,7 @@ async def exchange_token_with_server( ) else: refresh_ownership = None # rebind-ok: grant-specific branches assign one ownership value - if not code: + if not authorization_code: raise HTTPException( status_code=400, detail="code is required for authorization_code grant", @@ -1244,14 +1346,18 @@ async def exchange_token_with_server( # below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the # sealed server to this request so a code minted for one bridge server cannot be spent at another. # A raw upstream code (scripted path) opens to None and the code is used as-is. - bridge_identity = open_bridge_authorization_code(code) - if bridge_identity is not None: - if bridge_identity.mcp_server_id != resolved_server.server_id: - raise HTTPException( - status_code=400, - detail="Authorization code was issued for a different MCP server", - ) - code = bridge_identity.upstream_code + bridge_identity = open_bridge_authorization_code(authorization_code) + if bridge_identity is not None and ( + bridge_identity.mcp_server_id != resolved_server.server_id + or ( + bridge_identity.requested_scope is not None + and bridge_identity.expected_issuer != resolved_server.issuer + ) + ): + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) binding: Final = resolved_server.oauth_identity_binding if binding is not None and binding.mode == "enforce": if bridge_identity is None or not bridge_identity.oauth_nonce: @@ -1273,7 +1379,7 @@ async def exchange_token_with_server( resolved_redirect_uri: Final = redirect_uri if bridge_token_relay else f"{proxy_base_url}/callback" token_data = { "grant_type": "authorization_code", - "code": code, + "code": bridge_identity.upstream_code if bridge_identity else authorization_code, "redirect_uri": resolved_redirect_uri, **token_request.body, } @@ -1353,6 +1459,16 @@ async def exchange_token_with_server( else None ) + requested_scope: Final = ( + managed_code.requested_scope + if managed_code is not None + else bridge_identity.requested_scope + if bridge_identity is not None + else scope or " ".join(refresh_credential.get("scopes") or ()) + if refresh_credential + else None + ) + # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. # Errors are non-fatal: the token is still returned to the client. @@ -1379,6 +1495,7 @@ async def exchange_token_with_server( user_id=user_id, token_response=token_response, identity_binding_proof=binding_proof, + requested_scope=requested_scope, **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), ) else: @@ -1431,8 +1548,9 @@ async def exchange_token_with_server( result["expires_in"] = token_response["expires_in"] if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - if token_response.get("scope"): - result["scope"] = token_response["scope"] + effective_scope: Final = token_response.get("scope", requested_scope) + if effective_scope is not None: + result["scope"] = effective_scope # RFC 6749 §5.1: token responses must not be cached. return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS) @@ -2116,6 +2234,7 @@ async def authorize_mcp_session( code_challenge_method: str | None = None, response_type: str | None = None, resource: str | None = None, + scope: str | None = None, ) -> Response: return aggregate_authorize( request=request, @@ -2127,6 +2246,7 @@ async def authorize_mcp_session( response_type=response_type, session_user_id=_session_cookie_user_id(request), resource=resource, + scope=scope, ) @@ -2169,6 +2289,7 @@ async def authorize( response_type=response_type, session_user_id=_session_cookie_user_id(request), resource=resource, + scope=scope, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -2281,7 +2402,9 @@ async def token_endpoint( ) -async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: +async def _vendor_credential_state( + user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () +) -> VendorCredentialState: """Whether the gateway itself can see a live vendor credential for this user and server. The one reading of "authorized" the connect page displays and the finish step enforces, so @@ -2299,7 +2422,10 @@ async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCreden credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed return "unavailable" - return "absent" if oauth_grant_state(credential) == "absent" else "present" + if oauth_grant_state(credential) == "absent": + return "absent" + granted_scopes: Final = credential.get("scopes") if credential else None + return "present" if set(requested_scopes).issubset(granted_scopes or ()) else "absent" @router.get("/authorize/flow") @@ -2542,6 +2668,11 @@ async def callback( mcp_server_id: Final = state_data.get("mcp_server_id") dcr_client_id: Final = state_data.get("dcr_client_id") dcr_client_secret: Final = state_data.get("dcr_client_secret") + managed_scope: Final = issuer_state.get("managed_scope") + expected_issuer: Final = TypeAdapter[str | None](str | None).validate_python( + issuer_state.get("expected_issuer") + ) + scope_owner: Final = TypeAdapter[str | None](str | None).validate_python(issuer_state.get("scope_user_id")) forwarded_code = code if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_bridge_authorization_code( @@ -2549,6 +2680,8 @@ async def callback( litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id, oauth_nonce=state_data.get("oauth_nonce"), + requested_scope=managed_scope if isinstance(managed_scope, str) else None, + expected_issuer=expected_issuer, ) elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_passthrough_authorization_code( @@ -2561,7 +2694,23 @@ async def callback( ), ) - params = {"code": forwarded_code, "state": original_state} + managed_code: Final = ( + _ManagedAuthorizationCode( + upstream_code=forwarded_code, + mcp_server_id=mcp_server_id, + issuer=expected_issuer, + requested_scope=managed_scope, + scope_user_id=scope_owner, + ) + if isinstance(managed_scope, str) and isinstance(mcp_server_id, str) and not litellm_user_id + else None + ) + returned_code: Final = ( + _MANAGED_AUTH_CODE_PREFIX + encrypt_value_helper(managed_code.model_dump_json()) + if managed_code is not None + else forwarded_code + ) + params = {"code": returned_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) response = RedirectResponse(url=complete_returned_url, status_code=302) _clear_oauth_state_cookie(response, request, state) @@ -2807,11 +2956,16 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response + resource_scopes: Final = [ + value + for value in (mcp_server.scopes or () if mcp_server else ()) + if value != "offline_access" or not mcp_server or not mcp_server.is_gateway_managed_oauth2 + ] if mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, - "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), + "scopes_supported": resource_scopes, } if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: @@ -2822,7 +2976,7 @@ async def _build_oauth_protected_resource_response( (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") ], "resource": resource_url, - "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), + "scopes_supported": resource_scopes, } diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index c818f6b05bd..a720006cebf 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -37,10 +37,12 @@ class MCPUpstreamAuthError(Exception): status_code: int, www_authenticate: str | None, server_name: str, + required_scope: str | None = None, ) -> None: self.status_code = status_code self.www_authenticate = www_authenticate self.server_name = server_name + self.required_scope = required_scope super().__init__(f"Upstream MCP server {server_name!r} returned {status_code}") def to_http_exception( @@ -74,18 +76,24 @@ class MCPUpstreamAuthError(Exception): ``get_passthrough_resource_metadata_url`` in ``oauth_utils.py``. """ challenge: str | None = self.www_authenticate - if challenge is None and self.status_code == 401 and base_url: + if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url: prefix: Final = base_url.rstrip("/") if request_path and request_path.startswith(f"/{self.server_name}/mcp"): resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp" else: resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}" challenge = f'Bearer resource_metadata="{resource_metadata_url}"' + scoped_challenge: Final = ( + (f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"') + + (f', scope="{self.required_scope}"' if self.required_scope else "") + if self.required_scope is not None + else challenge + ) detail: Final = "Forbidden" if self.status_code == 403 else "Unauthorized" return HTTPException( status_code=self.status_code, detail=detail, - headers={"www-authenticate": challenge} if challenge else None, + headers={"www-authenticate": scoped_challenge} if scoped_challenge else None, ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index aaeb50981b6..b94ed1ee251 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -10,13 +10,15 @@ becomes an outcome, never a second failure. from __future__ import annotations +import re from collections.abc import Iterator from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias +from urllib.request import parse_http_list, parse_keqv_list import httpx import httpx2 from mcp.types import Tool as MCPTool -from pydantic import ConfigDict +from pydantic import ConfigDict, TypeAdapter from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.exceptions import ( @@ -98,6 +100,29 @@ def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None return None +def upstream_insufficient_scope(exc: BaseException) -> str | None: + auth: Final = upstream_auth_challenge(exc) + if auth is None or auth[0] != 403 or not auth[1]: + return None + scheme, _, parameters = auth[1].partition(" ") + if scheme.lower() != "bearer": + return None + fields: Final = parse_http_list(parameters) + if any("=" not in field or not field.split("=", 1)[1] for field in fields): + return None + keys: Final = tuple(field.split("=", 1)[0].strip().lower() for field in fields) + if len(set(keys)) != len(keys) or any(re.fullmatch(r"[!#$%&'*+.^_`|~0-9a-z-]+", key) is None for key in keys): + return None + parsed: Final = TypeAdapter(dict[str, str]).validate_python(parse_keqv_list(fields)) + values: Final = {key: value for key, value in zip(keys, parsed.values())} + if values.get("error") != "insufficient_scope": + return None + scope: Final = values.get("scope", "") + if scope and re.fullmatch(r"[!#-\[\]-~]+(?: [!#-\[\]-~]+)*", scope) is None: + return None + return scope + + def raise_classified_list_failure( exc: BaseException, server_name: str, diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 5b45c0dd859..f9909bc5d92 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -250,7 +250,9 @@ class ConsentTeam(LiteLLMBaseModel): class LookupVendorCredential(Protocol): """Injected read of a user's vendor credential for one server.""" - def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + def __call__( + self, user_id: str, server_id: str, requested_scopes: tuple[str, ...] = (), / + ) -> Awaitable[VendorCredentialState]: ... class LookupServerReachability(Protocol): @@ -273,7 +275,9 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject ) -async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: +async def _unavailable_vendor_credential( + user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () +) -> VendorCredentialState: return "unavailable" @@ -310,6 +314,7 @@ class _ConnectFlow(LiteLLMBaseModel): exp: int resource_server_id: str | None = None audience: SessionAudience | None = None + requested_scopes: tuple[str, ...] = () class _GatewayAuthCode(LiteLLMBaseModel): @@ -555,6 +560,7 @@ def aggregate_authorize( response_type: str | None, session_user_id: str | None, resource: str | None = None, + scope: str | None = 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 @@ -589,6 +595,7 @@ def aggregate_authorize( code_challenge=code_challenge or "", resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, + requested_scopes=tuple((scope or "").split()) if scoped_server is not None else (), ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) @@ -748,6 +755,7 @@ def _new_connect_flow( code_challenge: str, resource_server_id: str | None, audience: SessionAudience | None, + requested_scopes: tuple[str, ...] = (), ) -> _ConnectFlow: now: Final = datetime.now(timezone.utc) return _ConnectFlow( @@ -760,6 +768,7 @@ def _new_connect_flow( exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS, resource_server_id=resource_server_id, audience=audience, + requested_scopes=requested_scopes, ) @@ -847,6 +856,7 @@ class ConnectFlowDescription(TypedDict): server_id: ReadOnly[str | None] server_name: ReadOnly[str | None] connected: ReadOnly[bool | None] + requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]] async def _describe_opened_flow( @@ -856,7 +866,11 @@ async def _describe_opened_flow( ) -> ConnectFlowDescription | Response: state, server = await _flow_target(flow, lookup_server_reachability) if state == "interactive" and server is not None: - credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) + credential: Final = ( + await lookup_vendor_credential(flow.user_id, server.server_id, flow.requested_scopes) + if flow.requested_scopes + else await lookup_vendor_credential(flow.user_id, server.server_id) + ) if credential == "unavailable": return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) interactive_description: Final[ConnectFlowDescription] = { @@ -865,6 +879,7 @@ async def _describe_opened_flow( "server_id": server.server_id, "server_name": server.server_name or server.alias or server.name, "connected": credential == "present", + **({"requested_scopes": flow.requested_scopes} if flow.requested_scopes else {}), } return interactive_description described: Final[ConnectFlowDescription] = { diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e492af133ff..87b0f907641 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -38,6 +38,7 @@ import httpx2 from fastapi import HTTPException from httpx import HTTPStatusError from mcp import ReadResourceResult, Resource +from mcp.shared.auth import OAuthMetadata from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import ( CallToolResult, @@ -57,7 +58,7 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, Field, SecretStr, TypeAdapter +from pydantic import AnyUrl, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, assert_never import litellm @@ -96,6 +97,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ServerListFault, raise_classified_list_failure, upstream_auth_challenge, + upstream_insufficient_scope, ) from litellm.proxy._experimental.mcp_server.interactions import bind_target from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution @@ -881,6 +883,10 @@ def carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serve new_server.authorization_response_iss_parameter_supported = ( # rebind-ok: publish on the existing rebuild object previous_server.authorization_response_iss_parameter_supported ) + if may_carry and new_server.issuer == previous_server.issuer and new_server.authorization_server_metadata is None: + new_server.authorization_server_metadata = ( # rebind-ok: publish capability on the existing rebuild object + previous_server.authorization_server_metadata + ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url if may_carry and new_server.token_url is None and previous_server.token_url: @@ -916,7 +922,12 @@ def _restrict_discovery_to_corroborated_authorization_server( return metadata if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): return metadata - if not metadata.token_url and not metadata.registration_url and not metadata.client_id_metadata_document_supported: + if ( + not metadata.token_url + and not metadata.registration_url + and not metadata.client_id_metadata_document_supported + and metadata.authorization_server_metadata is None + ): return metadata bridge_note: Final = ( " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" @@ -935,7 +946,12 @@ def _restrict_discovery_to_corroborated_authorization_server( bridge_note, ) return metadata.model_copy( - update={"token_url": None, "registration_url": None, "client_id_metadata_document_supported": False} + update={ + "token_url": None, + "registration_url": None, + "client_id_metadata_document_supported": False, + "authorization_server_metadata": None, + } ) @@ -2069,6 +2085,7 @@ class MCPServerManager: return server.model_copy( update={ "scopes": server.scopes or metadata.scopes, + "authorization_server_metadata": metadata.authorization_server_metadata, "issuer": server.issuer or discovered_issuer, "client_id_metadata_document_supported": metadata.client_id_metadata_document_supported, "authorization_response_iss_parameter_supported": ( @@ -2661,6 +2678,9 @@ class MCPServerManager: scopes=resolved_scopes, configured_scopes=tuple(configured_scopes) if configured_scopes else None, issuer=effective_issuer, + authorization_server_metadata=( + gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None + ), client_id_metadata_document_supported=( gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None ), @@ -3242,6 +3262,9 @@ class MCPServerManager: scopes=resolved_scopes, configured_scopes=configured_scopes, issuer=effective_issuer, + authorization_server_metadata=( + gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None + ), client_id_metadata_document_supported=( gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None ), @@ -5366,7 +5389,16 @@ class MCPServerManager: continue scopes = self._extract_scopes(data.get("scopes_supported")) + try: + authorization_metadata = OAuthMetadata.model_validate(data) + except ValidationError: + authorization_metadata = None metadata = MCPOAuthMetadata( + authorization_server_metadata=( + TypeAdapter(dict[str, object]).validate_python(authorization_metadata.model_dump(mode="json")) + if authorization_metadata is not None + else None + ), scopes=scopes, authorization_url=data.get("authorization_endpoint"), token_url=data.get("token_endpoint"), @@ -6365,11 +6397,12 @@ class MCPServerManager: # oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not # added here even though the list path still relays for it. relays_upstream_auth: Final = mcp_server.is_client_forwarded_token + managed_per_user: Final = mcp_server.is_gateway_managed_oauth2 and mcp_server.needs_user_oauth_token server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or "" async def _call_tool_via_client(client, params): async with self._limit_outbound_concurrency(mcp_server): - if not relays_upstream_auth: + if not relays_upstream_auth and not managed_per_user: return await client.call_tool( params, host_progress_callback=host_progress_callback, @@ -6392,8 +6425,16 @@ class MCPServerManager: allow_input_required=allow_input_required, ) except Exception as e: + required_scope: Final = upstream_insufficient_scope(e) if managed_per_user else None + if required_scope is not None: + raise MCPUpstreamAuthError( + status_code=403, + www_authenticate=None, + server_name=mcp_server.alias or server_label, + required_scope=required_scope, + ) from e auth_info: Final = _extract_upstream_auth_failure(e) - if auth_info is None or auth_info[0] != 401: + if not relays_upstream_auth or auth_info is None or auth_info[0] != 401: # A genuine (non-auth or 403-forbidden) upstream/transport failure. # raise_on_error demoted the client-layer log to debug, so surface it here at # warning level to keep the outage visible; the caller still gets the graceful diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 69db4525a17..f7cfcb5d5f7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -35929,6 +35929,22 @@ ], "title": "Resource" } + }, + { + "in": "query", + "name": "scope", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Scope" + } } ], "responses": { diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index a7027102250..9b37a63adb1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from datetime import datetime from typing import Annotated, Any, Final, Literal @@ -29,6 +30,7 @@ MCPInfo = Annotated[dict[str, Any], AfterValidator(_validate_mcp_protocol_metada class MCPOAuthMetadata(LiteLLMBaseModel): + authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) scopes: list[str] | None = None """Resource-driven scopes for the authorization request: the RFC 9728 protected-resource ``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource @@ -117,6 +119,7 @@ class MCPServer(LiteLLMBaseModel): # scope=="user" values must be supplied per-user. env_vars: list[dict[str, Any]] | None = None # OAuth-specific fields + authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) client_id: str | None = None client_secret: str | None = None issuer: str | None = None diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index e95cdb902e8..63561059489 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -1,11 +1,13 @@ import base64 import hashlib +import json import re import secrets import textwrap import time import uuid from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Final @@ -22,6 +24,9 @@ from integration._support.mcp import ( EntryPoint, McpCaller, McpPeer, + ScriptedTool, + scripted_peer, + text_result, Outcome, _outcome_from_rpc, call_tool, @@ -32,6 +37,7 @@ from integration._support.mcp import ( from integration._support.mcp_grants import create_toolset from integration._support.oauth_server import AuthorizationServer, oauth_server from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server ADD: Final = {"a": 2, "b": 3} CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb" @@ -240,7 +246,8 @@ class _Pkce: def _authorize_through_gateway( - gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce + gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce, + scope: str = "tools.call", ) -> str: started: Final = gateway.client.get( f"/{alias}/authorize", @@ -251,7 +258,7 @@ def _authorize_through_gateway( "state": "client-state", "code_challenge": pkce.challenge, "code_challenge_method": "S256", - "scope": "tools.call", + "scope": scope, }, headers={"x-litellm-api-key": key}, ) @@ -575,3 +582,73 @@ def test_token_exchange_callers_with_different_subject_tokens_own_separate_listi assert second.list_tools().ok assert _echoed_description(second.call(f"{alias}-add", probe)) == "Add two integers" assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer" + + +@contextmanager +def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, Wire]]: + with scripted_peer( + ScriptedTool("read_op", lambda _: text_result("read completed")), + ScriptedTool("write_op", lambda _: text_result("write completed")), + ScriptedTool("forbid_op", lambda _: text_result("must not execute")), + ) as inner: + def enforce(request: Request) -> Reply: + body: Final = json.loads(request.body) if request.body else {} + if body.get("method") == "tools/call": + name: Final = body["params"]["name"] + if name == "forbid_op": + return Reply(status=403, body=b"denied") + required: Final = "tools.write" if name == "write_op" else "tools.read" + token: Final = request.headers.get("authorization", "").removeprefix("Bearer ") + with auth.lock: + granted: Final = auth.access_tokens.get(token, {}).get("scope", "") + if required not in granted.split(): + return Reply(status=403, headers={"WWW-Authenticate": f'Bearer error="insufficient_scope", scope="{required}"'}) + forwarded: Final = httpx.request(request.method, inner.url, content=request.body, headers=dict(request.headers)) + return Reply(status=forwarded.status_code, body=forwarded.content, content_type=forwarded.headers.get("content-type", "application/json")) + + with wire_server(enforce) as wire: + yield McpPeer(wire.url + "/mcp", inner.calls), wire + + +@pytest.mark.parametrize("entry", ["rest", "server_mcp"]) +def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: Gateway, entry: EntryPoint) -> None: + with oauth_server(scopes=("tools.read", "tools.write")) as auth, _scope_enforcing_peer(auth) as (peer, wire), gateway.scenario() as scenario: + alias: Final = "scope" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, peer, auth, alias, auth_type="oauth2", oauth2_flow="authorization_code", + credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]}, + ) + owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read") + first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce) + assert first.status_code == 200, first.text + caller: Final = McpCaller(gateway, owner, entry, alias) + read: Final = caller.call(f"{alias}-read_op", {}, identity) + assert read.ok and read.text == "read completed", read.raw + peer.drain() + wire.drain() + denied: Final = ( + call_tool(gateway, owner, identity, f"{alias}-write_op", {}) if entry == "rest" + else caller.rpc("tools/call", {"name": f"{alias}-write_op", "arguments": {}}) + ) + assert denied.status_code == 403, denied.text + challenge: Final = denied.headers["www-authenticate"] + assert 'error="insufficient_scope"' in challenge and 'scope="tools.write"' in challenge + assert f'resource_metadata="{_base(gateway)}/.well-known/oauth-protected-resource/' in challenge + attempts: Final = tuple(request for request in wire.drain() if request.body and json.loads(request.body).get("method") == "tools/call") + assert len(attempts) == 1 and tool_calls(peer.drain()) == (), "denied operation was retried or executed" + second_code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.write") + second: Final = _redeem(gateway, alias, owner, "scope-client", second_code, pkce) + assert second.status_code == 200, second.text + with auth.lock: + granted: Final = auth.access_tokens[_issued_token(second.json())]["scope"] + assert granted.split() == ["tools.read", "tools.write"] + for operation in ("read", "write"): + result: Final = caller.call(f"{alias}-{operation}_op", {}, identity) + assert result.ok and result.text == f"{operation} completed", result.raw + forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity) + assert forbidden.status_code == 200 and not forbidden.ok, forbidden.raw + unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity) + assert unauthorized.status_code == 401, unauthorized.raw diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index b41e786881d..fe2254fbd46 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -267,3 +267,32 @@ def test_pure_non_auth_response_still_classifies_upstream_error(): fault = classify_list_exception(exc) assert fault.tag == "upstream_error" assert fault.status_code == 502 + + +@pytest.mark.parametrize("status, challenge, expected", [ + (403, 'Bearer error="insufficient_scope", scope="tools.write"', "tools.write"), + (403, 'bEaReR error=insufficient_scope, scope="tools.read tools.write"', "tools.read tools.write"), + (403, 'Bearer error="insufficient_scope"', ""), + (403, 'Bearer error="insufficient_scope", scope=', None), + (403, 'Bearer error="insufficient_scope", scope', None), + (401, 'Bearer error="insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error="invalid_token", scope="tools.write"', None), + (403, '', None), + (403, 'Basic error="insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error_description="error=insufficient_scope", scope="tools.write"', None), + (403, 'Bearer error="insufficient_scope", error="invalid_token"', None), + (403, 'Bearer error="insufficient_scope", scope="tools.write\r\ninjected"', None), + (403, 'Bearer error="insufficient_scope", scope="tools.write", Basic realm="other"', None), +]) +@pytest.mark.parametrize("client", ["httpx", "httpx2"]) +def test_managed_scope_step_up_only_accepts_unambiguous_bearer_challenges(status, challenge, expected, client): + import httpx2 + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_insufficient_scope + + transport = httpx if client == "httpx" else httpx2 + response = transport.Response(status, headers={"www-authenticate": challenge}, request=transport.Request("POST", "https://peer.test/mcp")) + cause = transport.HTTPStatusError("upstream rejected", request=response.request, response=response) + wrapped = RuntimeError("transport failure") + wrapped.__cause__ = cause + assert upstream_insufficient_scope(wrapped) == expected diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 4f22200145e..0086410fc26 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11981,6 +11981,7 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", " monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) prisma: Final = MagicMock() prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) return handler, signing_key @@ -13732,7 +13733,7 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change( "matching", "unicode_grant", "foreign_refresh", "missing_refresh", "missing_grant", "database_missing", "database_outage", "static_client", "anonymous", ]) -async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant( +async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant( monkeypatch: pytest.MonkeyPatch, state: str, ) -> None: from types import SimpleNamespace @@ -13766,12 +13767,204 @@ async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant( table.find_unique.side_effect = RuntimeError("database unavailable") if state == "static_client": server.client_id = "configured-client" - resolved: Final = await endpoints._saved_cimd_refresh_client_id( + resolved: Final = await endpoints._saved_refresh_credential( server, None if state == "anonymous" else "alice", "foreign-refresh" if state == "foreign_refresh" else grant, ) - assert resolved == (identity if state in ("matching", "unicode_grant") else None) - if state in ("anonymous", "static_client", "database_missing"): + if state in ("matching", "unicode_grant", "static_client"): + assert resolved is not None and resolved["cimd_client_id"] == identity + assert resolved["refresh_token"] == grant + else: + assert resolved is None + if state in ("anonymous", "database_missing"): table.find_unique.assert_not_awaited() else: table.find_unique.assert_awaited_once_with(where={"user_id_server_id": {"user_id": "alice", "server_id": server.server_id}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prior_scopes, expected", [(["tools.read"], "tools.read tools.write"), (None, "tools.write"), ([], "tools.write")]) +async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_identity, monkeypatch, prior_scopes, expected): + from types import SimpleNamespace + from urllib.parse import parse_qs, urlsplit + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "scope-grant-test-salt") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client", "scopes": ["tools.read"]}) + global_mcp_server_manager.registry[server.server_id] = server + credential = db.encrypt_value_helper(json.dumps({"type": "oauth2", "access_token": "old-access", "scopes": prior_scopes})) + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + response = await endpoints.authorize_with_server( + request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.write", + code_challenge="challenge", code_challenge_method="S256", + ) + assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] + assert table.find_unique.call_args.kwargs["where"] == { + "user_id_server_id": {"user_id": "jwt-owner", "server_id": server.server_id} + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", None)]) +async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted( + jwt_oauth_identity, monkeypatch, respx_mock, returned_scope, expected, +): + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlsplit + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "requested-scope-test-salt") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + global_mcp_server_manager.registry[server.server_id] = server + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=None) + table.upsert = AsyncMock() + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + started = await endpoints.authorize_with_server( + request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read tools.write", + code_challenge="challenge", code_challenge_method="S256", + ) + state = parse_qs(urlsplit(started.headers["location"]).query)["state"][0] + cookies = SimpleCookie() + cookies.load(started.headers["set-cookie"]) + callback_request = _token_request({"Cookie": "; ".join(f"{name}={value.value}" for name, value in cookies.items())}) + returned = await endpoints.callback(callback_request, code="provider-code", state=state, error=None, iss=None) + code = parse_qs(urlsplit(returned.headers["location"]).query)["code"][0] + upstream = respx_mock.post(server.token_url).respond(200, json={ + "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", + **({"scope": returned_scope} if returned_scope is not None else {}), + }) + exchanged = await endpoints.exchange_token_with_server( + request, server, "authorization_code", client_id="scope-client", code=code, code_verifier="verifier", redirect_uri="http://127.0.0.1:9/cb", client_secret=None, + ) + assert parse_qs(upstream.calls[0].request.content.decode())["code"] == ["provider-code"] + saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert saved is not None and saved.get("scopes") == expected + assert json.loads(exchanged.body)["scope"] == " ".join(expected or ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("advertised, expected", [(True, "tools.read offline_access"), (False, "tools.read")]) +async def test_managed_authorize_requests_only_selected_scope_and_supported_refresh_access(monkeypatch, advertised, expected): + from urllib.parse import parse_qs, urlsplit + + from mcp.shared.auth import OAuthMetadata + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + monkeypatch.setenv("LITELLM_SALT_KEY", "metadata-scope-test-salt") + metadata = OAuthMetadata( + issuer="https://idp.example.com", authorization_endpoint="https://idp.example.com/authorize", + token_endpoint="https://idp.example.com/token", response_types_supported=["code"], + scopes_supported=["tools.read", "admin"] + (["offline_access"] if advertised else []), + ) + server = _cimd_oauth_server().model_copy(update={"authorization_server_metadata": metadata.model_dump(mode="json")}) + response = await endpoints.authorize_with_server( + _cimd_request(), server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read", + code_challenge="challenge", code_challenge_method="S256", + ) + assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mismatch", ["server", "issuer", "user", "mode", "tampered", "invalid_payload"]) +async def test_managed_scope_code_cannot_cross_oauth_boundaries(jwt_oauth_identity, monkeypatch, respx_mock, mismatch): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "scope-boundary-test-salt") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + payload = endpoints._ManagedAuthorizationCode( + upstream_code="provider-code", mcp_server_id="other-server" if mismatch == "server" else server.server_id, + issuer="https://other-issuer.test" if mismatch == "issuer" else server.issuer, + requested_scope="tools.read tools.write", scope_user_id="other-user" if mismatch == "user" else "jwt-owner", + ) + code = endpoints._MANAGED_AUTH_CODE_PREFIX + ("invalid" if mismatch == "tampered" else endpoints.encrypt_value_helper("{}" if mismatch == "invalid_payload" else payload.model_dump_json())) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + if mismatch == "mode": + server.auth_type = MCPAuth.true_passthrough + with pytest.raises(HTTPException) as raised: + await endpoints.exchange_token_with_server( + request, server, "authorization_code", code=code, redirect_uri="http://127.0.0.1:9/cb", + client_id="scope-client", client_secret=None, code_verifier="verifier", + ) + assert raised.value.status_code == 400 + assert respx_mock.calls.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", None)]) +async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch, respx_mock, scope, expected): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + _, signing_key = jwt_oauth_identity + monkeypatch.setenv("LITELLM_SALT_KEY", "refresh-scope-test-salt") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"}) + global_mcp_server_manager.registry[server.server_id] = server + credential = db.encrypt_value_helper(json.dumps({ + "type": "oauth2", "access_token": "old-access", "refresh_token": "owner-refresh", "scopes": ["tools.read", "tools.write"], + })) + table = proxy_server.prisma_client.db.litellm_mcpusercredentials + table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential)) + table.upsert = AsyncMock() + upstream = respx_mock.post(server.token_url).respond(200, json={ + "access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh", + **({"scope": scope} if scope is not None else {}), + }) + request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"}) + await endpoints.exchange_token_with_server( + request, server, "refresh_token", code=None, redirect_uri=None, client_id="scope-client", client_secret=None, + code_verifier=None, refresh_token="owner-refresh", + ) + saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert saved is not None and saved.get("scopes") == expected + assert upstream.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("per_server_relay", [False, True]) +async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch, per_server_relay): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _cimd_oauth_server().model_copy(update={ + "scopes": ["tools.read", "offline_access"], "per_server_oauth_discovery": per_server_relay, + }) + global_mcp_server_manager.registry[server.server_id] = server + result = await endpoints._build_oauth_protected_resource_response(_cimd_request(), server.server_id, True) + assert result["scopes_supported"] == ["tools.read"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("granted, required, expected", [ + (["tools.read"], ("tools.write",), "absent"), + (["tools.read", "tools.write"], ("tools.write",), "present"), + (None, ("tools.write",), "absent"), + (["tools.read"], (), "present"), +]) +async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): + from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + + credential = {**_stored_grant(), "scopes": granted} + read = AsyncMock(return_value=credential) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(db, "get_user_oauth_credential", read) + assert await endpoints._vendor_credential_state("owner", "server", required) == expected + read.assert_awaited_once_with(prisma, "owner", "server") diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 76d071b1852..ed24a56eba5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -991,7 +991,7 @@ SCOPED_RESOURCE = "https://llm.example.com/mcp/github" _MANAGER_PATCH = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" -def _scoped_authorize(client_id, resource, session_user_id="u1"): +def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs): return aggregate_authorize( request=_request(query=f"client_id={client_id}"), client_id=client_id, @@ -1002,6 +1002,7 @@ def _scoped_authorize(client_id, resource, session_user_id="u1"): response_type="code", session_user_id=session_user_id, resource=resource, + **kwargs, ) @@ -2578,3 +2579,25 @@ async def test_token_exchange_relays_a_mint_refusal(failure, status, error): response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) assert response.status_code == status assert json.loads(response.body)["error"] == error + + +@pytest.mark.asyncio +async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it(): + from unittest.mock import AsyncMock, patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope="tools.write") + vendor = AsyncMock(return_value="absent") + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + assert json.loads(described.body)["connected"] is False + vendor.assert_awaited_once_with("u1", server.server_id, ("tools.write",)) + premature = await _complete_page(response, scoped_server=server, vendor=vendor) + assert premature.status_code == 400 + vendor.return_value = "present" + complete = await _complete_page(response, scoped_server=server, vendor=vendor) + assert complete.status_code == 303 + assert parse_qs(urlparse(complete.headers["location"]).query)["code"][0].startswith(GATEWAY_AUTH_CODE_PREFIX) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 7290f0de86d..bea4263268b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -498,3 +498,15 @@ def test_passthrough_admission_recognizes_only_matching_authorization(oauth_head server = MCPServer(server_id="catalog", name="catalog", alias="catalog", transport=MCPTransport.http) assert _client_has_passthrough_authorization(server, oauth_headers, server_headers) is authorized + + +@pytest.mark.parametrize("base_url, request_path, scope, expected", [ + ("https://gateway.test/proxy", "/alias/mcp", "tools.write", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/alias/mcp", error="insufficient_scope", scope="tools.write"'), + ("https://gateway.test/proxy", "/mcp/alias", "", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/mcp/alias", error="insufficient_scope"'), + (None, None, "tools.write", 'Bearer error="insufficient_scope", scope="tools.write"'), +]) +def test_managed_step_up_challenge_uses_only_the_gateway_metadata(base_url, request_path, scope, expected): + error = MCPUpstreamAuthError(status_code=403, www_authenticate=None, server_name="alias", required_scope=scope) + response = error.to_http_exception(base_url=base_url, request_path=request_path) + assert response.status_code == 403 + assert response.headers == {"www-authenticate": expected} diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index c1d64bf3d3a..f84557fb9d0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2203,6 +2203,31 @@ class TestMCPServerManager: assert exc_info.value.www_authenticate == challenge assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is True + @pytest.mark.asyncio + async def test_managed_scope_challenge_requires_reauthorization_without_retry(self): + server = MCPServer( + server_id="managed-scope", name="managed-scope", alias="managed-scope", + url="https://up.example.com/mcp", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + ) + upstream_error = self._upstream_status_error( + 403, 'Bearer error="insufficient_scope", scope="tools.write", resource_metadata="https://untrusted.test/metadata"' + ) + manager = MCPServerManager() + client = AsyncMock() + client.call_tool = AsyncMock(side_effect=upstream_error) + manager.create_mcp_client = AsyncMock(return_value=client) + + with pytest.raises(MCPUpstreamAuthError) as raised: + await self._run_call_regular(manager, server) + + response = raised.value.to_http_exception("https://gateway.test", "/managed-scope/mcp") + assert response.status_code == 403 + assert response.headers == { + "www-authenticate": 'Bearer resource_metadata="https://gateway.test/.well-known/oauth-protected-resource/managed-scope/mcp", error="insufficient_scope", scope="tools.write"' + } + assert client.call_tool.await_count == 1 + @pytest.mark.asyncio @pytest.mark.parametrize("is_error", [False, True]) async def test_call_passthrough_returns_tool_result_unchanged(self, is_error): @@ -4525,6 +4550,7 @@ class TestMCPServerManager: f"{issuer}/.well-known/oauth-authorization-server", { "issuer": issuer, + "response_types_supported": ["code"], "authorization_endpoint": "https://idp.example.com/authorize", "token_endpoint": "https://idp.example.com/token", "scopes_supported": ["read", "write"], @@ -4546,6 +4572,11 @@ class TestMCPServerManager: assert result.client_id_metadata_document_supported is True + assert result.authorization_server_metadata is not None + assert result.authorization_server_metadata["scopes_supported"] == ["read", "write"] + merged = manager._merge_discovered_oauth_metadata(self._passthrough_call_server(MCPAuth.oauth2), result) + assert merged.authorization_server_metadata == result.authorization_server_metadata + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_rejects_issuer_mismatch(self): """RFC 8414 §3.3 fail-closed: a document self-attesting a DIFFERENT issuer than the one it was @@ -19252,17 +19283,23 @@ def test_oauth_rebuild_retains_only_corroborated_cimd_capability(capability: boo transport=MCPTransport.http, auth_type=MCPAuth.oauth2, authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", client_id_metadata_document_supported=capability, + authorization_server_metadata={"issuer": "https://idp.example.com", "response_types_supported": ["code"], "scopes_supported": ["offline_access"]}, ) rebuilt: Final = previous.model_copy(update={ "client_id_metadata_document_supported": not capability if rebuild == "fresh_discovery" else None, "authorization_url": "https://changed.example.com/authorize" if rebuild == "repointed" else previous.authorization_url, "token_url": None, + "authorization_server_metadata": {"scopes_supported": []} if rebuild == "fresh_discovery" else None, "issuer": "https://idp.example.com" if rebuild == "anchored" else None, "issuer_is_anchored": rebuild == "anchored", }) carry_forward_resolved_oauth_endpoints(rebuilt, previous) expected: Final = not capability if rebuild == "fresh_discovery" else None if rebuild in ("repointed", "anchored") else capability assert rebuilt.client_id_metadata_document_supported is expected + assert rebuilt.authorization_server_metadata == ( + {"scopes_supported": []} if rebuild == "fresh_discovery" + else previous.authorization_server_metadata if rebuild == "same" else None + ) @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx new file mode 100644 index 00000000000..ef0e8c464b0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx @@ -0,0 +1,50 @@ +import { webcrypto } from "node:crypto"; +import { afterEach, expect, it, vi } from "vitest"; +import { render, screen, waitFor } from "@testing-library/react"; +import ConnectFlowBanner from "./ConnectFlowBanner"; + +const locationDescriptor = Object.getOwnPropertyDescriptor(window, "location")!; + +afterEach(() => { + vi.unstubAllGlobals(); + window.sessionStorage.clear(); + Object.defineProperty(window, "location", locationDescriptor); +}); + +it("carries a challenged scope through the connect button into vendor authorization", async () => { + vi.stubGlobal("crypto", webcrypto); + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ client_id: "registered-client" })), + ); + Object.defineProperty(window, "location", { + configurable: true, + value: { + href: "https://gateway.example/ui/connect?connect_flow=flow-handle", + origin: "https://gateway.example", + pathname: "/ui/connect", + }, + }); + render( + , + ); + await waitFor(() => expect(window.location.href).toContain("scope=tools.write")); + expect(new URL(window.location.href, window.location.origin).pathname).toBe( + "/v1/mcp/server/oauth/server-1/authorize", + ); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index 0d6e708f734..acdab1215b7 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -79,6 +79,7 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onC server={vendorServer} accessToken={accessToken} onConnect={onConnected} + scopes={flow?.requested_scopes} variant="button" autoStartKey={`litellm-mcp-autostart:${flowHandle}`} /> diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 1fee8923e94..3c50557501c 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -31,6 +31,7 @@ interface OAuth2ConnectButtonProps { onConnect: (serverId: string) => void; variant?: "badge" | "button"; autoStartKey?: string | null; + scopes?: string[]; } export const OAuth2ConnectButton: React.FC = ({ @@ -39,12 +40,14 @@ export const OAuth2ConnectButton: React.FC = ({ onConnect, variant = "badge", autoStartKey = null, + scopes, }) => { const name = server.server_name ?? server.alias ?? server.server_id; const { startOAuthFlow, status } = useUserMcpOAuthFlow({ accessToken, serverId: server.server_id, serverAlias: name, + scopes, onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]), }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 2b38da75411..a61570727ee 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4776,6 +4776,7 @@ export interface ConnectFlowStatus { server_id: string | null; server_name: string | null; connected: boolean | null; + requested_scopes?: string[]; } /** diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx new file mode 100644 index 00000000000..6cc2977c44c --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx @@ -0,0 +1,68 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import * as networking from "@/components/networking"; +import { setSecureItem } from "@/utils/secureStorage"; +import { useUserMcpOAuthFlow } from "./useUserMcpOAuthFlow"; + +vi.mock("@/components/networking", () => ({ + exchangeMcpOAuthToken: vi.fn(), + storeMCPOAuthUserCredential: vi.fn(), + registerMcpOAuthClient: vi.fn(), + buildMcpOAuthAuthorizeUrl: vi.fn(), + getProxyBaseUrl: () => "", + serverRootPath: "", +})); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), error: vi.fn() } })); + +describe("user OAuth granted scopes", () => { + beforeEach(() => { + vi.clearAllMocks(); + window.sessionStorage.clear(); + setSecureItem( + "litellm-user-mcp-oauth-flow-state", + JSON.stringify({ + state: "callback-state", + codeVerifier: "verifier", + serverId: "server-1", + redirectUri: "https://app.example/callback", + scopes: ["tools.write"], + }), + ); + setSecureItem( + "litellm-user-mcp-oauth-result", + JSON.stringify({ + state: "callback-state", + code: "provider-code", + }), + ); + }); + + it.each([ + ["tools.read tools.write", ["tools.read", "tools.write"]], + ["tools.read", ["tools.read"]], + ["", []], + [undefined, ["tools.write"]], + ])("persists the actual grant %s after reconnect", async (scope, expected) => { + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({ + access_token: "vendor-token", + refresh_token: "vendor-refresh", + scope, + }); + const onSuccess = vi.fn(); + const { result } = renderHook(() => + useUserMcpOAuthFlow({ + accessToken: "caller", + serverId: "server-1", + onSuccess, + }), + ); + await waitFor(() => expect(result.current.status).toBe("success")); + expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledExactlyOnceWith("caller", "server-1", { + access_token: "vendor-token", + refresh_token: "vendor-refresh", + expires_in: undefined, + scopes: expected, + }); + expect(onSuccess).toHaveBeenCalledOnce(); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index d7684a76f05..e4cac9dcf7e 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -216,7 +216,7 @@ export const useUserMcpOAuthFlow = ({ access_token: token.access_token, refresh_token: token.refresh_token, expires_in: token.expires_in, - scopes: flowState.scopes, + scopes: typeof token.scope === "string" ? token.scope.split(" ").filter(Boolean) : flowState.scopes, }); setStatus("success"); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f36456628e5..d8d13022f68 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -51565,6 +51565,7 @@ export interface operations { code_challenge_method?: string | null; response_type?: string | null; resource?: string | null; + scope?: string | null; }; header?: never; path?: never; From 804b5fb83c7cbb9e22acf475218cfb38b6e26d1e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:44:01 -0700 Subject: [PATCH 2/8] fix(mcp): relay scope challenges before streaming responses --- litellm/experimental_mcp_client/client.py | 6 +- litellm/proxy/_experimental/mcp_server/db.py | 6 +- .../mcp_server/discoverable_endpoints.py | 5 +- .../mcp_server/gateway_dcr_flow.py | 7 +- .../mcp_server/legacy_callbacks.py | 23 ++ .../mcp_server/mcp_server_manager.py | 10 + .../_experimental/mcp_server/operations.py | 12 +- .../mcp_server/scope_challenge.py | 93 ++++++++ .../proxy/_experimental/mcp_server/server.py | 22 +- tests/integration/mcp/test_mcp_oauth_flows.py | 4 +- .../mcp_server/test_db_credentials.py | 12 + .../mcp_server/test_discoverable_endpoints.py | 4 +- .../test_mcp_oauth_passthrough_tools.py | 212 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 41 ++++ .../test_mcp_server_tool_calls_and_headers.py | 41 +++- .../src/hooks/useUserMcpOAuthFlow.test.tsx | 27 ++- .../src/hooks/useUserMcpOAuthFlow.tsx | 5 +- 17 files changed, 487 insertions(+), 43 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/scope_challenge.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 729a6acb20d..8b33aefd663 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -79,6 +79,7 @@ from litellm.constants import ( ) from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration +from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response from litellm.proxy._experimental.mcp_server.result_conversion import ( age_freshness, @@ -832,7 +833,10 @@ class MCPClient: verify=ssl_config, follow_redirects=True, event_hooks=MappingProxyType( - {"response": [capture_upstream_error_response], "request": [guard] if guard else []} + { + "response": [record_upstream_tool_authorization, capture_upstream_error_response], + "request": [guard] if guard else [], + } ), ) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 823f0ed675c..8ab0d0118f4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1790,7 +1790,7 @@ async def store_user_oauth_credential( payload["refresh_token"] = refresh_token if expires_at: payload["expires_at"] = expires_at - if scopes: + if scopes is not None: payload["scopes"] = scopes # Guard against silently overwriting a BYOK credential with an OAuth token. @@ -2123,9 +2123,7 @@ async def refresh_user_oauth_token( new_refresh_token: Final[str | None] = body.get("refresh_token") or refresh_token raw_scope: Final = body.get("scope") - scopes: list[str] | None = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get( - "scopes" - ) + scopes: Final = raw_scope.split() if isinstance(raw_scope, str) else cred.get("scopes") await store_user_oauth_credential( prisma_client=prisma_client, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index b1cd855ceb1..48475d09be2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -693,7 +693,7 @@ async def _store_per_user_token_server_side( refresh_token: Final[str | None] = token_response.get("refresh_token") or None raw_scope: Final = token_response.get("scope", requested_scope) - scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) else None try: prisma_client: Final = get_prisma_client_or_throw("Database not connected. Cannot store per-user OAuth token.") @@ -1548,7 +1548,8 @@ async def exchange_token_with_server( result["expires_in"] = token_response["expires_in"] if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - effective_scope: Final = token_response.get("scope", requested_scope) + token_fields: Final = TypeAdapter[dict[str, object]](dict[str, object]).validate_python(token_response) + effective_scope: Final = token_fields.get("scope", requested_scope) if effective_scope is not None: result["scope"] = effective_scope diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f9909bc5d92..8b65d02e6a2 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -879,9 +879,12 @@ async def _describe_opened_flow( "server_id": server.server_id, "server_name": server.server_name or server.alias or server.name, "connected": credential == "present", - **({"requested_scopes": flow.requested_scopes} if flow.requested_scopes else {}), } - return interactive_description + return ( + {**interactive_description, "requested_scopes": flow.requested_scopes} + if flow.requested_scopes + else interactive_description + ) described: Final[ConnectFlowDescription] = { "state": state, "client_origin": _origin_only(flow.redirect_uri), diff --git a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py index 17bb84e066a..c14293c6135 100644 --- a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final, Protocol +import httpx2 from mcp.client.session import ClientRequestContext from mcp.types import ( CreateMessageRequestParams, @@ -14,6 +15,11 @@ from mcp.types import ( from litellm.constants import MCP_CLIENT_TIMEOUT from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx +from litellm.proxy._experimental.mcp_server.scope_challenge import ( + OAuthScopeResponse, + is_tool_call_request, + scope_response_from_request, +) from litellm.proxy._types import UserAPIKeyAuth @@ -96,3 +102,20 @@ def create_elicitation_callback(timeout: float | None = None) -> ElicitationCall ) return callback + + +def get_scope_response() -> OAuthScopeResponse | None: + context: Final = get_active_mcp_request_ctx() + return scope_response_from_request(context.request) if context is not None else None + + +async def record_upstream_tool_authorization(response: httpx2.Response) -> None: + pending: Final = get_scope_response() + if pending is None or not pending.pending or not response.is_success or response.request.method != "POST": + return + try: + is_tool_call: Final = is_tool_call_request(response.request.content) + except httpx2.RequestNotRead: + return + if is_tool_call: + pending.allow() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 87b0f907641..64c1d4ce086 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -100,6 +100,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( upstream_insufficient_scope, ) from litellm.proxy._experimental.mcp_server.interactions import bind_target +from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, @@ -6779,6 +6780,15 @@ class MCPServerManager: ) tasks.append(during_hook_task) + scope_response: Final = get_scope_response() + if scope_response is not None and not ( + mcp_server.is_gateway_managed_oauth2 + and mcp_server.needs_user_oauth_token + and not mcp_server.spec_path + and mcp_server.transport in (MCPTransport.http, MCPTransport.sse) + ): + scope_response.allow() + caller_oauth2_headers: Final = oauth2_headers oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2f0b8e64c9e..5087bc6df09 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -94,6 +94,7 @@ from litellm.proxy._experimental.mcp_server.interactions import ( seal_continuation, target_digest, ) +from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: F401 # legacy module exports MCPServerManager, _caller_authorization_fans_out, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -3009,11 +3010,12 @@ async def _execute_mcp_server_tool_call( is_error=True, ) except MCPUpstreamAuthError as e: - # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a - # mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST - # call path and the connect-time preemptive check do. Return an explicit isError - # naming the upstream status (at info level, not a traceback) so the client still - # learns it must re-authenticate upstream and expected pass-through 401s don't spam. + scope_response: Final = get_scope_response() + if scope_response is not None: + scope_response.deny(e) + # The request-owned response adapter relays managed scope challenges before + # Streamable HTTP headers commit. Other auth failures retain the existing + # explicit tool error, including failures after streaming has begun. verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code) return CallToolResult( content=[ diff --git a/litellm/proxy/_experimental/mcp_server/scope_challenge.py b/litellm/proxy/_experimental/mcp_server/scope_challenge.py new file mode 100644 index 00000000000..c326418bee7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/scope_challenge.py @@ -0,0 +1,93 @@ +"""Keep upstream scope challenges ahead of Streamable HTTP response headers.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter, ValidationError +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.types import Message, Send + +from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + +if TYPE_CHECKING: + from mcp.server.context import CallNext, HandlerResult, ServerRequestContext + +SCOPE_RESPONSE_KEY: Final = "litellm.mcp.scope_response" +_REQUEST: Final = TypeAdapter[dict[str, object]](dict[str, object]) + + +def is_tool_call_request(body: bytes) -> bool: + try: + return _REQUEST.validate_json(body).get("method") == "tools/call" + except ValidationError: + return False + + +class OAuthScopeResponse: + def __init__(self, send: Send, base_url: str, request_path: str | None) -> None: + self._send = send + self._base_url = base_url + self._request_path = request_path + self._ready = asyncio.Event() + self._lock = asyncio.Lock() + self._start: Message | None = None + self._challenge: MCPUpstreamAuthError | None = None + self._refused = False + + @property + def pending(self) -> bool: + return not self._ready.is_set() + + def allow(self) -> None: + self._ready.set() + + def deny(self, error: MCPUpstreamAuthError) -> None: + if self.pending and error.status_code == 403 and error.required_scope is not None: + self._challenge = error + self._ready.set() + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self._start = message + if message["status"] != 200: + self.allow() + return + await self._ready.wait() + async with self._lock: + if self._refused: + return + if self._start is not None: + start: Final = self._start + self._start = None + if self._challenge is not None: + error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) + response: Final = JSONResponse( + status_code=error.status_code, content={"detail": error.detail}, headers=error.headers + ) + self._refused = True + await self._send( + {"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers} + ) + await self._send({"type": "http.response.body", "body": bytes(response.body)}) + return + await self._send(start) + await self._send(message) + + +def scope_response_from_request(request: object) -> OAuthScopeResponse | None: + if not isinstance(request, Request): + return None + response: Final = request.scope.get(SCOPE_RESPONSE_KEY) + return response if isinstance(response, OAuthScopeResponse) else None + + +async def finish_scope_response(context: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult: + response: Final = scope_response_from_request(context.request) + try: + return await call_next(context) + finally: + if response is not None: + response.allow() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a30bf30e45f..3f925ba3fb8 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -72,6 +72,12 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( redact_mcp_resource_url, well_known_root_suffix, ) +from litellm.proxy._experimental.mcp_server.scope_challenge import ( + SCOPE_RESPONSE_KEY, + OAuthScopeResponse, + finish_scope_response, + is_tool_call_request, +) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ActingUser, GrantedToolsetIds, @@ -604,7 +610,7 @@ if MCP_AVAILABLE: name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) - server.middleware.append(GatewayVersionPolicy()) + server.middleware.extend((finish_scope_response, GatewayVersionPolicy())) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/sse/messages") @@ -2346,6 +2352,18 @@ if MCP_AVAILABLE: await rejection(scope, receive, send) return + scope_response: Final = ( + OAuthScopeResponse( + send, + get_request_base_url(StarletteRequest(scope)), + scope.get("_original_path") or scope.get("path"), + ) + if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES + else None + ) + if scope_response is not None: + scope[SCOPE_RESPONSE_KEY] = scope_response + use_stateful: Final = bool(session_id or is_initialize) target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless @@ -2463,7 +2481,7 @@ if MCP_AVAILABLE: touch_last_seen=(scope.get("method") or "").upper() != "DELETE", copy_existing_session_auth_context=is_initialize, ) - local_send = send + local_send = scope_response or send if use_stateful and is_initialize: local_send = _wrap_send_with_stateful_session_auth_context( local_send, diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 63561059489..331e3198ff8 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -649,6 +649,6 @@ def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: G result: Final = caller.call(f"{alias}-{operation}_op", {}, identity) assert result.ok and result.text == f"{operation} completed", result.raw forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity) - assert forbidden.status_code == 200 and not forbidden.ok, forbidden.raw + assert forbidden.status == 200 and not forbidden.ok, forbidden.raw unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity) - assert unauthorized.status_code == 401, unauthorized.raw + assert unauthorized.status == 401, unauthorized.raw diff --git a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py index 10993280e89..6b7c4350426 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1856,3 +1856,15 @@ async def test_saved_cimd_grant_refreshes_from_encrypted_storage_on_a_fresh_repl assert parse_qs(upstream.calls[0].request.content.decode())["client_id"] == [identity] assert "client_secret" not in parse_qs(upstream.calls[0].request.content.decode()) assert server.client_id is None and server.client_id_metadata_document_supported is capability + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_scope, expected", [(None, ["a"]), ("b", ["b"]), ("", [])]) +async def test_silent_refresh_preserves_omitted_scope_but_accepts_explicit_grant(monkeypatch, provider_scope, expected): + import litellm.proxy._experimental.mcp_server.db as db_mod + + await _run_refresh(monkeypatch, _refresh_server(), { + "access_token": "new-access", + **({"scope": provider_scope} if provider_scope is not None else {}), + }) + assert db_mod.store_user_oauth_credential.await_args.kwargs["scopes"] == expected diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 0086410fc26..bc4ee3a62fb 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -13811,7 +13811,7 @@ async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_id @pytest.mark.asyncio -@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", None)]) +@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])]) async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted( jwt_oauth_identity, monkeypatch, respx_mock, returned_scope, expected, ): @@ -13904,7 +13904,7 @@ async def test_managed_scope_code_cannot_cross_oauth_boundaries(jwt_oauth_identi @pytest.mark.asyncio -@pytest.mark.parametrize("scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", None)]) +@pytest.mark.parametrize("scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])]) async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch, respx_mock, scope, expected): from types import SimpleNamespace diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index bea4263268b..15b43ccfdee 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,6 +1,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import asyncio import logging import sys from unittest.mock import AsyncMock, MagicMock @@ -510,3 +511,214 @@ def test_managed_step_up_challenge_uses_only_the_gateway_metadata(base_url, requ response = error.to_http_exception(base_url=base_url, request_path=request_path) assert response.status_code == 403 assert response.headers == {"www-authenticate": expected} + + +@pytest.mark.asyncio +async def test_scope_challenge_replaces_pending_stream_headers_once(): + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send = AsyncMock() + response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") + await response({"type": "http.response.start", "status": 200, "headers": []}) + body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"})) + await asyncio.sleep(0) + send.assert_not_awaited() + response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + await asyncio.wait_for(body, 1) + await response({"type": "http.response.body", "body": b"private tool error"}) + assert send.await_count == 2 + start, result = [call.args[0] for call in send.await_args_list] + assert start["status"] == 403 + challenge = dict(start["headers"])[b"www-authenticate"].decode() + assert challenge == ('Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp", ' + 'error="insufficient_scope", scope="tools.write"') + assert result["body"] == b'{"detail":"Forbidden"}' + + +@pytest.mark.asyncio +async def test_scope_authorization_releases_progress_without_waiting_for_tool_result(): + import httpx2 + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import ( + SCOPE_RESPONSE_KEY, + OAuthScopeResponse, + ) + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + request = Request({"type": "http", SCOPE_RESPONSE_KEY: gate}) + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request) + token = active_mcp_request_ctx_var.set(context) + try: + await gate({"type": "http.response.start", "status": 200, "headers": []}) + progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) + await asyncio.sleep(0) + for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]: + await record_upstream_tool_authorization(httpx2.Response(status, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": method}))) + assert gate.pending + unread = httpx2.Request("POST", "https://upstream.example/mcp", content=iter([b'{"method":"tools/call"}'])) + await record_upstream_tool_authorization(httpx2.Response(200, request=unread)) + assert gate.pending + # Use the actual HTTP response hook contract, before reading the tool response body. + from litellm.experimental_mcp_client.client import MCPClient + + upstream = MCPClient(server_url="https://upstream.example/mcp") + async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client: + async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}): + await asyncio.wait_for(progress, 1) + assert send.await_args_list[-1].args[0]["body"] == b"progress" + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + await gate({"type": "http.response.body", "body": b"result"}) + assert send.await_args_list[0].args[0]["status"] == 200 + assert send.await_args_list[-1].args[0]["body"] == b"result" + finally: + active_mcp_request_ctx_var.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure): + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + with pytest.raises(failure): + await finish_scope_response(context, AsyncMock(side_effect=failure)) + await gate({"type": "http.response.start", "status": 200, "headers": []}) + await gate({"type": "http.response.body", "body": b"error"}) + assert not gate.pending + assert send.await_args_list[-1].args[0]["body"] == b"error" + + +@pytest.mark.asyncio +async def test_scope_response_preserves_transport_validation_errors(): + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + await gate({"type": "http.response.start", "status": 400, "headers": []}) + await gate({"type": "http.response.body", "body": b"invalid request"}) + assert send.await_args_list[0].args[0]["status"] == 400 + assert send.await_args_list[1].args[0]["body"] == b"invalid request" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("denied", [False, True]) +@pytest.mark.parametrize("stateless", [False, True]) +@pytest.mark.parametrize("padding_size", [0, 131072]) +async def test_streamable_sdk_relays_scope_challenge_before_committing_http_response(denied, stateless, padding_size): + import httpx2 + from mcp.server.lowlevel import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.types import CallToolResult, TextContent + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response + + async def call_tool(context, params): + token = active_mcp_request_ctx_var.set(context) + try: + gate = get_scope_response() + assert gate is not None + if denied: + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + return CallToolResult(content=[TextContent(type="text", text="done")]) + finally: + active_mcp_request_ctx_var.reset(token) + + server = Server("scope-test", on_call_tool=call_tool) + server.middleware.append(finish_scope_response) + manager = StreamableHTTPSessionManager(server, stateless=stateless) + + async def app(scope, receive, send): + response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + scope[SCOPE_RESPONSE_KEY] = response + await manager.handle_request(scope, receive, response) + + async with manager.run(): + async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client: + headers = {"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"} + if not stateless: + initialized = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers, json={ + "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": { + "protocolVersion": "2025-11-25", "capabilities": {}, "clientInfo": {"name": "scope-test" + "x" * padding_size, "version": "1"}, + }, + }), 3) + assert initialized.status_code == 200 + headers["MCP-Session-Id"] = initialized.headers["mcp-session-id"] + notified = await client.post("https://gateway.example/mcp/docs", headers=headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + assert notified.status_code == 202 + response = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "write", "arguments": {"padding": "x" * padding_size}}}), 3) + assert response.status_code == (403 if denied else 200) + if denied: + assert 'error="insufficient_scope"' in response.headers["www-authenticate"] + assert response.json() == {"detail": "Forbidden"} + else: + assert '"text":"done"' in response.text + assert "www-authenticate" not in response.headers + + +@pytest.mark.asyncio +async def test_scope_challenges_are_isolated_between_concurrent_requests(): + import httpx2 + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, record_upstream_tool_authorization + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + async def request_scope(denied): + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + token = active_mcp_request_ctx_var.set(context) + try: + await gate({"type": "http.response.start", "status": 200, "headers": []}) + await asyncio.sleep(0) + assert get_scope_response() is gate + if denied: + gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) + else: + await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"}))) + await gate({"type": "http.response.body", "body": b"done"}) + return send.await_args_list[0].args[0]["status"] + finally: + active_mcp_request_ctx_var.reset(token) + + assert await asyncio.gather(request_scope(True), request_scope(False)) == [403, 200] + assert get_scope_response() is None + + +@pytest.mark.parametrize("body", [b"", b"{", b"[]", b'{"method":"initialize"}']) +def test_non_tool_payloads_do_not_hold_response_headers(body): + from litellm.proxy._experimental.mcp_server.scope_challenge import is_tool_call_request, scope_response_from_request + + assert not is_tool_call_request(body) + assert scope_response_from_request(None) is None + + +@pytest.mark.asyncio +async def test_scope_middleware_preserves_non_http_requests(): + from mcp.server.context import ServerRequestContext + from mcp.types import CallToolResult, TextContent + + from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response + + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call") + async def call_next(request): + return CallToolResult(content=[TextContent(type="text", text=request.method)]) + + returned = await finish_scope_response(context, call_next) + assert returned.content[0].text == "tools/call" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index f84557fb9d0..2c9644f4ae0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7273,6 +7273,47 @@ class TestMCPServerManager: def _unrestricted_auth() -> UserAPIKeyAuth: return UserAPIKeyAuth(api_key="sk-test") + @pytest.mark.asyncio + @pytest.mark.parametrize("managed", [False, True]) + async def test_tool_response_waits_only_for_managed_user_oauth_authorization(self, managed): + import asyncio + + from mcp.server.context import ServerRequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + auth = self._unrestricted_auth() + manager, logging = self._manager_ready_for_call_tool( + [MCPTool(name="test_tool", inputSchema={"type": "object"})], + caller=ListedToolsCaller(user_api_key_auth=auth), + ) + if managed: + manager.registry["test-server"] = manager.registry["test-server"].model_copy(update={ + "auth_type": MCPAuth.oauth2, "oauth2_flow": "authorization_code", + }) + send = AsyncMock() + gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/test-server") + context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + token = active_mcp_request_ctx_var.set(context) + await gate({"type": "http.response.start", "status": 200, "headers": []}) + progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) + + async def upstream_call(*args, **kwargs): + await asyncio.sleep(0) + assert send.await_count == (0 if managed else 2) + return CallToolResult(content=[], is_error=False) + + manager.create_mcp_client.return_value.call_tool.side_effect = upstream_call + try: + result = await manager.call_tool(server_name="test-server", name="test_tool", arguments={}, user_api_key_auth=auth, proxy_logging_obj=logging) + assert not result.is_error + finally: + gate.allow() + await asyncio.wait_for(progress, 1) + active_mcp_request_ctx_var.reset(token) + @pytest.mark.asyncio async def test_call_tool_hands_listed_tool_description_and_schema_to_pre_call_hooks(self): schema = {"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]} diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 4a214617142..89484b8fa51 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -264,11 +264,17 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_reques @pytest.mark.asyncio -async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx): - """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session - tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an - upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked - 500 or a raw traceback, so the client still learns it must re-authenticate upstream.""" +@pytest.mark.parametrize("required_scope, http_transport", [(None, False), (None, True), ("tools.write", True)]) +async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx, required_scope, http_transport): + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + + status_code = 403 if required_scope is not None else 401 + send = AsyncMock() + response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") + request = Request({"type": "http", SCOPE_RESPONSE_KEY: response}) + await response({"type": "http.response.start", "status": 200, "headers": []}) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( @@ -284,7 +290,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r return data async def mock_call_mcp_tool(*args, **kwargs): - raise MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="pt") + raise MCPUpstreamAuthError(status_code=status_code, www_authenticate=None, server_name="pt", required_scope=required_scope) mock_logger = MagicMock() with patch( @@ -298,20 +304,27 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): with patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): result = await mcp_server_tool_call( - _mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"}) + _mcp_request_ctx(request=request if http_transport else None), _call_tool_params("test_tool", {"param": "value"}) ) assert result.is_error is True # The dedicated MCPUpstreamAuthError branch (not the generic Exception fallthrough) produces this # specific message and logs at info, never a traceback via verbose_logger.exception. assert "upstream authentication required" in result.content[0].text - assert "401" in result.content[0].text + assert str(status_code) in result.content[0].text exception_calls = [str(c.args[0]) for c in mock_logger.exception.call_args_list if c.args] assert not any("mcp_server_tool_call" in m for m in exception_calls), ( "must not log a traceback for the expected re-auth" ) info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args] assert any("Upstream auth failure" in m for m in info_calls) + response.allow() + await response({"type": "http.response.body", "body": b"tool error"}) + start = send.await_args_list[0].args[0] + assert start["status"] == (403 if required_scope is not None else 200) + if required_scope is not None: + assert b'error="insufficient_scope"' in dict(start["headers"])[b"www-authenticate"] + def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): @@ -2464,9 +2477,12 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): (("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)), ) async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( - method: str, chunked: bool, character: str, bytes_before_cap: int + method: str, chunked: bool, character: str, bytes_before_cap: int, _mcp_request_ctx ) -> None: + from starlette.requests import Request + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response params: Final = ( { @@ -2490,9 +2506,14 @@ async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( send: Final = AsyncMock() received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future() - async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None: + async def handle_request(incoming: Scope, downstream_receive: Receive, outgoing: Send) -> None: assert receive.await_count == (2 if chunked else 1) received.set_result(await _drain_body(downstream_receive)) + + async def dispatch(context): + return CallToolResult(content=[]) + + await finish_scope_response(_mcp_request_ctx(request=Request(incoming)), dispatch) await outgoing({"type": "http.response.start", "status": 200, "headers": []}) await outgoing({"type": "http.response.body", "body": b"{}"}) diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx index 6cc2977c44c..8ce8c20b160 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx @@ -18,16 +18,14 @@ describe("user OAuth granted scopes", () => { beforeEach(() => { vi.clearAllMocks(); window.sessionStorage.clear(); - setSecureItem( - "litellm-user-mcp-oauth-flow-state", - JSON.stringify({ - state: "callback-state", - codeVerifier: "verifier", - serverId: "server-1", - redirectUri: "https://app.example/callback", - scopes: ["tools.write"], - }), - ); + const flowState = { + state: "callback-state", + codeVerifier: "verifier", + serverId: "server-1", + redirectUri: "https://app.example/callback", + scopes: ["tools.write"], + }; + setSecureItem("litellm-user-mcp-oauth-flow-state", JSON.stringify(flowState)); setSecureItem( "litellm-user-mcp-oauth-result", JSON.stringify({ @@ -57,12 +55,17 @@ describe("user OAuth granted scopes", () => { }), ); await waitFor(() => expect(result.current.status).toBe("success")); - expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledExactlyOnceWith("caller", "server-1", { + const expectedCredential = { access_token: "vendor-token", refresh_token: "vendor-refresh", expires_in: undefined, scopes: expected, - }); + }; + expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledExactlyOnceWith( + "caller", + "server-1", + expectedCredential, + ); expect(onSuccess).toHaveBeenCalledOnce(); }); }); diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index e4cac9dcf7e..ed6c842b57d 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -68,6 +68,9 @@ const getStorage = (key: string): string | null => { return getSecureItem(key); }; +const grantedScopes = (scope: unknown, requested: string[] | undefined): string[] | undefined => + typeof scope === "string" ? scope.split(" ").filter(Boolean) : requested; + export const useUserMcpOAuthFlow = ({ accessToken, serverId, @@ -216,7 +219,7 @@ export const useUserMcpOAuthFlow = ({ access_token: token.access_token, refresh_token: token.refresh_token, expires_in: token.expires_in, - scopes: typeof token.scope === "string" ? token.scope.split(" ").filter(Boolean) : flowState.scopes, + scopes: grantedScopes(token.scope, flowState.scopes), }); setStatus("success"); From 9db4806ec1012b18ef702fc211f364e45bff3247 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:01:32 -0700 Subject: [PATCH 3/8] fix(mcp): bind aggregate OAuth step-up scopes to upstream consent --- .../mcp_server/discoverable_endpoints.py | 4 +- .../_experimental/mcp_server/exceptions.py | 22 +- .../mcp_server/gateway_dcr_flow.py | 110 ++++++-- .../mcp_server/mcp_server_manager.py | 1 + .../_experimental/mcp_server/oauth_utils.py | 31 +++ .../proxy/_experimental/mcp_server/server.py | 2 +- tests/integration/mcp/test_mcp_oauth_flows.py | 73 +++++- .../mcp_server/test_gateway_dcr_flow.py | 243 +++++++++++++++++- .../test_mcp_oauth_passthrough_tools.py | 19 +- .../test_mcp_server_tool_calls_and_headers.py | 14 +- .../ConnectFlowBanner.integration.test.tsx | 41 +++ .../src/components/chat/ConnectFlowBanner.tsx | 3 +- .../src/components/chat/MCPAppsPanel.tsx | 5 +- 13 files changed, 524 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 48475d09be2..b9a1a73549a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -692,8 +692,8 @@ async def _store_per_user_token_server_side( expires_in = None refresh_token: Final[str | None] = token_response.get("refresh_token") or None - raw_scope: Final = token_response.get("scope", requested_scope) - scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) else None + raw_scope: Final[object] = token_response.get("scope", requested_scope) + scopes: Final[list[str] | None] = raw_scope.split() if isinstance(raw_scope, str) else None try: prisma_client: Final = get_prisma_client_or_throw("Database not connected. Cannot store per-user OAuth token.") diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index a720006cebf..9554be3ee0d 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -4,6 +4,8 @@ from typing import Final from fastapi import HTTPException +from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + class MCPServerURLCredentialsError(HTTPException): """A fixed, sanitized URL-credential migration error safe for operator previews.""" @@ -38,11 +40,13 @@ class MCPUpstreamAuthError(Exception): www_authenticate: str | None, server_name: str, required_scope: str | None = None, + server_id: str | None = None, ) -> None: self.status_code = status_code self.www_authenticate = www_authenticate self.server_name = server_name self.required_scope = required_scope + self.server_id = server_id super().__init__(f"Upstream MCP server {server_name!r} returned {status_code}") def to_http_exception( @@ -75,17 +79,31 @@ class MCPUpstreamAuthError(Exception): the client originally targeted, matching the path-aware behaviour of ``get_passthrough_resource_metadata_url`` in ``oauth_utils.py``. """ + aggregate_step_up: Final = ( + self.required_scope is not None + and request_path is not None + and (request_path.rstrip("/") == "/mcp" or request_path.startswith("/toolset/")) + ) + requested_scope: Final = ( + " ".join( + encode_upstream_scope(self.server_id or self.server_name, item) for item in self.required_scope.split() + ) + if aggregate_step_up and self.required_scope + else self.required_scope + ) challenge: str | None = self.www_authenticate if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url: prefix: Final = base_url.rstrip("/") - if request_path and request_path.startswith(f"/{self.server_name}/mcp"): + if aggregate_step_up: + resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource" + elif request_path and request_path.startswith(f"/{self.server_name}/mcp"): resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp" else: resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}" challenge = f'Bearer resource_metadata="{resource_metadata_url}"' scoped_challenge: Final = ( (f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"') - + (f', scope="{self.required_scope}"' if self.required_scope else "") + + (f', scope="{requested_scope}"' if requested_scope else "") if self.required_scope is not None else challenge ) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 8b65d02e6a2..ef1019a0ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -59,8 +59,11 @@ from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, + UPSTREAM_SCOPE_PREFIX, + UpstreamScope, canonical_resource_uri, canonicalize_url_identity, + decode_upstream_scope, get_request_base_url, is_loopback_redirect_host, validate_redirect_uri_shape, @@ -315,6 +318,7 @@ class _ConnectFlow(LiteLLMBaseModel): resource_server_id: str | None = None audience: SessionAudience | None = None requested_scopes: tuple[str, ...] = () + upstream_scopes: tuple[UpstreamScope, ...] = () class _GatewayAuthCode(LiteLLMBaseModel): @@ -586,6 +590,13 @@ def aggregate_authorize( if session_user_id is None: return _login_redirect(base_url, request) scoped_server: Final = resolve_scoped_resource_server(request, resource) + scope_tokens: Final = tuple(token for token in (scope or "").split() if token.startswith(UPSTREAM_SCOPE_PREFIX)) + decoded_scopes: Final = tuple(decode_upstream_scope(token) for token in scope_tokens) + if any(item is None for item in decoded_scopes): + return _oauth_error(400, "invalid_scope", "invalid upstream permission scope") + upstream_scopes: Final = tuple(item for item in decoded_scopes if item is not None) + if scoped_server is not None and any(item.server_id != scoped_server.server_id for item in upstream_scopes): + return _oauth_error(400, "invalid_scope", "upstream permission does not match the requested resource") handle: Final = secrets.token_urlsafe(24) flow: Final = _new_connect_flow( session_user_id=session_user_id, @@ -595,7 +606,8 @@ def aggregate_authorize( code_challenge=code_challenge or "", resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, - requested_scopes=tuple((scope or "").split()) if scoped_server is not None else (), + requested_scopes=tuple((scope or "").split()) if scoped_server is not None and not upstream_scopes else (), + upstream_scopes=upstream_scopes, ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) @@ -756,6 +768,7 @@ def _new_connect_flow( resource_server_id: str | None, audience: SessionAudience | None, requested_scopes: tuple[str, ...] = (), + upstream_scopes: tuple[UpstreamScope, ...] = (), ) -> _ConnectFlow: now: Final = datetime.now(timezone.utc) return _ConnectFlow( @@ -769,6 +782,7 @@ def _new_connect_flow( resource_server_id=resource_server_id, audience=audience, requested_scopes=requested_scopes, + upstream_scopes=upstream_scopes, ) @@ -824,20 +838,20 @@ def _open_flow_for( @catalog_operation(global_manager) async def _flow_target( - flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability + user_id: str, server_id: str | None, lookup_server_reachability: LookupServerReachability ) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: - if flow.resource_server_id is None: + if server_id is None: return "unscoped", None from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle MCPServerManager, global_mcp_server_manager, ) - server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) if ( server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server) - or not await lookup_server_reachability(flow.user_id, server.server_id) + or not await lookup_server_reachability(user_id, server.server_id) ): return "stale", None state: Final = ( @@ -864,27 +878,11 @@ async def _describe_opened_flow( lookup_vendor_credential: LookupVendorCredential, lookup_server_reachability: LookupServerReachability, ) -> ConnectFlowDescription | Response: - state, server = await _flow_target(flow, lookup_server_reachability) + if flow.upstream_scopes: + return await _describe_upstream_scope_flow(flow, lookup_vendor_credential, lookup_server_reachability) + state, server = await _flow_target(flow.user_id, flow.resource_server_id, lookup_server_reachability) if state == "interactive" and server is not None: - credential: Final = ( - await lookup_vendor_credential(flow.user_id, server.server_id, flow.requested_scopes) - if flow.requested_scopes - else await lookup_vendor_credential(flow.user_id, server.server_id) - ) - if credential == "unavailable": - return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) - interactive_description: Final[ConnectFlowDescription] = { - "state": state, - "client_origin": _origin_only(flow.redirect_uri), - "server_id": server.server_id, - "server_name": server.server_name or server.alias or server.name, - "connected": credential == "present", - } - return ( - {**interactive_description, "requested_scopes": flow.requested_scopes} - if flow.requested_scopes - else interactive_description - ) + return await _describe_interactive_target(flow, server, flow.requested_scopes, lookup_vendor_credential) described: Final[ConnectFlowDescription] = { "state": state, "client_origin": _origin_only(flow.redirect_uri), @@ -895,6 +893,68 @@ async def _describe_opened_flow( return described +async def _describe_interactive_target( + flow: _ConnectFlow, + server: MCPServer, + scopes: tuple[str, ...], + lookup_vendor_credential: LookupVendorCredential, +) -> ConnectFlowDescription | Response: + credential: Final = ( + await lookup_vendor_credential(flow.user_id, server.server_id, scopes) + if scopes + else await lookup_vendor_credential(flow.user_id, server.server_id) + ) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + described: Final[ConnectFlowDescription] = { + "state": "interactive", + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return {**described, "requested_scopes": scopes} if scopes else described + + +async def _describe_upstream_scope_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + server_ids: Final = tuple(dict.fromkeys(item.server_id for item in flow.upstream_scopes)) + targets: Final = tuple( + [await _flow_target(flow.user_id, server_id, lookup_server_reachability) for server_id in server_ids] + ) + if any(state != "interactive" or server is None for state, server in targets): + return _oauth_error(400, "invalid_scope", "an upstream permission target is unavailable") + descriptions: Final = tuple( + [ + await _describe_interactive_target( + flow, + server, + tuple( + dict.fromkeys( + item.permission for item in flow.upstream_scopes if item.server_id == server.server_id + ) + ), + lookup_vendor_credential, + ) + for _, server in targets + if server is not None + ] + ) + for description in descriptions: + if isinstance(description, Response) or not description["connected"]: + return description + return { + "state": "interactive", + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None, + "server_name": None, + "connected": True, + } + + async def describe_connect_flow( request: Request, flow_handle: str, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 64c1d4ce086..975bbb7a835 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6433,6 +6433,7 @@ class MCPServerManager: www_authenticate=None, server_name=mcp_server.alias or server_label, required_scope=required_scope, + server_id=mcp_server.server_id, ) from e auth_info: Final = _extract_upstream_auth_failure(e) if not relays_upstream_auth or auth_info is None or auth_info[0] != 401: diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 276cf250b30..b8675f749e6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -2,6 +2,8 @@ (BYOK + discoverable / pass-through OAuth proxy).""" import os +import re +from dataclasses import dataclass from ipaddress import ip_address from typing import TYPE_CHECKING, Final, NoReturn from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit @@ -819,3 +821,32 @@ def build_upstream_oauth2_token_request( if not resource: return client_auth return TokenEndpointClientAuth(headers=client_auth.headers, body={**client_auth.body, "resource": resource}) + + +UPSTREAM_SCOPE_PREFIX: Final = "litellm:mcp_scope:" + + +@dataclass(frozen=True, slots=True) +class UpstreamScope: + server_id: str + permission: str + + +def encode_upstream_scope(server_id: str, permission: str) -> str: + return f"{UPSTREAM_SCOPE_PREFIX}{server_id.encode().hex()}:{permission.encode().hex()}" + + +def decode_upstream_scope(token: str) -> UpstreamScope | None: + encoded: Final = token.removeprefix(UPSTREAM_SCOPE_PREFIX).split(":") + if not token.startswith(UPSTREAM_SCOPE_PREFIX) or len(encoded) != 2: + return None + try: + server_id: Final = bytes.fromhex(encoded[0]).decode() + permission: Final = bytes.fromhex(encoded[1]).decode() + except (ValueError, UnicodeError): + return None + if not server_id or re.fullmatch(r"[!#-\[\]-~]+", permission) is None: + return None + if encode_upstream_scope(server_id, permission) != token: + return None + return UpstreamScope(server_id, permission) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3f925ba3fb8..a13a54c4912 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2356,7 +2356,7 @@ if MCP_AVAILABLE: OAuthScopeResponse( send, get_request_base_url(StarletteRequest(scope)), - scope.get("_original_path") or scope.get("path"), + get_route_relative_request_path(scope), ) if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES else None diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 331e3198ff8..0bd299f3bbd 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -24,14 +24,14 @@ from integration._support.mcp import ( EntryPoint, McpCaller, McpPeer, - ScriptedTool, - scripted_peer, - text_result, Outcome, + ScriptedTool, _outcome_from_rpc, call_tool, mcp_peer, register_mcp, + scripted_peer, + text_result, tool_calls, ) from integration._support.mcp_grants import create_toolset @@ -652,3 +652,70 @@ def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: G assert forbidden.status == 200 and not forbidden.ok, forbidden.raw unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity) assert unauthorized.status == 401, unauthorized.raw + + +@pytest.mark.parametrize("route", ["/mcp", "/toolset/{toolset}/mcp"]) +def test_managed_step_up_advertises_sdk_compatible_metadata_and_bound_permissions(gateway: Gateway, route: str) -> None: + from mcp.shared.auth_utils import check_resource_allowed + + with ( + oauth_server(scopes=("tools.read", "tools.write")) as auth, + _scope_enforcing_peer(auth) as (peer, wire), + gateway.scenario() as scenario, + ): + alias: Final = "scope" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]}, + ) + toolset_name: Final = "scope" + uuid.uuid4().hex[:8] + toolset_id: Final = create_toolset( + scenario, ((identity, "read_op"), (identity, "write_op")), toolset_name=toolset_name + ) + owner: Final = scenario.key( + user_id=scenario.user(), object_permission={"mcp_servers": [identity], "mcp_toolsets": [toolset_id]} + ) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read") + first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce) + assert first.status_code == 200, first.text + path: Final = route.format(toolset=toolset_name) + headers: Final = {**ACCEPT, "Authorization": f"Bearer {owner}"} + initialized: Final = gateway.client.post( + path, headers=headers, json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)} + ) + assert initialized.status_code == 200, initialized.text + peer.drain() + wire.drain() + response: Final = gateway.client.post( + path, + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": f"{alias}-write_op", "arguments": {}}, + }, + ) + assert response.status_code == 403, response.text + challenge: Final = response.headers["www-authenticate"] + metadata_url: Final = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0] + metadata: Final = gateway.client.get(metadata_url) + assert metadata.status_code == 200, metadata.text + assert check_resource_allowed( + requested_resource=_base(gateway) + path, configured_resource=metadata.json()["resource"] + ) + encoded_scope: Final = challenge.split('scope="', 1)[1].split('"', 1)[0] + assert encoded_scope == f"litellm:mcp_scope:{identity.encode().hex()}:{b'tools.write'.hex()}" + assert tool_calls(peer.drain()) == (), "denied operation executed upstream" + attempts: Final = tuple( + request + for request in wire.drain() + if request.body and json.loads(request.body).get("method") == "tools/call" + ) + assert len(attempts) == 1, "denied operation was retried" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index ed24a56eba5..de4e0dc8723 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -1915,7 +1915,9 @@ def _redis_that(async_increment, get=None): cache.redis_cache.async_increment = async_increment cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) - cache.redis_cache.async_get_cache = AsyncMock(side_effect=AssertionError("peek must read the client, not the wrapper")) + cache.redis_cache.async_get_cache = AsyncMock( + side_effect=AssertionError("peek must read the client, not the wrapper") + ) cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) return cache @@ -2069,7 +2071,11 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) rotated = json.loads( - (await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)))).body + ( + await _refresh_native( + payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) + ) + ).body )["refresh_token"] minter = _Minter() @@ -2601,3 +2607,236 @@ async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it(): complete = await _complete_page(response, scoped_server=server, vendor=vendor) assert complete.status_code == 303 assert parse_qs(urlparse(complete.headers["location"]).query)["code"][0].startswith(GATEWAY_AUTH_CODE_PREFIX) + + +@pytest.mark.asyncio +async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_gateway_access(): + from unittest.mock import AsyncMock + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + encoded = f"litellm:mcp_scope:{server.server_id.encode().hex()}:{b'tools.write'.hex()}" + response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=encoded) + vendor = AsyncMock(return_value="absent") + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["connected"] is False + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + premature = await _complete_page(response, scoped_server=server, vendor=vendor) + assert premature.status_code == 400 + vendor.return_value = "present" + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + tokens = await _redeem(code, client_id) + assert tokens.status_code == 200 + assert _opened_principal(json.loads(tokens.body)).resource_server_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "scope", + [ + "litellm:mcp_scope:", + "litellm:mcp_scope:aa:bb", + "litellm:mcp_scope:6769746875622d6964:20", + "litellm:mcp_scope:6964:5A", + ], +) +async def test_aggregate_step_up_rejects_malformed_permissions(scope): + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=scope) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_scope" + + +@pytest.mark.asyncio +async def test_step_up_validates_every_target_before_credential_reads_and_preserves_each_grant(): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + first, second = _scoped_mcp_server("first"), _scoped_mcp_server("second") + scopes = " ".join( + [encode_upstream_scope(first.server_id, "first.write"), encode_upstream_scope(second.server_id, "second.write")] + ) + response = _scoped_authorize(client_id, "https://llm.example.com", scope=scopes) + handle, cookies = _flow_cookie_from(response) + vendor = AsyncMock(return_value="absent") + reachable = AsyncMock(side_effect=[True, False]) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.side_effect = lambda identity: { + first.server_id: first, + second.server_id: second, + }.get(identity) + refused = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable) + assert refused.status_code == 400 + vendor.assert_not_awaited() + reachable.side_effect = None + reachable.return_value = True + vendor.side_effect = ["present", "absent"] + pending = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable) + assert json.loads(pending.body)["server_id"] == second.server_id + assert json.loads(pending.body)["requested_scopes"] == ["second.write"] + assert [call.args for call in vendor.await_args_list] == [ + ("u1", first.server_id, ("first.write",)), + ("u1", second.server_id, ("second.write",)), + ] + vendor.reset_mock() + wrong_user = await describe_connect_flow(_request(cookies=cookies), handle, "u2", vendor, reachable) + assert wrong_user.status_code == 403 + vendor.assert_not_awaited() + vendor.side_effect = None + vendor.return_value = "present" + complete = await complete_connect_flow( + _request(cookies=cookies, method="POST"), + handle, + "u1", + DualCache(), + lookup_vendor_credential=vendor, + lookup_server_reachability=reachable, + ) + assert complete.status_code == 303 + assert vendor.await_count == 2 + + +@pytest.mark.asyncio +async def test_named_resource_rejects_another_upstreams_permission_before_consent(): + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope=encode_upstream_scope("other-id", "write")) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_scope" + + +@pytest.mark.parametrize( + "identity, permission", + [ + ("server:id/one", "tools.write"), + ("snowman-☃", "https://permissions.example/read"), + ("id", "!#$%&'()*+,-./:;<=>?@[]^_`{|}~"), + ], +) +def test_upstream_permission_scope_round_trip_preserves_identity_and_permission(identity, permission): + from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope, encode_upstream_scope + + encoded = encode_upstream_scope(identity, permission) + decoded = decode_upstream_scope(encoded) + assert decoded is not None + assert (decoded.server_id, decoded.permission) == (identity, permission) + assert " " not in encoded + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state, expected", [("absent", 400), ("unavailable", 503)]) +async def test_upstream_step_up_cannot_finish_with_an_unproven_grant(state, expected): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + response = _scoped_authorize( + client_id, "https://llm.example.com/mcp", scope=encode_upstream_scope(server.server_id, "tools.write") + ) + vendor = AsyncMock(return_value=state) + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + assert completed.status_code == expected + assert "location" not in completed.headers + vendor.return_value = "present" + retried = await _complete_page(response, scoped_server=server, vendor=vendor) + assert retried.status_code == 303 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/mcp", "/toolset/analysis/mcp"]) +async def test_sdk_completes_aggregate_step_up_discovery_consent_and_pkce(route): + from unittest.mock import AsyncMock + + import httpx2 + from mcp.client.auth import OAuthClientProvider + from mcp.shared.auth import AuthorizationCodeResult, OAuthClientMetadata + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_authorization_server_response, + oauth_protected_resource_root, + ) + from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + + server = _scoped_mcp_server() + storage = AsyncMock() + storage.get_tokens.return_value = None + storage.get_client_info.return_value = None + callback = AsyncMock() + vendor = AsyncMock(return_value="absent") + cache = DualCache() + + async def consent(url): + params = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()} + response = aggregate_authorize(request=_request(), session_user_id="u1", **params) + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body)["requested_scopes"] == ["tools.write"] + assert json.loads(described.body)["connected"] is False + premature = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert premature.status_code == 400 + vendor.return_value = "present" + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + returned = parse_qs(urlparse(completed.headers["location"]).query) + callback.return_value = AuthorizationCodeResult(code=returned["code"][0], state=returned["state"][0]) + + provider = OAuthClientProvider( + server_url=f"https://llm.example.com{route}", + client_metadata=OAuthClientMetadata(redirect_uris=[REDIRECT_URI], token_endpoint_auth_method="none"), + storage=storage, + redirect_handler=consent, + callback_handler=callback, + ) + flow = provider.async_auth_flow(httpx2.Request("POST", f"https://llm.example.com{route}")) + initial = await anext(flow) + challenge = MCPUpstreamAuthError( + status_code=403, + www_authenticate=None, + server_name=server.server_name, + server_id=server.server_id, + required_scope="tools.write", + ).to_http_exception(base_url="https://llm.example.com", request_path=route) + request = await flow.asend(httpx2.Response(challenge.status_code, headers=challenge.headers, request=initial)) + assert request.url.path == "/.well-known/oauth-protected-resource" + request = await flow.asend(httpx2.Response(200, json=oauth_protected_resource_root(_request()), request=request)) + assert request.url.path == "/.well-known/oauth-authorization-server/mcp" + request = await flow.asend( + httpx2.Response(200, json=_build_aggregate_authorization_server_response(_request(), False), request=request) + ) + assert request.url.path == "/register" + registered = await register_aggregate_client( + _request("/register", method="POST"), json.loads(request.content), token_exchange_available=False + ) + request = await flow.asend(httpx2.Response(registered.status_code, content=registered.body, request=request)) + assert request.url.path == "/token" + params = {key: values[0] for key, values in parse_qs(request.content.decode()).items()} + tokens = await _redeem(cache=cache, **params) + assert tokens.status_code == 200 + payload = json.loads(tokens.body) + assert _opened_principal(payload).resource_server_id is None + retry = await flow.asend(httpx2.Response(200, json=payload, request=request)) + assert retry.url == initial.url + assert retry.headers["authorization"] == f"Bearer {payload['access_token']}" + with pytest.raises(StopAsyncIteration): + await flow.asend(httpx2.Response(200, json={"result": "write completed"}, request=retry)) + callback.assert_awaited_once() + storage.set_tokens.assert_awaited_once() + + +@pytest.mark.parametrize( + "scope", + ["tools.write", "litellm:mcp_scope::7772697465", "litellm:mcp_scope:6964:5A", "litellm:mcp_scope:6964:7772 697465"], +) +def test_upstream_permission_rejects_noncanonical_tokens(scope): + from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope + + assert decode_upstream_scope(scope) is None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 15b43ccfdee..a937d85929a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,4 +1,5 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations + """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" import asyncio @@ -463,7 +464,6 @@ async def test_fetch_tools_logs_upstream_request_details_on_500(caplog): assert "upstream-token-0123456789" not in caplog.text - @pytest.mark.asyncio async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, caplog): manager = MCPServerManager() @@ -722,3 +722,20 @@ async def test_scope_middleware_preserves_non_http_requests(): returned = await finish_scope_response(context, call_next) assert returned.content[0].text == "tools/call" + + +@pytest.mark.parametrize("path", ["/mcp", "/toolset/demo/mcp"]) +def test_managed_step_up_metadata_matches_the_sdk_configured_resource(path): + from mcp.shared.auth_utils import check_resource_allowed + from urllib.parse import urlsplit + + error = MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write") + response = error.to_http_exception("https://gateway.test", path) + challenge = response.headers["www-authenticate"] + metadata = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0] + metadata_path = urlsplit(metadata).path + resource_path = metadata_path.removeprefix("/.well-known/oauth-protected-resource") + assert check_resource_allowed( + requested_resource=f"https://gateway.test{path}", + configured_resource=f"https://gateway.test{resource_path}", + ), "the SDK must accept the advertised resource before it can reauthorize" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 89484b8fa51..1c76d7ddb87 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -265,7 +265,9 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_reques @pytest.mark.asyncio @pytest.mark.parametrize("required_scope, http_transport", [(None, False), (None, True), ("tools.write", True)]) -async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx, required_scope, http_transport): +async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( + _mcp_request_ctx, required_scope, http_transport +): from starlette.requests import Request from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse @@ -273,7 +275,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r status_code = 403 if required_scope is not None else 401 send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") - request = Request({"type": "http", SCOPE_RESPONSE_KEY: response}) + request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response}) await response({"type": "http.response.start", "status": 200, "headers": []}) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError @@ -290,7 +292,9 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r return data async def mock_call_mcp_tool(*args, **kwargs): - raise MCPUpstreamAuthError(status_code=status_code, www_authenticate=None, server_name="pt", required_scope=required_scope) + raise MCPUpstreamAuthError( + status_code=status_code, www_authenticate=None, server_name="pt", required_scope=required_scope + ) mock_logger = MagicMock() with patch( @@ -304,7 +308,8 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): with patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): result = await mcp_server_tool_call( - _mcp_request_ctx(request=request if http_transport else None), _call_tool_params("test_tool", {"param": "value"}) + _mcp_request_ctx(request=request if http_transport else None), + _call_tool_params("test_tool", {"param": "value"}), ) assert result.is_error is True @@ -326,7 +331,6 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r assert b'error="insufficient_scope"' in dict(start["headers"])[b"www-authenticate"] - def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): try: from litellm.proxy._experimental.mcp_server.server import ( diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx index ef0e8c464b0..51647247f75 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx @@ -48,3 +48,44 @@ it("carries a challenged scope through the connect button into vendor authorizat expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); }); + +it("starts consent for the next upstream in the same connect flow", async () => { + vi.stubGlobal("crypto", webcrypto); + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ client_id: "registered-client" })), + ); + Object.defineProperty(window, "location", { + configurable: true, + value: { + href: "https://gateway.example/ui/connect?connect_flow=flow-handle", + origin: "https://gateway.example", + pathname: "/ui/connect", + }, + }); + const props = { + flowHandle: "flow-handle", + accessToken: "caller", + onConnected: vi.fn(), + failed: false, + }; + const first = { + state: "interactive" as const, + client_origin: "https://client.example", + server_id: "first", + server_name: "First", + connected: false, + requested_scopes: ["first.write"], + }; + const { rerender } = render(); + await waitFor(() => expect(window.location.href).toContain("scope=first.write")); + rerender( + , + ); + await waitFor(() => expect(window.location.href).toContain("scope=second.write")); + expect(new URL(window.location.href, window.location.origin).pathname).toBe("/v1/mcp/server/oauth/second/authorize"); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index acdab1215b7..e32ae028097 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -76,12 +76,13 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onC {vendorServer !== null && ( )} diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 3c50557501c..d8386afeb93 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -43,13 +43,14 @@ export const OAuth2ConnectButton: React.FC = ({ scopes, }) => { const name = server.server_name ?? server.alias ?? server.server_id; - const { startOAuthFlow, status } = useUserMcpOAuthFlow({ + const flowOptions = { accessToken, serverId: server.server_id, serverAlias: name, scopes, onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]), - }); + }; + const { startOAuthFlow, status } = useUserMcpOAuthFlow(flowOptions); useEffect(() => { if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return; From f9eea5890fc37a058b8d2f973c18ca2ec3c5c585 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:39:21 -0700 Subject: [PATCH 4/8] fix(mcp): describe gateway access accurately before final consent --- .../mcp_server/gateway_dcr_flow.py | 6 ++-- .../mcp_server/test_gateway_dcr_flow.py | 29 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index ef1019a0ba3..a895dfed162 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -946,12 +946,14 @@ async def _describe_upstream_scope_flow( for description in descriptions: if isinstance(description, Response) or not description["connected"]: return description + if flow.resource_server_id is not None: + return descriptions[0] return { - "state": "interactive", + "state": "unscoped", "client_origin": _origin_only(flow.redirect_uri), "server_id": None, "server_name": None, - "connected": True, + "connected": None, } diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index de4e0dc8723..c79b2efde5e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -2624,6 +2624,9 @@ async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_g premature = await _complete_page(response, scoped_server=server, vendor=vendor) assert premature.status_code == 400 vendor.return_value = "present" + ready = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(ready.body)["state"] == "unscoped" + assert json.loads(ready.body)["server_id"] is None completed = await _complete_page(response, scoped_server=server, vendor=vendor) code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] tokens = await _redeem(code, client_id) @@ -2840,3 +2843,29 @@ def test_upstream_permission_rejects_noncanonical_tokens(scope): from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope assert decode_upstream_scope(scope) is None + + +@pytest.mark.asyncio +async def test_named_step_up_finish_describes_and_retains_its_single_server_ceiling(): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize( + client_id, SCOPED_RESOURCE, scope=encode_upstream_scope(server.server_id, "tools.write") + ) + vendor = AsyncMock(return_value="present") + ready = await _describe_page(response, scoped_server=server, vendor=vendor) + payload = json.loads(ready.body) + assert payload["state"] == "interactive" + assert payload["server_id"] == server.server_id + assert payload["connected"] is True + completed = await _complete_page(response, scoped_server=server, vendor=vendor) + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + tokens = await _redeem(code, client_id) + assert tokens.status_code == 200 + assert _opened_principal(json.loads(tokens.body)).resource_server_id == server.server_id From 29c1a62fac60f28fe2ed8e116dfd943a6553c2a0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:09:15 -0700 Subject: [PATCH 5/8] fix(mcp): preserve standalone client optional dependencies --- litellm/experimental_mcp_client/client.py | 12 ++++++-- .../test_mcp_client.py | 29 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 8b33aefd663..fab08d7bc27 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -79,7 +79,7 @@ from litellm.constants import ( ) from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration -from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization +from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response from litellm.proxy._experimental.mcp_server.result_conversion import ( age_freshness, @@ -106,6 +106,14 @@ if TYPE_CHECKING: from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback +def _mcp_response_hooks() -> tuple[Callable[[httpx2.Response], Awaitable[None]], ...]: + if get_active_mcp_request_ctx() is None: + return (capture_upstream_error_response,) + from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization + + return (record_upstream_tool_authorization, capture_upstream_error_response) + + def to_basic_auth(auth_value: str) -> str: """Convert auth value to Basic Auth format.""" return base64.b64encode(auth_value.encode("utf-8")).decode() @@ -834,7 +842,7 @@ class MCPClient: follow_redirects=True, event_hooks=MappingProxyType( { - "response": [record_upstream_tool_authorization, capture_upstream_error_response], + "response": list(_mcp_response_hooks()), "request": [guard] if guard else [], } ), diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index ff192c9279c..645ad761664 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -2731,6 +2731,35 @@ async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, mon ) +def test_standalone_client_request_does_not_require_proxy_dependencies(): + import textwrap + + from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + + script = textwrap.dedent(""" + import asyncio + import sys + for name in ("fastapi", "orjson", "email_validator", "annotated_doc"): + sys.modules[name] = None + import httpx2 + from litellm.experimental_mcp_client.client import MCPClient + + async def request(): + upstream = MCPClient(server_url="https://upstream.example/mcp") + transport = httpx2.MockTransport(lambda request: httpx2.Response(200)) + async with upstream._create_httpx_client_factory(transport=transport)() as client: + response = await client.post(upstream.server_url, json={"method": "tools/call"}) + assert response.status_code == 200 + assert all(sys.modules[name] is None for name in ("fastapi", "orjson", "email_validator", "annotated_doc")) + + asyncio.run(request()) + print("standalone request succeeded") + """) + result = run_child_interpreter(script, timeout=60) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "standalone request succeeded" + + def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): import subprocess From 157657a98d8fe04f6ffb0fdfe8a819130108e82e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:12:46 -0700 Subject: [PATCH 6/8] fix(mcp): retain explicit empty grants during native refresh --- .../authz_code_refresher.py | 7 ++++--- .../test_authz_code_refresher.py | 20 ++++++++++--------- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index b926ca8a6a4..a653ef5e5ff 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -65,7 +65,7 @@ def _parse_expires_in(raw: object) -> int | None: def _parse_scopes(raw: object) -> tuple[str, ...] | None: - return tuple(raw.split()) if isinstance(raw, str) and raw else None + return tuple(raw.split()) if isinstance(raw, str) else None class AuthorizationCodeRefresher: @@ -156,7 +156,8 @@ class AuthorizationCodeRefresher: rotated: Final = body.get("refresh_token") new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token expires_in: Final = _parse_expires_in(body.get("expires_in")) - scopes: Final = _parse_scopes(body.get("scope")) or token.scopes + granted_scopes: Final = _parse_scopes(body.get("scope")) + scopes: Final = granted_scopes if granted_scopes is not None else token.scopes await self._persist( user_id, @@ -164,7 +165,7 @@ class AuthorizationCodeRefresher: access_token, new_refresh, expires_in, - scopes or None, + scopes, **({"identity_binding_proof": binding_proof} if binding_proof is not None else {}), **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), ) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index 350e7f1f910..3068e5a0116 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -210,16 +210,17 @@ async def test_unrotated_refresh_token_is_carried_forward(): @pytest.mark.asyncio -async def test_unrecorded_scope_is_carried_forward(): +@pytest.mark.parametrize("prior_scopes", [(), ("read", "write")]) +async def test_unrecorded_scope_is_carried_forward(prior_scopes): persisted = [] refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) token = await refresher.refresh( - "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) + "a", "s", OAuthToken("old", refresh_token="rt", scopes=prior_scopes) ) assert token is not None # response omitted "scope" -> the user's recorded grant is preserved, not dropped - assert token.scopes == ("read", "write") - assert persisted[0][5] == ("read", "write") + assert token.scopes == prior_scopes + assert persisted[0][5] == prior_scopes @pytest.mark.asyncio @@ -257,18 +258,19 @@ async def test_refresh_omits_resource_when_unset(): @pytest.mark.asyncio -async def test_returned_scope_overrides_prior_when_present(): +@pytest.mark.parametrize("scope, expected", [("read", ("read",)), ("", ()), (" ", ())]) +async def test_returned_scope_overrides_prior_when_present(scope, expected): persisted = [] refresher = _refresher( - body={"access_token": "new-at", "scope": "read"}, + body={"access_token": "new-at", "scope": scope}, persist_sink=persisted, ) token = await refresher.refresh( "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) ) assert token is not None - assert token.scopes == ("read",) # a present scope replaces the prior grant - assert persisted[0][5] == ("read",) + assert token.scopes == expected + assert persisted[0][5] == expected @pytest.mark.asyncio @@ -366,7 +368,7 @@ async def test_cimd_refresh_on_fresh_replica_preserves_user_and_client_identity( {}, ) ] - assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, None)] + assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, ())] assert server.client_id is None From fc9895c745f336bfacb5204570e94e3c01a553f0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:35:04 -0700 Subject: [PATCH 7/8] fix(mcp): preserve empty grants through credential persistence --- .../per_user_oauth_store.py | 2 +- .../test_per_user_oauth_store.py | 57 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index d8094d5facb..2e9f3232c8d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -88,7 +88,7 @@ async def _persist_credential( access_token=access_token, refresh_token=refresh_token, expires_in=expires_in, - scopes=list(scopes) if scopes else None, + scopes=list(scopes) if scopes is not None else None, skip_byok_guard=True, identity_binding_proof=identity_binding_proof, **({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}), diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index e7308884060..5e2bb0f3722 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -26,6 +26,63 @@ class _RecordingStore: self.invalidations.append((user_id, server_id)) +@pytest.fixture +def credential_database(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + + prisma = MagicMock() + table = prisma.db.litellm_mcpusercredentials + table.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setenv("LITELLM_SALT_KEY", "credential-persistence-test-salt") + return table + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes, expected", [(None, None), ((), []), (("read",), ["read"])]) +async def test_credential_persistence_preserves_known_and_unknown_scopes(credential_database, scopes, expected): + from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential + + await _persist_credential("alice", "srv", "access", "refresh", 60, scopes) + row = credential_database.upsert.call_args.kwargs["data"]["create"] + payload = decode_oauth_payload(row["credential_b64"]) + assert payload is not None and payload.get("scopes") == expected + assert row["user_id"] == "alice" and row["server_id"] == "srv" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prior_scopes, response_scope, expected", + [((), None, []), (("read", "write"), "", []), (("read", "write"), None, ["read", "write"])], +) +async def test_native_refresh_persists_actual_grant_through_storage_adapter( + credential_database, prior_scopes, response_scope, expected, +): + from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload + from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import AuthorizationCodeRefresher + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", name="srv", transport="http", auth_type="oauth2", + token_url="https://idp.example/token", client_id="client", client_secret="secret", + ) + + async def token_endpoint(url, form, headers): + assert url == server.token_url and form["grant_type"] == "refresh_token" + return {"access_token": "rotated", **({"scope": response_scope} if response_scope is not None else {})} + + refresher = AuthorizationCodeRefresher(lambda _: server, token_endpoint, _persist_credential) + refreshed = await refresher.refresh("alice", "srv", OAuthToken("old", refresh_token="refresh", scopes=prior_scopes)) + assert refreshed is not None and refreshed.scopes == tuple(expected) + payload = decode_oauth_payload(credential_database.upsert.call_args.kwargs["data"]["create"]["credential_b64"]) + assert payload is not None and payload.get("scopes") == expected + assert payload["access_token"] == "rotated" + + class _BlockingStore: def __init__(self, access_token: str) -> None: self._access_token = access_token From f524d452cb280e128989155179b3ec723c553b70 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:46:01 -0700 Subject: [PATCH 8/8] test(mcp): run scope regressions in OAuth CI shard --- .github/workflows/test-unit.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 18df865c988..fb01ab55e6b 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -805,6 +805,10 @@ jobs: tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py + tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py + tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py + tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py + tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py tests/unit/proxy/_experimental/mcp_server/outbound_credentials python-version: "3.12" workers: 4