fix: handle RoomIO microphone track replacement

livekit-agents/livekit/agents/voice/room_io/_input.py

@@ -3,7 +3,7 @@
 import asyncio
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterable
-from typing import Any, Generic, TypeVar, cast
+from typing import Any, Generic, NamedTuple, TypeVar, cast
 
 from typing_extensions import override
 
@@ -19,6 +19,11 @@
 T = TypeVar("T", bound=rtc.AudioFrame | rtc.VideoFrame)
 
 
+class _StreamReadResult(NamedTuple):
+    received_frame: bool
+    forwarded_frame: bool
+
+
 class _ParticipantInputStream(Generic[T], ABC):
     """
     A stream that dynamically transitions between new audio and video feeds from a connected
@@ -45,7 +50,7 @@ def __init__(
         self._participant_identity: str | None = None
         self._attached = True
 
-        self._forward_atask: asyncio.Task[None] | None = None
+        self._forward_atask: asyncio.Task[_StreamReadResult] | None = None
         self._tasks: set[asyncio.Task[Any]] = set()
 
         self._room.on("track_subscribed", self._on_track_available)
@@ -118,10 +123,11 @@ def set_participant(self, participant: rtc.RemoteParticipant | str | None) -> No
                 self._on_track_available(publication.track, publication, participant)
 
     async def aclose(self) -> None:
-        if self._stream:
-            await self._stream.aclose()
-            self._stream = None
+        stream = self._stream
+        self._stream = None
         self._publication = None
+        if stream:
+            await stream.aclose()
         if self._processor:
             self._processor._close()
             self._processor = None
@@ -135,11 +141,11 @@ async def aclose(self) -> None:
     @log_exceptions(logger=logger)
     async def _forward_task(
         self,
-        old_task: asyncio.Task[None] | None,
+        old_task: asyncio.Task[_StreamReadResult] | None,
         stream: rtc.VideoStream | rtc.AudioStream,
         publication: rtc.RemoteTrackPublication,
         participant: rtc.RemoteParticipant,
-    ) -> None:
+    ) -> _StreamReadResult:
         if old_task:
             await aio.cancel_and_wait(old_task)
 
@@ -148,15 +154,23 @@ async def _forward_task(
             "source": rtc.TrackSource.Name(publication.source),
         }
         logger.debug("start reading stream", extra=extra)
+        received_frame = False
+        forwarded_frame = False
         async for event in stream:
+            received_frame = True
             if not self._attached:
                 # drop frames if the stream is detached
                 continue
+            forwarded_frame = True
             frame = cast(T, event.frame)
             self._process_frame(frame)
             await self._data_ch.send(frame)
 
         logger.debug("stream closed", extra=extra)
+        return _StreamReadResult(
+            received_frame=received_frame,
+            forwarded_frame=forwarded_frame,
+        )
 
     def _process_frame(self, frame: T) -> None:
         """Hook for subclasses to process frames in-place before forwarding."""
@@ -290,11 +304,11 @@ def _create_stream(self, track: rtc.Track, participant: rtc.Participant) -> rtc.
     @override
     async def _forward_task(
         self,
-        old_task: asyncio.Task[None] | None,
+        old_task: asyncio.Task[_StreamReadResult] | None,
         stream: rtc.AudioStream,  # type: ignore[override]
         publication: rtc.RemoteTrackPublication,
         participant: rtc.RemoteParticipant,
-    ) -> None:
+    ) -> _StreamReadResult:
         if old_task:
             await aio.cancel_and_wait(old_task)
 
@@ -331,18 +345,47 @@ async def _forward_task(
                     "error reading pre-connect audio buffer", extra=logging_extra, exc_info=e
                 )
 
-        await super()._forward_task(old_task, stream, publication, participant)
-
-        # push a silent frame to flush the stt final result if any
+        is_initial_stream = True
+        consecutive_empty_streams = 0
         silent_samples = int(self._sample_rate * 0.5)
-        await self._data_ch.send(
-            rtc.AudioFrame(
-                b"\x00\x00" * silent_samples,
-                sample_rate=self._sample_rate,
-                num_channels=self._num_channels,
-                samples_per_channel=silent_samples,
-            )
-        )
+        while True:
+            read_result = await super()._forward_task(None, stream, publication, participant)
+
+            # push a silent frame to flush the stt final result if any
+            if is_initial_stream or read_result.forwarded_frame:
+                await self._data_ch.send(
+                    rtc.AudioFrame(
+                        b"\x00\x00" * silent_samples,
+                        sample_rate=self._sample_rate,
+                        num_channels=self._num_channels,
+                        samples_per_channel=silent_samples,
+                    )
+                )
+            is_initial_stream = False
+
+            if self._stream is not stream or self._publication is not publication:
+                return read_result
+
+            # An RTC stream may reach EOS while its track remains subscribed.
+            track = publication.track
+            self._close_stream()
+            if track is None or not publication.subscribed:
+                return read_result
+
+            if read_result.received_frame:
+                consecutive_empty_streams = 0
+            else:
+                consecutive_empty_streams += 1
+            if consecutive_empty_streams > 1:
+                logger.warning(
+                    "replacement audio stream closed before receiving frames; not reattaching",
+                    extra={"participant": participant.identity, "track_id": track.sid},
+                )
+                return read_result
+
+            stream = self._create_stream(track, participant)
+            self._stream = stream
+            self._publication = publication
 
     def _resample_frames(self, frames: Iterable[rtc.AudioFrame]) -> Iterable[rtc.AudioFrame]:
         resampler: rtc.AudioResampler | None = None

tests/test_room_io.py

@@ -58,14 +58,32 @@ def unregister_text_stream_handler(self, topic: str) -> None:
 
 
 class _MockAudioStream:
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+        self.ended = asyncio.Event()
+        self._events: asyncio.Queue[rtc.AudioFrame | None] = asyncio.Queue()
+
     def __aiter__(self):
         return self
 
     async def __anext__(self):
-        raise StopAsyncIteration
+        self.started.set()
+        frame = await self._events.get()
+        if frame is None:
+            raise StopAsyncIteration
+        return SimpleNamespace(frame=frame)
 
     async def aclose(self) -> None:
-        pass
+        self.end()
+
+    def end(self) -> None:
+        if self.ended.is_set():
+            return
+        self.ended.set()
+        self._events.put_nowait(None)
+
+    def push_frame(self, frame: rtc.AudioFrame) -> None:
+        self._events.put_nowait(frame)
 
 
 class _MockFrameProcessor(rtc.FrameProcessor[rtc.AudioFrame]):
@@ -133,8 +151,12 @@ def _make_track_available_args(
     publication = MagicMock()
     publication.source = rtc.TrackSource.SOURCE_MICROPHONE
     publication.sid = sid
+    publication.track = track
+    publication.subscribed = True
+    publication.audio_features = []
     participant = MagicMock()
     participant.identity = identity
+    participant.track_publications = {sid: publication}
     return track, publication, participant
 
 
@@ -259,6 +281,157 @@ async def _close_agent() -> None:
 # -- frame processor lifecycle tests ------------------------------------------
 
 
+@pytest.mark.asyncio
+async def test_audio_input_reattaches_when_stream_closes_while_track_is_subscribed() -> None:
+    room = _FakeRoom()
+    audio_input = _make_audio_input_stream(room, noise_cancellation=None)
+    audio_input.set_participant("test-user")
+    track, publication, participant = _make_track_available_args()
+    first_stream = _MockAudioStream()
+    replacement_stream = _MockAudioStream()
+    second_replacement_stream = _MockAudioStream()
+    final_replacement_stream = _MockAudioStream()
+    frame = rtc.AudioFrame(bytes(480 * 2), 24000, 1, 480)
+
+    with patch(
+        "livekit.rtc.AudioStream.from_track",
+        side_effect=[
+            first_stream,
+            replacement_stream,
+            second_replacement_stream,
+            final_replacement_stream,
+        ],
+    ) as create_stream:
+        assert audio_input._on_track_available(track, publication, participant)
+        await asyncio.wait_for(first_stream.started.wait(), timeout=1)
+
+        first_stream.push_frame(frame)
+        first_stream.end()
+        await asyncio.wait_for(replacement_stream.started.wait(), timeout=1)
+
+        replacement_stream.push_frame(frame)
+        replacement_stream.end()
+        await asyncio.wait_for(second_replacement_stream.started.wait(), timeout=1)
+
+        second_replacement_stream.end()
+        await asyncio.wait_for(final_replacement_stream.started.wait(), timeout=1)
+
+    assert create_stream.call_count == 4
+    assert audio_input._stream is final_replacement_stream
+    assert audio_input._publication is publication
+
+    await audio_input.aclose()
+
+
+@pytest.mark.asyncio
+async def test_audio_input_counts_frames_received_while_detached_for_reattach_limit() -> None:
+    room = _FakeRoom()
+    audio_input = _make_audio_input_stream(room, noise_cancellation=None)
+    audio_input.set_participant("test-user")
+    audio_input.on_detached()
+    track, publication, participant = _make_track_available_args()
+    first_stream = _MockAudioStream()
+    replacement_stream = _MockAudioStream()
+    second_replacement_stream = _MockAudioStream()
+    frame = rtc.AudioFrame(bytes(480 * 2), 24000, 1, 480)
+
+    with patch(
+        "livekit.rtc.AudioStream.from_track",
+        side_effect=[first_stream, replacement_stream, second_replacement_stream],
+    ) as create_stream:
+        assert audio_input._on_track_available(track, publication, participant)
+        await asyncio.wait_for(first_stream.started.wait(), timeout=1)
+
+        first_stream.push_frame(frame)
+        first_stream.end()
+        await asyncio.wait_for(replacement_stream.started.wait(), timeout=1)
+
+        replacement_stream.end()
+        await asyncio.wait_for(second_replacement_stream.started.wait(), timeout=1)
+
+    assert create_stream.call_count == 3
+    assert audio_input._data_ch.qsize() == 1
+    assert audio_input._stream is second_replacement_stream
+    assert audio_input._publication is publication
+
+    await audio_input.aclose()
+
+
+@pytest.mark.asyncio
+async def test_audio_input_stops_reattaching_when_replacement_closes_without_frames() -> None:
+    room = _FakeRoom()
+    audio_input = _make_audio_input_stream(room, noise_cancellation=None)
+    audio_input.set_participant("test-user")
+    track, publication, participant = _make_track_available_args()
+    first_stream = _MockAudioStream()
+    replacement_stream = _MockAudioStream()
+
+    with patch(
+        "livekit.rtc.AudioStream.from_track",
+        side_effect=[first_stream, replacement_stream],
+    ) as create_stream:
+        assert audio_input._on_track_available(track, publication, participant)
+        await asyncio.wait_for(first_stream.started.wait(), timeout=1)
+
+        first_stream.end()
+        await asyncio.wait_for(replacement_stream.started.wait(), timeout=1)
+
+        replacement_stream.end()
+        assert audio_input._forward_atask is not None
+        await audio_input._forward_atask
+
+    assert create_stream.call_count == 2
+    assert audio_input._data_ch.qsize() == 1
+    assert audio_input._stream is None
+    assert audio_input._publication is None
+
+    await audio_input.aclose()
+
+
+@pytest.mark.asyncio
+async def test_audio_input_does_not_reattach_after_unsubscribe() -> None:
+    room = _FakeRoom()
+    audio_input = _make_audio_input_stream(room, noise_cancellation=None)
+    audio_input.set_participant("test-user")
+    track, publication, participant = _make_track_available_args()
+    rtc_stream = _MockAudioStream()
+
+    with patch("livekit.rtc.AudioStream.from_track", return_value=rtc_stream) as create_stream:
+        assert audio_input._on_track_available(track, publication, participant)
+        await asyncio.wait_for(rtc_stream.started.wait(), timeout=1)
+
+        publication.track = None
+        publication.subscribed = False
+        rtc_stream.end()
+        assert audio_input._forward_atask is not None
+        await audio_input._forward_atask
+
+    create_stream.assert_called_once()
+    assert audio_input._stream is None
+    assert audio_input._publication is None
+
+    await audio_input.aclose()
+
+
+@pytest.mark.asyncio
+async def test_audio_input_does_not_reattach_during_close() -> None:
+    room = _FakeRoom()
+    audio_input = _make_audio_input_stream(room, noise_cancellation=None)
+    audio_input.set_participant("test-user")
+    track, publication, participant = _make_track_available_args()
+    rtc_stream = _MockAudioStream()
+
+    with patch("livekit.rtc.AudioStream.from_track", return_value=rtc_stream) as create_stream:
+        assert audio_input._on_track_available(track, publication, participant)
+        await asyncio.wait_for(rtc_stream.started.wait(), timeout=1)
+
+        await audio_input.aclose()
+
+    create_stream.assert_called_once()
+    assert audio_input._stream is None
+    assert audio_input._publication is None
+
+
 @pytest.mark.asyncio
 async def test_direct_processor_lifecycle() -> None:
     """Direct FrameProcessor survives track transitions and is only closed on aclose()."""
@@ -342,6 +515,7 @@ async def test_selector_processor_track_disappears() -> None:
     assert stream._processor is processor
 
     # track unpublished with no replacement
+    participant.track_publications.clear()
     stream._on_track_unavailable(publication, participant)
 
     assert processor.close_calls == 1