fix: prevent live audio loss during close

assemblyai/dictation/v1/client.py

@@ -3,6 +3,7 @@
 import concurrent.futures
 import os
 import queue
+import threading
 from typing import Any, Callable, Iterator, Optional
 
 import httpx
@@ -60,6 +61,7 @@ def __init__(
                 future that will hold its outcome. Supplied by the transcriber.
         """
         self._queue: "queue.Queue[Optional[bytes]]" = queue.Queue()
+        self._lock = threading.Lock()
         self._closed = False
         self._aborted = False
         self._future = start(self._chunks())
@@ -77,14 +79,16 @@ def _chunks(self) -> Iterator[bytes]:
     def closed(self) -> bool:
         """Whether the audio has ended, by `close()`, `result()` or `abort()`."""
 
-        return self._closed
+        with self._lock:
+            return self._closed
 
     def write(self, chunk: bytes) -> None:
         """
         Queues a piece of audio for upload.
 
         Safe to call from any thread, including an audio library's capture
-        callback; it never blocks.
+        callback. It can briefly wait for thread synchronization, but it does
+        not wait for network upload.
 
         Args:
             chunk: `bytes`, `bytearray` or `memoryview` of audio, in order.
@@ -93,21 +97,31 @@ def write(self, chunk: bytes) -> None:
             TypeError: if `chunk` is not bytes-like.
             RuntimeError: if the session is closed.
         """
-        if self._closed:
-            raise RuntimeError(
-                "the live session is closed; no more audio can be written"
-            )
+        with self._lock:
+            if self._closed:
+                raise RuntimeError(
+                    "the live session is closed; no more audio can be written"
+                )
+
+        chunk = _as_bytes(chunk)
 
-        self._queue.put(_as_bytes(chunk))
+        with self._lock:
+            if self._closed:
+                raise RuntimeError(
+                    "the live session is closed; no more audio can be written"
+                )
+            self._queue.put(chunk)
 
     def close(self) -> None:
         """
         Ends the audio. The server transcribes what was sent; `result()`
-        returns it. Idempotent, and never blocks.
+        returns it. Idempotent. It can briefly wait for thread synchronization,
+        but it does not wait for network upload.
         """
-        if not self._closed:
-            self._closed = True
-            self._queue.put(None)
+        with self._lock:
+            if not self._closed:
+                self._closed = True
+                self._queue.put(None)
 
     def result(self) -> DictationResponse:
         """
@@ -118,8 +132,11 @@ def result(self) -> DictationResponse:
                 server sent while the upload was still in flight.
             RuntimeError: if the session was aborted.
         """
-        if self._aborted:
-            raise RuntimeError("the live session was aborted; there is no result")
+        with self._lock:
+            if self._aborted:
+                raise RuntimeError(
+                    "the live session was aborted; there is no result"
+                )
 
         self.close()
 
@@ -135,13 +152,14 @@ def abort(self) -> None:
         the connection. After `close()` the upload is already complete, so
         aborting then waits out the in-flight response instead.
         """
-        if self._aborted or self._future.done():
-            return
+        with self._lock:
+            if self._aborted or self._future.done():
+                return
 
-        self._aborted = True
-        if not self._closed:
-            self._closed = True
-            self._queue.put(None)
+            self._aborted = True
+            if not self._closed:
+                self._closed = True
+                self._queue.put(None)
 
         try:
             self._future.result()

assemblyai/sync/v1/client.py

@@ -3,6 +3,7 @@
 import concurrent.futures
 import os
 import queue
+import threading
 from typing import Any, Callable, Iterator, Optional
 
 import httpx
@@ -65,6 +66,7 @@ def __init__(
                 future that will hold its outcome. Supplied by the transcriber.
         """
         self._queue: "queue.Queue[Optional[bytes]]" = queue.Queue()
+        self._lock = threading.Lock()
         self._closed = False
         self._aborted = False
         self._future = start(self._chunks())
@@ -82,14 +84,16 @@ def _chunks(self) -> Iterator[bytes]:
     def closed(self) -> bool:
         """Whether the audio has ended, by `close()`, `result()` or `abort()`."""
 
-        return self._closed
+        with self._lock:
+            return self._closed
 
     def write(self, chunk: bytes) -> None:
         """
         Queues a piece of audio for upload.
 
         Safe to call from any thread, including an audio library's capture
-        callback; it never blocks.
+        callback. It can briefly wait for thread synchronization, but it does
+        not wait for network upload.
 
         Args:
             chunk: `bytes`, `bytearray` or `memoryview` of audio, in order.
@@ -98,21 +102,31 @@ def write(self, chunk: bytes) -> None:
             TypeError: if `chunk` is not bytes-like.
             RuntimeError: if the session is closed.
         """
-        if self._closed:
-            raise RuntimeError(
-                "the live session is closed; no more audio can be written"
-            )
+        with self._lock:
+            if self._closed:
+                raise RuntimeError(
+                    "the live session is closed; no more audio can be written"
+                )
+
+        chunk = _as_bytes(chunk)
 
-        self._queue.put(_as_bytes(chunk))
+        with self._lock:
+            if self._closed:
+                raise RuntimeError(
+                    "the live session is closed; no more audio can be written"
+                )
+            self._queue.put(chunk)
 
     def close(self) -> None:
         """
         Ends the audio. The server transcribes what was sent; `result()`
-        returns it. Idempotent, and never blocks.
+        returns it. Idempotent. It can briefly wait for thread synchronization,
+        but it does not wait for network upload.
         """
-        if not self._closed:
-            self._closed = True
-            self._queue.put(None)
+        with self._lock:
+            if not self._closed:
+                self._closed = True
+                self._queue.put(None)
 
     def result(self) -> types.SyncTranscriptResponse:
         """
@@ -123,8 +137,11 @@ def result(self) -> types.SyncTranscriptResponse:
                 the server sent while the upload was still in flight.
             RuntimeError: if the session was aborted.
         """
-        if self._aborted:
-            raise RuntimeError("the live session was aborted; there is no result")
+        with self._lock:
+            if self._aborted:
+                raise RuntimeError(
+                    "the live session was aborted; there is no result"
+                )
 
         self.close()
 
@@ -140,13 +157,14 @@ def abort(self) -> None:
         the connection. After `close()` the upload is already complete, so
         aborting then waits out the in-flight response instead.
         """
-        if self._aborted or self._future.done():
-            return
+        with self._lock:
+            if self._aborted or self._future.done():
+                return
 
-        self._aborted = True
-        if not self._closed:
-            self._closed = True
-            self._queue.put(None)
+            self._aborted = True
+            if not self._closed:
+                self._closed = True
+                self._queue.put(None)
 
         try:
             self._future.result()

tests/unit/test_dictation.py

@@ -22,7 +22,7 @@
     _DICTATION_MAX_STT_PROMPT_LEN,
 )
 from assemblyai._multipart import _STREAM_READ_SIZE
-from assemblyai.dictation.v1 import api
+from assemblyai.dictation.v1 import api, client as dictation_client
 
 aai.settings.api_key = "test"
 
@@ -943,6 +943,50 @@ def test_live_session_rejects_writes_after_close(monkeypatch):
         session.result()
 
 
+def test_live_session_rejects_write_that_loses_race_with_close(monkeypatch):
+    # Given a writer paused during byte conversion after its initial state check
+    seen = _fake_live(monkeypatch)
+    conversion_started = threading.Event()
+    resume_conversion = threading.Event()
+    original_as_bytes = dictation_client._as_bytes
+
+    def pause_conversion(chunk):
+        conversion_started.set()
+        assert resume_conversion.wait(timeout=2.0)
+        return original_as_bytes(chunk)
+
+    monkeypatch.setattr(dictation_client, "_as_bytes", pause_conversion)
+    errors = []
+
+    with aai.DictationTranscriber() as transcriber:
+        session = transcriber.open_live()
+        assert seen["called"].wait(timeout=2.0)
+
+        def write_audio():
+            try:
+                session.write(b"final-audio-frame")
+            except Exception as error:
+                errors.append(error)
+
+        writer = threading.Thread(target=write_audio)
+        writer.start()
+        try:
+            assert conversion_started.wait(timeout=2.0)
+
+            # When close wins before the writer queues its audio
+            session.close()
+            session.result()
+        finally:
+            resume_conversion.set()
+            writer.join(timeout=2.0)
+
+    # Then the write reports that the session closed and no audio is stranded
+    assert not writer.is_alive()
+    assert len(errors) == 1
+    assert isinstance(errors[0], RuntimeError)
+    assert seen["chunks"] == []
+
+
 def test_live_session_rejects_text_at_write_time(monkeypatch):
     # Given a producer handing out str
     _fake_live(monkeypatch)

tests/unit/test_sync_live.py

@@ -18,7 +18,7 @@
 
 import assemblyai as aai
 from assemblyai._multipart import _STREAM_READ_SIZE
-from assemblyai.sync.v1 import api, async_api
+from assemblyai.sync.v1 import api, async_api, client as sync_client
 
 aai.settings.api_key = "test"
 
@@ -586,6 +586,50 @@ def test_live_session_rejects_writes_after_close(monkeypatch):
         session.result()
 
 
+def test_live_session_rejects_write_that_loses_race_with_close(monkeypatch):
+    # Given a writer paused during byte conversion after its initial state check
+    seen = _fake_live(monkeypatch, api, _drain)
+    conversion_started = threading.Event()
+    resume_conversion = threading.Event()
+    original_as_bytes = sync_client._as_bytes
+
+    def pause_conversion(chunk):
+        conversion_started.set()
+        assert resume_conversion.wait(timeout=2.0)
+        return original_as_bytes(chunk)
+
+    monkeypatch.setattr(sync_client, "_as_bytes", pause_conversion)
+    errors = []
+
+    with aai.SyncTranscriber() as transcriber:
+        session = transcriber.open_live()
+        assert seen["called"].wait(timeout=2.0)
+
+        def write_audio():
+            try:
+                session.write(b"final-audio-frame")
+            except Exception as error:
+                errors.append(error)
+
+        writer = threading.Thread(target=write_audio)
+        writer.start()
+        try:
+            assert conversion_started.wait(timeout=2.0)
+
+            # When close wins before the writer queues its audio
+            session.close()
+            session.result()
+        finally:
+            resume_conversion.set()
+            writer.join(timeout=2.0)
+
+    # Then the write reports that the session closed and no audio is stranded
+    assert not writer.is_alive()
+    assert len(errors) == 1
+    assert isinstance(errors[0], RuntimeError)
+    assert seen["chunks"] == []
+
+
 def test_live_session_rejects_text_at_write_time(monkeypatch):
     # Given a producer handing out str
     _fake_live(monkeypatch, api, _drain)