diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ce6f77f78a0..5bcde688521 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -359,6 +359,46 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): return {} +def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("budget_reservation") + return getattr(user_api_key_auth, "budget_reservation", None) + + +def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None: + stamped: Final = metadata.get("user_api_key_budget_reservation") + if isinstance(stamped, dict): + return stamped + on_auth_object: Final = _budget_reservation_on_auth_object(metadata.get("user_api_key_auth")) + return on_auth_object if isinstance(on_auth_object, dict) else None + + +def _stamp_budget_reservation_callback_bound(litellm_params: Mapping[str, object], callback_bound: bool) -> None: + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + budget_reservation = budget_reservation_from_metadata(metadata) + if budget_reservation is not None: + budget_reservation["callback_bound"] = callback_bound + + +def bind_budget_reservation_to_callbacks(litellm_params: Mapping[str, object]) -> None: + """Mark the request's budget reservation as owned by the success callbacks of this call. + + The proxy releases any reservation still unbound when the request ends; one bound here + is left for the cost callback, which may finish after the response has been sent. Bind + only where a success handler is guaranteed to run: a logging object merely existing is + not that, since the proxy builds one for every route before calling anything. + """ + _stamp_budget_reservation_callback_bound(litellm_params, True) + + +def unbind_budget_reservation_from_callbacks(litellm_params: Mapping[str, object]) -> None: + """Hand a failed call's reservation back to the request-end release: failure handlers never settle it.""" + _stamp_budget_reservation_callback_bound(litellm_params, False) + + def reconstruct_model_name( model_name: str, custom_llm_provider: str | None, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a486cdeff19..f5967185af0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5493,7 +5493,7 @@ def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. Detached internal sub-calls only inherit the identity keys, so the auth object is the - fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + fallback there, exactly as budget_reservation_from_metadata does for reservations. """ for metadata_variable_name in ("metadata", "litellm_metadata"): metadata = litellm_params.get(metadata_variable_name) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 08f64e610bd..a6c0792a86f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -650,6 +650,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str "type": "http", "headers": scope_headers, "path": ws_scope.get("path", ""), + "state": ws_scope.setdefault("state", {}), # mutable-ok: Starlette's socket state, shared with the request } for key in ("root_path", "app_root_path"): if key in ws_scope: @@ -3086,31 +3087,30 @@ async def _reserve_budget_after_common_checks( request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None - if skip_budget_checks: - return - if general_settings.get("disable_budget_reservation") is True: - return + if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True: + from litellm.proxy.spend_tracking.budget_reservation import ( + reserve_budget_for_request, + ) - from litellm.proxy.spend_tracking.budget_reservation import ( - reserve_budget_for_request, - ) - - user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( - request_body=request_data, - route=route, - llm_router=llm_router, - valid_token=user_api_key_auth_obj, - team_object=team_object, - user_object=user_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - end_user_id=end_user_id, - end_user_object=end_user_object, - apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, - fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, - raw_body=await read_raw_json_body(request=request), - ) + user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( + request_body=request_data, + route=route, + llm_router=llm_router, + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_id=end_user_id, + end_user_object=end_user_object, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, + raw_body=await read_raw_json_body(request=request), + ) + if request is not None: + reservation: Final = user_api_key_auth_obj.budget_reservation + request.state.budget_reservation = reservation # rebind-ok: read by the release middleware def _should_skip_budget_checks( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 08e8e4f8c10..81894a5ff12 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -11,6 +11,7 @@ from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -630,17 +631,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: - metadata_budget_reservation: Final = metadata.get("user_api_key_budget_reservation") - if isinstance(metadata_budget_reservation, dict): - return metadata_budget_reservation - - user_api_key_auth_obj: Final = metadata.get("user_api_key_auth") - if user_api_key_auth_obj is None: - return None - if isinstance(user_api_key_auth_obj, dict): - budget_reservation: Final = user_api_key_auth_obj.get("budget_reservation") - return budget_reservation if isinstance(budget_reservation, dict) else None - return getattr(user_api_key_auth_obj, "budget_reservation", None) + return budget_reservation_from_metadata(metadata) def _get_request_tags_for_cost_tracking( diff --git a/litellm/proxy/middleware/budget_reservation_release_middleware.py b/litellm/proxy/middleware/budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f7ac885274e --- /dev/null +++ b/litellm/proxy/middleware/budget_reservation_release_middleware.py @@ -0,0 +1,33 @@ +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +from starlette.types import ASGIApp, Receive, Scope, Send + +_SCOPES_AUTH_STAMPS: Final = frozenset({"http", "websocket"}) + + +class BudgetReservationReleaseMiddleware: + """Releases the budget reservation auth made for a request once no callback owns it. + + Auth stamps the reservation on the request or socket state; a call that starts + claims it for the cost callbacks, which settle it on success or failure. When the + response has been sent or the socket has closed and the reservation is still + unclaimed, nothing else ever would, so it is released here instead of pinning the + spend counter until its TTL. + """ + + def __init__(self, app: ASGIApp, release: Callable[[Mapping[str, object]], Awaitable[None]]) -> None: + self.app = app + self.release = release + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] not in _SCOPES_AUTH_STAMPS: + await self.app(scope, receive, send) + return + try: + await self.app(scope, receive, send) + finally: + state: Final = scope.get("state") + budget_reservation: Final = state.get("budget_reservation") if isinstance(state, Mapping) else None + if isinstance(budget_reservation, Mapping): + await self.release(budget_reservation) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bbbf5e4b94f..8902e599788 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.constants import ( from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) @@ -1637,6 +1638,7 @@ async def pass_through_request( **kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) ## CUSTOM HEADERS - `x-litellm-*` custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2561,6 +2563,7 @@ async def websocket_passthrough_request( **success_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) # Call the proxy logging success hook if proxy_logging_obj: @@ -2732,6 +2735,7 @@ async def _relay_passthrough_response_bytes( **success_handler_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) def _extract_model_from_vertex_ai_setup(setup_response: Mapping[str, object]) -> str | None: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index fe9e104789b..fae26b5a72d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -9,6 +9,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues @@ -218,6 +219,7 @@ class PassThroughStreamingHandler: and response.status_code < 400 ): logging_scheduled = True + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) @@ -250,6 +252,8 @@ class PassThroughStreamingHandler: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) + else: + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) @staticmethod async def _route_streaming_logging_to_handler( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87da8a6e44e..6fa7a7d761f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -654,6 +654,9 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, BillingRecorder, ) +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) from litellm.proxy.plugin_routes import ( register_plugins_from_config, ) @@ -729,7 +732,10 @@ from litellm.proxy.shutdown.scheduled_jobs import ( pause_scheduled_jobs, stop_in_flight_scheduler_jobs, ) -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.budget_reservation import ( + get_budget_window_start, + release_unbound_budget_reservation, +) from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( run_scheduled_daily_global_spend_reconcile, ) @@ -2358,6 +2364,7 @@ app.add_middleware( # it sees prisma_client as of the first request rather than import time. sink_factory=lambda: gateway_request_accumulator if prisma_client is not None else None, ) +app.add_middleware(BudgetReservationReleaseMiddleware, release=release_unbound_budget_reservation) app.add_middleware(InFlightRequestsMiddleware) app.add_middleware(SecurityHeadersMiddleware) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f9e5c4ff1e4..ac9b07a55f7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -366,6 +366,7 @@ async def reserve_budget_for_request( "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, + "callback_bound": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), "input_tokens": max(input_token_counts.values(), default=None), } @@ -474,6 +475,19 @@ async def release_or_invalidate_budget_reservation( budget_reservation["finalized"] = True +async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None: + """Release a reservation no logging callback took ownership of, once the request ended. + + A handler whose litellm call never builds a logging object (batch cancel, file + content, anything without the client decorator) runs no cost callback, so nothing + else would ever reconcile its reservation. A bound reservation is left alone: its + success or failure handler settles it, possibly after the response has been sent. + """ + if not isinstance(budget_reservation, dict) or budget_reservation.get("callback_bound") is True: + return + await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation) + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index 30aa1d97bfc..6bbf2ffed6b 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -59,9 +59,17 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments) + return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared) + return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous) + + +def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup: + from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks + + if asynchronous and not is_internal_call(): + bind_budget_reservation_to_callbacks(call_setup.logger.litellm_params) + return call_setup def check_limits(kwargs: Mapping[str, object]) -> None: @@ -96,6 +104,9 @@ def finalize( class LoggingSurface(Protocol): + @property + def litellm_params(self) -> Mapping[str, object]: ... + def update_from_kwargs( self, kwargs: dict[str, object], @@ -236,8 +247,12 @@ def sync_success_for_async_call( def failure_handler( logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool ) -> Coroutine[object, object, None] | None: + from litellm.litellm_core_utils.core_helpers import unbind_budget_reservation_from_callbacks + trace: Final = "".join(traceback.format_exception(error)) if asynchronous: + if not is_internal_call(): + unbind_budget_reservation_from_callbacks(logger.litellm_params) return logger.async_failure_handler(error, trace, start, end) logger.failure_handler(error, trace, start, end) return None diff --git a/litellm/utils.py b/litellm/utils.py index d96a1c5ae3f..812299560c7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -81,7 +81,11 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, + normalize_drop_params, + unbind_budget_reservation_from_callbacks, +) from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, @@ -1880,6 +1884,8 @@ def client(original_function): # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" + if not _is_litellm_internal_call: + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) kwargs["litellm_logging_obj"] = logging_obj modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) @@ -2081,6 +2087,7 @@ def client(original_function): # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( e, traceback_exception, start_time, end_time diff --git a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py index 498f0a734a3..ab5fd7f80ee 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py @@ -53,6 +53,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_enabled(): # Setup logging object with model info litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -132,6 +133,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_disabled(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -194,6 +196,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_no_usage_chunk(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -249,6 +252,7 @@ async def test_vertex_ai_anthropic_streaming_model_extraction(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b2ad13c205e..bca61a0e76f 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -6,6 +6,8 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + bind_budget_reservation_to_callbacks, + budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, @@ -13,7 +15,60 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, + unbind_budget_reservation_from_callbacks, ) +from litellm.proxy._types import UserAPIKeyAuth + + +class TestBudgetReservationBinding: + """The request-end release skips a reservation a cost callback has claimed, so the claim + must land on the one dict auth stamped, through whichever metadata field or auth object + carries it, and a failed call must be able to hand it back.""" + + @staticmethod + def _reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + @pytest.mark.parametrize("metadata_variable_name", ["metadata", "litellm_metadata"]) + def test_reservation_stamped_on_the_metadata_is_bound(self, metadata_variable_name: str): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({metadata_variable_name: {"user_api_key_budget_reservation": reservation}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_the_auth_object_is_bound(self): + reservation = self._reservation() + user_api_key_auth = UserAPIKeyAuth(token="hashed") + user_api_key_auth.budget_reservation = reservation + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": user_api_key_auth}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_a_dumped_auth_object_is_bound(self): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": {"budget_reservation": reservation}}}) + + assert reservation["callback_bound"] is True + + def test_unbind_hands_a_claimed_reservation_back(self): + reservation = self._reservation() + litellm_params = {"litellm_metadata": {"user_api_key_budget_reservation": reservation}} + bind_budget_reservation_to_callbacks(litellm_params) + + unbind_budget_reservation_from_callbacks(litellm_params) + + assert reservation["callback_bound"] is False + + def test_request_without_a_reservation_binds_nothing(self): + metadata = {"user_api_key_auth": UserAPIKeyAuth(token="hashed")} + + bind_budget_reservation_to_callbacks({"metadata": metadata, "litellm_metadata": None}) + + assert budget_reservation_from_metadata(metadata) is None + assert "user_api_key_budget_reservation" not in metadata class TestGetOrCreateMetadataBucket: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index a9bbb4992c6..1bfc55d1486 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -26,6 +26,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, @@ -7941,3 +7942,19 @@ def test_response_cost_calculator_prices_terminal_responses_event_from_its_respo assert event_cost is not None and event_cost > 0 assert event_cost == inner_cost assert logging_obj.cost_breakdown["input_cost"] is not None and logging_obj.cost_breakdown["input_cost"] > 0 + + +class TestBudgetReservationBinding: + """The proxy builds a logging object for every route before calling anything, so a + logging object seeing the reservation is no promise that a cost callback will settle + it: the claim belongs to the call wrapper, and this object must leave it unbound.""" + + def test_update_environment_variables_leaves_the_reservation_unbound(self, logging_obj): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + logging_obj.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + + assert logging_obj.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/test_litellm/llms/xai/test_xai_key_fallback.py index 092e4951547..cbc507c5ee3 100644 --- a/tests/test_litellm/llms/xai/test_xai_key_fallback.py +++ b/tests/test_litellm/llms/xai/test_xai_key_fallback.py @@ -12,6 +12,9 @@ from litellm.types.router import GenericLiteLLMParams class FakeLogging: + def __init__(self) -> None: + self.litellm_params: dict = {} + def update_from_kwargs(self, **kwargs): pass diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index da36071a5b4..f03abe8f124 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -59,6 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, get_api_key, user_api_key_auth, + user_api_key_auth_websocket_for_model, ) from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata @@ -9043,3 +9044,90 @@ async def test_router_settings_model_group_alias_authorizes_target_for_team(monk await authorize() assert (await request.json())["model"] == target assert get_client_requested_model(request) == "AgentX-LLM" + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_hands_the_reservation_to_the_request_state(): + from fastapi import Request + + request = Request(scope={"type": "http"}) + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/batches/batch_123/cancel", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + request=request, + ) + + assert user_api_key_auth_obj.budget_reservation is reservation + assert request.state.budget_reservation is reservation + assert request.scope["state"]["budget_reservation"] is reservation + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_clears_the_request_state_when_budget_checks_skip(): + from fastapi import Request + + request = Request(scope={"type": "http", "state": {"budget_reservation": {"reserved_cost": 0.5}}}) + + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=UserAPIKeyAuth(token="test_token"), + request_data={"model": "free-model"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=True, + general_settings={}, + request=request, + ) + + assert request.state.budget_reservation is None + + +@pytest.mark.asyncio +async def test_websocket_auth_hands_the_reservation_to_the_socket_state(): + from fastapi import WebSocket + + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + websocket = WebSocket( + scope={ + "type": "websocket", + "path": "/v1/realtime", + "headers": [(b"authorization", b"Bearer sk-1234")], + "query_string": b"model=gpt-realtime", + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + async def auth_that_reserves(request, api_key): + request.state.budget_reservation = reservation + return UserAPIKeyAuth(token="hashed", budget_reservation=reservation) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + new=AsyncMock(side_effect=auth_that_reserves), + ): + result = await user_api_key_auth_websocket_for_model(websocket, model="gpt-realtime") + + assert result.budget_reservation == reservation + assert websocket.state.budget_reservation is reservation + assert websocket.scope["state"]["budget_reservation"] is reservation diff --git a/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f37a20dff8b --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py @@ -0,0 +1,349 @@ +""" +Tests for BudgetReservationReleaseMiddleware. + +Auth reserves budget before the handler runs and hands the reservation to the +request or socket state. A litellm call made through the async client wrapper +claims it for the cost callback that runs after the call; anything still unclaimed +when the response is done or the socket has closed would keep the spend counter +pinned until its TTL, so the middleware releases it. +""" + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from datetime import datetime +from typing import Final + +import pytest +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route +from starlette.types import ASGIApp, Message, Receive, Scope, Send +from starlette.websockets import WebSocket + +import litellm +from litellm.caching import DualCache +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) +from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, + release_unbound_budget_reservation, + reserve_budget_for_request, +) +from litellm.proxy.utils import ProxyLogging +from litellm.utils import Rules, function_setup + +KEY_TOKEN: Final = "hashed-release-middleware-key" +COUNTER_KEY: Final = f"spend:key:{KEY_TOKEN}" +CHAT_BODY: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + + +@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.fixture +def no_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + for callback_list_name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, callback_list_name, []) + + +async def _reserve() -> dict: + reservation: Final = await reserve_budget_for_request( + request_body=CHAT_BODY, + route="/v1/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth(token=KEY_TOKEN, max_budget=1.0, spend=0.0), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + ) + assert reservation is not None + assert reservation["reserved_cost"] > 0 + return reservation + + +async def _chat(reservation: dict, **kwargs: object) -> object: + return await litellm.acompletion( + **CHAT_BODY, + metadata={"user_api_key_budget_reservation": reservation}, + **kwargs, + ) + + +def _proxy_pre_call_setup(route_type: str, reservation: dict) -> None: + function_setup( + original_function=route_type, + rules_obj=Rules(), + start_time=datetime.now(), + **CHAT_BODY, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + +def _app( + handler: Callable[[Request], Awaitable[Response]], + release: Callable[[Mapping[str, object]], Awaitable[None]] = release_unbound_budget_reservation, +) -> Starlette: + app: Final = Starlette(routes=[Route("/", handler, methods=["POST"])]) + app.add_middleware(BudgetReservationReleaseMiddleware, release=release) + return app + + +async def _post(app: ASGIApp) -> None: + scope: Final = { + "type": "http", + "method": "POST", + "path": "/", + "raw_path": b"/", + "headers": [], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "client": ("testclient", 1), + } + + body_delivered: Final = asyncio.Event() + client_never_disconnects: Final = asyncio.Event() + + async def receive() -> Message: + if body_delivered.is_set(): + await client_never_disconnects.wait() + body_delivered.set() + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + return None + + await app(scope, receive, send) + + +def _counter(spend_counter_cache: DualCache) -> float | None: + return spend_counter_cache.in_memory_cache.get_cache(key=COUNTER_KEY) + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_after_the_response(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + assert _counter(spend_counter_cache) == pytest.approx(reservation["reserved_cost"]) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({"id": "batch_123", "status": "cancelling"}) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + raise RuntimeError("upstream refused the cancel") + + with pytest.raises(RuntimeError): + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_released( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + + async def cancel_batch_without_a_client_wrapper() -> dict: + return {"id": "batch_123", "status": "cancelling"} + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acancel_batch", reservation) + return JSONResponse(await cancel_batch_without_a_client_wrapper()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_of_a_failed_call_is_released_after_the_error_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + refused: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + try: + await _chat(reservation, mock_response=refused) + except litellm.AuthenticationError: + return JSONResponse({"error": {"message": "bad key"}}, status_code=401) + raise AssertionError("the mocked call must fail") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + response: Final = await _chat(reservation, mock_response="ok") + return JSONResponse(response.model_dump()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_that_finishes_after_the_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + stream: Final = await _chat(reservation, mock_response="ok", stream=True) + + async def sse() -> AsyncIterator[bytes]: + async for chunk in stream: + yield f"data: {chunk.model_dump_json()}\n\n".encode() + yield b"data: [DONE]\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + actual_cost: Final = reserved_cost / 4 + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=actual_cost) + + assert _counter(spend_counter_cache) == pytest.approx(actual_cost) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_socket_closes( + spend_counter_cache: DualCache, +): + reservation: Final = await _reserve() + + async def listen_without_a_provider_key(scope: Scope, receive: Receive, send: Send) -> None: + websocket: Final = WebSocket(scope, receive, send) + websocket.state.budget_reservation = reservation + await websocket.close(code=1011, reason="Required 'DEEPGRAM_API_KEY' in environment") + + async def receive() -> Message: + return {"type": "websocket.connect"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware( + listen_without_a_provider_key, release=release_unbound_budget_reservation + ) + await middleware({"type": "websocket", "path": "/deepgram/v1/listen", "headers": []}, receive, send) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_runs_once_per_request_with_the_stamped_reservation(): + released: Final = [] + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({}) + + await _post(_app(handler, release=release)) + + assert released == [reservation] + assert released[0] is reservation + + +@pytest.mark.asyncio +async def test_request_without_a_reservation_releases_nothing(): + released: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def unauthenticated(request: Request) -> Response: + return JSONResponse({}) + + async def budget_checks_skipped(request: Request) -> Response: + request.state.budget_reservation = None + return JSONResponse({}) + + await _post(_app(unauthenticated, release=release)) + await _post(_app(budget_checks_skipped, release=release)) + + assert released == [] + + +@pytest.mark.asyncio +async def test_lifespan_scopes_pass_through(): + released: Final = [] + seen: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def inner(scope: Scope, receive: Receive, send: Send) -> None: + seen.append(scope["type"]) + + async def receive() -> Message: + return {"type": "lifespan.startup"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware(inner, release=release) + await middleware({"type": "lifespan", "state": {"budget_reservation": {"reserved_cost": 1.0}}}, receive, send) + + assert seen == ["lifespan"] + assert released == [] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 6cd537489ac..6ad850866b7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4256,6 +4256,112 @@ async def test_pass_through_request_non_streaming_success_unchanged(): mock_success_handler.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upstream_status_code, claimed_by_the_success_handler", + [(200, True), (500, False)], + ids=["success-claims-the-reservation", "upstream-error-leaves-it-for-the-request-end-release"], +) +async def test_pass_through_request_claims_the_budget_reservation_only_when_its_success_handler_runs( + upstream_status_code: int, claimed_by_the_success_handler: bool +): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=upstream_status_code, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + response = await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert response.status_code == upstream_status_code + assert reservation["callback_bound"] is claimed_by_the_success_handler + assert mock_worker.ensure_initialized_and_enqueue.call_count == int(claimed_by_the_success_handler) + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_the_budget_reservation_for_the_request_end_release_when_its_success_handler_cannot_be_enqueued(): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=refuse_to_enqueue) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + with pytest.raises(ProxyException): + await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert reservation["callback_bound"] is False + + @pytest.mark.asyncio async def test_pass_through_request_internal_failure_still_raises_proxy_exception(): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index ea6adc35b9a..88b82349c83 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -838,3 +838,72 @@ async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exceptio assert failure_payload["completion_tokens"] == 12 assert failure_payload["response_cost"] > 12 * 3.75e-06 assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deferred_dispatch_armed", + [False, True], + ids=["enqueued-at-end-of-stream", "parked-for-deferred-dispatch"], +) +async def test_chunk_processor_claims_the_budget_reservation_before_handing_it_to_the_cost_callback( + deferred_dispatch_armed: bool, +): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + if deferred_dispatch_armed: + logging_obj._on_deferred_stream_complete = AsyncMock() + claimed_when_the_callback_ran = [] + + async def cost_callback(**kwargs): + claimed_when_the_callback_ran.append(reservation["callback_bound"]) + + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=cost_callback, + ): + pass + + if deferred_dispatch_armed: + (parked_cost_callback,) = logging_obj._deferred_stream_complete_args + await parked_cost_callback + else: + await GLOBAL_LOGGING_WORKER.flush() + + assert reservation["callback_bound"] is True + assert claimed_when_the_callback_ran == [True] + + +@pytest.mark.asyncio +async def test_chunk_processor_leaves_the_budget_reservation_for_the_request_end_release_when_the_cost_callback_cannot_be_enqueued(): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=refuse_to_enqueue): + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=AsyncMock(), + ): + pass + + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..214b5cde7da 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -26,6 +26,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, count_request_input_tokens, estimate_request_max_cost, + release_unbound_budget_reservation, reserve_budget_for_request, ) from litellm.proxy.utils import ProxyLogging @@ -546,3 +547,41 @@ async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_d assert counter is not None assert counter.max_budget == expected_max_budget assert counter.fallback_spend == 0.5 + + +@pytest.mark.asyncio +async def test_reservation_starts_unbound_to_any_callback(): + reservation: Final = await _reserve("/v1/responses") + + assert reservation is not None + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + reservation["callback_bound"] = True + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + assert reservation["finalized"] is False diff --git a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py index 1f6a214398a..05f2d13a079 100644 --- a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py @@ -9,9 +9,10 @@ import pytest from pydantic import TypeAdapter import litellm +from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import callbacks_legacy_python as legacy -from litellm.rust_bridge.callbacks_legacy_python import check_limits, setup +from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup _OCR_KWARGS: Final = MappingProxyType( { @@ -81,6 +82,82 @@ def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Map assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] +def _budget_reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + +def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]: + return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}} + + +def test_setup_claims_the_budget_reservation_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_claims_the_budget_reservation_a_supplied_logger_already_saw() -> None: + reservation: Final = _budget_reservation() + supplied: Final = _supplied_logger() + supplied.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + assert reservation["callback_bound"] is False + + setup("aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_leaves_the_budget_reservation_alone_for_a_sync_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=False) + + assert reservation["callback_bound"] is False + + +def test_setup_leaves_the_budget_reservation_alone_for_an_internal_call() -> None: + reservation: Final = _budget_reservation() + token: Final = is_internal_call.set(True) + try: + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is False + + +def test_failure_handler_hands_the_budget_reservation_back_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + assert reservation["callback_bound"] is True + + pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True) + + assert reservation["callback_bound"] is False + assert pending is not None + pending.close() + + +def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation_claim_in_place() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + token: Final = is_internal_call.set(True) + try: + pending: Final = failure_handler(result.logger, RuntimeError("inner step failed"), now, now, asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is True + assert pending is not None + pending.close() + + CONTRACT_PATH: Final = ( Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1e59f4d878e..cf61a6d9f65 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4984,6 +4984,100 @@ async def test_wrapper_async_fires_post_call_failure_deployment_hook_on_internal assert isinstance(recorder.calls[0][1], litellm.AuthenticationError) +def _budget_reservation(callback_bound: bool = False) -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": callback_bound} + + +_BUDGET_RESERVATION_CALL_KWARGS: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} +_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_for_the_cost_callback() -> None: + reservation = _budget_reservation() + + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_before_the_stream_is_consumed() -> None: + reservation = _budget_reservation() + + stream = await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + stream=True, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + async for _ in stream: + pass + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_a_supplied_logging_object_already_saw() -> None: + reservation = _budget_reservation() + logging_obj, kwargs = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **_BUDGET_RESERVATION_CALL_KWARGS, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + assert reservation["callback_bound"] is False + + await litellm.acompletion(**kwargs, litellm_logging_obj=logging_obj, mock_response="ok") + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_hands_the_budget_reservation_back_when_the_call_fails() -> None: + reservation = _budget_reservation() + + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_wrapper_async_leaves_the_budget_reservation_alone_on_internal_calls() -> None: + claimed_by_the_outer_call = _budget_reservation(callback_bound=True) + never_claimed = _budget_reservation() + + token = is_internal_call.set(True) + try: + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": never_claimed}, + ) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": claimed_by_the_outer_call}, + ) + finally: + is_internal_call.reset(token) + + assert never_claimed["callback_bound"] is False + assert claimed_by_the_outer_call["callback_bound"] is True + + @pytest.mark.asyncio async def test_wrapper_async_does_not_fire_failure_hook_for_pre_call_budget_error( monkeypatch: pytest.MonkeyPatch,