fix(voice): preserve transcript flush during session close

livekit-agents/livekit/agents/voice/agent_session.py

@@ -1248,6 +1248,7 @@ async def _aclose_impl(
                     audio_recognition._commit_user_turn(
                         audio_detached=True,
                         transcript_timeout=self._opts.session_close_transcript_timeout,
+                        session_close=True,
                     )
 
                 await activity.aclose()

livekit-agents/livekit/agents/voice/audio_recognition.py

@@ -7,7 +7,7 @@
 from collections import deque
 from collections.abc import AsyncIterable, Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
 
 from opentelemetry import trace
 from opentelemetry.sdk.trace import ReadableSpan
@@ -249,9 +249,11 @@ def __init__(
         self._hooks = hooks
         self._audio_input_atask: asyncio.Task[None] | None = None
         self._commit_user_turn_atask: asyncio.Task[None] | None = None
+        self._session_close_commit_user_turn_atask: asyncio.Task[None] | None = None
         self._stt_consumer_atask: asyncio.Task[None] | None = None
         self._vad_atask: asyncio.Task[None] | None = None
         self._end_of_turn_task: asyncio.Task[None] | None = None
+        self._session_close_end_of_turn_atask: asyncio.Task[None] | None = None
         self._endpointing: BaseEndpointing = endpointing
         self._turn_detector = turn_detection if not isinstance(turn_detection, str) else None
         self._stt = stt
@@ -792,42 +794,96 @@ def _push_audio(
 
     async def _aclose(self) -> None:
         self._closing.set()
-        try:
-            if self._commit_user_turn_atask is not None:
-                await aio.cancel_and_wait(self._commit_user_turn_atask)
 
-            if self._stt_pipeline is not None:
-                await self._stt_pipeline.aclose()
-                self._stt_pipeline = None
+        async def _cleanup() -> None:
+            flush_error: BaseException | None = None
 
-            await aio.cancel_and_wait(*self._tasks)
+            async def _finish_close_flush(task: asyncio.Task[None]) -> None:
+                nonlocal flush_error
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                except BaseException as exc:
+                    if flush_error is None:
+                        flush_error = exc
 
-            if self._stt_consumer_atask is not None:
-                await aio.cancel_and_wait(self._stt_consumer_atask)
+            try:
+                commit_user_turn_atask = self._commit_user_turn_atask
+                if commit_user_turn_atask is not None:
+                    session_close_commit_user_turn_atask = cast(
+                        asyncio.Task[None] | None,
+                        getattr(self, "_session_close_commit_user_turn_atask", None),
+                    )
+                    if commit_user_turn_atask is session_close_commit_user_turn_atask:
+                        await _finish_close_flush(commit_user_turn_atask)
+                    else:
+                        await aio.cancel_and_wait(commit_user_turn_atask)
 
-            if self._vad_atask is not None:
-                await aio.cancel_and_wait(self._vad_atask)
+                if self._stt_pipeline is not None:
+                    await self._stt_pipeline.aclose()
+                    self._stt_pipeline = None
 
-            if self._interruption_atask is not None:
-                await aio.cancel_and_wait(self._interruption_atask)
+                await aio.cancel_and_wait(*self._tasks)
 
-            if self._end_of_turn_task is not None:
-                await aio.cancel_and_wait(self._end_of_turn_task)
+                if self._stt_consumer_atask is not None:
+                    await aio.cancel_and_wait(self._stt_consumer_atask)
 
-            if self._turn_detector_stream is not None:
-                await self._turn_detector_stream.aclose()
-                self._turn_detector_stream = None
-            self._turn_detector_prediction_fut = None
+                if self._vad_atask is not None:
+                    await aio.cancel_and_wait(self._vad_atask)
 
-            if self._backchannel_boundary_timer is not None:
-                self._backchannel_boundary_timer.cancel()
-                self._backchannel_boundary_timer = None
-                self._backchannel_boundary_callback = None
-        finally:
-            self._cancel_transcription_timeout()
-            # EOU normally ends this span, but teardown cancels EOU before a
-            # pending speech segment necessarily produces a transcript.
-            self._end_user_turn_span()
+                if self._interruption_atask is not None:
+                    await aio.cancel_and_wait(self._interruption_atask)
+
+                end_of_turn_task = self._end_of_turn_task
+                if end_of_turn_task is not None:
+                    session_close_end_of_turn_atask = cast(
+                        asyncio.Task[None] | None,
+                        getattr(self, "_session_close_end_of_turn_atask", None),
+                    )
+                    if end_of_turn_task is session_close_end_of_turn_atask:
+                        await _finish_close_flush(end_of_turn_task)
+                    else:
+                        await aio.cancel_and_wait(end_of_turn_task)
+
+                if self._turn_detector_stream is not None:
+                    await self._turn_detector_stream.aclose()
+                    self._turn_detector_stream = None
+                self._turn_detector_prediction_fut = None
+
+                if self._backchannel_boundary_timer is not None:
+                    self._backchannel_boundary_timer.cancel()
+                    self._backchannel_boundary_timer = None
+                    self._backchannel_boundary_callback = None
+
+                if flush_error is not None:
+                    raise flush_error
+            finally:
+                self._cancel_transcription_timeout()
+                # EOU normally ends this span, but teardown cancels EOU before a
+                # pending speech segment necessarily produces a transcript.
+                self._end_user_turn_span()
+
+        cleanup_task = asyncio.create_task(_cleanup())
+        outer_cancelled = False
+        while not cleanup_task.done():
+            try:
+                await asyncio.shield(cleanup_task)
+            except asyncio.CancelledError:
+                outer_cancelled = True
+            except BaseException:
+                break
+
+        try:
+            cleanup_task.result()
+        except BaseException as exc:
+            if outer_cancelled:
+                logger.error("audio recognition cleanup failed after cancellation", exc_info=exc)
+                raise asyncio.CancelledError from exc
+            raise
+
+        if outer_cancelled:
+            raise asyncio.CancelledError
 
     def _update_stt(
         self,
@@ -1043,6 +1099,7 @@ def _commit_user_turn(
         transcript_timeout: float,
         stt_flush_duration: float = 2.0,
         skip_reply: bool = False,
+        session_close: bool = False,
     ) -> asyncio.Future[str]:
         loop = asyncio.get_running_loop()
         fut: asyncio.Future[str] = loop.create_future()
@@ -1109,6 +1166,8 @@ async def _commit_user_turn() -> None:
                 skip_reply=skip_reply,
                 trigger="manual",
             )
+            if session_close:
+                self._session_close_end_of_turn_atask = self._end_of_turn_task
             self._user_turn_committed = True
             if not fut.done():
                 fut.set_result(transcript)
@@ -1126,6 +1185,9 @@ def _on_task_done(task: asyncio.Task[None]) -> None:
 
         self._commit_user_turn_atask = asyncio.create_task(_commit_user_turn())
         self._commit_user_turn_atask.add_done_callback(_on_task_done)
+        if session_close:
+            self._session_close_commit_user_turn_atask = self._commit_user_turn_atask
+            fut.add_done_callback(lambda done: None if done.cancelled() else done.exception())
         return fut
 
     @property

tests/test_agent_session.py

@@ -49,6 +49,7 @@
 from livekit.agents.voice.io import PlaybackFinishedEvent
 
 from .fake_session import FakeActions, create_session, run_session
+from .fake_stt import FakeSTT
 
 pytestmark = [pytest.mark.unit, pytest.mark.virtual_time, pytest.mark.no_concurrent]
 
@@ -242,6 +243,50 @@ async def test_events_and_metrics() -> None:
     check_timestamp(metrics_events[2].metrics.audio_duration, 2.0, speed_factor=speed)
 
 
+async def test_aclose_flushes_pending_interim_transcript() -> None:
+    """A normal session close commits an STT interim transcript that is still flushing."""
+    session = create_session(
+        FakeActions(),
+        turn_handling={"endpointing": {"min_delay": 0.0, "max_delay": 0.0}},
+        extra_kwargs={"session_close_transcript_timeout": 0.01},
+    )
+    final_transcripts: list[UserInputTranscribedEvent] = []
+    interim_received = asyncio.Event()
+    agent = MyAgent()
+
+    def on_transcript(event: UserInputTranscribedEvent) -> None:
+        if event.transcript == "last words":
+            if event.is_final:
+                final_transcripts.append(event)
+            else:
+                interim_received.set()
+
+    session.on("user_input_transcribed", on_transcript)
+
+    synchronizer = session.output.audio._synchronizer
+    try:
+        await session.start(agent)
+        stt = session.stt
+        assert isinstance(stt, FakeSTT)
+        stream = await asyncio.wait_for(stt.stream_ch.recv(), timeout=1)
+        stream.send_fake_transcript("last words", is_final=False)
+        await asyncio.wait_for(interim_received.wait(), timeout=1)
+
+        await session.aclose()
+
+        assert [event.transcript for event in final_transcripts] == ["last words"]
+        for chat_ctx in (agent.chat_ctx, session.history):
+            assert [
+                item.text_content
+                for item in chat_ctx.items
+                if item.type == "message" and item.role == "user"
+            ] == ["last words"]
+    finally:
+        if session._started:
+            await session.aclose()
+        await synchronizer.aclose()
+
+
 async def test_tts_node_ttfb_excludes_upstream_latency() -> None:
     # the LLM stream stays open for its full duration and the fake TTS only starts
     # synthesizing once its input is flushed. tts_node_ttfb must anchor on the text

tests/test_audio_recognition_aclose.py

@@ -154,6 +154,96 @@ async def long_running_task():
         assert commit_task.done()
         assert end_of_turn_task.done()
 
+    @pytest.mark.asyncio
+    async def test_aclose_propagates_outer_cancellation_after_finishing_cleanup(self) -> None:
+        """Outer cancellation waits for the close flush and all remaining cleanup."""
+        audio_recognition = self._create_audio_recognition()
+        flush_started = asyncio.Event()
+        release_flush = asyncio.Event()
+
+        async def flush() -> None:
+            flush_started.set()
+            await release_flush.wait()
+
+        flush_task = asyncio.create_task(flush())
+        audio_recognition._commit_user_turn_atask = flush_task
+        audio_recognition._session_close_commit_user_turn_atask = flush_task
+        stt_pipeline = MagicMock()
+        stt_pipeline.aclose = AsyncMock()
+        audio_recognition._stt_pipeline = stt_pipeline
+        turn_detector_stream = MagicMock()
+        turn_detector_stream.aclose = AsyncMock()
+        audio_recognition._turn_detector_stream = turn_detector_stream
+        background_tasks = [asyncio.create_task(asyncio.Event().wait()) for _ in range(4)]
+        audio_recognition._tasks = {background_tasks[0]}
+        audio_recognition._stt_consumer_atask = background_tasks[1]
+        audio_recognition._vad_atask = background_tasks[2]
+        audio_recognition._interruption_atask = background_tasks[3]
+        close_task = asyncio.create_task(audio_recognition._aclose())
+
+        try:
+            await flush_started.wait()
+            close_task.cancel()
+            await asyncio.sleep(0)
+
+            assert not close_task.done()
+            assert not flush_task.cancelled()
+
+            release_flush.set()
+            with pytest.raises(asyncio.CancelledError):
+                await close_task
+
+            assert flush_task.done()
+            assert not flush_task.cancelled()
+            stt_pipeline.aclose.assert_awaited_once()
+            turn_detector_stream.aclose.assert_awaited_once()
+            assert all(task.cancelled() for task in background_tasks)
+        finally:
+            release_flush.set()
+            await asyncio.gather(close_task, flush_task, *background_tasks, return_exceptions=True)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("flush_task_attr", ["_commit_user_turn_atask", "_end_of_turn_task"])
+    async def test_aclose_cancels_non_session_close_flushes(self, flush_task_attr: str) -> None:
+        """Missing close markers leave ordinary commit and EOU work cancellable."""
+        audio_recognition = self._create_audio_recognition()
+        flush_task = asyncio.create_task(asyncio.Event().wait())
+        setattr(audio_recognition, flush_task_attr, flush_task)
+
+        await audio_recognition._aclose()
+
+        assert flush_task.cancelled()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "flush_task_attr",
+        ["_session_close_commit_user_turn_atask", "_session_close_end_of_turn_atask"],
+    )
+    async def test_aclose_propagates_session_close_flush_errors_after_cleanup(
+        self, flush_task_attr: str
+    ) -> None:
+        """A failed close-owned flush is observed only after the remaining teardown finishes."""
+        audio_recognition = self._create_audio_recognition()
+
+        async def fail_flush() -> None:
+            raise RuntimeError("transcript flush failed")
+
+        flush_task = asyncio.create_task(fail_flush())
+        setattr(audio_recognition, flush_task_attr, flush_task)
+        if flush_task_attr == "_session_close_commit_user_turn_atask":
+            audio_recognition._commit_user_turn_atask = flush_task
+        else:
+            audio_recognition._end_of_turn_task = flush_task
+
+        turn_detector_stream = MagicMock()
+        turn_detector_stream.aclose = AsyncMock()
+        audio_recognition._turn_detector_stream = turn_detector_stream
+
+        with pytest.raises(RuntimeError, match="transcript flush failed"):
+            await audio_recognition._aclose()
+
+        turn_detector_stream.aclose.assert_awaited_once()
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(("is_recording", "expected_end_count"), [(True, 1), (False, 0)])
     async def test_aclose_finalizes_user_turn_span(