diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index f2b23d599f9..71c35e62ac0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -13,7 +13,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind from litellm.types.utils import LlmProviders, ModelResponse, Usage @@ -43,18 +43,11 @@ class _ModelNamingEvent(BaseModel): response: _NamedModel | None = None -class _EndedResponse(BaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - model: str | None = None - usage: ResponseAPIUsage | None = None - - class _ResponsesTurnEndEvent(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) type: Literal["response.completed", "response.incomplete", "response.failed"] - response: _EndedResponse + response: ResponsesAPIResponse _MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) @@ -70,6 +63,7 @@ class _Billing: @dataclass(frozen=True, slots=True) class _CompletedTurn: model: str + response: ResponsesAPIResponse usage: Usage @@ -136,6 +130,7 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet return None return _CompletedTurn( model=ended.response.model or session_model, + response=ended.response, usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter ended.response.usage ), @@ -149,23 +144,24 @@ def _responses_billing(messages: WebsocketMessages, model: str, requested_model: usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) return _Billing( usage=usage, - cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns), + cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns), ) -def _model_cost(usage: Usage, model: str) -> float | None: +def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None: try: return litellm.completion_cost( - completion_response=ModelResponse(model=model, usage=usage), + completion_response=response, model=model, custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + call_type="responses", ) except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log return None -def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float: - costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None) +def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None) cost: Final = next((cost for cost in costs if cost), None) if cost is None: verbose_proxy_logger.warning( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 393b193c73f..5b2789eded9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2743,6 +2743,14 @@ def _session_end_needs_a_row( return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"})) +def _turn_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]] +) -> bool: + return upstream_frame_kind is not None and any( + upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames + ) + + async def _relay_client_frame( websocket: WebSocket, upstream_ws: ClientConnection, @@ -2765,6 +2773,7 @@ async def _relay_client_frame( _EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) +WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0 async def websocket_passthrough_request( @@ -2779,6 +2788,7 @@ async def websocket_passthrough_request( setup_model_rewriter: Callable[[str], str] | None = None, client_frame_model_gate: WebsocketFrameModelGate | None = None, upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, + walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS, ): """ WebSocket passthrough request handler. @@ -2795,6 +2805,8 @@ async def websocket_passthrough_request( client_frame_model_gate: Optional per-frame check of every model a client frame names upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones that open the session, so each turn logs its own spend and keeps the session's provider-named model + walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client + left, so the provider's usage for it still reaches spend """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -3019,7 +3031,8 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - await upstream_ws.close() + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() break text_data: str | None = message.get("text") @@ -3147,11 +3160,17 @@ async def websocket_passthrough_request( async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) if not _client_socket_is_open(websocket): + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() return - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + try: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + except WebSocketDisconnect: + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + raise async def forward_upstream_to_client() -> Close | None: try: @@ -3176,12 +3195,15 @@ async def websocket_passthrough_request( raise # Create tasks for bidirectional message forwarding - tasks: Final = [ - asyncio.create_task(forward_client_to_upstream()), - asyncio.create_task(forward_upstream_to_client()), - ] + client_task: Final = asyncio.create_task(forward_client_to_upstream()) + upstream_task: Final = asyncio.create_task(forward_upstream_to_client()) + tasks: Final = (client_task, upstream_task) - done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if client_task.done() and _turn_in_flight(upstream_frame_kind, websocket_messages): + await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds) + done: Final = tuple(task for task in tasks if task.done()) + pending: Final = tuple(task for task in tasks if not task.done()) # Cancel remaining tasks for task in pending: diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index 7494c913a4d..3390b426f02 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -4,9 +4,13 @@ from types import MappingProxyType from typing import Final from unittest.mock import MagicMock +import httpx import pytest import litellm +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( OpenAIWebsocketPassthroughLoggingHandler, openai_websocket_frame_kind, @@ -45,14 +49,22 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou def _responses_completed( - model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" + model: str, + input_tokens: int, + output_tokens: int, + event: str = "response.completed", + served: Mapping[str, object] = MappingProxyType({}), ) -> dict[str, object]: return { "type": event, "response": { "id": "resp_fake", + "object": "response", + "created_at": 1761700000, "model": model, "status": "completed", + "output": [], + **served, "usage": { "input_tokens": input_tokens, "output_tokens": output_tokens, @@ -201,10 +213,7 @@ def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): handled = _handle( [ _session_created(RESPONSES_MODEL), - { - "type": "response.completed", - "response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}}, - }, + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}), ], url_route="/openai_passthrough/v1/responses", ) @@ -357,6 +366,39 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha assert handled["kwargs"]["model"] == turn_model +def _http_passthrough_cost(response: Mapping[str, object]) -> float: + http_response = httpx.Response( + 200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses") + ) + _, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( + RESPONSES_MODEL, http_response, _logging_obj(), "openai" + ) + return cost + + +@pytest.mark.parametrize( + "served", + [ + {"service_tier": "priority"}, + {"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]}, + ], + ids=["priority_tier", "hosted_web_search"], +) +def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served): + """ + The provider bills a turn by the tier that served it and the hosted tools it + called, not only its tokens, so a websocket turn has to cost exactly what the + same response costs over HTTP, never the cheaper token-only price. + """ + turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served) + plain_turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5) + + handled = _handle([turn], url_route="/openai_passthrough/v1/responses") + + assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"])) + assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"]) + + @pytest.mark.parametrize( ("frame_type", "kind"), [ 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 c4b08685d74..bae193d0da2 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 @@ -8103,6 +8103,7 @@ async def _relay_turns( endpoint: str = "/openai_passthrough/v1/responses", target: str = "wss://api.openai.com/v1/responses", opening_access_groups: Sequence[str] = (), + walkout_turn_drain_seconds: float = 0.05, ) -> SimpleNamespace: with ( _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, @@ -8121,6 +8122,7 @@ async def _relay_turns( accept_websocket=False, client_frame_model_gate=gate, upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=walkout_turn_drain_seconds, ), timeout=10, ) @@ -8297,6 +8299,97 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou ] == [["response.created", "response.completed"], ["response.created"]] +class UpstreamFinishingTheTurnAfterTheClientLeft: + """Starts a turn for each client frame and sends the rest of it only after the client left, and stops once closed""" + + def __init__(self, client_left: asyncio.Event, started: str, rest_of_the_turn: Sequence[str]): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + self._client_left = client_left + self._started = started + self._rest_of_the_turn = rest_of_the_turn + self._closed = False + self._finishing: list[asyncio.Task[None]] = [] + self.send = AsyncMock(side_effect=self._start_turn) + self.close = AsyncMock(side_effect=self._close) + + def _close(self, *args: object, **kwargs: object) -> None: + self._closed = True + self._frames.put_nowait(None) + + async def _start_turn(self, frame: str | bytes) -> None: + self._frames.put_nowait(self._started) + self._finishing.append(asyncio.create_task(self._finish_once_the_client_left())) + + async def _finish_once_the_client_left(self) -> None: + await self._client_left.wait() + for frame in self._rest_of_the_turn: + await asyncio.sleep(0) + if self._closed: + return + self._frames.put_nowait(frame) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_the_provider_ends_it( + hang_up_read_before_the_next_send: bool, +): + """ + Leaving mid-answer must not make a turn free. The provider keeps running and + bills it, so the gateway keeps reading until the turn ends and logs its usage, + whether it reads the hang-up first or finds it by failing to relay a frame. + """ + from starlette.websockets import WebSocketDisconnect, WebSocketState + + client_left = asyncio.Event() + send_failed = asyncio.Event() + client = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"response.created"' + ) + receive = client.receive + send_text = client.send_text + + async def receive_until_leaving() -> Mapping[str, object]: + message = await receive() + if message["type"] == "websocket.disconnect": + client_left.set() + if not hang_up_read_before_the_next_send: + await send_failed.wait() + client.client_state = WebSocketState.DISCONNECTED + return message + + async def send_text_until_leaving(text: str) -> None: + if client_left.is_set(): + send_failed.set() + raise WebSocketDisconnect(code=1006) + await send_text(text) + + client.receive = receive_until_leaving + client.send_text = AsyncMock(side_effect=send_text_until_leaving) + upstream_ws = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + _responses_turn_started("gpt-5.4"), + (json.dumps({"type": "response.output_text.delta", "delta": "Hi"}), _responses_turn_ended("gpt-5.4")), + ) + + relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.output_text.delta", "response.completed"]] + upstream_ws.close.assert_awaited() + + @pytest.mark.asyncio async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn(): """