diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 39e1c14a6e6..a850c50248c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -506,6 +506,14 @@ async def _fetch_global_spend_with_event_coordination( Reads the proxy budget aggregate user row, which accrues proxy-wide spend per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``. + + With Redis configured, the shared scalar is read straight from Redis: the + global spend is incremented by every worker, so a per-worker in-memory copy + can go stale across resets and falsely reject requests at auth time. Reading + the Redis value directly keeps the budget check on the shared total. On a + Redis miss the coordinator loads the authoritative DB spend once and + repopulates Redis. Without Redis (single process), the normal DualCache + coordinator path is used. """ async def _load_global_spend() -> float | None: @@ -514,9 +522,12 @@ async def _fetch_global_spend_with_event_coordination( ) return float(proxy_budget_row.spend) if proxy_budget_row is not None else None + spend_cache: Final = ( + user_api_key_cache.redis_cache if user_api_key_cache.redis_cache is not None else user_api_key_cache + ) return await _global_spend_coordinator.get_or_load( cache_key=cache_key, - cache=user_api_key_cache, # pyright: ignore[reportArgumentType] + cache=spend_cache, # pyright: ignore[reportArgumentType] load_fn=_load_global_spend, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d00614cc8ed..e977727d68d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -262,19 +262,19 @@ class ResetBudgetJob: @staticmethod async def _invalidate_global_proxy_spend_cache() -> None: - """Zero the cached global-proxy spend accumulator after the proxy - budget aggregate row is reset so atomic increments start from a known floor. + """Drop the cached global-proxy spend accumulator after the proxy + budget aggregate row is reset so the next auth-time read reloads the + zeroed DB row. Reads and writes are atomic now (increment-if-present on + Redis), so deleting cannot resurrect a stale pre-reset value, and it + cannot clobber a concurrent increment the way a plain SET to zero could. """ try: from litellm.proxy.proxy_server import user_api_key_cache - await user_api_key_cache.async_set_cache( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, - value=0.0, - ) - except Exception as e: # noqa: BLE001 # cache zero failure must not break budget reset + await user_api_key_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) + except Exception as e: # noqa: BLE001 # cache invalidation failure must not break budget reset verbose_proxy_logger.warning( - "Failed to zero global proxy spend cache after budget reset: %s", + "Failed to invalidate global proxy spend cache after budget reset: %s", e, ) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2cd207cd6bd..5e4a83cfe46 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -459,6 +459,24 @@ async def test_async_increment_if_exists_increments_present_key( assert "if cur == false then return nil end" in lua +@pytest.mark.asyncio +async def test_async_increment_if_exists_decodes_bytes_result( + monkeypatch, redis_no_ping +): + """Redis returns bytes; the method decodes them to a float.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=None) + mock_redis_instance = AsyncMock() + mock_redis_instance.eval = AsyncMock(return_value=b"7.0") + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + result = await redis_cache.async_increment_if_exists(key="k", value=2.0) + + assert result == 7.0 + + @pytest.mark.asyncio async def test_async_increment_if_exists_leaves_missing_key_absent( monkeypatch, redis_no_ping 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 eea556b9a0d..1555bc34a06 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 @@ -6121,6 +6121,36 @@ async def test_global_proxy_spend_none_when_proxy_budget_row_missing(): assert result is None +@pytest.mark.asyncio +async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured(): + """With Redis, the global spend read goes straight to the shared Redis + scalar instead of a per-worker in-memory copy, which can go stale across + resets and falsely reject requests at auth time.""" + from litellm.proxy.auth.user_api_key_auth import ( + _fetch_global_spend_with_event_coordination, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=8.5) + cache.redis_cache = redis_cache + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + side_effect=AssertionError("a Redis hit must not hit the DB") + ) + + result = await _fetch_global_spend_with_event_coordination( + cache_key="default_user_id:spend", + user_api_key_cache=cache, + prisma_client=prisma_client, + ) + + assert result == 8.5 + redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend") + + @pytest.mark.asyncio async def test_temp_budget_increase_applied_for_cached_key(): """ diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index a0e9c157025..77a39225b26 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1092,8 +1092,8 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache( reset_budget_job, mock_prisma_client, monkeypatch ): - """Resetting the proxy-wide budget aggregate row zeroes the cached - global-spend accumulator so the next period starts from a known floor.""" + """Resetting the proxy-wide budget aggregate row invalidates the cached + global-spend accumulator so the next read reloads the zeroed DB row.""" counter_cache = _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) @@ -1113,15 +1113,15 @@ def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - counter_cache.user_api_key_cache.async_set_cache.assert_any_call( - key="default_user_id:spend", value=0.0 + counter_cache.user_api_key_cache.async_delete_cache.assert_any_call( + key="default_user_id:spend" ) def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( reset_budget_job, mock_prisma_client, monkeypatch ): - """The global-spend accumulator must only be zeroed when the proxy + """The global-spend accumulator must only be invalidated when the proxy budget aggregate row itself resets, not on every user reset.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1152,6 +1152,33 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( ) +def test_reset_budget_for_proxy_budget_row_swallows_global_spend_cache_failure( + reset_budget_job, mock_prisma_client, monkeypatch +): + """A failure invalidating the global spend cache must not break the reset.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + counter_cache.user_api_key_cache.async_delete_cache = AsyncMock( + side_effect=RuntimeError("redis down") + ) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 150.0, + "budget_duration": "30d", + "budget_reset_at": now, + "id": "row-1", + "user_id": "litellm-proxy-budget", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Team budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) 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 8bf29be240e..277e6e083fd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1400,3 +1400,37 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke ttl=ps.get_management_object_ttl(cache), ) assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) == 3.5 + + +@pytest.mark.asyncio +async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatch): + """A Redis failure on the global spend increment must not fail the request.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + + redis_cache = MagicMock() + redis_cache.async_increment_if_exists = AsyncMock( + side_effect=RuntimeError("redis down") + ) + cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300) + monkeypatch.setattr(ps, "user_api_key_cache", cache) + + await cache.async_set_cache( + key="user-redis-down", + value={"user_id": "user-redis-down", "spend": 0.0}, + ) + await ps.update_cache( + token=None, + user_id="user-redis-down", + end_user_id=None, + team_id=None, + response_cost=1.0, + parent_otel_span=None, + ) + + redis_cache.async_increment_if_exists.assert_awaited_once_with( + key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, + value=1.0, + ttl=ps.get_management_object_ttl(cache), + )