From 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 01/13] fix(auth): inherit organization_alias from the org for JWT and team-linked keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 42 ++++++- .../proxy/auth/test_user_api_key_auth.py | 104 +++++++++++++++++- 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20ab9904f46..110c524ecdf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_end_user_object, get_jwt_key_mapping_object, get_object_permission, + get_org_object, get_project_object, get_team_object, get_user_object, @@ -2398,6 +2399,37 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() +async def _inherit_org_identity( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTableCachedObj | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> None: + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + if ( + user_api_key_auth_obj.org_id is None + or user_api_key_auth_obj.organization_alias is not None + or prisma_client is None + ): + return + try: + org_object: Final = await get_org_object( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.debug("org alias lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + return + if org_object is not None: + user_api_key_auth_obj.organization_alias = org_object.organization_alias + + @tracer.wrap() async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, @@ -2622,8 +2654,14 @@ async def _run_centralized_common_checks( ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result - if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: - user_api_key_auth_obj.org_id = team_object.organization_id + await _inherit_org_identity( + user_api_key_auth_obj=user_api_key_auth_obj, + team_object=cast(LiteLLM_TeamTableCachedObj | None, team_object), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and 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 6cce6d0316b..03efbfa7185 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 @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -31,7 +32,7 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object +from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -5293,6 +5294,107 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_id,team_org_id,existing_alias,lookup_mode,expected_org_id,expected_alias", + [ + (None, "t1", "org-from-team", None, "success", "org-from-team", "acme-org"), + ("org-jwt", None, None, None, "success", "org-jwt", "acme-org"), + ("org-pinned", None, None, "preset", "success", "org-pinned", "preset"), + ("org-missing", None, None, None, "missing", "org-missing", None), + ], +) +async def test_centralized_common_checks_inherits_org_alias( + key_org_id, + team_id, + team_org_id, + existing_alias, + lookup_mode, + expected_org_id, + expected_alias, +): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id=team_id, + org_id=key_org_id, + organization_alias=existing_alias, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = ( + LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None + ) + organization = LiteLLM_OrganizationTable( + organization_id=expected_org_id, + organization_alias="acme-org", + budget_id="budget-id", + models=[], + created_by="test", + updated_by="test", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["prisma_client"] = MagicMock() + 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", + new_callable=AsyncMock, + return_value=fetched_team, + ) as mock_get_team_object, + patch( + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + return_value=organization, + ) as mock_get_org_object, + 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") + + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert token.organization_alias == expected_alias + assert identity_seen_by_common_checks == [(expected_org_id, expected_alias)] + if team_id is None: + mock_get_team_object.assert_not_awaited() + else: + mock_get_team_object.assert_awaited_once() + if existing_alias is not None: + mock_get_org_object.assert_not_awaited() + else: + mock_get_org_object.assert_awaited_once() + assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted From cfd8c186161068bef2ed5faae002dbfb3b2ab63e Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 21:24:36 +0000 Subject: [PATCH 02/13] fix(auth): inherit org budget, tpm and rpm limits for JWT and team-linked keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 31 ++++++++++++----- .../proxy/auth/test_user_api_key_auth.py | 34 +++++++++++++++---- 2 files changed, 50 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 110c524ecdf..df5257908ce 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2409,11 +2409,17 @@ async def _inherit_org_identity( ) -> None: if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id - if ( - user_api_key_auth_obj.org_id is None - or user_api_key_auth_obj.organization_alias is not None - or prisma_client is None - ): + already_populated: Final = any( + value is not None + for value in ( + user_api_key_auth_obj.organization_alias, + user_api_key_auth_obj.organization_max_budget, + user_api_key_auth_obj.organization_tpm_limit, + user_api_key_auth_obj.organization_rpm_limit, + user_api_key_auth_obj.organization_metadata, + ) + ) + if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None: return try: org_object: Final = await get_org_object( @@ -2422,12 +2428,21 @@ async def _inherit_org_identity( user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, ) except Exception: - verbose_proxy_logger.debug("org alias lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + verbose_proxy_logger.debug("org lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) return - if org_object is not None: - user_api_key_auth_obj.organization_alias = org_object.organization_alias + if org_object is None: + return + user_api_key_auth_obj.organization_alias = org_object.organization_alias + user_api_key_auth_obj.organization_metadata = org_object.metadata + budget: Final = org_object.litellm_budget_table + if budget is None: + return + user_api_key_auth_obj.organization_max_budget = budget.max_budget + user_api_key_auth_obj.organization_tpm_limit = budget.tpm_limit + user_api_key_auth_obj.organization_rpm_limit = budget.rpm_limit @tracer.wrap() 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 03efbfa7185..fd9b6f09678 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 @@ -5296,22 +5296,26 @@ 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,lookup_mode,expected_org_id,expected_alias", + "key_org_id,team_id,team_org_id,existing_alias,existing_rpm,lookup_mode,expected_org_id,expected_alias,expected_limits", [ - (None, "t1", "org-from-team", None, "success", "org-from-team", "acme-org"), - ("org-jwt", None, None, None, "success", "org-jwt", "acme-org"), - ("org-pinned", None, None, "preset", "success", "org-pinned", "preset"), - ("org-missing", None, None, None, "missing", "org-missing", None), + (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)), ], ) -async def test_centralized_common_checks_inherits_org_alias( +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, ): import litellm.proxy.proxy_server as _proxy_server_mod from fastapi import Request @@ -5325,6 +5329,7 @@ async def test_centralized_common_checks_inherits_org_alias( team_id=team_id, org_id=key_org_id, organization_alias=existing_alias, + organization_rpm_limit=existing_rpm, ) request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") @@ -5336,9 +5341,15 @@ async def test_centralized_common_checks_inherits_org_alias( organization_id=expected_org_id, organization_alias="acme-org", budget_id="budget-id", + metadata={"model_rpm_limit": {"gpt-4o": 2}}, models=[], created_by="test", updated_by="test", + litellm_budget_table=( + None + if lookup_mode == "no_budget" + else LiteLLM_BudgetTable(budget_id="budget-id", max_budget=12.5, tpm_limit=700, rpm_limit=7) + ), ) attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) @@ -5380,16 +5391,25 @@ async def test_centralized_common_checks_inherits_org_alias( 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)] if team_id is None: mock_get_team_object.assert_not_awaited() else: mock_get_team_object.assert_awaited_once() - if existing_alias is not None: + if existing_alias is not None or existing_rpm is not None: mock_get_org_object.assert_not_awaited() + assert token.organization_metadata is None else: 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": + assert token.organization_metadata == {"model_rpm_limit": {"gpt-4o": 2}} finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) From cc875a6eb38a2737a172da9a97ecf9f960c0750f Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 21:51:06 +0000 Subject: [PATCH 03/13] fix(auth): exempt org lookup fallback from strict lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ac1cc3cfb35..17498467485 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1309,7 +1309,7 @@ def _ensure_litellm_received_at_on_request_state(request: Request) -> datetime: received_at: Final = datetime.now(timezone.utc) try: request.state.litellm_received_at = received_at - except Exception: + except Exception: # noqa: BLE001 # organization lookup must not fail authentication pass return received_at From 87c00bf47b7ef0c0dcc8aba29bb7b4e2c68ad94e Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 21:51:18 +0000 Subject: [PATCH 04/13] fix(auth): place strict lint exemption on org lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 17498467485..711fa50f93d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1309,7 +1309,7 @@ def _ensure_litellm_received_at_on_request_state(request: Request) -> datetime: received_at: Final = datetime.now(timezone.utc) try: request.state.litellm_received_at = received_at - except Exception: # noqa: BLE001 # organization lookup must not fail authentication + except Exception: pass return received_at @@ -2634,7 +2634,7 @@ async def _inherit_org_identity( proxy_logging_obj=proxy_logging_obj, include_budget_table=True, ) - except Exception: + 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) return if org_object is None: From 4a13ebbc5b3c62cf50f184d2e26f017caa86c35d Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:01:50 +0000 Subject: [PATCH 05/13] 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(): From 79b6cd29172ae2827259051b0b45b8af12c51a60 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:03:16 +0000 Subject: [PATCH 06/13] fix(auth): treat a missing org row as no org limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 7 ++++--- tests/test_litellm/proxy/auth/test_user_api_key_auth.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 89d543f9ccb..10924cf62bd 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -41,6 +41,7 @@ from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, + OrganizationNotFoundError, TeamNotFoundError, _cache_key_object, _can_object_call_model, @@ -2634,13 +2635,13 @@ async def _inherit_org_identity( proxy_logging_obj=proxy_logging_obj, include_budget_table=True, ) - except Exception: # noqa: BLE001 # organization lookup must not fail authentication + except OrganizationNotFoundError: + return + except Exception: # noqa: BLE001 # DB outage handling is decided by allow_requests_on_db_unavailable 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 user_api_key_auth_obj.organization_alias = org_object.organization_alias user_api_key_auth_obj.organization_metadata = org_object.metadata budget: Final = org_object.litellm_budget_table 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 377e2d9342d..0bf49523869 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,6 +35,7 @@ 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, @@ -5814,7 +5815,7 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, ("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-missing", None, None, None, None, "missing", False, 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)), @@ -5892,7 +5893,7 @@ async def test_centralized_common_checks_inherits_org_identity( ) as mock_checks, ): if lookup_mode == "missing": - mock_get_org_object.return_value = None + mock_get_org_object.side_effect = OrganizationNotFoundError("x") elif lookup_mode == "db_failure": mock_get_org_object.side_effect = RuntimeError("db unavailable") From ee9294af53f2e681fce2f95a80ae266766f19ce8 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:08:45 +0000 Subject: [PATCH 07/13] test(auth): annotate centralized auth mocks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_user_api_key_auth.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 0bf49523869..bc2370579d4 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 @@ -5877,17 +5877,17 @@ async def test_centralized_common_checks_inherits_org_identity( for k, v in attrs.items(): setattr(_proxy_server_mod, k, v) with ( - patch( + patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists "litellm.proxy.auth.user_api_key_auth.get_team_object", new_callable=AsyncMock, return_value=fetched_team, ) as mock_get_team_object, - patch( + patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists "litellm.proxy.auth.user_api_key_auth.get_org_object", new_callable=AsyncMock, return_value=organization, ) as mock_get_org_object, - patch( + patch( # test-quality-ok: capture downstream token state without invoking unrelated common checks "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock, ) as mock_checks, From c4ad6194a0aa2009b12d09fb4f5cd8f671c5a423 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:28:11 +0000 Subject: [PATCH 08/13] fix(auth): only fail closed on DB outages during org lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 9 +++++++-- tests/test_litellm/proxy/auth/test_user_api_key_auth.py | 9 ++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 10924cf62bd..41888cb9a64 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2637,11 +2637,16 @@ async def _inherit_org_identity( ) except OrganizationNotFoundError: return - except Exception: # noqa: BLE001 # DB outage handling is decided by allow_requests_on_db_unavailable - if not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): + except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits + if ( + PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) + and 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 user_api_key_auth_obj.organization_alias = org_object.organization_alias user_api_key_auth_obj.organization_metadata = org_object.metadata budget: Final = org_object.litellm_budget_table 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 bc2370579d4..ce8310c8aa3 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 @@ -5818,6 +5818,7 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, ("org-missing", None, None, None, None, "missing", False, 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-bad-row", None, None, None, None, "bad_row", False, False, "org-bad-row", None, (None, None, None)), ("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)), ], ) @@ -5895,10 +5896,12 @@ async def test_centralized_common_checks_inherits_org_identity( if lookup_mode == "missing": mock_get_org_object.side_effect = OrganizationNotFoundError("x") elif lookup_mode == "db_failure": - mock_get_org_object.side_effect = RuntimeError("db unavailable") + mock_get_org_object.side_effect = ConnectionRefusedError("db unavailable") + elif lookup_mode == "bad_row": + mock_get_org_object.side_effect = ValueError("row failed validation") if expect_lookup_error: - with pytest.raises(RuntimeError, match="db unavailable"): + with pytest.raises(ConnectionRefusedError, match="db unavailable"): await _run_centralized_common_checks( user_api_key_auth_obj=token, request=request, @@ -5943,7 +5946,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 not in {"missing", "db_failure"}: + if lookup_mode not in {"missing", "db_failure", "bad_row"}: assert token.organization_metadata == {"model_rpm_limit": {"gpt-4o": 2}} finally: for k, v in originals.items(): From 1b69a5b0a45012408794d1b8aa95043c4f2ae945 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 22:34:02 +0000 Subject: [PATCH 09/13] test(proxy): model missing organizations in MCP auth fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/test_user_api_key_auth_mcp.py | 6 ++++++ .../_experimental/mcp_server/test_discoverable_endpoints.py | 6 +++++- 2 files changed, 11 insertions(+), 1 deletion(-) 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 90ce821d62e..a0fb76349b2 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 @@ -5739,9 +5739,15 @@ class TestMCPDcrBridgeDelegateAdmission: prisma and are swallowed (``_safe_fetch`` / the SCIM gate's fail-open), so their checks skip. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key.""" + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) + get_org_object = AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")) patchers = [ patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), + patch( # test-quality-ok: central auth now resolves org limits; this fixture models a missing org row + "litellm.proxy.auth.user_api_key_auth.get_org_object", get_org_object + ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ] 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 aa45b2f6793..200df078e00 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 @@ -11870,9 +11870,13 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.bridge_token_flow import authorize_oauth_credential_request from litellm.proxy._types import UserAPIKeyAuth, hash_token - from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, jwt_key_mapping_cache_key handler, signing_key = jwt_oauth_identity + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth.get_org_object", + AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")), + ) key: Final = "sk-oauth-permission-test" hashed: Final = hash_token(key) credential: Final = UserAPIKeyAuth( From 8983eefea57ea39d143f22103b7fa259d5518269 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 21:30:32 +0000 Subject: [PATCH 10/13] fix(auth): drop redundant cast on team_object in centralized checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ced86318d7b..b4d8648c8a9 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2901,7 +2901,7 @@ async def _run_centralized_common_checks( await _inherit_org_identity( user_api_key_auth_obj=user_api_key_auth_obj, - team_object=cast(LiteLLM_TeamTableCachedObj | None, team_object), + team_object=team_object, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, From 82e3f3980d44f3822fa30ae089d0a034335a402c Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 23:59:33 +0000 Subject: [PATCH 11/13] refactor(auth): resolve org identity through an auth_checks helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 28 ++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 29 +++++-------------- .../auth/test_user_api_key_auth_mcp.py | 2 +- .../mcp_server/test_discoverable_endpoints.py | 2 +- .../proxy/auth/test_user_api_key_auth.py | 2 +- 5 files changed, 39 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index cdada970956..ba37eed037f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4012,6 +4012,34 @@ async def get_org_object( return _org_obj +async def get_org_object_for_request( + org_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> LiteLLM_OrganizationTable | None: + try: + return await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, + ) + except OrganizationNotFoundError: + return None + except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits + if ( + PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) + and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ): + raise + verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) + return None + + async def _get_resources_from_access_groups( access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b4d8648c8a9..7ef1c2775ab 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -41,7 +41,6 @@ from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, - OrganizationNotFoundError, TeamNotFoundError, _cache_key_object, _can_object_call_model, @@ -59,7 +58,7 @@ from litellm.proxy.auth.auth_checks import ( get_jwt_key_mapping_object, get_key_end_user_budget_id, get_object_permission, - get_org_object, + get_org_object_for_request, get_project_object, get_team_membership, get_team_object, @@ -2630,25 +2629,13 @@ async def _inherit_org_identity( ) if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None: return - try: - org_object: Final = await get_org_object( - org_id=user_api_key_auth_obj.org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - include_budget_table=True, - ) - except OrganizationNotFoundError: - return - except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits - if ( - PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) - and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - ): - raise - verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) - return + org_object: Final = await get_org_object_for_request( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) if org_object is None: return user_api_key_auth_obj.organization_alias = org_object.organization_alias 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 a0fb76349b2..4380df194ed 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 @@ -5746,7 +5746,7 @@ class TestMCPDcrBridgeDelegateAdmission: patchers = [ patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), patch( # test-quality-ok: central auth now resolves org limits; this fixture models a missing org row - "litellm.proxy.auth.user_api_key_auth.get_org_object", get_org_object + "litellm.proxy.auth.auth_checks.get_org_object", get_org_object ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), 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 200df078e00..3556722ff6e 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 @@ -11874,7 +11874,7 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions( handler, signing_key = jwt_oauth_identity monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.get_org_object", + "litellm.proxy.auth.auth_checks.get_org_object", AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")), ) key: Final = "sk-oauth-permission-test" 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 761bd454eaf..8593be751fa 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 @@ -6065,7 +6065,7 @@ async def test_centralized_common_checks_inherits_org_identity( return_value=fetched_team, ) as mock_get_team_object, patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists - "litellm.proxy.auth.user_api_key_auth.get_org_object", + "litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock, return_value=organization, ) as mock_get_org_object, From 9075cafb98e3b22c0bedce288217039ce3058698 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 12:55:14 -0700 Subject: [PATCH 12/13] fix(auth): serve the last-known org through a database outage A JWT whose team sits in an org resolves the org on every request, and the org row is cached for only DEFAULT_IN_MEMORY_TTL seconds while the team and user rows ride the 60s management-object TTL. A few seconds into a database outage the org lookup failed closed and that traffic got 503s while the same request through a virtual key kept succeeding on its cached team. get_org_object now also keeps a last-known copy of the org row under the management-object TTL, and get_org_object_for_request serves that copy when the database is unreachable, so JWT traffic degrades the same way the team lookup does. A missing copy keeps the previous behaviour: fail closed unless allow_requests_on_db_unavailable is set. --- litellm/proxy/auth/auth_checks.py | 30 +++++++++--- .../proxy/auth/test_auth_checks.py | 48 +++++++++++++++++++ 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c92d8a1a543..161a91f648d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4008,10 +4008,21 @@ async def get_org_object( model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) + if include_budget_table: + await user_api_key_cache.async_set_cache( + key=_last_known_org_cache_key(org_id), + value=_org_obj, + model_type=LiteLLM_OrganizationTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) return _org_obj +def _last_known_org_cache_key(org_id: str) -> str: + return f"org_id:{org_id}:with_budget:last_known" + + async def get_org_object_for_request( org_id: str, prisma_client: PrismaClient, @@ -4031,13 +4042,18 @@ async def get_org_object_for_request( except OrganizationNotFoundError: return None except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits - if ( - PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e) - and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - ): - raise - verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) - return None + if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): + verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) + return None + last_known_org: Final = await user_api_key_cache.async_get_cache( + key=_last_known_org_cache_key(org_id), + model_type=LiteLLM_OrganizationTable, + ) + if last_known_org is not None: + return last_known_org + if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): + return None + raise async def _get_resources_from_access_groups( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1ae986db23b..08764ad5b18 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6087,6 +6087,54 @@ async def test_organization_budget_check_carries_org_state_on_the_token(): assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) +@pytest.mark.asyncio +async def test_get_org_object_for_request_serves_last_known_org_through_db_outage(): + """A JWT whose team sits in an org resolves the org on every request, and the org row + is cached for only DEFAULT_IN_MEMORY_TTL seconds while the team and user rows ride the + 60s management-object TTL. Without a last-known copy, a DB outage a few seconds old + turned that traffic into 503s while the same request through a virtual key kept + succeeding on its cached team.""" + from litellm.proxy.auth.auth_checks import get_org_object_for_request + + org_row = MagicMock() + org_row.model_dump = lambda: { + "organization_id": "org-1", + "organization_alias": "platform-org", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + "litellm_budget_table": {"budget_id": "b1", "max_budget": 50.0, "tpm_limit": 700, "rpm_limit": 7}, + } + prisma_client = MagicMock() + prisma_client.db.litellm_organizationtable.find_unique = AsyncMock( + side_effect=[org_row, ConnectionRefusedError("db unavailable")] + ) + user_api_key_cache = UserApiKeyCache() + + async def _lookup(): + return await get_org_object_for_request( + org_id="org-1", + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists + warm = await _lookup() + assert warm is not None and warm.organization_alias == "platform-org" + await user_api_key_cache.async_delete_cache("org_id:org-1:with_budget") + + during_outage = await _lookup() + + assert prisma_client.db.litellm_organizationtable.find_unique.await_count == 2 + assert during_outage is not None + assert during_outage.organization_alias == "platform-org" + assert during_outage.litellm_budget_table is not None + assert during_outage.litellm_budget_table.rpm_limit == 7 + assert during_outage.litellm_budget_table.max_budget == 50.0 + + @pytest.mark.parametrize( "max_budget, spend, expect_blocked", [ From cae6634192dbad73ef089dbf8a1f28a3df7a56bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 13:36:37 -0700 Subject: [PATCH 13/13] fix(auth): keep the last-known org copy when the auth prefetch warmed the org row The last-known org copy was written only on get_org_object's DB-read path. The virtual-key auth prefetch fills the same 5s org entry directly, so with keys and JWTs of one org on the same worker the JWT lookup always hit the cache, never wrote the copy, and a DB outage turned that JWT traffic into 503s again. get_org_object_for_request now writes the copy itself whenever this worker holds none, under the management-object TTL, and get_org_object is back to its shape on main. --- litellm/proxy/auth/auth_checks.py | 30 ++++++--- .../proxy/auth/test_auth_checks.py | 65 +++++++++++++++++-- 2 files changed, 81 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 161a91f648d..65795e09976 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4008,13 +4008,6 @@ async def get_org_object( model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) - if include_budget_table: - await user_api_key_cache.async_set_cache( - key=_last_known_org_cache_key(org_id), - value=_org_obj, - model_type=LiteLLM_OrganizationTable, - ttl=get_management_object_ttl(user_api_key_cache), - ) return _org_obj @@ -4023,6 +4016,23 @@ def _last_known_org_cache_key(org_id: str) -> str: return f"org_id:{org_id}:with_budget:last_known" +async def _keep_last_known_org( + org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache +) -> None: + cache_key: Final = _last_known_org_cache_key(org_id) + held_locally: Final = await user_api_key_cache.async_get_cache( + key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable + ) + if held_locally is not None: + return + await user_api_key_cache.async_set_cache( + key=cache_key, + value=org, + model_type=LiteLLM_OrganizationTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + + async def get_org_object_for_request( org_id: str, prisma_client: PrismaClient, @@ -4031,7 +4041,7 @@ async def get_org_object_for_request( proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_OrganizationTable | None: try: - return await get_org_object( + org: Final = await get_org_object( org_id=org_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -4054,6 +4064,10 @@ async def get_org_object_for_request( if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): return None raise + if org is None: + return None + await _keep_last_known_org(org, org_id, user_api_key_cache) + return org async def _get_resources_from_access_groups( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 08764ad5b18..b64e4d6ae6c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6087,17 +6087,19 @@ async def test_organization_budget_check_carries_org_state_on_the_token(): assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) +@pytest.mark.parametrize("warmed_by_auth_prefetch", [False, True]) @pytest.mark.asyncio -async def test_get_org_object_for_request_serves_last_known_org_through_db_outage(): +async def test_get_org_object_for_request_serves_last_known_org_through_db_outage(warmed_by_auth_prefetch): """A JWT whose team sits in an org resolves the org on every request, and the org row is cached for only DEFAULT_IN_MEMORY_TTL seconds while the team and user rows ride the 60s management-object TTL. Without a last-known copy, a DB outage a few seconds old turned that traffic into 503s while the same request through a virtual key kept - succeeding on its cached team.""" + succeeding on its cached team. The copy must exist whoever filled the short-lived entry: + this lookup's own DB read, or the virtual-key auth prefetch warming it for the same org.""" + from litellm.proxy._types import LiteLLM_OrganizationTable from litellm.proxy.auth.auth_checks import get_org_object_for_request - org_row = MagicMock() - org_row.model_dump = lambda: { + org_columns = { "organization_id": "org-1", "organization_alias": "platform-org", "budget_id": "b1", @@ -6105,11 +6107,20 @@ async def test_get_org_object_for_request_serves_last_known_org_through_db_outag "updated_by": "admin", "litellm_budget_table": {"budget_id": "b1", "max_budget": 50.0, "tpm_limit": 700, "rpm_limit": 7}, } + org_row = MagicMock() + org_row.model_dump = lambda: org_columns + db_outage = ConnectionRefusedError("db unavailable") prisma_client = MagicMock() prisma_client.db.litellm_organizationtable.find_unique = AsyncMock( - side_effect=[org_row, ConnectionRefusedError("db unavailable")] + side_effect=[db_outage] if warmed_by_auth_prefetch else [org_row, db_outage] ) user_api_key_cache = UserApiKeyCache() + if warmed_by_auth_prefetch: + await user_api_key_cache.async_set_cache( + key="org_id:org-1:with_budget", + value=LiteLLM_OrganizationTable.model_validate(org_columns), + model_type=LiteLLM_OrganizationTable, + ) async def _lookup(): return await get_org_object_for_request( @@ -6127,7 +6138,7 @@ async def test_get_org_object_for_request_serves_last_known_org_through_db_outag during_outage = await _lookup() - assert prisma_client.db.litellm_organizationtable.find_unique.await_count == 2 + assert prisma_client.db.litellm_organizationtable.find_unique.await_count == (1 if warmed_by_auth_prefetch else 2) assert during_outage is not None assert during_outage.organization_alias == "platform-org" assert during_outage.litellm_budget_table is not None @@ -6135,6 +6146,48 @@ async def test_get_org_object_for_request_serves_last_known_org_through_db_outag assert during_outage.litellm_budget_table.max_budget == 50.0 +@pytest.mark.asyncio +async def test_get_org_object_for_request_writes_the_last_known_org_only_when_absent(): + """The last-known copy is written when this worker holds none, never per request: + with Redis attached, a write on every cached org hit would cost one SET per JWT request.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import get_org_object_for_request + + class _WriteRecordingCache(UserApiKeyCache): + def __init__(self): + super().__init__() + self.written_keys = [] + + async def async_set_cache(self, key, value, local_only=False, **kwargs): + self.written_keys.append(key) + return await super().async_set_cache(key=key, value=value, local_only=local_only, **kwargs) + + user_api_key_cache = _WriteRecordingCache() + await user_api_key_cache.async_set_cache( + key="org_id:org-1:with_budget", + value=LiteLLM_OrganizationTable( + organization_id="org-1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + ), + model_type=LiteLLM_OrganizationTable, + ) + + for _ in range(3): + org = await get_org_object_for_request( + org_id="org-1", + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert org is not None and org.organization_alias == "platform-org" + + assert user_api_key_cache.written_keys.count("org_id:org-1:with_budget:last_known") == 1 + + @pytest.mark.parametrize( "max_budget, spend, expect_blocked", [