fix(realtime): bound background task cleanup

src/agents/realtime/session.py

@@ -82,6 +82,7 @@ class _RealtimeSessionClosedSentinel:
 
 
 _REALTIME_SESSION_CLOSED_SENTINEL = _RealtimeSessionClosedSentinel()
+_BACKGROUND_TASK_CLEANUP_TIMEOUT = 1.0
 
 
 def _serialize_tool_output(output: Any) -> str:
@@ -192,6 +193,7 @@ def __init__(
             asyncio.Queue()
         )
         self._event_iterator_waiters = 0
+        self._closing = False
         self._closed = False
         self._stored_exception: BaseException | None = None
         self._pending_tool_calls: dict[str, _PendingToolCall] = {}
@@ -1265,6 +1267,8 @@ async def _run_output_guardrails(self, text: str, response_id: str) -> bool:
 
     def _enqueue_guardrail_task(self, text: str, response_id: str) -> None:
         # Runs the guardrails in a separate task to avoid blocking the main loop
+        if self._closing or self._closed:
+            return
 
         task = asyncio.create_task(self._run_output_guardrails(text, response_id))
         self._guardrail_tasks.add(task)
@@ -1277,6 +1281,11 @@ def _on_guardrail_task_done(self, task: asyncio.Task[Any]) -> None:
         # Remove from tracking set
         self._guardrail_tasks.discard(task)
 
+        if self._closing or self._closed:
+            if not task.cancelled():
+                task.exception()
+            return
+
         # Check for exceptions and propagate as events
         if not task.cancelled():
             exception = task.exception()
@@ -1291,11 +1300,8 @@ def _on_guardrail_task_done(self, task: asyncio.Task[Any]) -> None:
                     )
                 )
 
-    def _cleanup_guardrail_tasks(self) -> None:
-        for task in self._guardrail_tasks:
-            if not task.done():
-                task.cancel()
-        self._guardrail_tasks.clear()
+    async def _cleanup_guardrail_tasks(self, caller_task: asyncio.Task[Any] | None) -> None:
+        await self._cancel_and_wait_for_tasks(self._guardrail_tasks, "guardrail", caller_task)
 
     def _enqueue_tool_call_task(
         self,
@@ -1307,6 +1313,9 @@ def _enqueue_tool_call_task(
         call_id_reserved: bool = False,
     ) -> None:
         """Run tool calls in the background to avoid blocking realtime transport."""
+        if self._closing or self._closed:
+            return
+
         handle_kwargs: dict[str, Any] = {"agent_snapshot": agent_snapshot}
         if dispatch_snapshot is not None:
             handle_kwargs["dispatch_snapshot"] = dispatch_snapshot
@@ -1322,6 +1331,11 @@ def _enqueue_tool_call_task(
     def _on_tool_call_task_done(self, task: asyncio.Task[Any]) -> None:
         self._tool_call_tasks.discard(task)
 
+        if self._closing or self._closed:
+            if not task.cancelled():
+                task.exception()
+            return
+
         if task.cancelled():
             return
 
@@ -1364,11 +1378,38 @@ def _on_tool_call_task_done(self, task: asyncio.Task[Any]) -> None:
             )
         )
 
-    def _cleanup_tool_call_tasks(self) -> None:
-        for task in self._tool_call_tasks:
+    async def _cleanup_tool_call_tasks(self, caller_task: asyncio.Task[Any] | None) -> None:
+        await self._cancel_and_wait_for_tasks(self._tool_call_tasks, "tool call", caller_task)
+
+    async def _cancel_and_wait_for_tasks(
+        self,
+        tasks: set[asyncio.Task[Any]],
+        label: str,
+        caller_task: asyncio.Task[Any] | None,
+    ) -> None:
+        tasks_to_wait: list[asyncio.Task[Any]] = []
+
+        for task in list(tasks):
+            if task is caller_task:
+                tasks.discard(task)
+                continue
             if not task.done():
                 task.cancel()
-        self._tool_call_tasks.clear()
+            tasks_to_wait.append(task)
+
+        if tasks_to_wait:
+            _done, pending = await asyncio.wait(
+                tasks_to_wait,
+                timeout=_BACKGROUND_TASK_CLEANUP_TIMEOUT,
+            )
+            if pending:
+                logger.warning(
+                    "Timed out waiting for %d realtime %s background task(s) to stop.",
+                    len(pending),
+                    label,
+                )
+
+        tasks.difference_update(tasks_to_wait)
 
     def _wake_event_iterators(self) -> None:
         for _ in range(self._event_iterator_waiters):
@@ -1379,24 +1420,38 @@ async def _cleanup(self) -> None:
         if self._closed:
             self._wake_event_iterators()
             return
+        if self._closing:
+            self._wake_event_iterators()
+            return
+
+        self._closing = True
+        caller_task = asyncio.current_task()
 
-        # Cancel and cleanup guardrail tasks
-        self._cleanup_guardrail_tasks()
-        self._cleanup_tool_call_tasks()
+        try:
+            # Cancel and cleanup guardrail tasks
+            await asyncio.gather(
+                self._cleanup_guardrail_tasks(caller_task),
+                self._cleanup_tool_call_tasks(caller_task),
+            )
 
-        # Remove ourselves as a listener
-        self._model.remove_listener(self)
+            # Remove ourselves as a listener
+            self._model.remove_listener(self)
 
-        # Close the model connection
-        await self._model.close()
+            # Close the model connection
+            await self._model.close()
 
-        # Clear pending approval tracking
-        self._pending_tool_calls.clear()
-        self._pending_tool_outputs.clear()
+            # Clear pending approval tracking
+            self._pending_tool_calls.clear()
+            self._pending_tool_outputs.clear()
 
-        # Mark as closed
-        self._closed = True
-        self._wake_event_iterators()
+            # Mark as closed
+            self._closed = True
+        except BaseException:
+            self._closing = False
+            raise
+        else:
+            self._closing = False
+            self._wake_event_iterators()
 
     def _dispatch_snapshot_from_settings(
         self,

tests/realtime/test_session.py

@@ -11,6 +11,7 @@
 from agents.exceptions import ToolTimeoutError, UserError
 from agents.guardrail import GuardrailFunctionOutput, OutputGuardrail
 from agents.handoffs import Handoff
+from agents.realtime import session as session_module
 from agents.realtime.agent import RealtimeAgent
 from agents.realtime.config import RealtimeRunConfig, RealtimeSessionModelSettings
 from agents.realtime.events import (
@@ -208,6 +209,156 @@ async def test_aiter_exits_waiting_iterators_when_session_closes():
             task.result()
 
 
+@pytest.mark.asyncio
+async def test_cleanup_awaits_cancelled_task_finalizers_before_model_close():
+    close_order: list[str] = []
+
+    class _CloseRecordingModel(_DummyModel):
+        async def close(self):
+            close_order.append("model_close")
+
+    model = _CloseRecordingModel()
+    agent = RealtimeAgent(name="agent")
+    session = RealtimeSession(model, agent, None)
+    guardrail_started = asyncio.Event()
+    tool_started = asyncio.Event()
+
+    async def tracked_task(label: str, started: asyncio.Event) -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            await asyncio.sleep(0)
+            close_order.append(label)
+
+    guardrail = asyncio.create_task(tracked_task("guardrail", guardrail_started))
+    tool_call = asyncio.create_task(tracked_task("tool", tool_started))
+    session._guardrail_tasks.add(guardrail)
+    session._tool_call_tasks.add(tool_call)
+
+    await guardrail_started.wait()
+    await tool_started.wait()
+
+    await session._cleanup()
+
+    try:
+        assert close_order[-1] == "model_close"
+        assert set(close_order[:2]) == {"guardrail", "tool"}
+        assert len(session._guardrail_tasks) == 0
+        assert len(session._tool_call_tasks) == 0
+    finally:
+        await asyncio.gather(guardrail, tool_call, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+async def test_cleanup_bounds_wait_for_cancellation_resistant_tasks(monkeypatch):
+    monkeypatch.setattr(session_module, "_BACKGROUND_TASK_CLEANUP_TIMEOUT", 0.01, raising=False)
+
+    model = _DummyModel()
+    agent = RealtimeAgent(name="agent")
+    session = RealtimeSession(model, agent, None)
+    started = asyncio.Event()
+    cancel_seen = asyncio.Event()
+    release = asyncio.Event()
+
+    async def cancellation_resistant_task() -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            cancel_seen.set()
+            await release.wait()
+
+    task = asyncio.create_task(cancellation_resistant_task())
+    session._guardrail_tasks.add(task)
+    await started.wait()
+
+    try:
+        await asyncio.wait_for(session._cleanup(), timeout=1)
+        assert cancel_seen.is_set()
+        assert session._closed is True
+        assert not task.done()
+        assert len(session._guardrail_tasks) == 0
+    finally:
+        release.set()
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+async def test_tracked_task_can_close_session_without_awaiting_itself():
+    class _CloseCountingModel(_DummyModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.close_count = 0
+
+        async def close(self):
+            self.close_count += 1
+
+    model = _CloseCountingModel()
+    agent = RealtimeAgent(name="agent")
+    session = RealtimeSession(model, agent, None)
+    started = asyncio.Event()
+    finished = asyncio.Event()
+
+    async def close_from_tracked_task() -> None:
+        started.set()
+        await session.close()
+        finished.set()
+
+    task = asyncio.create_task(close_from_tracked_task())
+    session._tool_call_tasks.add(task)
+    await started.wait()
+
+    await asyncio.wait_for(task, timeout=1)
+
+    assert finished.is_set()
+    assert session._closed is True
+    assert model.close_count == 1
+    assert task not in session._tool_call_tasks
+
+
+@pytest.mark.asyncio
+async def test_late_background_task_failures_after_cleanup_do_not_mutate_closed_session(
+    monkeypatch,
+):
+    monkeypatch.setattr(session_module, "_BACKGROUND_TASK_CLEANUP_TIMEOUT", 0.01, raising=False)
+
+    model = _DummyModel()
+    agent = RealtimeAgent(name="agent")
+    session = RealtimeSession(model, agent, None)
+    guardrail_started = asyncio.Event()
+    tool_started = asyncio.Event()
+    release = asyncio.Event()
+
+    async def fail_after_cleanup_timeout(started: asyncio.Event) -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            await release.wait()
+            raise RuntimeError("late background failure") from None
+
+    guardrail_task = asyncio.create_task(fail_after_cleanup_timeout(guardrail_started))
+    tool_task = asyncio.create_task(fail_after_cleanup_timeout(tool_started))
+    session._guardrail_tasks.add(guardrail_task)
+    session._tool_call_tasks.add(tool_task)
+    guardrail_task.add_done_callback(session._on_guardrail_task_done)
+    tool_task.add_done_callback(session._on_tool_call_task_done)
+    await guardrail_started.wait()
+    await tool_started.wait()
+
+    await asyncio.wait_for(session._cleanup(), timeout=1)
+    release.set()
+    await asyncio.gather(guardrail_task, tool_task, return_exceptions=True)
+    await asyncio.sleep(0)
+
+    assert session._stored_exception is None
+    assert session._event_queue.empty()
+    assert guardrail_task not in session._guardrail_tasks
+    assert tool_task not in session._tool_call_tasks
+
+
 @pytest.mark.asyncio
 async def test_transcription_completed_adds_new_user_item():
     model = _DummyModel()