From 25e5587457406a27e6fb2eb349f64566cc19294d Mon Sep 17 00:00:00 2001 From: markstuart-oai <323302876+markstuart-oai@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:38:30 +0000 Subject: [PATCH] refactor: share WebSocket connection handling Castiron-Internal-PR: https://github.com/openai/openai-python-internal/pull/256 Castiron-Source-SHA: 46ea721b3232d5eec49eb7f22ec07a08e1e657b7 Castiron-Public-Base-SHA: 95dbee9ee406587053a0a21296eb82418073f2aa --- src/openai/_websocket_connection.py | 300 ++++++++++++++++++ .../resources/beta/responses/responses.py | 249 ++------------- src/openai/resources/live/forks.py | 249 ++------------- src/openai/resources/live/live.py | 249 ++------------- src/openai/resources/live/sideband.py | 249 ++------------- src/openai/resources/realtime/realtime.py | 249 ++------------- src/openai/resources/responses/responses.py | 262 ++------------- tests/test_send_queue_reconnect.py | 30 ++ 8 files changed, 499 insertions(+), 1338 deletions(-) create mode 100644 src/openai/_websocket_connection.py diff --git a/src/openai/_websocket_connection.py b/src/openai/_websocket_connection.py new file mode 100644 index 0000000000..d26769acb8 --- /dev/null +++ b/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") diff --git a/src/openai/resources/beta/responses/responses.py b/src/openai/resources/beta/responses/responses.py index 773c23dd54..0199b766ae 100644 --- a/src/openai/resources/beta/responses/responses.py +++ b/src/openai/resources/beta/responses/responses.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import time -import random import logging from types import TracebackType from typing import TYPE_CHECKING, Any, Dict, List, Union, Callable, Iterable, Iterator, Optional, Awaitable, cast @@ -48,10 +46,11 @@ from ...._send_queue import SendQueue from ...._base_client import _merge_mappings, make_request_options from ...._event_handler import EventHandlerRegistry +from ...._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ....types.beta.beta_response import BetaResponse from ....types.beta.beta_tool_param import BetaToolParam from ....types.beta.beta_service_tier import BetaServiceTier -from ....types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ....types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ....types.beta.beta_compacted_response import BetaCompactedResponse from ....types.websocket_connection_options import WebSocketConnectionOptions from ....types.beta.beta_response_includable import BetaResponseIncludable @@ -4236,7 +4235,7 @@ def input_tokens(self) -> AsyncInputTokensWithStreamingResponse: return AsyncInputTokensWithStreamingResponse(self._responses.input_tokens) -class AsyncResponsesConnection: +class AsyncResponsesConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Responses API""" response: AsyncResponsesResponseResource @@ -4256,18 +4255,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.response = AsyncResponsesResponseResource(self) @@ -4344,17 +4343,6 @@ async def send(self, event: BetaResponsesClientEvent | BetaResponsesClientEventP return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> BetaResponsesServerEvent: """ Converts a raw `str` or `bytes` message into a `BetaResponsesServerEvent` object. @@ -4366,94 +4354,6 @@ def parse_event(self, data: str | bytes) -> BetaResponsesServerEvent: construct_type_unchecked(value=json.loads(data), type_=cast(Any, BetaResponsesServerEvent)), ) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncResponsesConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -4729,7 +4629,7 @@ async def __aexit__( await self.__connection.close() -class ResponsesConnection: +class ResponsesConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Responses API""" response: ResponsesResponseResource @@ -4749,18 +4649,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.response = ResponsesResponseResource(self) @@ -4839,17 +4739,6 @@ def send(self, event: BetaResponsesClientEvent | BetaResponsesClientEventParam) return self._connection.send(data) - 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 - self._connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> BetaResponsesServerEvent: """ Converts a raw `str` or `bytes` message into a `BetaResponsesServerEvent` object. @@ -4861,88 +4750,6 @@ def parse_event(self, data: str | bytes) -> BetaResponsesServerEvent: construct_type_unchecked(value=json.loads(data), type_=cast(Any, BetaResponsesServerEvent)), ) - 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: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - time.sleep(delay) - - if self._intentionally_closed: - self._is_reconnecting = False - return False - - try: - self._connection = self._make_ws(self._extra_query, self._extra_headers) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[ResponsesConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/src/openai/resources/live/forks.py b/src/openai/resources/live/forks.py index 8ee48e89d6..fa65e1be92 100644 --- a/src/openai/resources/live/forks.py +++ b/src/openai/resources/live/forks.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import time -import random import logging from types import TracebackType from typing import TYPE_CHECKING, Any, Union, Callable, Iterator, Optional, Awaitable, cast @@ -22,10 +20,11 @@ from ..._send_queue import SendQueue from ..._base_client import _merge_mappings from ..._event_handler import EventHandlerRegistry +from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ...types.live.error_event import ErrorEvent from ...types.live.fork_client_event import ForkClientEvent from ...types.live.fork_server_event import ForkServerEvent -from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ...types.live.fork_client_event_param import ForkClientEventParam from ...types.websocket_connection_options import WebSocketConnectionOptions from ...types.live.fork_session_config_param import ForkSessionConfigParam @@ -107,7 +106,7 @@ def connect( ) -class AsyncForksConnection: +class AsyncForksConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Forks API""" session: AsyncForksSessionResource @@ -128,18 +127,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.session = AsyncForksSessionResource(self) @@ -217,17 +216,6 @@ async def send(self, event: ForkClientEvent | ForkClientEventParam) -> None: return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> ForkServerEvent: """ Converts a raw `str` or `bytes` message into a `ForkServerEvent` object. @@ -236,94 +224,6 @@ def parse_event(self, data: str | bytes) -> ForkServerEvent: """ return cast(ForkServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ForkServerEvent))) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncForksConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -602,7 +502,7 @@ async def __aexit__( await self.__connection.close() -class ForksConnection: +class ForksConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Forks API""" session: ForksSessionResource @@ -623,18 +523,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.session = ForksSessionResource(self) @@ -712,17 +612,6 @@ def send(self, event: ForkClientEvent | ForkClientEventParam) -> None: return self._connection.send(data) - 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 - self._connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> ForkServerEvent: """ Converts a raw `str` or `bytes` message into a `ForkServerEvent` object. @@ -731,88 +620,6 @@ def parse_event(self, data: str | bytes) -> ForkServerEvent: """ return cast(ForkServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ForkServerEvent))) - 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: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - time.sleep(delay) - - if self._intentionally_closed: - self._is_reconnecting = False - return False - - try: - self._connection = self._make_ws(self._extra_query, self._extra_headers) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[ForksConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/src/openai/resources/live/live.py b/src/openai/resources/live/live.py index 16f625cdcf..f89cac978a 100644 --- a/src/openai/resources/live/live.py +++ b/src/openai/resources/live/live.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import time -import random import logging from types import TracebackType from typing import TYPE_CHECKING, Any, Union, Callable, Iterator, Optional, Awaitable, cast @@ -35,10 +33,11 @@ from ..._send_queue import SendQueue from ..._base_client import _merge_mappings, make_request_options from ..._event_handler import EventHandlerRegistry +from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ...types.live.error_event import ErrorEvent from ...types.live.client_event import ClientEvent from ...types.live.server_event import ServerEvent -from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ...types.live.client_event_param import ClientEventParam from ...types.live.live_create_response import LiveCreateResponse from ...types.live.session_config_param import SessionConfigParam @@ -328,7 +327,7 @@ def sessions(self) -> AsyncSessionsWithStreamingResponse: return AsyncSessionsWithStreamingResponse(self._live.sessions) -class AsyncLiveConnection: +class AsyncLiveConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Live API""" session: AsyncLiveSessionResource @@ -349,18 +348,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.session = AsyncLiveSessionResource(self) @@ -438,17 +437,6 @@ async def send(self, event: ClientEvent | ClientEventParam) -> None: return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> ServerEvent: """ Converts a raw `str` or `bytes` message into a `ServerEvent` object. @@ -457,94 +445,6 @@ def parse_event(self, data: str | bytes) -> ServerEvent: """ return cast(ServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ServerEvent))) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncLiveConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -820,7 +720,7 @@ async def __aexit__( await self.__connection.close() -class LiveConnection: +class LiveConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Live API""" session: LiveSessionResource @@ -841,18 +741,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.session = LiveSessionResource(self) @@ -930,17 +830,6 @@ def send(self, event: ClientEvent | ClientEventParam) -> None: return self._connection.send(data) - 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 - self._connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> ServerEvent: """ Converts a raw `str` or `bytes` message into a `ServerEvent` object. @@ -949,88 +838,6 @@ def parse_event(self, data: str | bytes) -> ServerEvent: """ return cast(ServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ServerEvent))) - 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: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - time.sleep(delay) - - if self._intentionally_closed: - self._is_reconnecting = False - return False - - try: - self._connection = self._make_ws(self._extra_query, self._extra_headers) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[LiveConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/src/openai/resources/live/sideband.py b/src/openai/resources/live/sideband.py index 9d14fba2ef..9cd1c68340 100644 --- a/src/openai/resources/live/sideband.py +++ b/src/openai/resources/live/sideband.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import time -import random import logging from types import TracebackType from typing import TYPE_CHECKING, Any, Union, Callable, Iterator, Optional, Awaitable, cast @@ -22,8 +20,9 @@ from ..._send_queue import SendQueue from ..._base_client import _merge_mappings from ..._event_handler import EventHandlerRegistry +from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ...types.live.error_event import ErrorEvent -from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ...types.live.connect_client_event import ConnectClientEvent from ...types.live.connect_server_event import ConnectServerEvent from ...types.websocket_connection_options import WebSocketConnectionOptions @@ -110,7 +109,7 @@ def connect( ) -class AsyncSidebandConnection: +class AsyncSidebandConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Sideband API""" session: AsyncSidebandSessionResource @@ -131,18 +130,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.session = AsyncSidebandSessionResource(self) @@ -220,17 +219,6 @@ async def send(self, event: ConnectClientEvent | ConnectClientEventParam) -> Non return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> ConnectServerEvent: """ Converts a raw `str` or `bytes` message into a `ConnectServerEvent` object. @@ -241,94 +229,6 @@ def parse_event(self, data: str | bytes) -> ConnectServerEvent: ConnectServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ConnectServerEvent)) ) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncSidebandConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -617,7 +517,7 @@ async def __aexit__( await self.__connection.close() -class SidebandConnection: +class SidebandConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Sideband API""" session: SidebandSessionResource @@ -638,18 +538,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.session = SidebandSessionResource(self) @@ -727,17 +627,6 @@ def send(self, event: ConnectClientEvent | ConnectClientEventParam) -> None: return self._connection.send(data) - 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 - self._connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> ConnectServerEvent: """ Converts a raw `str` or `bytes` message into a `ConnectServerEvent` object. @@ -748,88 +637,6 @@ def parse_event(self, data: str | bytes) -> ConnectServerEvent: ConnectServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, ConnectServerEvent)) ) - 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: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - time.sleep(delay) - - if self._intentionally_closed: - self._is_reconnecting = False - return False - - try: - self._connection = self._make_ws(self._extra_query, self._extra_headers) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[SidebandConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/src/openai/resources/realtime/realtime.py b/src/openai/resources/realtime/realtime.py index 7898f1b325..7a90858c72 100644 --- a/src/openai/resources/realtime/realtime.py +++ b/src/openai/resources/realtime/realtime.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import time -import random import logging from types import TracebackType from typing import TYPE_CHECKING, Any, Union, Callable, Iterator, Awaitable, cast @@ -47,6 +45,7 @@ ) from ..._event_handler import EventHandlerRegistry from ...types.realtime import session_update_event_param +from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from .translations.translations import ( Translations, AsyncTranslations, @@ -55,7 +54,7 @@ TranslationsWithStreamingResponse, AsyncTranslationsWithStreamingResponse, ) -from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ...types.websocket_connection_options import WebSocketConnectionOptions from ...types.realtime.realtime_error_event import RealtimeErrorEvent from ...types.realtime.realtime_client_event import RealtimeClientEvent @@ -291,7 +290,7 @@ def translations(self) -> AsyncTranslationsWithStreamingResponse: return AsyncTranslationsWithStreamingResponse(self._realtime.translations) -class AsyncRealtimeConnection: +class AsyncRealtimeConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Realtime API""" session: AsyncRealtimeSessionResource @@ -315,18 +314,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.session = AsyncRealtimeSessionResource(self) @@ -407,17 +406,6 @@ async def send(self, event: RealtimeClientEvent | RealtimeClientEventParam) -> N return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> RealtimeServerEvent: """ Converts a raw `str` or `bytes` message into a `RealtimeServerEvent` object. @@ -428,94 +416,6 @@ def parse_event(self, data: str | bytes) -> RealtimeServerEvent: RealtimeServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, RealtimeServerEvent)) ) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncRealtimeConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -818,7 +718,7 @@ async def __aexit__( await self.__connection.close() -class RealtimeConnection: +class RealtimeConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Realtime API""" session: RealtimeSessionResource @@ -842,18 +742,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.session = RealtimeSessionResource(self) @@ -934,17 +834,6 @@ def send(self, event: RealtimeClientEvent | RealtimeClientEventParam) -> None: return self._connection.send(data) - 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 - self._connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> RealtimeServerEvent: """ Converts a raw `str` or `bytes` message into a `RealtimeServerEvent` object. @@ -955,88 +844,6 @@ def parse_event(self, data: str | bytes) -> RealtimeServerEvent: RealtimeServerEvent, construct_type_unchecked(value=json.loads(data), type_=cast(Any, RealtimeServerEvent)) ) - 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: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - time.sleep(delay) - - if self._intentionally_closed: - self._is_reconnecting = False - return False - - try: - self._connection = self._make_ws(self._extra_query, self._extra_headers) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[RealtimeConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/src/openai/resources/responses/responses.py b/src/openai/resources/responses/responses.py index 15e00d2d02..5c3bdc1252 100644 --- a/src/openai/resources/responses/responses.py +++ b/src/openai/resources/responses/responses.py @@ -3,9 +3,7 @@ from __future__ import annotations import json -import random import logging -import threading from copy import copy from types import TracebackType from typing import ( @@ -64,6 +62,7 @@ response_retrieve_params, responses_client_event_param, ) +from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ...lib._parsing._responses import ( TextFormatT, parse_response, @@ -73,7 +72,7 @@ from ...types.responses.tool_param import ToolParam, ParseableToolParam from ...types.responses.service_tier import ServiceTier from ...types.shared_params.metadata import Metadata -from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close +from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ...types.shared_params.reasoning import Reasoning from ...types.responses.parsed_response import ParsedResponse from ...lib.streaming.responses._responses import ResponseStreamManager, AsyncResponseStreamManager @@ -4089,7 +4088,7 @@ def _make_tools(tools: Iterable[ParseableToolParam] | Omit) -> List[ToolParam] | return converted_tools -class AsyncResponsesConnection: +class AsyncResponsesConnection(AsyncWebSocketConnectionBase): """Represents a live WebSocket connection to the Responses API""" response: AsyncResponsesResponseResource @@ -4109,18 +4108,18 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=False) self.response = AsyncResponsesResponseResource(self) @@ -4197,17 +4196,6 @@ async def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> return await self._connection.send(data) - 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) - def parse_event(self, data: str | bytes) -> ResponsesServerEvent: """ Converts a raw `str` or `bytes` message into a `ResponsesServerEvent` object. @@ -4219,94 +4207,6 @@ def parse_event(self, data: str | bytes) -> ResponsesServerEvent: construct_type_unchecked(value=json.loads(data), type_=cast(Any, ResponsesServerEvent)), ) - 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"] - - log.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) - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[AsyncResponsesConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: @@ -4582,7 +4482,7 @@ async def __aexit__( await self.__connection.close() -class ResponsesConnection: +class ResponsesConnection(WebSocketConnectionBase): """Represents a live WebSocket connection to the Responses API""" response: ResponsesResponseResource @@ -4602,20 +4502,19 @@ def __init__( extra_headers: Headers = {}, send_queue: SendQueue | None = None, ) -> 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._close_event = threading.Event() - self._connection_lock = threading.Lock() - self._is_reconnecting = False - self._send_queue = send_queue if send_queue is not None else SendQueue() + super().__init__( + connection, + make_ws=make_ws, + on_reconnecting=on_reconnecting, + max_retries=max_retries, + initial_delay=initial_delay, + max_delay=max_delay, + extra_query=extra_query, + extra_headers=extra_headers, + send_queue=send_queue, + logger=log, + interruptible_reconnect=True, + ) self._event_handler_registry = EventHandlerRegistry(use_lock=True) self.response = ResponsesResponseResource(self) @@ -4694,20 +4593,6 @@ def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: return self._connection.send(data) - 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 - self._close_event.set() - with self._connection_lock: - connection = self._connection - connection.close(code=code, reason=reason) - def parse_event(self, data: str | bytes) -> ResponsesServerEvent: """ Converts a raw `str` or `bytes` message into a `ResponsesServerEvent` object. @@ -4719,95 +4604,6 @@ def parse_event(self, data: str | bytes) -> ResponsesServerEvent: construct_type_unchecked(value=json.loads(data), type_=cast(Any, ResponsesServerEvent)), ) - 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._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 - - with self._connection_lock: - 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"] - - log.info( - "Reconnecting to WebSocket API (attempt %d/%d) after %.1fs delay", - attempt, - self._max_retries, - delay, - ) - if self._close_event.wait(delay): - self._is_reconnecting = False - return False - - try: - connection = self._make_ws(self._extra_query, self._extra_headers) - with self._connection_lock: - closed = self._intentionally_closed - if not closed: - self._connection = connection - if closed: - connection.close() - self._is_reconnecting = False - return False - log.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: - log.warning("Failed to flush send queue after reconnect") - def on( self, event_type: str, handler: Callable[..., Any] | None = None ) -> Union[ResponsesConnection, Callable[[Callable[..., Any]], Callable[..., Any]]]: diff --git a/tests/test_send_queue_reconnect.py b/tests/test_send_queue_reconnect.py index de6140098c..0eacc81a04 100644 --- a/tests/test_send_queue_reconnect.py +++ b/tests/test_send_queue_reconnect.py @@ -11,6 +11,9 @@ from openai import OpenAI, AsyncOpenAI from openai._exceptions import WebSocketQueueFullError from openai._send_queue import SendQueue +from openai.resources.live.live import LiveConnection +from openai.resources.live.forks import ForksConnection +from openai.resources.live.sideband import SidebandConnection from openai.resources.realtime.realtime import RealtimeConnection, AsyncRealtimeConnection from openai.resources.responses.responses import ResponsesConnection, AsyncResponsesConnection from openai.resources.beta.responses.responses import ( @@ -21,6 +24,33 @@ from .lib.test_websocket_redirects import options, resource, reconnect, unexpected_http, async_http_client +@pytest.mark.parametrize( + ("connection_type", "recovers"), + [ + (ResponsesConnection, False), + (RealtimeConnection, True), + (LiveConnection, True), + (ForksConnection, True), + (SidebandConnection, True), + (BetaResponsesConnection, True), + ], +) +def test_close_during_replacement_preserves_connection_behavior(connection_type: type[Any], recovers: bool) -> None: + original, replacement = MagicMock(), MagicMock() + + def open_replacement(*_args: Any) -> MagicMock: + connection.close() + return replacement + + connection = connection_type(original, make_ws=open_replacement, on_reconnecting=reconnect, initial_delay=0) + assert connection._reconnect(RuntimeError("fake disconnect")) is recovers + original.close.assert_called_once_with(code=1000, reason="") + if recovers: + replacement.close.assert_not_called() + else: + replacement.close.assert_called_once_with() + + @pytest.mark.parametrize("name", ["realtime", "responses", "beta.responses"]) @pytest.mark.parametrize("prequeue", [False, True], ids=["empty", "prequeued"]) def test_manager_keeps_queue_limit_during_reconnect(monkeypatch: pytest.MonkeyPatch, name: str, prequeue: bool) -> None: