diff --git a/.castiron.stats.yml b/.castiron.stats.yml index 50c57a64d0..3e08e91b71 100644 --- a/.castiron.stats.yml +++ b/.castiron.stats.yml @@ -1,7 +1,7 @@ schema_version: 1 -generation_id: 92191688-2814-4240-a392-408609958bf7 +generation_id: 518e3d93-186c-4c88-86e8-55cf698e7f68 openapi_spec_hash: c3c2cc244c107fdd6e624a7b961015a9 openapi_transformed_spec_hash: 5b121d4b6cca6573c41db323a7399527 config_hash: b707b209c892562880c415693d111ad6 -codegen_sha: 877cf0eac8dabd1578d463ff47f140c493bd4d1a -codegen_hash: 78d280346819510f8ed577aad3d2df4df8dab155df63a48d6ae06e50a98111aa +codegen_sha: fa973427eb9189f3d853d3facf90d0e9378564d1 +codegen_hash: 808925bbb2d1383e24adc2561d78688239e15f45bc83c0f1994fe04f05611980 diff --git a/scripts/castiron/generated-files.json b/scripts/castiron/generated-files.json index d1eb8b59d3..71cbeccd2f 100644 --- a/scripts/castiron/generated-files.json +++ b/scripts/castiron/generated-files.json @@ -379,7 +379,7 @@ }, "src/openai/resources/beta/responses/responses.py": { "mode": "100644", - "sha256": "c5a3b0cb69c3f990dc2ae76db0d2c5fb2df54b8fc312d330cb33d223d1389d26" + "sha256": "29ff4679a33f28cde16c0237daba278183374263a482a251f77eeabe5b010f94" }, "src/openai/resources/beta/threads/__init__.py": { "mode": "100644", @@ -631,7 +631,7 @@ }, "src/openai/resources/responses/responses.py": { "mode": "100644", - "sha256": "86130a834c60c1b7df85118cbc918e84d0a23c122099f1f4e931c1a39b35abfa" + "sha256": "ce72d2035cffd146d598fb548ed36ff17e47cf758ea1faa784bc245095b8a2d8" }, "src/openai/resources/safety/__init__.py": { "mode": "100644", diff --git a/src/openai/resources/beta/responses/responses.py b/src/openai/resources/beta/responses/responses.py index 64016c7ae6..a2a96b22fc 100644 --- a/src/openai/resources/beta/responses/responses.py +++ b/src/openai/resources/beta/responses/responses.py @@ -16,7 +16,7 @@ from ...._utils import is_given, path_template, maybe_transform, strip_not_given, async_maybe_transform from ...._compat import cached_property from ...._httpx2 import normalize_httpx_url -from ...._models import FinalRequestOptions, construct_type_unchecked +from ...._models import construct_type_unchecked from .input_items import ( InputItems, AsyncInputItems, @@ -42,11 +42,21 @@ response_retrieve_params, beta_responses_client_event_param, ) -from ...._exceptions import OpenAIError, WebSocketConnectionClosedError +from ...._exceptions import OpenAIError from ...._send_queue import SendQueue -from ...._base_client import _merge_mappings, make_request_options +from ...._base_client import make_request_options from ...._event_handler import EventHandlerRegistry -from ...._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase +from ...._websocket_connection import ( + WebSocketConnectionBase, + AsyncWebSocketConnectionBase, + _exit_websocket, + _connect_websocket, + _async_exit_websocket, + _websocket_entry_guard, + _async_connect_websocket, + _record_websocket_progress, + _async_websocket_entry_guard, +) from ....types.beta.beta_response import BetaResponse from ....types.beta.beta_tool_param import BetaToolParam from ....types.beta.beta_service_tier import BetaServiceTier @@ -4897,34 +4907,12 @@ def __init__( self.response = AsyncResponsesResponseResource(self) - async def __aiter__(self) -> AsyncIterator[BetaResponsesServerEvent]: + def __aiter__(self) -> AsyncIterator[BetaResponsesServerEvent]: """ An infinite-iterator that will continue to yield events until the connection is closed. """ - from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError - - while True: - try: - yield await self.recv() - except ConnectionClosed as exc: - # Older SDK bases read close codes only from ConnectionClosedError. - # Adapt the policy input; keep the original close for EOF and errors. - reconnect_exc = ( - ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent) - if isinstance(exc, ConnectionClosedOK) - else exc - ) - if not await self._reconnect(reconnect_exc): - unsent = self._send_queue.drain() - if unsent: - raise WebSocketConnectionClosedError( - "WebSocket connection closed with unsent messages", - unsent_messages=unsent, - ) from exc - if isinstance(exc, ConnectionClosedOK): - return - raise + return self._iter_events() async def recv(self) -> BetaResponsesServerEvent: """ @@ -4933,15 +4921,7 @@ async def recv(self) -> BetaResponsesServerEvent: Canceling this method is safe. There's no risk of losing data. """ event = self.parse_event(await self.recv_bytes()) - event_type = ( - cast("dict[str, object]", event).get("type") - if isinstance(cast(object, event), dict) - else getattr(event, "type", None) - ) - # A successful upgrade can still be followed by an admission error. - # Reset the budget only after receiving a non-error application event. - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, event) return event async def recv_bytes(self) -> bytes: @@ -4954,15 +4934,7 @@ async def recv_bytes(self) -> bytes: """ message = await self._connection.recv(decode=False) log.debug("Received WebSocket message: %i bytes", len(message)) - if self._reconnect_attempt: - # Account for raw application progress without changing frame delivery. - try: - event_data: object = json.loads(message) - except (ValueError, RecursionError): - return message - event_type = cast("dict[str, object]", event_data).get("type") if isinstance(event_data, dict) else None - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, message, raw=True) return message async def send(self, event: BetaResponsesClientEvent | BetaResponsesClientEventParam) -> None: @@ -5175,7 +5147,7 @@ async def __aenter__(self) -> AsyncResponsesConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - try: + async with _async_websocket_entry_guard(ws, self.__clear_connection): self.__connection = AsyncResponsesConnection( ws, make_ws=self._connect_ws if self.__on_reconnecting is not None else None, @@ -5190,74 +5162,23 @@ async def __aenter__(self) -> AsyncResponsesConnection: 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 + def __clear_connection(self) -> None: + self.__connection = None + enter = __aenter__ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> AsyncWebSocketConnection: - try: - from ....lib._websocket import _WebSocketConnect as connect - except ImportError as exc: - raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc - - url = self._prepare_url() - try: - # Older runtimes do not declare this optional helper. - query_params = getattr(self.__client, "_websocket_query_params") # noqa: B009 - except AttributeError: - query_params = self.__client._websocket_query_params_fallback - url = url.copy_with( - params=query_params( - url, - extra_query, - ), - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - options = await self.__client._prepare_options( - FinalRequestOptions.construct( - method="get", - url=str(url), - headers=dict(extra_headers), - security={"bearer_auth": True}, - ) - ) - url = self.__client._prepare_url(options.url).copy_merge_params( - self.__client.qs.stringify(cast(Any, options.params)) - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - headers = { - key.lower(): (key, value) - for header_set in ( - self.__client.auth_headers, - {}, - self.__client.default_headers, - options.headers if is_given(options.headers) else {}, - ) - for key, value in header_set.items() - } - log.debug("Connecting to WebSocket API") - 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()), {}), - **websocket_options, + return await _async_connect_websocket( + self.__client, + self._prepare_url, + extra_query, + extra_headers, + self.__websocket_connection_options, + {}, + log, ) def _prepare_url(self) -> httpx2.URL: @@ -5275,12 +5196,7 @@ def _prepare_url(self) -> httpx2.URL: async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - if self.__connection is not None: - try: - await self.__connection.close() - except Exception: - if exc is None: - raise + await _async_exit_websocket(self.__connection, exc) class ResponsesConnection(WebSocketConnectionBase): @@ -5324,29 +5240,7 @@ def __iter__(self) -> Iterator[BetaResponsesServerEvent]: An infinite-iterator that will continue to yield events until the connection is closed. """ - from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError - - while True: - try: - yield self.recv() - except ConnectionClosed as exc: - # Older SDK bases read close codes only from ConnectionClosedError. - # Adapt the policy input; keep the original close for EOF and errors. - reconnect_exc = ( - ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent) - if isinstance(exc, ConnectionClosedOK) - else exc - ) - if not self._reconnect(reconnect_exc): - unsent = self._send_queue.drain() - if unsent: - raise WebSocketConnectionClosedError( - "WebSocket connection closed with unsent messages", - unsent_messages=unsent, - ) from exc - if isinstance(exc, ConnectionClosedOK): - return - raise + return self._iter_events() def recv(self) -> BetaResponsesServerEvent: """ @@ -5355,15 +5249,7 @@ def recv(self) -> BetaResponsesServerEvent: Canceling this method is safe. There's no risk of losing data. """ event = self.parse_event(self.recv_bytes()) - event_type = ( - cast("dict[str, object]", event).get("type") - if isinstance(cast(object, event), dict) - else getattr(event, "type", None) - ) - # A successful upgrade can still be followed by an admission error. - # Reset the budget only after receiving a non-error application event. - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, event) return event def recv_bytes(self) -> bytes: @@ -5378,15 +5264,7 @@ def recv_bytes(self) -> bytes: message = _recv_bytes(self._connection) log.debug("Received WebSocket message: %i bytes", len(message)) - if self._reconnect_attempt: - # Account for raw application progress without changing frame delivery. - try: - event_data: object = json.loads(message) - except (ValueError, RecursionError): - return message - event_type = cast("dict[str, object]", event_data).get("type") if isinstance(event_data, dict) else None - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, message, raw=True) return message def send(self, event: BetaResponsesClientEvent | BetaResponsesClientEventParam) -> None: @@ -5593,7 +5471,7 @@ def __enter__(self) -> ResponsesConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - try: + with _websocket_entry_guard(ws, self.__clear_connection): self.__connection = ResponsesConnection( ws, make_ws=self._connect_ws if self.__on_reconnecting is not None else None, @@ -5608,74 +5486,23 @@ def __enter__(self) -> ResponsesConnection: 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 + def __clear_connection(self) -> None: + self.__connection = None + enter = __enter__ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketConnection: - try: - from websockets.sync.client import connect - except ImportError as exc: - raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc - - url = self._prepare_url() - try: - # Older runtimes do not declare this optional helper. - query_params = getattr(self.__client, "_websocket_query_params") # noqa: B009 - except AttributeError: - query_params = self.__client._websocket_query_params_fallback - url = url.copy_with( - params=query_params( - url, - extra_query, - ), - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - options = self.__client._prepare_options( - FinalRequestOptions.construct( - method="get", - url=str(url), - headers=dict(extra_headers), - security={"bearer_auth": True}, - ) - ) - url = self.__client._prepare_url(options.url).copy_merge_params( - self.__client.qs.stringify(cast(Any, options.params)) - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - headers = { - key.lower(): (key, value) - for header_set in ( - self.__client.auth_headers, - {}, - self.__client.default_headers, - options.headers if is_given(options.headers) else {}, - ) - for key, value in header_set.items() - } - log.debug("Connecting to WebSocket API") - 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()), {}), - **websocket_options, + return _connect_websocket( + self.__client, + self._prepare_url, + extra_query, + extra_headers, + self.__websocket_connection_options, + {}, + log, ) def _prepare_url(self) -> httpx2.URL: @@ -5693,12 +5520,7 @@ def _prepare_url(self) -> httpx2.URL: def __exit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - if self.__connection is not None: - try: - self.__connection.close() - except Exception: - if exc is None: - raise + _exit_websocket(self.__connection, exc) class BaseResponsesConnectionResource: diff --git a/src/openai/resources/responses/responses.py b/src/openai/resources/responses/responses.py index 854e429823..53fd2f7034 100644 --- a/src/openai/resources/responses/responses.py +++ b/src/openai/resources/responses/responses.py @@ -13,10 +13,10 @@ from ... import _legacy_response from ..._types import Body, Omit, Query, Headers, NoneType, NotGiven, omit, not_given -from ..._utils import is_given, path_template, maybe_transform, strip_not_given, async_maybe_transform +from ..._utils import path_template, maybe_transform, strip_not_given, async_maybe_transform from ..._compat import cached_property from ..._httpx2 import normalize_httpx_url -from ..._models import FinalRequestOptions, construct_type_unchecked +from ..._models import construct_type_unchecked from ..._resource import SyncAPIResource, AsyncAPIResource from ..._response import to_streamed_response_wrapper, async_to_streamed_response_wrapper from .input_items import ( @@ -36,9 +36,9 @@ InputTokensWithStreamingResponse, AsyncInputTokensWithStreamingResponse, ) -from ..._exceptions import OpenAIError, WebSocketConnectionClosedError +from ..._exceptions import OpenAIError from ..._send_queue import SendQueue -from ..._base_client import _merge_mappings, make_request_options +from ..._base_client import make_request_options from ..._event_handler import EventHandlerRegistry from ...types.responses import ( response_create_params, @@ -46,7 +46,17 @@ response_retrieve_params, responses_client_event_param, ) -from ..._websocket_connection import WebSocketConnectionBase, AsyncWebSocketConnectionBase +from ..._websocket_connection import ( + WebSocketConnectionBase, + AsyncWebSocketConnectionBase, + _exit_websocket, + _connect_websocket, + _async_exit_websocket, + _websocket_entry_guard, + _async_connect_websocket, + _record_websocket_progress, + _async_websocket_entry_guard, +) from ...lib._parsing._responses import TextFormatT from ...lib._response_resources import parse, stream, async_parse, async_stream from ...types.responses.response import Response @@ -4086,34 +4096,12 @@ def __init__( self.response = AsyncResponsesResponseResource(self) - async def __aiter__(self) -> AsyncIterator[ResponsesServerEvent]: + def __aiter__(self) -> AsyncIterator[ResponsesServerEvent]: """ An infinite-iterator that will continue to yield events until the connection is closed. """ - from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError - - while True: - try: - yield await self.recv() - except ConnectionClosed as exc: - # Older SDK bases read close codes only from ConnectionClosedError. - # Adapt the policy input; keep the original close for EOF and errors. - reconnect_exc = ( - ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent) - if isinstance(exc, ConnectionClosedOK) - else exc - ) - if not await self._reconnect(reconnect_exc): - unsent = self._send_queue.drain() - if unsent: - raise WebSocketConnectionClosedError( - "WebSocket connection closed with unsent messages", - unsent_messages=unsent, - ) from exc - if isinstance(exc, ConnectionClosedOK): - return - raise + return self._iter_events() async def recv(self) -> ResponsesServerEvent: """ @@ -4122,15 +4110,7 @@ async def recv(self) -> ResponsesServerEvent: Canceling this method is safe. There's no risk of losing data. """ event = self.parse_event(await self.recv_bytes()) - event_type = ( - cast("dict[str, object]", event).get("type") - if isinstance(cast(object, event), dict) - else getattr(event, "type", None) - ) - # A successful upgrade can still be followed by an admission error. - # Reset the budget only after receiving a non-error application event. - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, event) return event async def recv_bytes(self) -> bytes: @@ -4143,15 +4123,7 @@ async def recv_bytes(self) -> bytes: """ message = await self._connection.recv(decode=False) log.debug("Received WebSocket message: %i bytes", len(message)) - if self._reconnect_attempt: - # Account for raw application progress without changing frame delivery. - try: - event_data: object = json.loads(message) - except (ValueError, RecursionError): - return message - event_type = cast("dict[str, object]", event_data).get("type") if isinstance(event_data, dict) else None - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, message, raw=True) return message async def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: @@ -4364,7 +4336,7 @@ async def __aenter__(self) -> AsyncResponsesConnection: """ ws = await self._connect_ws(self.__extra_query, self.__extra_headers) - try: + async with _async_websocket_entry_guard(ws, self.__clear_connection): self.__connection = AsyncResponsesConnection( ws, make_ws=self._connect_ws if self.__on_reconnecting is not None else None, @@ -4379,74 +4351,23 @@ async def __aenter__(self) -> AsyncResponsesConnection: 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 + def __clear_connection(self) -> None: + self.__connection = None + enter = __aenter__ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> AsyncWebSocketConnection: - try: - from ...lib._websocket import _WebSocketConnect as connect - except ImportError as exc: - raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc - - url = self._prepare_url() - try: - # Older runtimes do not declare this optional helper. - query_params = getattr(self.__client, "_websocket_query_params") # noqa: B009 - except AttributeError: - query_params = self.__client._websocket_query_params_fallback - url = url.copy_with( - params=query_params( - url, - extra_query, - ), - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - options = await self.__client._prepare_options( - FinalRequestOptions.construct( - method="get", - url=str(url), - headers=dict(extra_headers), - security={"bearer_auth": True}, - ) - ) - url = self.__client._prepare_url(options.url).copy_merge_params( - self.__client.qs.stringify(cast(Any, options.params)) - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - headers = { - key.lower(): (key, value) - for header_set in ( - self.__client.auth_headers, - {}, - self.__client.default_headers, - options.headers if is_given(options.headers) else {}, - ) - for key, value in header_set.items() - } - log.debug("Connecting to WebSocket API") - 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()), {}), - **websocket_options, + return await _async_connect_websocket( + self.__client, + self._prepare_url, + extra_query, + extra_headers, + self.__websocket_connection_options, + {}, + log, ) def _prepare_url(self) -> httpx2.URL: @@ -4464,12 +4385,7 @@ def _prepare_url(self) -> httpx2.URL: async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - if self.__connection is not None: - try: - await self.__connection.close() - except Exception: - if exc is None: - raise + await _async_exit_websocket(self.__connection, exc) class ResponsesConnection(WebSocketConnectionBase): @@ -4514,29 +4430,7 @@ def __iter__(self) -> Iterator[ResponsesServerEvent]: An infinite-iterator that will continue to yield events until the connection is closed. """ - from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError - - while True: - try: - yield self.recv() - except ConnectionClosed as exc: - # Older SDK bases read close codes only from ConnectionClosedError. - # Adapt the policy input; keep the original close for EOF and errors. - reconnect_exc = ( - ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent) - if isinstance(exc, ConnectionClosedOK) - else exc - ) - if not self._reconnect(reconnect_exc): - unsent = self._send_queue.drain() - if unsent: - raise WebSocketConnectionClosedError( - "WebSocket connection closed with unsent messages", - unsent_messages=unsent, - ) from exc - if isinstance(exc, ConnectionClosedOK): - return - raise + return self._iter_events() def recv(self) -> ResponsesServerEvent: """ @@ -4545,15 +4439,7 @@ def recv(self) -> ResponsesServerEvent: Canceling this method is safe. There's no risk of losing data. """ event = self.parse_event(self.recv_bytes()) - event_type = ( - cast("dict[str, object]", event).get("type") - if isinstance(cast(object, event), dict) - else getattr(event, "type", None) - ) - # A successful upgrade can still be followed by an admission error. - # Reset the budget only after receiving a non-error application event. - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, event) return event def recv_bytes(self) -> bytes: @@ -4568,15 +4454,7 @@ def recv_bytes(self) -> bytes: message = _recv_bytes(self._connection) log.debug("Received WebSocket message: %i bytes", len(message)) - if self._reconnect_attempt: - # Account for raw application progress without changing frame delivery. - try: - event_data: object = json.loads(message) - except (ValueError, RecursionError): - return message - event_type = cast("dict[str, object]", event_data).get("type") if isinstance(event_data, dict) else None - if isinstance(event_type, str) and event_type and event_type != "error": - self._reconnect_attempt = 0 + _record_websocket_progress(self, message, raw=True) return message def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: @@ -4783,7 +4661,7 @@ def __enter__(self) -> ResponsesConnection: """ ws = self._connect_ws(self.__extra_query, self.__extra_headers) - try: + with _websocket_entry_guard(ws, self.__clear_connection): self.__connection = ResponsesConnection( ws, make_ws=self._connect_ws if self.__on_reconnecting is not None else None, @@ -4798,74 +4676,23 @@ def __enter__(self) -> ResponsesConnection: 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 + def __clear_connection(self) -> None: + self.__connection = None + enter = __enter__ def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> WebSocketConnection: - try: - from websockets.sync.client import connect - except ImportError as exc: - raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc - - url = self._prepare_url() - try: - # Older runtimes do not declare this optional helper. - query_params = getattr(self.__client, "_websocket_query_params") # noqa: B009 - except AttributeError: - query_params = self.__client._websocket_query_params_fallback - url = url.copy_with( - params=query_params( - url, - extra_query, - ), - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - options = self.__client._prepare_options( - FinalRequestOptions.construct( - method="get", - url=str(url), - headers=dict(extra_headers), - security={"bearer_auth": True}, - ) - ) - url = self.__client._prepare_url(options.url).copy_merge_params( - self.__client.qs.stringify(cast(Any, options.params)) - ) - url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme)) - headers = { - key.lower(): (key, value) - for header_set in ( - self.__client.auth_headers, - {}, - self.__client.default_headers, - options.headers if is_given(options.headers) else {}, - ) - for key, value in header_set.items() - } - log.debug("Connecting to WebSocket API") - 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()), {}), - **websocket_options, + return _connect_websocket( + self.__client, + self._prepare_url, + extra_query, + extra_headers, + self.__websocket_connection_options, + {}, + log, ) def _prepare_url(self) -> httpx2.URL: @@ -4883,12 +4710,7 @@ def _prepare_url(self) -> httpx2.URL: def __exit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None ) -> None: - if self.__connection is not None: - try: - self.__connection.close() - except Exception: - if exc is None: - raise + _exit_websocket(self.__connection, exc) class BaseResponsesConnectionResource: diff --git a/tests/lib/responses/test_websocket_auth.py b/tests/lib/responses/test_websocket_auth.py index 9bba47455d..559111d23d 100644 --- a/tests/lib/responses/test_websocket_auth.py +++ b/tests/lib/responses/test_websocket_auth.py @@ -114,7 +114,8 @@ def reconnect(_event: ReconnectingEvent) -> None: @pytest.mark.parametrize("mode", ["sync", "async"]) -async def test_handshake_honors_prepared_url_query_and_headers(mode: str) -> None: +@pytest.mark.parametrize("beta", [False, True], ids=["stable", "beta"]) +async def test_handshake_honors_prepared_url_query_and_headers(mode: str, beta: bool) -> None: def script(socket: ServerConnection) -> None: assert socket.request is not None assert socket.request.path == "/prepared?before=kept&prepared=yes" @@ -151,7 +152,7 @@ async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOp default_headers={"X-Removed": "client value"}, http_client=httpx2.Client(trust_env=False), ) as client: - with client.responses.connect() as connection: + with (client.beta.responses if beta else client.responses).connect() as connection: assert connection.recv().type == "response.future" else: async with AsyncPreparedClient( @@ -160,7 +161,9 @@ async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOp default_headers={"X-Removed": "client value"}, http_client=httpx2.AsyncClient(trust_env=False), ) as async_client: - async with async_client.responses.connect() as async_connection: + async with ( + async_client.beta.responses if beta else async_client.responses + ).connect() as async_connection: assert (await async_connection.recv()).type == "response.future" diff --git a/tests/lib/test_websocket_entry_cleanup.py b/tests/lib/test_websocket_entry_cleanup.py index c0ff2e0ba9..b9f44454f7 100644 --- a/tests/lib/test_websocket_entry_cleanup.py +++ b/tests/lib/test_websocket_entry_cleanup.py @@ -2,6 +2,7 @@ import json import asyncio +import importlib from typing import Any import pytest @@ -57,6 +58,7 @@ async def test_enter_owns_socket_through_initial_flush( class Socket: closed = False + close_count = 0 sent: list[str] = [] async def send(self, value: str) -> None: @@ -66,6 +68,10 @@ async def send(self, value: str) -> None: async def close(self, **_kwargs: Any) -> None: self.closed = True + self.close_count += 1 + if cancel and product in {"responses", "beta.responses"} and self.close_count == 1: + # Reentrant exit must not close a failed entry's socket again. + await session.__aexit__(None, None, None) sock = Socket() @@ -93,6 +99,8 @@ async def open_and_close() -> None: await asyncio.wait_for(entering, timeout=10) assert [json.loads(value)["type"] for value in sock.sent] == ["synthetic.first"] assert sock.closed + if cancel and product in {"responses", "beta.responses"}: + assert sock.close_count == 1 finally: release.set() if not entering.done(): @@ -109,12 +117,17 @@ def test_interrupted_sync_entry_closes_socket_and_preserves_error( class Socket: close_called = False + close_count = 0 def send(self, _value: str) -> None: raise original def close(self, **_kwargs: Any) -> None: self.close_called = True + self.close_count += 1 + if product in {"responses", "beta.responses"} and self.close_count == 1: + # Reentrant exit must not close a failed entry's socket again. + session.__exit__(None, None, None) if close_fails: raise RuntimeError("cleanup also failed") @@ -132,11 +145,14 @@ def connect(*_args: Any, **_kwargs: Any) -> Socket: pytest.fail("entry must not succeed") assert raised.value is original assert sock.close_called + if product in {"responses", "beta.responses"}: + assert sock.close_count == 1 @pytest.mark.asyncio -async def test_async_constructor_failure_closes_acquired_socket(monkeypatch: pytest.MonkeyPatch) -> None: - import openai.resources.realtime.realtime as module +@pytest.mark.parametrize("product", ["realtime", "responses", "beta.responses"]) +async def test_async_constructor_failure_closes_acquired_socket(monkeypatch: pytest.MonkeyPatch, product: str) -> None: + module = importlib.import_module(f"openai.resources.{product}.{product.rsplit('.', 1)[-1]}") original = RuntimeError("cannot initialize connection") @@ -157,8 +173,9 @@ def fail_construct(*_args: Any, **_kwargs: Any) -> None: 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) + session = manager(client, product) + connection_type = "AsyncRealtimeConnection" if product == "realtime" else "AsyncResponsesConnection" + monkeypatch.setattr(module, connection_type, fail_construct) with pytest.raises(RuntimeError) as raised: await session.enter() assert raised.value is original diff --git a/tests/test_websocket_recovery_regressions.py b/tests/test_websocket_recovery_regressions.py index def2b21edf..888e299316 100644 --- a/tests/test_websocket_recovery_regressions.py +++ b/tests/test_websocket_recovery_regressions.py @@ -159,6 +159,49 @@ async def test_received_close_policy_on_async_iterator( assert [e.type for e in events] == (["error"] if expected else []) +@pytest.mark.parametrize( + "connection_type,mode", + [ + (ResponsesConnection, "sync"), + (BetaResponsesConnection, "sync"), + (AsyncResponsesConnection, "async"), + (AsyncBetaResponsesConnection, "async"), + ], + ids=["stable-sync", "beta-sync", "stable-async", "beta-async"], +) +@pytest.mark.asyncio +async def test_responses_public_iterator_keeps_recv_override_and_generator_controls( + monkeypatch: pytest.MonkeyPatch, connection_type: type[Any], mode: str +) -> None: + connection = connection_type(MagicMock()) + monkeypatch.setattr( + connection, "recv", AsyncMock(return_value="first") if mode == "async" else MagicMock(return_value="first") + ) + events: Any = aiter(connection) if mode == "async" else iter(connection) + assert (await anext(events) if mode == "async" else next(events)) == "first" + monkeypatch.setattr( + connection, "recv", AsyncMock(return_value="second") if mode == "async" else MagicMock(return_value="second") + ) + assert (await anext(events) if mode == "async" else next(events)) == "second" + error = asyncio.CancelledError("caller") if mode == "async" else ValueError("caller") + with pytest.raises(type(error)) as caught: + if mode == "async": + await events.athrow(error) + else: + events.throw(error) + assert caught.value is error + closing_events: Any = aiter(connection) if mode == "async" else iter(connection) + assert (await anext(closing_events) if mode == "async" else next(closing_events)) == "second" + if mode == "async": + await closing_events.aclose() + with pytest.raises(StopAsyncIteration): + await anext(closing_events) + else: + closing_events.close() + with pytest.raises(StopIteration): + next(closing_events) + + def register_handlers(conn: Any, mode: str, seen: list[str]) -> None: def handler(event: Any) -> None: seen.append(event.type)