Merge 53f41e40e9 into 784754ca48

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -6807,7 +6807,7 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio
 
 
 @with_service_target(SPEND_COUNTERS_TARGET)
-async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None:
+async def set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None:
     """
     Set a Redis-backed spend counter to `value`, mirror it into the short-lived
     spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both
@@ -6943,7 +6943,7 @@ async def _reset_key_budget_windows(
         duration = window.get("budget_duration")
         if isinstance(duration, str) and duration:
             counter_key = f"spend:key:{hashed_api_key}:window:{duration}"
-            await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
+            await set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
 
 
 @router.post(
@@ -7017,7 +7017,7 @@ async def reset_key_spend_fn(
         # amount for every enforcement check immediately instead of the stale
         # pre-reset value.
         _counter_key: Final = f"spend:key:{hashed_api_key}"
-        await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
+        await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
         await _reset_key_budget_windows(
             prisma_client=prisma_client,
             hashed_api_key=hashed_api_key,

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -1,4 +1,5 @@
 #### SPEND MANAGEMENT #####
+import asyncio
 import collections
 import json
 import os
@@ -295,6 +296,8 @@ class _TeamTable(Protocol):
 class _VerificationTokenTable(Protocol):
     """The subset of the Prisma verification token table API this module uses."""
 
+    async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SupportsModelDump]: ...
+
     async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ...
 
 
@@ -310,6 +313,28 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
     return VerificationTokenRepository(prisma_client).table
 
 
+_SPEND_COUNTER_RESET_BATCH_SIZE: Final = 100
+
+
+async def _reset_spend_counters(prisma_client: PrismaClient) -> None:
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        set_spend_counter_with_floor_and_broadcast,
+    )
+
+    for rows, identifier_field, prefix in (
+        (await _verification_token_table(prisma_client).find_many(where={}), "token", "spend:key:"),
+        (await _team_table(prisma_client).find_many(where={}), "team_id", "spend:team:"),
+    ):
+        for start in range(0, len(rows), _SPEND_COUNTER_RESET_BATCH_SIZE):
+            await asyncio.gather(
+                *(
+                    set_spend_counter_with_floor_and_broadcast(counter_key=f"{prefix}{identifier}", value=0.0)
+                    for row in rows[start : start + _SPEND_COUNTER_RESET_BATCH_SIZE]
+                    if isinstance(identifier := row.model_dump().get(identifier_field), str) and identifier
+                )
+            )
+
+
 def _spend_logs_daily_summary_sql(
     *,
     start_date_iso: str,
@@ -3629,6 +3654,8 @@ async def global_spend_reset():
     await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={})
     await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={})
 
+    await _reset_spend_counters(prisma_client)
+
     return {
         "message": "Spend for all API Keys and Teams reset successfully",
         "status": "success",

tests/integration/spend/test_global_spend_reset.py

@@ -132,10 +132,6 @@ def test_global_spend_reset_zeroes_token_and_team_spend_but_keeps_logs(gateway:
 
 
 def test_reset_exhausted_key_is_served_again(gateway: Gateway) -> None:
-    pytest.skip(
-        "BUG: /global/spend/reset zeroes DB spend but leaves Redis spend:key counters, "
-        "so an exhausted key stays refused with budget_exceeded"
-    )
     with (
         gateway.scenario() as scenario,
         httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,

tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -8347,3 +8347,66 @@ def test_spend_calculate_rejection_body_is_openai_shaped(client):
     assert error["type"] == "invalid_request_error"
     assert error["param"] is None
     assert error["code"] == "400"
+
+
+@pytest.mark.parametrize("key_count", [0, 2, 201])
+def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch, key_count: int):
+    """Regression test for https://github.com/BerriAI/litellm/issues/46363
+
+    Budget enforcement reads the Redis spend counters before the DB rows, so
+    a global reset that only zeroes the DB leaves an exhausted key refused
+    with the stale pre-reset counter. The reset must zero every live
+    spend:key: and spend:team: counter, not just the rows.
+    """
+    import litellm.proxy.management_endpoints.key_management_endpoints as kme
+    from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset
+
+    class Row:
+        def __init__(self, data):
+            self._data = data
+
+        def model_dump(self):
+            return self._data
+
+    class MockTable:
+        def __init__(self, rows):
+            self.rows = rows
+
+        async def find_many(self, *args, **kwargs):
+            return self.rows
+
+        async def update_many(self, *args, **kwargs):
+            return len(self.rows)
+
+    class MockDB:
+        def __init__(self):
+            self.litellm_verificationtoken = MockTable(
+                [Row({"token": f"key-{index}"}) for index in range(1, key_count + 1)]
+            )
+            self.litellm_teamtable = MockTable([Row({"team_id": "team-1"})])
+
+    class MockPrismaClient:
+        def __init__(self):
+            self.db = MockDB()
+
+    reset_calls: list[tuple[str, float]] = []
+
+    async def mock_set_counter(counter_key: str, value: float) -> None:
+        reset_calls.append((counter_key, value))
+
+    monkeypatch.setattr(
+        kme,
+        "set_spend_counter_with_floor_and_broadcast",
+        mock_set_counter,
+    )
+
+    monkeypatch.setattr(ps, "prisma_client", MockPrismaClient())
+    response: Final = asyncio.run(global_spend_reset())
+    assert response == {
+        "message": "Spend for all API Keys and Teams reset successfully",
+        "status": "success",
+    }
+    assert sorted(reset_calls) == sorted(
+        [(f"spend:key:key-{index}", 0.0) for index in range(1, key_count + 1)]
+        + [("spend:team:team-1", 0.0)]
+    ), reset_calls