fix(avatar): recover DataStreamAudioReceiver from a lost stream trailer

livekit-agents/livekit/agents/voice/avatar/_datastream_io.py

@@ -21,6 +21,10 @@
 RPC_PLAYBACK_STARTED = "lk.playback_started"
 AUDIO_STREAM_TOPIC = "lk.audio_stream"
 
+# Fallback for the last segment: with no chunk or trailer for this long, assume the
+# trailer was lost and end the segment. Generous — audio is pushed faster than realtime.
+STREAM_IDLE_TIMEOUT = 10.0
+
 
 @dataclass
 class _PlaybackStartedEvent:
@@ -363,6 +367,11 @@ def __init__(
 
         self._current_reader: rtc.ByteStreamReader | None = None
         self._current_reader_cleared: bool = False
+        # Set when the current reader's trailer no longer matters: a new stream header
+        # arrived with a reader still open (the sender never overlaps streams, so the
+        # trailer was lost — the "avatar stops receiving audio" wedge), or clear_buffer
+        # discarded the segment. Ends the reader instead of blocking forever.
+        self._current_reader_superseded: asyncio.Event = asyncio.Event()
 
         self._rpc_send_ch = utils.aio.Chan[PlaybackFinishedEvent | _PlaybackStartedEvent]()
         self._rpc_max_retries = rpc_max_retries
@@ -394,6 +403,8 @@ def _handle_clear_buffer(data: rtc.RpcInvocationData) -> str:
 
             if self._current_reader:
                 self._current_reader_cleared = True
+                # post-clear data is discarded, so don't wait for a trailer that may be lost
+                self._current_reader_superseded.set()
 
             # clear the audio internal buffer
             while not self._data_ch.empty():
@@ -412,6 +423,9 @@ def _handle_stream_received(
                 return
 
             self._stream_readers.append(reader)
+            if self._current_reader is not None:
+                # new segment started → the open reader's trailer is not coming
+                self._current_reader_superseded.set()
             self._stream_reader_changed.set()
 
         self._register_clear_buffer_rpc(
@@ -491,6 +505,7 @@ async def _recv_task(self) -> None:
 
             while self._stream_readers:
                 self._current_reader = self._stream_readers.pop(0)
+                self._current_reader_superseded.clear()
 
                 if (
                     not (attrs := self._current_reader.info.attributes)
@@ -508,7 +523,7 @@ async def _recv_task(self) -> None:
                 )
 
                 try:
-                    async for data in self._current_reader:
+                    async for data in self._iter_reader(self._current_reader):
                         if self._current_reader_cleared:
                             # ignore the rest data of the current reader if clear_buffer was called
                             while not self._data_ch.empty():
@@ -534,6 +549,57 @@ async def _recv_task(self) -> None:
 
             self._stream_reader_changed.clear()
 
+    async def _iter_reader(self, reader: rtc.ByteStreamReader) -> AsyncIterator[bytes]:
+        """Yield the reader's chunks, but end the segment if its trailer is lost.
+
+        A queued chunk or trailer always wins the race, so a healthy segment is
+        unchanged. The segment ends early only when the reader is superseded by a newer
+        stream or a clear_buffer, or stays idle for ``STREAM_IDLE_TIMEOUT``
+        (the last-segment case).
+        """
+        superseded = asyncio.ensure_future(self._current_reader_superseded.wait())
+        try:
+            while True:
+                next_chunk = asyncio.ensure_future(reader.__anext__())
+                try:
+                    done, _ = await asyncio.wait(
+                        {next_chunk, superseded},
+                        timeout=STREAM_IDLE_TIMEOUT,
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                except asyncio.CancelledError:
+                    next_chunk.cancel()
+                    raise
+
+                if next_chunk in done:
+                    try:
+                        yield next_chunk.result()
+                    except StopAsyncIteration:
+                        return
+                    continue
+
+                next_chunk.cancel()
+                if self._current_reader_cleared:
+                    # expected after clear_buffer: the rest of the stream is discarded anyway
+                    logger.debug(
+                        "audio stream ended early after clear_buffer",
+                        extra={"stream_id": reader.info.stream_id},
+                    )
+                    return
+                reason = (
+                    "superseded by a newer stream"
+                    if superseded in done
+                    else (f"idle for {STREAM_IDLE_TIMEOUT:g}s")
+                )
+                logger.warning(
+                    "audio stream ended before its trailer arrived (%s)",
+                    reason,
+                    extra={"stream_id": reader.info.stream_id},
+                )
+                return
+        finally:
+            superseded.cancel()
+
     def __aiter__(self) -> AsyncIterator[rtc.AudioFrame | AudioSegmentEnd]:
         return self
 

tests/test_datastream_receiver_lost_trailer.py

@@ -0,0 +1,213 @@
+"""DataStreamAudioReceiver must survive a lost stream trailer.
+
+If a segment's trailer is lost, the receiver used to block on it forever, wedging
+all later segments. A new stream header (or an idle timeout) must end the segment.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from livekit import rtc
+from livekit.agents.voice.avatar import AudioSegmentEnd, _datastream_io as ds
+
+pytestmark = pytest.mark.unit
+
+SAMPLE_RATE = 24000
+FRAME_MS = 10
+CHUNK = b"\x00\x00" * int(SAMPLE_RATE * FRAME_MS / 1000)  # exactly one 10 ms frame
+SENDER = "agent"
+
+
+class _FakeReader:
+    """Stands in for rtc.ByteStreamReader: chunks are pushed in, the trailer is
+    an explicit close() — which the lost-trailer tests never call."""
+
+    def __init__(self, stream_id: str) -> None:
+        self.info = SimpleNamespace(
+            stream_id=stream_id,
+            attributes={"sample_rate": str(SAMPLE_RATE), "num_channels": "1"},
+        )
+        self._queue: asyncio.Queue[bytes | None] = asyncio.Queue()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        item = await self._queue.get()
+        if item is None:
+            raise StopAsyncIteration
+        return item
+
+    def push(self, data: bytes) -> None:
+        self._queue.put_nowait(data)
+
+    def close(self) -> None:
+        self._queue.put_nowait(None)
+
+
+async def _start_receiver():
+    room = MagicMock()
+    handlers: dict = {}
+    room.register_byte_stream_handler.side_effect = lambda topic, h: handlers.__setitem__(topic, h)
+    room.local_participant._rpc_handlers = {}
+    rpc_methods: dict = {}
+    room.local_participant.register_rpc_method.side_effect = lambda name, h: (
+        rpc_methods.__setitem__(name, h)
+    )
+    sender = SimpleNamespace(identity=SENDER)
+    with patch.object(ds.utils, "wait_for_participant", AsyncMock(return_value=sender)):
+        receiver = ds.DataStreamAudioReceiver(room, sender_identity=SENDER, frame_size_ms=FRAME_MS)
+        await receiver.start()
+    return receiver, handlers[ds.AUDIO_STREAM_TOPIC], rpc_methods[ds.RPC_CLEAR_BUFFER]
+
+
+async def _next(receiver, timeout: float = 1.0):
+    return await asyncio.wait_for(receiver.__anext__(), timeout)
+
+
+async def test_normal_segments_are_unaffected():
+    receiver, on_stream, _clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        first.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        first.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()
+
+
+async def test_lost_trailer_is_recovered_by_the_next_stream():
+    receiver, on_stream, _clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        first.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        # first.close() never happens: the trailer was lost.
+
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+        # The first segment is ended so the second one can be read.
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()
+
+
+async def test_clear_buffer_ends_the_segment_without_waiting_for_the_trailer():
+    """The cleared segment's data is discarded anyway, so AudioSegmentEnd must
+    arrive right away — no next stream header and no idle timeout needed, even
+    if the trailer was lost."""
+    receiver, on_stream, clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        first.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+
+        # interruption: clear_buffer lands, then the stream's trailer is lost
+        assert clear(SimpleNamespace(caller_identity=SENDER)) == "ok"
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+
+        # the receiver is ready for the next utterance
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()
+
+
+async def test_clear_buffer_with_a_healthy_trailer_still_ends_one_segment():
+    """A trailer that does arrive after clear_buffer must not produce a second
+    AudioSegmentEnd or disturb the next segment."""
+    receiver, on_stream, clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        first.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+
+        assert clear(SimpleNamespace(caller_identity=SENDER)) == "ok"
+        first.close()  # healthy interruption: the trailer still arrives
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()
+
+
+async def test_last_stream_lost_trailer_is_recovered_by_idle_timeout(monkeypatch):
+    """No following stream to supersede it: an open reader that goes silent for
+    STREAM_IDLE_TIMEOUT while a segment is owed is ended anyway, so the receiver
+    is ready for the next utterance instead of wedged on the lost trailer."""
+    monkeypatch.setattr(ds, "STREAM_IDLE_TIMEOUT", 0.2)
+    receiver, on_stream, _clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        first.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        # trailer lost, and no next stream arrives.
+        assert isinstance(await _next(receiver, timeout=1.0), AudioSegmentEnd)
+
+        # the receiver is not wedged: a later utterance is read normally.
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()
+
+
+async def test_queued_trailer_still_wins_over_the_superseded_signal():
+    """Ordered channel: a header can only follow its predecessor's trailer, so
+    a trailer that is already queued must end the segment normally (frames
+    flushed) — the superseded signal is a fallback, never a preemption."""
+    receiver, on_stream, _clear = await _start_receiver()
+    try:
+        first = _FakeReader("s1")
+        on_stream(first, SENDER)
+        # a partial (half) frame that only a normal flush would emit
+        first.push(CHUNK[: len(CHUNK) // 2])
+        first.close()
+        second = _FakeReader("s2")
+        on_stream(second, SENDER)
+
+        frame = await _next(receiver)
+        assert isinstance(frame, rtc.AudioFrame)
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+
+        second.push(CHUNK)
+        assert isinstance(await _next(receiver), rtc.AudioFrame)
+        second.close()
+        assert isinstance(await _next(receiver), AudioSegmentEnd)
+    finally:
+        await receiver.aclose()