From 68717ecebf436677531cb8e0be6fb516d82cc515 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 30 Oct 2026 10:49:09 -0700 Subject: [PATCH] feat(passthrough): let a WebSocket relay bill a session from its client text websocket_passthrough_request takes an optional meter. The relay measures each client text frame once it reaches the upstream, and when the session ends it logs the model, provider and cost the meter prices for the total. A route can then bill an upstream that reports no usage of its own, without the relay knowing what is counted or the client's text being kept Routes that pass no meter behave exactly as before --- .../pass_through_endpoints.py | 17 ++++ .../pass_through_endpoints.py | 16 ++- .../test_pass_through_endpoints.py | 97 +++++++++++++++++++ 3 files changed, 129 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cfac28db1a9..647e145aac1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -148,6 +148,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebSocketSessionMeter, + WebSocketSessionUsage, ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage @@ -2595,6 +2597,13 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None: + logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object + logging_obj.model_call_details["custom_llm_provider"] = usage.custom_llm_provider # rebind-ok: same + logging_obj.model_call_details["response_cost"] = usage.response_cost # rebind-ok: same + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2607,6 +2616,7 @@ async def websocket_passthrough_request( setup_model_rewriter: Callable[[str], str] | None = None, builtin_provider_route: bool = False, model: str | None = None, + meter: WebSocketSessionMeter | None = None, ): """ WebSocket passthrough request handler. @@ -2622,6 +2632,7 @@ async def websocket_passthrough_request( setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives + meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2632,6 +2643,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() websocket_messages: Final[list[dict[str, object]]] = [] + metered_total = 0 # rebind-ok: forward_client_to_upstream adds the measure of each client text frame it delivers litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -2750,6 +2762,7 @@ async def websocket_passthrough_request( async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" + nonlocal metered_total try: while True: message = await websocket.receive() @@ -2819,6 +2832,8 @@ 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 meter is not None: + metered_total += meter.measure(text_data) elif bytes_data is not None: await upstream_ws.send(bytes_data) except asyncio.CancelledError: @@ -2945,6 +2960,8 @@ async def websocket_passthrough_request( self.url = target_url mock_response: Final = MockWebSocketResponse(target) + if meter is not None: + _log_session_usage(logging_obj, meter.price(metered_total)) # Use the same success handler as HTTP passthrough endpoints GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..2f269dc87d9 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,5 +1,6 @@ +from dataclasses import dataclass from enum import Enum -from typing import Final +from typing import Final, Protocol from pydantic import Field from typing_extensions import TypedDict @@ -75,6 +76,19 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ +@dataclass(frozen=True, slots=True) +class WebSocketSessionUsage: + model: str + custom_llm_provider: str + response_cost: float | None + + +class WebSocketSessionMeter(Protocol): + def measure(self, client_text: str) -> int: ... + + def price(self, total: int) -> WebSocketSessionUsage: ... + + class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ Settings for a specific guardrail on a passthrough endpoint. 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 a9098ac3d59..b7e76ce0068 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 @@ -77,6 +77,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebSocketSessionUsage, ) from tests._master_key import MASTER_KEY from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -7736,6 +7737,102 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text upstream_ws.close.assert_awaited_once() +@dataclass(frozen=True, slots=True) +class _LengthMeter: + rate: float + + def measure(self, client_text: str) -> int: + return len(client_text) + + def price(self, total: int) -> WebSocketSessionUsage: + return WebSocketSessionUsage(model="voice-1", custom_llm_provider="tts-provider", response_cost=total * self.rate) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_meters_price_for_the_client_text_it_delivered(): + """Text frames reach the upstream unchanged and are measured, binary frames are not, and the session is logged at + the model, provider and cost the meter prices for the total.""" + frames = ("hello", b"\x00\x01binary", "world!") + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock( + side_effect=[ + *( + {"type": "websocket.receive", "bytes": frame} + if isinstance(frame, bytes) + else {"type": "websocket.receive", "text": frame} + for frame in frames + ), + {"type": "websocket.disconnect"}, + ] + ) + ) + + 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 websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example/v1/stream", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/provider/v1/stream", + accept_websocket=False, + meter=_LengthMeter(rate=0.5), + ) + + assert [call.args[0] for call in upstream_ws.send.await_args_list] == list(frames) + logging_obj = success_handler.call_args.kwargs["logging_obj"] + assert logging_obj.model == "voice-1" + assert logging_obj.model_call_details["model"] == "voice-1" + assert logging_obj.model_call_details["custom_llm_provider"] == "tts-provider" + assert logging_obj.model_call_details["response_cost"] == (len("hello") + len("world!")) * 0.5 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_meters_only_client_text_the_upstream_accepted(): + upstream_ws = RecordingUpstreamWebSocket() + upstream_ws.send = AsyncMock(side_effect=[None, ConnectionError("upstream went away")]) + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": "hello"}, + {"type": "websocket.receive", "text": "never delivered"}, + {"type": "websocket.disconnect"}, + ] + ) + ) + + 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 websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example/v1/stream", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/provider/v1/stream", + accept_websocket=False, + meter=_LengthMeter(rate=0.5), + ) + + logging_obj = success_handler.call_args.kwargs["logging_obj"] + assert [call.args[0] for call in upstream_ws.send.await_args_list] == ["hello", "never delivered"] + assert logging_obj.model_call_details["response_cost"] == len("hello") * 0.5 + + @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