From 6a42c661fd6cd4ae9c86341c0a7b81d6486d8467 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 22:05:51 +0000 Subject: [PATCH] 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> --- .../proxy/hooks/proxy_track_cost_callback.py | 17 ++++ .../spend_tracking/budget_reservation.py | 5 +- .../hooks/test_proxy_track_cost_callback.py | 96 +++++++++++++++++++ .../proxy/test_budget_reservation.py | 58 +++++++++++ 4 files changed, 172 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..591df37a7d1 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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 diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index d9f6e33c43c..1e659c4ee4f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/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) 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 eff892f2d80..9f8473f90e0 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 @@ -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(): """ diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..7313b4b2806 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/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,