fix(realtime): end iteration after clean server close (#4461)

Co-authored-by: ayaangazali <ayaangazali.work@gmail.com>

Author: seratch

src/agents/realtime/__init__.py

@@ -62,6 +62,7 @@
     RealtimeModelAudioInterruptedEvent,
     RealtimeModelCachedTokensDetails,
     RealtimeModelConnectionStatusEvent,
+    RealtimeModelEndOfStreamEvent,
     RealtimeModelErrorEvent,
     RealtimeModelEvent,
     RealtimeModelExceptionEvent,
@@ -166,6 +167,7 @@
     "RealtimeModelAudioInterruptedEvent",
     "RealtimeModelCachedTokensDetails",
     "RealtimeModelConnectionStatusEvent",
+    "RealtimeModelEndOfStreamEvent",
     "RealtimeModelErrorEvent",
     "RealtimeModelEvent",
     "RealtimeModelExceptionEvent",

src/agents/realtime/model_events.py

@@ -144,6 +144,13 @@ class RealtimeModelConnectionStatusEvent:
     type: Literal["connection_status"] = "connection_status"
 
 
+@dataclass
+class RealtimeModelEndOfStreamEvent:
+    """The model event stream ended permanently and will emit no further events."""
+
+    type: Literal["end_of_stream"] = "end_of_stream"
+
+
 @dataclass
 class RealtimeModelTurnStartedEvent:
     """Triggered when the model starts generating a response for a turn."""
@@ -248,6 +255,7 @@ class RealtimeModelRawServerEvent:
     | RealtimeModelItemUpdatedEvent
     | RealtimeModelItemDeletedEvent
     | RealtimeModelConnectionStatusEvent
+    | RealtimeModelEndOfStreamEvent
     | RealtimeModelTurnStartedEvent
     | RealtimeModelUsageEvent
     | RealtimeModelTurnEndedEvent

src/agents/realtime/openai_realtime.py

@@ -126,6 +126,8 @@
     RealtimeModelAudioEvent,
     RealtimeModelAudioInterruptedEvent,
     RealtimeModelCachedTokensDetails,
+    RealtimeModelConnectionStatusEvent,
+    RealtimeModelEndOfStreamEvent,
     RealtimeModelErrorEvent,
     RealtimeModelEvent,
     RealtimeModelExceptionEvent,
@@ -552,6 +554,7 @@ def __init__(self, *, transport_config: TransportConfig | None = None) -> None:
         self._websocket: ClientConnection | None = None
         self._websocket_task: asyncio.Task[None] | None = None
         self._connection_attempt_active = False
+        self._close_requested = False
         self._response_create_tasks: set[asyncio.Task[None]] = set()
         self._user_input_lock = asyncio.Lock()
         self._listeners: list[RealtimeModelListener] = []
@@ -639,6 +642,7 @@ async def connect(self, options: RealtimeModelConfig) -> None:
                 headers=headers,
                 transport_config=self._transport_config,
             )
+            self._close_requested = False
             try:
                 self._websocket_task = asyncio.create_task(self._listen_for_messages())
                 await self._update_session_config(model_settings)
@@ -724,6 +728,12 @@ async def _emit_event(self, event: RealtimeModelEvent) -> None:
         for listener in list(self._listeners):
             await listener.on_event(event)
 
+    async def _emit_normal_disconnect(self) -> None:
+        logger.debug("WebSocket connection closed normally")
+        if not self._close_requested:
+            await self._emit_event(RealtimeModelConnectionStatusEvent(status="disconnected"))
+            await self._emit_event(RealtimeModelEndOfStreamEvent())
+
     async def _listen_for_messages(self):
         assert self._websocket is not None, "Not connected"
 
@@ -746,8 +756,7 @@ async def _listen_for_messages(self):
                     )
 
         except websockets.exceptions.ConnectionClosedOK:
-            # Normal connection closure - no exception event needed
-            logger.debug("WebSocket connection closed normally")
+            await self._emit_normal_disconnect()
         except websockets.exceptions.ConnectionClosed as e:
             await self._emit_event(
                 RealtimeModelExceptionEvent(
@@ -760,6 +769,9 @@ async def _listen_for_messages(self):
                     exception=e, context="WebSocket error in message listener"
                 )
             )
+        else:
+            # ClientConnection.__aiter__ consumes ConnectionClosedOK and returns normally.
+            await self._emit_normal_disconnect()
         finally:
             await self._cancel_response_create_tasks()
             await self._release_response_waiters()
@@ -1249,6 +1261,7 @@ async def close(self) -> None:
             cleanup_error: BaseException | None = None
 
             if self._websocket:
+                self._close_requested = True
                 try:
                     await self._websocket.close()
                 except BaseException as exc:

src/agents/realtime/session.py

@@ -80,6 +80,7 @@
 )
 from .model import RealtimeModel, RealtimeModelConfig, RealtimeModelListener
 from .model_events import (
+    RealtimeModelEndOfStreamEvent,
     RealtimeModelEvent,
     RealtimeModelInputAudioTranscriptionCompletedEvent,
     RealtimeModelOutputTextDeltaEvent,
@@ -227,6 +228,7 @@ def __init__(
         self._event_iterator_waiters = 0
         self._closing = False
         self._closed = False
+        self._model_stream_ended = False
         self._cleanup_task: asyncio.Task[None] | None = None
         self._stored_exception: BaseException | None = None
         self._pending_tool_calls: dict[str, _PendingToolCall] = {}
@@ -310,7 +312,7 @@ async def __aexit__(self, _exc_type: Any, _exc_val: Any, _exc_tb: Any) -> None:
     async def __aiter__(self) -> AsyncIterator[RealtimeSessionEvent]:
         """Iterate over events from the session."""
         while True:
-            if self._closed and self._event_queue.empty():
+            if (self._closed or self._model_stream_ended) and self._event_queue.empty():
                 return
 
             # Check if there's a stored exception to raise
@@ -576,6 +578,10 @@ async def on_event(self, event: RealtimeModelEvent) -> None:
             )
         elif event.type == "connection_status":
             pass
+        elif event.type == "end_of_stream":
+            assert isinstance(event, RealtimeModelEndOfStreamEvent)
+            self._model_stream_ended = True
+            self._wake_event_iterators()
         elif event.type == "turn_started":
             is_late_start_for_active_response = (
                 event.response_id is not None

tests/realtime/test_model_events.py

@@ -27,6 +27,13 @@ def test_usage_event_types_are_publicly_exported() -> None:
         assert getattr(realtime, name) is not None
 
 
+def test_end_of_stream_event_is_publicly_exported() -> None:
+    assert "RealtimeModelEndOfStreamEvent" in realtime.__all__
+
+    event = realtime.RealtimeModelEndOfStreamEvent()
+    assert event.type == "end_of_stream"
+
+
 def test_custom_model_can_construct_typed_usage_without_openai_types() -> None:
     event = realtime.RealtimeModelUsageEvent(
         usage=Usage(requests=1, input_tokens=8, output_tokens=5, total_tokens=13),

tests/realtime/test_openai_realtime.py

@@ -14,11 +14,12 @@
 from agents import Agent, WebSearchTool, function_tool
 from agents.exceptions import UserError
 from agents.handoffs import handoff
-from agents.realtime import RealtimeSessionModelSettings
+from agents.realtime import RealtimeAgent, RealtimeSession, RealtimeSessionModelSettings
 from agents.realtime.model import RealtimeModelConfig, RealtimePlaybackTracker
 from agents.realtime.model_events import (
     RealtimeModelAudioEvent,
     RealtimeModelAudioInterruptedEvent,
+    RealtimeModelConnectionStatusEvent,
     RealtimeModelErrorEvent,
     RealtimeModelOutputTextDeltaEvent,
     RealtimeModelRawServerEvent,
@@ -40,6 +41,10 @@
 )
 
 
+async def _collect_session_events(session: RealtimeSession) -> list[Any]:
+    return [event async for event in session]
+
+
 class TestOpenAIRealtimeWebSocketModel:
     """Test suite for OpenAIRealtimeWebSocketModel connection and event handling."""
 
@@ -3261,6 +3266,168 @@ async def handler(websocket):
             await model.close()
             assert model._websocket is None
 
+    @pytest.mark.asyncio
+    async def test_normal_server_close_ends_session_iteration(self):
+        """A clean server close must end session iteration without an exception."""
+
+        async def handler(websocket):
+            await websocket.recv()
+            await websocket.close(code=1000, reason="session ended")
+
+        async with websockets.serve(handler, "127.0.0.1", 0) as server:
+            sockets = list(server.sockets)
+            port = sockets[0].getsockname()[1]
+            session = RealtimeSession(
+                OpenAIRealtimeWebSocketModel(),
+                RealtimeAgent(name="agent"),
+                None,
+                model_config={
+                    "api_key": "test-key",
+                    "url": f"ws://127.0.0.1:{port}/v1/realtime",
+                    "initial_model_settings": {"model_name": "gpt-realtime"},
+                },
+            )
+
+            await session.__aenter__()
+            try:
+                events = await asyncio.wait_for(
+                    _collect_session_events(session),
+                    timeout=1,
+                )
+                assert session._closed is False
+            finally:
+                await session.close()
+
+        disconnects = [
+            event.data
+            for event in events
+            if event.type == "raw_model_event" and event.data.type == "connection_status"
+        ]
+        assert disconnects == [RealtimeModelConnectionStatusEvent(status="disconnected")]
+        assert any(event.type == "history_updated" for event in events)
+
+    @pytest.mark.asyncio
+    async def test_client_close_does_not_emit_server_disconnect(self):
+        """Caller-owned close must not look like a clean server disconnect."""
+
+        connection_count = 0
+
+        async def handler(websocket):
+            nonlocal connection_count
+            connection_count += 1
+            if connection_count == 1:
+                await websocket.wait_closed()
+            else:
+                await websocket.recv()
+                await websocket.close(code=1000, reason="session ended")
+
+        async with websockets.serve(handler, "127.0.0.1", 0) as server:
+            sockets = list(server.sockets)
+            port = sockets[0].getsockname()[1]
+            model = OpenAIRealtimeWebSocketModel()
+            listener = AsyncMock()
+            model.add_listener(listener)
+
+            await model.connect(
+                {
+                    "api_key": "test-key",
+                    "url": f"ws://127.0.0.1:{port}/v1/realtime",
+                    "initial_model_settings": {"model_name": "gpt-realtime"},
+                }
+            )
+            await model.close()
+
+            first_connection_events = [call.args[0] for call in listener.on_event.await_args_list]
+            assert not any(event.type == "connection_status" for event in first_connection_events)
+
+            await model.connect(
+                {
+                    "api_key": "test-key",
+                    "url": f"ws://127.0.0.1:{port}/v1/realtime",
+                    "initial_model_settings": {"model_name": "gpt-realtime"},
+                }
+            )
+            assert model._websocket_task is not None
+            await asyncio.wait_for(model._websocket_task, timeout=1)
+            await model.close()
+
+        emitted_events = [call.args[0] for call in listener.on_event.await_args_list]
+        disconnects = [event for event in emitted_events if event.type == "connection_status"]
+        assert disconnects == [RealtimeModelConnectionStatusEvent(status="disconnected")]
+
+    @pytest.mark.asyncio
+    async def test_cancelled_close_before_websocket_handshake_preserves_server_disconnect(self):
+        """A cancelled preliminary close must not claim transport-close ownership."""
+        allow_server_close = asyncio.Event()
+
+        async def handler(websocket):
+            await websocket.recv()
+            await allow_server_close.wait()
+            await websocket.close(code=1000, reason="session ended")
+
+        async with websockets.serve(handler, "127.0.0.1", 0) as server:
+            sockets = list(server.sockets)
+            port = sockets[0].getsockname()[1]
+            model = OpenAIRealtimeWebSocketModel()
+            listener = AsyncMock()
+            model.add_listener(listener)
+
+            await model.connect(
+                {
+                    "api_key": "test-key",
+                    "url": f"ws://127.0.0.1:{port}/v1/realtime",
+                    "initial_model_settings": {"model_name": "gpt-realtime"},
+                }
+            )
+
+            cancellation_started = asyncio.Event()
+
+            async def wait_until_cancelled():
+                cancellation_started.set()
+                await asyncio.Future()
+
+            with patch.object(model, "_cancel_response_create_tasks", wait_until_cancelled):
+                close_task = asyncio.create_task(model.close())
+                await asyncio.wait_for(cancellation_started.wait(), timeout=1)
+                close_task.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await close_task
+
+            allow_server_close.set()
+            assert model._websocket_task is not None
+            await asyncio.wait_for(model._websocket_task, timeout=1)
+            await model.close()
+
+        emitted_events = [call.args[0] for call in listener.on_event.await_args_list]
+        disconnects = [event for event in emitted_events if event.type == "connection_status"]
+        assert disconnects == [RealtimeModelConnectionStatusEvent(status="disconnected")]
+
+    @pytest.mark.asyncio
+    async def test_abnormal_server_close_still_raises(self):
+        """An abnormal server close must retain the existing exception behavior."""
+
+        async def handler(websocket):
+            await websocket.recv()
+            await websocket.close(code=1011, reason="server failure")
+
+        async with websockets.serve(handler, "127.0.0.1", 0) as server:
+            sockets = list(server.sockets)
+            port = sockets[0].getsockname()[1]
+            session = RealtimeSession(
+                OpenAIRealtimeWebSocketModel(),
+                RealtimeAgent(name="agent"),
+                None,
+                model_config={
+                    "api_key": "test-key",
+                    "url": f"ws://127.0.0.1:{port}/v1/realtime",
+                    "initial_model_settings": {"model_name": "gpt-realtime"},
+                },
+            )
+
+            with pytest.raises(websockets.exceptions.ConnectionClosedError):
+                async with session:
+                    await asyncio.wait_for(_collect_session_events(session), timeout=1)
+
     @pytest.mark.asyncio
     async def test_ping_timeout_success_when_server_responds_quickly(self):
         """Test that connection stays alive when server responds to pings within timeout."""