refactor: share WebSocket connection handling

src/openai/_websocket_connection.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import time
+import random
+import logging
+import threading
+from typing import TYPE_CHECKING, Callable, Awaitable
+
+from ._types import Query, Headers
+from ._send_queue import SendQueue
+from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close
+
+if TYPE_CHECKING:
+    from websockets.sync.client import ClientConnection as WebSocketConnection
+    from websockets.asyncio.client import ClientConnection as AsyncWebSocketConnection
+
+
+class WebSocketConnectionBase:
+    def __init__(
+        self,
+        connection: WebSocketConnection,
+        *,
+        make_ws: Callable[[Query, Headers], WebSocketConnection] | None = None,
+        on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None,
+        max_retries: int = 5,
+        initial_delay: float = 0.5,
+        max_delay: float = 8.0,
+        extra_query: Query = {},
+        extra_headers: Headers = {},
+        send_queue: SendQueue | None = None,
+        logger: logging.Logger,
+        interruptible_reconnect: bool = False,
+    ) -> None:
+        self._connection = connection
+        self._make_ws = make_ws
+        self._on_reconnecting = on_reconnecting
+        self._max_retries = max_retries
+        self._reconnect_attempt = 0
+        self._initial_delay = initial_delay
+        self._max_delay = max_delay
+        self._extra_query = extra_query
+        self._extra_headers = extra_headers
+        self._intentionally_closed = False
+        self._is_reconnecting = False
+        self._send_queue = send_queue if send_queue is not None else SendQueue()
+        self._logger = logger
+        self._interruptible_reconnect = interruptible_reconnect
+        if interruptible_reconnect:
+            self._close_event = threading.Event()
+            self._connection_lock = threading.Lock()
+
+    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)
+            return
+        self._connection.send(data)
+
+    def close(self, *, code: int = 1000, reason: str = "") -> None:
+        self._intentionally_closed = True
+        if self._interruptible_reconnect:
+            self._close_event.set()
+            with self._connection_lock:
+                connection = self._connection
+        else:
+            connection = self._connection
+        connection.close(code=code, reason=reason)
+
+    def _reconnect(self, exc: Exception) -> bool:
+        """Attempt to reconnect after a connection failure.
+
+        Returns ``True`` if a new connection was established, ``False`` if the
+        caller should re-raise the original exception.
+        """
+        if (
+            self._on_reconnecting is None
+            or self._make_ws is None
+            or (self._interruptible_reconnect and self._close_event.is_set())
+        ):
+            return False
+
+        from websockets.exceptions import ConnectionClosedError
+
+        close_code = 1006
+        if isinstance(exc, ConnectionClosedError) and exc.rcvd is not None:
+            close_code = exc.rcvd.code
+
+        if not is_recoverable_close(close_code):
+            return False
+
+        if self._interruptible_reconnect:
+            with self._connection_lock:
+                self._is_reconnecting = True
+        else:
+            self._is_reconnecting = True
+
+        for attempt in range(self._reconnect_attempt + 1, self._max_retries + 1):
+            self._reconnect_attempt = attempt
+            base_delay = min(self._initial_delay * (2 ** (attempt - 1)), self._max_delay)
+            jitter = 0.75 + random.random() * 0.25
+            delay = base_delay * jitter
+
+            event = ReconnectingEvent(
+                attempt=attempt,
+                max_attempts=self._max_retries,
+                delay=delay,
+                close_code=close_code,
+                extra_query=self._extra_query,
+                extra_headers=self._extra_headers,
+            )
+
+            try:
+                result = self._on_reconnecting(event)
+            except Exception:
+                self._is_reconnecting = False
+                return False
+
+            if result is not None and result.get("abort"):
+                self._is_reconnecting = False
+                return False
+
+            if result is not None:
+                if "extra_query" in result:
+                    self._extra_query = result["extra_query"]
+                if "extra_headers" in result:
+                    self._extra_headers = result["extra_headers"]
+
+            self._logger.info(
+                "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay",
+                attempt,
+                self._max_retries,
+                delay,
+            )
+            if self._interruptible_reconnect:
+                closed = self._close_event.wait(delay)
+            else:
+                time.sleep(delay)
+                closed = self._intentionally_closed
+            if closed:
+                self._is_reconnecting = False
+                return False
+
+            try:
+                connection = self._make_ws(self._extra_query, self._extra_headers)
+                if self._interruptible_reconnect:
+                    with self._connection_lock:
+                        closed = self._intentionally_closed
+                        if not closed:
+                            self._connection = connection
+                    if closed:
+                        connection.close()
+                        self._is_reconnecting = False
+                        return False
+                else:
+                    self._connection = connection
+                self._logger.info("Reconnected to WebSocket API")
+                self._is_reconnecting = False
+                self._flush_send_queue()
+                return True
+            except Exception:
+                pass
+
+        self._is_reconnecting = False
+        return False
+
+    def _flush_send_queue(self) -> None:
+        """Send all queued messages over the current connection."""
+        try:
+            self._send_queue.flush_sync(lambda data: self._connection.send(data), requeue_failed=False)
+        except Exception:
+            self._logger.warning("Failed to flush send queue after reconnect")
+
+
+class AsyncWebSocketConnectionBase:
+    def __init__(
+        self,
+        connection: AsyncWebSocketConnection,
+        *,
+        make_ws: Callable[[Query, Headers], Awaitable[AsyncWebSocketConnection]] | None = None,
+        on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None,
+        max_retries: int = 5,
+        initial_delay: float = 0.5,
+        max_delay: float = 8.0,
+        extra_query: Query = {},
+        extra_headers: Headers = {},
+        send_queue: SendQueue | None = None,
+        logger: logging.Logger,
+    ) -> None:
+        self._connection = connection
+        self._make_ws = make_ws
+        self._on_reconnecting = on_reconnecting
+        self._max_retries = max_retries
+        self._reconnect_attempt = 0
+        self._initial_delay = initial_delay
+        self._max_delay = max_delay
+        self._extra_query = extra_query
+        self._extra_headers = extra_headers
+        self._intentionally_closed = False
+        self._is_reconnecting = False
+        self._send_queue = send_queue if send_queue is not None else SendQueue()
+        self._logger = logger
+
+    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)
+            return
+        await self._connection.send(data)
+
+    async def close(self, *, code: int = 1000, reason: str = "") -> None:
+        self._intentionally_closed = True
+        await self._connection.close(code=code, reason=reason)
+
+    async def _reconnect(self, exc: Exception) -> bool:
+        """Attempt to reconnect after a connection failure.
+
+        Returns ``True`` if a new connection was established, ``False`` if the
+        caller should re-raise the original exception.
+        """
+        import asyncio
+
+        if self._on_reconnecting is None or self._make_ws is None:
+            return False
+
+        from websockets.exceptions import ConnectionClosedError
+
+        close_code = 1006
+        if isinstance(exc, ConnectionClosedError) and exc.rcvd is not None:
+            close_code = exc.rcvd.code
+
+        if not is_recoverable_close(close_code):
+            return False
+
+        self._is_reconnecting = True
+
+        for attempt in range(self._reconnect_attempt + 1, self._max_retries + 1):
+            self._reconnect_attempt = attempt
+            base_delay = min(self._initial_delay * (2 ** (attempt - 1)), self._max_delay)
+            jitter = 0.75 + random.random() * 0.25
+            delay = base_delay * jitter
+
+            event = ReconnectingEvent(
+                attempt=attempt,
+                max_attempts=self._max_retries,
+                delay=delay,
+                close_code=close_code,
+                extra_query=self._extra_query,
+                extra_headers=self._extra_headers,
+            )
+
+            try:
+                result = self._on_reconnecting(event)
+            except Exception:
+                self._is_reconnecting = False
+                return False
+
+            if result is not None and result.get("abort"):
+                self._is_reconnecting = False
+                return False
+
+            if result is not None:
+                if "extra_query" in result:
+                    self._extra_query = result["extra_query"]
+                if "extra_headers" in result:
+                    self._extra_headers = result["extra_headers"]
+
+            self._logger.info(
+                "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay",
+                attempt,
+                self._max_retries,
+                delay,
+            )
+            await asyncio.sleep(delay)
+
+            if self._intentionally_closed:
+                self._is_reconnecting = False
+                return False
+
+            try:
+                self._connection = await self._make_ws(self._extra_query, self._extra_headers)
+                self._logger.info("Reconnected to WebSocket API")
+                self._is_reconnecting = False
+                await self._flush_send_queue()
+                return True
+            except Exception:
+                pass
+
+        self._is_reconnecting = False
+        return False
+
+    async def _flush_send_queue(self) -> None:
+        """Send all queued messages over the current connection."""
+
+        async def _send(data: str) -> None:
+            await self._connection.send(data)
+
+        try:
+            await self._send_queue.flush_async(_send, requeue_failed=False)
+        except Exception:
+            self._logger.warning("Failed to flush send queue after reconnect")