fix(proxy): fall back to direct spend increments once the early reconcile has finalized the reservation

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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:

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 (

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
 # ---------------------------------------------------------------------------