From fb4b14358f8c5eccd21f34170a6724880dfeda0e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:26:46 -0700 Subject: [PATCH] fix(passthrough): bill websocket sessions the client walked out of A realtime client closes as soon as it has its answer, so the next frame the provider sends lands on a socket that is already gone. That raised a WebSocketDisconnect out of the upstream forwarding task, which the relay re-raised before it reached its cost block, so the session never turned into spend. Treat the client leaving as the end of the session and record the model on the logging object up front, so the spend row names the model that ran instead of unknown. --- ...i_websocket_passthrough_logging_handler.py | 11 +-- .../pass_through_endpoints.py | 6 ++ ...i_websocket_passthrough_logging_handler.py | 6 +- .../test_pass_through_endpoints.py | 68 ++++++++++++++++++- 4 files changed, 83 insertions(+), 8 deletions(-) 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,