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 <amasen02@users.noreply.github.com>
(cherry picked from commit daced81f20)

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)),
             ),
         )
 

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)