diff --git a/assemblyai/dictation/v1/client.py b/assemblyai/dictation/v1/client.py index ec025697..63b3d119 100644 --- a/assemblyai/dictation/v1/client.py +++ b/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() diff --git a/assemblyai/sync/v1/client.py b/assemblyai/sync/v1/client.py index 5a688506..5c81f663 100644 --- a/assemblyai/sync/v1/client.py +++ b/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() diff --git a/tests/unit/test_dictation.py b/tests/unit/test_dictation.py index ba8fafe4..7a61769a 100644 --- a/tests/unit/test_dictation.py +++ b/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) diff --git a/tests/unit/test_sync_live.py b/tests/unit/test_sync_live.py index c7cbaebf..a850629c 100644 --- a/tests/unit/test_sync_live.py +++ b/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)