feat(passthrough): let a WebSocket relay bill a session from its client text (#49281)

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(

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.

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