From 633827722473eb6d3e2091560cda3c2853cf9045 Mon Sep 17 00:00:00 2001 From: er-s-an Date: Sun, 11 Oct 2026 09:27:46 +0800 Subject: [PATCH 1/2] fix(proxy): reconcile budget reservations when no database logger Without a prisma client, cost_tracking registered nothing while auth still reserved spend and bound the reservation to success callbacks. Nothing reconciled on success, so the in-memory reservation stuck until TTL and later requests failed with BudgetExceededError under custom_auth_run_common_checks. Register a success-only BudgetReservationReconcileLogger when prisma is absent so reservations settle to response_cost (or release on reconcile failure). Fixes #45895. --- .../budget_reservation_reconcile_logger.py | 61 ++++++++ litellm/proxy/proxy_server.py | 11 ++ ...est_budget_reservation_reconcile_logger.py | 146 ++++++++++++++++++ .../unit/proxy/proxy_server/test_lifecycle.py | 11 +- 4 files changed, 226 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/hooks/budget_reservation_reconcile_logger.py create mode 100644 tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py diff --git a/litellm/proxy/hooks/budget_reservation_reconcile_logger.py b/litellm/proxy/hooks/budget_reservation_reconcile_logger.py new file mode 100644 index 00000000000..b63d17ddf19 --- /dev/null +++ b/litellm/proxy/hooks/budget_reservation_reconcile_logger.py @@ -0,0 +1,61 @@ +"""Reconcile in-memory budget reservations when no database logger is registered. + +See https://github.com/BerriAI/litellm/issues/45895. +""" + +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + budget_reservation_from_metadata, + get_litellm_metadata_from_kwargs, +) + +class BudgetReservationReconcileLogger(CustomLogger): + """Settle auth's spend-counter reservation on success when ProxyDBLogger is absent. + + With ``custom_auth_run_common_checks`` and no ``database_url``, auth still reserves + spend in the in-memory counter and ``bind_budget_reservation_to_callbacks`` marks + ``callback_bound=True``. ``ProxyDBLogger`` (the only success path that reconciled) + is registered only when ``prisma_client is not None``, so the reservation stuck + until TTL and every later request on that key failed with ``BudgetExceededError``. + """ + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + del response_obj, start_time, end_time + metadata: Final = get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)) + if not isinstance(metadata, Mapping): + return + budget_reservation: Final = budget_reservation_from_metadata(metadata) + if not isinstance(budget_reservation, dict) or budget_reservation.get("finalized") is True: + return + + sl_object = kwargs.get("standard_logging_object") + response_cost: object + if isinstance(sl_object, Mapping): + response_cost = sl_object.get("response_cost") + else: + response_cost = kwargs.get("response_cost") + actual_cost: Final = float(response_cost or 0.0) + + from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + + try: + await reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=actual_cost, + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to reconcile budget reservation without a database logger; releasing instead" + ) + from litellm.proxy.spend_tracking.budget_reservation import ( + release_or_invalidate_budget_reservation, + ) + + await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 089647c6297..8f3596811c8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2902,6 +2902,17 @@ def cost_tracking(): litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer)) litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer)) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) + else: + # No DB logger means nothing reconciled auth's in-memory reservation after + # success (callback_bound=True skipped the request-end release). Register a + # success-only reconciler so no-database custom_auth budgets do not stick. #45895 + from litellm.proxy.hooks.budget_reservation_reconcile_logger import ( + BudgetReservationReconcileLogger, + ) + + litellm.logging_callback_manager.add_litellm_async_success_callback( + BudgetReservationReconcileLogger() + ) async def _drain_spend_event_producer_on_shutdown() -> None: diff --git a/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py b/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py new file mode 100644 index 00000000000..35024d9a7d9 --- /dev/null +++ b/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py @@ -0,0 +1,146 @@ +"""Unit tests for BudgetReservationReconcileLogger (#45895).""" + +from datetime import UTC, datetime +from unittest.mock import AsyncMock + +import pytest + +from litellm.proxy.hooks.budget_reservation_reconcile_logger import ( + BudgetReservationReconcileLogger, +) + + +def _reservation(callback_bound: bool = True) -> dict: + return { + "reserved_cost": 0.00255, + "entries": [{"counter_key": "spend:key:hashed", "reserved_cost": 0.00255}], + "finalized": False, + "callback_bound": callback_bound, + "input_cost": 0.0, + "input_tokens": None, + } + + +@pytest.mark.asyncio +async def test_success_reconciles_reservation_to_response_cost(monkeypatch: pytest.MonkeyPatch) -> None: + reservation = _reservation() + reconcile = AsyncMock(return_value=()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + reconcile, + ) + logger = BudgetReservationReconcileLogger() + + await logger.async_log_success_event( + { + "litellm_params": { + "metadata": {"user_api_key_budget_reservation": reservation}, + }, + "response_cost": 0.0, + }, + response_obj=object(), + start_time=datetime.now(UTC), + end_time=datetime.now(UTC), + ) + + reconcile.assert_awaited_once() + assert reconcile.await_args.kwargs["budget_reservation"] is reservation + assert reconcile.await_args.kwargs["actual_cost"] == 0.0 + + +@pytest.mark.asyncio +async def test_success_uses_standard_logging_response_cost(monkeypatch: pytest.MonkeyPatch) -> None: + reservation = _reservation() + reconcile = AsyncMock(return_value=()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + reconcile, + ) + logger = BudgetReservationReconcileLogger() + + await logger.async_log_success_event( + { + "litellm_params": { + "litellm_metadata": {"user_api_key_budget_reservation": reservation}, + }, + "standard_logging_object": {"response_cost": 0.001}, + "response_cost": 9.9, # ignored when SL payload present + }, + response_obj=object(), + start_time=datetime.now(UTC), + end_time=datetime.now(UTC), + ) + + assert reconcile.await_args.kwargs["actual_cost"] == 0.001 + + +@pytest.mark.asyncio +async def test_success_noop_when_no_reservation(monkeypatch: pytest.MonkeyPatch) -> None: + reconcile = AsyncMock(return_value=()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + reconcile, + ) + logger = BudgetReservationReconcileLogger() + + await logger.async_log_success_event( + {"litellm_params": {"metadata": {}}}, + response_obj=object(), + start_time=datetime.now(UTC), + end_time=datetime.now(UTC), + ) + + reconcile.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_success_noop_when_already_finalized(monkeypatch: pytest.MonkeyPatch) -> None: + reservation = _reservation() + reservation["finalized"] = True + reconcile = AsyncMock(return_value=()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + reconcile, + ) + logger = BudgetReservationReconcileLogger() + + await logger.async_log_success_event( + { + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + "response_cost": 0.0, + }, + response_obj=object(), + start_time=datetime.now(UTC), + end_time=datetime.now(UTC), + ) + + reconcile.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_success_falls_back_to_release_when_reconcile_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + reservation = _reservation() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + AsyncMock(side_effect=RuntimeError("counter down")), + ) + release = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.budget_reservation.release_or_invalidate_budget_reservation", + release, + ) + logger = BudgetReservationReconcileLogger() + + await logger.async_log_success_event( + { + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + "response_cost": 0.0, + }, + response_obj=object(), + start_time=datetime.now(UTC), + end_time=datetime.now(UTC), + ) + + release.assert_awaited_once_with(budget_reservation=reservation) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 69f9c72c9c8..c9174858e8e 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -654,18 +654,23 @@ def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeyp } -def test_cost_tracking_no_op_when_prisma_missing(monkeypatch): - """Without a prisma_client cost_tracking is a no-op — not an error.""" +def test_cost_tracking_registers_budget_reconciler_when_prisma_missing(monkeypatch): + """Without a prisma_client, still register a success reconciler for in-memory reservations (#45895).""" import litellm + from litellm.proxy.hooks.budget_reservation_reconcile_logger import ( + BudgetReservationReconcileLogger, + ) monkeypatch.setattr(ps, "prisma_client", None, raising=False) monkeypatch.setattr(litellm, "callbacks", [], raising=False) monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False) cost_tracking() + cost_tracking() # dedupe by logger class key assert litellm.callbacks == [] - assert litellm._async_success_callback == [] + assert len(litellm._async_success_callback) == 1 + assert isinstance(litellm._async_success_callback[0], BudgetReservationReconcileLogger) # --------------------------------------------------------------------------- From 5bdb4e6fcddfa13eb1f054f9831284418777cce3 Mon Sep 17 00:00:00 2001 From: er-s-an <92594860+er-s-an@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:28:45 +0800 Subject: [PATCH 2/2] fix(proxy): defer background reservation settlement until usage arrives --- .../budget_reservation_reconcile_logger.py | 41 +++-- litellm/proxy/proxy_server.py | 4 +- ...est_budget_reservation_reconcile_logger.py | 142 ++++++++++++++++++ 3 files changed, 170 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/hooks/budget_reservation_reconcile_logger.py b/litellm/proxy/hooks/budget_reservation_reconcile_logger.py index b63d17ddf19..16ff6d2d2a3 100644 --- a/litellm/proxy/hooks/budget_reservation_reconcile_logger.py +++ b/litellm/proxy/hooks/budget_reservation_reconcile_logger.py @@ -8,11 +8,15 @@ from datetime import datetime from typing import Final from litellm._logging import verbose_proxy_logger +from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger +from litellm.interactions.background_cost_polling import is_pollable_background_interaction from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, - get_litellm_metadata_from_kwargs, + get_litellm_metadata_from_kwargs, # pyright: ignore[reportUnknownVariableType] # shared helper has a legacy untyped signature ) +from litellm.types.interactions import InteractionsAPIResponse + class BudgetReservationReconcileLogger(CustomLogger): """Settle auth's spend-counter reservation on success when ProxyDBLogger is absent. @@ -27,35 +31,44 @@ class BudgetReservationReconcileLogger(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: - del response_obj, start_time, end_time - metadata: Final = get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)) + metadata: Final = get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # shared helper returns untyped metadata if not isinstance(metadata, Mapping): return - budget_reservation: Final = budget_reservation_from_metadata(metadata) + budget_reservation: Final = budget_reservation_from_metadata(metadata) # pyright: ignore[reportUnknownArgumentType] # helper only reads known keys from the checked mapping if not isinstance(budget_reservation, dict) or budget_reservation.get("finalized") is True: return - sl_object = kwargs.get("standard_logging_object") - response_cost: object - if isinstance(sl_object, Mapping): - response_cost = sl_object.get("response_cost") - else: - response_cost = kwargs.get("response_cost") - actual_cost: Final = float(response_cost or 0.0) + sl_object: Final = kwargs.get("standard_logging_object") + response_cost: Final[object] = ( # pyright: ignore[reportUnknownVariableType] # callback payloads are legacy untyped mappings + sl_object.get("response_cost") # pyright: ignore[reportUnknownMemberType] # float below validates the untyped cost value + if isinstance(sl_object, Mapping) + else kwargs.get("response_cost") + ) + if ( + response_cost is None + and BACKGROUND_INTERACTION_COST_POLLING_ENABLED + and isinstance(response_obj, InteractionsAPIResponse) + and response_obj.usage is None + and is_pollable_background_interaction(response_obj) + ): + return + actual_cost: Final = float(response_cost or 0.0) # pyright: ignore[reportArgumentType] # float validates legacy callback values at runtime - from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, # pyright: ignore[reportUnknownVariableType] # shared reconciler accepts legacy reservation dictionaries + ) try: await reconcile_budget_reservation( budget_reservation=budget_reservation, actual_cost=actual_cost, ) - except Exception: + except Exception: # noqa: BLE001 # a counter backend failure must still release the reservation verbose_proxy_logger.exception( "Failed to reconcile budget reservation without a database logger; releasing instead" ) from litellm.proxy.spend_tracking.budget_reservation import ( - release_or_invalidate_budget_reservation, + release_or_invalidate_budget_reservation, # pyright: ignore[reportUnknownVariableType] # shared release accepts legacy reservation dictionaries ) await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f3596811c8..73db8cf9526 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2910,9 +2910,7 @@ def cost_tracking(): BudgetReservationReconcileLogger, ) - litellm.logging_callback_manager.add_litellm_async_success_callback( - BudgetReservationReconcileLogger() - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(BudgetReservationReconcileLogger()) async def _drain_spend_event_producer_on_shutdown() -> None: diff --git a/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py b/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py index 35024d9a7d9..006678a5739 100644 --- a/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py +++ b/tests/unit/proxy/hooks/test_budget_reservation_reconcile_logger.py @@ -1,13 +1,18 @@ """Unit tests for BudgetReservationReconcileLogger (#45895).""" +from collections.abc import Mapping from datetime import UTC, datetime +from typing import Final from unittest.mock import AsyncMock import pytest +from litellm.caching import DualCache +from litellm.proxy import proxy_server from litellm.proxy.hooks.budget_reservation_reconcile_logger import ( BudgetReservationReconcileLogger, ) +from litellm.types.interactions import InteractionsAPIResponse def _reservation(callback_bound: bool = True) -> dict: @@ -144,3 +149,140 @@ async def test_success_falls_back_to_release_when_reconcile_fails( ) release.assert_awaited_once_with(budget_reservation=reservation) + + +@pytest.fixture +def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache: + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", None) + return cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ("in_progress", "queued")) +@pytest.mark.parametrize("standard_logging", (False, True)) +async def test_background_interaction_keeps_reservation_until_final_cost( + spend_counter_cache: DualCache, status: str, standard_logging: bool +) -> None: + reservation: Final = _reservation() + counter_key: Final = "spend:key:hashed" + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=reservation["reserved_cost"]) + logger: Final = BudgetReservationReconcileLogger() + now: Final = datetime(2026, 1, 1, tzinfo=UTC) + kwargs: Final = { + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + **({"standard_logging_object": {"response_cost": None}, "response_cost": 9.9} if standard_logging else {}), + } + pending: Final = InteractionsAPIResponse(id="interactions/background", status=status) + + await logger.async_log_success_event(kwargs, pending, now, now) + + assert reservation["finalized"] is False + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + + completed: Final = pending.model_copy(update={"status": "completed", "usage": {"total_tokens": 12}}) + final_kwargs: Final = {**kwargs, "standard_logging_object": {"response_cost": 0.004}} + await logger.async_log_success_event(final_kwargs, completed, now, now) + + assert reservation["finalized"] is True + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.004) + + await logger.async_log_success_event(final_kwargs, completed, now, now) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.004) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "interaction_id", "polling_enabled", "has_usage"), + ( + ("in_progress", "interactions/background", False, False), + ("queued", "interactions/background", False, False), + ("in_progress", None, True, False), + ("queued", "", True, False), + ("completed", "interactions/background", True, False), + ("failed", "interactions/background", True, False), + ("cancelled", "interactions/background", True, False), + ("requires_action", "interactions/background", True, False), + ("in_progress", "interactions/background", True, True), + ), +) +async def test_missing_cost_releases_reservation_when_no_usage_poll_will_settle_it( + spend_counter_cache: DualCache, + monkeypatch: pytest.MonkeyPatch, + status: str, + interaction_id: str | None, + polling_enabled: bool, + has_usage: bool, +) -> None: + monkeypatch.setattr( + "litellm.proxy.hooks.budget_reservation_reconcile_logger.BACKGROUND_INTERACTION_COST_POLLING_ENABLED", + polling_enabled, + ) + reservation: Final = _reservation() + counter_key: Final = "spend:key:hashed" + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=reservation["reserved_cost"]) + logger: Final = BudgetReservationReconcileLogger() + now: Final = datetime(2026, 1, 1, tzinfo=UTC) + + await logger.async_log_success_event( + {"litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}}, + InteractionsAPIResponse(id=interaction_id, status=status, usage={"total_tokens": 12} if has_usage else None), + now, + now, + ) + + assert reservation["finalized"] is True + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response_cost", "expected_cost"), + ((0.0, 0.0), (0.004, 0.004), ("0.004", 0.004), (b"0.004", 0.004), (memoryview(b"0.004"), 0.004)), +) +async def test_reported_interaction_cost_is_settled_immediately( + spend_counter_cache: DualCache, response_cost: object, expected_cost: float +) -> None: + reservation: Final = _reservation() + counter_key: Final = "spend:key:hashed" + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=reservation["reserved_cost"]) + logger: Final = BudgetReservationReconcileLogger() + now: Final = datetime(2026, 1, 1, tzinfo=UTC) + + await logger.async_log_success_event( + { + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + "response_cost": response_cost, + }, + InteractionsAPIResponse(id="interactions/background", status="in_progress"), + now, + now, + ) + + assert reservation["finalized"] is True + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(expected_cost) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kwargs", + ( + {}, + {"litellm_params": {"metadata": None}}, + {"litellm_params": {"metadata": "invalid"}}, + {"litellm_params": {"metadata": {1: "ignored"}}}, + ), +) +async def test_missing_or_invalid_metadata_leaves_existing_spend_unchanged( + spend_counter_cache: DualCache, kwargs: Mapping[str, object] +) -> None: + counter_key: Final = "spend:key:hashed" + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.00255) + logger: Final = BudgetReservationReconcileLogger() + now: Final = datetime(2026, 1, 1, tzinfo=UTC) + + await logger.async_log_success_event(kwargs, object(), now, now) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.00255)