From d130954867db7c5f1646c8cf3e1d2698c579b402 Mon Sep 17 00:00:00 2001 From: CoralGarden52 <2193436736@qq.com> Date: Sun, 30 Aug 2026 09:49:14 +0000 Subject: [PATCH] fix(models): retry pre-request websocket handshake failures --- src/agents/models/openai_responses.py | 47 ++++++++++++-------- tests/models/test_openai_responses.py | 64 +++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 18 deletions(-) diff --git a/src/agents/models/openai_responses.py b/src/agents/models/openai_responses.py index b9c6c75013..9b20302299 100644 --- a/src/agents/models/openai_responses.py +++ b/src/agents/models/openai_responses.py @@ -448,11 +448,18 @@ def _did_start_websocket_response(error: Exception) -> bool: return bool(getattr(error, "_openai_agents_ws_response_started", False)) +def _is_websocket_disconnect_error(error: Exception) -> bool: + exc_module = error.__class__.__module__ + exc_name = error.__class__.__name__ + # websockets reports a peer closing before a valid HTTP upgrade as InvalidMessage. + return exc_module.startswith("websockets") and ( + exc_name.startswith("ConnectionClosed") or exc_name == "InvalidMessage" + ) + + def _is_never_sent_websocket_error(error: Exception) -> bool: for candidate in _iter_retry_error_chain(error): - if candidate.__class__.__module__.startswith( - "websockets" - ) and candidate.__class__.__name__.startswith("ConnectionClosed"): + if _is_websocket_disconnect_error(candidate): if "client closed" not in str(candidate).lower(): return True return False @@ -1343,17 +1350,19 @@ async def _iter_websocket_response_events( ) retry_pre_event_disconnect = _should_retry_pre_event_websocket_disconnect() while True: - connection = await self._await_websocket_with_timeout( - self._ensure_websocket_connection( - ws_url, request_headers, connect_timeout=request_timeouts.connect - ), - request_timeouts.connect, - "connect", - ) + connection: Any = None received_any_event = False yielded_terminal_event = False sent_request_frame = False try: + connection = await self._await_websocket_with_timeout( + self._ensure_websocket_connection( + ws_url, request_headers, connect_timeout=request_timeouts.connect + ), + request_timeouts.connect, + "connect", + ) + # Once we begin awaiting `send()`, treat the request as potentially # transmitted to avoid replaying it on send/close races. sent_request_frame = True @@ -1410,11 +1419,15 @@ async def _iter_websocket_response_events( is_non_terminal_generator_exit = ( isinstance(exc, GeneratorExit) and not yielded_terminal_event ) - if isinstance(exc, asyncio.CancelledError) or is_non_terminal_generator_exit: - self._force_abort_websocket_connection(connection) - self._clear_websocket_connection_state() - elif not (yielded_terminal_event and isinstance(exc, GeneratorExit)): - await self._drop_websocket_connection() + if connection is not None: + if ( + isinstance(exc, asyncio.CancelledError) + or is_non_terminal_generator_exit + ): + self._force_abort_websocket_connection(connection) + self._clear_websocket_connection_state() + elif not (yielded_terminal_event and isinstance(exc, GeneratorExit)): + await self._drop_websocket_connection() if ( isinstance(exc, Exception) @@ -1472,9 +1485,7 @@ def _should_wrap_pre_event_websocket_disconnect(self, exc: Exception) -> bool: "Responses websocket connection closed before a terminal response event." ) - exc_module = exc.__class__.__module__ - exc_name = exc.__class__.__name__ - return exc_module.startswith("websockets") and exc_name.startswith("ConnectionClosed") + return _is_websocket_disconnect_error(exc) def _get_websocket_request_timeouts(self, timeout: Any) -> _WebsocketRequestTimeouts: if timeout is None or _is_openai_omitted_value(timeout): diff --git a/tests/models/test_openai_responses.py b/tests/models/test_openai_responses.py index d6b787df9b..7c9602767b 100644 --- a/tests/models/test_openai_responses.py +++ b/tests/models/test_openai_responses.py @@ -2925,6 +2925,46 @@ async def fake_open( assert model._ws_connection is ws2 +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_websocket_model_retries_if_handshake_fails_before_request(monkeypatch): + client = DummyWSClient() + + class InvalidMessage(Exception): + pass + + InvalidMessage.__module__ = "websockets.exceptions" + + ws = DummyWSConnection([_response_completed_frame("resp-retried", 1)]) + model = OpenAIResponsesWSModel(model="gpt-4", openai_client=client) # type: ignore[arg-type] + open_calls = 0 + + async def fake_open( + ws_url: str, headers: dict[str, str], *, connect_timeout: float | None = None + ) -> DummyWSConnection: + nonlocal open_calls + open_calls += 1 + if open_calls == 1: + raise InvalidMessage("did not receive a valid HTTP response") + return ws + + monkeypatch.setattr(model, "_open_websocket_connection", fake_open) + + response = await model.get_response( + system_instructions=None, + input="hi", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + ) + + assert response.response_id == "resp-retried" + assert open_calls == 2 + assert len(ws.sent_messages) == 1 + + @pytest.mark.allow_call_model_methods @pytest.mark.asyncio async def test_websocket_model_does_not_retry_if_send_raises_after_writing_on_reused_connection( @@ -4214,6 +4254,30 @@ def test_websocket_get_retry_advice_marks_connect_timeout_replay_safe() -> None: assert advice.replay_safety == "safe" +@pytest.mark.allow_call_model_methods +def test_websocket_get_retry_advice_marks_handshake_failure_replay_safe() -> None: + model = OpenAIResponsesWSModel(model="gpt-4", openai_client=cast(Any, DummyWSClient())) + + class InvalidMessage(Exception): + pass + + InvalidMessage.__module__ = "websockets.exceptions" + error = InvalidMessage("did not receive a valid HTTP response") + + advice = model.get_retry_advice( + ModelRetryAdviceRequest( + error=error, + attempt=1, + stream=True, + previous_response_id="resp_prev", + ) + ) + + assert advice is not None + assert advice.suggested is True + assert advice.replay_safety == "safe" + + @pytest.mark.allow_call_model_methods def test_websocket_get_retry_advice_marks_request_lock_timeout_replay_safe() -> None: model = OpenAIResponsesWSModel(model="gpt-4", openai_client=cast(Any, DummyWSClient()))