feat: add RobustMicrophone to handle hardware disconnects (#6076)

livekit-agents/livekit/agents/utils/__init__.py

@@ -10,6 +10,7 @@
 from .misc import is_dev_mode, is_given, is_hosted, nodename, shortuuid, time_ms
 from .moving_average import MovingAverage
 from .participant import wait_for_agent, wait_for_participant, wait_for_track_publication
+from .robust_microphone import RobustMicrophone
 
 EventEmitter = rtc.EventEmitter
 
@@ -41,6 +42,7 @@
     "wait_for_participant",
     "wait_for_track_publication",
     "resolve_env_var",
+    "RobustMicrophone",
 ]
 
 # Cleanup docs of unexported modules

livekit-agents/livekit/agents/utils/robust_microphone.py

@@ -0,0 +1,110 @@
+import asyncio
+import time
+from typing import Any
+
+from livekit import rtc
+from livekit.agents.log import logger
+from livekit.agents.utils import aio
+
+
+class RobustMicrophone:
+    """
+    A robust microphone capture utility that wraps rtc.MediaDevices().open_input()
+    and automatically restarts the audio stream if it stalls (e.g. if the microphone
+    cable becomes loose).
+
+    This class exposes an `rtc.AudioSource` as `.source` which you can use to
+    create your LocalAudioTrack.
+    """
+
+    def __init__(
+        self,
+        *,
+        sample_rate: int = 48000,
+        num_channels: int = 1,
+        stall_timeout: float = 2.0,
+        **open_input_kwargs: Any,
+    ) -> None:
+        self._sample_rate = sample_rate
+        self._num_channels = num_channels
+        self._stall_timeout = stall_timeout
+        self._kwargs = open_input_kwargs
+
+        self._devices = rtc.MediaDevices()
+        self._source = rtc.AudioSource(self._sample_rate, self._num_channels)
+
+        self._mic_obj: Any = None
+        self._mic_track: rtc.LocalAudioTrack | None = None
+        self._mic_stream: rtc.AudioStream | None = None
+
+        self._running = False
+        self._monitor_task: asyncio.Task[None] | None = None
+        self._last_frame_time: float = time.monotonic()
+
+    @property
+    def source(self) -> rtc.AudioSource:
+        return self._source
+
+    def start(self) -> None:
+        if self._running:
+            return
+        self._running = True
+        self._monitor_task = asyncio.create_task(self._monitor_loop())
+
+    async def aclose(self) -> None:
+        if not self._running:
+            return
+        self._running = False
+        if self._monitor_task:
+            await aio.cancel_and_wait(self._monitor_task)
+        await self._close_internal()
+
+    async def _close_internal(self) -> None:
+        if self._mic_stream:
+            await self._mic_stream.aclose()
+            self._mic_stream = None
+        self._mic_track = None
+        self._mic_obj = None
+
+    def _start_internal(self) -> None:
+        # Provide defaults for sample_rate and num_channels if not provided by user
+        kwargs = dict(self._kwargs)
+        if "sample_rate" not in kwargs:
+            kwargs["sample_rate"] = self._sample_rate
+        if "num_channels" not in kwargs:
+            kwargs["num_channels"] = self._num_channels
+
+        self._mic_obj = self._devices.open_input(**kwargs)
+        self._mic_track = rtc.LocalAudioTrack.create_audio_track("robust-mic-internal", self._mic_obj.source)
+        self._mic_stream = rtc.AudioStream.from_track(self._mic_track)
+        self._last_frame_time = time.monotonic()
+
+    async def _monitor_loop(self) -> None:
+        self._start_internal()
+
+        while self._running:
+            try:
+                assert self._mic_stream is not None
+
+                # Wait for the next audio event with a timeout
+                event = await asyncio.wait_for(
+                    self._mic_stream.__anext__(), timeout=self._stall_timeout
+                )
+                self._last_frame_time = time.monotonic()
+
+                # Forward the captured frame to our own source
+                await self._source.capture_frame(event.frame)
+
+            except asyncio.TimeoutError:
+                # Stall detected!
+                logger.warning(
+                    f"RobustMicrophone: No audio frames received for {self._stall_timeout}s. Restarting microphone..."
+                )
+                await self._close_internal()
+                await asyncio.sleep(0.5)  # Brief pause before reconnecting
+                self._start_internal()
+            except Exception as e:
+                logger.error(f"RobustMicrophone error in monitor loop: {e}", exc_info=e)
+                await asyncio.sleep(1.0)
+                await self._close_internal()
+                self._start_internal()

tests/test_robust_microphone.py

@@ -0,0 +1,70 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from livekit import rtc
+from livekit.agents.utils.robust_microphone import RobustMicrophone
+
+pytestmark = pytest.mark.unit
+
+
+@pytest.mark.asyncio
+async def test_robust_microphone_startup_and_shutdown():
+    with patch("livekit.rtc.MediaDevices", autospec=True) as mock_media_devices, \
+         patch("livekit.rtc.AudioSource", autospec=True) as mock_audio_source, \
+         patch("livekit.rtc.LocalAudioTrack.create_audio_track", autospec=True) as mock_create_track, \
+         patch("livekit.rtc.AudioStream.from_track", autospec=True) as mock_from_track:
+
+        # Mock the stream iterator to just hang (we'll shut down before it times out)
+        mock_stream = AsyncMock()
+        mock_stream.__anext__.side_effect = asyncio.TimeoutError
+        mock_from_track.return_value = mock_stream
+
+        mic = RobustMicrophone(stall_timeout=10.0)
+        mic.start()
+
+        # Give it a moment to start the internal stream
+        await asyncio.sleep(0.1)
+
+        assert mock_media_devices.called
+        assert mock_create_track.called
+        assert mock_from_track.called
+
+        await mic.aclose()
+
+
+@pytest.mark.asyncio
+async def test_robust_microphone_restarts_on_stall():
+    with patch("livekit.rtc.MediaDevices", autospec=True) as mock_media_devices, \
+         patch("livekit.rtc.AudioSource", autospec=True) as mock_audio_source, \
+         patch("livekit.rtc.LocalAudioTrack.create_audio_track", autospec=True) as mock_create_track, \
+         patch("livekit.rtc.AudioStream.from_track", autospec=True) as mock_from_track:
+
+        mock_stream = AsyncMock()
+        mock_stream.aclose = AsyncMock()
+
+        # First call times out immediately, causing a restart
+        # Second call hangs so we can shut down
+        call_count = 0
+        async def mock_anext():
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise asyncio.TimeoutError()
+            else:
+                await asyncio.sleep(10)
+                return MagicMock()
+
+        mock_stream.__anext__ = mock_anext
+        mock_from_track.return_value = mock_stream
+
+        mic = RobustMicrophone(stall_timeout=0.1)
+        mic.start()
+
+        # Wait for the timeout and restart to happen
+        await asyncio.sleep(0.5)
+
+        # MediaDevices.open_input should have been called twice (initial + restart)
+        assert mock_from_track.call_count >= 2
+
+        await mic.aclose()