fix(proxy): reset a deleted customer's admission spend counter so a re-created one is not refused (#51224)

* fix(proxy): reset a deleted customer's admission spend counter so a re-created one is not refused

/customer/delete only dropped the committed twin, so spend:end_user:<id> kept the old spend and a
re-created customer with the same id stayed over budget. Zero the counter and its DB floor marker on
every worker after the delete commits, the same way a key spend reset does.

* fix(proxy): drop a deleted customer's spend counter before zeroing it

A Redis that rejects writes but still serves reads, such as one at its memory limit, kept the old
counter when the zeroing SET failed, so the re-created customer stayed refused. Delete the counter
first, which Redis still allows in that state, then write zero and the floor.

* test(proxy): type the customer delete spend counter test parameters

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>

litellm/proxy/management_endpoints/customer_endpoints.py

@@ -815,7 +815,12 @@ async def delete_end_user(
     See below for all params 
     ```
     """
-    from litellm.proxy.proxy_server import prisma_client, spend_counter_store
+    from litellm.proxy.proxy_server import (
+        invalidate_spend_counter,
+        prisma_client,
+        set_spend_counter_with_floor_and_broadcast,
+        spend_counter_store,
+    )
 
     try:
         if prisma_client is None:
@@ -845,9 +850,11 @@ async def delete_end_user(
             verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
 
             await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids))
-            await drop_committed_spend(
-                spend_counter_store, tuple(end_user_counter(user_id) for user_id in data.user_ids)
-            )
+            counter_keys: Final = tuple(end_user_counter(user_id) for user_id in data.user_ids)
+            for counter_key in counter_keys:
+                await invalidate_spend_counter(counter_key=counter_key)
+                await set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
+            await drop_committed_spend(spend_counter_store, counter_keys)
 
             return DeleteCustomersResponse(
                 deleted_customers=response,

tests/unit/proxy/management_endpoints/test_customer_endpoints.py

@@ -10,11 +10,15 @@ from fastapi.routing import APIRoute
 from fastapi.testclient import TestClient
 from pydantic import BaseModel
 
+import litellm
+from litellm.caching.caching import DualCache
 from litellm.proxy._types import (
+    LiteLLM_BudgetTable,
     LiteLLM_EndUserTable,
     LitellmUserRoles,
     ProxyException,
 )
+from litellm.proxy.auth.auth_checks import check_end_user_budget
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, router
 from litellm.types.proxy.management_endpoints.common_daily_activity import (
@@ -26,6 +30,7 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import (
     DeleteCustomersResponse,
     UnblockUsersResponse,
 )
+from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
 
 app = FastAPI()
 
@@ -250,20 +255,47 @@ def test_delete_customer_not_found(mock_prisma_client, mock_user_api_key_auth):
     assert response_json["error"]["code"] == "404"
 
 
-def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_starts_from_its_own_row(
-    monkeypatch, mock_prisma_client, mock_user_api_key_auth
-):
-    from types import SimpleNamespace
-
-    from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
-
-    spend: Final = FakeSpendCounters(
-        redis=None,
-        memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0},
-    )
+@pytest.fixture(autouse=True)
+def spend_counters(monkeypatch: pytest.MonkeyPatch) -> FakeSpendCounters:
+    """A fresh spend counter store and floor cache per test, so no counter a delete writes outlives its test."""
+    spend: Final = FakeSpendCounters(redis={})
     monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
+    monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", DualCache())
+    return spend
+
+
+def _over_budget_customers(spend: FakeSpendCounters) -> None:
+    from litellm.proxy.proxy_server import spend_counter_cache
+
+    assert spend.redis is not None
+    for user_id, spent in (("cust-gone", 7.0), ("cust-stays", 3.0)):
+        spend.redis[f"spend:end_user:{user_id}"] = spent
+        spend.redis[f"spend:committed:end_user:{user_id}"] = spent
+        spend_counter_cache.in_memory_cache.set_cache(key=f"spend_db_floor:spend:end_user:{user_id}", value=spent)
+
+
+async def _check_budget_of_new_customer(user_id: str) -> None:
+    await check_end_user_budget(
+        end_user_obj=LiteLLM_EndUserTable(
+            user_id=user_id, blocked=False, spend=0.0, litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0)
+        ),
+        route="/chat/completions",
+    )
+
+
+@pytest.mark.parametrize("fail_closed", [False, True])
+@pytest.mark.asyncio
+async def test_a_deleted_customer_recreated_under_the_same_id_is_admitted_while_others_stay_over_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    mock_prisma_client: MagicMock,
+    mock_user_api_key_auth: None,
+    spend_counters: FakeSpendCounters,
+    fail_closed: bool,
+) -> None:
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"fail_closed_budget_enforcement": fail_closed})
+    _over_budget_customers(spend_counters)
     mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
-        return_value=[SimpleNamespace(user_id="cust-gone")]
+        return_value=[LiteLLM_EndUserTable(user_id="cust-gone", blocked=False)]
     )
     mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1)
 
@@ -271,8 +303,79 @@ def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_start
         "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"}
     )
 
-    assert response.status_code == 200
-    assert spend.memory == {"spend:committed:end_user:cust-stays": 3.0}
+    assert response.status_code == 200, response.text
+    await _check_budget_of_new_customer("cust-gone")
+    with pytest.raises(litellm.BudgetExceededError) as exc_info:
+        await _check_budget_of_new_customer("cust-stays")
+    assert exc_info.value.current_cost == 3.0
+    assert spend_counters.redis is not None
+    assert "spend:committed:end_user:cust-gone" not in spend_counters.redis
+    assert spend_counters.redis["spend:committed:end_user:cust-stays"] == 3.0
+
+
+@pytest.mark.parametrize(
+    ("existing", "delete_error"),
+    [
+        pytest.param([], None, id="customer-missing"),
+        pytest.param(
+            [LiteLLM_EndUserTable(user_id="cust-gone", blocked=False)], RuntimeError("db down"), id="db-fails"
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_a_failed_customer_delete_keeps_its_spend_enforced(
+    mock_prisma_client: MagicMock,
+    mock_user_api_key_auth: None,
+    spend_counters: FakeSpendCounters,
+    existing: list[LiteLLM_EndUserTable],
+    delete_error: Exception | None,
+) -> None:
+    _over_budget_customers(spend_counters)
+    mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=existing)
+    mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(side_effect=delete_error, return_value=1)
+
+    response: Final = client.post(
+        "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"}
+    )
+
+    assert response.status_code >= 400
+    with pytest.raises(litellm.BudgetExceededError) as exc_info:
+        await _check_budget_of_new_customer("cust-gone")
+    assert exc_info.value.current_cost == 7.0
+
+
+@pytest.mark.parametrize(
+    ("failing", "redis_down"),
+    [
+        pytest.param(frozenset({"set"}), False, id="redis-rejects-writes"),
+        pytest.param(frozenset(), True, id="redis-down"),
+    ],
+)
+@pytest.mark.parametrize("fail_closed", [False, True])
+@pytest.mark.asyncio
+async def test_a_recreated_customer_is_admitted_when_redis_rejects_the_counter_reset(
+    monkeypatch: pytest.MonkeyPatch,
+    mock_prisma_client: MagicMock,
+    mock_user_api_key_auth: None,
+    failing: frozenset[str],
+    redis_down: bool,
+    fail_closed: bool,
+) -> None:
+    spend: Final = FakeSpendCounters(redis={}, failing=failing, fail=redis_down)
+    monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"fail_closed_budget_enforcement": fail_closed})
+    _over_budget_customers(spend)
+    mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
+        return_value=[LiteLLM_EndUserTable(user_id="cust-gone", blocked=False)]
+    )
+    mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1)
+
+    response: Final = client.post(
+        "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"}
+    )
+
+    assert response.status_code == 200, response.text
+    await _check_budget_of_new_customer("cust-gone")
 
 
 def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth):