diff --git a/litellm/proxy/_experimental/mcp_server/idp_consent.py b/litellm/proxy/_experimental/mcp_server/idp_consent.py new file mode 100644 index 00000000000..2ec63c63a22 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/idp_consent.py @@ -0,0 +1,203 @@ +"""Pure logic for the delegated-OBO consent-capture flow (Path B, item 2.2). + +A server-terminal ``authorization_code + offline_access`` flow against the user's IdP: unlike the +per-user MCP-server OAuth flow (which relays the code back to a browser client that finishes the +exchange), here the gateway is the OAuth client end to end. It generates its own PKCE, seals the user +and the target IdP into the state, and on callback exchanges the code for the user's tokens itself, +then stores the refresh_token as the user's IdP grant for the mint arm to use. + +This module is the pure core (PKCE, state sealing, authorize-URL building, code exchange); the HTTP +POST, the crypto, and the persistence are injected so it is testable without a live IdP or DB. The +endpoints that wire it onto the request path live with the other ``/v1/mcp`` management routes. +""" + +from __future__ import annotations + +import base64 +import hashlib +import secrets +from collections.abc import Awaitable, Callable +from urllib.parse import urlencode, urlsplit, urlunsplit + +from pydantic import BaseModel, ConfigDict + +from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( + TokenEndpointAuthConfigError, + build_token_endpoint_client_auth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( + IdpOAuthProvider, +) + +# POSTs an OAuth form to a token endpoint and returns the parsed JSON body, or None on any failure. +TokenEndpointPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable["dict[str, object] | None"]] + +_PKCE_VERIFIER_BYTES = 64 +_CODE_CHALLENGE_METHOD = "S256" + + +_STATE_MAX_AGE_SECONDS = 600.0 + + +class CaptureState(BaseModel): + """The sealed state carried through the IdP round-trip: who is consenting, for which IdP, and the + PKCE verifier to replay at the token exchange. Sealed (encrypted) so the client cannot tamper with + the bound user or forge a verifier, and stamped with ``issued_at`` so a captured state cannot be + replayed indefinitely.""" + + model_config = ConfigDict(frozen=True) + + user_id: str + token_url: str + code_verifier: str + issued_at: float + + +def state_is_fresh(state: CaptureState, *, now: float, max_age_seconds: float = _STATE_MAX_AGE_SECONDS) -> bool: + """Whether the sealed state is within its lifetime, bounding replay of a captured state.""" + return 0 <= (now - state.issued_at) <= max_age_seconds + + +class CapturedGrant(BaseModel): + """The user's IdP grant captured from the authorization_code exchange.""" + + model_config = ConfigDict(frozen=True) + + access_token: str + refresh_token: str | None + expires_in: int | None + scopes: tuple[str, ...] + + +def generate_pkce() -> tuple[str, str]: + """Return a fresh ``(code_verifier, code_challenge)`` pair (RFC 7636 S256).""" + verifier = secrets.token_urlsafe(_PKCE_VERIFIER_BYTES) + digest = hashlib.sha256(verifier.encode()).digest() + challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=") + return verifier, challenge + + +def seal_capture_state(state: CaptureState, encrypt: Callable[[str], str]) -> str: + return encrypt(state.model_dump_json()) + + +def unseal_capture_state(blob: str, decrypt: Callable[[str], str | None]) -> CaptureState | None: + decrypted = decrypt(blob) + if decrypted is None: + return None + try: + return CaptureState.model_validate_json(decrypted) + except ValueError: + return None + + +def build_authorize_url( + provider: IdpOAuthProvider, + *, + redirect_uri: str, + state: str, + code_challenge: str, +) -> str: + """Build the IdP authorize URL for the consent redirect, merging the OAuth params onto any query + the configured ``authorize_url`` already carries.""" + params = { + "response_type": "code", + "client_id": provider.client_id, + "redirect_uri": redirect_uri, + "state": state, + "scope": " ".join(provider.scopes), + "code_challenge": code_challenge, + "code_challenge_method": _CODE_CHALLENGE_METHOD, + } + parts = urlsplit(provider.authorize_url) + merged_query = "&".join(q for q in (parts.query, urlencode(params)) if q) + return urlunsplit((parts.scheme, parts.netloc, parts.path, merged_query, parts.fragment)) + + +def _parse_expires_in(raw: object) -> int | None: + if isinstance(raw, bool): + return None + if isinstance(raw, (int, float)): + return int(raw) + if isinstance(raw, str): + try: + return int(float(raw)) + except ValueError: + return None + return None + + +def _parse_scopes(raw: object, fallback: tuple[str, ...]) -> tuple[str, ...]: + if isinstance(raw, str) and raw: + return tuple(raw.split()) + return fallback + + +async def default_token_post(url: str, form: dict[str, str], headers: dict[str, str]) -> dict[str, object] | None: + """The real httpx POST to an IdP token endpoint for the consent exchange (Oauth2Check provider). + + Returns the parsed JSON body, or None on any failure so the exchange fails closed. Mirrors the + per-provider token-endpoint POST helpers; extracting one shared helper across them is a follow-up. + """ + from litellm._logging import verbose_logger # noqa: PLC0415 # lazy import; avoids cycle + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # lazy import; avoids cycle + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # httpx handler untyped + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # lazy import + + request_headers = {"Accept": "application/json", **headers} + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response = await client.post(url, headers=request_headers, data=form) # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post untyped + if response is None: + return None + response.raise_for_status() + body: dict[str, object] = response.json() # pyright: ignore[reportAny] # untyped JSON body, validated below + except Exception as exc: # noqa: BLE001 # any IdP/transport error is a capture miss, not a 500 + verbose_logger.warning("MCP IdP consent code exchange failed: %s", exc) + return None + else: + return body + + +async def exchange_code_for_grant( + provider: IdpOAuthProvider, + *, + code: str, + redirect_uri: str, + code_verifier: str, + post: TokenEndpointPost, +) -> CapturedGrant | None: + """Exchange the authorization code for the user's IdP grant (server-side), or None on failure. + + Fails closed (None) rather than raising: a missing access_token, a bad client-auth config, or any + transport error is a capture miss the caller surfaces as an error, never a partial grant. + """ + try: + client_auth = build_token_endpoint_client_auth( + auth_method=None, + client_id=provider.client_id, + client_secret=provider.client_secret.get_secret_value(), + ) + except TokenEndpointAuthConfigError: + return None + form = { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "code_verifier": code_verifier, + **client_auth.body, + } + body = await post(provider.token_url, form, client_auth.headers) + if body is None: + return None + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return None + refresh_token = body.get("refresh_token") + return CapturedGrant( + access_token=access_token, + refresh_token=refresh_token if isinstance(refresh_token, str) and refresh_token else None, + expires_in=_parse_expires_in(body.get("expires_in")), + scopes=_parse_scopes(body.get("scope"), provider.scopes), + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_oauth_config.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_oauth_config.py new file mode 100644 index 00000000000..db4453716f6 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_oauth_config.py @@ -0,0 +1,91 @@ +"""The IdP OAuth provider config for delegated-OBO consent capture (Path B, item 2.2). + +The mint arm (item 3) exchanges a user's stored IdP grant; this is where that grant comes from: a +first-time consent flow runs an ``authorization_code + offline_access`` grant against the user's IdP +(e.g. Okta) so the gateway captures and stores the user's refresh token. The token_exchange server +config carries the IdP's token endpoint and the gateway's client credentials, but not the authorize +endpoint or the offline_access scope the capture needs, so those live here. + +One provider serves every ``token_exchange`` upstream fronted by the same IdP, so a provider is keyed +by its token endpoint - the same anchor the mint arm's grant store uses (``idp_grant_key``) - and the +consent flow stores the captured grant under that key, ready for the mint arm to read. +""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter + +from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_subject_source import ( + idp_grant_key, +) + + +class IdpOAuthProvider(BaseModel): + """One IdP's OAuth config for the consent-capture ``authorization_code`` flow. + + ``token_url`` is the IdP's token endpoint; it is also the anchor the captured grant is keyed by + (via ``idp_grant_key``), so it must match the ``token_exchange_endpoint`` of the servers this IdP + fronts. ``scopes`` must include an offline-access scope for the IdP to return a refresh_token; the + default covers the OIDC + Okta case. + """ + + model_config = ConfigDict(frozen=True) + + token_url: str + authorize_url: str + client_id: str + client_secret: SecretStr + scopes: tuple[str, ...] = ("openid", "offline_access") + + @property + def grant_key(self) -> str: + """The (user, idp) storage key's idp component the captured grant is stored under.""" + return idp_grant_key(self.token_url) + + +class IdpOAuthProviderRegistry: + """Immutable lookup of the configured IdP providers, keyed by their grant key. + + Built once from config at startup. A provider is resolved by the ``idp`` the consent flow targets + (the token endpoint), so the captured grant lands under the exact key the mint arm reads. + """ + + def __init__(self, providers: tuple[IdpOAuthProvider, ...]) -> None: + self._by_key: dict[str, IdpOAuthProvider] = {p.grant_key: p for p in providers} + + def get(self, grant_key: str) -> IdpOAuthProvider | None: + return self._by_key.get(grant_key) + + def get_by_token_url(self, token_url: str) -> IdpOAuthProvider | None: + return self._by_key.get(idp_grant_key(token_url)) + + def __len__(self) -> int: + return len(self._by_key) + + +def load_idp_oauth_providers(raw_providers: object) -> IdpOAuthProviderRegistry: + """Build the registry from the ``mcp_idp_oauth_providers`` config block (a list of dicts). + + Each entry is validated into an ``IdpOAuthProvider``; a malformed entry raises at load time (fail + fast at startup) rather than surfacing as a broken consent flow later. A missing/empty block yields + an empty registry, so the consent endpoints simply 404 until an IdP is configured. + """ + if not isinstance(raw_providers, list): + return IdpOAuthProviderRegistry(()) + return IdpOAuthProviderRegistry(_PROVIDERS_ADAPTER.validate_python(raw_providers)) + + +_PROVIDERS_ADAPTER: TypeAdapter[tuple[IdpOAuthProvider, ...]] = TypeAdapter(tuple[IdpOAuthProvider, ...]) + + +_registry = IdpOAuthProviderRegistry(()) + + +def set_idp_oauth_registry(registry: IdpOAuthProviderRegistry) -> None: + """Install the process-wide registry (called once from config load at startup).""" + global _registry + _registry = registry + + +def get_idp_oauth_registry() -> IdpOAuthProviderRegistry: + return _registry diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_provider.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_provider.py index 5898b74ce32..aae2102e37e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_provider.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_provider.py @@ -54,7 +54,13 @@ async def _persist_credential( refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, -) -> None: +) -> bool: + """Persist the grant, returning whether it was actually stored. + + A refresh treats the skipped save as best-effort (it still returns the fresh token), but the + first-time consent capture must surface a False so the callback never reports a connection that + was never stored. + """ from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # lazy: avoids import cycle store_user_idp_grant, ) @@ -65,7 +71,7 @@ async def _persist_credential( # persist observable, since a refresh that rotated the IdP refresh_token then failed to save # it strands the user until re-consent. verbose_logger.warning("MCP IdP grant persist skipped: database not connected; a rotated grant may be lost") - return + return False await store_user_idp_grant( prisma_client=prisma_client, user_id=user_id, @@ -75,6 +81,7 @@ async def _persist_credential( expires_in=expires_in, scopes=list(scopes) if scopes else None, ) + return True async def _post_token_endpoint(url: str, form: dict[str, str], headers: dict[str, str]) -> dict[str, object] | None: @@ -134,14 +141,17 @@ async def capture_user_idp_grant( refresh_token: str | None = None, expires_in: int | None = None, scopes: tuple[str, ...] | None = None, -) -> None: +) -> bool: """Persist a user's IdP grant for on-behalf-of exchange, keyed by the IdP (the AS token endpoint). The store-back path the first-time consent flow calls once it has captured the user's IdP grant (``authorization_code + offline_access`` against the IdP). One grant serves every ``token_exchange`` upstream that IdP fronts, since it is keyed by the IdP endpoint rather than an upstream server. + + Returns whether the grant was stored, so the caller can distinguish a real connection from a + skipped save (database unavailable) rather than reporting a connection that never persisted. """ - await _persist_credential( + return await _persist_credential( user_id, idp_grant_key(token_exchange_endpoint), access_token, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_source.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_source.py index 340162d8c68..59c52f38b9d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_source.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/idp_subject_source.py @@ -62,7 +62,7 @@ class PersistIdpGrant(Protocol): refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, - ) -> None: ... + ) -> bool: ... def _parse_expires_in(raw: object) -> int | None: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 757835fd05a..2d38acfc184 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,6 +19,7 @@ import functools import importlib import json import os +import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Literal, Optional, Set @@ -34,7 +35,7 @@ from fastapi import ( Response, status, ) -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, RedirectResponse try: from prisma.errors import RecordNotFoundError, UniqueViolationError @@ -1729,6 +1730,136 @@ if MCP_AVAILABLE: scope=scope, ) + # Module-scope Depends singleton so the param default is a name, not a call (B008); same dep the + # relay OAuth routes use (Authorization header or the SSO session cookie). + _idp_consent_auth_dep = Depends(_mcp_oauth_user_api_key_auth) + + def _idp_callback_redirect_uri(request: Request) -> str: + from litellm.proxy._experimental.mcp_server.oauth_utils import ( # noqa: PLC0415 # lazy: avoids cycle + get_request_base_url, + ) + + return f"{get_request_base_url(request)}/v1/mcp/idp/callback" + + @router.get("/idp/authorize", include_in_schema=False) + async def mcp_idp_authorize( + request: Request, + user_api_key_dict: UserAPIKeyAuth = _idp_consent_auth_dep, + token_url: str = "", + ): + """Start the delegated-OBO consent capture: redirect the signed-in user to their IdP to + authorize the gateway, so the callback can capture and store their IdP grant (Path B, 2.2).""" + from litellm.proxy._experimental.mcp_server.idp_consent import ( # noqa: PLC0415 # lazy: avoids cycle + CaptureState, + build_authorize_url, + generate_pkce, + seal_capture_state, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( # noqa: PLC0415 # lazy: avoids cycle + get_idp_oauth_registry, + ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, # noqa: PLC0415 # lazy: avoids cycle + ) + + provider = get_idp_oauth_registry().get_by_token_url(token_url) + if provider is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": "No IdP provider configured for this token endpoint"}, + ) + user_id = user_api_key_dict.user_id + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "No user id on the session; sign in first"} + ) + + code_verifier, code_challenge = generate_pkce() + state = seal_capture_state( + CaptureState( + user_id=user_id, token_url=provider.token_url, code_verifier=code_verifier, issued_at=time.time() + ), + encrypt_value_helper, + ) + authorize_url = build_authorize_url( + provider, + redirect_uri=_idp_callback_redirect_uri(request), + state=state, + code_challenge=code_challenge, + ) + return RedirectResponse(authorize_url) + + @router.get("/idp/callback", include_in_schema=False) + async def mcp_idp_callback( + request: Request, + code: str | None = None, + state: str | None = None, + error: str | None = None, + ): + """Terminal callback for the consent capture: the gateway itself exchanges the code for the + user's IdP grant and stores it (server-terminal, unlike the relay callback). The user is + recovered from the sealed state, so this endpoint needs no session.""" + from litellm.proxy._experimental.mcp_server.idp_consent import ( # noqa: PLC0415 # lazy: avoids cycle + default_token_post, + exchange_code_for_grant, + state_is_fresh, + unseal_capture_state, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( # noqa: PLC0415 # lazy: avoids cycle + get_idp_oauth_registry, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_subject_provider import ( # noqa: PLC0415 # lazy: avoids cycle + capture_user_idp_grant, + ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, # noqa: PLC0415 # lazy: avoids cycle + ) + + if error: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail={"error": f"IdP returned an error: {error}"} + ) + if not code or not state: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "Missing code or state"}) + + capture_state = unseal_capture_state(state, lambda blob: decrypt_value_helper(blob, "oauth_state")) + if capture_state is None or not state_is_fresh(capture_state, now=time.time()): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "Invalid, tampered, or expired state"} + ) + provider = get_idp_oauth_registry().get_by_token_url(capture_state.token_url) + if provider is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail={"error": "IdP provider is no longer configured"} + ) + + grant = await exchange_code_for_grant( + provider, + code=code, + redirect_uri=_idp_callback_redirect_uri(request), + code_verifier=capture_state.code_verifier, + post=default_token_post, + ) + if grant is None: + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, detail={"error": "IdP did not return a usable grant"} + ) + + persisted = await capture_user_idp_grant( + capture_state.user_id, + provider.token_url, + grant.access_token, + refresh_token=grant.refresh_token, + expires_in=grant.expires_in, + scopes=grant.scopes, + ) + if not persisted: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": "Could not store your connection; the database is unavailable. Please try again"}, + ) + return JSONResponse({"status": "connected", "offline_access": grant.refresh_token is not None}) + @router.post( "/server/oauth/{server_id}/register", include_in_schema=False, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0a2b65f927..107ec35dcce 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4992,6 +4992,16 @@ class ProxyConfig: await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases) + ## MCP IdP OAUTH PROVIDERS (delegated-OBO consent capture) + mcp_idp_oauth_providers = config.get("mcp_idp_oauth_providers", None) + if mcp_idp_oauth_providers: + from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( + load_idp_oauth_providers, + set_idp_oauth_registry, + ) + + set_idp_oauth_registry(load_idp_oauth_providers(mcp_idp_oauth_providers)) + ## VECTOR STORES vector_store_registry_config = config.get("vector_store_registry", None) if vector_store_registry_config: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_consent.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_consent.py new file mode 100644 index 00000000000..5ce93187f21 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_consent.py @@ -0,0 +1,160 @@ +"""Tests for the delegated-OBO consent-capture flow logic (Path B, item 2.2). + +Covers the pure core: PKCE generation, sealed-state round-trip and tamper rejection, authorize-URL +building (offline_access + S256), and the server-side code->grant exchange with its fail-closed paths. +The token POST and the crypto are injected, so no live IdP is needed. +""" + +import base64 +import hashlib + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.idp_consent import ( + CaptureState, + build_authorize_url, + exchange_code_for_grant, + generate_pkce, + seal_capture_state, + unseal_capture_state, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( + IdpOAuthProvider, + IdpOAuthProviderRegistry, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_subject_source import ( + idp_grant_key, +) + +_PROVIDER = IdpOAuthProvider( + token_url="https://idp.example.com/oauth2/v1/token", + authorize_url="https://idp.example.com/oauth2/v1/authorize", + client_id="gateway-client", + client_secret=SecretStr("gateway-secret"), + scopes=("openid", "offline_access"), +) + + +def test_generate_pkce_is_s256_of_the_verifier(): + verifier, challenge = generate_pkce() + expected = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=") + assert challenge == expected + assert "=" not in challenge # base64url, unpadded + # Two calls produce different verifiers (not a constant). + assert generate_pkce()[0] != verifier + + +def test_sealed_state_round_trips(): + store: dict[str, str] = {} + + def encrypt(plain): + token = f"sealed::{plain}" + store[token] = plain + return token + + def decrypt(blob): + return store.get(blob) + + state = CaptureState(user_id="alice", token_url=_PROVIDER.token_url, code_verifier="verifier-xyz", issued_at=1000.0) + sealed = seal_capture_state(state, encrypt) + assert unseal_capture_state(sealed, decrypt) == state + + +def test_state_is_fresh_bounds_replay(): + from litellm.proxy._experimental.mcp_server.idp_consent import state_is_fresh + + state = CaptureState(user_id="a", token_url=_PROVIDER.token_url, code_verifier="v", issued_at=1000.0) + assert state_is_fresh(state, now=1000.0, max_age_seconds=600) is True + assert state_is_fresh(state, now=1599.0, max_age_seconds=600) is True + assert state_is_fresh(state, now=1601.0, max_age_seconds=600) is False # expired + assert state_is_fresh(state, now=900.0, max_age_seconds=600) is False # issued in the future (clock skew / forged) + + +def test_unseal_rejects_tampered_or_undecryptable_state(): + # decrypt returning None (bad/forged blob) -> None, never a partial/forged CaptureState. + assert unseal_capture_state("garbage", lambda _b: None) is None + # decrypts to non-CaptureState JSON -> None. + assert unseal_capture_state("x", lambda _b: '{"not":"a state"}') is None + + +def test_build_authorize_url_carries_offline_access_and_s256(): + url = build_authorize_url( + _PROVIDER, + redirect_uri="https://gw.example.com/v1/mcp/idp/callback", + state="sealed-state", + code_challenge="chal", + ) + assert url.startswith("https://idp.example.com/oauth2/v1/authorize?") + assert "response_type=code" in url + assert "client_id=gateway-client" in url + assert "code_challenge=chal" in url + assert "code_challenge_method=S256" in url + assert "offline_access" in url + assert "redirect_uri=https%3A%2F%2Fgw.example.com%2Fv1%2Fmcp%2Fidp%2Fcallback" in url + + +def test_build_authorize_url_preserves_existing_query(): + provider = _PROVIDER.model_copy(update={"authorize_url": "https://idp.example.com/authorize?tenant=acme"}) + url = build_authorize_url(provider, redirect_uri="https://gw/cb", state="s", code_challenge="c") + assert "tenant=acme" in url + assert "response_type=code" in url + + +@pytest.mark.asyncio +async def test_exchange_code_for_grant_parses_the_grant(): + captured_form: dict[str, str] = {} + + async def post(url, form, headers): + captured_form.update(form) + return {"access_token": "user-at", "refresh_token": "user-rt", "expires_in": 3600, "scope": "openid offline_access"} + + grant = await exchange_code_for_grant( + _PROVIDER, code="auth-code", redirect_uri="https://gw/cb", code_verifier="ver", post=post + ) + assert grant is not None + assert grant.access_token == "user-at" + assert grant.refresh_token == "user-rt" + assert grant.expires_in == 3600 + assert grant.scopes == ("openid", "offline_access") + # The authorization_code grant is posted with the PKCE verifier and the client id. + assert captured_form["grant_type"] == "authorization_code" + assert captured_form["code"] == "auth-code" + assert captured_form["code_verifier"] == "ver" + assert captured_form["client_id"] == "gateway-client" + + +@pytest.mark.asyncio +async def test_exchange_code_falls_closed_on_missing_access_token_or_failure(): + async def no_access_token(url, form, headers): + return {"refresh_token": "rt"} # no access_token + + async def transport_failure(url, form, headers): + return None + + assert await exchange_code_for_grant(_PROVIDER, code="c", redirect_uri="r", code_verifier="v", post=no_access_token) is None + assert await exchange_code_for_grant(_PROVIDER, code="c", redirect_uri="r", code_verifier="v", post=transport_failure) is None + + +@pytest.mark.asyncio +async def test_exchange_code_carries_forward_scopes_when_response_omits_scope(): + async def post(url, form, headers): + return {"access_token": "at"} # no scope, no refresh_token, no expires_in + + grant = await exchange_code_for_grant(_PROVIDER, code="c", redirect_uri="r", code_verifier="v", post=post) + assert grant is not None + assert grant.scopes == ("openid", "offline_access") + assert grant.refresh_token is None + assert grant.expires_in is None + + +def test_provider_grant_key_matches_the_mint_arm_key(): + # The captured grant must land under the exact key the mint arm reads. + assert _PROVIDER.grant_key == idp_grant_key("https://idp.example.com/oauth2/v1/token") + + +def test_registry_resolves_by_grant_key_and_token_url(): + registry = IdpOAuthProviderRegistry((_PROVIDER,)) + assert registry.get(_PROVIDER.grant_key) is _PROVIDER + assert registry.get_by_token_url("https://idp.example.com/oauth2/v1/token/") is _PROVIDER # normalized + assert registry.get("idp::https://other/token") is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_idp_consent_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_idp_consent_endpoints.py new file mode 100644 index 00000000000..2242b1f5456 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_idp_consent_endpoints.py @@ -0,0 +1,171 @@ +"""Tests for the delegated-OBO consent-capture endpoints (Path B, item 2.2). + +These pin the endpoint wiring on top of the (separately tested) flow logic: the authorize route seals +the signed-in user into the redirect and points at the configured IdP, and the terminal callback +unseals the user, exchanges the code, and stores the grant under the mint arm's key. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server import idp_consent +from litellm.proxy._experimental.mcp_server.idp_consent import CaptureState, seal_capture_state +from litellm.proxy._experimental.mcp_server.outbound_credentials import idp_subject_provider +from litellm.proxy._experimental.mcp_server.outbound_credentials.idp_oauth_config import ( + IdpOAuthProvider, + IdpOAuthProviderRegistry, + set_idp_oauth_registry, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_idp_authorize, + mcp_idp_callback, +) + +_TOKEN_URL = "https://idp.example.com/oauth2/v1/token" +_PROVIDER = IdpOAuthProvider( + token_url=_TOKEN_URL, + authorize_url="https://idp.example.com/oauth2/v1/authorize", + client_id="gateway-client", + client_secret=SecretStr("gateway-secret"), +) + + +@pytest.fixture(autouse=True) +def _salt(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key-for-idp-consent-endpoint-tests") + + +@pytest.fixture(autouse=True) +def _registry(): + set_idp_oauth_registry(IdpOAuthProviderRegistry((_PROVIDER,))) + yield + set_idp_oauth_registry(IdpOAuthProviderRegistry(())) + + +@pytest.fixture(autouse=True) +def _base_url(monkeypatch): + monkeypatch.setenv("PROXY_BASE_URL", "https://gw.example.com") + + +@pytest.mark.asyncio +async def test_authorize_redirects_to_the_idp_sealing_the_user(): + resp = await mcp_idp_authorize( + MagicMock(), user_api_key_dict=UserAPIKeyAuth(user_id="alice"), token_url=_TOKEN_URL + ) + location = resp.headers["location"] + assert location.startswith("https://idp.example.com/oauth2/v1/authorize?") + assert "code_challenge_method=S256" in location + assert "offline_access" in location + assert "redirect_uri=https%3A%2F%2Fgw.example.com%2Fv1%2Fmcp%2Fidp%2Fcallback" in location + + +@pytest.mark.asyncio +async def test_authorize_404_for_an_unconfigured_idp(): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await mcp_idp_authorize(MagicMock(), user_api_key_dict=UserAPIKeyAuth(user_id="alice"), token_url="https://other/token") + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_authorize_400_without_a_user(): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await mcp_idp_authorize(MagicMock(), user_api_key_dict=UserAPIKeyAuth(user_id=None), token_url=_TOKEN_URL) + assert exc.value.status_code == 400 + + +def _seal(user_id: str, *, issued_at: float | None = None) -> str: + import time + + return seal_capture_state( + CaptureState(user_id=user_id, token_url=_TOKEN_URL, code_verifier="the-verifier", issued_at=issued_at or time.time()), + encrypt_value_helper, + ) + + +@pytest.mark.asyncio +async def test_callback_exchanges_and_stores_the_grant_for_the_sealed_user(monkeypatch): + exchanges: list[dict[str, str]] = [] + captures: list[tuple] = [] + + async def fake_post(url, form, headers): + exchanges.append(form) + return {"access_token": "alice-idp-at", "refresh_token": "alice-idp-rt", "expires_in": 3600} + + async def fake_capture(user_id, token_exchange_endpoint, access_token, *, refresh_token=None, expires_in=None, scopes=None): + captures.append((user_id, token_exchange_endpoint, access_token, refresh_token)) + return True + + monkeypatch.setattr(idp_consent, "default_token_post", fake_post) + monkeypatch.setattr(idp_subject_provider, "capture_user_idp_grant", fake_capture) + + resp = await mcp_idp_callback(MagicMock(), code="auth-code", state=_seal("alice")) + + # The code + PKCE verifier from the sealed state are exchanged, and the grant is stored under the + # sealed user + the IdP token endpoint (the mint arm's key), never a caller-supplied identity. + assert exchanges[0]["code"] == "auth-code" + assert exchanges[0]["code_verifier"] == "the-verifier" + assert captures == [("alice", _TOKEN_URL, "alice-idp-at", "alice-idp-rt")] + assert resp.status_code == 200 + + +@pytest.mark.asyncio +async def test_callback_503_when_the_grant_could_not_be_stored(monkeypatch): + from fastapi import HTTPException + + async def fake_post(url, form, headers): + return {"access_token": "alice-idp-at", "refresh_token": "alice-idp-rt", "expires_in": 3600} + + async def fake_capture(user_id, token_exchange_endpoint, access_token, *, refresh_token=None, expires_in=None, scopes=None): + return False # the store was skipped (e.g. database not connected) + + monkeypatch.setattr(idp_consent, "default_token_post", fake_post) + monkeypatch.setattr(idp_subject_provider, "capture_user_idp_grant", fake_capture) + + # A grant that did not persist must not be reported as connected; the user would otherwise think + # they linked their IdP while every later delegated call fails closed with no stored grant. + with pytest.raises(HTTPException) as exc: + await mcp_idp_callback(MagicMock(), code="auth-code", state=_seal("alice")) + assert exc.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_callback_rejects_a_tampered_state(monkeypatch): + from fastapi import HTTPException + + captured = [] + monkeypatch.setattr(idp_subject_provider, "capture_user_idp_grant", lambda *a, **k: captured.append(a)) + with pytest.raises(HTTPException) as exc: + await mcp_idp_callback(MagicMock(), code="auth-code", state="not-a-valid-sealed-state") + assert exc.value.status_code == 400 + assert captured == [] # nothing stored on a bad state + + +@pytest.mark.asyncio +async def test_callback_rejects_an_expired_state(monkeypatch): + from fastapi import HTTPException + + captured = [] + monkeypatch.setattr(idp_subject_provider, "capture_user_idp_grant", lambda *a, **k: captured.append(a)) + # A state sealed well beyond the 600s lifetime must be rejected, bounding replay. + with pytest.raises(HTTPException) as exc: + await mcp_idp_callback(MagicMock(), code="auth-code", state=_seal("alice", issued_at=1.0)) + assert exc.value.status_code == 400 + assert captured == [] + + +@pytest.mark.asyncio +async def test_callback_400_on_idp_error(): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await mcp_idp_callback(MagicMock(), error="access_denied") + assert exc.value.status_code == 400