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