From 8e3742a5f3dee525548152f7379d789074cc915d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:39:02 -0700 Subject: [PATCH] fix(proxy): answer 503 temporarily_unavailable when the token exchange cannot verify the subject token A subject token JWT auth could not check, because the IdP's JWKS was unreachable with no cached copy or the auth database was down, came back as 400 invalid_request with the same fixed message a bad token gets, so clients re-logged in instead of retrying the way they already do for a mint-time 503. Those checks now answer 503 temporarily_unavailable and log the reason, while real rejections stay 400 invalid_request. --- .../mcp_server/gateway_dcr_flow.py | 14 +++++- .../mcp_server/idp_token_exchange.py | 43 ++++++++++++++++--- .../mcp_server/test_gateway_dcr_flow.py | 9 ++-- .../mcp_server/test_idp_token_exchange.py | 35 ++++++++++++++- 4 files changed, 88 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index ba24e861d6e..e66504af47a 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -221,7 +221,7 @@ class SubjectIdentity(BaseModel): class SubjectTokenRefusal(BaseModel): model_config = ConfigDict(frozen=True) - error: Literal["unsupported_grant_type", "invalid_request"] + error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"] description: str = Field(min_length=1) @@ -1136,6 +1136,16 @@ def _reload_failure_response(failure: ReloadUserFailure) -> Response: assert_never(failure) +def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response: + match refusal.error: + case "temporarily_unavailable": + return _oauth_error(503, refusal.error, refusal.description) + case "unsupported_grant_type" | "invalid_request": + return _oauth_error(400, refusal.error, refusal.description) + case _: + assert_never(refusal.error) + + def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response: match failure: case "not_a_member": @@ -1314,7 +1324,7 @@ class _GrantIssuer: return target_refusal identity: Final = await exchange_subject_token(subject_token, self._request) if isinstance(identity, SubjectTokenRefusal): - return _oauth_error(400, identity.error, identity.description) + return _subject_token_refusal_response(identity) principal: Final = SessionPrincipal( user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id ) diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 7a453e85cce..7ab9dc034bf 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -15,9 +15,13 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" +SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = ( + "the gateway could not verify subject_token because its identity provider or database is unavailable; retry" +) @dataclass(frozen=True, slots=True) @@ -141,9 +145,12 @@ async def identity_from_subject_token( ) -> SubjectIdentity | SubjectTokenRefusal: """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which - RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token. The - reason stays in the proxy log: this endpoint is public and JWT auth's own wording can - name the JWKS URL it fetched or quote the IdP's response.""" + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token, and a + check the gateway could not complete (the IdP's JWKS unreachable with no cached copy, + the auth database down) as ``temporarily_unavailable``, so the client retries instead + of treating a valid token as bad. The reason stays in the proxy log: this endpoint is + public and JWT auth's own wording can name the JWKS URL it fetched or quote the IdP's + response.""" unmet: Final = prerequisites.refusal() if unmet is not None: return unmet @@ -152,17 +159,39 @@ async def identity_from_subject_token( try: result: Final = await authorize(subject_token, request_headers) except HTTPException as denied: - return _rejected_by_jwt_auth(denied.detail) + return _refusal_for(denied, denied.detail) except ProxyException as denied: - return _rejected_by_jwt_auth(denied.message) + return _refusal_for(denied, denied.message) except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures - return _rejected_by_jwt_auth(denied) + return _refusal_for(denied, denied) user_id: Final = result["user_id"] if user_id is None: return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) -def _rejected_by_jwt_auth(reason: object) -> SubjectTokenRefusal: +def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal: + if _gateway_could_not_verify(denied): + verbose_proxy_logger.error("token exchange could not verify a subject_token, retryable: %s", reason) + return SubjectTokenRefusal(error="temporarily_unavailable", description=SUBJECT_TOKEN_CHECK_UNAVAILABLE) verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + + +def _gateway_could_not_verify(denied: Exception) -> bool: + """A 5xx from JWT auth (the IdP's JWKS unreachable with no cached copy) or a database + outage anywhere in the chain (``get_user_object`` wraps prisma failures in a bare + ``ValueError``) is the gateway failing, not the token.""" + if _is_server_error(denied): + return True + return PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied) is not None + + +def _is_server_error(denied: Exception) -> bool: + match denied: + case HTTPException(status_code=status_code): + return status_code >= 500 + case ProxyException(code=code): + return code.isdigit() and int(code) >= 500 + case _: + return False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index c42f8763e74..2943ff4b74a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -2324,13 +2324,16 @@ async def test_token_exchange_refuses_a_malformed_request_before_touching_the_id @pytest.mark.asyncio -@pytest.mark.parametrize("error", ["unsupported_grant_type", "invalid_request"]) -async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error): +@pytest.mark.parametrize( + "error, status", + [("unsupported_grant_type", 400), ("invalid_request", 400), ("temporarily_unavailable", 503)], +) +async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error, status): client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] minter = _Minter() exchanger = _Exchanger(SubjectTokenRefusal(error=error, description="subject_token was rejected: bad signature")) response = await _exchange_native(client_id, minter, exchanger) - assert response.status_code == 400 + assert response.status_code == status body = json.loads(response.body) assert (body["error"], body["error_description"]) == (error, "subject_token was rejected: bad signature") assert minter.calls == [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index e12c8823f99..d86ef137576 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -2,17 +2,19 @@ import logging import pytest from fastapi import HTTPException +from prisma.errors import DataError from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, + SUBJECT_TOKEN_CHECK_UNAVAILABLE, TokenExchangePrerequisites, identity_from_subject_token, token_exchange_available, ) from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException -from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.handle_jwt import JWKSUnreachableError, JWTHandler, jwks_unavailable_exception IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} @@ -23,6 +25,7 @@ EVERY_GATE_HOLDS = { "maps_jwts_to_virtual_keys": False, } JWKS_URL = "https://idp.example.com/.well-known/jwks.json" +JWKS_DOWN = jwks_unavailable_exception(JWKSUnreachableError(f"ConnectError fetching {JWKS_URL} after 3 attempts")) def _authorized(user_id="u1", team_id="team-b"): @@ -162,6 +165,7 @@ def test_availability_is_read_from_the_running_proxy( (Exception("Validation fails: signature verification failed"), "signature verification failed"), (Exception("Invalid JWT Submitted"), "Invalid JWT"), (Exception(f"Failed to fetch keys from {JWKS_URL}: 502 Bad Gateway from the IdP"), JWKS_URL), + (ValueError("User doesn't exist in db. 'user_id'=u1. Got error - not found"), "not found"), ], ) async def test_a_jwt_the_proxy_rejects_is_refused_with_the_reason_kept_in_the_log(raised, reason, caplog): @@ -179,3 +183,32 @@ async def test_a_jwt_that_resolves_no_user_cannot_be_exchanged(): assert refusal == SubjectTokenRefusal( error="invalid_request", description="subject_token names no user the gateway knows" ) + + +def _user_lookup_wrapping_a_database_outage(): + p1001 = DataError( + data={"user_facing_error": {"message": "Can't reach database server at `127.0.0.1`:`5432`", "meta": {}}} + ) + try: + raise p1001 + except DataError as outage: + try: + raise ValueError(f"User doesn't exist in db. 'user_id'=u1. Got error - {outage}") + except ValueError as wrapped: + return wrapped + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised, reason", + [ + (JWKS_DOWN, JWKS_URL), + (HTTPException(status_code=503, detail="the auth database is not reachable"), "not reachable"), + (_user_lookup_wrapping_a_database_outage(), "Can't reach database server"), + ], +) +async def test_an_idp_or_gateway_outage_is_reported_as_retryable_not_as_a_bad_token(raised, reason, caplog): + caplog.set_level(logging.ERROR, logger="LiteLLM Proxy") + refusal = await _identity(_Authorizer(raises=raised)) + assert refusal == SubjectTokenRefusal(error="temporarily_unavailable", description=SUBJECT_TOKEN_CHECK_UNAVAILABLE) + assert reason in caplog.text