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

Connect your MCP servers to {clientLabel}

+

+ Authorize the servers you want to use below, then finish connecting to return to {clientLabel}. Closing + this tab finishes for you. +

+
+
+
(finishedRef.current = true)}> + + +
+
+
+ ); +}; + +export default ConnectFlowBanner; diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 867522090d2d..ac26e90dfb63 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -13,7 +13,7 @@ import { getMCPOAuthUserCredentialStatus, listMCPTools, } from "../networking"; -import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types"; +import { AUTH_TYPE, MCPServer, MCPTool, handleTransport, isUnsupportedOnGatewayConnect } from "../mcp_tools/types"; import MessageManager from "@/components/molecules/message_manager"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; @@ -70,6 +70,7 @@ interface Props { accessToken: string; selectedServers: string[]; onChange: (servers: string[]) => void; + connectMode?: boolean; } const AVATAR_COLORS = [ @@ -95,7 +96,7 @@ type TabKey = "all" | "connected"; const TOOLS_FETCH_CONCURRENCY = 5; -const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange }) => { +const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, connectMode }) => { const [servers, setServers] = useState([]); const [loading, setLoading] = useState(true); const [query, setQuery] = useState(""); @@ -105,6 +106,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const [toolCounts, setToolCounts] = useState>({}); const [loadingCounts, setLoadingCounts] = useState(false); const [oauthConnected, setOauthConnected] = useState>(new Set()); + const [oauthChecking, setOauthChecking] = useState>(new Set()); const serversRef = useRef([]); useEffect(() => { @@ -147,6 +149,14 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange } } catch { // ignore + } finally { + if (!fetchLoadCancelledRef.current) { + setOauthChecking((prev) => { + const next = new Set(prev); + next.delete(server.server_id); + return next; + }); + } } }, [accessToken], @@ -159,9 +169,13 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange .then(async (serverData) => { if (fetchLoadCancelledRef.current) return; const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? []; + const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); setServers(list); + setOauthChecking(new Set(oauthServers.map((s) => s.server_id))); setLoading(false); + oauthServers.forEach((s) => checkOauthCredential(s)); + setLoadingCounts(true); const chunks = Array.from({ length: Math.ceil(list.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) => list.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY), @@ -171,9 +185,6 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange await Promise.allSettled(chunk.map((s) => fetchToolCount(s))); } if (!fetchLoadCancelledRef.current) setLoadingCounts(false); - - const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); - oauthServers.forEach((s) => checkOauthCredential(s)); }) .catch(() => { if (!fetchLoadCancelledRef.current) { @@ -230,6 +241,36 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange } }; + const renderConnectionIndicator = (server: MCPServer) => { + if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) { + return ( + + Not supported on this connection + + ); + } + if (server.auth_type === AUTH_TYPE.OAUTH2) { + if (oauthConnected.has(server.server_id)) { + return ; + } + if (oauthChecking.has(server.server_id)) { + return ; + } + return ( + setOauthConnected((prev) => new Set(prev).add(id))} + variant="badge" + /> + ); + } + if (selectedServers.includes(nameOf(server))) { + return ; + } + return null; + }; + const { data: detailToolsResult, isLoading: loadingTools } = useQuery({ queryKey: ["mcp-apps-panel-detail-tools", detailServer?.server_id], queryFn: () => listMCPTools(accessToken, detailServer!.server_id), @@ -396,24 +437,30 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange

MCP Servers

- - Beta - -
-
-

Browse tools, authenticate once, use in chat

- {loadingCounts ? ( - - - Loading tools... + {!connectMode && ( + + Beta - ) : totalTools > 0 ? ( - - - {totalTools} tool{totalTools !== 1 ? "s" : ""} available - - ) : null} + )}
+ {connectMode ? ( +

Click a server to see its tools and connect

+ ) : ( +
+

Browse tools, authenticate once, use in chat

+ {loadingCounts ? ( + + + Loading tools... + + ) : totalTools > 0 ? ( + + + {totalTools} tool{totalTools !== 1 ? "s" : ""} available + + ) : null} +
+ )}
@@ -464,10 +511,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange
{filtered.map((server, idx) => { const name = nameOf(server); - const isConnected = selectedServers.includes(name); const color = getAvatarColor(name); const isLeftCol = idx % 2 === 0; const count = toolCounts[name]; + const unsupported = !!connectMode && isUnsupportedOnGatewayConnect(server.auth_type); return (
= ({ accessToken, selectedServers, onChange onClick={() => setDetailServer(server)} className={`flex items-center gap-3 p-4 bg-card cursor-pointer transition-colors hover:bg-accent/30 min-w-0 ${ isLeftCol ? "border-r" : "" - } ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""}`} + } ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""} ${ + unsupported ? "opacity-50" : "" + }`} > {server.mcp_info?.logo_url ? ( = ({ accessToken, selectedServers, onChange ) : null}
- {server.auth_type === AUTH_TYPE.OAUTH2 ? ( - oauthConnected.has(server.server_id) ? ( - - ) : ( - { - setOauthConnected((prev) => new Set(prev).add(id)); - }} - variant="badge" - /> - ) - ) : isConnected ? ( - - ) : null} + {renderConnectionIndicator(server)}
); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx index 7ce803d583c9..175df00fd665 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx @@ -13,6 +13,7 @@ import { preservedDeclaredAppCredentials, withoutMintedTokenCredentials, credentialAuthClass, + isUnsupportedOnGatewayConnect, } from "./types"; describe("getOAuthAuthorizationIdentity", () => { @@ -231,3 +232,23 @@ describe("credentialAuthClass", () => { expect(credentialAuthClass(null)).toBeNull(); }); }); + +describe("isUnsupportedOnGatewayConnect", () => { + it("flags the modes that need a caller-supplied upstream token or subject", () => { + // client-forwarded: caller presents the upstream Authorization per call + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.TRUE_PASSTHROUGH)).toBe(true); + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH_DELEGATE)).toBe(true); + // OBO: caller's own IdP token is the exchange subject, which the session bearer is not + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE)).toBe(true); + }); + + it("does not flag modes the gateway can serve from server-side state or interactive vaulting", () => { + // interactive authorization_code is the one mode the connect grid vaults per user + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH2)).toBe(false); + // server-configured credentials need no per-user connect + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.API_KEY)).toBe(false); + expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.NONE)).toBe(false); + expect(isUnsupportedOnGatewayConnect(null)).toBe(false); + expect(isUnsupportedOnGatewayConnect(undefined)).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index dd7ed2bfbe4c..2078f6441af2 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -52,6 +52,15 @@ export const AUTH_TYPE = { export const isClientForwardedTokenMode = (authType?: string | null): boolean => authType === AUTH_TYPE.TRUE_PASSTHROUGH || authType === AUTH_TYPE.OAUTH_DELEGATE; +// Auth modes that cannot be used through the gateway aggregate connect flow, where the client holds +// only an identity-only session bearer and upstream credentials are resolved server-side from the +// per-user vault. The vault is only populated by interactive authorization_code (oauth2). The +// client-forwarded modes need the caller to present the upstream Authorization per call, and +// oauth2_token_exchange (OBO) needs the caller's own IdP token as the subject to exchange; the +// session bearer is neither, so none of these can complete a tool call on this connection. +export const isUnsupportedOnGatewayConnect = (authType?: string | null): boolean => + isClientForwardedTokenMode(authType) || authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE; + export const OAUTH_FLOW = { INTERACTIVE: "interactive", M2M: "m2m", diff --git a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts index 83b2cb21bede..37046ca977a8 100644 --- a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts +++ b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts @@ -16,6 +16,14 @@ import { getProxyBaseUrl, serverRootPath } from "@/components/networking"; */ export const TOOLS_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-tools-state"; +/** + * sessionStorage flag set by useUserMcpOAuthFlow right before it navigates the whole page + * to the upstream IdP to authorize one server. ConnectFlowBanner's auto-finish-on-close + * handler skips while this is set, so authorizing a server is not mistaken for the user + * leaving the gateway DCR connect flow. + */ +export const PERSERVER_CONNECTING_KEY = "litellm-mcp-perserver-connecting"; + /** * Build the OAuth callback URL for the current UI deployment. * diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index 13bad2272971..0c54a0c2a479 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -23,7 +23,7 @@ import NotificationsManager from "@/components/molecules/notifications_manager"; import { extractErrorMessage } from "@/utils/errorUtils"; import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; -import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils"; +import { buildCallbackUrl, clearStorage, PERSERVER_CONNECTING_KEY } from "./mcpOAuthUtils"; export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; @@ -135,6 +135,7 @@ export const useUserMcpOAuthFlow = ({ returnUrl.searchParams.set("mcpOauthReturn", "apps"); setStorage(RETURN_URL_KEY, returnUrl.toString()); + sessionStorage.setItem(PERSERVER_CONNECTING_KEY, "1"); window.location.href = authorizeUrl; } catch (err) { const msg = extractErrorMessage(err);