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