diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6d2ce73624f..ca35be52fad 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -557,6 +557,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # real parent span. _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span + _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation _metadata.update( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) 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 9bddeda0723..6681558f8da 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 @@ -4877,3 +4877,119 @@ async def test_unusable_upstream_cost_records_zero_not_the_flat_estimate(): assert len(payloads) == 1 assert payloads[0]["response_cost"] == 0.0 assert payloads[0]["total_tokens"] == 1874 + + +def _passthrough_kwargs_for_reservation( + user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None +) -> dict: + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = ( + "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" + ) + mock_request.headers = Headers({}) + + return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=user_api_key_dict, + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body=parsed_body if parsed_body is not None else {}, + litellm_call_id="lit-5425-call-id", + ) + + +async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: + from datetime import datetime + + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + + callback_kwargs = { + **kwargs, + "stream": False, + "standard_logging_object": { + "response_cost": 0.002, + "request_tags": None, + }, + } + + increment_spend_counters = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + increment_spend_counters, + ), + patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await _ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=callback_kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + return increment_spend_counters + + +@pytest.mark.asyncio +async def test_passthrough_success_reconciles_budget_reservation(): + """ + A successful pass-through request must hand its pre-call budget reservation + to the spend-counter update so the reserved amount is reconciled down to the + actual cost. Without it the reservation stays in the shared Redis counter and + the actual cost is added on top, so the counter drifts above real spend until + the key falsely trips BudgetExceededError. + """ + budget_reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:hashed-token", "reserved_cost": 0.5}], + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + budget_reservation=budget_reservation, + ) + + reservation = user_api_key_dict.budget_reservation + kwargs = _passthrough_kwargs_for_reservation(user_api_key_dict) + assert ( + kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] + is reservation + ) + + increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) + + increment_spend_counters.assert_awaited_once() + assert increment_spend_counters.await_args.kwargs["budget_reservation"] is reservation + assert increment_spend_counters.await_args.kwargs["budget_reservation"] == budget_reservation + + +@pytest.mark.asyncio +async def test_passthrough_body_cannot_forge_budget_reservation(): + """ + The reservation is an internal counter handle: a client-supplied metadata + field naming arbitrary counter keys must never reach the spend-counter + update, or a caller could decrement another entity's Redis counter. + """ + forged = { + "reserved_cost": 99.0, + "entries": [{"counter_key": "spend:team:victim", "reserved_cost": 99.0}], + } + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-token", user_id="u1") + + kwargs = _passthrough_kwargs_for_reservation( + user_api_key_dict, + parsed_body={"litellm_metadata": {"user_api_key_budget_reservation": forged}}, + ) + assert ( + kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None + ) + + increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) + + increment_spend_counters.assert_awaited_once() + assert increment_spend_counters.await_args.kwargs["budget_reservation"] is None