diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c8895c01f7b..d8d63440234 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/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, diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcb9eb4a280..a2d41eda46f 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/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", diff --git a/tests/integration/spend/test_global_spend_reset.py b/tests/integration/spend/test_global_spend_reset.py index cac66471dea..72b09b11788 100644 --- a/tests/integration/spend/test_global_spend_reset.py +++ b/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, diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index e119acd5ffc..d05accf5208 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/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