From 2bad61d8ecd22eaf4a9445925409228134f06849 Mon Sep 17 00:00:00 2001 From: Mohammad Ali Farhan Date: Sun, 9 Aug 2026 02:56:08 +0530 Subject: [PATCH 1/2] fix(proxy): keep counters reserved when reservation cleanup fails When post-call reconciliation raises, the fallback invalidates the reserved counters and returns an empty set so the caller applies the actual cost directly. That is only safe if invalidation succeeded. The inner handler swallowed its exception and fell through to the same `return set()`, so a failed invalidation was indistinguishable from a successful one. The counter then kept its reservation and took the full actual cost on top, leaving previous_spend + reserved_cost + actual_cost. That over-counts, so it shows up as false budget exhaustion and 429s that persist until a durable reseed repairs the counter. Return the reserved keys from the invalidation handler instead, so the caller skips the direct increment and the reservation already on the counter stands in for the cost. Fixes #35569 --- litellm/proxy/proxy_server.py | 5 +- .../proxy/proxy_server/test_spend_counters.py | 74 +++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 90daaaeae6b..be410adde51 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2637,8 +2637,11 @@ async def _reconcile_budget_reservation_for_counter_update( await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( - "Failed to invalidate reserved counters after reservation reconciliation failed" + "Failed to invalidate reserved counters after reservation reconciliation failed; " + "treating them as still reserved so the caller does not add actual cost on top of " + "a reservation that is still on the counter" ) + return reserved_counter_keys return set() return reserved_counter_keys diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..a8e53018857 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -672,6 +672,80 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat assert fake_invalidate.called is True +@pytest.mark.asyncio +async def test_reconcile_budget_reservation_keeps_keys_reserved_when_invalidation_fails( + monkeypatch, +): + """When invalidation also raises, the reservation is still on the counter. Returning + an empty set would tell the caller to add the full actual cost on top of it, leaving + previous_spend + reserved_cost + actual_cost and a counter that can 429 forever.""" + import litellm.proxy.spend_tracking.budget_reservation as br + + monkeypatch.setattr( + br, "get_reserved_counter_keys", MagicMock(return_value={"spend:key:abc"}) + ) + monkeypatch.setattr( + br, "reconcile_budget_reservation", AsyncMock(side_effect=RuntimeError("boom")) + ) + monkeypatch.setattr( + br, + "invalidate_budget_reservation_counters", + AsyncMock(side_effect=RuntimeError("delete failed")), + ) + + result = await ps._reconcile_budget_reservation_for_counter_update( + budget_reservation={"foo": "bar"}, response_cost=1.0 + ) + + assert result == {"spend:key:abc"} + + +@pytest.mark.asyncio +async def test_increment_spend_counters_does_not_double_count_when_cleanup_fails( + monkeypatch, +): + """End-to-end guard: with both reconcile and invalidation failing, the key counter + still holds its reservation, so increment_spend_counters must not increment it a + second time with the actual cost.""" + import litellm.proxy.spend_tracking.budget_reservation as br + + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) + fake_user_cache = _make_user_api_key_cache(get_value=None) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) + + key_counter = "spend:key:hashed-tok" + monkeypatch.setattr( + br, "get_reserved_counter_keys", MagicMock(return_value={key_counter}) + ) + monkeypatch.setattr( + br, "reconcile_budget_reservation", AsyncMock(side_effect=RuntimeError("boom")) + ) + monkeypatch.setattr( + br, + "invalidate_budget_reservation_counters", + AsyncMock(side_effect=RuntimeError("delete failed")), + ) + + await ps.increment_spend_counters( + token="hashed-tok", + team_id="t1", + user_id="u1", + response_cost=5.0, + budget_reservation={"entries": [{"counter_key": key_counter}]}, + ) + + incremented = [ + call.kwargs.get("key", call.args[0] if call.args else None) + for call in fake_cache.redis_cache.async_increment.call_args_list + ] + assert key_counter not in incremented, ( + f"reserved counter was incremented on top of its reservation: {incremented}" + ) + + # --------------------------------------------------------------------------- # _increment_end_user_and_tag_spend_counters # --------------------------------------------------------------------------- From aa7fb3b6e46a1590c0a0b5a123c8d63061f0da47 Mon Sep 17 00:00:00 2001 From: Mohammad Ali Farhan Date: Sun, 9 Aug 2026 03:04:26 +0530 Subject: [PATCH 2/2] test(proxy): drop the unnecessary class patch from the reserve-on-cleanup-failure test The test copied a `SpendCounterReseed.coalesced` patch from its neighbours, but it sets `prisma_client` to None, which already makes the reseed path a no-op. Patching a class attribute couples the test to global state for nothing. Both new tests still fail on the parent commit. --- tests/test_litellm/proxy/proxy_server/test_spend_counters.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index a8e53018857..ecc60ec52d5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -714,7 +714,6 @@ async def test_increment_spend_counters_does_not_double_count_when_cleanup_fails monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) key_counter = "spend:key:hashed-tok" monkeypatch.setattr(