5353from ..items import ItemHelpers , ModelResponse , TResponseInputItem
5454from ..logger import log_model_action_debug , log_model_action_error , logger
5555from ..model_settings import MCPToolChoice
56- from ..retry import ModelRetryAdvice , ModelRetryAdviceRequest
56+ from ..retry import ModelRetryAdvice , ModelRetryAdviceRequest , ModelRetryNormalizedError
5757from ..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+
447464def _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+
451479def _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 ,
0 commit comments