From f8d187785d82deb8d9b74b87aea73bc537e0019a Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 16:52:09 -0700 Subject: [PATCH] finalize invalidated budget reservations --- litellm/proxy/hooks/proxy_track_cost_callback.py | 1 + tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index a61fda91fd..86f7849bb1 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -489,6 +489,7 @@ async def _update_database_and_spend_counters( await _invalidate_budget_reservation_counters( budget_reservation=budget_reservation ) + budget_reservation["finalized"] = True except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after spend counter update failed" 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 9f414f560c..351d3059d4 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 @@ -450,6 +450,7 @@ async def test_update_database_and_spend_counters_invalidates_reservation_when_c mock_invalidate_budget_reservation_counters.assert_awaited_once_with( budget_reservation=budget_reservation, ) + assert budget_reservation["finalized"] is True proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()