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 2aae72f5e1d..eec6ae102d2 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 @@ -155,6 +155,11 @@ class OpenAIWebsocketPassthroughLoggingHandler: unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} return unpriced + logging_obj.model = model + logging_obj.model_call_details.update( + MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}) + ) + billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( websocket_messages, model ) @@ -166,11 +171,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: } return unbilled - logging_obj.model_call_details.update( - MappingProxyType( - {"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, "response_cost": billing.cost} - ) - ) + logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) priced: Final[PassThroughEndpointLoggingTypedDict] = { "result": ModelResponse( id=f"openai-websocket-{start_time.timestamp()}", diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7a79ac6cdf3..30ef0c535a9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -22,6 +22,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -2513,6 +2514,11 @@ async def websocket_passthrough_request( except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd + except WebSocketDisconnect: + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while upstream was still sending", endpoint + ) + return None except asyncio.CancelledError: verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index 3a8e1f7e8fa..81f8436335d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -167,10 +167,14 @@ def test_a_session_nobody_named_a_model_for_is_left_unpriced(): def test_a_session_that_reported_no_usage_still_records_its_model(): - handled = _handle([_session_created(REALTIME_MODEL)]) + logging_obj = _logging_obj() + + handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) assert handled["result"] is None assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" def test_a_responses_turn_without_readable_usage_is_not_billed(): 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 e4d9c6f17cc..ae139d2a70b 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 @@ -2,7 +2,7 @@ import asyncio import json import logging import os -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from io import BytesIO from types import MappingProxyType, SimpleNamespace @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest -from fastapi import Request, UploadFile +from fastapi import Request, UploadFile, WebSocketDisconnect from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile @@ -5283,6 +5283,70 @@ async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame() assert "model" not in relayed.success_handler.call_args.kwargs +class StreamingUpstreamWebSocket: + def __init__(self, setup_frame: str, streamed_frames: Sequence[str]): + self._setup_frame = setup_frame + self._streamed_frames = iter(streamed_frames) + self.close = AsyncMock() + self.send = AsyncMock() + + async def recv(self, decode: bool = True) -> str: + return self._setup_frame + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + streamed_frame: Final = next(self._streamed_frames, None) + if streamed_frame is None: + await asyncio.Event().wait() + return streamed_frame or "" + + +REALTIME_SESSION_CREATED: Final = json.dumps( + {"type": "session.created", "session": {"id": "sess_lit7014", "model": "gpt-realtime-2.1"}} +) +REALTIME_RESPONSE_DONE: Final = json.dumps( + { + "type": "response.done", + "response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}}, + } +) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of(): + """ + A realtime client closes the moment it has the answer it asked for, so the + next frame the provider sends lands on a socket that is already gone. That + is how a session ends, not a gateway failure, and the turns the caller + already used still have to reach spend. + """ + upstream_ws = StreamingUpstreamWebSocket( + REALTIME_SESSION_CREATED, + [REALTIME_RESPONSE_DONE, json.dumps({"type": "response.output_text.delta", "delta": "too late"})], + ) + websocket = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == ["session.created", "response.done"] + + def _passthrough_kwargs_for_reservation( user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None,