From 66c0fe23da36666373918ebdb6f8dd342d0bdba5 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 22:55:26 -0700 Subject: [PATCH] handle bad reservation counters after spend write --- litellm/proxy/proxy_server.py | 25 ++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 48 +++++++++++++++++++ 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e77561a5e59..1a9a53c9840 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1994,17 +1994,32 @@ async def _reconcile_budget_reservation_for_counter_update( from litellm.proxy.spend_tracking.budget_reservation import ( get_reserved_counter_keys, + invalidate_budget_reservation_counters, reconcile_budget_reservation, ) reserved_counter_keys = get_reserved_counter_keys( budget_reservation=budget_reservation ) - await reconcile_budget_reservation( - budget_reservation=budget_reservation, - actual_cost=response_cost or 0.0, - finalize=False, - ) + try: + await reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=response_cost or 0.0, + finalize=False, + ) + except Exception: + verbose_proxy_logger.warning( + "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing", + exc_info=True, + ) + try: + await invalidate_budget_reservation_counters( + budget_reservation=budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate reserved counters after reservation reconciliation failed" + ) return reserved_counter_keys diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2acb5493775..669bc85856d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5532,6 +5532,54 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): ps.spend_counter_cache = orig_counter +@pytest.mark.asyncio +async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import increment_spend_counters + + counter_cache = DualCache() + budget_reservation = { + "reserved_cost": 0.5, + "entries": [ + { + "counter_key": "spend:key:key-bad-reserved-counter", + "entity_type": "Key", + "entity_id": "key-bad-reserved-counter", + "reserved_cost": 0.5, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + import litellm.proxy.proxy_server as ps + + orig_counter = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + try: + with patch( + "litellm.proxy.proxy_server.verbose_proxy_logger.warning" + ) as mock_warning: + await increment_spend_counters( + token="key-bad-reserved-counter", + team_id=None, + user_id=None, + response_cost=0.25, + budget_reservation=budget_reservation, + ) + + mock_warning.assert_called_once() + assert budget_reservation["finalized"] is True + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-bad-reserved-counter" + ) + is None + ) + finally: + ps.spend_counter_cache = orig_counter + + @pytest.mark.asyncio async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure(): from litellm.caching.dual_cache import DualCache