mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
feat(mcp): consent capture of the user's IdP grant for delegated OBO
The mint arm exchanges a user's stored IdP grant, but nothing captured it, so a delegated token_exchange call always failed closed. This adds the first-time consent flow that captures the grant. A dedicated per-IdP OAuth config (mcp_idp_oauth_providers) carries the authorize URL, token URL, client credentials, and scopes (offline_access), keyed by the IdP token endpoint so a captured grant lands under the exact key the mint arm reads. Two /v1/mcp/idp endpoints run a server-terminal authorization_code flow: authorize authenticates the signed-in user, seals them plus a gateway-generated PKCE verifier into an encrypted, time-stamped state, and redirects to the IdP; the terminal callback recovers the user from the sealed state, exchanges the code itself (the gateway is the OAuth client end to end, unlike the relay flow), and stores the refresh token via capture_user_idp_grant. The state is the encrypted CaptureState, so it doubles as CSRF binding and, with issued_at, bounds replay; the redirect target is the configured authorize_url, never caller input. Part of LIT-4448 (build item 2.2).
This commit is contained in:
Tin Chi Lo
parent
8ee8e2b097
commit
ceb8845679
8 changed files with 782 additions and 6 deletions
203
litellm/proxy/_experimental/mcp_server/idp_consent.py
Normal file
203
litellm/proxy/_experimental/mcp_server/idp_consent.py
Normal file
|
|
@ -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),
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue