From 4a13ebbc5b3c62cf50f184d2e26f017caa86c35d Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:01:50 +0000 Subject: [PATCH] fix(auth): fail closed on org lookup errors when DB is required Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 4 +- .../proxy/auth/test_user_api_key_auth.py | 85 ++++++++++++------- 2 files changed, 56 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 711fa50f93d..89d543f9ccb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2635,7 +2635,9 @@ async def _inherit_org_identity( include_budget_table=True, ) except Exception: # noqa: BLE001 # organization lookup must not fail authentication - verbose_proxy_logger.debug("org lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + if not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): + raise + verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) return if org_object is None: return diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index b5200de115e..377e2d9342d 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -35,7 +35,6 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.auth_checks import ( - OrganizationNotFoundError, TeamNotFoundError, UserNotFoundError, get_key_object, @@ -5809,27 +5808,31 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, @pytest.mark.asyncio @pytest.mark.parametrize( - "key_org_id,team_id,team_org_id,existing_alias,existing_rpm,lookup_mode,expected_org_id,expected_alias,expected_limits", + "key_org_id,team_id,team_org_id,existing_alias,existing_rpm,lookup_mode,allow_db_unavailable,expect_lookup_error,expected_org_id,expected_alias,expected_limits", [ - (None, "t1", "org-from-team", None, None, "success", "org-from-team", "acme-org", (12.5, 700, 7)), - ("org-jwt", None, None, None, None, "success", "org-jwt", "acme-org", (12.5, 700, 7)), - ("org-pinned", None, None, "preset", None, "success", "org-pinned", "preset", (None, None, None)), - ("org-view", None, None, None, 3, "success", "org-view", None, (None, None, 3)), - ("org-missing", None, None, None, None, "missing", "org-missing", None, (None, None, None)), - ("org-nobudget", None, None, None, None, "no_budget", "org-nobudget", "acme-org", (None, None, None)), + (None, "t1", "org-from-team", None, None, "success", False, False, "org-from-team", "acme-org", (12.5, 700, 7)), + ("org-jwt", None, None, None, None, "success", False, False, "org-jwt", "acme-org", (12.5, 700, 7)), + ("org-pinned", None, None, "preset", None, "success", False, False, "org-pinned", "preset", (None, None, None)), + ("org-view", None, None, None, 3, "success", False, False, "org-view", None, (None, None, 3)), + ("org-missing", None, None, None, None, "missing", True, False, "org-missing", None, (None, None, None)), + ("org-db-failure-allowed", None, None, None, None, "db_failure", True, False, "org-db-failure-allowed", None, (None, None, None)), + ("org-db-failure-denied", None, None, None, None, "db_failure", False, True, "org-db-failure-denied", None, (None, None, None)), + ("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)), ], ) async def test_centralized_common_checks_inherits_org_identity( - key_org_id, - team_id, - team_org_id, - existing_alias, - existing_rpm, - lookup_mode, - expected_org_id, - expected_alias, - expected_limits, -): + key_org_id: str | None, + team_id: str | None, + team_org_id: str | None, + existing_alias: str | None, + existing_rpm: int | None, + lookup_mode: str, + allow_db_unavailable: bool, + expect_lookup_error: bool, + expected_org_id: str | None, + expected_alias: str | None, + expected_limits: tuple[float | None, int | None, int | None], +) -> None: import litellm.proxy.proxy_server as _proxy_server_mod from fastapi import Request from starlette.datastructures import URL @@ -5867,11 +5870,11 @@ async def test_centralized_common_checks_inherits_org_identity( attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) attrs["prisma_client"] = MagicMock() + attrs["general_settings"] = {"allow_requests_on_db_unavailable": allow_db_unavailable} originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} try: for k, v in attrs.items(): setattr(_proxy_server_mod, k, v) - identity_seen_by_common_checks = [] with ( patch( "litellm.proxy.auth.user_api_key_auth.get_team_object", @@ -5886,30 +5889,48 @@ async def test_centralized_common_checks_inherits_org_identity( patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock, - side_effect=lambda **kw: identity_seen_by_common_checks.append( - (kw["valid_token"].org_id, kw["valid_token"].organization_alias) - ), ) as mock_checks, ): if lookup_mode == "missing": - mock_get_org_object.side_effect = OrganizationNotFoundError("x") + mock_get_org_object.return_value = None + elif lookup_mode == "db_failure": + mock_get_org_object.side_effect = RuntimeError("db unavailable") - await _run_centralized_common_checks( - user_api_key_auth_obj=token, - request=request, - request_data={"model": "gpt-4o"}, - route="/chat/completions", - ) + if expect_lookup_error: + with pytest.raises(RuntimeError, match="db unavailable"): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + else: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + assert token.org_id == expected_org_id + if expect_lookup_error: + mock_checks.assert_not_awaited() + assert token.organization_alias is None + assert token.organization_max_budget is None + assert token.organization_tpm_limit is None + assert token.organization_rpm_limit is None + return mock_checks.assert_awaited_once() - assert token.org_id == expected_org_id assert token.organization_alias == expected_alias assert ( token.organization_max_budget, token.organization_tpm_limit, token.organization_rpm_limit, ) == expected_limits - assert identity_seen_by_common_checks == [(expected_org_id, expected_alias)] + checked_token = mock_checks.await_args.kwargs["valid_token"] + assert checked_token.org_id == expected_org_id + assert checked_token.organization_alias == expected_alias if team_id is None: mock_get_team_object.assert_not_awaited() else: @@ -5921,7 +5942,7 @@ async def test_centralized_common_checks_inherits_org_identity( mock_get_org_object.assert_awaited_once() assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id assert mock_get_org_object.await_args.kwargs["include_budget_table"] is True - if lookup_mode != "missing": + if lookup_mode not in {"missing", "db_failure"}: assert token.organization_metadata == {"model_rpm_limit": {"gpt-4o": 2}} finally: for k, v in originals.items():