scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(mcp): make the dcr_bridge refresh path fail correctly on outages, dead tokens, and revoked owners

Browse source 
Four fixes to the refresh_token grant for dcr_bridge oauth_delegate, surfaced by an adversarial pass over the exchange path

Route the user-subject re-validation's outage check through the chain-aware classifier, so a transient DB outage (which get_user_object wraps in a bare ValueError) reports as unavailable (a retryable 503) rather than collapsing to no_active_key and an invalid_grant, matching how admission now handles the same wrapper

When the upstream reports its own refresh token as already elapsed (refresh_expires_in non-positive), do not seal it into a full-TTL refresh envelope; return no refresh so the exchange degrades to an access-only response, mirroring how the access grant refuses an already-elapsed access token instead of capping it

When the upstream rejects the sealed refresh token with 400 invalid_grant (revoked or expired at the IdP), return 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

Gate key-subject renewal on the owner's SCIM state, mirroring admission's _reject_if_admitted_owner_scim_deactivated, so an offboarded user cannot keep refreshing a still-active key; the check fails open on a missing owner or a DB blip so a key that outlives its owner record does not get wrongly revoked

Each fix has a mutation-checked regression test
This commit is contained in:
Tin Chi Lo 2026-07-13 13:27:27 -07:00
parent 52df186e3f
commit 1f1628d85c
2 changed files with 238 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

78
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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()

187
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -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"