From 45fed6a50a231822aeb616c99d4b6f17ffd48da0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:44:10 -0700 Subject: [PATCH 1/5] feat(mcp): generalize the bridge envelope identity to a key_hash or user_id subject The scripted two-header client mints under a virtual key it presents at the token endpoint (key_hash), but the interactive DCR client authenticates via SSO at the bridged authorize, which yields a user, not a key. Make EnvelopeIdentity a discriminated subject (subject_type key_hash | user_id) with key_hash_identity / user_identity constructors, and dispatch admission on it: a key_hash reloads the key, a user_id reloads the user and admits them as themselves (user-level budget and SCIM enforced via the same centralized gate; no team bound, since a user belongs to many teams or none). The interactive producer that mints a user_id envelope lands in the follow-up commit. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 57 +++++++++- .../mcp_server/discoverable_endpoints.py | 4 +- .../outbound_credentials/envelope.py | 48 ++++++-- .../auth/test_user_api_key_auth_mcp.py | 105 +++++++++++++++++- .../test_bridge_credentials.py | 7 +- .../outbound_credentials/test_envelope.py | 35 ++++-- .../mcp_server/test_discoverable_endpoints.py | 3 +- 7 files changed, 230 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index e300a22e5db3..faec35db41a0 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 @@ -18,6 +18,9 @@ is_bridge_envelope_shaped, resolve_bridge_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -543,7 +546,7 @@ async def _admit_dcr_bridge_delegate( header_key = server.alias or server.server_name if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") - admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) + admitted = await MCPRequestHandler._reload_admitted_principal(result.identity) await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} new_headers = {**(mcp_server_auth_headers or {}), **injected} @@ -572,6 +575,58 @@ async def _run_pre_db_read_auth_checks(request: Request, route: str) -> None: route=route, ) + @staticmethod + async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth: + """Reload the live litellm record the envelope's subject references. + + Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that + minted the envelope (the scripted two-header client that presents a litellm key at the + token endpoint), a ``user_id`` reloads the user that authenticated interactively (the + DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both + return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so + team/project/org/budget/SCIM enforcement is identical to the principal presenting + itself directly.""" + match identity.subject_type: + case "key_hash": + return await MCPRequestHandler._reload_admitted_key(identity.subject) + case "user_id": + return await MCPRequestHandler._reload_admitted_user(identity.subject) + case _: + assert_never(identity.subject_type) + + @staticmethod + async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: + """Reload the live user an interactively-minted envelope references and admit them as + themselves. + + 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`` rides on the returned ``UserAPIKeyAuth`` and the + caller's centralized policy gate then enforces the user's live budget and org state, + and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No + team is bound; a user may belong to many teams or none, so the envelope grants the + user's own access rather than silently selecting one team's scope. A missing user + fails closed with a 401 rather than admitting an unresolved identity.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") + 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, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if user_object is None: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + return UserAPIKeyAuth(user_id=user_object.user_id) + @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: """Reload the live key record an admitted envelope references and re-check live policy. diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 8d1713a5911e..3e727ce95bbb 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -961,15 +961,15 @@ def _finish_bridge_mint( build_bridge_token_response, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, ) grant = _bridge_grant_from_token_response(token_response) if not isinstance(grant, UpstreamTokenGrant): return _upstream_rejection_to_mint_error(grant) - identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=ready.key_hash) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=ready.key_hash) sealed = build_bridge_token_response(identity, grant, ready.keys, now) if not isinstance(sealed, SealedEnvelope): return "too_large" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 517c2ef5c8f2..783e64d13e27 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -67,20 +67,42 @@ _ENVELOPE_JWT_ALGORITHM = "HS256" +EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"] +"""Discriminator for what litellm principal the envelope binds the grant to. + +``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it +presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR +client mints under the SSO-authenticated user, which is the only identity that browser login +yields). Admission reloads a key record for the first and a user record for the second, then +runs both through the same live-policy gate, so team/org/budget/revocation enforcement is +identical either way.""" + + class EnvelopeIdentity(BaseModel): - """The litellm identity the envelope binds the inner grant to. - - ``key_hash`` is the hashed litellm key that authorized the mint, never a raw - credential (and the edge rejects a bare hash presented as a bearer). Admission - reloads the live key record by it, so the key's current team/org/object-permission - restrictions and its revocation state are enforced at use time rather than frozen at - mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed - across a server boundary. + """The litellm principal the envelope binds the inner grant to. + + ``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a + hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw + credential (and the edge rejects a bare hash or id presented as a bearer). Admission + reloads the live record by it, so the principal's current team/org restrictions and its + revocation state are enforced at use time rather than frozen at mint time. ``server_id`` + binds the envelope to one MCP server so it cannot be replayed across a server boundary. """ model_config = ConfigDict(frozen=True) server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) + + +def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity: + """The identity for the scripted client that mints under a presented virtual key.""" + return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash) + + +def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity: + """The identity for the interactive DCR client that mints under its SSO user subject.""" + return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id) class UpstreamTokenGrant(BaseModel): @@ -200,7 +222,8 @@ class _EnvelopeClaims(BaseModel): iat: int exp: int server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) grant: str = Field(min_length=1) @@ -236,7 +259,8 @@ def mint_envelope( iat=int(now.timestamp()), exp=int(expires_at.timestamp()), server_id=identity.server_id, - key_hash=identity.key_hash, + subject_type=identity.subject_type, + subject=identity.subject, grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), ) token = ENVELOPE_PREFIX + jwt.encode( @@ -281,7 +305,7 @@ def open_envelope( if not isinstance(grant, UpstreamTokenGrant): return grant return OpenedEnvelope( - identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash), + identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject), grant=grant, ) 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 c785ac577f7d..a6affe5496c0 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 @@ -4910,6 +4910,7 @@ def _mint_bridge_envelope( cls, *, key_hash=None, + user_id=None, server_id="bridge-server-id", access_token="inner-upstream-access-token", token_type="Bearer", @@ -4921,17 +4922,23 @@ def _mint_bridge_envelope( envelope_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, + user_identity, ) from pydantic import SecretStr + identity = ( + user_identity(server_id=server_id, user_id=user_id) + if user_id is not None + else key_hash_identity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH) + ) keys = envelope_keys_from_master_key(master_key or cls._MASTER_KEY) now = minted_at or datetime.now(timezone.utc) sealed = mint_envelope( - identity=EnvelopeIdentity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), + identity=identity, grant=UpstreamTokenGrant( access_token=SecretStr(access_token), token_type=token_type, @@ -4999,6 +5006,22 @@ def _patch_key_reload(*, return_value=None, side_effect=None, team_blocked=False stack.enter_context(patcher) yield get_key_object + @staticmethod + @contextlib.contextmanager + def _patch_user_reload(*, return_value=None, side_effect=None): + """Patch the user-subject reload path an interactively-minted envelope takes: the + ``get_user_object`` lookup ``_reload_admitted_user`` runs (which also drives the SCIM gate), + plus the ``prisma_client`` / ``user_api_key_cache`` globals. The centralized gate's own + fetches fail-safe to None under the MagicMock prisma, so an unblocked user admits. Yields the + ``get_user_object`` mock so a caller can assert the sealed user_id was the reload key.""" + get_user_object = AsyncMock(return_value=return_value, side_effect=side_effect) + 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 + @staticmethod def _mcp_request(path="/mcp/bridge_delegate_server"): """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how @@ -5060,6 +5083,84 @@ async def test_valid_envelope_reloads_live_key_and_admits_its_authorization_cont "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_user_subject_envelope_admits_under_the_reloaded_user(self): + """An interactively-minted (user_id) envelope admits under the reloaded USER, not a key: the + reload is keyed by the sealed user_id, the admitted auth carries that user_id, the raw-key + pipeline is never invoked, and the inner upstream token is injected for egress. This is the + interactive-DCR admission the whole flow exists for.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + 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="sso-user-7", metadata={"scim_active": True}) + ) as get_user_object, + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, mcp_server_auth_headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert get_user_object.await_args.kwargs["user_id"] == "sso-user-7" + assert auth_result.user_id == "sso-user-7" + mock_auth.assert_not_called() + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + + async def test_user_subject_envelope_missing_user_fails_closed_401(self): + """A user_id envelope whose user has since been deleted must fail closed: get_user_object + resolves None, so admission 401s instead of admitting an unresolved identity.""" + envelope = self._mint_bridge_envelope(user_id="ghost-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + 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=None), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): + """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries + scim_active False, so admission 401s rather than letting an offboarded user keep tool access + until the envelope expires.""" + envelope = self._mint_bridge_envelope(user_id="offboarded-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + 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}) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + async def test_revoked_key_envelope_fails_closed_401(self): """An envelope whose key has since been deleted must fail closed: ``get_key_object`` raises for the missing row, so admission 401s instead of admitting the caller as an unrestricted diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 82e8e2aae89e..ecea86bbed40 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -28,13 +28,14 @@ EnvelopeTooLarge, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, ) _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") _SERVER_ID = _IDENTITY.server_id @@ -138,7 +139,7 @@ def test_resolve_envelope_minted_for_another_server_is_invalid(): captured or misrouted envelope cannot forward one server's upstream credential to another. The valid access token stays sealed; the mismatch alone fails the resolve.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - other_server_identity = EnvelopeIdentity(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash) + other_server_identity = key_hash_identity(server_id="srv-OTHER", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=other_server_identity) result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) @@ -155,7 +156,7 @@ def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - unicode_identity = EnvelopeIdentity(server_id="srv-café", key_hash=_IDENTITY.key_hash) + unicode_identity = key_hash_identity(server_id="srv-café", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=unicode_identity) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index b44f3f84cc95..7a2b51c2a954 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -36,8 +36,10 @@ SealedEnvelope, UpstreamTokenGrant, is_envelope, + key_hash_identity, mint_envelope, open_envelope, + user_identity, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value @@ -51,7 +53,7 @@ _WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") def _full_grant() -> UpstreamTokenGrant: @@ -137,12 +139,13 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims(): def test_claim_layout_and_no_plaintext_token_in_envelope(): token = _sealed_token(_full_grant()) claims = _unverified_claims(token) - assert set(claims) == {"iss", "iat", "exp", "server_id", "key_hash", "grant"} + assert set(claims) == {"iss", "iat", "exp", "server_id", "subject_type", "subject", "grant"} assert claims["iss"] == ENVELOPE_ISSUER assert claims["iat"] == int(_NOW.timestamp()) assert claims["exp"] == int(_NOW.timestamp()) + 600 assert claims["server_id"] == "srv-456" - assert claims["key_hash"] == "hashed-key-123" + assert claims["subject_type"] == "key_hash" + assert claims["subject"] == "hashed-key-123" assert _ACCESS_TOKEN not in token assert _ACCESS_TOKEN not in json.dumps(claims) assert _REFRESH_TOKEN not in json.dumps(claims) @@ -226,11 +229,11 @@ def test_wrong_issuer_is_malformed_payload(): def test_missing_identity_claim_is_malformed_payload(): claims = _unverified_claims(_sealed_token(_full_grant())) - forged = _forge({key: value for key, value in claims.items() if key != "key_hash"}) + forged = _forge({key: value for key, value in claims.items() if key != "subject"}) assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) -@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"]) +@pytest.mark.parametrize("identity_claim", ["server_id", "subject"]) def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): claims = _unverified_claims(_sealed_token(_full_grant())) forged = _forge({**claims, identity_claim: ""}) @@ -463,9 +466,11 @@ def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): def test_empty_identity_and_key_fields_are_rejected_at_construction(): with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="", key_hash="hashed-key-123") + EnvelopeIdentity(server_id="", subject_type="key_hash", subject="hashed-key-123") with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="srv-456", key_hash="") + EnvelopeIdentity(server_id="srv-456", subject_type="key_hash", subject="") + with pytest.raises(ValidationError): + EnvelopeIdentity(server_id="srv-456", subject_type="not-a-subject-type", subject="x") with pytest.raises(ValidationError): EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) with pytest.raises(ValidationError): @@ -474,6 +479,20 @@ def test_empty_identity_and_key_fields_are_rejected_at_construction(): UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer") +def test_user_subject_identity_round_trips(): + """The user_id subject variant seals and opens with its discriminator intact, so the edge can + tell an interactively-minted (user) envelope from a scripted (key_hash) one and reload the right + kind of record.""" + identity = user_identity(server_id="srv-456", user_id="user-42") + sealed = mint_envelope(identity, _full_grant(), _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity.server_id == "srv-456" + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "user-42" + + def test_public_models_are_frozen(): sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW) assert isinstance(sealed, SealedEnvelope) @@ -484,4 +503,4 @@ def test_public_models_are_frozen(): with pytest.raises(ValidationError): opened.grant = _minimal_grant() with pytest.raises(ValidationError): - _IDENTITY.key_hash = "someone-elses-hash" + _IDENTITY.subject = "someone-elses-hash" 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 68466e624ec7..e43312e400e3 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 @@ -4441,7 +4441,8 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) assert isinstance(opened, BridgeEnvelopeAdmitted) - assert opened.identity.key_hash == "hashed-litellm-key-77" + assert opened.identity.subject_type == "key_hash" + assert opened.identity.subject == "hashed-litellm-key-77" assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" From 02e9c5631a88d0bdb52d1d4ccb1de21e9c29bede Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:58:52 -0700 Subject: [PATCH 2/5] feat(mcp): interactive SSO sign-in for dcr_bridge oauth_delegate DCR clients Completes the oauth_delegate bridge for real DCR clients (Claude Code, Claude Desktop), which send no litellm key and cannot use the scripted two-header path. On the short-circuit bridge arm the gateway now captures the SSO-authenticated litellm user from the browser session at /authorize and seals it into the OAuth state; at /callback it seals that user plus the upstream code into a gateway authorization code the client echoes back; at /token it recovers the user, exchanges the real upstream code, and mints a user-subject envelope. The user identity captured in the browser thus rides to the back-channel token call with nothing stored server-side, and admission opens the envelope under that user. The scripted key_hash path is unchanged (raw upstream code, key from the request); without a session the browser is sent through login first. --- .../mcp_server/discoverable_endpoints.py | 178 ++++++++++++++-- .../mcp_server/test_discoverable_endpoints.py | 200 +++++++++++++++++- 2 files changed, 353 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3e727ce95bbb..bd26abe0b26d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,7 +12,7 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, SecretStr, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -41,6 +41,7 @@ if TYPE_CHECKING: from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, EnvelopeKeys, UpstreamTokenGrant, ) @@ -98,6 +99,8 @@ def encode_state_with_base_url( code_challenge: Optional[str] = None, code_challenge_method: Optional[str] = None, client_redirect_uri: Optional[str] = None, + litellm_user_id: str | None = None, + mcp_server_id: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -108,6 +111,11 @@ def encode_state_with_base_url( code_challenge: PKCE code challenge from client code_challenge_method: PKCE code challenge method from client client_redirect_uri: Original redirect_uri from client + litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize + (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway + authorization code so the token mint can bind the envelope to this user + mcp_server_id: The bridge server the interactive flow targets, sealed alongside + litellm_user_id so the gateway code cannot be replayed against another server Returns: An encrypted string that encodes all values @@ -118,6 +126,8 @@ def encode_state_with_base_url( "code_challenge": code_challenge, "code_challenge_method": code_challenge_method, "client_redirect_uri": client_redirect_uri, + "litellm_user_id": litellm_user_id, + "mcp_server_id": mcp_server_id, } state_json = json.dumps(state_data, sort_keys=True) encrypted_state = encrypt_value_helper(state_json) @@ -145,6 +155,68 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +_BRIDGE_AUTH_CODE_PREFIX = "llm_bcode_" + + +class _BridgeAuthorizationCode(BaseModel): + """The identity and upstream code the gateway seals into the authorization code it hands a DCR + client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint.""" + + model_config = ConfigDict(frozen=True) + upstream_code: str = Field(min_length=1) + litellm_user_id: str = Field(min_length=1) + mcp_server_id: str = Field(min_length=1) + + +def is_bridge_authorization_code(code: str) -> bool: + """Cheap prefix check that ``code`` is a gateway-sealed bridge authorization code rather than a + raw upstream code, so the token endpoint can route without decrypting.""" + return code.startswith(_BRIDGE_AUTH_CODE_PREFIX) + + +def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str: + """Seal the upstream authorization code and the SSO-captured litellm user into a gateway + authorization code. The DCR client only echoes this opaque value back at the token endpoint; the + gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to + exchange with the upstream), so a litellm identity captured in the browser at authorize survives + to the back-channel token call with nothing stored server-side. Encrypted with the repo's + authenticated symmetric helper (the same family the OAuth state uses), so the client can neither + read nor forge it.""" + payload = json.dumps( + {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id}, + sort_keys=True, + ) + return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload) + + +def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None: + """Recover the sealed identity and upstream code, or ``None`` when ``code`` is not a gateway + bridge code or does not decrypt / validate. Total over hostile input: a raw upstream code (the + scripted two-header path) returns ``None`` and the caller falls through to the existing + behavior.""" + if not is_bridge_authorization_code(code): + return None + decrypted = decrypt_value_helper( + code[len(_BRIDGE_AUTH_CODE_PREFIX) :], "bridge_authorization_code", return_original_value=False + ) + if not isinstance(decrypted, str): + return None + try: + return _BridgeAuthorizationCode.model_validate_json(decrypted) + except ValidationError: + return None + + +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).""" + base_url = get_request_base_url(request) + return RedirectResponse(f"{base_url}/sso/key/generate") + + # LIT-4197: some upstream authorization servers reject an over-long ``state`` # (the encrypted OAuth session blob routinely exceeds their limit). The upstream # only needs an opaque value it echoes back on ``/callback``, so we forward a @@ -697,12 +769,31 @@ async def authorize_with_server( parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) + + # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in + # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and + # carry it to the back-channel token mint. Seal the SSO user and the target server into the state; + # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a + # litellm key, so the browser session is the only identity source; without one there is nothing to + # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. + litellm_user_id: str | None = None + if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate: + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import + _user_id_from_session_cookie, + ) + + litellm_user_id = _user_id_from_session_cookie(request) + if litellm_user_id is None: + return _redirect_to_litellm_login(request) + encoded_state = encode_state_with_base_url( base_url=base_url, original_state=state, code_challenge=code_challenge, code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, + litellm_user_id=litellm_user_id, + mcp_server_id=mcp_server.server_id if litellm_user_id else None, ) relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) @@ -824,11 +915,14 @@ def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenG @dataclass(frozen=True, slots=True) class _BridgeMintReady: - """Everything the seal needs, resolved once before the exchange: the authorizing key hash and the - master-key-derived envelope keys. Passing this forward means identity resolution and key derivation - happen exactly once, and ``_finish_bridge_mint`` has no preconditions left that could fail.""" - - key_hash: str + """Everything the seal needs, resolved once before the exchange: the identity to bind the envelope + to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted + two-header client (resolved from the litellm key it presents) or a user_id subject for the + interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal + serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to + fail.""" + + identity: "EnvelopeIdentity" keys: "EnvelopeKeys" @@ -844,8 +938,8 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: status, code, desc = ( 400, "invalid_request", - "this server issues a gateway-bound credential; send a litellm credential " - "(x-litellm-api-key or Authorization) on the token request", + "this server issues a gateway-bound credential; complete the interactive sign-in, or " + "send a litellm credential (x-litellm-api-key or Authorization) on the token request", ) case "unsupported_grant": status, code, desc = ( @@ -923,18 +1017,30 @@ def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _Br assert_never(rejection) -async def _prepare_bridge_mint(request: Request, grant_type: str) -> "_BridgeMintReady | _BridgeMintError": +async def _prepare_bridge_mint( + request: Request, + grant_type: str, + mcp_server: MCPServer, + bridge_identity: _BridgeAuthorizationCode | None = None, +) -> "_BridgeMintReady | _BridgeMintError": """Phase 1, BEFORE the upstream exchange: reject a grant this mint does not support, confirm the gateway can mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready context or a precise failure value. Running before the exchange is what makes every - failure here fail closed without consuming the single-use code or rotating a refresh token. A bridge - server issues only envelopes and seals no upstream refresh_token, so the client holds none to - present: the refresh_token grant is rejected up front rather than exchanged (which could rotate the - upstream credential) and its result then discarded. Identity-resolution failures keep their origin - so the mapper statuses each truthfully.""" + failure here fail closed without consuming the single-use code. + + Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged + authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway + authorization code) and mints a user subject. The scripted two-header client presents a litellm key + on the token request instead, so its identity is the active key's hash and mints a key_hash subject. + A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; + neither source present is ``no_identity``.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + key_hash_identity, + user_identity, + ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import master_key, ) @@ -943,16 +1049,21 @@ async def _prepare_bridge_mint(request: Request, grant_type: str) -> "_BridgeMin return "unsupported_grant" if not master_key: return "not_configured" + keys = envelope_keys_from_master_key(master_key) + if bridge_identity is not None: + identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) + return _BridgeMintReady(identity=identity, keys=keys) resolved = await _resolve_active_litellm_key(request) if not isinstance(resolved, _ResolvedKey): return _key_resolution_failure_to_mint_error(resolved) - return _BridgeMintReady(key_hash=resolved.key_hash, keys=envelope_keys_from_master_key(master_key)) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) + return _BridgeMintReady(identity=identity, keys=keys) def _finish_bridge_mint( ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime ) -> "JSONResponse | _BridgeMintError": - """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope using + """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope under the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the upstream token with nothing stored server-side. The only failures here are properties of the upstream response (no usable token, an already-expired lifetime, or a token too large to seal), @@ -963,14 +1074,12 @@ def _finish_bridge_mint( from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import SealedEnvelope, UpstreamTokenGrant, - key_hash_identity, ) grant = _bridge_grant_from_token_response(token_response) if not isinstance(grant, UpstreamTokenGrant): return _upstream_rejection_to_mint_error(grant) - identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=ready.key_hash) - sealed = build_bridge_token_response(identity, grant, ready.keys, now) + sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now) if not isinstance(sealed, SealedEnvelope): return "too_large" # Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the @@ -1014,6 +1123,7 @@ async def exchange_token_with_server( except TokenEndpointAuthConfigError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + bridge_identity: _BridgeAuthorizationCode | None = None if grant_type == "refresh_token": if not refresh_token: raise HTTPException( @@ -1033,6 +1143,19 @@ async def exchange_token_with_server( status_code=400, detail="code is required for authorization_code grant", ) + # Interactive dcr_bridge oauth_delegate: the client presents the gateway authorization code the + # callback sealed. Recover the SSO user and the real upstream code from it; the upstream exchange + # below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the + # sealed server to this request so a code minted for one bridge server cannot be spent at another. + # A raw upstream code (scripted path) opens to None and the code is used as-is. + bridge_identity = open_bridge_authorization_code(code) + if bridge_identity is not None: + if bridge_identity.mcp_server_id != mcp_server.server_id: + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) + code = bridge_identity.upstream_code bridge_token_relay = _dcr_bridge_relays_client_registration(mcp_server) if bridge_token_relay and not redirect_uri: raise HTTPException( @@ -1058,7 +1181,7 @@ async def exchange_token_with_server( # phase 3. A failure here returns without ever touching the upstream credential. bridge_mint_ready: _BridgeMintReady | None = None if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: - prepared = await _prepare_bridge_mint(request, grant_type) + prepared = await _prepare_bridge_mint(request, grant_type, mcp_server, bridge_identity) if not isinstance(prepared, _BridgeMintReady): return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared @@ -1706,7 +1829,20 @@ async def callback( # states while permitting same-origin / allowlisted clients. redirect_uri = _get_validated_client_redirect_uri(request, state_data) - params = {"code": code, "state": original_state} + # Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step + # captured. Instead of forwarding the raw upstream code (which the client would present at the + # token endpoint with no way to prove who signed in), seal the user and the upstream code into a + # gateway authorization code and forward THAT. The token endpoint decrypts it to bind the + # envelope to this user. Every other flow forwards the raw code unchanged. + litellm_user_id = state_data.get("litellm_user_id") + mcp_server_id = state_data.get("mcp_server_id") + forwarded_code = code + if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: + forwarded_code = seal_bridge_authorization_code( + upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id + ) + + params = {"code": forwarded_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) response = RedirectResponse(url=complete_returned_url, status_code=302) _clear_oauth_state_cookie(response, request, state) 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 e43312e400e3..966619ee6b88 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 @@ -4365,7 +4365,7 @@ async def test_register_bridge_relay_never_persists(): _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" -async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_client_out=None): +async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _ResolvedKey, exchange_token_with_server, @@ -4398,13 +4398,17 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie request=_bridge_mock_request(), mcp_server=server, grant_type="authorization_code", - code="auth-code", + code=code, redirect_uri="https://claude.ai/api/mcp/auth_callback", client_id="dcr-client-123", client_secret=None, code_verifier="verifier", ) - if server.is_oauth_delegate and server.is_dcr_bridge: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import is_bridge_authorization_code + + # The key_hash path resolves the presented litellm key; the interactive SSO path recovers identity + # from the gateway authorization code instead, so it never awaits the resolver. + if server.is_oauth_delegate and server.is_dcr_bridge and not is_bridge_authorization_code(code): key_resolver.assert_awaited_once() else: key_resolver.assert_not_awaited() @@ -4446,6 +4450,193 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" +def test_bridge_authorization_code_round_trips_and_rejects_hostile_input(): + """The gateway authorization code seals and recovers the upstream code and the SSO user, and is + total over hostile input: a raw upstream code (scripted path) opens to None, and a tampered or + non-gateway value opens to None rather than raising.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + is_bridge_authorization_code, + open_bridge_authorization_code, + seal_bridge_authorization_code, + ) + + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + sealed = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-9", mcp_server_id="srv-1" + ) + assert is_bridge_authorization_code(sealed) + opened = open_bridge_authorization_code(sealed) + assert opened is not None + assert opened.upstream_code == "up-code" + assert opened.litellm_user_id == "sso-user-9" + assert opened.mcp_server_id == "srv-1" + assert open_bridge_authorization_code("raw-upstream-code") is None + assert open_bridge_authorization_code(sealed[:-4] + "aaaa") is None + + +@pytest.mark.asyncio +async def test_interactive_bridge_token_exchange_mints_user_subject_envelope(): + """An interactive dcr_bridge oauth_delegate exchange (the client presents the gateway code the + callback sealed, and NO litellm key) mints an envelope bound to the SSO-captured user: it opens + to a user_id subject, and the upstream exchange used the real upstream code recovered from the + gateway code, not the sealed wrapper.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="REAL-UPSTREAM-CODE", litellm_user_id="sso-user-42", mcp_server_id=server.server_id + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _exchange_for_bridge_server( + server, upstream, key_hash=None, code=gateway_code, fake_client_out=captured + ) + + token = json.loads(response.body)["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "sso-user-42" + assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" + assert captured["client"].post.call_args.kwargs["data"]["code"] == "REAL-UPSTREAM-CODE" + + +@pytest.mark.asyncio +async def test_interactive_bridge_gateway_code_for_another_server_is_rejected_400(): + """A gateway authorization code is bound to the server it was minted for: presenting it at another + server's token endpoint is a 400, so a code cannot be replayed across a server boundary.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-42", mcp_server_id="a-different-server-id" + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + with pytest.raises(HTTPException) as exc: + await _exchange_for_bridge_server(server, upstream, key_hash=None, code=gateway_code) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_seals_sso_user_into_state(): + """On the short-circuit bridge oauth_delegate arm, authorize captures the SSO user from the UI + session cookie and seals it (and the target server) into the encrypted OAuth state, so the + callback can later mint a user-bound gateway code; it still proceeds to the upstream redirect.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + captured: dict = {} + + def _capture(**kwargs): + captured.update(kwargs) + return "mocked_encrypted_state" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value="sso-user-42", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encode_state_with_base_url", + side_effect=_capture, + ), + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + + assert captured["litellm_user_id"] == "sso-user-42" + assert captured["mcp_server_id"] == server.server_id + assert "/sso/key/generate" not in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_without_session_redirects_to_login(): + """Without a UI session there is no identity to bind, so the short-circuit bridge oauth_delegate + authorize sends the browser through litellm login instead of proceeding to the upstream.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + with patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value=None, + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + assert "/sso/key/generate" in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_callback_seals_user_into_gateway_code(): + """When the OAuth state carries the captured SSO user, the callback forwards a gateway + authorization code (sealing the user and upstream code) to the client instead of the raw upstream + code, so the client's later token call can prove who signed in.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + is_bridge_authorization_code, + ) + + state_data = { + "original_state": "client-state", + "client_redirect_uri": "http://127.0.0.1:60108/cb", + "base_url": "http://127.0.0.1:60108/cb", + "litellm_user_id": "sso-user-42", + "mcp_server_id": "bridge_srv", + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_encoded_oauth_state", + return_value="enc", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash", + return_value=state_data, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._get_validated_client_redirect_uri", + return_value="http://127.0.0.1:60108/cb", + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await callback(request=_bridge_mock_request(), code="REAL-UPSTREAM-CODE", state="relay") + + forwarded_code = parse_qs(urlparse(response.headers["location"]).query)["code"][0] + assert is_bridge_authorization_code(forwarded_code) + + @pytest.mark.asyncio async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm_identity(): """Without a resolvable litellm identity on the token request, the exchange must not mint an @@ -4727,10 +4918,11 @@ def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity from litellm.types.mcp import MCPAuth ready = _BridgeMintReady( - key_hash="hashed-litellm-key-77", + identity=key_hash_identity(server_id="bridge_srv", key_hash="hashed-litellm-key-77"), keys=envelope_keys_from_master_key(_BRIDGE_MASTER_KEY), ) response = _finish_bridge_mint( From f96899ae2b793f049a01a91db648980cd85021d2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 11:08:08 -0700 Subject: [PATCH 3/5] fix(mcp): classify the user-subject reload's errors like the key path (503 outage, 401 missing) _reload_admitted_user mirrored only part of _reload_admitted_key's error contract: it caught ProxyException and HTTPException but had no arm for anything else, so a transient DB outage surfaced as an opaque 500 instead of the retryable 503 the key path guarantees, and a missing user surfaced as a 500 too. The missing-user case is the subtle one: get_user_object raises a bare Exception for a deleted user (not a ProxyException like get_key_object does for a missing key), so the ProxyException/HTTPException clause never caught it and the user_object-is-None branch it was supposed to hit is unreachable on the production path. Add the same except-Exception arm the key path uses, with the one deliberate difference the differing get_user_object contract requires: a database-service-unavailable error still raises the retryable 503, while a missing user or any other non-outage resolution failure fails closed as a 401 rather than propagating as a 500. The regression tests now drive the real behavior (get_user_object raising) rather than a None return that never happens in production, and cover both the 503 outage and the 401 missing-user paths. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 13 +++++-- .../auth/test_user_api_key_auth_mcp.py | 34 +++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index faec35db41a0..01861abbbc06 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 @@ -605,8 +605,14 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: caller's centralized policy gate then enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No team is bound; a user may belong to many teams or none, so the envelope grants the - user's own access rather than silently selecting one team's scope. A missing user - fails closed with a 401 rather than admitting an unresolved identity.""" + user's own access rather than silently selecting one team's scope. + + Error handling mirrors the key path's retryable-503 contract, with one deliberate + difference: ``get_key_object`` raises a ``ProxyException`` for a missing key, but + ``get_user_object`` raises a bare ``Exception`` for a missing user (it does not surface as a + ``ProxyException``/``HTTPException``). So a transient DB outage still surfaces as a retryable + 503 via ``_raise_503_if_db_unavailable``, while a missing user, or any other non-outage + resolution failure, fails closed as a 401 rather than propagating as an opaque 500.""" from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -621,6 +627,9 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: ) except (ProxyException, HTTPException): raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + except Exception as e: # noqa: BLE001 # DB outage -> retryable 503; a missing user (bare Exception) or any other resolution failure -> fail closed 401, never an opaque 500 + MCPRequestHandler._raise_503_if_db_unavailable(e) + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None if user_object is None: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: 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 a6affe5496c0..a441e6a154bd 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 @@ -5117,8 +5117,10 @@ async def test_user_subject_envelope_admits_under_the_reloaded_user(self): } async def test_user_subject_envelope_missing_user_fails_closed_401(self): - """A user_id envelope whose user has since been deleted must fail closed: get_user_object - resolves None, so admission 401s instead of admitting an unresolved identity.""" + """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. + get_user_object raises a bare Exception for a missing user (it does not return None on the + production path), so the reload must catch it and fail closed rather than let it propagate as an + opaque 500. Regression for the missing-user path surfacing as a 500.""" envelope = self._mint_bridge_envelope(user_id="ghost-user") scope = { "type": "http", @@ -5129,7 +5131,7 @@ async def test_user_subject_envelope_missing_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=None), + self._patch_user_reload(side_effect=Exception("user not found")), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: @@ -5137,6 +5139,28 @@ async def test_user_subject_envelope_missing_user_fails_closed_401(self): assert exc_info.value.status_code == 401 + async def test_user_subject_envelope_db_outage_is_retryable_503(self): + """A transient database outage while reloading the envelope's user is a retryable 503, not an + opaque 500, matching the key path's contract so an interactive DCR client retries instead of + treating a live identity as invalid. Regression for the user reload dropping the 503 arm.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + 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(side_effect=ConnectionError("auth database unreachable")), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 503 + async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries scim_active False, so admission 401s rather than letting an offboarded user keep tool access @@ -5151,9 +5175,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", metadata={"scim_active": False})), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: From c46863b0e64a4962b84ddf41dc1a9faf7faac3dd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 11:18:53 -0700 Subject: [PATCH 4/5] fix(mcp): admit a user-subject envelope with the user's own MCP object permission _reload_admitted_user returned a bare UserAPIKeyAuth(user_id=...), so the shared get_allowed_mcp_servers found no key/team/object-permission grants and an interactive SSO client could admit successfully yet see zero tools on a normal (allow_all_keys=False) server. The key path returns the full key record whose object permission drives that computation; the user path dropped it. Resolve the user's own MCP object permission and put it on the returned auth, so the same get_allowed_mcp_servers the key path uses grants the user their litellm-granted servers and access groups. This reuses get_object_permission (the id-to-grants resolver keys and teams already use) and does not duplicate any permission logic; get_user_object does not load object_permission, so it is resolved from the user's object_permission_id the same way the key and team paths 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 get_allowed_mcp_servers does not do off one auth object, and faking one here would be the kind of half-measure that spawns more bugs. Tests cover the user's object permission riding onto the admitted auth, and the existing admit/SCIM/missing-user/503 cases still hold. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 32 ++++++++++--- .../auth/test_user_api_key_auth_mcp.py | 46 ++++++++++++++++++- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 01861abbbc06..ac3e4439fa1b 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 @@ -601,11 +601,14 @@ 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`` rides on the returned ``UserAPIKeyAuth`` and the - caller's centralized policy gate then enforces the user's live budget and org state, - and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No - team is bound; a user may belong to many teams or none, so the envelope grants the - user's own access rather than silently selecting one team's scope. + 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 + ``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. Error handling mirrors the key path's retryable-503 contract, with one deliberate difference: ``get_key_object`` raises a ``ProxyException`` for a missing key, but @@ -613,7 +616,7 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: ``ProxyException``/``HTTPException``). So a transient DB outage still surfaces as a retryable 503 via ``_raise_503_if_db_unavailable``, while a missing user, or any other non-outage resolution failure, fails closed as a 401 rather than propagating as an opaque 500.""" - from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client is None: @@ -634,7 +637,22 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: raise HTTPException(status_code=401, detail="Invalid or expired credential") - return UserAPIKeyAuth(user_id=user_object.user_id) + # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared + # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same + # get_object_permission resolver the key and team paths use; no permission logic is duplicated. + object_permission = user_object.object_permission + if user_object.object_permission_id and object_permission is None: + object_permission = await get_object_permission( + object_permission_id=user_object.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return UserAPIKeyAuth( + user_id=user_object.user_id, + user_role=user_object.user_role, + object_permission=object_permission, + object_permission_id=user_object.object_permission_id, + ) @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: 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 a441e6a154bd..0b1905689ac6 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 @@ -5103,7 +5103,13 @@ async def test_user_subject_envelope_admits_under_the_reloaded_user(self): 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="sso-user-7", metadata={"scim_active": True}) + return_value=MagicMock( + user_id="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=None, + object_permission_id=None, + ) ) as get_user_object, ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() @@ -5116,6 +5122,44 @@ async def test_user_subject_envelope_admits_under_the_reloaded_user(self): "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_user_subject_envelope_carries_the_users_mcp_object_permission(self): + """The admitted user's own MCP object permission rides on the returned auth so the shared + get_allowed_mcp_servers grants the user their litellm-granted servers, rather than admitting a + bare user with no MCP access. Regression for the signed-in SSO client getting zero tools because + the reload dropped the user's object permission.""" + object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-user-7", mcp_servers=["bridge_delegate_server"] + ) + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + 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="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=object_permission, + object_permission_id="op-user-7", + ) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, _headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert auth_result.object_permission is not None + assert auth_result.object_permission.mcp_servers == ["bridge_delegate_server"] + async def test_user_subject_envelope_missing_user_fails_closed_401(self): """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. get_user_object raises a bare Exception for a missing user (it does not return None on the From 61c7e706dd24db17131053445c36bd00ad8fa228 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 13:25:41 -0700 Subject: [PATCH 5/5] fix(mcp): classify get_user_object's wrapped DB outage across the exception chain get_user_object catches every DB failure in a broad except and re-raises a bare ValueError (litellm/proxy/auth/auth_checks.py), so a real outage and a missing user look identical and the original error survives only as __context__. The dcr_bridge admission path keyed its 503-vs-401 decision on the exception type, so a transient outage during a user-subject reload surfaced as a 401 rather than a retryable 503, and the regression test injected a raw ConnectionError, a shape get_user_object never produces, so it passed on a fiction Add PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain, which walks __cause__/__context__ (bounded and cycle-safe) the PEP 3134 way, and route _raise_503_if_db_unavailable through it. Move the user's object_permission resolution inside the single classified try so an outage there is a 503 too, never an opaque 500. Pin get_user_object's wrapping with a contract test that drives the real function, and drive the reload tests with that same faithful shape so a chain-blind regression fails them --- .../mcp_server/auth/user_api_key_auth_mcp.py | 46 +++++++++++-------- litellm/proxy/db/exception_handler.py | 30 ++++++++++++ .../auth/test_user_api_key_auth_mcp.py | 33 ++++++++++--- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++ .../proxy/db/test_exception_handler.py | 40 ++++++++++++++++ 5 files changed, 156 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index ac3e4439fa1b..421f1dcfbeab 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 @@ -610,12 +610,16 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: ``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. - Error handling mirrors the key path's retryable-503 contract, with one deliberate - difference: ``get_key_object`` raises a ``ProxyException`` for a missing key, but - ``get_user_object`` raises a bare ``Exception`` for a missing user (it does not surface as a - ``ProxyException``/``HTTPException``). So a transient DB outage still surfaces as a retryable - 503 via ``_raise_503_if_db_unavailable``, while a missing user, or any other non-outage - resolution failure, fails closed as a 401 rather than propagating as an opaque 500.""" + Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a + type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key + and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a + bare ``ValueError``, so a missing user and a real outage look identical and the original error + survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause + chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any + other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The + object-permission load shares this one boundary, so an outage there is classified the same + way (``get_object_permission`` itself swallows a failed load to ``None``, matching how + ``get_key_object`` best-effort-loads a key's object permission).""" from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -628,25 +632,25 @@ async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: user_api_key_cache=user_api_key_cache, user_id_upsert=False, ) + # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared + # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same + # get_object_permission resolver the key and team paths use; no permission logic is duplicated. + object_permission = user_object.object_permission if user_object is not None else None + if user_object is not None and object_permission is None and user_object.object_permission_id: + object_permission = await get_object_permission( + object_permission_id=user_object.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) except (ProxyException, HTTPException): raise HTTPException(status_code=401, detail="Invalid or expired credential") from None - except Exception as e: # noqa: BLE001 # DB outage -> retryable 503; a missing user (bare Exception) or any other resolution failure -> fail closed 401, never an opaque 500 + except Exception as e: # noqa: BLE001 # a DB outage anywhere in the resolution is a retryable 503, not an opaque 500; anything else fails closed as 401 MCPRequestHandler._raise_503_if_db_unavailable(e) raise HTTPException(status_code=401, detail="Invalid or expired credential") from None if user_object is None: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: raise HTTPException(status_code=401, detail="Invalid or expired credential") - # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared - # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same - # get_object_permission resolver the key and team paths use; no permission logic is duplicated. - object_permission = user_object.object_permission - if user_object.object_permission_id and object_permission is None: - object_permission = await get_object_permission( - object_permission_id=user_object.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) return UserAPIKeyAuth( user_id=user_object.user_id, user_role=user_object.user_role, @@ -697,10 +701,14 @@ def _raise_503_if_db_unavailable(e: Exception) -> None: """Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure (401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``, - which renders a service-unavailable database error as 503 on the standard pipeline.""" + which renders a service-unavailable database error as 503 on the standard pipeline. + + Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object`` + re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception + would miss a real outage wrapped inside it.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): raise HTTPException( status_code=503, detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 3a93896a2064..e4c565e44648 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,10 @@ ) from litellm.secret_managers.main import str_to_bool +# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain. +# Real exception chains are a few links deep; the cap also makes the walk cycle-safe. +_MAX_EXCEPTION_CHAIN_DEPTH = 20 + class PrismaDBExceptionHandler: """ @@ -218,6 +222,32 @@ def is_database_service_unavailable_error(e: Exception) -> bool: ), ) + @staticmethod + def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool: + """Like ``is_database_service_unavailable_error`` but also walks the + ``__cause__`` / ``__context__`` chain. + + ``is_database_service_unavailable_error`` classifies a single exception + by type, which a caller that catches a raw DB failure and re-raises a + domain exception of a different type defeats. ``get_user_object`` in + ``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps + every DB error, a genuine outage included, in a bare ``ValueError`` + whose original error survives only as ``__context__``. A type check on + the ``ValueError`` misses the outage, so the caller would mistake an + infrastructure fault for an auth failure. Walking the chain recovers the + real signal, which is the PEP 3134 way to inspect a wrapped cause. + + The walk is depth-bounded, which also makes it cycle-safe. + """ + current: BaseException | None = e + for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): + if not isinstance(current, Exception): + return False + if PrismaDBExceptionHandler.is_database_service_unavailable_error(current): + return True + current = current.__cause__ or current.__context__ + return False + @staticmethod def handle_db_exception(e: Exception): """ 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 0b1905689ac6..6f132aaae9cd 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 @@ -5022,6 +5022,22 @@ def _patch_user_reload(*, return_value=None, side_effect=None): ): yield get_user_object + @staticmethod + def _wrapped_user_lookup_error(original: BaseException) -> ValueError: + """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it + catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the + original error (a missing-user Exception or a real outage) survives only as ``__context__``. + Injecting a raw ConnectionError/Exception instead would exercise a shape production never + produces and let a chain-blind outage classifier pass. That wrapping fidelity is itself pinned by + test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + try: + raise original + except BaseException: + try: + raise ValueError(f"User doesn't exist in db. Got error - {original}") + except ValueError as wrapped: + return wrapped + @staticmethod def _mcp_request(path="/mcp/bridge_delegate_server"): """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how @@ -5162,9 +5178,10 @@ async def test_user_subject_envelope_carries_the_users_mcp_object_permission(sel async def test_user_subject_envelope_missing_user_fails_closed_401(self): """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. - get_user_object raises a bare Exception for a missing user (it does not return None on the - production path), so the reload must catch it and fail closed rather than let it propagate as an - opaque 500. Regression for the missing-user path surfacing as a 500.""" + get_user_object catches the missing row and re-raises a bare ValueError (it does not return None + on the production path), so the reload must fail closed rather than let it propagate as an opaque + 500, and must not mistake the wrapped ValueError for a DB outage. Regression for the missing-user + path surfacing as a 500.""" envelope = self._mint_bridge_envelope(user_id="ghost-user") scope = { "type": "http", @@ -5175,7 +5192,7 @@ async def test_user_subject_envelope_missing_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(side_effect=Exception("user not found")), + self._patch_user_reload(side_effect=self._wrapped_user_lookup_error(Exception())), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: @@ -5186,7 +5203,9 @@ async def test_user_subject_envelope_missing_user_fails_closed_401(self): async def test_user_subject_envelope_db_outage_is_retryable_503(self): """A transient database outage while reloading the envelope's user is a retryable 503, not an opaque 500, matching the key path's contract so an interactive DCR client retries instead of - treating a live identity as invalid. Regression for the user reload dropping the 503 arm.""" + treating a live identity as invalid. get_user_object wraps the outage in a bare ValueError, so this + exercises the chain-aware classifier; a raw ConnectionError would falsely pass even the old + chain-blind check because it is an OSError. Regression for the user reload dropping the 503 arm.""" envelope = self._mint_bridge_envelope(user_id="sso-user-7") scope = { "type": "http", @@ -5197,7 +5216,9 @@ async def test_user_subject_envelope_db_outage_is_retryable_503(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(side_effect=ConnectionError("auth database unreachable")), + self._patch_user_reload( + side_effect=self._wrapped_user_lookup_error(ConnectionError("auth database unreachable")) + ), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d12ff20ee5b0..3433d7dc2d3f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -701,6 +701,38 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch): assert creation_args["user_role"] == "internal_user" +@pytest.mark.asyncio +async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(): + """Pin get_user_object's exception contract: it catches every DB failure in a broad except and + re-raises a bare ValueError, so a real outage survives only as __context__ rather than as the + exception type. The MCP dcr_bridge admission and refresh paths depend on this to tell a transient + outage (retry, 503) from a missing user (fail closed), which is why they classify across the cause + chain instead of the top exception's type. If this wrapping ever changes, that classification must + change with it, so this test guards the contract the callers rely on.""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma_client = MagicMock() + mock_prisma_client.db = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + side_effect=ConnectionError("can't reach database server") + ) + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(ValueError) as exc_info: + await get_user_object( + user_id="outage-contract-probe-user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=False, + proxy_logging_obj=None, + ) + + assert isinstance(exc_info.value.__context__, ConnectionError) + + @pytest.mark.asyncio async def test_get_user_object_upsert_includes_user_email(): """Test that user_email is included when creating a new user via get_user_object upsert""" diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 0634a01326c6..23099177812e 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -286,6 +286,46 @@ def test_is_database_service_unavailable_error_excludes_non_infra(error): ) +def _wrapped_like_get_user_object(original): + """Reproduce get_user_object's exception contract (litellm/proxy/auth/auth_checks.py): it catches + every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original error + survives only as ``__context__``. Building it by raising inside an ``except`` sets ``__context__`` + exactly as production does.""" + try: + raise original + except BaseException: + try: + raise ValueError("User doesn't exist in db. Got error - x") + except ValueError as wrapped: + return wrapped + + +def test_is_database_service_unavailable_error_in_chain_sees_through_wrapping(): + """The chain-aware classifier must see a real outage that a caller wrapped in a different type. + get_user_object turns a connection error into a bare ValueError whose type check reads as non-infra, + so the single-exception check returns False and only the chain walk recovers the outage. A missing + user (whose wrapped cause is a plain Exception) must stay non-infra on both.""" + outage = _wrapped_like_get_user_object(ConnectionError("can't reach database server")) + missing_user = _wrapped_like_get_user_object(Exception()) + + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(outage) is False + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(outage) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(missing_user) is False + # parity: a raw outage with no wrapper is still an outage, and a plain ValueError is not + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ConnectionError("boom")) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ValueError("nope")) is False + + +def test_is_database_service_unavailable_error_in_chain_terminates_on_a_cause_cycle(): + """The walk must terminate on a pathological __cause__ cycle rather than hang. Neither link is an + outage, so the bounded walk returns False instead of looping forever.""" + first = ValueError("first") + second = ValueError("second") + first.__cause__ = second + second.__cause__ = first + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(first) is False + + def test_is_database_service_unavailable_error_asyncpg(monkeypatch): """asyncpg connection/interface errors map to service-unavailable. asyncpg is not a hard dependency, so inject a stand-in module to exercise the