fix(websocket): preserve queued binary frames during reconnect

src/openai/_exceptions.py

@@ -198,13 +198,32 @@ class InvalidWebhookSignatureError(ValueError):
 
 
 class WebSocketConnectionClosedError(OpenAIError):
-    """Raised when a WebSocket connection closes with unsent messages."""
+    """Raised when a WebSocket connection closes with unsent messages.
+
+    `unsent_frames` preserves the text and binary frames in their original order.
+    `unsent_messages` retains the text messages, including UTF-8 binary frames
+    represented as text for compatibility.
+    """
 
     unsent_messages: list[str]
+    unsent_frames: list[str] | list[str | bytes]
 
-    def __init__(self, message: str, *, unsent_messages: list[str]) -> None:
+    def __init__(self, message: str, *, unsent_messages: list[str] | list[str | bytes]) -> None:
         super().__init__(message)
-        self.unsent_messages = unsent_messages
+        self.unsent_frames = unsent_messages
+        if all(isinstance(data, str) for data in unsent_messages):
+            self.unsent_messages = cast("list[str]", unsent_messages)
+        else:
+            self.unsent_messages = []
+            for data in unsent_messages:
+                if isinstance(data, str):
+                    self.unsent_messages.append(data)
+                else:
+                    try:
+                        self.unsent_messages.append(data.decode("utf-8"))
+                    except UnicodeDecodeError:
+                        # Opaque binary remains available in unsent_frames.
+                        pass
 
 
 class WebSocketQueueFullError(OpenAIError):

src/openai/_send_queue.py

@@ -12,32 +12,32 @@
 class SendQueue:
     """Bounded byte-size queue for outgoing WebSocket messages.
 
-    Messages are stored as pre-serialized strings. The queue enforces a
+    Messages retain their text or binary frame type. The queue enforces a
     maximum byte budget so that unbounded buffering cannot occur during
     reconnection windows.
     """
 
     def __init__(self, max_bytes: int = 1_048_576) -> None:
-        self._queue: list[tuple[str, int]] = []  # (data, byte_length)
+        self._queue: list[tuple[str | bytes, int]] = []  # (data, byte_length)
         self._bytes: int = 0
         self._max_bytes = max_bytes
         self._lock = threading.Lock()
         self._flush_done: threading.Event | None = None
 
-    def enqueue(self, data: str) -> None:
+    def enqueue(self, data: str | bytes) -> None:
         """Append *data* to the queue.
 
         Raises :class:`WebSocketQueueFullError` if the message would
         exceed the byte-size limit.
         """
-        byte_length = len(data.encode("utf-8"))
+        byte_length = len(data.encode("utf-8") if isinstance(data, str) else data)
         with self._lock:
             if self._bytes + byte_length > self._max_bytes:
                 raise WebSocketQueueFullError("send queue is full, message discarded")
             self._queue.append((data, byte_length))
             self._bytes += byte_length
 
-    def flush_sync(self, send: typing.Callable[[str], object], *, requeue_failed: bool = True) -> None:
+    def flush_sync(self, send: typing.Callable[[str | bytes], object], *, requeue_failed: bool = True) -> None:
         """Send every queued message via *send*.
 
         If *send* raises, the failing message and all subsequent messages
@@ -63,7 +63,7 @@ def flush_sync(self, send: typing.Callable[[str], object], *, requeue_failed: bo
             self._end_flush(pending)
 
     async def flush_async(
-        self, send: typing.Callable[[str], typing.Awaitable[object]], *, requeue_failed: bool = True
+        self, send: typing.Callable[[str | bytes], typing.Awaitable[object]], *, requeue_failed: bool = True
     ) -> None:
         """Async variant of :meth:`flush_sync`."""
         while isinstance(pending := self._begin_flush(), threading.Event):
@@ -86,7 +86,7 @@ async def flush_async(
         finally:
             self._end_flush(pending)
 
-    def _begin_flush(self) -> deque[tuple[str, int]] | threading.Event:
+    def _begin_flush(self) -> deque[tuple[str | bytes, int]] | threading.Event:
         with self._lock:
             if self._flush_done is not None:
                 return self._flush_done
@@ -96,14 +96,14 @@ def _begin_flush(self) -> deque[tuple[str, int]] | threading.Event:
             # Pending messages remain charged until their sends succeed.
             return pending
 
-    def _end_flush(self, pending: deque[tuple[str, int]]) -> None:
+    def _end_flush(self, pending: deque[tuple[str | bytes, int]]) -> None:
         with self._lock:
             self._queue = list(pending) + self._queue
             assert self._flush_done is not None
             self._flush_done.set()
             self._flush_done = None
 
-    def drain(self) -> list[str]:
+    def drain(self) -> list[str | bytes]:
         """Remove and return all queued messages."""
         with self._lock:
             items = [data for data, _ in self._queue]

src/openai/_websocket_connection.py

@@ -51,8 +51,7 @@ def __init__(
 
     def send_raw(self, data: bytes | str) -> None:
         if self._is_reconnecting:
-            raw = data if isinstance(data, str) else data.decode("utf-8")
-            self._send_queue.enqueue(raw)
+            self._send_queue.enqueue(data)
             return
         self._connection.send(data)
 
@@ -206,8 +205,7 @@ def __init__(
 
     async def send_raw(self, data: bytes | str) -> None:
         if self._is_reconnecting:
-            raw = data if isinstance(data, str) else data.decode("utf-8")
-            self._send_queue.enqueue(raw)
+            self._send_queue.enqueue(data)
             return
         await self._connection.send(data)
 
@@ -306,7 +304,7 @@ async def _attempt_reconnect(self, close_code: int) -> bool:
     async def _flush_send_queue(self) -> None:
         """Send all queued messages over the current connection."""
 
-        async def _send(data: str) -> None:
+        async def _send(data: str | bytes) -> None:
             await self._connection.send(data)
 
         try:

tests/test_send_queue.py

@@ -47,7 +47,7 @@ def test_flush_sync(self) -> None:
         q.enqueue("b")
         q.enqueue("c")
 
-        sent: list[str] = []
+        sent: list[str | bytes] = []
         q.flush_sync(sent.append)
         assert sent == ["a", "b", "c"]
         assert len(q) == 0
@@ -58,9 +58,9 @@ def test_flush_sync_requeues_on_failure(self) -> None:
         q.enqueue("b")
         q.enqueue("c")
 
-        sent: list[str] = []
+        sent: list[str | bytes] = []
 
-        def failing_send(data: str) -> None:
+        def failing_send(data: str | bytes) -> None:
             if data == "b":
                 raise RuntimeError("send failed")
             sent.append(data)
@@ -79,9 +79,9 @@ async def test_flush_async(self) -> None:
         q.enqueue("a")
         q.enqueue("b")
 
-        sent: list[str] = []
+        sent: list[str | bytes] = []
 
-        async def async_send(data: str) -> None:
+        async def async_send(data: str | bytes) -> None:
             sent.append(data)
 
         await q.flush_async(async_send)
@@ -95,9 +95,9 @@ async def test_flush_async_requeues_on_failure(self) -> None:
         q.enqueue("b")
         q.enqueue("c")
 
-        sent: list[str] = []
+        sent: list[str | bytes] = []
 
-        async def failing_send(data: str) -> None:
+        async def failing_send(data: str | bytes) -> None:
             if data == "b":
                 raise RuntimeError("send failed")
             sent.append(data)
@@ -116,7 +116,7 @@ def test_flush_sync_preserves_new_items_on_failure(self) -> None:
         q.enqueue("a")
         q.enqueue("b")
 
-        def failing_send(data: str) -> None:
+        def failing_send(data: str | bytes) -> None:
             if data == "b":
                 # Simulate another thread enqueuing during flush
                 q.enqueue("new")
@@ -130,10 +130,10 @@ def failing_send(data: str) -> None:
         assert remaining == ["b", "new"]
 
 
-async def flush_in_background(q: SendQueue, send: Callable[[str], None], asynchronous: bool) -> None:
+async def flush_in_background(q: SendQueue, send: Callable[[str | bytes], None], asynchronous: bool) -> None:
     if asynchronous:
 
-        async def async_send(data: str) -> None:
+        async def async_send(data: str | bytes) -> None:
             await asyncio.to_thread(send, data)
 
         await q.flush_async(async_send)
@@ -155,7 +155,7 @@ async def test_repeated_failed_flushes_remain_bounded(asynchronous: bool) -> Non
     for attempt in range(3):
         entered, release = threading.Event(), threading.Event()
 
-        def send(data: str, entered: threading.Event = entered, release: threading.Event = release) -> None:
+        def send(data: str | bytes, entered: threading.Event = entered, release: threading.Event = release) -> None:
             assert data == "éé"
             entered.set()
             assert release.wait(5)
@@ -174,7 +174,7 @@ def send(data: str, entered: threading.Event = entered, release: threading.Event
                 await task
         assert q._bytes == 6
 
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     await flush_in_background(q, sent.append, asynchronous)
     assert sent == ["éé", "a", "b"]
     assert q._bytes == 0
@@ -189,9 +189,9 @@ async def test_partial_flush_releases_only_successful_bytes(asynchronous: bool,
     q.enqueue("é")
     q.enqueue("bb")
     entered, release = threading.Event(), threading.Event()
-    sent: list[str] = []
+    sent: list[str | bytes] = []
 
-    def send(data: str) -> None:
+    def send(data: str | bytes) -> None:
         if data == "bb":
             entered.set()
             assert release.wait(5)
@@ -225,7 +225,7 @@ async def test_drain_during_flush_keeps_pending_bytes_charged(asynchronous: bool
     q.enqueue("aaa")
     entered, release = threading.Event(), threading.Event()
 
-    def send(_data: str) -> None:
+    def send(_data: str | bytes) -> None:
         entered.set()
         assert release.wait(5)
         raise RuntimeError("fake send failure")
@@ -256,17 +256,17 @@ async def test_overlapping_flushes_preserve_order(first_async: bool, second_asyn
     q.enqueue("a")
     q.enqueue("b")
     entered, release, second_sent = threading.Event(), threading.Event(), threading.Event()
-    sent: list[str] = []
+    sent: list[str | bytes] = []
 
-    def first_send(data: str) -> None:
+    def first_send(data: str | bytes) -> None:
         if data == "a":
             entered.set()
             assert release.wait(5)
             if fail:
                 raise RuntimeError("fake send failure")
         sent.append(data)
 
-    def second_send(data: str) -> None:
+    def second_send(data: str | bytes) -> None:
         second_sent.set()
         sent.append(data)
 
@@ -296,7 +296,7 @@ async def test_cancelled_flush_restores_unsent_messages() -> None:
     q.enqueue("b")
     entered = asyncio.Event()
 
-    async def send(data: str) -> None:
+    async def send(data: str | bytes) -> None:
         if data == "b":
             entered.set()
             await asyncio.Event().wait()
@@ -309,7 +309,7 @@ async def send(data: str) -> None:
         await task
     with pytest.raises(WebSocketQueueFullError):
         q.enqueue("d")
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     q.flush_sync(sent.append)
     assert sent == ["b", "cc"]
     assert q._bytes == 0
@@ -320,14 +320,14 @@ async def test_cancelled_flush_waiter_does_not_strand_ownership() -> None:
     q = SendQueue(max_bytes=2)
     q.enqueue("a")
     entered, release = asyncio.Event(), asyncio.Event()
-    sent: list[str] = []
+    sent: list[str | bytes] = []
 
-    async def first_send(data: str) -> None:
+    async def first_send(data: str | bytes) -> None:
         entered.set()
         await release.wait()
         sent.append(data)
 
-    async def second_send(data: str) -> None:
+    async def second_send(data: str | bytes) -> None:
         sent.append(data)
 
     first = asyncio.create_task(q.flush_async(first_send))
@@ -350,7 +350,7 @@ def test_interrupted_sync_flush_restores_unsent_messages() -> None:
     q.enqueue("a")
     q.enqueue("b")
 
-    def send(data: str) -> None:
+    def send(data: str | bytes) -> None:
         if data == "b":
             raise KeyboardInterrupt
 
@@ -359,7 +359,7 @@ def send(data: str) -> None:
     q.enqueue("c")
     with pytest.raises(WebSocketQueueFullError):
         q.enqueue("d")
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     q.flush_sync(sent.append)
     assert sent == ["b", "c"]
     assert q._bytes == 0

tests/test_send_queue_reconnect.py

@@ -136,7 +136,7 @@ def test_reconnect_retries_bounded_send_queue(
     ws = MagicMock()
     attempts = 0
 
-    def failing_send(data: str) -> None:
+    def failing_send(data: str | bytes) -> None:
         nonlocal attempts
         assert data == ("b" if attempts == 1 else "aaa")
         if attempts == 0:
@@ -167,7 +167,7 @@ def failing_send(data: str) -> None:
     ws.recv.return_value = '{"type": "response.created"}'
     connection.recv()
 
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     ws.send.side_effect = sent.append
     assert connection._reconnect(RuntimeError("fake disconnect"))
     assert sent == []
@@ -191,7 +191,7 @@ async def test_async_reconnect_retries_bounded_send_queue(
     ws = MagicMock()
     attempts = 0
 
-    async def failing_send(data: str) -> None:
+    async def failing_send(data: str | bytes) -> None:
         nonlocal attempts
         assert data == ("b" if attempts == 1 else "aaa")
         if attempts == 0:
@@ -222,7 +222,7 @@ async def failing_send(data: str) -> None:
     ws.recv = AsyncMock(return_value='{"type": "response.created"}')
     await connection.recv()
 
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     ws.send.side_effect = sent.append
     assert await connection._reconnect(RuntimeError("fake disconnect"))
     assert sent == []
@@ -247,7 +247,7 @@ async def test_cancelled_flush_does_not_replay_active_send(
     active = asyncio.Event()
     socket = MagicMock()
 
-    async def suspended_send(data: str) -> None:
+    async def suspended_send(data: str | bytes) -> None:
         assert data == "é"
         active.set()
         await asyncio.Event().wait()
@@ -268,7 +268,7 @@ async def suspended_send(data: str) -> None:
     # Only the active send was attempted; its bytes are released, and the
     # queued tail and concurrent enqueue remain in order for the next socket.
     q.enqueue("dd")
-    sent: list[str] = []
+    sent: list[str | bytes] = []
     socket.send.side_effect = sent.append
     await connection._flush_send_queue()
     assert sent == ["b", "c", "dd"]