From 363cc0da042fb0710ca83a4b21534c19ccb69b21 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 26 Oct 2026 13:38:29 -0700 Subject: [PATCH] fix(passthrough): bill a WebSocket session whose client hangs up mid-stream When a client closed its socket while the upstream was still sending, the relay's next send to the client raised, the session took the unexpected error path, the client got a 1011 close and the proxy wrote a failure row with no cost. The upstream had already done, and billed, the work A send to a client that has gone now ends the session the way a client disconnect does: the upstream is closed, the frames received so far reach the success handler and spend is logged once. Each upstream frame is recorded before it is sent, so the frame whose send failed is billed too. For Deepgram listen that can be the closing Metadata frame, which carries the session's duration. Starlette raises WebSocketDisconnect when the socket fails and uvicorn raises RuntimeError once the transport is closed, so both count as the client leaving Every WebSocket passthrough (Deepgram listen, Vertex Live, realtime) shares this relay. Text-to-speech makes it common, since callers hang up or barge in while audio is still streaming Refs LIT-10070 --- .../pass_through_endpoints.py | 34 ++++-- .../test_pass_through_endpoints.py | 114 +++++++++++++++++- 2 files changed, 140 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fc0e5be3f4c..f8953c44ba2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -29,7 +29,7 @@ from fastapi.responses import StreamingResponse from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route -from starlette.websockets import WebSocketState +from starlette.websockets import WebSocketDisconnect, WebSocketState from websockets.asyncio.client import connect from websockets.exceptions import ( ConnectionClosedError, @@ -2429,6 +2429,22 @@ def _client_socket_is_open(websocket: WebSocket) -> bool: ) +async def _send_to_client(websocket: WebSocket, frame: str | bytes) -> bool: + """ + False once the client has gone, so the relay ends the session like a client disconnect and still bills what the + upstream produced. Starlette raises WebSocketDisconnect when the transport errors. RuntimeError comes from uvicorn + after the client's close has arrived, or from the event loop once the transport is closed + """ + try: + if isinstance(frame, bytes): + await websocket.send_bytes(frame) + else: + await websocket.send_text(frame) + except (WebSocketDisconnect, RuntimeError): + return False + return True + + def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: """ The upstream close worth telling the client about: anything other than a plain, reasonless normal close. @@ -2705,11 +2721,7 @@ async def websocket_passthrough_request( is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) json_frame_ordinal: Final = count() - async def relay_upstream_frame(upstream_message: str | bytes) -> None: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + def record_upstream_frame(upstream_message: str | bytes) -> None: message_data: Final = _json_object_frame(upstream_message) if message_data is None: return @@ -2718,10 +2730,18 @@ async def websocket_passthrough_request( return websocket_messages.append(message_data) + async def relay_upstream_frame(upstream_message: str | bytes) -> bool: + record_upstream_frame(upstream_message) + return await _send_to_client(websocket, upstream_message) + async def forward_upstream_to_client() -> Close | None: try: while True: - await relay_upstream_frame(await upstream_ws.recv()) + if not await relay_upstream_frame(await upstream_ws.recv()): + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while the upstream was sending", endpoint + ) + return None except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd 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 db0a74d360a..466c3be5114 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 @@ -7,7 +7,7 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Coroutine, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime @@ -27,6 +27,7 @@ from fastapi.testclient import TestClient from pydantic import TypeAdapter, ValidationError from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile +from starlette.websockets import WebSocketDisconnect import litellm from litellm._logging import verbose_proxy_logger @@ -69,6 +70,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) +from litellm.types.utils import StandardLoggingPayload from tests._master_key import MASTER_KEY from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -7617,6 +7619,116 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text upstream_ws.close.assert_awaited_once() +RELAY_HANG_GUARD_SECONDS: Final = 5.0 + + +def _final_results_frame(start: float, duration: float, transcript: str) -> str: + return json.dumps( + { + "type": "Results", + "start": start, + "duration": duration, + "is_final": True, + "channel": {"alternatives": [{"transcript": transcript}]}, + } + ) + + +class _CapturingSpendLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.payloads.append(kwargs["standard_logging_object"]) + + +@dataclass(frozen=True, slots=True) +class _RelayedSession: + spend_logs: tuple[StandardLoggingPayload, ...] + failure_hook: AsyncMock + + +async def _relay_and_capture_spend( + monkeypatch: pytest.MonkeyPatch, + upstream_ws: FakeUpstreamWebSocket, + websocket: MagicMock, + target: str, + endpoint: str, +) -> _RelayedSession: + """Runs the relay with the real success handler and returns what a spend logger receives""" + capturing_logger: Final = _CapturingSpendLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + queued: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: the fake worker keeps each coroutine + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + return_value=FakeUpstreamConnect(upstream_ws), + ), + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + ): + proxy_logging.pre_call_hook = AsyncMock(return_value={}) + proxy_logging.post_call_success_hook = AsyncMock() + proxy_logging.post_call_failure_hook = AsyncMock() + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: queued.append(async_coroutine) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=websocket, + target=target, + custom_headers={"Authorization": "Token dg-provider-key"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + ), + timeout=RELAY_HANG_GUARD_SECONDS, + ) + for coroutine in queued: + await coroutine + return _RelayedSession( + spend_logs=tuple(capturing_logger.payloads), failure_hook=proxy_logging.post_call_failure_hook + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "client_send_error", + [ + pytest.param(RuntimeError("unable to perform operation on "), id="transport closed"), + pytest.param(WebSocketDisconnect(code=1006), id="socket failed"), + ], +) +async def test_websocket_passthrough_bills_a_session_whose_client_hangs_up_while_upstream_is_sending( + monkeypatch: pytest.MonkeyPatch, client_send_error: Exception +) -> None: + """A client that hangs up mid-stream has still used what the upstream produced, so the session is billed like a + normal disconnect, up to the end of the frame whose send failed and not past it""" + upstream_ws: Final = FakeUpstreamWebSocket( + _final_results_frame(start=0.0, duration=1.0, transcript="relayed"), + b"\x00\x01audio", + _final_results_frame(start=1.0, duration=1.5, transcript="sent as the client hung up"), + _final_results_frame(start=2.5, duration=5.0, transcript="never relayed"), + ) + websocket: Final = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, client_send_error]) + + session: Final = await _relay_and_capture_spend( + monkeypatch, upstream_ws, websocket, DEEPGRAM_LISTEN_TARGET, "/deepgram/v1/listen" + ) + + nova_3_usd_per_second: Final = litellm.model_cost["deepgram/streaming/nova-3"]["input_cost_per_second"] + assert [(log["status"], log["model"]) for log in session.spend_logs] == [("success", "nova-3")] + assert session.spend_logs[0]["response_cost"] == pytest.approx(2.5 * nova_3_usd_per_second) + session.failure_hook.assert_not_awaited() + assert all(call.kwargs.get("code") != 1011 for call in websocket.close.await_args_list) + + @pytest.mark.asyncio async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage(): """Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the