From 791ea571192b0c628bef4fdffe32e17ab33eeb7c Mon Sep 17 00:00:00 2001 From: markstuart-oai <323302876+markstuart-oai@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:43:08 +0000 Subject: [PATCH] fix: clean up failed WebSocket setup and honor native options Castiron-Internal-PR: https://github.com/openai/openai-python-internal/pull/295 Castiron-Source-SHA: 5fbf16077384d56a8645730663af13b8d2834ff6 Castiron-Public-Base-SHA: 0e5794dc48b8f8d51f6047f8a0b1814eb1e65850 --- .castiron.stats.yml | 8 +- scripts/castiron/generated-files.json | 16 +- src/openai/lib/realtime_translation.py | 4 +- .../resources/beta/realtime/realtime.py | 4 + .../resources/beta/responses/responses.py | 90 ++++-- src/openai/resources/live/forks.py | 90 ++++-- src/openai/resources/live/live.py | 90 ++++-- src/openai/resources/live/sideband.py | 90 ++++-- src/openai/resources/realtime/realtime.py | 90 ++++-- .../realtime/translations/translations.py | 90 ++++-- src/openai/resources/responses/responses.py | 90 ++++-- src/openai/types/__init__.py | 1 + .../types/websocket_connection_options.py | 7 +- tests/lib/responses/test_websocket_session.py | 3 + .../test_beta_realtime_websocket_options.py | 90 ++++++ .../test_generated_websocket_compression.py | 282 ++++++++++++++++++ tests/lib/test_websocket_entry_cleanup.py | 165 ++++++++++ tests/test_example_security.py | 6 +- tests/test_websocket_connection_options.py | 43 +++ 19 files changed, 1026 insertions(+), 233 deletions(-) create mode 100644 tests/lib/test_generated_websocket_compression.py create mode 100644 tests/lib/test_websocket_entry_cleanup.py diff --git a/.castiron.stats.yml b/.castiron.stats.yml index 07cc92d5db..4d33e03f35 100644 --- a/.castiron.stats.yml +++ b/.castiron.stats.yml @@ -1,8 +1,8 @@ schema_version: 1 -generation_id: bb452d19-8d5f-4013-bfd3-3edecb25c37e +generation_id: e5b5d7d9-1b97-4d6a-a1ca-8cd00ff2fa0b openapi_spec_hash: 72d34e76f62e434dc0b71e0b4ac0ce8a openapi_transformed_spec_hash: bca4c6179df519f310841c4730b72c18 config_hash: 8a49d8603252fe24c2dc96f50eb0d05b -codegen_sha: 016ee2b2da744118cf209752f57c8616ad20b842 -codegen_hash: 82cc4b7aca374db5d389b1ac06369a1ffa5946795a9ab86c87f57a5f5ec00390 -public_codegen_sha: e23db4894b46c4da1d4d9625eb0876d0ab512e78 +codegen_sha: 865a21f405fcf114702b18a248a26dd5ae716da5 +codegen_hash: 3b7d96954238e7ffaf3fd7e28623a73939510e435fb58ef4c9a3914be6f857fb +public_codegen_sha: d60ab2223decd81f675989eed4b14fa5363ad9a9 diff --git a/scripts/castiron/generated-files.json b/scripts/castiron/generated-files.json index 1233e098ff..e4dbe8ba20 100644 --- a/scripts/castiron/generated-files.json +++ b/scripts/castiron/generated-files.json @@ -411,7 +411,7 @@ }, "src/openai/resources/beta/responses/responses.py": { "mode": "100644", - "sha256": "e631b84ca7b00615a9e2f232d45747ec39f284ffa1ceb71517032e969bb386bb" + "sha256": "c7e705185228ea5a1662dc0fe569e0c16bb7eadf38beff925f6067643d98feb6" }, "src/openai/resources/beta/threads/__init__.py": { "mode": "100644", @@ -591,11 +591,11 @@ }, "src/openai/resources/live/forks.py": { "mode": "100644", - "sha256": "f3395e5c09d1e690b3c4fd0e2adf04990362e217df3b1ab2a046657d6e1b54b2" + "sha256": "ff834c298a4566869e81e769be92bcbed216510557ff4eb344fd840bab2f2cd5" }, "src/openai/resources/live/live.py": { "mode": "100644", - "sha256": "472c140e576cb42c4965e1f623fc7b29e50b01dae6b474c2e9db2f0190289f40" + "sha256": "a0ae2ae3eaaeccf152edca7999656be8cf10e5da49fc6335d479d71ba259b30d" }, "src/openai/resources/live/sessions.py": { "mode": "100644", @@ -603,7 +603,7 @@ }, "src/openai/resources/live/sideband.py": { "mode": "100644", - "sha256": "073e10b3edc936de7d35776069354c4272a9d7636070cb3ce4804dca978915f5" + "sha256": "37028100a5301b90b58cb6ed4ea2e20e0277fa476796d5d932cf848ad2a294c1" }, "src/openai/resources/models.py": { "mode": "100644", @@ -631,7 +631,7 @@ }, "src/openai/resources/realtime/realtime.py": { "mode": "100644", - "sha256": "a1af3c8791a4f29b66e3401348b6c6dea6dc658a8a6bb88843cdb30b4626119c" + "sha256": "29cf5fe34140d2f985151b264d3438897ed1579d7db1ef523a5831dd89ac7936" }, "src/openai/resources/realtime/translations/__init__.py": { "mode": "100644", @@ -643,7 +643,7 @@ }, "src/openai/resources/realtime/translations/translations.py": { "mode": "100644", - "sha256": "0da275b10d6f42a82cc83dde4ed659efc7b75ce1d8191365c22a3914fbb840e4" + "sha256": "2ed29f4fa975a92d0067c7c28e30fefad8070fc4bd6a7b35ff124b0bbaeb94b5" }, "src/openai/resources/responses/__init__.py": { "mode": "100644", @@ -663,7 +663,7 @@ }, "src/openai/resources/responses/responses.py": { "mode": "100644", - "sha256": "2239f2e9fe19aa019d2580a0f74a5d2efcedc2c284650ee1c2567af334a045e6" + "sha256": "77b3bd59e28c76c8c54d21d556718d30d6be437f1b19507e40e3708a24986e77" }, "src/openai/resources/safety/__init__.py": { "mode": "100644", @@ -755,7 +755,7 @@ }, "src/openai/types/__init__.py": { "mode": "100644", - "sha256": "5caaa29a79510a0a3e4a92f7b59d3ccb3f8f6651706f658738cb5e2a22b91481" + "sha256": "47c98e4ffd0eb3a28a9d60b78facf0bdd165d6d35b24ebb6bea5834a0165eeea" }, "src/openai/types/admin/__init__.py": { "mode": "100644", diff --git a/src/openai/lib/realtime_translation.py b/src/openai/lib/realtime_translation.py index 141d262c5e..ea79e5fa17 100644 --- a/src/openai/lib/realtime_translation.py +++ b/src/openai/lib/realtime_translation.py @@ -23,7 +23,7 @@ RealtimeTranslationClientEventParam, RealtimeTranslationSessionClosedEvent, ) -from ..types.websocket_connection_options import WebSocketConnectionOptions +from ..types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions __all__ = ["TranslationSession", "AsyncTranslationSession", "TranslationConnection", "AsyncTranslationConnection"] @@ -158,7 +158,7 @@ def __init__( model: str, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, ) -> None: self._resource = resource self._client = resource._client diff --git a/src/openai/resources/beta/realtime/realtime.py b/src/openai/resources/beta/realtime/realtime.py index ee70d639b3..7dd4fdc165 100644 --- a/src/openai/resources/beta/realtime/realtime.py +++ b/src/openai/resources/beta/realtime/realtime.py @@ -383,6 +383,8 @@ async def __aenter__(self) -> AsyncRealtimeConnection: log.debug("Custom WebSocket connection options provided") options: WebsocketConnectionOptions = {"max_size": None, **self.__websocket_connection_options} + if options["max_size"] is None: + options.setdefault("compression", None) self.__connection = AsyncRealtimeConnection( await connect( str(url), @@ -577,6 +579,8 @@ def __enter__(self) -> RealtimeConnection: log.debug("Custom WebSocket connection options provided") options: SyncWebSocketConnectionOptions = {"max_size": None, **self.__websocket_connection_options} + if options["max_size"] is None: + options.setdefault("compression", None) self.__connection = RealtimeConnection( connect( str(url), diff --git a/src/openai/resources/beta/responses/responses.py b/src/openai/resources/beta/responses/responses.py index 0199b766ae..6e61b51ebc 100644 --- a/src/openai/resources/beta/responses/responses.py +++ b/src/openai/resources/beta/responses/responses.py @@ -52,7 +52,7 @@ from ....types.beta.beta_service_tier import BetaServiceTier from ....types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides from ....types.beta.beta_compacted_response import BetaCompactedResponse -from ....types.websocket_connection_options import WebSocketConnectionOptions +from ....types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ....types.beta.beta_response_includable import BetaResponseIncludable from ....types.beta.beta_response_input_param import BetaResponseInputParam from ....types.beta.beta_response_prompt_param import BetaResponsePromptParam @@ -2069,7 +2069,7 @@ def connect( self, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -4545,20 +4545,29 @@ async def __aenter__(self) -> AsyncResponsesConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncResponsesConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncResponsesConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -4603,11 +4612,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -4851,7 +4865,7 @@ def __init__( client: OpenAI, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -4935,20 +4949,29 @@ def __enter__(self) -> ResponsesConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = ResponsesConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = ResponsesConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -4993,11 +5016,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/live/forks.py b/src/openai/resources/live/forks.py index b792516d39..1cece08732 100644 --- a/src/openai/resources/live/forks.py +++ b/src/openai/resources/live/forks.py @@ -28,7 +28,7 @@ from ...types.live.fork_server_event import ForkServerEvent 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.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ...types.live.fork_session_config_param import ForkSessionConfigParam from ...types.live.session_update_config_param import SessionUpdateConfigParam from ...types.responses.response_input_item_param import ResponseInputItemParam @@ -51,7 +51,7 @@ def connect( session_id: str, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -423,20 +423,29 @@ async def __aenter__(self) -> AsyncForksConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncForksConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncForksConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -481,11 +490,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -734,7 +748,7 @@ def __init__( session_id: str, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -819,20 +833,29 @@ def __enter__(self) -> ForksConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = ForksConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = ForksConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -877,11 +900,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/live/live.py b/src/openai/resources/live/live.py index b93468b96c..422057288e 100644 --- a/src/openai/resources/live/live.py +++ b/src/openai/resources/live/live.py @@ -42,7 +42,7 @@ from ...types.live.client_event_param import ClientEventParam from ...types.live.live_create_response import LiveCreateResponse from ...types.live.session_config_param import SessionConfigParam -from ...types.websocket_connection_options import WebSocketConnectionOptions +from ...types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ...types.live.media_session_config_param import MediaSessionConfigParam from ...types.live.session_update_config_param import SessionUpdateConfigParam from ...types.responses.response_input_item_param import ResponseInputItemParam @@ -143,7 +143,7 @@ def connect( self, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -641,20 +641,29 @@ async def __aenter__(self) -> AsyncLiveConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncLiveConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncLiveConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -699,11 +708,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -949,7 +963,7 @@ def __init__( client: OpenAI, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -1033,20 +1047,29 @@ def __enter__(self) -> LiveConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = LiveConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = LiveConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -1091,11 +1114,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/live/sideband.py b/src/openai/resources/live/sideband.py index 329ce7a13a..a3fb2ebc73 100644 --- a/src/openai/resources/live/sideband.py +++ b/src/openai/resources/live/sideband.py @@ -27,7 +27,7 @@ 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 +from ...types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ...types.live.connect_client_event_param import ConnectClientEventParam from ...types.live.session_update_config_param import SessionUpdateConfigParam from ...types.responses.response_input_item_param import ResponseInputItemParam @@ -51,7 +51,7 @@ def connect( graceful_close: bool | Omit = omit, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -430,20 +430,29 @@ async def __aenter__(self) -> AsyncSidebandConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncSidebandConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncSidebandConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -496,11 +505,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -752,7 +766,7 @@ def __init__( graceful_close: bool | Omit = omit, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -838,20 +852,29 @@ def __enter__(self) -> SidebandConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = SidebandConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = SidebandConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -904,11 +927,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/realtime/realtime.py b/src/openai/resources/realtime/realtime.py index 1575803283..2ab8ced9f3 100644 --- a/src/openai/resources/realtime/realtime.py +++ b/src/openai/resources/realtime/realtime.py @@ -55,7 +55,7 @@ AsyncTranslationsWithStreamingResponse, ) from ...types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides -from ...types.websocket_connection_options import WebSocketConnectionOptions +from ...types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ...types.realtime.realtime_error_event import RealtimeErrorEvent from ...types.realtime.realtime_client_event import RealtimeClientEvent from ...types.realtime.realtime_server_event import RealtimeServerEvent @@ -115,7 +115,7 @@ def connect( model: str | Omit = omit, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -611,20 +611,29 @@ async def __aenter__(self) -> AsyncRealtimeConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncRealtimeConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncRealtimeConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -692,11 +701,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -949,7 +963,7 @@ def __init__( model: str | Omit = omit, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -1035,20 +1049,29 @@ def __enter__(self) -> RealtimeConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = RealtimeConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = RealtimeConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -1114,11 +1137,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/realtime/translations/translations.py b/src/openai/resources/realtime/translations/translations.py index ec50d97eda..8fb860f720 100644 --- a/src/openai/resources/realtime/translations/translations.py +++ b/src/openai/resources/realtime/translations/translations.py @@ -30,7 +30,7 @@ from ...._event_handler import EventHandlerRegistry from ...._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase from ....types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides -from ....types.websocket_connection_options import WebSocketConnectionOptions +from ....types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ....types.realtime.realtime_error_event import RealtimeErrorEvent from ....types.realtime.realtime_translation_client_event import RealtimeTranslationClientEvent from ....types.realtime.realtime_translation_server_event import RealtimeTranslationServerEvent @@ -79,7 +79,7 @@ def connect( model: str, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -503,20 +503,29 @@ async def __aenter__(self) -> AsyncTranslationsConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncTranslationsConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncTranslationsConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -562,11 +571,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -811,7 +825,7 @@ def __init__( model: str, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -899,20 +913,29 @@ def __enter__(self) -> TranslationsConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = TranslationsConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = TranslationsConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -958,11 +981,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/resources/responses/responses.py b/src/openai/resources/responses/responses.py index 5c3bdc1252..54300192eb 100644 --- a/src/openai/resources/responses/responses.py +++ b/src/openai/resources/responses/responses.py @@ -77,7 +77,7 @@ from ...types.responses.parsed_response import ParsedResponse from ...lib.streaming.responses._responses import ResponseStreamManager, AsyncResponseStreamManager from ...types.responses.compacted_response import CompactedResponse -from ...types.websocket_connection_options import WebSocketConnectionOptions +from ...types.websocket_connection_options import WebSocketConnectionOptions, SyncWebSocketConnectionOptions from ...types.responses.response_includable import ResponseIncludable from ...types.shared_params.responses_model import ResponsesModel from ...types.responses.response_input_param import ResponseInputParam @@ -1986,7 +1986,7 @@ def connect( self, extra_query: Query = {}, extra_headers: Headers = {}, - websocket_connection_options: WebSocketConnectionOptions = {}, + websocket_connection_options: SyncWebSocketConnectionOptions = {}, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -4398,20 +4398,29 @@ async def __aenter__(self) -> AsyncResponsesConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = AsyncResponsesConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = AsyncResponsesConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - await self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + await self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + await ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -4456,11 +4465,16 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return await connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: @@ -4705,7 +4719,7 @@ def __init__( client: OpenAI, extra_query: Query, extra_headers: Headers, - websocket_connection_options: WebSocketConnectionOptions, + websocket_connection_options: SyncWebSocketConnectionOptions, on_reconnecting: Callable[[ReconnectingEvent], ReconnectingOverrides | None] | None = None, max_retries: int = 5, initial_delay: float = 0.5, @@ -4789,20 +4803,29 @@ def __enter__(self) -> ResponsesConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - self.__connection = ResponsesConnection( - ws, - make_ws=self._connect_ws if self.__on_reconnecting is not None else None, - on_reconnecting=self.__on_reconnecting, - max_retries=self.__max_retries, - initial_delay=self.__initial_delay, - max_delay=self.__max_delay, - extra_query=self.__extra_query, - extra_headers=self.__extra_headers, - send_queue=self.__send_queue, - ) + try: + self.__connection = ResponsesConnection( + ws, + make_ws=self._connect_ws if self.__on_reconnecting is not None else None, + on_reconnecting=self.__on_reconnecting, + max_retries=self.__max_retries, + initial_delay=self.__initial_delay, + max_delay=self.__max_delay, + extra_query=self.__extra_query, + extra_headers=self.__extra_headers, + send_queue=self.__send_queue, + ) - self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) - self.__connection._flush_send_queue() + self.__event_handler_registry.merge_into(self.__connection._event_handler_registry) + self.__connection._flush_send_queue() + except BaseException: + self.__connection = None + try: + ws.close() + except BaseException: + # Preserve the exception that prevented successful entry. + pass + raise return self.__connection @@ -4847,11 +4870,16 @@ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketCo if self.__websocket_connection_options: log.debug("Custom WebSocket connection options provided") + websocket_options = self.__websocket_connection_options.copy() + websocket_options.setdefault("max_size", None) + if websocket_options.get("max_size") is None: + websocket_options.setdefault("compression", None) + return connect( str(url), user_agent_header=None, additional_headers=_merge_mappings(dict(headers.values()), {}), - **self.__websocket_connection_options, + **websocket_options, ) def _prepare_url(self) -> httpx2.URL: diff --git a/src/openai/types/__init__.py b/src/openai/types/__init__.py index 24bc681ca5..65ef9dbda1 100644 --- a/src/openai/types/__init__.py +++ b/src/openai/types/__init__.py @@ -119,6 +119,7 @@ from .websocket_connection_options import ( WebSocketConnectionOptions as WebSocketConnectionOptions, WebsocketConnectionOptions as WebsocketConnectionOptions, + SyncWebSocketConnectionOptions as SyncWebSocketConnectionOptions, ) from .image_create_variation_params import ImageCreateVariationParams as ImageCreateVariationParams from .image_gen_partial_image_event import ImageGenPartialImageEvent as ImageGenPartialImageEvent diff --git a/src/openai/types/websocket_connection_options.py b/src/openai/types/websocket_connection_options.py index d8bfabad18..050a049475 100644 --- a/src/openai/types/websocket_connection_options.py +++ b/src/openai/types/websocket_connection_options.py @@ -25,7 +25,12 @@ class SyncWebSocketConnectionOptions(TypedDict, total=False): """List of supported subprotocols, in order of decreasing preference.""" compression: str | None - """The “permessage-deflate” extension is enabled by default. Set compression to None to disable it. See the [compression guide](https://websockets.readthedocs.io/en/stable/topics/compression.html) for details.""" + """With unlimited messages (`max_size=None`, the SDK default), compression + defaults to `None`. With a finite `max_size`, including zero, the native + `websockets` default enables permessage-deflate. Set `compression="deflate"` + or `compression=None` to choose explicitly. Factories provided in + `extensions` can still negotiate compression when this option is `None`. + """ # limits max_size: int | None diff --git a/tests/lib/responses/test_websocket_session.py b/tests/lib/responses/test_websocket_session.py index f5fe69dd58..db616621ae 100644 --- a/tests/lib/responses/test_websocket_session.py +++ b/tests/lib/responses/test_websocket_session.py @@ -1335,6 +1335,7 @@ def script(socket: ServerConnection) -> None: async def test_sync_close_progresses_while_command_iterable_is_paused() -> None: entered = threading.Event() release = threading.Event() + peer_ready = threading.Event() def tools() -> Generator[dict[str, Any], None, None]: entered.set() @@ -1343,11 +1344,13 @@ def tools() -> Generator[dict[str, Any], None, None]: def script(socket: ServerConnection) -> None: # Closing during serialization must not send a request. + peer_ready.set() with pytest.raises(ConnectionClosedOK): socket.recv(timeout=10) with script_server(script) as url: async with session_for("sync", url) as driver: + assert await asyncio.to_thread(peer_ready.wait, 5), "WebSocket peer didn't start" send = asyncio.create_task( driver.call(driver.session.default, "send", {"type": "response.create", "tools": tools()}) ) diff --git a/tests/lib/test_beta_realtime_websocket_options.py b/tests/lib/test_beta_realtime_websocket_options.py index 259b9b43bd..09a318b276 100644 --- a/tests/lib/test_beta_realtime_websocket_options.py +++ b/tests/lib/test_beta_realtime_websocket_options.py @@ -6,6 +6,7 @@ import pytest from websockets.exceptions import ConnectionClosed, ConnectionClosedError from websockets.sync.server import ServerConnection, serve +from websockets.extensions.permessage_deflate import ClientPerMessageDeflateFactory from openai import OpenAI, AsyncOpenAI from openai.types.websocket_connection_options import WebSocketConnectionOptions @@ -82,3 +83,92 @@ def receive() -> bytes: assert exc.value.sent is not None and exc.value.sent.code == 1009 else: assert await asyncio.wait_for(conn.recv_bytes(), 10) == expected + + +@pytest.fixture +def compression_peer() -> Iterator[tuple[str, bytes, list[tuple[list[str], list[str]]]]]: + wire = b'{"type":"future.event","text":"ok"}' + observed: list[tuple[list[str], list[str]]] = [] + + def accept(socket: ServerConnection) -> None: + assert socket.request is not None + observed.append( + ( + socket.request.headers.get_all("Sec-WebSocket-Extensions"), + [extension.name for extension in socket.protocol.extensions], + ) + ) + socket.send(wire) + try: + socket.recv(timeout=5) + except ConnectionClosed: + # The client closes after consuming or rejecting the test event. + pass + + with serve(accept, "127.0.0.1", 0, close_timeout=1) as server: + thread = Thread(target=server.serve_forever) + thread.start() + try: + yield f"ws://127.0.0.1:{server.socket.getsockname()[1]}", wire, observed + finally: + server.shutdown() + thread.join(5) + assert not thread.is_alive() + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("options", "compressed"), + [ + pytest.param({}, False, id="absent-limit"), + pytest.param({"max_size": None}, False, id="explicit-unlimited"), + pytest.param({"max_size": 0}, True, id="zero-limit-native-default"), + pytest.param({"max_size": 512}, True, id="finite-limit-native-default"), + pytest.param({"compression": "deflate"}, True, id="explicit-deflate-unlimited"), + pytest.param({"compression": None}, False, id="explicit-off-unlimited"), + pytest.param({"max_size": 512, "compression": None}, False, id="explicit-off-finite"), + pytest.param({"max_size": 512, "compression": "deflate"}, True, id="explicit-deflate-finite"), + pytest.param( + {"extensions": [ClientPerMessageDeflateFactory()]}, + True, + id="explicit-extension-unlimited", + ), + ], +) +async def test_beta_realtime_negotiates_compression_by_caller_choice( + compression_peer: tuple[str, bytes, list[tuple[list[str], list[str]]]], + mode: Literal["sync", "async"], + options: WebSocketConnectionOptions, + compressed: bool, +) -> None: + url, wire, observed = compression_peer + original = options.copy() + reject = options.get("max_size") == 0 + if mode == "sync": + + def receive() -> bytes: + with OpenAI(api_key="local", websocket_base_url=url) as client: + with client.beta.realtime.connect(model="test", websocket_connection_options=options) as conn: + return conn.recv_bytes() + + if reject: + with pytest.raises(ConnectionClosedError) as exc: + await asyncio.wait_for(asyncio.to_thread(receive), 10) + assert exc.value.sent is not None and exc.value.sent.code == 1009 + else: + assert await asyncio.wait_for(asyncio.to_thread(receive), 10) == wire + else: + async with AsyncOpenAI(api_key="local", websocket_base_url=url) as client: + async with client.beta.realtime.connect(model="test", websocket_connection_options=options) as conn: + if reject: + with pytest.raises(ConnectionClosedError) as exc: + await asyncio.wait_for(conn.recv_bytes(), 10) + assert exc.value.sent is not None and exc.value.sent.code == 1009 + else: + assert await asyncio.wait_for(conn.recv_bytes(), 10) == wire + + assert options == original + assert len(observed) == 1 + offered, negotiated = observed[0] + assert ("permessage-deflate" in ", ".join(offered)) is compressed + assert ("permessage-deflate" in negotiated) is compressed diff --git a/tests/lib/test_generated_websocket_compression.py b/tests/lib/test_generated_websocket_compression.py new file mode 100644 index 0000000000..79578d62e3 --- /dev/null +++ b/tests/lib/test_generated_websocket_compression.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +import json +import asyncio +from typing import Any, Literal, Callable, Generator +from threading import Event, Thread +from contextlib import contextmanager + +import httpx2 +import pytest +from websockets.exceptions import ConnectionClosed, ConnectionClosedError +from websockets.sync.server import ServerConnection, serve +from websockets.extensions.permessage_deflate import ClientPerMessageDeflateFactory + +from openai import OpenAI, AsyncOpenAI +from openai.types.websocket_reconnection import ReconnectingEvent +from tests.lib.test_websocket_entry_cleanup import PRODUCTS +from openai.types.websocket_connection_options import WebSocketConnectionOptions + + +@pytest.fixture(autouse=True) +def no_proxies(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NO_PROXY", "*") + monkeypatch.setenv("no_proxy", "*") + + +def manager( + client: Any, + product: str, + options: WebSocketConnectionOptions, + on_reconnecting: Callable[[ReconnectingEvent], None] | None = None, +) -> Any: + settings: dict[str, Any] = { + "websocket_connection_options": options, + "on_reconnecting": on_reconnecting, + "initial_delay": 0, + "max_retries": 1, + } + if product == "realtime": + return client.realtime.connect(model="test", **settings) + if product == "transcription": + return client.realtime.connect(extra_query={"intent": "transcription"}, **settings) + if product == "sideband": + return client.realtime.connect(call_id="call_1", **settings) + if product == "responses": + return client.responses.connect(**settings) + if product == "beta.responses": + return client.beta.responses.connect(**settings) + if product == "live": + return client.live.connect(**settings) + if product == "live-sideband": + return client.live.sideband.connect(session_id="live_1", **settings) + if product == "live-forks": + return client.live.forks.connect(session_id="live_1", **settings) + assert product == "translation" + return client.realtime.translations.connect("test", **settings) + + +@contextmanager +def compression_peer( + product: str, *, reconnect: bool = False, wire: bytes | None = None +) -> Generator[tuple[str, bytes, list[tuple[list[str], list[str]]], list[tuple[int, str | bytes]]], None, None]: + if wire is None: + if product == "live": + wire = b'{"type":"session.output_transcript.delta","event_id":"recovered","delta":"go"}' + else: + wire = b'{"type":"input_audio_buffer.cleared","event_id":"recovered"}' + handshakes: list[tuple[list[str], list[str]]] = [] + frames: list[tuple[int, str | bytes]] = [] + failures: list[Exception] = [] + finished: list[Event] = [] + + def accept(socket: ServerConnection) -> None: + done = Event() + finished.append(done) + try: + assert socket.request is not None + handshakes.append( + ( + socket.request.headers.get_all("Sec-WebSocket-Extensions"), + [extension.name for extension in socket.protocol.extensions], + ) + ) + current = len(handshakes) + if reconnect and current == 1: + frames.append((current, socket.recv(timeout=5))) + socket.close(1011, "synthetic restart") + return + socket.send(wire) + try: + while True: + message = socket.recv(timeout=5) + frames.append((current, message)) + if reconnect and current == 2 and message == "after-reconnect": + socket.send(b"test.peer.received") + except ConnectionClosed: + # Caller either consumed the message and closed, or rejected size 0. + pass + except Exception as exc: + failures.append(exc) + finally: + done.set() + + with serve(accept, "127.0.0.1", 0, close_timeout=1) as server: + thread = Thread(target=server.serve_forever) + thread.start() + try: + yield f"http://127.0.0.1:{server.socket.getsockname()[1]}/v1", wire, handshakes, frames + finally: + server.shutdown() + thread.join(5) + assert not thread.is_alive() + for done in finished: + assert done.wait(5), "peer handler did not finish" + assert not failures, failures + + +def assert_compression(handshakes: list[tuple[list[str], list[str]]], compressed: bool, count: int = 1) -> None: + assert len(handshakes) == count + for offered, negotiated in handshakes: + assert ("permessage-deflate" in ", ".join(offered)) is compressed + assert ("permessage-deflate" in negotiated) is compressed + + +async def receive_once( + product: str, + mode: Literal["sync", "async"], + options: WebSocketConnectionOptions, + compressed: bool, + *, + wire: bytes | None = None, +) -> None: + with compression_peer(product, wire=wire) as (url, wire, handshakes, frames): + reject = options.get("max_size") == 0 + if mode == "sync": + + def read() -> bytes: + with OpenAI( + api_key="local", + base_url=url, + websocket_base_url=url.replace("http:", "ws:"), + http_client=httpx2.Client(trust_env=False), + ) as client: + with manager(client, product, options) as connection: + return connection.recv_bytes() + + if reject: + with pytest.raises(ConnectionClosedError) as exc: + await asyncio.wait_for(asyncio.to_thread(read), 10) + assert exc.value.sent is not None and exc.value.sent.code == 1009 + else: + assert await asyncio.wait_for(asyncio.to_thread(read), 10) == wire + else: + async with AsyncOpenAI( + api_key="local", + base_url=url, + websocket_base_url=url.replace("http:", "ws:"), + http_client=httpx2.AsyncClient(trust_env=False), + ) as async_client: + async with manager(async_client, product, options) as connection: + if reject: + with pytest.raises(ConnectionClosedError) as exc: + await asyncio.wait_for(connection.recv_bytes(), 10) + assert exc.value.sent is not None and exc.value.sent.code == 1009 + else: + assert await asyncio.wait_for(connection.recv_bytes(), 10) == wire + assert frames == [] + assert_compression(handshakes, compressed) + + +@pytest.mark.parametrize("product", PRODUCTS) +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("options", [{}, {"max_size": None}], ids=["absent-limit", "explicit-unlimited"]) +async def test_generated_unlimited_default( + product: str, mode: Literal["sync", "async"], options: WebSocketConnectionOptions +) -> None: + await receive_once(product, mode, options, compressed=False) + + +@pytest.mark.parametrize("product", PRODUCTS) +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_generated_default_receives_more_than_native_message_limit( + product: str, mode: Literal["sync", "async"] +) -> None: + event: dict[str, object] = {"event_id": "event_1", "delta": "x" * (2 * 1024 * 1024)} + if product.startswith("live") or product == "translation": + event.update(type="session.output_transcript.delta", start_ms=0, end_ms=200) + elif product == "transcription": + event.update(type="conversation.item.input_audio_transcription.delta", item_id="item_1", content_index=0) + else: + event.update( + type="response.output_text.delta", + response_id="resp_1", + item_id="item_1", + output_index=0, + content_index=0, + sequence_number=1, + logprobs=[], + ) + await receive_once(product, mode, {}, compressed=False, wire=json.dumps(event).encode()) + + +@pytest.mark.parametrize("product", ["realtime", "live"]) +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("options", "compressed"), + [ + pytest.param({"max_size": 0}, True, id="zero-limit-native-default"), + pytest.param({"max_size": 512}, True, id="finite-limit-native-default"), + pytest.param({"compression": "deflate"}, True, id="explicit-deflate-unlimited"), + pytest.param({"compression": None}, False, id="explicit-off-unlimited"), + pytest.param({"max_size": 512, "compression": None}, False, id="explicit-off-finite"), + pytest.param({"max_size": 512, "compression": "deflate"}, True, id="explicit-deflate-finite"), + pytest.param( + {"extensions": [ClientPerMessageDeflateFactory()]}, + True, + id="explicit-extension-unlimited", + ), + ], +) +async def test_generated_respects_compression_options( + product: str, mode: Literal["sync", "async"], options: WebSocketConnectionOptions, compressed: bool +) -> None: + await receive_once(product, mode, options, compressed) + + +@pytest.mark.parametrize("product", ["realtime", "live"]) +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("options", "compressed"), + [ + pytest.param({}, False, id="absent-limit"), + pytest.param({"max_size": None}, False, id="explicit-unlimited"), + pytest.param({"max_size": 512}, True, id="finite-limit-native-default"), + pytest.param({"compression": "deflate"}, True, id="explicit-deflate-unlimited"), + pytest.param({"compression": None}, False, id="explicit-off-unlimited"), + pytest.param({"extensions": [ClientPerMessageDeflateFactory()]}, True, id="explicit-extension-unlimited"), + ], +) +async def test_generated_reconnect_preserves_compression_without_replay( + product: str, mode: Literal["sync", "async"], options: WebSocketConnectionOptions, compressed: bool +) -> None: + attempts: list[int] = [] + + def on_reconnecting(event: ReconnectingEvent) -> None: + assert event.close_code == 1011 + attempts.append(event.attempt) + + with compression_peer(product, reconnect=True) as (url, _wire, handshakes, frames): + if mode == "sync": + + def recover() -> None: + with OpenAI( + api_key="local", + base_url=url, + websocket_base_url=url.replace("http:", "ws:"), + http_client=httpx2.Client(trust_env=False), + ) as client: + with manager(client, product, options, on_reconnecting) as connection: + connection.send_raw("before-close") + event = next(iter(connection)) + assert event.event_id == "recovered" + connection.send_raw("after-reconnect") + assert connection.recv_bytes() == b"test.peer.received" + + await asyncio.wait_for(asyncio.to_thread(recover), 10) + else: + async with AsyncOpenAI( + api_key="local", + base_url=url, + websocket_base_url=url.replace("http:", "ws:"), + http_client=httpx2.AsyncClient(trust_env=False), + ) as async_client: + async with manager(async_client, product, options, on_reconnecting) as connection: + await connection.send_raw("before-close") + event = await asyncio.wait_for(anext(aiter(connection)), 10) + assert event.event_id == "recovered" + await connection.send_raw("after-reconnect") + assert await asyncio.wait_for(connection.recv_bytes(), 5) == b"test.peer.received" + assert attempts == [1] + assert frames == [(1, "before-close"), (2, "after-reconnect")] + assert_compression(handshakes, compressed, count=2) diff --git a/tests/lib/test_websocket_entry_cleanup.py b/tests/lib/test_websocket_entry_cleanup.py new file mode 100644 index 0000000000..7c1248c2af --- /dev/null +++ b/tests/lib/test_websocket_entry_cleanup.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import json +import asyncio +from typing import Any + +import pytest + +from openai import OpenAI, AsyncOpenAI + +PRODUCTS = [ + "realtime", + "transcription", + "sideband", + "responses", + "beta.responses", + "live", + "live-sideband", + "live-forks", + "translation", +] + + +def manager(client: Any, product: str) -> Any: + if product == "realtime": + return client.realtime.connect(model="test") + if product == "transcription": + return client.realtime.connect(extra_query={"intent": "transcription"}) + if product == "sideband": + return client.realtime.connect(call_id="call_1") + if product == "responses": + return client.responses.connect() + if product == "beta.responses": + return client.beta.responses.connect() + if product == "live": + return client.live.connect() + if product == "live-sideband": + return client.live.sideband.connect(session_id="live_1") + if product == "live-forks": + return client.live.forks.connect(session_id="live_1") + return client.realtime.translations.connect("test") + + +@pytest.fixture(autouse=True) +def no_proxies(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NO_PROXY", "*") + monkeypatch.setenv("no_proxy", "*") + + +@pytest.mark.parametrize("product", PRODUCTS) +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.asyncio +async def test_enter_owns_socket_through_initial_flush( + monkeypatch: pytest.MonkeyPatch, product: str, cancel: bool +) -> None: + sending, release = asyncio.Event(), asyncio.Event() + + class Socket: + closed = False + sent: list[str] = [] + + async def send(self, value: str) -> None: + sending.set() + await release.wait() + self.sent.append(value) + + async def close(self, **_kwargs: Any) -> None: + self.closed = True + + sock = Socket() + + async def connect(*_args: Any, **_kwargs: Any) -> Socket: + return sock + + monkeypatch.setattr("openai.lib._websocket._WebSocketConnect", connect) + async with AsyncOpenAI(api_key="fake-key") as client: + session = manager(client, product) + session.send({"type": "synthetic.first"}) + + async def open_and_close() -> None: + async with session: + pass + + entering = asyncio.create_task(open_and_close()) + try: + await asyncio.wait_for(sending.wait(), timeout=10) + if cancel: + entering.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(entering, timeout=10) + else: + release.set() + await asyncio.wait_for(entering, timeout=10) + assert [json.loads(value)["type"] for value in sock.sent] == ["synthetic.first"] + assert sock.closed + finally: + release.set() + if not entering.done(): + entering.cancel() + await asyncio.gather(entering, return_exceptions=True) + + +@pytest.mark.parametrize("product", PRODUCTS) +@pytest.mark.parametrize("close_fails", [False, True]) +def test_interrupted_sync_entry_closes_socket_and_preserves_error( + monkeypatch: pytest.MonkeyPatch, product: str, close_fails: bool +) -> None: + original = KeyboardInterrupt("initial send interrupted") + + class Socket: + close_called = False + + def send(self, _value: str) -> None: + raise original + + def close(self, **_kwargs: Any) -> None: + self.close_called = True + if close_fails: + raise RuntimeError("cleanup also failed") + + sock = Socket() + + def connect(*_args: Any, **_kwargs: Any) -> Socket: + return sock + + monkeypatch.setattr("websockets.sync.client.connect", connect) + with OpenAI(api_key="fake-key") as client: + session = manager(client, product) + session.send({"type": "synthetic.first"}) + with pytest.raises(KeyboardInterrupt) as raised: + with session: + pytest.fail("entry must not succeed") + assert raised.value is original + assert sock.close_called + + +@pytest.mark.asyncio +async def test_async_constructor_failure_closes_acquired_socket(monkeypatch: pytest.MonkeyPatch) -> None: + import openai.resources.realtime.realtime as module + + original = RuntimeError("cannot initialize connection") + + class Socket: + close_called = False + + async def close(self, **_kwargs: Any) -> None: + self.close_called = True + raise ValueError("cannot close cleanly") + + sock = Socket() + + async def connect(*_args: Any, **_kwargs: Any) -> Socket: + return sock + + def fail_construct(*_args: Any, **_kwargs: Any) -> None: + raise original + + monkeypatch.setattr("openai.lib._websocket._WebSocketConnect", connect) + async with AsyncOpenAI(api_key="fake-key") as client: + session = client.realtime.connect(model="test") + monkeypatch.setattr(module, "AsyncRealtimeConnection", fail_construct) + with pytest.raises(RuntimeError) as raised: + await session.enter() + assert raised.value is original + assert sock.close_called diff --git a/tests/test_example_security.py b/tests/test_example_security.py index e366ad8017..55085150fa 100644 --- a/tests/test_example_security.py +++ b/tests/test_example_security.py @@ -34,7 +34,10 @@ def test_interrupted_example_stops_descendants( worker.write_text( "import subprocess, sys, time\nfrom pathlib import Path\n" "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])\n" - "Path(__file__).with_suffix('.pid').write_text(str(child.pid))\n" + "pid = Path(__file__).with_suffix('.pid')\n" + "temporary = pid.with_suffix('.pid.tmp')\n" + "temporary.write_text(str(child.pid))\n" + "temporary.replace(pid)\n" "time.sleep(60)\n" ) command = [sys.executable, str(worker)] @@ -51,6 +54,7 @@ def wait(process: subprocess.Popen[bytes], timeout: float | None = None) -> int: while not worker.with_suffix(".pid").exists() and time.monotonic() < deadline: time.sleep(0.01) interrupted = True + assert worker.with_suffix(".pid").exists(), "Worker didn't publish its descendant PID" raise KeyboardInterrupt() return original_wait(process, timeout=timeout) diff --git a/tests/test_websocket_connection_options.py b/tests/test_websocket_connection_options.py index 9f95ce9672..0a1f496840 100644 --- a/tests/test_websocket_connection_options.py +++ b/tests/test_websocket_connection_options.py @@ -26,6 +26,7 @@ def test_beta_realtime_import_still_works_with_old_alias() -> None: if TYPE_CHECKING: from openai import OpenAI, AsyncOpenAI + from openai.lib.realtime_translation import TranslationSession, AsyncTranslationSession from openai.types.websocket_connection_options import SyncWebSocketConnectionOptions def native_options_contract( @@ -54,3 +55,45 @@ def beta_realtime_options_contract(client: OpenAI, async_client: AsyncOpenAI) -> client.beta.realtime.connect(model="test", websocket_connection_options={"max_size": 512}) client.beta.realtime.connect(model="test", websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] async_client.beta.realtime.connect(model="test", websocket_connection_options={"write_limit": 65536}) + + def translation_options_contract( + client: OpenAI, async_client: AsyncOpenAI + ) -> tuple[TranslationSession, TranslationSession, AsyncTranslationSession]: + valid_sync = TranslationSession( + client.realtime.translations, model="test", websocket_connection_options={"max_size": None, "max_queue": 8} + ) + invalid_sync = TranslationSession( + client.realtime.translations, + model="test", + websocket_connection_options={"write_limit": 65536}, # pyright: ignore[reportArgumentType] + ) + valid_async = AsyncTranslationSession( + async_client.realtime.translations, model="test", websocket_connection_options={"write_limit": 65536} + ) + return valid_sync, invalid_sync, valid_async + + def synchronous_options_contract(client: OpenAI) -> None: + # Native sync websockets accepts max_queue on version 14 and above. + client.realtime.connect(model="test", websocket_connection_options={"max_queue": 8}) + client.responses.connect(websocket_connection_options={"max_size": None, "compression": None}) + client.live.connect(websocket_connection_options={"extensions": None, "subprotocols": None}) + + # Native sync clients do not implement write_limit on any supported version. + client.realtime.connect(model="test", websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.responses.connect(websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.beta.responses.connect(websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.live.connect(websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.live.sideband.connect(session_id="live_1", websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.live.forks.connect(session_id="live_1", websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + client.realtime.translations.connect("test", websocket_connection_options={"write_limit": 65536}) # pyright: ignore[reportArgumentType] + + def asynchronous_options_contract(client: AsyncOpenAI) -> None: + original: types.WebSocketConnectionOptions = {"write_limit": 65536, "max_queue": 8} + legacy: types.WebsocketConnectionOptions = {"write_limit": (65536, 16384)} + client.realtime.connect(model="test", websocket_connection_options=original) + client.responses.connect(websocket_connection_options=legacy) + client.beta.responses.connect(websocket_connection_options={"write_limit": 65536}) + client.live.connect(websocket_connection_options={"write_limit": 65536}) + client.live.sideband.connect(session_id="live_1", websocket_connection_options={"write_limit": 65536}) + client.live.forks.connect(session_id="live_1", websocket_connection_options={"write_limit": 65536}) + client.realtime.translations.connect("test", websocket_connection_options={"write_limit": 65536})