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

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):

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()))