diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 13db808a934..8346ed6b3aa 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -547,9 +547,10 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No the egress side. No DB connection is a gateway fault (``unresolvable``) and a database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / - ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user, because - ``get_user_object`` raises a bare ``Exception`` for a deleted user rather than a ``ProxyException``, - so a missing user must not be misclassified as an opaque gateway fault.""" + ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` + catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look + identical, the original error surviving only as ``__context__``), so the outage check walks the cause + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -575,8 +576,8 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No ) except (ProxyException, HTTPException): return "no_active_key" - except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (bare Exception) or any other resolution failure fails closed as no_active_key, never a 500 - if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): return "unavailable" verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) return "no_active_key" @@ -587,15 +588,52 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No return None +async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool: + """True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an + offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``. + A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``), + matching admission and the standard builder: a key may outlive its owner record, and a transient DB + blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal.""" + if key.user_id is None: + return False + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return False + try: + owner = await get_user_object( + user_id=key.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key + verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__) + return False + return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False + + async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None": """Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type: a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is - active or a precise failure otherwise, so revocation (a blocked key, a deactivated user) gates renewal - for either identity source, the same way admission gates the egress.""" + active or a precise failure otherwise, so revocation gates renewal for either identity source the same + way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring + admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a + deactivated or deleted user all fail closed to ``no_active_key``.""" match identity.subject_type: case "key_hash": reloaded = await _reload_active_key_by_hash(identity.subject) - return None if isinstance(reloaded, _ResolvedKey) else reloaded + if not isinstance(reloaded, _ResolvedKey): + return reloaded + if await _key_owner_scim_deactivated(reloaded.key): + return "no_active_key" + return None case "user_id": return await _reload_active_user_by_id(identity.subject) case _: @@ -1242,7 +1280,11 @@ def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | """Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal. Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in`` (the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and - bounds the refresh envelope's TTL.""" + bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed + (``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead + token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to + an access-only response (the client re-authenticates at access expiry), mirroring how + :func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import RefreshCredential, ) @@ -1252,8 +1294,10 @@ def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | refresh = token_response.get("refresh_token") if not isinstance(refresh, str) or not refresh: return None - scope = token_response.get("scope") lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) + if lifetime == "expired": + return None + scope = token_response.get("scope") return RefreshCredential( refresh_token=SecretStr(refresh), scope=scope if isinstance(scope, str) and scope else None, @@ -1422,6 +1466,20 @@ async def exchange_token_with_server( "does not send yet (tracked as LIT-4339)", mcp_server.server_id, ) + upstream_rejected_bridge_refresh = ( + is_bridge + and grant_type == "refresh_token" + and exc.response.status_code == 400 + and "invalid_grant" in exc.response.text + ) + if upstream_rejected_bridge_refresh: + verbose_logger.info( + "bridge refresh: the upstream rejected the sealed refresh token for server=%s with " + "invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client " + "re-runs authorization_code rather than an opaque upstream error", + mcp_server.server_id, + ) + return _bridge_mint_error_response("invalid_refresh") raise token_response = response.json() 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 d82c3b8896f..8571e2a7174 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 @@ -5047,6 +5047,137 @@ async def test_revalidate_active_subject_dispatches_on_subject_type(): key_reload2.assert_not_awaited() +def test_upstream_refresh_credential_expired_refresh_token_is_not_sealed(): + """An upstream that reports its refresh token already elapsed (refresh_expires_in non-positive) must + not be sealed: _upstream_refresh_credential returns None so the exchange degrades to an access-only + response, mirroring how the access grant refuses an already-elapsed access token rather than capping a + dead token to the full refresh TTL. A live or unspecified lifetime still yields a credential.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _upstream_refresh_credential + + assert _upstream_refresh_credential({"access_token": "A", "refresh_token": "R", "refresh_expires_in": 0}) is None + assert _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": -5}) is None + live = _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": 1800}) + assert live is not None and live.expires_in == 1800 + unspecified = _upstream_refresh_credential({"refresh_token": "R"}) + assert unspecified is not None and unspecified.expires_in is None + + +@pytest.mark.asyncio +async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant(): + """When the sealed upstream refresh token has been revoked or expired at the IdP, the upstream returns + 400 invalid_grant. The bridge refresh path maps that to an RFC 6749 invalid_grant response so the OAuth + client re-runs authorization_code, rather than surfacing the opaque upstream error it cannot act on.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="LIVE-ENVELOPE-REFRESH") + + error_response = MagicMock() + error_response.status_code = 400 + error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}' + error_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) + ) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=error_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + new=AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_env, + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_globals): + """A key_hash refresh envelope whose key is still active but whose OWNING user was SCIM-deactivated must + fail closed to no_active_key, mirroring how admission's _reject_if_admitted_owner_scim_deactivated + revokes an offboarded owner's key. Without this, an offboarded user keeps renewing a live key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="offboarded-owner")) + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": False})), + ), + ): + result = await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) + + assert result == "no_active_key" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fails_open(proxy_globals): + """The key-owner SCIM gate blocks only an explicit scim_active False: an active owner renews (None), and + a missing owner (get_user_object's wrapped ValueError) fails OPEN, since a key may outlive its owner + record and a transient blip must not revoke a live key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="live-owner")) + identity = key_hash_identity(server_id="s", key_hash="kh") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": True})), + ), + ): + assert await _revalidate_active_subject(identity) is None + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ), + ): + assert await _revalidate_active_subject(identity) is None + + @pytest.mark.asyncio async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): """master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a @@ -5857,25 +5988,28 @@ async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_glob assert await _resolve_active_litellm_key(request) == "unresolvable" +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. The wrapping fidelity is 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 + + @pytest.mark.asyncio async def test_reload_active_user_by_id_missing_user_is_no_active_key(proxy_globals): """A user_id refresh envelope whose user has been deleted must fail closed to no_active_key (the - refresh path maps it to invalid_grant), not unresolvable/500. get_user_object raises a bare Exception - for a missing user, so a missing user must not be misclassified as an opaque gateway fault.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - - proxy_globals.user_api_key_cache = UserApiKeyCache() - proxy_globals.prisma_client = object() - - with patch("litellm.proxy.auth.auth_checks.get_user_object", new=AsyncMock(side_effect=Exception("no user"))): - assert await _reload_active_user_by_id("gone-user") == "no_active_key" - - -@pytest.mark.asyncio -async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): - """A transient DB outage while re-validating the user on refresh is a retryable outage, distinct from - a missing user, so the refresh path can surface a 503 rather than blaming the caller.""" + refresh path maps it to invalid_grant), not unresolvable/500. get_user_object catches the missing row + and re-raises a bare ValueError, so a missing user must not be misclassified as a DB outage or an + opaque gateway fault.""" from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5884,7 +6018,26 @@ async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): with patch( "litellm.proxy.auth.auth_checks.get_user_object", - new=AsyncMock(side_effect=ConnectionError("user database unreachable")), + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ): + assert await _reload_active_user_by_id("gone-user") == "no_active_key" + + +@pytest.mark.asyncio +async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): + """A transient DB outage while re-validating the user on refresh is a retryable outage, distinct from + a missing user, so the refresh path surfaces "unavailable" (a 503) rather than blaming the caller. + get_user_object wraps the outage in a bare ValueError, so this exercises the chain-aware classifier; a + raw ConnectionError would falsely pass even a chain-blind check because it is an OSError.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(ConnectionError("user database unreachable"))), ): assert await _reload_active_user_by_id("sso-user-7") == "unavailable"