fix(proxy): reconcile budget reservation before enqueuing spend to the DB

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

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -595,6 +595,10 @@ async def _update_database_and_spend_counters(
     request_tags: list[str] | None = None,
     model_access_groups: Sequence[str] | None = None,
 ) -> bool:
+    if budget_reservation is not None:
+        await _reconcile_budget_reservation_before_db_update(
+            budget_reservation=budget_reservation, response_cost=response_cost
+        )
     try:
         charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
             token=user_api_key,
@@ -652,6 +656,19 @@ async def _update_database_and_spend_counters(
     return True
 
 
+async def _reconcile_budget_reservation_before_db_update(budget_reservation: dict, response_cost: float) -> None:
+    from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
+
+    try:
+        await reconcile_budget_reservation(
+            budget_reservation=budget_reservation, actual_cost=response_cost, finalize=False
+        )
+    except Exception:
+        verbose_proxy_logger.debug(
+            "Budget reservation reconcile before DB update failed; deferring to counter update", exc_info=True
+        )
+
+
 async def _release_budget_reservation(budget_reservation: dict | None) -> None:
     if budget_reservation is None:
         return

litellm/proxy/spend_tracking/budget_reservation.py

@@ -905,10 +905,7 @@ async def _set_reserved_entry_actual_cost(
             increment=adjustment,
         )
     elif reseed_on_inconsistent:
-        # Post-call reconcile / release: the counter was flushed, expired or reseeded
-        # between reservation and reconcile, so the optimistic delta no longer applies.
-        # Reseed from the DB floor (which cannot include this request's cost yet) and
-        # add the settled cost, since increment_spend_counters skips reserved keys.
+        # The reconcile runs before this request's spend is enqueued to the DB, so the reseeded floor excludes it.
         reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key)
         if reseeded and actual_cost > 0:
             await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost)

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -667,6 +667,102 @@ async def test_update_database_and_spend_counters_preserves_counter_exception_wh
     proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_update_database_and_spend_counters_reconciles_reservation_before_db_update():
+    call_order: list[str] = []
+    proxy_logging_obj = MagicMock()
+
+    async def _update_database(**kwargs):
+        call_order.append("update_database")
+        return True
+
+    proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=_update_database)
+    increment_spend_counters = AsyncMock()
+    budget_reservation = {"reserved_cost": 0.5, "entries": []}
+
+    async def _reconcile(**kwargs):
+        call_order.append("reconcile")
+
+    with patch(  # test-quality-ok: the helper imports reconcile_budget_reservation in its body, no injection seam
+        "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
+        new_callable=AsyncMock,
+        side_effect=_reconcile,
+    ) as mock_reconcile_budget_reservation:
+        charged = await _update_database_and_spend_counters(
+            proxy_logging_obj=proxy_logging_obj,
+            increment_spend_counters=increment_spend_counters,
+            user_api_key="test_api_key",
+            user_id="test_user_id",
+            end_user_id=None,
+            team_id="test_team_id",
+            org_id="test_org_id",
+            kwargs={},
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+            response_cost=0.2,
+            budget_reservation=budget_reservation,
+        )
+
+    assert charged is True
+    assert call_order == ["reconcile", "update_database"]
+    mock_reconcile_budget_reservation.assert_awaited_once_with(
+        budget_reservation=budget_reservation,
+        actual_cost=0.2,
+        finalize=False,
+    )
+    increment_spend_counters.assert_awaited_once()
+    assert increment_spend_counters.await_args.kwargs["budget_reservation"] is budget_reservation
+
+
+@pytest.mark.asyncio
+async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails_after_early_reconcile():
+    proxy_logging_obj = MagicMock()
+    db_exception = RuntimeError("db unavailable")
+    proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception)
+    increment_spend_counters = AsyncMock()
+    budget_reservation = {"reserved_cost": 0.5, "entries": []}
+
+    with (
+        patch(  # test-quality-ok: the helper imports reconcile_budget_reservation in its body, no injection seam
+            "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
+            new_callable=AsyncMock,
+        ) as mock_reconcile_budget_reservation,
+        patch(  # test-quality-ok: _release_budget_reservation imports the release in its body, no injection seam
+            "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
+            new_callable=AsyncMock,
+        ) as mock_release_budget_reservation,
+    ):
+        with pytest.raises(RuntimeError) as exc_info:
+            await _update_database_and_spend_counters(
+                proxy_logging_obj=proxy_logging_obj,
+                increment_spend_counters=increment_spend_counters,
+                user_api_key="test_api_key",
+                user_id="test_user_id",
+                end_user_id=None,
+                team_id="test_team_id",
+                org_id="test_org_id",
+                kwargs={},
+                completion_response=None,
+                start_time=datetime.now(),
+                end_time=datetime.now(),
+                response_cost=0.2,
+                budget_reservation=budget_reservation,
+            )
+
+        assert exc_info.value is db_exception
+        mock_reconcile_budget_reservation.assert_awaited_once_with(
+            budget_reservation=budget_reservation,
+            actual_cost=0.2,
+            finalize=False,
+        )
+        mock_release_budget_reservation.assert_awaited_once_with(
+            budget_reservation=budget_reservation,
+        )
+
+    increment_spend_counters.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_track_cost_callback_skips_when_no_standard_logging_object():
     """

tests/test_litellm/proxy/test_budget_reservation.py

@@ -2266,6 +2266,64 @@ async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced(
     assert reservation["finalized"] is True
 
 
+@pytest.mark.asyncio
+async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands_between_passes(
+    spend_counter_state,
+):
+    """The early reconcile (before the spend row is enqueued) reseeds from a DB
+    floor that cannot yet include this request. When the periodic flush commits
+    the row before increment_spend_counters runs its second reconcile, the
+    applied_adjustment early-return must keep the counter from adding the cost
+    a second time."""
+    import litellm.proxy.proxy_server as ps
+    from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
+
+    counter_cache, _ = spend_counter_state
+    counter_key = "spend:team_member:user-flush:team-flush"
+    redis_cache = _ExpiringRedisCache()
+    counter_cache.redis_cache = redis_cache
+    counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6)
+
+    reservation = {
+        "reserved_cost": 0.6,
+        "entries": [
+            {
+                "counter_key": counter_key,
+                "entity_type": "TeamMember",
+                "entity_id": "user-flush:team-flush",
+                "reserved_cost": 0.6,
+                "applied_adjustment": 0.0,
+            }
+        ],
+        "finalized": False,
+    }
+
+    with patch.object(  # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for
+        ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.3)
+    ):
+        await reconcile_budget_reservation(
+            budget_reservation=reservation, actual_cost=0.05, finalize=False
+        )
+
+    assert redis_cache.store[counter_key] == pytest.approx(0.35)
+    assert reservation["entries"][0]["applied_adjustment"] == pytest.approx(-0.55)
+    assert reservation["finalized"] is False
+
+    with patch.object(  # test-quality-ok: the flush landing between the passes makes the DB floor include this request
+        ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.35)
+    ):
+        await ps.increment_spend_counters(
+            token="key-flush",
+            team_id="team-flush",
+            user_id="user-flush",
+            response_cost=0.05,
+            budget_reservation=reservation,
+        )
+
+    assert redis_cache.store[counter_key] == pytest.approx(0.35)
+    assert reservation["finalized"] is True
+
+
 @pytest.mark.asyncio
 async def test_should_invalidate_reserved_counters_after_persisted_spend_failure(
     spend_counter_state,