diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 05949ed90ca..1ae106be390 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -734,7 +734,7 @@ async def _reconcile_budget_reservation_before_db_update( "Failed to invalidate budget reservation counters after pre-persist reconcile failed" ) finally: - budget_reservation["finalized"] = True + budget_reservation["finalized"] = True # rebind-ok: the counter update reads the stamp off the shared dict async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1c931863a2f..b531a6b0757 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3072,7 +3072,7 @@ async def _reconcile_budget_reservation_for_counter_update( budget_reservation: dict | None, response_cost: float | None, ) -> set[str]: - if budget_reservation is None: + if budget_reservation is None or budget_reservation.get("finalized") is True: return set() from litellm.proxy.spend_tracking.budget_reservation import ( 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 68462065393..0731c233fef 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -921,6 +921,30 @@ 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_for_counter_update_finalized_reservation_falls_back_to_direct_increment( + monkeypatch, +): + """A reservation already finalized before the counter update (the pre-persist + reconcile failed and dropped its counters) must not shield its keys from the + direct increment, or the settled cost is never added back after the drop.""" + import litellm.proxy.spend_tracking.budget_reservation as br + + fake_reconcile = AsyncMock() + monkeypatch.setattr(br, "reconcile_budget_reservation", fake_reconcile) + + result = await ps._reconcile_budget_reservation_for_counter_update( + budget_reservation={ + "finalized": True, + "entries": [{"counter_key": "spend:key:abc"}], + }, + response_cost=1.0, + ) + + assert result == set() + fake_reconcile.assert_not_awaited() + + # --------------------------------------------------------------------------- # _prepare_end_user_and_tag_spend_increments # ---------------------------------------------------------------------------