From 25ebb9458c29843438a9da9e2d2cc4af03fbef56 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:52:41 -0700 Subject: [PATCH] fix(proxy): surface a database outage from the user read as 503 no_db_connection (#42399) get_user_object wrapped every failed read, a refused connection included, in ValueError("User doesn't exist in db ..."), so JWT callers got a 401 naming a missing user while Postgres was down and virtual-key callers got 503 no_db_connection for the same outage. A connection or transport error now propagates as-is and the auth exception mapper answers 503 no_db_connection; a genuinely missing row and query-level errors still answer 401. The MCP auth and token-exchange docstrings and the exception-chain helper's docstring described the old wrap and are updated to the new contract. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 14 +-- .../mcp_server/bridge_token_flow.py | 6 +- .../mcp_server/idp_token_exchange.py | 5 +- litellm/proxy/auth/auth_checks.py | 12 ++- litellm/proxy/db/exception_handler.py | 7 +- .../test_user_api_key_auth.py | 19 ++-- .../auth/test_user_api_key_auth_mcp.py | 13 +-- .../proxy/auth/test_auth_checks.py | 90 ++++++++++++++----- 8 files changed, 107 insertions(+), 59 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 2e5cad29b0a..60dc91a69cc 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 @@ -966,10 +966,11 @@ class MCPRequestHandler: on top of these direct grants, each source bounded by ITS OWN org, so a user spanning organizations cannot leak one org's servers past another's ceiling. - Error handling: ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError``, so a - missing user and a real outage look identical (the cause survives only as ``__context__``). - ``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 while any - other failure fails closed as 401, not an opaque 500; the object-permission load shares that boundary.""" + Error handling: ``get_user_object`` lets a database outage propagate as-is and re-raises every other + DB failure as a bare ``ValueError`` (the cause surviving only as ``__context__``). + ``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 whichever + shape it arrives in, while any other failure fails closed as 401, not an opaque 500; the + object-permission load shares that boundary.""" 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 @@ -1116,9 +1117,8 @@ class MCPRequestHandler: (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. - 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.""" + Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself, so an outage a + caller re-raised inside a domain exception is still recognized.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(e) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 37a893973e3..6ae33cc1629 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -276,9 +276,9 @@ async def load_active_user_by_id( 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. ``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. + lets a real outage propagate as-is and re-raises any other DB failure as a bare ``ValueError`` (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. ``source="database"`` reads the row from the database, never the cache, so the credential mint refuses a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh row in the cache for the requests the credential makes next. Every other caller keeps the cache read, diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 80868296b50..a437df17e6a 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -196,9 +196,8 @@ def _check_unavailable_description(outage: GatewayOutage) -> str: def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None: - """A database fault anywhere in the chain (``get_user_object`` wraps prisma failures in a - bare ``ValueError``) or a 5xx from JWT auth (the IdP's JWKS unreachable with no cached - copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or + """A database fault anywhere in the chain or a 5xx from JWT auth (the IdP's JWKS + unreachable with no cached copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or version-skewed query engine) is named as such, the way the mint path words it, so the client is not told to wait on a deployment that needs repair.""" fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b3efde7f05f..f0588b3fa79 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2697,9 +2697,15 @@ async def get_user_object( raise except Exception as e: _log_budget_lookup_failure("user", e) - raise ValueError( - f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}" - ) + raise _user_read_failure(user_id=user_id, error=e) + + +def _user_read_failure(user_id: str, error: Exception) -> Exception: + if PrismaDBExceptionHandler.is_database_service_unavailable_error(error): + return error + return ValueError( + f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {error}" + ) async def _cache_management_object( diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f25c2787252..99167c1275c 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -398,11 +398,8 @@ class PrismaDBExceptionHandler: ``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 + domain exception of a different type defeats. A type check on the + wrapper 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. diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index f5e8d861d79..9cdac341b1f 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -1540,18 +1540,17 @@ async def test_user_budget_lookup_is_also_unenforced_when_the_database_is_down() """ KNOWN LIMITATION, pinned deliberately rather than discovered later. - `get_user_object` cannot tell "row absent" from "database unreachable": the - absent case raises inside its own try (auth_checks.py:2177) and the handler - at :2213 rewrites every exception into the same - `ValueError("User doesn't exist in db...")`. A connection error, a query - timeout and a malformed row all reach us as that one type and message. + `get_user_object` lets a connection-level outage propagate as-is and rewrites + every other read failure (a query-level Prisma error, a malformed row) into + the same `ValueError("User doesn't exist in db...")` as an absent row, and + `_read_user_model_max_budget` swallows every exception either way. - So tolerating the absent case, which the test above requires, unavoidably - tolerates an outage too, and a user who DOES have a per-model budget goes + So tolerating the absent case, which the test above requires, also + tolerates an outage, and a user who DOES have a per-model budget goes unenforced while the DB is unreachable. This is pre-existing behaviour of - `get_user_object` that the virtual-key path inherits identically; it is not - introduced here. Distinguishing them needs a dedicated exception type for - the absent case and a change to both auth paths. + the virtual-key path; it is not introduced here. Distinguishing them needs + `_read_user_model_max_budget` to let an outage through the way the JWT + path does. """ from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget 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 5c885a3168d..35e7055bbc0 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 @@ -5936,12 +5936,13 @@ class TestMCPDcrBridgeDelegateAdmission: @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.""" + """Reproduce get_user_object's exception contract (litellm/proxy/auth/auth_checks.py): a read + failure that is not a database outage is re-raised as a bare ``ValueError`` with the original + error only as ``__context__``, while an outage propagates raw (pinned by + test_get_user_object_surfaces_a_db_outage_as_503_not_as_a_missing_user and + test_get_user_object_still_reports_a_non_outage_read_failure_as_a_missing_user in + test_auth_checks). The wrapped shape is the harder one for the outage classifier, so injecting + it here keeps a chain-blind classifier from passing.""" try: raise original except BaseException: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0cab6535f3d..14b739e60ba 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -75,6 +75,8 @@ from litellm.constants import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from prisma.errors import DataError from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, @@ -892,36 +894,80 @@ async def test_get_user_object_upsert_sets_budget_reset_at(monkeypatch, has_budg assert "budget_reset_at" not in creation_args -@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 +def _user_read_raising(error: Exception) -> tuple[MagicMock, MagicMock]: + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + return prisma_client, cache - 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() + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outage", + [ + httpx.ConnectError("All connection attempts failed"), + httpx.ReadTimeout("timed out"), + DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1:41071`", + "error_code": "P1001", + } + } + ), + ], + ids=["connect_error", "read_timeout", "p1001_as_data_error"], +) +async def test_get_user_object_surfaces_a_db_outage_as_503_not_as_a_missing_user(outage): + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + prisma_client, cache = _user_read_raising(outage) with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): - with pytest.raises(ValueError, match="User doesn't exist in db\\.") as exc_info: + with pytest.raises(type(outage)) as raised: await get_user_object( - user_id="outage-contract-probe-user", - prisma_client=mock_prisma_client, - user_api_key_cache=mock_cache, + user_id="outage-probe-user", + prisma_client=prisma_client, + user_api_key_cache=cache, user_id_upsert=False, proxy_logging_obj=None, ) - assert isinstance(exc_info.value.__context__, ConnectionError) + assert raised.value is outage + assert PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(raised.value) is outage + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", + [ + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("row validation failed"), + ], + ids=["query_level_data_error", "runtime_error"], +) +async def test_get_user_object_still_reports_a_non_outage_read_failure_as_a_missing_user(failure): + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + prisma_client, cache = _user_read_raising(failure) + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(ValueError, match="User doesn't exist in db\\.") as raised: + await get_user_object( + user_id="data-error-probe-user", + prisma_client=prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + ) + + assert raised.value.__context__ is failure + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("401", ProxyErrorTypes.auth_error) @pytest.mark.asyncio