fix(realtime): relay the upstream close even when a client message hit the closed socket first

When the upstream closes while the proxy is forwarding a client message,
the client loop ends before the backend relay sees the close, and the
relay skipped closing the client because it read the client loop's exit
as the client hanging up. The client loop now reports why it stopped, so
a close observed on the backend send still reaches the client with the
error event and the upstream close code

litellm/litellm_core_utils/realtime_streaming.py

@@ -3,6 +3,7 @@ import json
 import traceback
 from collections.abc import Coroutine, Mapping, Sequence
 from dataclasses import dataclass
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast
 
 from typing_extensions import ReadOnly
@@ -46,6 +47,11 @@ class BackendClose:
         return f"upstream websocket closed with code {self.code}: {self.reason}"
 
 
+class ClientLoopExit(Enum):
+    CLIENT_DISCONNECTED = auto()
+    BACKEND_CLOSED = auto()
+
+
 def backend_close_from(error: "ConnectionClosed") -> BackendClose:
     if error.rcvd is None:
         return BackendClose(code=1006, reason=str(error))
@@ -1291,7 +1297,9 @@ class RealTimeStreaming:
         item["content"] = new_content
         return item
 
-    async def client_ack_messages(self):
+    async def client_ack_messages(self) -> ClientLoopExit:
+        import websockets
+
         client_event: _ClientEventFrame
         try:
             while True:
@@ -1529,16 +1537,21 @@ class RealTimeStreaming:
                 if guardrail_turn_detection_injected and sent:
                     self._guardrail_turn_detection_update_sent = True
 
+        except websockets.exceptions.ConnectionClosed as e:
+            verbose_logger.debug("Backend closed while forwarding a client message: %s", e)
+            return ClientLoopExit.BACKEND_CLOSED
         except Exception as e:
             verbose_logger.debug("Error in client ack messages: %s", e)
+            return ClientLoopExit.CLIENT_DISCONNECTED
 
     async def bidirectional_forward(self) -> None:
         forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages())
         client_task: Final = asyncio.create_task(self.client_ack_messages())
         try:
             await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED)
-            if not client_task.done():
-                await self._close_client(forward_task.result())
+            if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED:
+                return
+            await self._close_client(await forward_task)
         finally:
             forward_task.cancel()
             client_task.cancel()

tests/test_litellm/litellm_core_utils/test_realtime_streaming.py

@@ -3134,9 +3134,13 @@ class _InlineLoggingWorker:
 
 class _RecordingLogging:
     def __init__(self) -> None:
+        self.model_call_details: dict[str, object] = {}
         self.logged_sessions: tuple[tuple[dict, ...], ...] = ()
         self.logged_failures: tuple[Exception, ...] = ()
 
+    def pre_call(self, input: str | dict, api_key: str) -> None:
+        return None
+
     async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None:
         self.logged_sessions = (*self.logged_sessions, tuple(result))
 
@@ -3257,6 +3261,39 @@ async def test_upstream_close_after_relayed_events_still_logs_the_session_as_suc
     client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
 
 
+@pytest.mark.asyncio
+async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client():
+    upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
+    backend_closed: Final = asyncio.Event()
+    client_messages: Final = iter((json.dumps({"type": "response.create"}),))
+
+    async def receive_text() -> str:
+        message = next(client_messages, None)
+        return message if message is not None else await _wait_forever()
+
+    async def send_to_backend(_message: str) -> None:
+        backend_closed.set()
+        raise upstream_close
+
+    async def recv_from_backend() -> bytes:
+        await backend_closed.wait()
+        raise upstream_close
+
+    client_ws: Final = _client_ws_that_never_sends()
+    client_ws.receive_text = receive_text
+    backend_ws: Final = MagicMock()
+    backend_ws.send = send_to_backend
+    backend_ws.recv = recv_from_backend
+    session: Final = _relay_session(client_ws, backend_ws)
+
+    await session.run()
+
+    (error_event,) = _error_events_sent_to(client_ws)
+    assert _UPSTREAM_REFUSAL in error_event["error"]["message"]
+    client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
+    assert session.logging.logged_failures == (upstream_close,)
+
+
 @pytest.mark.asyncio
 async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close():
     client_ws: Final = _client_ws_that_never_sends()