From 6e6dd3f5bcb1869131e6ba7a46599f7a46e594ed Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 26 Oct 2026 13:41:38 -0700 Subject: [PATCH] feat(passthrough): let a WebSocket route count client text for billing websocket_passthrough_request takes an optional count_client_text callback. It is applied to each client frame, text or binary, after the frame reaches the upstream, and only the summed integer reaches the success handler, under WEBSOCKET_CLIENT_TEXT_COUNT_KEY in request_body. The frames are never kept, so a route can bill on client text it must not store. Binary frames are counted too, since an upstream may parse them the same way as text frames Frames are relayed unchanged in both directions. Callers that pass no counter keep request_body={}. custom_headers is typed as the read-only Mapping[str, str] the relay actually uses Refs LIT-10070 --- .../pass_through_endpoints.py | 17 +- .../pass_through_endpoints.py | 2 + .../test_pass_through_endpoints.py | 146 ++++++++++++++++++ 3 files changed, 163 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f8953c44ba2..22a44de6bd9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -137,6 +137,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + WEBSOCKET_CLIENT_TEXT_COUNT_KEY, EndpointType, PassthroughStandardLoggingPayload, ) @@ -2476,13 +2477,14 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict async def websocket_passthrough_request( websocket: WebSocket, target: str, - custom_headers: dict, + custom_headers: Mapping[str, str], user_api_key_dict: UserAPIKeyAuth, forward_headers: bool | None = False, endpoint: str | None = None, cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + count_client_text: Callable[[str | bytes], int] | None = None, ): """ WebSocket passthrough request handler. @@ -2496,6 +2498,9 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + count_client_text: Optional read-only count of each client frame, text or binary, after it reaches the + upstream, such as billable characters. Only the sum reaches the success handler, under WEBSOCKET_CLIENT_TEXT_COUNT_KEY + in request_body, so routes whose client text must not be stored can still bill it """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2506,6 +2511,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() websocket_messages: Final[list[dict[str, object]]] = [] + client_text_count: int = 0 # rebind-ok: forward_client_to_upstream adds each client frame's count litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -2624,6 +2630,7 @@ async def websocket_passthrough_request( async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" + nonlocal client_text_count try: while True: message = await websocket.receive() @@ -2693,8 +2700,12 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + if count_client_text is not None: + client_text_count += count_client_text(text_data) elif bytes_data is not None: await upstream_ws.send(bytes_data) + if count_client_text is not None: + client_text_count += count_client_text(bytes_data) except asyncio.CancelledError: raise except Exception: @@ -2835,7 +2846,9 @@ async def websocket_passthrough_request( end_time=end_time, logging_obj=logging_obj, cache_hit=False, - request_body={}, + request_body={} + if count_client_text is None + else {WEBSOCKET_CLIENT_TEXT_COUNT_KEY: client_text_count}, **success_kwargs, ) ) diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..40a78d79194 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -27,6 +27,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" +WEBSOCKET_CLIENT_TEXT_COUNT_KEY: Final = "litellm_websocket_client_text_count" + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" 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 466c3be5114..8bdfe7648f4 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 @@ -67,6 +67,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + WEBSOCKET_CLIENT_TEXT_COUNT_KEY, EndpointType, PassthroughStandardLoggingPayload, ) @@ -7822,6 +7823,151 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +DEEPGRAM_SPEAK_TARGET = "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en" +SPEAK_AUDIO_CHUNK = bytes(range(256)) * 8 +SPOKEN_TEXT = "The quick brown fox jumps over the lazy dog." +SPEAK_FRAME = json.dumps({"type": "Speak", "text": SPOKEN_TEXT}) +FLUSH_FRAME = json.dumps({"type": "Flush"}) +FLUSHED_FRAME = json.dumps({"type": "Flushed", "sequence_id": 0}) +SPEAK_CLOSE_FRAME = json.dumps({"type": "Close"}) +CLIENT_BINARY_FRAME = b"\x00\xffclient-binary" +SPEAK_SESSION_FRAMES = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME}) +SPEAK_SESSION_CLIENT_LENGTH = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME))) +RELAY_HANG_GUARD_SECONDS = 5.0 + + +class SpeakingUpstreamWebSocket: + """Answers each Speak with audio as soon as it arrives and ends the session on Close, the way Deepgram speak + does, so a relay that held audio back until the end would never see the Close it waits for""" + + def __init__(self, close_code: int = 1000, close_reason: str = ""): + self.received: list[str | bytes] = [] + self.close = AsyncMock() + self._close_code = close_code + self._close_reason = close_reason + self._replies: asyncio.Queue[str | bytes | None] = asyncio.Queue() + + async def send(self, message: str | bytes) -> None: + self.received.append(message) + match json.loads(message).get("type") if isinstance(message, str) else None: + case "Speak": + self._replies.put_nowait(SPEAK_AUDIO_CHUNK) + case "Flush": + self._replies.put_nowait(FLUSHED_FRAME) + case "Close": + self._replies.put_nowait(None) + + async def recv(self, decode: bool | None = None) -> str | bytes: + from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + from websockets.frames import Close + + reply = await self._replies.get() + if reply is not None: + return reply + close = Close(self._close_code, self._close_reason) + closed = ConnectionClosedOK if self._close_code == 1000 else ConnectionClosedError + raise closed(rcvd=close, sent=close, rcvd_then_sent=True) + + +def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRAMES): + """A client socket that sends ``frames``, then Close once the first audio has reached it""" + first_audio = asyncio.Event() + + async def script(): + for frame in frames: + yield {"type": "websocket.receive", **frame} + await first_audio.wait() + yield {"type": "websocket.receive", "text": SPEAK_CLOSE_FRAME} + await asyncio.Event().wait() + + steps = script() + + async def receive(): + return await anext(steps) + + websocket = _client_websocket(receive) + websocket.send_bytes = AsyncMock(side_effect=lambda data: first_audio.set()) + return websocket + + +async def _relay_deepgram_speak(upstream_ws, websocket, count_client_text): + """Runs the generic relay the way the Deepgram speak route does and returns the success handler mock""" + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=websocket, + target=DEEPGRAM_SPEAK_TARGET, + custom_headers={"Authorization": "Token dg-provider-key"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/deepgram/v1/speak", + accept_websocket=False, + count_client_text=count_client_text, + ), + timeout=RELAY_HANG_GUARD_SECONDS, + ) + return success_handler + + +@pytest.mark.asyncio +async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed(): + """Audio reaches the client while the session is open, every client frame reaches the upstream unchanged, and + the success handler gets only the summed count of the text and binary frames, never the frames""" + upstream_ws = SpeakingUpstreamWebSocket() + websocket = _speak_client() + + success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + + assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] + assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [SPEAK_AUDIO_CHUNK] + assert [call.args[0] for call in websocket.send_text.await_args_list] == [FLUSHED_FRAME] + assert success_handler.call_args.kwargs["request_body"] == { + WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH + } + websocket.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body(): + """Listen, realtime and every other caller pass no counter, so their success handler sees what it saw before""" + upstream_ws = SpeakingUpstreamWebSocket() + + success_handler = await _relay_deepgram_speak(upstream_ws, _speak_client(), count_client_text=None) + + assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] + assert success_handler.call_args.kwargs["request_body"] == {} + assert success_handler.call_args.kwargs["response_body"] == [json.loads(FLUSHED_FRAME)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("close_code", "close_reason"), + [ + pytest.param(1008, "DATA-0000: invalid text", id="policy"), + pytest.param(1011, "NET-0001: internal", id="internal"), + ], +) +async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_with_an_error_close( + close_code: int, close_reason: str +): + upstream_ws = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason) + websocket = _speak_client() + + success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + + websocket.close.assert_awaited_once_with(code=close_code, reason=close_reason) + assert success_handler.call_args.kwargs["request_body"] == { + WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH + } + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """