Skip to content

Commit 265f16f

Browse files
fix(core): retry pre-request WebSocket handshake failures (#4780)
1 parent a34b396 commit 265f16f

3 files changed

Lines changed: 512 additions & 21 deletions

File tree

‎src/agents/models/openai_responses.py‎

Lines changed: 79 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@
5353
from ..items import ItemHelpers, ModelResponse, TResponseInputItem
5454
from ..logger import log_model_action_debug, log_model_action_error, logger
5555
from ..model_settings import MCPToolChoice
56-
from ..retry import ModelRetryAdvice, ModelRetryAdviceRequest
56+
from ..retry import ModelRetryAdvice, ModelRetryAdviceRequest, ModelRetryNormalizedError
5757
from ..tool import (
5858
ApplyPatchTool,
5959
CodeInterpreterTool,
@@ -444,15 +444,41 @@ def _get_wrapped_websocket_replay_safety(error: Exception) -> str | None:
444444
return replay_safety if replay_safety in {"safe", "unsafe"} else None
445445

446446

447+
def _mark_websocket_close_invalidation(error: Exception) -> None:
448+
setattr(error, "_openai_agents_ws_close_invalidated", True) # noqa: B010
449+
450+
451+
def _did_websocket_close_invalidate(error: Exception) -> bool:
452+
return any(
453+
getattr(candidate, "_openai_agents_ws_close_invalidated", False)
454+
for candidate in _iter_retry_error_chain(error)
455+
)
456+
457+
458+
def _websocket_close_invalidation_error() -> RuntimeError:
459+
error = RuntimeError("Responses websocket connection closed while establishing a connection.")
460+
_mark_websocket_close_invalidation(error)
461+
return error
462+
463+
447464
def _did_start_websocket_response(error: Exception) -> bool:
448465
return bool(getattr(error, "_openai_agents_ws_response_started", False))
449466

450467

468+
def _is_websocket_disconnect_error(error: Exception) -> bool:
469+
exc_module = error.__class__.__module__
470+
exc_name = error.__class__.__name__
471+
# websockets reports a peer closing before a valid HTTP upgrade as InvalidMessage. Only an
472+
# InvalidMessage caused by EOFError is transient according to websockets' retry policy.
473+
return exc_module.startswith("websockets") and (
474+
exc_name.startswith("ConnectionClosed")
475+
or (exc_name == "InvalidMessage" and isinstance(error.__cause__, EOFError))
476+
)
477+
478+
451479
def _is_never_sent_websocket_error(error: Exception) -> bool:
452480
for candidate in _iter_retry_error_chain(error):
453-
if candidate.__class__.__module__.startswith(
454-
"websockets"
455-
) and candidate.__class__.__name__.startswith("ConnectionClosed"):
481+
if _is_websocket_disconnect_error(candidate):
456482
if "client closed" not in str(candidate).lower():
457483
return True
458484
return False
@@ -1153,6 +1179,13 @@ def _supports_default_prompt_cache_key(self) -> bool:
11531179
return super()._supports_default_prompt_cache_key()
11541180

11551181
def get_retry_advice(self, request: ModelRetryAdviceRequest) -> ModelRetryAdvice | None:
1182+
if _did_websocket_close_invalidate(request.error):
1183+
return ModelRetryAdvice(
1184+
suggested=False,
1185+
reason=str(request.error),
1186+
normalized=ModelRetryNormalizedError(is_abort=True),
1187+
)
1188+
11561189
stateful_request = bool(request.previous_response_id or request.conversation_id)
11571190
wrapped_replay_safety = _get_wrapped_websocket_replay_safety(request.error)
11581191
if wrapped_replay_safety == "unsafe":
@@ -1343,17 +1376,26 @@ async def _iter_websocket_response_events(
13431376
)
13441377
retry_pre_event_disconnect = _should_retry_pre_event_websocket_disconnect()
13451378
while True:
1346-
connection = await self._await_websocket_with_timeout(
1347-
self._ensure_websocket_connection(
1348-
ws_url, request_headers, connect_timeout=request_timeouts.connect
1349-
),
1350-
request_timeouts.connect,
1351-
"connect",
1352-
)
1379+
connection: Any = None
13531380
received_any_event = False
13541381
yielded_terminal_event = False
13551382
sent_request_frame = False
13561383
try:
1384+
connection = await self._await_websocket_with_timeout(
1385+
self._ensure_websocket_connection(
1386+
ws_url,
1387+
request_headers,
1388+
connect_timeout=request_timeouts.connect,
1389+
request_close_generation=request_close_generation,
1390+
),
1391+
request_timeouts.connect,
1392+
"connect",
1393+
)
1394+
if self._ws_client_close_generation != request_close_generation:
1395+
await self._drop_websocket_connection()
1396+
connection = None
1397+
raise _websocket_close_invalidation_error()
1398+
13571399
# Once we begin awaiting `send()`, treat the request as potentially
13581400
# transmitted to avoid replaying it on send/close races.
13591401
sent_request_frame = True
@@ -1410,11 +1452,15 @@ async def _iter_websocket_response_events(
14101452
is_non_terminal_generator_exit = (
14111453
isinstance(exc, GeneratorExit) and not yielded_terminal_event
14121454
)
1413-
if isinstance(exc, asyncio.CancelledError) or is_non_terminal_generator_exit:
1414-
self._force_abort_websocket_connection(connection)
1415-
self._clear_websocket_connection_state()
1416-
elif not (yielded_terminal_event and isinstance(exc, GeneratorExit)):
1417-
await self._drop_websocket_connection()
1455+
if connection is not None:
1456+
if (
1457+
isinstance(exc, asyncio.CancelledError)
1458+
or is_non_terminal_generator_exit
1459+
):
1460+
self._force_abort_websocket_connection(connection)
1461+
self._clear_websocket_connection_state()
1462+
elif not (yielded_terminal_event and isinstance(exc, GeneratorExit)):
1463+
await self._drop_websocket_connection()
14181464

14191465
if (
14201466
isinstance(exc, Exception)
@@ -1435,10 +1481,12 @@ async def _iter_websocket_response_events(
14351481
is_pre_event_disconnect and not sent_request_frame
14361482
)
14371483
if (
1438-
is_pre_event_disconnect
1484+
isinstance(exc, Exception)
14391485
and self._ws_client_close_generation != request_close_generation
14401486
):
1441-
raise
1487+
_mark_websocket_close_invalidation(exc)
1488+
if is_pre_event_disconnect:
1489+
raise
14421490
if retry_pre_event_disconnect and is_retryable_pre_event_disconnect:
14431491
retry_pre_event_disconnect = False
14441492
continue
@@ -1472,9 +1520,7 @@ def _should_wrap_pre_event_websocket_disconnect(self, exc: Exception) -> bool:
14721520
"Responses websocket connection closed before a terminal response event."
14731521
)
14741522

1475-
exc_module = exc.__class__.__module__
1476-
exc_name = exc.__class__.__name__
1477-
return exc_module.startswith("websockets") and exc_name.startswith("ConnectionClosed")
1523+
return _is_websocket_disconnect_error(exc)
14781524

14791525
def _get_websocket_request_timeouts(self, timeout: Any) -> _WebsocketRequestTimeouts:
14801526
if timeout is None or _is_openai_omitted_value(timeout):
@@ -1593,13 +1639,20 @@ async def _ensure_websocket_connection(
15931639
headers: Mapping[str, str],
15941640
*,
15951641
connect_timeout: float | None,
1642+
request_close_generation: int | None = None,
15961643
) -> Any:
15971644
running_loop = asyncio.get_running_loop()
15981645
identity = (
15991646
ws_url,
16001647
tuple(sorted((str(key).lower(), str(value)) for key, value in headers.items())),
16011648
)
16021649

1650+
if (
1651+
request_close_generation is not None
1652+
and self._ws_client_close_generation != request_close_generation
1653+
):
1654+
raise _websocket_close_invalidation_error()
1655+
16031656
if self._ws_connection is not None and self._ws_connection_identity == identity:
16041657
if (
16051658
self._ws_connection_loop_ref is not None
@@ -1609,6 +1662,11 @@ async def _ensure_websocket_connection(
16091662
return self._ws_connection
16101663
if self._ws_connection is not None:
16111664
await self._drop_websocket_connection()
1665+
if (
1666+
request_close_generation is not None
1667+
and self._ws_client_close_generation != request_close_generation
1668+
):
1669+
raise _websocket_close_invalidation_error()
16121670
self._ws_connection = await self._open_websocket_connection(
16131671
ws_url,
16141672
headers,

‎tests/models/test_model_retry.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
should_disable_provider_managed_retries,
1919
should_disable_websocket_pre_event_retries,
2020
)
21+
from agents.models.openai_responses import OpenAIResponsesWSModel
2122
from agents.retry import (
2223
ModelRetryAdvice,
2324
ModelRetryAdviceRequest,
@@ -1630,6 +1631,82 @@ async def get_response() -> ModelResponse:
16301631
assert calls == 1
16311632

16321633

1634+
def _ws_close_invalidated_error() -> RuntimeError:
1635+
error = RuntimeError("Responses websocket connection closed while establishing a connection.")
1636+
setattr(error, "_openai_agents_ws_close_invalidated", True) # noqa: B010
1637+
return error
1638+
1639+
1640+
@pytest.mark.asyncio
1641+
async def test_get_response_with_retry_does_not_replay_websocket_close_invalidated_request() -> (
1642+
None
1643+
):
1644+
model = OpenAIResponsesWSModel(model="gpt-4", openai_client=cast(Any, object()))
1645+
calls = 0
1646+
1647+
async def get_response() -> ModelResponse:
1648+
nonlocal calls
1649+
calls += 1
1650+
raise _ws_close_invalidated_error()
1651+
1652+
async def rewind() -> None:
1653+
raise AssertionError("A close-invalidated request must not be rewound for retry")
1654+
1655+
with pytest.raises(RuntimeError, match="closed while establishing"):
1656+
await get_response_with_retry(
1657+
get_response=get_response,
1658+
rewind=rewind,
1659+
retry_settings=ModelRetrySettings(
1660+
max_retries=1,
1661+
backoff={"initial_delay": 0},
1662+
policy=retry_policies.network_error(),
1663+
),
1664+
get_retry_advice=model.get_retry_advice,
1665+
previous_response_id=None,
1666+
conversation_id=None,
1667+
)
1668+
1669+
assert calls == 1
1670+
1671+
1672+
@pytest.mark.asyncio
1673+
async def test_stream_response_with_retry_does_not_replay_websocket_close_invalidated_request() -> (
1674+
None
1675+
):
1676+
model = OpenAIResponsesWSModel(model="gpt-4", openai_client=cast(Any, object()))
1677+
calls = 0
1678+
1679+
def get_stream() -> AsyncIterator[TResponseStreamEvent]:
1680+
nonlocal calls
1681+
calls += 1
1682+
1683+
async def iterator() -> AsyncIterator[TResponseStreamEvent]:
1684+
raise _ws_close_invalidated_error()
1685+
yield # pragma: no cover
1686+
1687+
return iterator()
1688+
1689+
async def rewind() -> None:
1690+
raise AssertionError("A close-invalidated request must not be rewound for retry")
1691+
1692+
with pytest.raises(RuntimeError, match="closed while establishing"):
1693+
async for _event in stream_response_with_retry(
1694+
get_stream=get_stream,
1695+
rewind=rewind,
1696+
retry_settings=ModelRetrySettings(
1697+
max_retries=1,
1698+
backoff={"initial_delay": 0},
1699+
policy=retry_policies.network_error(),
1700+
),
1701+
get_retry_advice=model.get_retry_advice,
1702+
previous_response_id=None,
1703+
conversation_id=None,
1704+
):
1705+
pass
1706+
1707+
assert calls == 1
1708+
1709+
16331710
@pytest.mark.asyncio
16341711
async def test_get_response_with_retry_allows_custom_policy_to_override_provider_veto(
16351712
monkeypatch,

0 commit comments

Comments
 (0)