feat(extras): add pause()/resume() to MicrophoneStream

assemblyai/extras.py

@@ -1,3 +1,4 @@
+import threading
 import time
 from typing import BinaryIO, Generator, Optional
 from warnings import warn
@@ -22,6 +23,26 @@ def __init__(
 
 
 class MicrophoneStream:
+    """
+    A synchronous iterator of raw microphone audio chunks (~100ms each).
+
+    Pass it to ``StreamingClient.stream()``. Don't pass it directly to
+    ``AsyncStreamingClient.stream()`` — each read blocks for ~100ms, which
+    stalls the event loop and starves the read task (the same problem as
+    ``stream_file``). Wrap the blocking reads in a thread instead::
+
+        async def mic_chunks(mic: MicrophoneStream):
+            while True:
+                chunk = await asyncio.to_thread(next, mic, None)
+                if chunk is None:
+                    return
+                yield chunk
+
+    :meth:`pause`, :meth:`resume`, and :meth:`close` are thread-safe: they may
+    be called from another thread (or the event loop) while a read is in
+    flight.
+    """
+
     def __init__(self, sample_rate: int = 44_100, device_index: Optional[int] = None):
         """
         Creates a stream of audio from the microphone.
@@ -49,6 +70,12 @@ def __init__(self, sample_rate: int = 44_100, device_index: Optional[int] = None
         )
 
         self._open = True
+        self._paused = False
+        self._closed = False
+        # Serializes device reads against teardown: closing the PortAudio
+        # stream while another thread is blocked in read() deadlocks, so
+        # close() waits for any in-flight read to finish first.
+        self._lock = threading.Lock()
 
     def __iter__(self):
         """
@@ -60,27 +87,85 @@ def __iter__(self):
     def __next__(self):
         """
         Reads a chunk of audio from the microphone.
+
+        While paused (see :meth:`pause`), the microphone is still drained so its
+        input buffer doesn't overflow, but silence is yielded in place of the
+        captured audio. This keeps the streaming session alive (avoiding an idle
+        timeout) without forwarding what the mic picks up.
         """
         if not self._open:
             raise StopIteration
 
-        try:
-            return self._stream.read(self._chunk_size)
-        except KeyboardInterrupt:
-            raise StopIteration
+        with self._lock:
+            # Re-check after acquiring: a concurrent close() may have torn the
+            # stream down while this thread waited for the lock.
+            if not self._open:
+                raise StopIteration
+
+            try:
+                data = self._stream.read(self._chunk_size)
+            except KeyboardInterrupt:
+                raise StopIteration
+
+        if self._paused:
+            return b"\x00" * len(data)
+
+        return data
+
+    def pause(self) -> None:
+        """
+        Pause forwarding microphone audio.
+
+        The stream stays open and keeps reading from the microphone so its
+        internal buffer doesn't overflow, but :meth:`__next__` yields silence
+        instead of the captured audio. Useful for muting the mic while a voice
+        agent is speaking so its own output (e.g. TTS played back through the
+        speakers) isn't transcribed and fed into a feedback loop.
+
+        Thread-safe: may be called from any thread (or an asyncio event loop)
+        while another thread is reading from the stream. Takes effect within
+        one chunk (~100ms).
+        """
+        self._paused = True
+
+    def resume(self) -> None:
+        """
+        Resume forwarding live microphone audio after :meth:`pause`.
+
+        Thread-safe, like :meth:`pause`.
+        """
+        self._paused = False
+
+    @property
+    def paused(self) -> bool:
+        """Whether the stream is currently yielding silence (see :meth:`pause`)."""
+        return self._paused
 
     def close(self):
         """
         Closes the stream.
+
+        Thread-safe and idempotent: may be called from any thread, including
+        while another thread is blocked in a read. Teardown waits for the
+        in-flight chunk (~100ms) to finish first — closing the PortAudio
+        stream mid-read deadlocks the reading thread. After close(), the
+        iterator raises ``StopIteration``.
         """
 
+        # Stop new reads before waiting on the in-flight one, so the reader
+        # can't re-acquire the lock ahead of teardown.
         self._open = False
 
-        if self._stream.is_active():
-            self._stream.stop_stream()
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+
+            if self._stream.is_active():
+                self._stream.stop_stream()
 
-        self._stream.close()
-        self._pyaudio.terminate()
+            self._stream.close()
+            self._pyaudio.terminate()
 
 
 def stream_file(

tests/unit/test_extras.py

@@ -1,5 +1,7 @@
 from unittest.mock import mock_open, patch
 
+import pytest_mock
+
 import assemblyai as aai
 
 
@@ -72,3 +74,97 @@ def test_stream_file_exact_file():
 
     # Expecting one chunk
     assert len(chunks) == 1
+
+
+def test_microphone_stream_pause_resume(mocker: pytest_mock.MockerFixture):
+    """
+    A paused MicrophoneStream keeps draining the device but yields silence in
+    place of captured audio, and resumes forwarding live audio afterwards.
+    """
+    import pyaudio
+
+    live_chunk = b"\x01\x02\x03\x04" * 8
+    fake_stream = mocker.MagicMock()
+    fake_stream.read.return_value = live_chunk
+    mocker.patch.object(pyaudio.PyAudio, "open", return_value=fake_stream)
+
+    mic = aai.extras.MicrophoneStream(sample_rate=16000)
+
+    # Live: captured audio is forwarded unchanged.
+    assert mic.paused is False
+    assert next(mic) == live_chunk
+
+    # Paused: yields silence of the same length, but still reads from the device
+    # (so the input buffer keeps draining and the session stays alive).
+    mic.pause()
+    assert mic.paused is True
+    reads_before = fake_stream.read.call_count
+    assert next(mic) == b"\x00" * len(live_chunk)
+    assert fake_stream.read.call_count == reads_before + 1
+
+    # Resume: live audio is forwarded again.
+    mic.resume()
+    assert mic.paused is False
+    assert next(mic) == live_chunk
+
+
+def test_microphone_stream_close_during_read_is_thread_safe(
+    mocker: pytest_mock.MockerFixture,
+):
+    """
+    close() may be called from another thread while a read is in flight (the
+    natural companion to pause()/resume() in voice agents). Teardown must wait
+    for the in-flight chunk instead of closing the PortAudio stream mid-read,
+    which deadlocks the reading thread.
+    """
+    import threading
+    import time
+
+    import pyaudio
+
+    chunk = b"\x00" * 16
+    read_started = threading.Event()
+
+    def slow_read(num_frames, *args, **kwargs):
+        read_started.set()
+        time.sleep(0.05)  # simulate the ~100ms blocking device read
+        return chunk
+
+    fake_stream = mocker.MagicMock()
+    fake_stream.read.side_effect = slow_read
+    fake_stream.is_active.return_value = True
+    mocker.patch.object(pyaudio.PyAudio, "open", return_value=fake_stream)
+    terminate = mocker.patch.object(pyaudio.PyAudio, "terminate")
+
+    mic = aai.extras.MicrophoneStream(sample_rate=16000)
+
+    def consume():
+        for _ in mic:
+            pass
+
+    reader = threading.Thread(target=consume)
+    reader.start()
+    assert read_started.wait(timeout=2), "reader never reached the device read"
+
+    started = time.monotonic()
+    mic.close()  # concurrent with an in-flight read
+    close_duration = time.monotonic() - started
+
+    reader.join(timeout=2)
+    assert not reader.is_alive(), "reader thread did not exit after close()"
+    assert close_duration < 2, "close() blocked far longer than one chunk read"
+    # close() was called mid-read (after read_started, during the 50ms sleep),
+    # so a thread-safe close must have waited out the in-flight read rather
+    # than tearing the stream down underneath it.
+    assert close_duration >= 0.03, "close() did not wait for the in-flight read"
+
+    # Teardown ran exactly once, and only after the in-flight read finished.
+    fake_stream.stop_stream.assert_called_once()
+    fake_stream.close.assert_called_once()
+    terminate.assert_called_once()
+
+    # Idempotent: a second close() is a no-op, and iteration stays ended.
+    mic.close()
+    fake_stream.close.assert_called_once()
+    terminate.assert_called_once()
+    assert next(mic, None) is None