fix(core): retry pre-request WebSocket handshake failures

src/agents/models/openai_responses.py

@@ -448,11 +448,20 @@ 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. Only an
+    # InvalidMessage caused by EOFError is transient according to websockets' retry policy.
+    return exc_module.startswith("websockets") and (
+        exc_name.startswith("ConnectionClosed")
+        or (exc_name == "InvalidMessage" and isinstance(error.__cause__, EOFError))
+    )
+
+
 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 +1352,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 +1421,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 +1487,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):

tests/models/test_openai_responses.py

@@ -244,6 +244,17 @@ class ConnectionClosedError(Exception):
     return ConnectionClosedError(message)
 
 
+def _invalid_message_error(message: str, *, cause: BaseException | None = None) -> Exception:
+    class InvalidMessage(Exception):
+        pass
+
+    InvalidMessage.__module__ = "websockets.exceptions"
+    error = InvalidMessage(message)
+    if cause is not None:
+        error.__cause__ = cause
+    return error
+
+
 @pytest.mark.parametrize("parallel_tool_calls", [True, False, None])
 @pytest.mark.parametrize("tool_source", ["none", "function", "handoff"])
 def test_parallel_tool_calls_follow_converted_responses_tools(
@@ -2925,6 +2936,158 @@ 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()
+
+    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 _invalid_message_error(
+                "did not receive a valid HTTP response",
+                cause=EOFError("connection closed while reading HTTP status line"),
+            )
+        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_malformed_handshake(monkeypatch):
+    client = DummyWSClient()
+    error = _invalid_message_error("malformed HTTP status line")
+    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
+        raise error
+
+    monkeypatch.setattr(model, "_open_websocket_connection", fake_open)
+
+    with pytest.raises(type(error)) as exc_info:
+        await model.get_response(
+            system_instructions=None,
+            input="hi",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=ModelTracing.DISABLED,
+        )
+
+    assert exc_info.value is error
+    assert open_calls == 1
+
+
+@pytest.mark.allow_call_model_methods
+@pytest.mark.asyncio
+async def test_websocket_model_exhausts_one_retry_after_repeated_eof_handshake_failures(
+    monkeypatch,
+):
+    client = DummyWSClient()
+    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
+        raise _invalid_message_error(
+            "did not receive a valid HTTP response",
+            cause=EOFError("connection closed while reading HTTP status line"),
+        )
+
+    monkeypatch.setattr(model, "_open_websocket_connection", fake_open)
+
+    with pytest.raises(RuntimeError, match="before any response events were received"):
+        await model.get_response(
+            system_instructions=None,
+            input="hi",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=ModelTracing.DISABLED,
+        )
+
+    assert open_calls == 2
+    assert model._ws_connection is None
+
+
+@pytest.mark.allow_call_model_methods
+@pytest.mark.asyncio
+async def test_websocket_model_close_during_failing_handshake_prevents_retry(monkeypatch):
+    client = DummyWSClient()
+    model = OpenAIResponsesWSModel(model="gpt-4", openai_client=client)  # type: ignore[arg-type]
+    handshake_started = asyncio.Event()
+    release_handshake = asyncio.Event()
+    error = _invalid_message_error(
+        "did not receive a valid HTTP response",
+        cause=EOFError("connection closed while reading HTTP status line"),
+    )
+    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
+        handshake_started.set()
+        await release_handshake.wait()
+        raise error
+
+    monkeypatch.setattr(model, "_open_websocket_connection", fake_open)
+
+    request_task = asyncio.create_task(
+        model.get_response(
+            system_instructions=None,
+            input="hi",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=ModelTracing.DISABLED,
+        )
+    )
+    await handshake_started.wait()
+    await model.close()
+    release_handshake.set()
+
+    with pytest.raises(type(error)) as exc_info:
+        await request_task
+
+    assert exc_info.value is error
+    assert open_calls == 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 +4377,45 @@ 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()))
+    error = _invalid_message_error(
+        "did not receive a valid HTTP response",
+        cause=EOFError("connection closed while reading HTTP status line"),
+    )
+
+    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_ignores_malformed_handshake() -> None:
+    model = OpenAIResponsesWSModel(model="gpt-4", openai_client=cast(Any, DummyWSClient()))
+    error = _invalid_message_error("malformed HTTP status line")
+
+    advice = model.get_retry_advice(
+        ModelRetryAdviceRequest(
+            error=error,
+            attempt=1,
+            stream=True,
+            previous_response_id="resp_prev",
+        )
+    )
+
+    assert advice is None
+
+
 @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()))