From 54fe3f5522b29b1a055fd62952c8fd21a8442f00 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.98.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 an inline "end_user_id:" key, matching how this line writes it in auth_checks. Signed-off-by: amasen02 (cherry picked from commit daced81f20a6b98c01beecf66fa997310d90bfb4) --- .../proxy/common_utils/reset_budget_job.py | 13 +++++++++- .../common_utils/test_reset_budget_job.py | 26 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bf760a92d88..cfbb8bd4a2b 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -115,6 +115,14 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (f"tag:{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 (f"end_user_id:{row.user_id}",) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -346,6 +354,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, @@ -357,18 +366,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 608dc8cb5c8..9fee24c76c1 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 @@ -1341,6 +1341,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)