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()