From 9cd7cde1f9a4661ddca95d3c8d0e027710fbde8f Mon Sep 17 00:00:00 2001 From: amasen02 Date: Fri, 4 Sep 2026 15:37:53 +0530 Subject: [PATCH] fix(proxy): invalidate end-user spend counter and cache on budget reset (#39726) Adapted for stable/1.99.x: this line predates budget rollover (#38514), so each end user's spend counter is zeroed through the existing counter_keys list instead of counter_resets, and its cache entry is evicted via end_user_cache_key. Signed-off-by: amasen02 (cherry picked from commit daced81f20a6b98c01beecf66fa997310d90bfb4) --- .../proxy/common_utils/reset_budget_job.py | 15 +++++++++-- .../common_utils/test_reset_budget_job.py | 26 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 8fcb184b26a..d8ff61d713d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -32,7 +32,7 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) -from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -122,6 +122,14 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (tag_cache_key(row.tag_name),) +def _enduser_counter_key(row: _EndUserRow) -> str: + return f"spend:end_user:{row.user_id}" + + +def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]: + return (end_user_cache_key(row.user_id),) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -522,6 +530,7 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="tags", ) + endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids) return _BudgetCascade( budgets=tuple(budgets_to_reset), budget_ids=budget_ids, @@ -533,18 +542,20 @@ class ResetBudgetJob: for b in budgets_to_reset if b.budget_id is not None and b.budget_duration is not None ), - endusers=await self._collect_endusers_to_reset(budget_ids), + endusers=endusers, counter_keys=( *(_team_membership_counter_key(row) for row in team_memberships), *(_key_counter_key(row) for row in keys), *(_org_counter_key(row) for row in orgs), *(_tag_counter_key(row) for row in tags), + *(_enduser_counter_key(row) for row in endusers), ), cache_keys=( *(key for row in team_memberships for key in _team_membership_cache_keys(row)), *(key for row in keys for key in _key_cache_keys(row)), *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), + *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) 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 25c177a308d..a6d0c3b5fa9 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 @@ -1346,6 +1346,32 @@ def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budge assert deleted == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} +def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_job, mock_prisma_client, monkeypatch): + """When an end user's budget resets, its Redis spend counter is zeroed and its management cache is evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + budget = _budget_row(budget_id="budget-1") + mock_prisma_client.data["budget"] = [budget] + test_enduser = type( + "LiteLLM_EndUserTable", + (), + { + "spend": 20.0, + "litellm_budget_table": budget, + "budget_id": "budget-1", + "user_id": "customer-42", + }, + ) + mock_prisma_client.data["enduser"] = [test_enduser] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60) + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert "end_user_id:customer-42" in deleted + + + def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch): """Eviction runs after the commit, so a broken cache cannot undo the write.""" counter_cache = _make_counter_invalidation_job(monkeypatch)