diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e7252b03188..cf25e65598c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -581,14 +581,17 @@ async def _fetch_global_spend_with_event_coordination( if db_spend is not None and db_spend > numeric_value: # Use a monotonic write so a concurrent increment that already # pushed Redis above the DB value is not overwritten by a repair - # carrying a slightly-stale DB total. + # carrying a slightly-stale DB total. Use the resulting Redis value + # so the enforced total reflects any concurrent increment. + repaired_value: Final[float | None] = ( + await user_api_key_cache.redis_cache.async_set_max( + key=cache_key, value=db_spend + ) + ) + numeric_value = repaired_value if repaired_value is not None else db_spend user_api_key_cache.in_memory_cache.set_cache( - key=cache_key, value=db_spend + key=cache_key, value=numeric_value ) - await user_api_key_cache.redis_cache.async_set_max( - key=cache_key, value=db_spend - ) - numeric_value = db_spend return numeric_value diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5cbdbaaad67..7780a68c83f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3127,20 +3127,15 @@ async def update_cache( return try: # Atomically accumulate the shared global-proxy scalar in Redis. - # The returned value is authoritative across workers, so mirror it - # into this worker's in-memory tier to avoid accumulating a stale - # pre-reset local value. + # Authorization reads this Redis scalar directly, so there is no + # per-worker in-memory copy that can go stale across resets. if user_api_key_cache.redis_cache is not None: - current = await user_api_key_cache.redis_cache.async_increment( + await user_api_key_cache.redis_cache.async_increment( key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=response_cost, ttl=get_management_object_ttl(user_api_key_cache), refresh_ttl=True, ) - await user_api_key_cache.in_memory_cache.async_set_cache( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, - value=current, - ) else: # No Redis: single-process, so an in-memory increment is coherent. await user_api_key_cache.async_increment_cache( 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 4cc56785844..121eca303fe 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 @@ -6168,7 +6168,7 @@ async def test_global_proxy_spend_floors_stale_low_redis_to_db(): cache = UserApiKeyCache() redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock() + redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) cache.redis_cache = redis_cache proxy_budget_row = MagicMock() @@ -6204,7 +6204,7 @@ async def test_global_proxy_spend_floor_does_not_overwrite_concurrent_increment( cache = UserApiKeyCache() redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock() + redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) cache.redis_cache = redis_cache proxy_budget_row = MagicMock() @@ -6238,7 +6238,7 @@ async def test_global_proxy_spend_marker_cached_db_value_used_without_db_query() cache = UserApiKeyCache() redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock() + redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) cache.redis_cache = redis_cache # Seed the marker cache with a DB value greater than Redis. diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 70a1d6f80cf..065f3c83692 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1335,9 +1335,12 @@ async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset @pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monkeypatch): - """In-memory mirrors the authoritative Redis result instead of accruing a - stale pre-reset value, so budget checks do not reject off stale memory.""" +async def test_update_cache_global_proxy_spend_does_not_mirror_to_in_memory( + monkeypatch, +): + """Authorization reads the shared Redis scalar directly, so the increment + path must not maintain a per-worker in-memory copy that could go stale + across resets and confuse budget checks.""" from unittest.mock import AsyncMock, MagicMock from litellm.caching.caching import DualCache @@ -1348,12 +1351,12 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke monkeypatch.setattr(ps, "user_api_key_cache", cache) await cache.async_set_cache( - key="user-mirror", - value={"user_id": "user-mirror", "spend": 0.0}, + key="user-no-mirror", + value={"user_id": "user-no-mirror", "spend": 0.0}, ) await ps.update_cache( token=None, - user_id="user-mirror", + user_id="user-no-mirror", end_user_id=None, team_id=None, response_cost=1.5, @@ -1366,7 +1369,7 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke ttl=ps.get_management_object_ttl(cache), refresh_ttl=True, ) - assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) == 3.5 + assert cache.in_memory_cache.get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) is None @pytest.mark.asyncio