scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

fix(passthrough): bill a WebSocket session whose client hangs up mid-stream

Browse source 
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
This commit is contained in:
shrey kharbanda 2026-10-26 13:38:29 -07:00
parent c905cd065d
commit 363cc0da04
2 changed files with 140 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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

114
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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 <TCPTransport closed=True>"), 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