diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1d180d9be51..07828b45eee 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -557,6 +557,18 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def _logged_response_id( + response: Mapping[str, object], litellm_call_id: object, falls_back_on_empty_id: bool +) -> str | None: + call_id: Final = litellm_call_id if isinstance(litellm_call_id, str) else None + if falls_back_on_empty_id: + return _provider_response_id(response) or call_id + if "id" not in response: + return call_id + response_id: Final = response["id"] + return None if response_id is None else str(response_id) + + def mask_api_base_credentials(api_base: str) -> str: if "key=" not in api_base: return api_base @@ -6688,7 +6700,12 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = _provider_response_id(response_obj) or kwargs.get("litellm_call_id") + litellm_call_id: Final[object] = kwargs.get("litellm_call_id") + id = _logged_response_id( + response_obj, + litellm_call_id, + falls_back_on_empty_id=isinstance(init_response_obj, LiteLLMRealtimeStreamLoggingObject), + ) _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 8b98f490a69..12e72f559d1 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_set_request_parsed_body, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.spend_tracking.budget_reservation import share_budget_reservation_across_calls from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1622,6 +1623,7 @@ async def _responses_websocket_session( await websocket.close(code=1008, reason="Pre-call error") return + share_budget_reservation_across_calls(user_api_key_dict.budget_reservation) routed_data: Final = dict( data, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8e8f4def110..cd0c4a8433d 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -378,18 +378,28 @@ async def reconcile_budget_reservation( return pending +_SETTLING_CALL_ID: Final = "settling_call_id" + + +def share_budget_reservation_across_calls(budget_reservation: dict[str, object] | None) -> None: + if budget_reservation is not None: + budget_reservation.setdefault(_SETTLING_CALL_ID, None) + + def claim_budget_reservation( budget_reservation: dict[str, object] | None, litellm_call_id: str | None ) -> dict[str, object] | None: - """Hand the reservation to the first call that settles it. + """Hand a shared reservation to the first call that settles it. - A native Responses WebSocket connection shares its one reservation with every turn it logs, each under its - own call id. Two turns settling it at once would both move the counter from the estimate to their own cost, - so the first claimant keeps the reservation and every later call sees none and charges its cost directly.""" - if budget_reservation is None or litellm_call_id is None: + A Responses WebSocket connection shares its one reservation with every turn it logs, each under its own + call id. Two turns settling it at once would both move the counter from the estimate to their own cost, + so the first claimant keeps the reservation and every later call sees none and charges its cost directly. + A reservation no connection shared belongs to its one call and is returned unchanged.""" + if budget_reservation is None or litellm_call_id is None or _SETTLING_CALL_ID not in budget_reservation: return budget_reservation - settling_call_id: Final = budget_reservation.setdefault("settling_call_id", litellm_call_id) - return budget_reservation if settling_call_id == litellm_call_id else None + if budget_reservation[_SETTLING_CALL_ID] is None: + budget_reservation[_SETTLING_CALL_ID] = litellm_call_id + return budget_reservation if budget_reservation[_SETTLING_CALL_ID] == litellm_call_id else None def stamp_budget_reservation_actual_cost(budget_reservation: dict | None, actual_cost: float | None) -> None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 21b6fd8ee5e..78e6c4bf0d5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1938,6 +1938,31 @@ def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> di return {**msg_obj, **top_fields, **restored_nested} +_RESPONSE_CREATE_MASKED_FIELDS: Final = ("input", "instructions") + + +def _with_masked_fields(client: Mapping[str, object], forwarded: Mapping[str, object]) -> Mapping[str, object]: + masked: Final = { + key: forwarded[key] for key in _RESPONSE_CREATE_MASKED_FIELDS if key in client and key in forwarded + } + return {**client, **masked} + + +def _client_frame_with_masked_text(client_frame: str, forwarded_frame: str) -> str: + client: Final = _json_object_or_none(client_frame) + forwarded: Final = _json_object_or_none(forwarded_frame) + if client is None or forwarded is None: + return client_frame + client_nested: Final = client.get("response") + forwarded_nested: Final = forwarded.get("response") + nested: Final = ( + {"response": _with_masked_fields(client_nested, forwarded_nested)} + if _is_json_object(client_nested) and _is_json_object(forwarded_nested) + else {} + ) + return json.dumps({**_with_masked_fields(client, forwarded), **nested}) + + def _wrap_output_item_encrypted_content( event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object] ) -> dict[str, object] | None: @@ -2015,6 +2040,11 @@ class ResponsesWebSocketStreaming: turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) + def _logged_client_frame(self, client_frame: str, forwarded_frame: str) -> str: + if not self.guardrail_callbacks or forwarded_frame is client_frame: + return client_frame + return _client_frame_with_masked_text(client_frame, forwarded_frame) + def _record_client_frame(self, client_frame: str) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": @@ -2504,7 +2534,7 @@ class ResponsesWebSocketStreaming: try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) - self._record_client_frame(masked_first) + self._record_client_frame(self._logged_client_frame(self.first_message, masked_first)) await self.backend_ws.send(masked_first) while True: @@ -2512,7 +2542,7 @@ class ResponsesWebSocketStreaming: if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) - self._record_client_frame(masked) + self._record_client_frame(self._logged_client_frame(message, masked)) await self.backend_ws.send(masked) except Exception as e: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index ccc4f00ed47..fb5f87c3016 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -38,6 +38,7 @@ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, _get_status_fields, + get_standard_logging_object_payload, set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -8263,6 +8264,28 @@ def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw assert payload["id"] == expected_request_id +@pytest.mark.parametrize( + ("response_obj", "expected_id"), + (({"id": ""}, ""), ({"id": "chatcmpl-1"}, "chatcmpl-1"), ({}, "http-call")), +) +def test_non_websocket_payload_id_is_the_response_id_field_as_returned( + logging_obj: LitellmLogging, response_obj: dict[str, str], expected_id: str +): + now = datetime.datetime.now() + + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "http-call", "model": "gpt-4o", "messages": []}, + init_response_obj=response_obj, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["id"] == expected_id + + def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): reservation = {"reserved_cost": 0.5, "finalized": False} connection = _responses_ws_logging_obj() diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 3ae4af4a4bc..89418a0f7a8 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2724,7 +2724,10 @@ async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservat budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, ) - from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + from litellm.proxy.spend_tracking.budget_reservation import ( + claim_budget_reservation, + share_budget_reservation_across_calls, + ) from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args handler = _RecordingHandler() @@ -2736,6 +2739,9 @@ async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservat ) logger = ProxyDBLogger(producer) connection_kwargs = _offload_kwargs() + share_budget_reservation_across_calls( + budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(connection_kwargs)) + ) for turn_call_id in ("turn-1", "turn-2", "turn-3"): await logger.async_log_success_event( @@ -3088,7 +3094,7 @@ async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_ """Every turn of a Responses WebSocket connection carries the connection's one reservation under its own call id. The first turn settles it; a later turn reconciling the same estimate again would under-count the key, so it must charge its cost without the reservation.""" - reservation = {"reserved_cost": 0.5, "entries": []} + reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None} first_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-1", reservation)) second_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-2", reservation)) @@ -3100,7 +3106,7 @@ async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_ @pytest.mark.asyncio async def test_async_post_call_failure_hook_releases_a_shared_reservation_for_the_first_turn_only(): logger = ProxyDBLogger() - budget_reservation = {"reserved_cost": 0.5, "entries": []} + budget_reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None} user_api_key_dict = UserAPIKeyAuth( api_key="test_api_key", request_route="/custom/route", diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 14040a7206d..eb31797d460 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7711,10 +7711,7 @@ async def test_passthrough_success_reconciles_budget_reservation(): 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, - "settling_call_id": "lit-5425-call-id", - } + assert increment_spend_counters.await_args.kwargs["budget_reservation"] == budget_reservation @pytest.mark.asyncio diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 302d2d28843..439d4026af0 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3,7 +3,7 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from collections.abc import Mapping +from collections.abc import Coroutine, Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -885,6 +885,65 @@ class TestResponsesWSFirstFrameModelAuth: assert turn_request["proxy_server_request"] == turn_proxy_server_request proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) + @pytest.mark.asyncio + async def test_endpoint_hands_the_connection_reservation_to_its_first_settling_turn_only(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock( + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"}) + ) + ws.close = AsyncMock() + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-token", budget_reservation={"reserved_cost": 0.5, "entries": []} + ) + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock() + + async def relay() -> None: + return None + + def route(**_: object) -> Coroutine[object, object, None]: + return relay() + + with ( + patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; the reservation it leaves is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing starts the relay, which needs a live upstream socket + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + side_effect=route, + ), + patch( # test-quality-ok: the disconnect release reads proxy_logging_obj off proxy_server + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_obj, + ), + ): + await responses_websocket_endpoint(websocket=ws, model=None, user_api_key_dict=user_api_key_dict) + + reservation = user_api_key_dict.budget_reservation + assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") is reservation + assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") is None + @pytest.mark.asyncio async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self): from litellm.proxy.response_api_endpoints.endpoints import ( diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 4591daaded6..d902a0e7a47 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -1,4 +1,5 @@ import asyncio +import copy import threading import time from collections.abc import Mapping, Sequence @@ -46,6 +47,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( invalidate_budget_reservation_counters, release_budget_reservation, release_budget_reservation_on_cancel, + share_budget_reservation_across_calls, reserve_budget_for_request, ) from litellm.proxy.spend_tracking.input_tokens import ( @@ -3723,6 +3725,19 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one ) == pytest.approx(4.2) +def test_a_reservation_no_connection_shared_is_left_as_its_call_made_it(): + reservation: dict[str, object] = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:k", "reserved_cost": 0.5}], + } + made_by_its_call = copy.deepcopy(reservation) + + claims = [claim_budget_reservation(budget_reservation=reservation, litellm_call_id=call) for call in ("a", "b")] + + assert claims == [reservation, reservation] + assert reservation == made_by_its_call + + @pytest.mark.asyncio async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at_once( spend_counter_state, @@ -3756,6 +3771,7 @@ async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at from litellm.proxy.proxy_server import increment_spend_counters + share_budget_reservation_across_calls(reservation) first_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") second_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") assert first_turn is reservation diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 39bcf6710aa..f4b1b8c86d1 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -3343,6 +3343,54 @@ class TestNativeWebSocketPerTurnLogging: for turn in recorder.turns: turn.dispatch_failure_handlers.assert_not_awaited() + @pytest.mark.parametrize( + ("guardrail_callbacks", "logged_text"), + [([], "mail alice@example.com"), ([_FakeWSGuardrail()], "mail ")], + ) + @pytest.mark.asyncio + async def test_each_turn_logs_the_frame_the_client_sent_not_the_frame_sent_upstream( + self, guardrail_callbacks: list[_FakeWSGuardrail], logged_text: str + ): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream_prev" + ) + client_frame = { + "type": "response.create", + "model": "client-model-group", + "input": "mail alice@example.com", + "previous_response_id": wrapped_previous, + } + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([json.dumps(client_frame)], recorder.logged_turns) + backend = _ScriptedBackend([_completed_turn("resp_1", 10, 1)]) + handler = _make_streaming( + websocket=client, + backend_ws=backend, + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini-deployment", + request_defaults=_deployment_defaults(), + guardrail_callbacks=guardrail_callbacks, + ) + + await handler.bidirectional_forward() + + assert json.loads(backend.sent[0]) == { + "type": "response.create", + "model": "gpt-5.4-mini-deployment", + "input": logged_text, + "previous_response_id": "resp_upstream_prev", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + assert [turn.request_body for turn in recorder.turns] == [{**client_frame, "input": logged_text}] + assert recorder.turns[0].model_call_details["messages"] == [{"role": "user", "content": logged_text}] + @pytest.mark.asyncio async def test_a_turn_whose_input_is_nested_under_response_logs_its_user_message(self): import websockets.exceptions # noqa: F401 (lazy submodule must be importable)