fix(proxy): invalidate end-user spend counter and cache on budget reset (#39726)

Signed-off-by: amasen02 <amasen02@users.noreply.github.com>

litellm/proxy/common_utils/reset_budget_job.py

@@ -38,6 +38,7 @@ from litellm.proxy.common_utils.timezone_utils import (
     get_budget_reset_settings,
 )
 from litellm.proxy.common_utils.user_api_key_cache import (
+    end_user_cache_key,
     model_access_group_cache_key,
     model_access_group_spend_counter_key,
     tag_cache_key,
@@ -177,6 +178,21 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
     return (model_access_group_cache_key(row.access_group_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 _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float:
+    if not caps:
+        return 0.0
+    effective_budget_id = row.budget_id or litellm.max_end_user_budget_id
+    return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None)
+
+
 def _budget_link_where(
     budget_ids: Sequence[str],
     extra: Mapping[str, object] = MappingProxyType({}),
@@ -650,6 +666,7 @@ class ResetBudgetJob:
             if _rollover_enabled()
             else {}  # mutable-ok: empty sentinel immediately frozen by MappingProxyType
         )
+        endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids)
         return _BudgetCascade(
             budgets=tuple(budgets_to_reset),
             budget_ids=budget_ids,
@@ -661,7 +678,7 @@ 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_resets=(
                 *(
                     (_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps))
@@ -674,6 +691,10 @@ class ResetBudgetJob:
                     (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
                     for row in model_access_groups
                 ),
+                *(
+                    (_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps))
+                    for row in endusers
+                ),
             ),
             rollover_caps=rollover_caps,
             cache_keys=(
@@ -682,6 +703,7 @@ class ResetBudgetJob:
                 *(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 model_access_groups for key in _model_access_group_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

@@ -1495,6 +1495,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)