fix: reattach RoomIO audio input after stream closes

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
 
@@ -18,6 +18,13 @@
 
 T = TypeVar("T", bound=rtc.AudioFrame | rtc.VideoFrame)
 
+_AUDIO_STREAM_REATTACH_DELAY = 0.5
+
+
+class _StreamReadResult(NamedTuple):
+    received_frame: bool
+    forwarded_frame: bool
+
 
 class _ParticipantInputStream(Generic[T], ABC):
     """
@@ -45,7 +52,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 +125,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 +143,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 +156,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 +306,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 +347,58 @@ 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
+            if track is None or not publication.subscribed:
+                self._close_stream()
+                return read_result
+
+            if read_result.received_frame:
+                consecutive_empty_streams = 0
+            else:
+                consecutive_empty_streams += 1
+            if consecutive_empty_streams > 1:
+                self._close_stream()
+                logger.warning(
+                    "replacement audio stream closed before receiving frames; not reattaching",
+                    extra={"participant": participant.identity, "track_id": track.sid},
+                )
+                return read_result
+
+            await asyncio.sleep(_AUDIO_STREAM_REATTACH_DELAY)
+
+            if self._stream is not stream or self._publication is not publication:
+                return read_result
+
+            track = publication.track
+            self._close_stream()
+            if track is None or not publication.subscribed:
+                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