From db8dfe93a5a74fad5c4ba3f9fabaa5de39ad09ca Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:10:43 +0000 Subject: [PATCH] fix(auth): wait for the in-flight membership load before evicting its cache key on invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 10 +- .../proxy/auth/test_auth_checks.py | 95 ++++++------------- 2 files changed, 36 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 22644bcd207..13605d7dc7b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2204,8 +2204,6 @@ async def _fetch_team_membership_from_db( ) membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - return membership if membership is None: await user_api_key_cache.async_set_cache( key=_key, @@ -2218,8 +2216,6 @@ async def _fetch_team_membership_from_db( value=membership, model_type=LiteLLM_TeamMembership, ) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - await user_api_key_cache.async_delete_cache(key=_key) return membership @@ -2765,7 +2761,11 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) - _team_membership_inflight.pop(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None) + inflight: Final[object] = _team_membership_inflight.pop( + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None + ) + if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task(): + await asyncio.wait((inflight,)) if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 96e9958a55d..5f87f2def93 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6494,7 +6494,7 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): @pytest.mark.asyncio -async def test_get_team_membership_invalidation_mid_flight_discards_stale_load(): +async def test_get_team_membership_invalidation_waits_for_in_flight_load_then_evicts_it(): from litellm.proxy._types import LiteLLM_TeamMembership from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -6502,20 +6502,21 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() started = asyncio.Event() release_stale = asyncio.Event() - release_fresh = asyncio.Event() - loads = iter((("budget-old", release_stale), ("budget-new", release_fresh))) + rows = iter(("budget-old", "budget-new")) async def _find_unique(*args, **kwargs): - budget_id, release = next(loads) + budget_id = next(rows) row = MagicMock() row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} - started.set() - await release.wait() + if budget_id == "budget-old": + started.set() + await release_stale.wait() return row mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) cache = UserApiKeyCache() + _key = team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv") async def _load(): return await get_team_membership( @@ -6523,24 +6524,28 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() ) stale = asyncio.create_task(_load()) - await started.wait() - await invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) - started.clear() - fresh = asyncio.create_task(_load()) await asyncio.wait_for(started.wait(), timeout=2) - release_fresh.set() - fresh_result = await fresh - release_stale.set() - stale_result = await stale + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_stale.set() + await asyncio.wait_for(invalidation, timeout=2) + stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=_key) is None + + fresh_result = await _load() assert fresh_result is not None and fresh_result.budget_id == "budget-new" assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 - cached = CacheCodec.deserialize( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")), - model_type=LiteLLM_TeamMembership, - ) + cached = CacheCodec.deserialize(await cache.async_get_cache(key=_key), model_type=LiteLLM_TeamMembership) assert cached is not None and cached.budget_id == "budget-new" + again = await _load() + assert again is not None and again.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 @pytest.mark.asyncio @@ -6567,59 +6572,21 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) ) await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_write.set() + await asyncio.wait_for(invalidation, timeout=2) stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None -@pytest.mark.asyncio -async def test_get_team_membership_stale_write_finishing_after_fresh_load_never_serves_old_row(): - from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state - - write_started = asyncio.Event() - release_stale_write = asyncio.Event() - stale_writes = iter((release_stale_write,)) - - class _SlowFirstWriteCache(UserApiKeyCache): - async def async_set_cache(self, key, value, local_only=False, **kwargs): - release = next(stale_writes, None) - if release is not None: - write_started.set() - await release.wait() - return await super().async_set_cache(key, value, local_only=local_only, **kwargs) - - rows = iter(("budget-old", "budget-new", "budget-new")) - - async def _find_unique(*args, **kwargs): - row = MagicMock() - row.dict = lambda: {"user_id": "u-sw", "team_id": "t-sw", "spend": 1.0, "budget_id": next(rows)} - return row - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) - cache = _SlowFirstWriteCache() - - async def _load(): - return await get_team_membership( - user_id="u-sw", team_id="t-sw", prisma_client=mock_prisma_client, user_api_key_cache=cache - ) - - stale = asyncio.create_task(_load()) - await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-sw", team_id="t-sw", user_api_key_cache=cache) - fresh_result = await _load() - release_stale_write.set() - stale_result = await stale - after_result = await _load() - - assert fresh_result is not None and fresh_result.budget_id == "budget-new" - assert stale_result is not None and stale_result.budget_id == "budget-old" - assert after_result is not None and after_result.budget_id == "budget-new" - - @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request