feat(client): allow enqueuing to websockets even when not connected

Author: stainless-app[bot]

src/openai/__init__.py

@@ -29,10 +29,12 @@
     InternalServerError,
     PermissionDeniedError,
     LengthFinishReasonError,
+    WebSocketQueueFullError,
     UnprocessableEntityError,
     APIResponseValidationError,
     InvalidWebhookSignatureError,
     ContentFilterFinishReasonError,
+    WebSocketConnectionClosedError,
 )
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
@@ -87,6 +89,8 @@
     "DefaultAioHttpClient",
     "ReconnectingEvent",
     "ReconnectingOverrides",
+    "WebSocketQueueFullError",
+    "WebSocketConnectionClosedError",
 ]
 
 if not _t.TYPE_CHECKING:

src/openai/_event_handler.py

@@ -70,3 +70,16 @@ def has_handlers(self, event_type: str) -> bool:
             return bool(handlers)
         finally:
             self._release()
+
+    def merge_into(self, target: EventHandlerRegistry) -> None:
+        """Move all handlers from this registry into *target*, then clear self."""
+        self._acquire()
+        try:
+            for event_type, handlers in self._handlers.items():
+                for handler in handlers:
+                    once = id(handler) in self._once_ids
+                    target.add(event_type, handler, once=once)
+            self._handlers.clear()
+            self._once_ids.clear()
+        finally:
+            self._release()

src/openai/_exceptions.py

@@ -28,6 +28,8 @@
     "ContentFilterFinishReasonError",
     "InvalidWebhookSignatureError",
     "SubjectTokenProviderError",
+    "WebSocketConnectionClosedError",
+    "WebSocketQueueFullError",
 ]
 
 
@@ -187,3 +189,19 @@ def __init__(self) -> None:
 
 class InvalidWebhookSignatureError(ValueError):
     """Raised when a webhook signature is invalid, meaning the computed signature does not match the expected signature."""
+
+
+class WebSocketConnectionClosedError(OpenAIError):
+    """Raised when a WebSocket connection closes with unsent messages."""
+
+    unsent_messages: list[str]
+
+    def __init__(self, message: str, *, unsent_messages: list[str]) -> None:
+        super().__init__(message)
+        self.unsent_messages = unsent_messages
+
+
+class WebSocketQueueFullError(OpenAIError):
+    """Raised when the outgoing WebSocket message queue exceeds its byte-size limit."""
+
+    pass

src/openai/_send_queue.py

@@ -0,0 +1,90 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from __future__ import annotations
+
+import typing
+import threading
+
+from ._exceptions import WebSocketQueueFullError
+
+
+class SendQueue:
+    """Bounded byte-size queue for outgoing WebSocket messages.
+
+    Messages are stored as pre-serialized strings. 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._bytes: int = 0
+        self._max_bytes = max_bytes
+        self._lock = threading.Lock()
+
+    def enqueue(self, data: str) -> 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"))
+        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]) -> None:
+        """Send every queued message via *send*.
+
+        If *send* raises, the failing message and all subsequent messages
+        are re-queued and the error is re-raised.
+        """
+        with self._lock:
+            pending = list(self._queue)
+            self._queue.clear()
+            self._bytes = 0
+
+        for i, (data, _byte_length) in enumerate(pending):
+            try:
+                send(data)
+            except Exception:
+                with self._lock:
+                    remaining = pending[i:]
+                    self._queue = remaining + self._queue
+                    self._bytes = sum(bl for _, bl in self._queue)
+                raise
+
+    async def flush_async(self, send: typing.Callable[[str], typing.Awaitable[object]]) -> None:
+        """Async variant of :meth:`flush_sync`."""
+        with self._lock:
+            pending = list(self._queue)
+            self._queue.clear()
+            self._bytes = 0
+
+        for i, (data, _byte_length) in enumerate(pending):
+            try:
+                await send(data)
+            except Exception:
+                with self._lock:
+                    remaining = pending[i:]
+                    self._queue = remaining + self._queue
+                    self._bytes = sum(bl for _, bl in self._queue)
+                raise
+
+    def drain(self) -> list[str]:
+        """Remove and return all queued messages."""
+        with self._lock:
+            items = [data for data, _ in self._queue]
+            self._queue.clear()
+            self._bytes = 0
+            return items
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._queue)
+
+    def __bool__(self) -> bool:
+        with self._lock:
+            return len(self._queue) > 0