From e91c6ca78964c2e8967462cbba7246ff4c36d042 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Mon, 14 Sep 2026 23:22:01 +0000
Subject: [PATCH] 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>
---
 .../proxy/hooks/proxy_track_cost_callback.py  |  2 +-
 litellm/proxy/proxy_server.py                 |  2 +-
 .../proxy/proxy_server/test_spend_counters.py | 24 +++++++++++++++++++
 3 files changed, 26 insertions(+), 2 deletions(-)

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