diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3a61f773001..d0efdc23ea4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -411,22 +411,28 @@ class _ProxyDBLogger(CustomLogger): "that no poll task will settle" ) return - await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls. - if sl_object is None and ( + skippable_non_model_call = sl_object is None and ( not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime") - ): + ) + if skippable_non_model_call: + await _release_budget_reservation(budget_reservation=budget_reservation) verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object for call_type=%s", kwargs.get("call_type", "unknown"), ) return - if kwargs.get("stream") is not True or ( - kwargs.get("stream") is True and "complete_streaming_response" in kwargs - ): + if _is_completed_proxy_cost_call(kwargs): + # Releasing to $0 treats the call as free. Leaving the hold + # open is also wrong: the next priced request only + # reconciles its own reservation, so this one would keep + # blocking shared counters until TTL. Settle at the + # admission estimate instead. No spend-log row. There is + # no real cost to write. + await _settle_unpriced_success_reservation(budget_reservation=budget_reservation) if sl_object is not None: cost_tracking_failure_debug_info: dict | str = ( sl_object["response_cost_failure_debug_info"] @@ -438,6 +444,7 @@ class _ProxyDBLogger(CustomLogger): raise Exception( f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" ) + await _release_budget_reservation(budget_reservation=budget_reservation) except Exception as e: error_msg = f"Error in tracking cost callback - {e}\n Traceback:{traceback.format_exc()}" model = kwargs.get("model", "") @@ -709,6 +716,33 @@ async def _update_database_and_spend_counters( return True +def _is_completed_proxy_cost_call(kwargs: dict) -> bool: + return kwargs.get("stream") is not True or ( + kwargs.get("stream") is True + and ("complete_streaming_response" in kwargs or "async_complete_streaming_response" in kwargs) + ) + + +async def _settle_unpriced_success_reservation(budget_reservation: dict | None) -> None: + reserved_cost = float(budget_reservation.get("reserved_cost") or 0.0) if budget_reservation else 0.0 + try: + await _reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=reserved_cost, + ) + except Exception: # noqa: BLE001 # settle can fail on cache/redis; still raise cost-tracking after invalidating + verbose_proxy_logger.exception("Failed to settle budget reservation after unpriced successful call") + try: + await _invalidate_budget_reservation_counters( + budget_reservation=budget_reservation, + ) + except Exception: # noqa: BLE001 # invalidate is best-effort so the outer cost-tracking error still surfaces + verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after settle failed") + finally: + if budget_reservation is not None: + budget_reservation["finalized"] = True + + async def _release_budget_reservation(budget_reservation: dict | None) -> None: if budget_reservation is None: return @@ -722,6 +756,23 @@ async def _release_budget_reservation(budget_reservation: dict | None) -> None: ) +async def _reconcile_budget_reservation( + budget_reservation: dict | None, # mutable-ok: same reservation payload _release_budget_reservation takes + actual_cost: float, +) -> None: + if budget_reservation is None: + return + + from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, + ) + + await reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=actual_cost, + ) + + async def _invalidate_budget_reservation_counters( budget_reservation: dict | None, ) -> None: diff --git a/test-quality-budget.json b/test-quality-budget.json index 3c12371f02f..f382f2479a9 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 737 + "limit": 736 }, "TQ003": { "limit": 62 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 395ce68ec54..a2c15653cc8 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 @@ -390,7 +390,13 @@ async def test_track_cost_callback_releases_budget_reservation_when_spend_tracki @pytest.mark.asyncio -async def test_track_cost_callback_releases_budget_reservation_when_response_cost_missing(): +async def test_track_cost_callback_settles_budget_reservation_when_response_cost_missing(): + """A successful unpriced model call must not be refunded to $0. + + Settling at the admission estimate converts the hold into budget spend + without inventing a spend-log row. Health checks still release, covered + separately. + """ logger = _ProxyDBLogger() budget_reservation = {"reserved_cost": 0.5, "entries": []} user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) @@ -412,13 +418,21 @@ async def test_track_cost_callback_releases_budget_reservation_when_response_cos } with ( - patch( + patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam "litellm.proxy.proxy_server.proxy_logging_obj", ) as mock_proxy_logging, - patch( + patch( # test-quality-ok: assert the hold is settled, not refunded through release_budget_reservation "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock, ) as mock_release_budget_reservation, + patch( # test-quality-ok: settle is a proxy-internal reservation call, not an HTTP boundary + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + ) as mock_reconcile_budget_reservation, + patch( # test-quality-ok: unpriced settle must not write a spend-log row + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, ): mock_proxy_logging.failed_tracking_alert = AsyncMock() @@ -429,9 +443,175 @@ async def test_track_cost_callback_releases_budget_reservation_when_response_cos end_time=datetime.now(), ) + mock_release_budget_reservation.assert_not_awaited() + mock_reconcile_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + actual_cost=0.5, + ) + mock_update_database.assert_not_called() + mock_proxy_logging.failed_tracking_alert.assert_called() + # The hold is converted, not dropped: later traffic still sees the reserved cost. + assert budget_reservation["reserved_cost"] == 0.5 + assert budget_reservation.get("finalized") is not True + + +@pytest.mark.asyncio +async def test_track_cost_callback_settles_async_stream_when_response_cost_missing(): + """Async streams record completion on async_complete_streaming_response.""" + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) + + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "stream": True, + "async_complete_streaming_response": {"usage": {"total_tokens": 10}}, + "litellm_params": { + "metadata": { + "user_api_key_auth": user_api_key_auth, + }, + }, + "standard_logging_object": { + "response_cost": None, + "response_cost_failure_debug_info": "missing custom price", + "request_tags": None, + }, + } + + with ( + patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + patch( # test-quality-ok: async-complete streams must settle, not release + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + patch( # test-quality-ok: settle is a proxy-internal reservation call, not an HTTP boundary + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + ) as mock_reconcile_budget_reservation, + patch( # test-quality-ok: unpriced settle must not write a spend-log row + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_release_budget_reservation.assert_not_awaited() + mock_reconcile_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + actual_cost=0.5, + ) + mock_update_database.assert_not_called() + mock_proxy_logging.failed_tracking_alert.assert_called() + assert budget_reservation["reserved_cost"] == 0.5 + assert budget_reservation.get("finalized") is not True + + +@pytest.mark.asyncio +async def test_track_cost_callback_invalidates_reservation_when_settle_fails(): + """A failed settle must not leave the hold pinning later traffic until TTL.""" + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) + + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key_auth": user_api_key_auth, + }, + }, + "standard_logging_object": { + "response_cost": None, + "response_cost_failure_debug_info": "missing custom price", + "request_tags": None, + }, + "stream": False, + } + + with ( + patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + patch( # test-quality-ok: a failed settle must not refund through release_budget_reservation + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + patch( # test-quality-ok: force reconcile to fail so the invalidate path can be observed + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + side_effect=RuntimeError("redis down"), + ), + patch( # test-quality-ok: invalidate is the only way to unpin counters after settle fails + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new_callable=AsyncMock, + ) as mock_invalidate_budget_reservation_counters, + patch( # test-quality-ok: failed settle must not write a spend-log row + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_release_budget_reservation.assert_not_awaited() + mock_invalidate_budget_reservation_counters.assert_awaited_once() + settled = mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] + assert settled["reserved_cost"] == 0.5 + assert settled["finalized"] is True + mock_update_database.assert_not_called() + mock_proxy_logging.failed_tracking_alert.assert_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_budget_reservation_for_non_model_calls(): + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) + + kwargs = { + "call_type": "health", + "litellm_params": { + "metadata": { + "user_api_key_auth": user_api_key_auth, + }, + }, + "stream": False, + } + + with patch( # test-quality-ok: health checks have no cost row; release is the observable contract + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation: + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + mock_release_budget_reservation.assert_awaited_once_with( budget_reservation=budget_reservation, ) + # Health checks refund the hold; they must not stamp it settled. + assert budget_reservation.get("finalized") is not True + assert budget_reservation["reserved_cost"] == 0.5 def test_get_budget_reservation_from_metadata_handles_dict_auth_object():