fix(passthrough): bill websocket sessions the client walked out of

A realtime client closes as soon as it has its answer, so the next frame
the provider sends lands on a socket that is already gone. That raised a
WebSocketDisconnect out of the upstream forwarding task, which the relay
re-raised before it reached its cost block, so the session never turned
into spend. Treat the client leaving as the end of the session and record
the model on the logging object up front, so the spend row names the model
that ran instead of unknown.

litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py

@@ -155,6 +155,11 @@ class OpenAIWebsocketPassthroughLoggingHandler:
             unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}}
             return unpriced
 
+        logging_obj.model = model
+        logging_obj.model_call_details.update(
+            MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER})
+        )
+
         billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing(
             websocket_messages, model
         )
@@ -166,11 +171,7 @@ class OpenAIWebsocketPassthroughLoggingHandler:
             }
             return unbilled
 
-        logging_obj.model_call_details.update(
-            MappingProxyType(
-                {"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, "response_cost": billing.cost}
-            )
-        )
+        logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost}))
         priced: Final[PassThroughEndpointLoggingTypedDict] = {
             "result": ModelResponse(
                 id=f"openai-websocket-{start_time.timestamp()}",

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -22,6 +22,7 @@ from fastapi import (
     Response,
     UploadFile,
     WebSocket,
+    WebSocketDisconnect,
     status,
 )
 from fastapi.responses import StreamingResponse
@@ -2513,6 +2514,11 @@ async def websocket_passthrough_request(
                 except (ConnectionClosedOK, ConnectionClosedError) as e:
                     verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
                     return e.rcvd
+                except WebSocketDisconnect:
+                    verbose_proxy_logger.debug(
+                        "WebSocket passthrough (%s): client left while upstream was still sending", endpoint
+                    )
+                    return None
                 except asyncio.CancelledError:
                     verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client")
                     raise

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py

@@ -167,10 +167,14 @@ def test_a_session_nobody_named_a_model_for_is_left_unpriced():
 
 
 def test_a_session_that_reported_no_usage_still_records_its_model():
-    handled = _handle([_session_created(REALTIME_MODEL)])
+    logging_obj = _logging_obj()
+
+    handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
 
     assert handled["result"] is None
     assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"}
+    assert logging_obj.model_call_details["model"] == REALTIME_MODEL
+    assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
 
 
 def test_a_responses_turn_without_readable_usage_is_not_billed():

tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 import logging
 import os
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
 from contextlib import ExitStack, contextmanager
 from io import BytesIO
 from types import MappingProxyType, SimpleNamespace
@@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
 
 import httpx
 import pytest
-from fastapi import Request, UploadFile
+from fastapi import Request, UploadFile, WebSocketDisconnect
 from starlette.datastructures import FormData, Headers, QueryParams
 from starlette.datastructures import UploadFile as StarletteUploadFile
 
@@ -5283,6 +5283,70 @@ async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame()
     assert "model" not in relayed.success_handler.call_args.kwargs
 
 
+class StreamingUpstreamWebSocket:
+    def __init__(self, setup_frame: str, streamed_frames: Sequence[str]):
+        self._setup_frame = setup_frame
+        self._streamed_frames = iter(streamed_frames)
+        self.close = AsyncMock()
+        self.send = AsyncMock()
+
+    async def recv(self, decode: bool = True) -> str:
+        return self._setup_frame
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> str:
+        streamed_frame: Final = next(self._streamed_frames, None)
+        if streamed_frame is None:
+            await asyncio.Event().wait()
+        return streamed_frame or ""
+
+
+REALTIME_SESSION_CREATED: Final = json.dumps(
+    {"type": "session.created", "session": {"id": "sess_lit7014", "model": "gpt-realtime-2.1"}}
+)
+REALTIME_RESPONSE_DONE: Final = json.dumps(
+    {
+        "type": "response.done",
+        "response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}},
+    }
+)
+
+
+@pytest.mark.asyncio
+async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of():
+    """
+    A realtime client closes the moment it has the answer it asked for, so the
+    next frame the provider sends lands on a socket that is already gone. That
+    is how a session ends, not a gateway failure, and the turns the caller
+    already used still have to reach spend.
+    """
+    upstream_ws = StreamingUpstreamWebSocket(
+        REALTIME_SESSION_CREATED,
+        [REALTIME_RESPONSE_DONE, json.dumps({"type": "response.output_text.delta", "delta": "too late"})],
+    )
+    websocket = _client_websocket(_pending_receive)
+    websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)])
+
+    with (
+        _patched_websocket_passthrough_environment(upstream_ws),
+        patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler,
+    ):
+        await websocket_passthrough_request(
+            websocket=websocket,
+            target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
+            custom_headers={"Authorization": "Bearer sk-provider"},
+            user_api_key_dict=UserAPIKeyAuth(),
+            forward_headers=False,
+            endpoint="/openai_passthrough/v1/realtime",
+            accept_websocket=False,
+        )
+
+    relayed_frames: Final = success_handler.call_args.kwargs["response_body"]
+    assert [frame["type"] for frame in relayed_frames] == ["session.created", "response.done"]
+
+
 def _passthrough_kwargs_for_reservation(
     user_api_key_dict: UserAPIKeyAuth,
     parsed_body: Optional[dict] = None,