fix(voice): wake stream when producer is cancelled

src/agents/voice/pipeline.py

@@ -112,6 +112,9 @@ async def stream_events():
                         await output._add_text(text_event)
                     await output._turn_done()
                     await output._done()
+                except asyncio.CancelledError:
+                    await output._cancel()
+                    raise
                 except Exception as e:
                     log_model_and_tool_action_error(logger, "Error processing single voice turn", e)
                     await output._add_error(e)
@@ -135,6 +138,8 @@ async def process_turns():
             ):
                 transcription_session = None
                 reported_error = False
+                primary_exception: BaseException | None = None
+                close_exception: Exception | None = None
                 try:
                     try:
                         emitted_intro = False
@@ -165,33 +170,73 @@ async def process_turns():
                             async for text_event in result:
                                 await output._add_text(text_event)
                             await output._turn_done()
+                    except asyncio.CancelledError as e:
+                        # A transcription producer can be cancelled independently of the stream
+                        # consumer. Stop pending synthesis and publish the terminal event through
+                        # the ordered dispatcher before preserving that cancellation.
+                        primary_exception = e
+                        await output._cancel()
                     except Exception as e:
                         # Report before closing the session below. A `close()` that also fails
                         # would otherwise replace this exception on its way out and the consumer
                         # would see only the cleanup error.
                         log_model_and_tool_action_error(logger, "Error processing voice turns", e)
                         await output._add_error(e)
                         reported_error = True
-                        raise
+                        primary_exception = e
                 finally:
                     if transcription_session is not None:
+                        close_task = asyncio.create_task(transcription_session.close())
+                        close_cancellation: asyncio.CancelledError | None = None
+                        close_result: BaseException | None = None
+
+                        async def wait_for_close() -> None:
+                            nonlocal close_result
+                            close_result = (
+                                await asyncio.gather(close_task, return_exceptions=True)
+                            )[0]
+
                         try:
-                            await transcription_session.close()
+                            # Shield the provider cleanup so cancellation cannot interrupt its
+                            # resource release. A cancellation-resistant wait also handles a
+                            # second cancellation while the close operation is suspended.
+                            await asyncio.shield(close_task)
+                        except asyncio.CancelledError as e:
+                            close_cancellation = e
+                            await output._await_cleanup(wait_for_close())
                         except Exception as e:
+                            close_result = e
+
+                        if isinstance(close_result, asyncio.CancelledError):
+                            close_cancellation = close_result
+                        elif isinstance(close_result, Exception):
                             log_model_and_tool_action_error(
-                                logger, "Error closing voice transcription session", e
+                                logger, "Error closing voice transcription session", close_result
                             )
-                            # Report only if nothing else has, which keeps the turn error's
-                            # precedence. Clean runs and cancelled producers both arrive here
-                            # with no terminal event queued and no other way to be released.
-                            if not reported_error:
-                                await output._add_error(e)
-                            raise
-
-                # Only a clean run reaches here. The error path above has already queued its
-                # terminal event, and a cancelled producer has no consumer left to serve, so
-                # neither should start TTS work or wait on it.
-                await output._done()
+                            if primary_exception is None:
+                                # Report only if nothing else has, which keeps the turn error's
+                                # precedence.
+                                if not reported_error:
+                                    await output._add_error(close_result)
+                                close_exception = close_result
+                            # Keep cancellation or the turn error as the producer outcome; the
+                            # close failure is already logged as secondary information.
+
+                        if close_cancellation is not None and primary_exception is None:
+                            primary_exception = close_cancellation
+                            await output._cancel()
+
+                if primary_exception is not None:
+                    raise primary_exception
+                if close_exception is not None:
+                    raise close_exception
+                # Only a clean run reaches here. Error and cancellation paths have already queued
+                # their terminal events, so neither should start TTS work or wait on it.
+                try:
+                    await output._done()
+                except asyncio.CancelledError:
+                    await output._cancel()
+                    raise
 
         output._set_task(asyncio.create_task(process_turns()))
         return output

src/agents/voice/result.py

@@ -3,7 +3,7 @@
 import asyncio
 import base64
 from collections import deque
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Awaitable
 from typing import Any
 
 from ..exceptions import UserError
@@ -58,6 +58,10 @@ def __init__(
         self._turn_text_buffer = ""
         self._queue: asyncio.Queue[VoiceStreamEvent] = asyncio.Queue()
         self._tasks: list[asyncio.Task[Any]] = []
+        self._audio_task_queues: dict[
+            asyncio.Task[Any], asyncio.Queue[VoiceStreamEvent | None]
+        ] = {}
+        self._started_audio_tasks: set[asyncio.Task[Any]] = set()
         self._ordered_tasks: deque[asyncio.Queue[VoiceStreamEvent | None]] = (
             deque()
         )  # New: deque to hold local queues for each text segment
@@ -72,6 +76,8 @@ def __init__(
         self._first_byte_received = False
         self._generation_start_time: str | None = None
         self._completed_session = False
+        self._terminal_event_queued = False
+        self._terminal_event_enqueued = False
         self._stored_exception: BaseException | None = None
         self._tracing_span: Span[SpeechGroupSpanData] | None = None
 
@@ -96,6 +102,16 @@ def _enqueue_audio_segment(self, local_queue: asyncio.Queue[VoiceStreamEvent | N
         self._ordered_tasks.append(local_queue)
         self._dispatcher_event.set()
 
+    def _create_audio_task(
+        self,
+        text: str,
+        local_queue: asyncio.Queue[VoiceStreamEvent | None],
+        finish_turn: bool = False,
+    ) -> None:
+        task = asyncio.create_task(self._stream_audio(text, local_queue, finish_turn))
+        self._tasks.append(task)
+        self._audio_task_queues[task] = local_queue
+
     def _transform_audio_buffer(
         self, buffer: list[bytes], output_dtype: npt.DTypeLike
     ) -> npt.NDArray[np.int16 | np.float32]:
@@ -119,6 +135,10 @@ async def _stream_audio(
         local_queue: asyncio.Queue[VoiceStreamEvent | None],
         finish_turn: bool = False,
     ):
+        current_task = asyncio.current_task()
+        if current_task is not None:
+            self._started_audio_tasks.add(current_task)
+
         with speech_span(
             model=self.tts_model.model_name,
             input=text if self._voice_pipeline_config.trace_include_sensitive_data else "",
@@ -192,6 +212,10 @@ async def _stream_audio(
                     await local_queue.put(VoiceStreamEventLifecycle(event="turn_ended"))
                 else:
                     await local_queue.put(None)  # Signal completion for this segment
+            except asyncio.CancelledError:
+                # Let the ordered dispatcher advance past a segment cancelled during shutdown.
+                local_queue.put_nowait(None)
+                raise
             except Exception as e:
                 tts_span.set_error(
                     {
@@ -226,21 +250,15 @@ async def _add_text(self, text: str):
         if combined_sentences:
             local_queue: asyncio.Queue[VoiceStreamEvent | None] = asyncio.Queue()
             self._enqueue_audio_segment(local_queue)
-            self._tasks.append(
-                asyncio.create_task(self._stream_audio(combined_sentences, local_queue))
-            )
+            self._create_audio_task(combined_sentences, local_queue)
             if self._dispatcher_task is None:
                 self._dispatcher_task = asyncio.create_task(self._dispatch_audio())
 
     async def _turn_done(self):
         if self._text_buffer:
             local_queue: asyncio.Queue[VoiceStreamEvent | None] = asyncio.Queue()
             self._enqueue_audio_segment(local_queue)
-            self._tasks.append(
-                asyncio.create_task(
-                    self._stream_audio(self._text_buffer, local_queue, finish_turn=True)
-                )
-            )
+            self._create_audio_task(self._text_buffer, local_queue, finish_turn=True)
             self._text_buffer = ""
         elif self._started_processing_turn:
             local_queue = asyncio.Queue()
@@ -273,6 +291,54 @@ async def _done(self):
             self._dispatcher_task = asyncio.create_task(self._dispatch_audio())
         await self._wait_for_completion()
 
+    async def _cancel(self) -> None:
+        """Stop pending synthesis and enqueue an ordered terminal event."""
+        current_task = asyncio.current_task()
+        if self._completed_session and self._terminal_event_enqueued:
+            return
+
+        # Publish the terminal queue before awaiting any cancellable synthesis cleanup. A second
+        # cancellation can interrupt that await, but the dispatcher must still have a terminal
+        # event to release a consumer that is waiting independently of the producer.
+        self._completed_session = True
+        if not self._terminal_event_queued and not self._terminal_event_enqueued:
+            terminal_queue: asyncio.Queue[VoiceStreamEvent | None] = asyncio.Queue()
+            self._enqueue_audio_segment(terminal_queue)
+            terminal_queue.put_nowait(VoiceStreamEventLifecycle(event="session_ended"))
+            self._terminal_event_queued = True
+
+        if self._dispatcher_task is None or self._dispatcher_task.done():
+            self._dispatcher_task = asyncio.create_task(self._dispatch_audio())
+
+        tasks = [task for task in self._tasks if task is not current_task and not task.done()]
+        for task in tasks:
+            if task not in self._started_audio_tasks:
+                local_queue = self._audio_task_queues.get(task)
+                if local_queue is not None:
+                    local_queue.put_nowait(None)
+            task.cancel()
+
+        try:
+            if tasks:
+                await self._await_cleanup(asyncio.gather(*tasks, return_exceptions=True))
+            if self._dispatcher_task is not None:
+                await self._await_cleanup(
+                    asyncio.gather(self._dispatcher_task, return_exceptions=True)
+                )
+        finally:
+            # This must happen while the producer is still inside its TraceCtxManager, so the
+            # speech-group span closes before the enclosing trace emits trace_end.
+            self._finish_turn()
+
+    async def _await_cleanup(self, awaitable: Awaitable[Any]) -> None:
+        """Wait for cleanup to finish even if a caller cancels the producer again."""
+        while True:
+            try:
+                await asyncio.shield(awaitable)
+            except asyncio.CancelledError:
+                continue
+            return
+
     async def _dispatch_audio(self):
         # Dispatch audio chunks from each segment in the order they were added
         while True:
@@ -296,8 +362,12 @@ async def _dispatch_audio(self):
                         self._finish_turn()
                         break
                     if chunk.event == "session_ended":
+                        self._terminal_event_queued = True
+                        self._terminal_event_enqueued = True
                         return
+        self._terminal_event_queued = True
         await self._queue.put(VoiceStreamEventLifecycle(event="session_ended"))
+        self._terminal_event_enqueued = True
 
     async def _wait_for_completion(self):
         tasks: list[asyncio.Task[Any]] = self._tasks
@@ -382,6 +452,13 @@ async def stream(self) -> AsyncIterator[VoiceStreamEvent]:
                     await asyncio.shield(self.text_generation_task)
                 except BaseException as exc:
                     producer_exception = exc
+                    if isinstance(exc, asyncio.CancelledError) and self.text_generation_task.done():
+                        try:
+                            self.text_generation_task.exception()
+                        except asyncio.CancelledError as task_exception:
+                            # asyncio.shield() drops the cancellation message when the producer
+                            # is already cancelled; recover the original exception from the task.
+                            producer_exception = task_exception
 
             try:
                 await self._cleanup_tasks()