From 31161749e19ab6614a33484ea7bcc00d4a305116 Mon Sep 17 00:00:00 2001 From: libaojiang <2461443384@qq.com> Date: Wed, 14 Oct 2026 17:18:37 +0800 Subject: [PATCH 1/5] fix(proxy): zero live spend counters on global spend reset Budget enforcement reads the Redis spend counters before the DB rows, so /global/spend/reset zeroing only the DB left an exhausted key refused with the stale pre-reset counter. The reset now also sets every spend:key: and spend:team: counter to zero through the same set-with- floor-and-broadcast helper /key/{key}/reset_spend uses. Fixes #46363 --- .../spend_management_endpoints.py | 33 +++++++++ .../test_spend_management_endpoints.py | 70 +++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5482661c3de..3d0a8cc5b88 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -279,6 +279,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: ... @@ -294,6 +296,35 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken return VerificationTokenRepository(prisma_client).table +async def _reset_spend_counters(prisma_client: PrismaClient) -> None: + """Zero every live spend counter a budget check reads. + + Budget enforcement reads the Redis spend counters first, so zeroing only + the DB rows leaves an exhausted key refused with the stale pre-reset + counter. Mirrors /key/{key}/reset_spend, setting each counter to 0 (with + floor marker and broadcast) rather than deleting it. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _set_spend_counter_with_floor_and_broadcast, + ) + + key_rows: Final = await _verification_token_table(prisma_client).find_many(where={}) + for row in key_rows: + token: Final[str | None] = row.model_dump().get("token") + if token: + await _set_spend_counter_with_floor_and_broadcast( + counter_key=f"spend:key:{token}", value=0.0 + ) + + team_rows: Final = await _team_table(prisma_client).find_many(where={}) + for row in team_rows: + team_id: Final[str | None] = row.model_dump().get("team_id") + if team_id: + await _set_spend_counter_with_floor_and_broadcast( + counter_key=f"spend:team:{team_id}", value=0.0 + ) + + def _spend_logs_daily_summary_sql( *, start_date_iso: str, @@ -3697,6 +3728,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/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 25180038e2a..0d0f1fe059c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8059,3 +8059,73 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + +def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): + """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": "key-1"}), Row({"token": "key-2"})] + ) + 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, + ) + + async def run_reset(): + import litellm.proxy.proxy_server as ps + + ps.prisma_client = MockPrismaClient() + try: + return await global_spend_reset() + finally: + ps.prisma_client = None + + response: Final = asyncio.run(run_reset()) + assert response == { + "message": "Spend for all API Keys and Teams reset successfully", + "status": "success", + } + assert sorted(reset_calls) == [ + ("spend:key:key-1", 0.0), + ("spend:key:key-2", 0.0), + ("spend:team:team-1", 0.0), + ], reset_calls From b58821dbe97df75f0b4f938238ec70b6fba05140 Mon Sep 17 00:00:00 2001 From: libaojiang <2461443384@qq.com> Date: Wed, 14 Oct 2026 17:37:12 +0800 Subject: [PATCH 2/5] style: ruff format --- .../proxy/spend_tracking/spend_management_endpoints.py | 8 ++------ .../spend_tracking/test_spend_management_endpoints.py | 5 ++--- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3d0a8cc5b88..6ad2284717f 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -312,17 +312,13 @@ async def _reset_spend_counters(prisma_client: PrismaClient) -> None: for row in key_rows: token: Final[str | None] = row.model_dump().get("token") if token: - await _set_spend_counter_with_floor_and_broadcast( - counter_key=f"spend:key:{token}", value=0.0 - ) + await _set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:key:{token}", value=0.0) team_rows: Final = await _team_table(prisma_client).find_many(where={}) for row in team_rows: team_id: Final[str | None] = row.model_dump().get("team_id") if team_id: - await _set_spend_counter_with_floor_and_broadcast( - counter_key=f"spend:team:{team_id}", value=0.0 - ) + await _set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:team:{team_id}", value=0.0) def _spend_logs_daily_summary_sql( 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 0d0f1fe059c..fbefe0c916a 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8060,6 +8060,7 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): """Regression test for https://github.com/BerriAI/litellm/issues/46363 @@ -8090,9 +8091,7 @@ def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): class MockDB: def __init__(self): - self.litellm_verificationtoken = MockTable( - [Row({"token": "key-1"}), Row({"token": "key-2"})] - ) + self.litellm_verificationtoken = MockTable([Row({"token": "key-1"}), Row({"token": "key-2"})]) self.litellm_teamtable = MockTable([Row({"team_id": "team-1"})]) class MockPrismaClient: From e384bff79bda532961a92a0d483056027730544a Mon Sep 17 00:00:00 2001 From: libaojiang <2461443384@qq.com> Date: Wed, 14 Oct 2026 18:56:52 +0800 Subject: [PATCH 3/5] chore(proxy): use a public counter helper name and narrow row fields The basedpyright ratchet counts private cross-module imports and loose object assignments, so the shared counter helper drops its underscore prefix and the row reads narrow through isinstance. --- .../key_management_endpoints.py | 6 +++--- .../spend_tracking/spend_management_endpoints.py | 14 +++++++------- .../test_spend_management_endpoints.py | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 896f12cb0c7..dd93ca5fba0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6399,7 +6399,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 @@ -6515,7 +6515,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( @@ -6589,7 +6589,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 6ad2284717f..3ef8ba150c4 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -305,20 +305,20 @@ async def _reset_spend_counters(prisma_client: PrismaClient) -> None: floor marker and broadcast) rather than deleting it. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( - _set_spend_counter_with_floor_and_broadcast, + set_spend_counter_with_floor_and_broadcast, ) key_rows: Final = await _verification_token_table(prisma_client).find_many(where={}) for row in key_rows: - token: Final[str | None] = row.model_dump().get("token") - if token: - await _set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:key:{token}", value=0.0) + token: Final[object] = row.model_dump().get("token", "") + if isinstance(token, str) and token: + await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:key:{token}", value=0.0) team_rows: Final = await _team_table(prisma_client).find_many(where={}) for row in team_rows: - team_id: Final[str | None] = row.model_dump().get("team_id") - if team_id: - await _set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:team:{team_id}", value=0.0) + team_id: Final[object] = row.model_dump().get("team_id", "") + if isinstance(team_id, str) and team_id: + await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:team:{team_id}", value=0.0) def _spend_logs_daily_summary_sql( 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 fbefe0c916a..a8988a8271c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8105,7 +8105,7 @@ def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): monkeypatch.setattr( kme, - "_set_spend_counter_with_floor_and_broadcast", + "set_spend_counter_with_floor_and_broadcast", mock_set_counter, ) From f010e34888a5d1b2a5b2067a33af4349232640c5 Mon Sep 17 00:00:00 2001 From: libaojiang <2461443384@qq.com> Date: Wed, 14 Oct 2026 19:34:05 +0800 Subject: [PATCH 4/5] chore(proxy): drop Final on loop-local counter reads basedpyright flags Final assignment inside a loop, and the loop-local reads are implicitly final anyway. --- litellm/proxy/spend_tracking/spend_management_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3ef8ba150c4..43d400fb5b8 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -310,13 +310,13 @@ async def _reset_spend_counters(prisma_client: PrismaClient) -> None: key_rows: Final = await _verification_token_table(prisma_client).find_many(where={}) for row in key_rows: - token: Final[object] = row.model_dump().get("token", "") + token = row.model_dump().get("token", "") if isinstance(token, str) and token: await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:key:{token}", value=0.0) team_rows: Final = await _team_table(prisma_client).find_many(where={}) for row in team_rows: - team_id: Final[object] = row.model_dump().get("team_id", "") + team_id = row.model_dump().get("team_id", "") if isinstance(team_id, str) and team_id: await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:team:{team_id}", value=0.0) From 69354fa1e7f2d7cb4c98e01a06c7a4b18e9560ea Mon Sep 17 00:00:00 2001 From: libaojiang <2461443384@qq.com> Date: Thu, 22 Oct 2026 16:53:22 +0800 Subject: [PATCH 5/5] fix(spend): batch live counter resets and restore integration regression --- .../spend_management_endpoints.py | 34 +++++++++---------- .../spend/test_global_spend_reset.py | 4 --- .../test_spend_management_endpoints.py | 28 ++++++--------- 3 files changed, 27 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f7acc8180a0..4f3cf729c9a 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 @@ -297,29 +298,26 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken return VerificationTokenRepository(prisma_client).table -async def _reset_spend_counters(prisma_client: PrismaClient) -> None: - """Zero every live spend counter a budget check reads. +_SPEND_COUNTER_RESET_BATCH_SIZE: Final = 100 - Budget enforcement reads the Redis spend counters first, so zeroing only - the DB rows leaves an exhausted key refused with the stale pre-reset - counter. Mirrors /key/{key}/reset_spend, setting each counter to 0 (with - floor marker and broadcast) rather than deleting it. - """ + +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, ) - key_rows: Final = await _verification_token_table(prisma_client).find_many(where={}) - for row in key_rows: - token = row.model_dump().get("token", "") - if isinstance(token, str) and token: - await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:key:{token}", value=0.0) - - team_rows: Final = await _team_table(prisma_client).find_many(where={}) - for row in team_rows: - team_id = row.model_dump().get("team_id", "") - if isinstance(team_id, str) and team_id: - await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:team:{team_id}", value=0.0) + 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( diff --git a/tests/integration/spend/test_global_spend_reset.py b/tests/integration/spend/test_global_spend_reset.py index c400d90e184..19241c27579 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 b4102dee1fd..40b66b04997 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8182,7 +8182,8 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] -def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): +@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 @@ -8212,7 +8213,9 @@ def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): class MockDB: def __init__(self): - self.litellm_verificationtoken = MockTable([Row({"token": "key-1"}), Row({"token": "key-2"})]) + 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: @@ -8230,22 +8233,13 @@ def test_global_spend_reset_zeroes_live_spend_counters(monkeypatch): mock_set_counter, ) - async def run_reset(): - import litellm.proxy.proxy_server as ps - - ps.prisma_client = MockPrismaClient() - try: - return await global_spend_reset() - finally: - ps.prisma_client = None - - response: Final = asyncio.run(run_reset()) + 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) == [ - ("spend:key:key-1", 0.0), - ("spend:key:key-2", 0.0), - ("spend:team:team-1", 0.0), - ], reset_calls + 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