From 8311456cfc61d888bffa72ddbced57d1aa2f9e7f Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 14:48:55 -0700 Subject: [PATCH] invalidate reservations after release cleanup failure --- litellm/proxy/hooks/proxy_track_cost_callback.py | 8 ++++++++ .../proxy/hooks/test_proxy_track_cost_callback.py | 14 +++++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 9ea6c65eeb5..c1e039aff51 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -443,6 +443,14 @@ async def _update_database_and_spend_counters( verbose_proxy_logger.exception( "Failed to release budget reservation after database update failed" ) + try: + await _invalidate_budget_reservation_counters( + budget_reservation=budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after release failed" + ) raise try: diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index c0ae488d47b..614697a489f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -338,6 +338,11 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re patch( "litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception", ) as mock_log_exception, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback._invalidate_budget_reservation_counters", + new_callable=AsyncMock, + side_effect=RuntimeError("invalidate unavailable"), + ) as mock_invalidate_budget_reservation_counters, ): with pytest.raises(RuntimeError) as exc_info: await _update_database_and_spend_counters( @@ -360,9 +365,16 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re mock_release_budget_reservation.assert_awaited_once_with( budget_reservation=budget_reservation, ) - mock_log_exception.assert_called_once_with( + mock_invalidate_budget_reservation_counters.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + assert mock_log_exception.call_count == 2 + mock_log_exception.assert_any_call( "Failed to release budget reservation after database update failed" ) + mock_log_exception.assert_any_call( + "Failed to invalidate budget reservation counters after release failed" + ) increment_spend_counters.assert_not_awaited()