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", [