From 1cc5a4124e7ab914e7ad5f4812b771edcb02710a Mon Sep 17 00:00:00 2001 From: Ben Gotthold Date: Wed, 21 Oct 2026 22:54:43 +0000 Subject: [PATCH 01/21] feat(tts): WebSocket session client Adds `aai.TTSSession` and `aai.AsyncTTSSession` for the Text-to-Speech WebSocket API (`/v1/ws`): one connection, text in as it is written, audio out. `with` connects and waits for `Begin`; `speak(text)` yields one request's audio; `generate()` / `flush()` / `audio(flush_id)` stream text in and pipeline requests; `cancel()` is the barge-in; `on(TTSEvents.X, handler)` delivers the other frames. Exit sends `Terminate` and waits for `Termination`. Refusals and failures are `TTSError`. Temporary tokens go in the query (`token=`). Reuses `TTSConfig`, `TTSError` and the settings host. Co-Authored-By: Claude Fable 5.1 --- CLAUDE.md | 8 + README.md | 42 +++ assemblyai/__init__.py | 20 ++ assemblyai/__version__.py | 2 +- assemblyai/tts/__init__.py | 20 ++ assemblyai/tts/_base_ws.py | 92 ++++++ assemblyai/tts/v1/__init__.py | 22 ++ assemblyai/tts/v1/async_session.py | 377 +++++++++++++++++++++++ assemblyai/tts/v1/session.py | 444 +++++++++++++++++++++++++++ assemblyai/tts/v1/ws_models.py | 175 +++++++++++ tests/unit/test_tts_session.py | 425 +++++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 215 +++++++++++++ 12 files changed, 1841 insertions(+), 1 deletion(-) create mode 100644 assemblyai/tts/_base_ws.py create mode 100644 assemblyai/tts/v1/async_session.py create mode 100644 assemblyai/tts/v1/session.py create mode 100644 assemblyai/tts/v1/ws_models.py create mode 100644 tests/unit/test_tts_session.py create mode 100644 tests/unit/test_tts_session_async.py diff --git a/CLAUDE.md b/CLAUDE.md index a5c99eb2..4f0209b3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -54,6 +54,8 @@ aai.settings.api_key = "your-key" - `aai.DictationResponse` — Dictation result: `.text`, `.words` (`DictationWord` with `text` + `confidence`), `.confidence`, `.llm_response`, `.llm_error`, `.audio_duration_ms`, `.session_id`, `.request_time_ms`, `.sync_time_ms`, and `.final_text` (the LLM rewrite, falling back to `.text`) - `aai.TextToSpeech` — Text-to-Speech HTTP API: one text per request, audio out. Methods: `synthesize()` (complete audio as a `TTSResponse`), `stream()` (a `TTSStream` context manager yielding bytes as the API generates them), `close()` / `with` for a client built from `api_key=`. Host from `aai.settings.tts_base_url` (`https://tts.us.assemblyai.com`; set `https://tts.eu.assemblyai.com` for the EU zone) - `aai.AsyncTextToSpeech` — Asyncio counterpart of `TextToSpeech`; `synthesize()` is a coroutine and `stream()` returns an `AsyncTTSStream`. Owns an HTTP pool: use `async with` or `await aclose()` +- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin`; `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining; `cancel()` for barge-in; `keep_alive()`; `on(TTSEvents.X, handler)` for `Audio`, `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination`. Refusals raise `TTSError`. `token=` for temporary tokens +- `aai.AsyncTTSSession` — Asyncio counterpart: `async with`, coroutine sends, `async for chunk in session.speak(text)` - `aai.TTSConfig` — Text-to-Speech options: `voice_id` (a library voice such as `aai-jane`, or a cloned voice by UUID or slug), `model_id` (`assemblyai-tts-1`), `language`, `sample_rate`, `encoding` (`aai.TTSEncoding`: `pcm_s16le` default | `pcm_mulaw` | `pcm_alaw` | `mp3`), `container` (`aai.TTSContainer`: `raw` default | `wav`), `bit_rate`, `extend_biometric_consent` (owner-only consent renewal of a cloned voice). Rejects unknown fields - `aai.TTSResponse` — Text-to-Speech result: `.audio` (bytes) plus the response headers: `.request_id`, `.region`, `.model`, `.voice_id`, `.language`, `.encoding`, `.sample_rate`, `.container`, `.bit_rate`, `.content_type`, `.voice_expires_at`, `.billed_characters`, `.audio_duration_ms`. `TTSStream.info` carries the same fields without `audio`; `audio_duration_ms` is `None` on a stream with audio (`0` when the text had nothing to speak) - `aai.TTSError` — Raised on any Text-to-Speech refusal: `.status_code`, `.error_code` (stable API code; `biometric_terms_not_accepted` 403 until the account accepts the biometric data terms in the dashboard, `auth_unavailable` 503 retryable), `.param`, `.request_id`, `.retryable`, `.retry_after`, `.errors`, `.body` (the whole envelope) @@ -479,6 +481,12 @@ with tts.stream("Hello.", telephony) as speech: # A cloned or library voice by id tts.synthesize("Hello.", aai.TTSConfig(voice_id="aai-alba")) + +# A WebSocket session for a voice agent: no handshake per sentence, barge-in with cancel() +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in reply_sentences(): + for chunk in session.speak(sentence): + play(chunk) ``` ## LLM Gateway diff --git a/README.md b/README.md index f6168a32..594faa44 100644 --- a/README.md +++ b/README.md @@ -1013,6 +1013,48 @@ asyncio.run(main()) +
+ Stream text in and audio out over a WebSocket session (voice agents) + +`aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in llm_reply_sentences(): + for chunk in session.speak(sentence): # raw 24 kHz 16-bit mono PCM + play(chunk) + +print(session.termination.audio_duration_seconds, "s spoken") +``` + +Streaming the text in as it is generated, with a barge-in: + +```python +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for token in llm_tokens(): + session.generate(token) # joined exactly as sent: include the spaces + request = session.flush() # the request's flush_id; audio may start before this + + for chunk in session.audio(request): + if user_started_talking(): + session.cancel() # stops this request; the iteration ends + play(chunk) +``` + +Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. A `Generate` frame takes at most 2,000 characters; split longer text across calls. The session ends after an hour, or after `inactivity_timeout` seconds without a frame; send `keep_alive()` between sparse requests. + +Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries` (word timings per request, with `word_boundaries=True`), `TTSEvents.Cancelled`, `TTSEvents.Termination`, `TTSEvents.Error`. An `Audio` handler receives the audio frames instead of the pull iterators. A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the audio iteration and reaches the `Error` handler. + +For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`). + +The asyncio twin is `aai.AsyncTTSSession`: `async with`, coroutine `generate()`/`flush()`/`cancel()`, and `async for chunk in session.speak(text)`. + +
+
Renew a cloned voice's consent diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index f625c97c..a7653532 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -31,15 +31,25 @@ from .transcriber import Transcriber, Transcript, TranscriptGroup from .tts.v1 import ( AsyncTextToSpeech, + AsyncTTSSession, AsyncTTSStream, TextToSpeech, + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, TTSConfig, TTSContainer, TTSEncoding, TTSError, + TTSEvents, + TTSFlushDoneEvent, TTSResponse, TTSResponseInfo, + TTSSession, TTSStream, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, ) from .types import ( AssemblyAIError, @@ -140,6 +150,7 @@ def __getattr__(name: str): "AsyncDictationTranscriber", "AsyncLLMGateway", "AsyncSyncTranscriber", + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", "AsyncTranscriber", @@ -211,13 +222,22 @@ def __getattr__(name: str): "SyncTranscriptResponse", "SyncWord", "Timestamp", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", "TTSContainer", "TTSEncoding", "TTSError", + "TTSEvents", + "TTSFlushDoneEvent", "TTSResponse", "TTSResponseInfo", + "TTSSession", "TTSStream", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TextToSpeech", "Transcriber", "TranscriptionConfig", diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 14d9d2f5..29654eec 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.7.0" +__version__ = "1.8.0" diff --git a/assemblyai/tts/__init__.py b/assemblyai/tts/__init__.py index 4666a761..7c8d64bd 100644 --- a/assemblyai/tts/__init__.py +++ b/assemblyai/tts/__init__.py @@ -7,21 +7,41 @@ from .v1 import ( AsyncTextToSpeech, + AsyncTTSSession, AsyncTTSStream, TextToSpeech, + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, TTSConfig, TTSContainer, TTSEncoding, TTSError, + TTSEvents, + TTSFlushDoneEvent, TTSResponse, TTSResponseInfo, + TTSSession, TTSStream, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, ) __all__ = [ + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSSession", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TTSContainer", "TTSEncoding", "TTSError", diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py new file mode 100644 index 00000000..9ea66a46 --- /dev/null +++ b/assemblyai/tts/_base_ws.py @@ -0,0 +1,92 @@ +"""Helpers shared by the sync and async Text-to-Speech sessions.""" + +import platform +from typing import Any, Dict, Optional, Tuple +from urllib.parse import urlencode, urlsplit + +from .. import _audio +from ..__version__ import __version__ +from ..types import Settings +from .v1.models import TTSConfig + +# The route of the WebSocket API; the host comes from `settings.tts_base_url`. +WS_PATH = "/v1/ws" +# Characters a single `Generate` frame may carry (Unicode code points). +MAX_GENERATE_CHARS = 2000 +# Audio frames queued for pulling before the oldest is dropped: about 160 s +# at the API's ~80 ms per frame. Bounds memory when nothing reads them. +PULL_QUEUE_FRAMES = 2000 + + +def check_text_frame(text: object) -> None: + if not isinstance(text, str): + raise TypeError(f"text must be a str, not {type(text).__name__}") + if len(text) > MAX_GENERATE_CHARS: + raise ValueError( + f"text is {len(text)} characters, over the {MAX_GENERATE_CHARS}-character " + "limit of one Generate frame; split it across several generate() calls" + ) + + +def resolve_settings( + owner: str, base: Optional[Settings], api_key: Optional[str], token: Optional[str] +) -> Settings: + """ + The settings a session connects with: a copy, so nothing shared changes. + + `base` is a client's settings, else the global `assemblyai.settings`; + `api_key` replaces the copy's key. Without a key or a `token`, raises + the SDK's usual message. No HTTP client is built: a session needs none. + """ + + if base is None: + from .. import settings as _global + + base = _global + settings = base.model_copy() if hasattr(base, "model_copy") else base.copy() + if api_key is not None: + settings.api_key = api_key + if token is None and not settings.api_key: + raise ValueError( + "Please provide an API key: set the ASSEMBLYAI_API_KEY environment " + f"variable, set aai.settings.api_key, or pass api_key= or token= to {owner}." + ) + return settings + + +def _user_agent() -> str: + return ( + f"AssemblyAI/1.0 (sdk=Python/{__version__} " + f"runtime_env=Python/{platform.python_version()})" + ) + + +def build_uri( + base_url: str, + config: TTSConfig, + *, + token: Optional[str], + api_key: Optional[str], + word_boundaries: bool, + inactivity_timeout: Optional[int], +) -> Tuple[str, Dict[str, Optional[str]]]: + """The session URI and handshake headers from the settings and config.""" + + parts = urlsplit(base_url if "://" in base_url else f"https://{base_url}") + scheme = "ws" if parts.scheme in ("http", "ws") else "wss" + params: Dict[str, Any] = ( + _audio._config_to_json(config, exclude=("container", "bit_rate")) or {} + ) + params["word_boundaries"] = word_boundaries + if inactivity_timeout is not None: + params["inactivity_timeout"] = inactivity_timeout + if token is not None: + params["token"] = token + # The API reads booleans case-insensitively; send the canonical form. + query = urlencode( + {k: (str(v).lower() if isinstance(v, bool) else v) for k, v in params.items()} + ) + headers: Dict[str, Optional[str]] = {"User-Agent": _user_agent()} + if token is None and api_key: + headers["Authorization"] = api_key + return f"{scheme}://{parts.netloc}{WS_PATH}?{query}", headers diff --git a/assemblyai/tts/v1/__init__.py b/assemblyai/tts/v1/__init__.py index fd8224ed..ea39654a 100644 --- a/assemblyai/tts/v1/__init__.py +++ b/assemblyai/tts/v1/__init__.py @@ -1,6 +1,7 @@ from .api import TTSStream from .async_api import AsyncTTSStream from .async_client import AsyncTextToSpeech +from .async_session import AsyncTTSSession from .client import TextToSpeech from .models import ( TTSConfig, @@ -10,11 +11,32 @@ TTSResponse, TTSResponseInfo, ) +from .session import TTSSession +from .ws_models import ( + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, + TTSEvents, + TTSFlushDoneEvent, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, +) __all__ = [ + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSSession", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TTSContainer", "TTSEncoding", "TTSError", diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py new file mode 100644 index 00000000..603053c5 --- /dev/null +++ b/assemblyai/tts/v1/async_session.py @@ -0,0 +1,377 @@ +import asyncio +import json +import logging +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Tuple, Union + +from ... import _audio +from ... import async_client as _async_client +from .._base_ws import PULL_QUEUE_FRAMES, build_uri, check_text_frame, resolve_settings +from .models import TTSConfig, TTSError +from .ws_models import ( + TTSBeginEvent, + TTSEvents, + TTSTerminationEvent, + error_from_close, + error_from_frame, + parse_event, +) + +__all__ = ["AsyncTTSSession"] + +logger = logging.getLogger(__name__) + +_Handler = Callable[[Any], Any] +_Item = Tuple[str, Any, Any] + + +def websocket_connect_async( + uri: str, additional_headers: Dict[str, Optional[str]] +) -> Any: + """Opens an asyncio websocket; the streaming client's seam, imported + lazily so `import assemblyai` does not load websockets. Tests patch this.""" + + from ...streaming.v3.async_client import websocket_connect_async as _connect + + return _connect(uri, additional_headers=additional_headers) + + +class AsyncTTSSession: + """ + A Text-to-Speech WebSocket session for asyncio: the counterpart of + `TTSSession`. + + Open it with `async with`, which connects and waits for `Begin`. + `generate()`, `flush()`, `cancel()` and `keep_alive()` are coroutines; + `speak()` and `audio()` are async iterators of audio bytes. Handlers + registered with `on()` may be plain functions or coroutines. + + Example: + ```python + async with aai.AsyncTTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + async for chunk in session.speak("Hello there."): + await play(chunk) + ``` + """ + + def __init__( + self, + config: Optional[TTSConfig] = None, + *, + client: Optional[_async_client.AsyncClient] = None, + api_key: Optional[str] = None, + token: Optional[str] = None, + word_boundaries: bool = False, + inactivity_timeout: Optional[int] = None, + connect_timeout: float = 10.0, + terminate_timeout: float = 10.0, + ) -> None: + """Creates a session; see `TTSSession.__init__` for the arguments.""" + + _audio.check_config(type(self).__name__, config, TTSConfig) + self.config = config or TTSConfig() + self._settings = resolve_settings( + type(self).__name__, client.settings if client else None, api_key, token + ) + self._token = token + self._word_boundaries = word_boundaries + self._inactivity_timeout = inactivity_timeout + self._connect_timeout = connect_timeout + self._terminate_timeout = terminate_timeout + + self._handlers: Dict[str, List[_Handler]] = {} + self._ws: Any = None + # Created in `connect()`, so they bind to the running loop. + self._reader: Optional["asyncio.Task[None]"] = None + self._begun: Any = None + self._stopped: Any = None + self._queue: Any = None + self._next_flush_id = 0 + self._cancel_ids: List[int] = [] + self._cancelled_before = 0 + self._drop_warned = False + self.begin: Optional[TTSBeginEvent] = None + self.termination: Optional[TTSTerminationEvent] = None + self.error: Optional[TTSError] = None + + def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: + """Registers `handler(event)`, a function or coroutine function.""" + + self._handlers.setdefault(str(event), []).append(handler) + + async def connect(self) -> TTSBeginEvent: + """Opens the socket and waits for `Begin`; see `TTSSession.connect`.""" + + if self._ws is not None: + raise RuntimeError("the session is already connected") + self._begun = asyncio.Event() + self._stopped = asyncio.Event() + self._queue = asyncio.Queue(maxsize=PULL_QUEUE_FRAMES) + self._next_flush_id = 0 + self._cancel_ids = [] + self._cancelled_before = 0 + self._drop_warned = False + self.begin = None + self.termination = None + self.error = None + + uri, headers = build_uri( + self._settings.tts_base_url, + self.config, + token=self._token, + api_key=self._settings.api_key, + word_boundaries=self._word_boundaries, + inactivity_timeout=self._inactivity_timeout, + ) + try: + self._ws = await asyncio.wait_for( + websocket_connect_async(uri, headers), timeout=self._connect_timeout + ) + except Exception as exc: + status = getattr(getattr(exc, "response", None), "status_code", None) + raise TTSError( + f"WebSocket handshake failed: {exc}", status_code=status + ) from exc + + self._reader = asyncio.create_task(self._read_loop()) + try: + await asyncio.wait_for(self._begun.wait(), timeout=self._connect_timeout) + except asyncio.TimeoutError: + await self.close() + raise TTSError("The session did not begin in time") from None + if self.error is not None: + await self.close() + raise self.error + if self.begin is None: + await self.close() + raise TTSError("The session closed before it began") + return self.begin + + async def terminate(self) -> Optional[TTSTerminationEvent]: + """Ends the session cleanly; see `TTSSession.terminate`.""" + + if self._ws is None: + return None + if not self._stopped.is_set(): + try: + await self._send({"type": "Terminate"}) + except Exception: + pass + try: + await asyncio.wait_for( + self._stopped.wait(), timeout=self._terminate_timeout + ) + except asyncio.TimeoutError: + pass + await self.close() + return self.termination + + async def close(self) -> None: + """Closes the socket without `Terminate`.""" + + ws, self._ws = self._ws, None + if self._stopped is not None: + self._stopped.set() + if self._begun is not None: + self._begun.set() + if ws is not None: + try: + await ws.close() + except Exception: + pass + if self._queue is not None: + self._put(("closed", None, None)) + reader = self._reader + if reader is not None and reader is not asyncio.current_task(): + try: + await asyncio.wait_for(reader, timeout=self._terminate_timeout) + except (asyncio.TimeoutError, asyncio.CancelledError): + reader.cancel() + + async def __aenter__(self) -> "AsyncTTSSession": + await self.connect() + return self + + async def __aexit__(self, exc_type: Any, *exc: object) -> None: + if exc_type is None: + await self.terminate() + else: + await self.close() + + async def generate(self, text: str) -> None: + """Appends `text` to the current request; see `TTSSession.generate`.""" + + check_text_frame(text) + await self._send({"type": "Generate", "text": text}) + + async def flush(self) -> int: + """Concludes the current request and returns its `flush_id`.""" + + flush_id = self._next_flush_id + self._next_flush_id += 1 + await self._send({"type": "Flush"}) + return flush_id + + async def cancel(self) -> None: + """Barge-in; see `TTSSession.cancel`.""" + + self._cancel_ids.append(self._next_flush_id) + self._next_flush_id += 1 + await self._send({"type": "Cancel"}) + + async def keep_alive(self) -> None: + await self._send({"type": "KeepAlive"}) + + async def speak(self, text: str) -> AsyncIterator[bytes]: + """Speaks one text and yields its audio; see `TTSSession.speak`.""" + + await self.generate(text) + flush_id = await self.flush() + async for chunk in self.audio(flush_id): + yield chunk + + async def audio(self, flush_id: int) -> AsyncIterator[bytes]: + """Yields the audio of request `flush_id`; see `TTSSession.audio`.""" + + if self._queue is None: + raise RuntimeError("the session is not open") + while True: + if flush_id < self._cancelled_before: # dropped by a Cancel + return + try: + kind, ident, payload = await asyncio.wait_for( + self._queue.get(), timeout=0.5 + ) + except asyncio.TimeoutError: + if self._stopped.is_set(): + if self.error is not None: + raise self.error + return + continue + if kind == "audio": + if ident == flush_id: + yield payload + elif kind == "done": + if ident == flush_id: + return + elif kind == "cancelled": + if ident > flush_id: + return + elif kind == "error": + raise payload + elif kind == "closed": + if self.error is not None: + raise self.error + return + + def _put(self, item: _Item) -> None: + """Queues an item for `audio()`, dropping the oldest audio when full.""" + + while True: + try: + self._queue.put_nowait(item) + return + except asyncio.QueueFull: + try: + dropped = self._queue.get_nowait() + except asyncio.QueueEmpty: + continue + if dropped[0] == "audio" and not self._drop_warned: + self._drop_warned = True + logger.warning( + "AsyncTTSSession: dropping audio nobody reads; iterate " + "audio() or register an Audio handler" + ) + + async def _fail(self, error: TTSError) -> None: + if self.error is None: + self.error = error + await self._emit(TTSEvents.Error, error) + self._put(("error", None, self.error)) + self._stopped.set() + self._begun.set() + + async def _send(self, frame: Dict[str, Any]) -> None: + if self._ws is None or self._stopped.is_set(): + if self.error is not None: + raise self.error + raise RuntimeError("the session is not open") + await self._ws.send(json.dumps(frame)) + + async def _read_loop(self) -> None: + from websockets.exceptions import ConnectionClosed + + ws = self._ws + while not self._stopped.is_set() and ws is not None: + try: + raw = await ws.recv() + except ConnectionClosed as exc: + await self._on_closed(exc) + return + except asyncio.CancelledError: + raise + except Exception as exc: + if not self._stopped.is_set(): + await self._on_closed(exc) + return + if isinstance(raw, bytes): + continue + try: + data = json.loads(raw) + except ValueError: + logger.warning("ignoring a non-JSON frame") + continue + if isinstance(data, dict): + try: + await self._dispatch(data) + except Exception as exc: # a frame this SDK cannot read + await self._fail(TTSError(f"Unreadable frame from the API: {exc}")) + return + + async def _dispatch(self, data: Dict[str, Any]) -> None: + kind = data.get("type") + if kind == "Error": + await self._fail(error_from_frame(data)) + return + event = parse_event(data) + if event is None: + return + if kind == "Begin": + self.begin = event + self._begun.set() + elif kind == "Audio": + if TTSEvents.Audio.value not in self._handlers: + self._put(("audio", event.flush_id, event.audio)) + elif kind == "FlushDone": + self._put(("done", event.flush_id, None)) + elif kind == "Cancelled": + cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 + self._cancelled_before = max(self._cancelled_before, cancel_id) + self._put(("cancelled", cancel_id, None)) + elif kind == "Termination": + self.termination = event + await self._emit(str(kind), event) + if kind == "Termination": + self._stopped.set() + self._put(("closed", None, None)) + + async def _on_closed(self, exc: Exception) -> None: + if self._stopped.is_set(): + return + # `.code` is 1006 for an abnormal close, where `.rcvd` is None. + code = getattr(exc, "code", None) + reason = getattr(exc, "reason", "") or "" + if code != 1000 and self.error is None: + self.error = error_from_close(code, reason) + await self._emit(TTSEvents.Error, self.error) + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + async def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: + for handler in self._handlers.get(str(event), ()): + try: + result = handler(payload) + if asyncio.iscoroutine(result): + await result + except Exception: + logger.exception("AsyncTTSSession handler for %s raised", event) diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py new file mode 100644 index 00000000..66c404a2 --- /dev/null +++ b/assemblyai/tts/v1/session.py @@ -0,0 +1,444 @@ +import json +import logging +import queue +import threading +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union + +from ... import _audio +from ... import client as _client +from .._base_ws import PULL_QUEUE_FRAMES, build_uri, check_text_frame, resolve_settings +from .models import TTSConfig, TTSError +from .ws_models import ( + TTSBeginEvent, + TTSEvents, + TTSTerminationEvent, + error_from_close, + error_from_frame, + parse_event, +) + +__all__ = ["TTSSession"] + +logger = logging.getLogger(__name__) + +_Handler = Callable[[Any], None] +_Item = Tuple[str, Any, Any] # (kind, flush_id or cancel id, payload) + + +def websocket_connect(uri: str, **kwargs: Any) -> Any: + """Opens a sync websocket; the streaming client's seam, imported lazily so + `import assemblyai` does not load websockets. Tests patch this name.""" + + from ...streaming.v3.client import websocket_connect as _connect + + return _connect(uri, **kwargs) + + +class TTSSession: + """ + A Text-to-Speech WebSocket session: text in as it is written, audio out. + + Open it with `with`, which connects and waits for `Begin`. Send text with + `generate()` as it is produced, conclude each request with `flush()`, and + read its audio with `audio()`; `speak()` does all three for one text. + `cancel()` is the barge-in: it stops the request in flight and drops + whatever is queued. Leaving the `with` block sends `Terminate` and waits + for the session's `Termination` summary. + + Audio is delivered one of two ways: pulled with `speak()` / `audio()`, or + pushed to a handler registered with `on(TTSEvents.Audio, ...)`. Register + handlers for the other events (`FlushDone`, `WordBoundaries`, + `Cancelled`, `Termination`, `Error`) as needed; a handler takes the event + and runs on the session's reader thread. + + Example: + ```python + import assemblyai as aai + + aai.settings.api_key = "your-key" + + with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in reply_sentences(): + for chunk in session.speak(sentence): + play(chunk) # 24 kHz s16le mono + ``` + """ + + def __init__( + self, + config: Optional[TTSConfig] = None, + *, + client: Optional[_client.Client] = None, + api_key: Optional[str] = None, + token: Optional[str] = None, + word_boundaries: bool = False, + inactivity_timeout: Optional[int] = None, + connect_timeout: float = 10.0, + terminate_timeout: float = 10.0, + ) -> None: + """ + Creates a session; nothing connects until `with` or `connect()`. + + Args: + config: The voice and audio options (`voice_id`, `model_id`, + `language`, `sample_rate`, `encoding`). `container` and + `bit_rate` do not apply: a session streams raw PCM or G.711. + client: A `Client` whose settings (API key, `tts_base_url`) to + use. Defaults to the shared default client. + api_key: The API key to authenticate with; takes precedence over + `client` and the global settings. + token: A temporary token instead of an API key, for a session + opened on a user's device. + word_boundaries: Whether the server sends `WordBoundaries` frames. + inactivity_timeout: Seconds (5 to 120) after which an idle session + is closed; `None` for the account's default. + connect_timeout: Seconds to wait for the handshake and `Begin`. + terminate_timeout: Seconds to wait for `Termination` on exit. + + Raises: + TypeError: if `config` is not a `TTSConfig`. + """ + _audio.check_config(type(self).__name__, config, TTSConfig) + self.config = config or TTSConfig() + self._settings = resolve_settings( + type(self).__name__, client.settings if client else None, api_key, token + ) + self._token = token + self._word_boundaries = word_boundaries + self._inactivity_timeout = inactivity_timeout + self._connect_timeout = connect_timeout + self._terminate_timeout = terminate_timeout + + self._handlers: Dict[str, List[_Handler]] = {} + self._ws: Any = None + self._reader: Optional[threading.Thread] = None + self._send_lock = threading.Lock() + self.begin: Optional[TTSBeginEvent] = None + "The `Begin` frame, once connected." + self.termination: Optional[TTSTerminationEvent] = None + "The `Termination` frame, once the session ended cleanly." + self.error: Optional[TTSError] = None + "The error that ended the session, if one did." + self._reset() + + def _reset(self) -> None: + """Fresh per-connection state, so a session object can be reused.""" + + self._begun = threading.Event() + self._stopped = threading.Event() + self._queue: "queue.Queue[_Item]" = queue.Queue(maxsize=PULL_QUEUE_FRAMES) + self._next_flush_id = 0 + self._cancel_ids: List[int] = [] + # Requests with a lower flush_id than this were dropped by a Cancel. + self._cancelled_before = 0 + self._drop_warned = False + self.begin = None + self.termination = None + self.error = None + + # -- lifecycle ----------------------------------------------------------- + + def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: + """Registers `handler(event)` for a server frame type.""" + + self._handlers.setdefault(str(event), []).append(handler) + + def connect(self) -> TTSBeginEvent: + """ + Opens the socket and waits for `Begin`. + + Raises: + TTSError: if the API refuses the session (its `Error` frame), the + handshake is rejected, or `Begin` does not arrive in time. + """ + if self._ws is not None: + raise RuntimeError("the session is already connected") + self._reset() + + uri, headers = build_uri( + self._settings.tts_base_url, + self.config, + token=self._token, + api_key=self._settings.api_key, + word_boundaries=self._word_boundaries, + inactivity_timeout=self._inactivity_timeout, + ) + try: + self._ws = websocket_connect( + uri, additional_headers=headers, open_timeout=self._connect_timeout + ) + except Exception as exc: # the handshake failed + status = getattr(getattr(exc, "response", None), "status_code", None) + raise TTSError( + f"WebSocket handshake failed: {exc}", status_code=status + ) from exc + + self._reader = threading.Thread( + target=self._read_loop, name="assemblyai-tts-session", daemon=True + ) + self._reader.start() + + if not self._begun.wait(self._connect_timeout): + self.close() + raise TTSError("The session did not begin in time") + if self.error is not None: + self.close() + raise self.error + if self.begin is None: + self.close() + raise TTSError("The session closed before it began") + return self.begin + + def terminate(self) -> Optional[TTSTerminationEvent]: + """ + Ends the session cleanly: `Terminate`, then wait for `Termination`. + + Text sent and not flushed is spoken first. Returns the summary, or + `None` if the session had already ended or the wait timed out. + """ + if self._ws is None: + return None + if not self._stopped.is_set(): + try: + self._send({"type": "Terminate"}) + except Exception: # already closed + pass + self._stopped.wait(self._terminate_timeout) + self.close() + return self.termination + + def close(self) -> None: + """Closes the socket without `Terminate`; in-flight audio is dropped.""" + + ws, self._ws = self._ws, None + self._stopped.set() + self._begun.set() + if ws is not None: + try: + ws.close() + except Exception: # already closed + pass + self._put(("closed", None, None)) + if self._reader is not None and self._reader is not threading.current_thread(): + self._reader.join(timeout=self._terminate_timeout) + + def __enter__(self) -> "TTSSession": + self.connect() + return self + + def __exit__(self, exc_type: Any, *exc: object) -> None: + if exc_type is None: + self.terminate() + else: + self.close() + + # -- sending --------------------------------------------------------------- + + def generate(self, text: str) -> None: + """ + Appends `text` to the current request. + + Consecutive texts are joined exactly as sent, so include the spaces. + Audio may start before `flush()` once enough complete sentences have + arrived; `flush()` is what concludes the request. + + Raises: + ValueError: for a text over 2,000 characters (split it). + """ + check_text_frame(text) + self._send({"type": "Generate", "text": text}) + + def flush(self) -> int: + """ + Concludes the current request and returns its `flush_id`. + + Everything sent so far is spoken, then `FlushDone` arrives for this + id. Flushes can be pipelined; the audio comes in order. + """ + flush_id = self._next_flush_id + self._next_flush_id += 1 + self._send({"type": "Flush"}) + return flush_id + + def cancel(self) -> None: + """ + Barge-in: stops the request in flight and drops queued requests. + + The server answers `Cancelled`; audio already in transit may still + arrive before it. Text sent and not flushed is discarded. An + `audio()` iteration of a cancelled request ends. + """ + self._cancel_ids.append(self._next_flush_id) + self._next_flush_id += 1 + self._send({"type": "Cancel"}) + + def keep_alive(self) -> None: + """Resets the inactivity timer between sparse requests.""" + + self._send({"type": "KeepAlive"}) + + def speak(self, text: str) -> Iterator[bytes]: + """ + Speaks one text: `generate()`, `flush()`, then its audio. + + Returns an iterator of the request's audio bytes that ends at its + `FlushDone`, or early if the request is cancelled. + + Raises: + TTSError: from the iteration, if the session fails. + """ + self.generate(text) + return self.audio(self.flush()) + + def audio(self, flush_id: int) -> Iterator[bytes]: + """ + Yields the audio of request `flush_id` until its `FlushDone`. + + Requests are served in order, so read them in order. Audio is queued + for pulling only while no `TTSEvents.Audio` handler is registered. + + Raises: + TTSError: if the session fails before the request completes. + """ + while True: + if flush_id < self._cancelled_before: # dropped by a Cancel + return + try: + kind, ident, payload = self._queue.get(timeout=0.5) + except queue.Empty: + if self._stopped.is_set(): + if self.error is not None: + raise self.error + return + continue + if kind == "audio": + if ident == flush_id: + yield payload + elif kind == "done": + if ident == flush_id: + return + elif kind == "cancelled": + if ident > flush_id: # a Cancel sent after this request + return + elif kind == "error": + raise payload + elif kind == "closed": + if self.error is not None: + raise self.error + return + + # -- internals --------------------------------------------------------------- + + def _put(self, item: _Item) -> None: + """Queues an item for `audio()`, dropping the oldest audio when full.""" + + while True: + try: + self._queue.put_nowait(item) + return + except queue.Full: + try: + dropped = self._queue.get_nowait() + except queue.Empty: + continue + if dropped[0] == "audio" and not self._drop_warned: + self._drop_warned = True + logger.warning( + "TTSSession: dropping audio nobody reads; iterate audio() " + "or register an Audio handler" + ) + + def _send(self, frame: Dict[str, Any]) -> None: + if self._ws is None or self._stopped.is_set(): + if self.error is not None: + raise self.error + raise RuntimeError("the session is not open") + with self._send_lock: + self._ws.send(json.dumps(frame)) + + def _read_loop(self) -> None: + from websockets.exceptions import ConnectionClosed + + ws = self._ws + while not self._stopped.is_set() and ws is not None: + try: + raw = ws.recv(timeout=1) + except TimeoutError: + continue + except ConnectionClosed as exc: + self._on_closed(exc) + return + except Exception as exc: # the socket was closed under us + if not self._stopped.is_set(): + self._on_closed(exc) + return + if isinstance(raw, bytes): + continue + try: + data = json.loads(raw) + except ValueError: + logger.warning("ignoring a non-JSON frame") + continue + if isinstance(data, dict): + try: + self._dispatch(data) + except Exception as exc: # a frame this SDK cannot read + self._fail(TTSError(f"Unreadable frame from the API: {exc}")) + return + + def _fail(self, error: TTSError) -> None: + """Ends the session on an error: handlers, waiters, state.""" + + if self.error is None: + self.error = error + self._emit(TTSEvents.Error, error) + self._put(("error", None, self.error)) + self._stopped.set() + self._begun.set() + + def _dispatch(self, data: Dict[str, Any]) -> None: + kind = data.get("type") + if kind == "Error": + self._fail(error_from_frame(data)) + return + event = parse_event(data) + if event is None: + return + if kind == "Begin": + self.begin = event + self._begun.set() + elif kind == "Audio": + if TTSEvents.Audio.value not in self._handlers: + self._put(("audio", event.flush_id, event.audio)) + elif kind == "FlushDone": + self._put(("done", event.flush_id, None)) + elif kind == "Cancelled": + cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 + self._cancelled_before = max(self._cancelled_before, cancel_id) + self._put(("cancelled", cancel_id, None)) + elif kind == "Termination": + self.termination = event + self._emit(str(kind), event) + if kind == "Termination": + self._stopped.set() + self._put(("closed", None, None)) + + def _on_closed(self, exc: Exception) -> None: + if self._stopped.is_set(): + return + # `.code` is 1006 for an abnormal close, where `.rcvd` is None. + code = getattr(exc, "code", None) + reason = getattr(exc, "reason", "") or "" + if code != 1000 and self.error is None: + self.error = error_from_close(code, reason) + self._emit(TTSEvents.Error, self.error) + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: + for handler in self._handlers.get(str(event), ()): + try: + handler(payload) + except Exception: + logger.exception("TTSSession handler for %s raised", event) diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py new file mode 100644 index 00000000..8f273cdd --- /dev/null +++ b/assemblyai/tts/v1/ws_models.py @@ -0,0 +1,175 @@ +"""Models of the Text-to-Speech WebSocket API: the frames a session receives. + +The client frames (`Generate`, `Flush`, `Cancel`, `KeepAlive`, `Terminate`) are +plain dicts; the server frames are the events below, dispatched to handlers +registered with `TTSSession.on()`. +""" + +import base64 +from enum import Enum +from typing import Any, Dict, List, Optional, Type, TypeVar + +try: + # pydantic v2 import (`model_validator` exists only in v2) + from pydantic import BaseModel, Field, model_validator # noqa: F401 + + pydantic_v2 = True +except ImportError: + # pydantic v1 import (fallback for Python < 3.14) + from pydantic import BaseModel, Field + + pydantic_v2 = False + +from .models import TTSError + +__all__ = [ + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", +] + +# Close codes after which the same session parameters can succeed on a new +# connection: lost connection, a transient check failure, a server restart. +_RETRYABLE_CLOSE_CODES = frozenset({1006, 1011, 1012}) + +_M = TypeVar("_M", bound=BaseModel) + + +def _parse(model: Type[_M], data: Dict[str, Any]) -> _M: + if hasattr(model, "model_validate"): + return model.model_validate(data) + return model.parse_obj(data) + + +class TTSEvents(str, Enum): + """The server frames of a session, as `TTSSession.on()` names them.""" + + Begin = "Begin" + Audio = "Audio" + FlushDone = "FlushDone" + WordBoundaries = "WordBoundaries" + Cancelled = "Cancelled" + Termination = "Termination" + Error = "Error" + + def __str__(self): + return self.value + + +class TTSBeginEvent(BaseModel): + """The first frame of a session: it is open and configured.""" + + id: str + "The session id. Quote it in support requests." + + expires_at: int + "Unix time in seconds at which the session's maximum duration ends." + + configuration: Dict[str, Any] + """The session as created: the effective connect parameters, `model`, + `voice_id`, `session_ttl_sec`, `limits` and more. An open set.""" + + +class TTSAudioEvent(BaseModel): + """A piece of a request's audio.""" + + flush_id: int + "The request the audio belongs to." + + audio: bytes = Field(..., repr=False) + "Raw codec bytes (no header), decoded from the frame's base64." + + +class TTSFlushDoneEvent(BaseModel): + """A request's audio has all been delivered.""" + + flush_id: int + audio_duration_ms: int + "The playback duration of the request's audio; `0` for nothing to speak." + + +class TTSWord(BaseModel): + word: str + start: int + "Milliseconds on the session's audio timeline (not the request's)." + end: int + confidence: float + + +class TTSWordBoundariesEvent(BaseModel): + """Word timings of a request, best effort, after its `FlushDone`.""" + + flush_id: int + audio_start_ms: int + "Where the request's audio starts on the session's timeline." + words: List[TTSWord] + + +class TTSCancelledEvent(BaseModel): + """The answer to a `Cancel`.""" + + audio_duration_seconds: float + "Audio streamed for the cancelled request before it stopped." + + +class TTSTerminationEvent(BaseModel): + """The session's summary, after `Terminate`.""" + + session_duration_seconds: int + total_input_char_length: int + "The billed characters." + audio_duration_seconds: float + voice_expires_at: Optional[int] = None + + +def parse_event(data: Dict[str, Any]) -> Any: + """The event model for a server frame, or `None` for an unknown type.""" + + kind = data.get("type") + if kind == "Audio": + return TTSAudioEvent( + flush_id=data["flush_id"], audio=base64.b64decode(data["audio"]) + ) + if kind == "Begin": + return _parse(TTSBeginEvent, data) + if kind == "FlushDone": + return _parse(TTSFlushDoneEvent, data) + if kind == "WordBoundaries": + return _parse(TTSWordBoundariesEvent, data) + if kind == "Cancelled": + return _parse(TTSCancelledEvent, data) + if kind == "Termination": + return _parse(TTSTerminationEvent, data) + return None + + +def error_from_frame(data: Dict[str, Any]) -> TTSError: + """The `TTSError` for an `Error` frame; `body` keeps the frame.""" + + retry_after = data.get("retry_after_seconds") + return TTSError( + data.get("error") or "text-to-speech session failed", + error_code=data.get("error_type"), + param=data.get("param"), + retryable=data.get("retryable"), + retry_after=int(retry_after) if isinstance(retry_after, (int, float)) else None, + body=data, + ) + + +def error_from_close(code: Optional[int], reason: str) -> TTSError: + """The `TTSError` for a connection closed without a clean `Termination`.""" + + text = f"Connection closed (code={code})" + if reason: + text = f"{text}: {reason}" + return TTSError( + text, + retryable=code in _RETRYABLE_CLOSE_CODES, + body={"close_code": code, "reason": reason}, + ) diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py new file mode 100644 index 00000000..13bf0b47 --- /dev/null +++ b/tests/unit/test_tts_session.py @@ -0,0 +1,425 @@ +"""Tests for the threaded Text-to-Speech WebSocket session. + +`_FakeServer` plays the API: it answers `Generate`/`Flush` with the text's +bytes as one `Audio` frame and a `FlushDone`, `Cancel` with `Cancelled`, and +`Terminate` with `Termination` and a clean close. +""" + +import base64 +import json +import queue +import time +from typing import Any, List, Optional +from urllib.parse import parse_qs, urlsplit + +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +import assemblyai as aai + +aai.settings.api_key = "test" + +_BEGIN = { + "type": "Begin", + "id": "sess-1", + "expires_at": 1789556741, + "configuration": { + "voice_id": "aai-jane", + "model": "assemblyai-tts-1", + "sample_rate": 24000, + }, +} +_TERMINATION = { + "type": "Termination", + "session_duration_seconds": 3, + "total_input_char_length": 20, + "audio_duration_seconds": 1.5, +} + + +def _closed_ok() -> ConnectionClosedOK: + return ConnectionClosedOK(Close(1000, "Session ended"), None) + + +class _FakeServer: + def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): + self.sent: List[dict] = [] + self.inbox: "queue.Queue[Any]" = queue.Queue() + self.closed = False + self.text = "" + self.flush_id = 0 + self.answer_flush = answer_flush + if begin is not None: + self.push(begin) + + def push(self, item: Any) -> None: + self.inbox.put(json.dumps(item) if isinstance(item, dict) else item) + + def send(self, raw: str) -> None: + frame = json.loads(raw) + self.sent.append(frame) + kind = frame["type"] + if kind == "Generate": + self.text += frame["text"] + elif kind == "Flush": + if self.answer_flush: + audio = base64.b64encode(self.text.encode()).decode() + self.push({"type": "Audio", "audio": audio, "flush_id": self.flush_id}) + self.push( + { + "type": "FlushDone", + "flush_id": self.flush_id, + "audio_duration_ms": 80, + } + ) + self.flush_id += 1 + self.text = "" + elif kind == "Cancel": + self.flush_id += 1 + self.text = "" + self.push({"type": "Cancelled", "audio_duration_seconds": 0.1}) + elif kind == "Terminate": + self.push(_TERMINATION) + self.push(_closed_ok()) + + def recv(self, timeout: Optional[float] = None) -> str: + try: + item = self.inbox.get(timeout=0.02) + except queue.Empty: + raise TimeoutError() + if isinstance(item, Exception): + raise item + return item + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def server(mocker): + fake = _FakeServer() + connect = mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", return_value=fake + ) + fake.connect = connect + return fake + + +def _frames(server: _FakeServer, kind: str) -> List[dict]: + return [f for f in server.sent if f["type"] == kind] + + +def test_with_connects_waits_for_begin_and_terminates_on_exit(server): + with aai.TTSSession( + config=aai.TTSConfig(voice_id="aai-alba", sample_rate=8000) + ) as session: + assert session.begin is not None + assert session.begin.id == "sess-1" + assert session.begin.configuration["model"] == "assemblyai-tts-1" + + uri = server.connect.call_args.args[0] + headers = server.connect.call_args.kwargs["additional_headers"] + parts = urlsplit(uri) + assert parts.scheme == "wss" and parts.netloc == "tts.us.assemblyai.com" + assert parts.path == "/v1/ws" + assert parse_qs(parts.query) == { + "voice_id": ["aai-alba"], + "sample_rate": ["8000"], + "word_boundaries": ["false"], + } + assert headers["Authorization"] == "test" + assert "AssemblyAI" in headers["User-Agent"] + assert [f["type"] for f in server.sent] == ["Terminate"] + assert session.termination is not None + assert session.termination.total_input_char_length == 20 + assert server.closed + + +def test_speak_sends_generate_and_flush_and_yields_the_audio(server): + with aai.TTSSession() as session: + chunks = list(session.speak("Hello there.")) + + assert b"".join(chunks) == b"Hello there." + assert server.sent[:2] == [ + {"type": "Generate", "text": "Hello there."}, + {"type": "Flush"}, + ] + + +def test_requests_are_pipelined_and_read_in_flush_id_order(server): + with aai.TTSSession() as session: + session.generate("One. ") + first = session.flush() + session.generate("Two.") + second = session.flush() + assert (first, second) == (0, 1) + assert b"".join(session.audio(first)) == b"One. " + assert b"".join(session.audio(second)) == b"Two." + + +def test_cancel_ends_the_current_request_but_not_a_later_one(server): + server.answer_flush = False # the request stays in flight + with aai.TTSSession() as session: + session.generate("A long reply that gets interrupted.") + flush_id = session.flush() + stream = session.audio(flush_id) + session.cancel() # barge-in + assert list(stream) == [] + assert session._next_flush_id == 2 + + server.answer_flush = True + assert b"".join(session.speak("Next.")) == b"Next." + assert [f["type"] for f in server.sent] == [ + "Generate", + "Flush", + "Cancel", + "Generate", + "Flush", + "Terminate", + ] + + +def test_a_cancel_before_a_request_does_not_end_it(server): + with aai.TTSSession() as session: + session.cancel() # nothing in flight; Cancelled arrives later + assert b"".join(session.speak("Hi.")) == b"Hi." + + +def test_an_error_before_begin_raises_from_with(mocker): + fake = _FakeServer(begin=None) + fake.push( + { + "type": "Error", + "error_code": 3012, + "error": "The requested voice was not found.", + "error_type": "voice_not_found", + "retryable": False, + "param": "voice_id", + "retry_after_seconds": None, + } + ) + fake.push(ConnectionClosedError(Close(3012, "See Error message for details"), None)) + mocker.patch("assemblyai.tts.v1.session.websocket_connect", return_value=fake) + + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-nobody")): + pass + + err = exc_info.value + assert err.error_code == "voice_not_found" + assert err.param == "voice_id" + assert err.retryable is False + assert err.body["error_code"] == 3012 + assert fake.closed + + +def test_an_error_mid_session_raises_from_the_audio_and_reaches_the_handler(server): + errors = [] + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, errors.append) + server.answer_flush = False + flush_id = session.flush() + server.push( + { + "type": "Error", + "error_code": 3005, + "error": "No room", + "error_type": "at_capacity", + "retryable": True, + "param": None, + "retry_after_seconds": 2, + } + ) + list(session.audio(flush_id)) + + assert exc_info.value.error_code == "at_capacity" + assert exc_info.value.retry_after == 2 + assert errors == [exc_info.value] + assert session.error is exc_info.value + + +def test_a_lost_connection_is_a_retryable_error(server): + errors = [] + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, errors.append) + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(None, None)) # abnormal: no close frame + list(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + assert len(errors) == 1 + + +def test_a_send_after_a_failure_raises_the_stored_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.push(ConnectionClosedError(None, None)) + session._stopped.wait(2) + session.generate("too late") + assert exc_info.value.body["close_code"] == 1006 + + +def test_a_cancel_drops_every_pipelined_request(server): + server.answer_flush = False + with aai.TTSSession() as session: + first = session.flush() + second = session.flush() + session.cancel() + assert list(session.audio(first)) == [] + assert list(session.audio(second)) == [] # also dropped, must not hang + + +def test_an_unreadable_frame_ends_the_session_with_an_error(server): + with pytest.raises(aai.TTSError, match="Unreadable frame"): + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push({"type": "Audio", "flush_id": "zero"}) # no audio, wrong type + list(session.audio(flush_id)) + + +def test_a_session_object_can_be_reused(mocker): + fakes = [_FakeServer(), _FakeServer()] + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=fakes) + session = aai.TTSSession() + with session: + assert b"".join(session.speak("One.")) == b"One." + with session: + assert b"".join(session.speak("Two.")) == b"Two." + assert [f["type"] for f in fakes[1].sent] == ["Generate", "Flush", "Terminate"] + + +def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): + from assemblyai.tts._base_ws import PULL_QUEUE_FRAMES + + with aai.TTSSession() as session: + # audio of a request nobody reads floods the pull queue + for _ in range(PULL_QUEUE_FRAMES + 5): + server.push({"type": "Audio", "audio": "AA==", "flush_id": 7}) + deadline = time.monotonic() + 5 + while not server.inbox.empty() and time.monotonic() < deadline: + time.sleep(0.01) # let the reader queue the flood before anyone reads + assert session._queue.full() + session.generate("y") + flush_id = session.flush() + assert b"".join(session.audio(flush_id)) == b"y" # still served after the flood + assert session._queue.maxsize == PULL_QUEUE_FRAMES + assert "dropping audio" in caplog.text + + +def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): + received = [] + done = [] + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Audio, received.append) + session.on(aai.TTSEvents.FlushDone, done.append) + session.generate("Pushed.") + flush_id = session.flush() + assert list(session.audio(flush_id)) == [] # nothing queued for pulling + + assert [e.audio for e in received] == [b"Pushed."] + assert received[0].flush_id == 0 + assert done[0].audio_duration_ms == 80 + + +def test_word_boundaries_and_cancelled_reach_their_handlers(server): + words = [] + cancelled = [] + with aai.TTSSession(word_boundaries=True) as session: + session.on(aai.TTSEvents.WordBoundaries, words.append) + session.on(aai.TTSEvents.Cancelled, cancelled.append) + list(session.speak("Hello there.")) + server.push( + { + "type": "WordBoundaries", + "flush_id": 0, + "audio_start_ms": 0, + "words": [ + {"word": "Hello", "start": 265, "end": 530, "confidence": 0.9} + ], + } + ) + session.cancel() + session.keep_alive() + list(session.speak("Again.")) + + assert parse_qs(urlsplit(server.connect.call_args.args[0]).query)[ + "word_boundaries" + ] == ["true"] + assert words[0].words[0].word == "Hello" + assert cancelled[0].audio_duration_seconds == 0.1 + assert [f["type"] for f in server.sent].count("KeepAlive") == 1 + + +def test_a_temporary_token_goes_in_the_query_and_not_the_header(server, monkeypatch): + monkeypatch.setattr(aai.settings, "api_key", None) # a token needs no key + with aai.TTSSession(token="tmp-token", inactivity_timeout=30): + pass + + uri = server.connect.call_args.args[0] + headers = server.connect.call_args.kwargs["additional_headers"] + query = parse_qs(urlsplit(uri).query) + assert query["token"] == ["tmp-token"] + assert query["inactivity_timeout"] == ["30"] + assert "Authorization" not in headers + + +def test_generate_refuses_text_over_one_frame(server): + with aai.TTSSession() as session: + with pytest.raises(ValueError, match="2000"): + session.generate("x" * 2001) + with pytest.raises(TypeError): + session.generate(b"bytes") # type: ignore[arg-type] + + +def test_the_eu_host_comes_from_the_settings(server): + client = aai.Client(api_key="test") + client.settings.tts_base_url = "https://tts.eu.assemblyai.com" + with aai.TTSSession(client=client): + pass + assert urlsplit(server.connect.call_args.args[0]).netloc == "tts.eu.assemblyai.com" + + +def test_a_plain_http_host_opens_a_plain_websocket(server): + client = aai.Client(api_key="test") + client.settings.tts_base_url = "http://localhost:8080" + with aai.TTSSession(client=client): + pass + assert server.connect.call_args.args[0].startswith("ws://localhost:8080/v1/ws?") + + +def test_api_key_argument_is_used_without_touching_the_shared_settings(server): + with aai.TTSSession( + api_key="other", config=aai.TTSConfig(extend_biometric_consent=True) + ): + pass + headers = server.connect.call_args.kwargs["additional_headers"] + query = parse_qs(urlsplit(server.connect.call_args.args[0]).query) + assert headers["Authorization"] == "other" + assert query["extend_biometric_consent"] == ["true"] + assert aai.settings.api_key == "test" + + +def test_missing_key_is_a_clear_error(monkeypatch): + monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) + monkeypatch.setattr(aai.settings, "api_key", None) + with pytest.raises(ValueError, match="api_key= or token="): + aai.TTSSession() + + +def test_an_exception_inside_with_closes_without_terminate(server): + with pytest.raises(RuntimeError): + with aai.TTSSession(): + raise RuntimeError("boom") + assert server.sent == [] + assert server.closed + + +def test_wrong_config_type_is_a_type_error(): + with pytest.raises(TypeError): + aai.TTSSession(config=aai.DictationConfig()) # type: ignore[arg-type] diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py new file mode 100644 index 00000000..972513ac --- /dev/null +++ b/tests/unit/test_tts_session_async.py @@ -0,0 +1,215 @@ +"""Tests for the asyncio Text-to-Speech WebSocket session.""" + +import asyncio +import base64 +import json +from typing import Any, List, Optional +from urllib.parse import parse_qs, urlsplit + +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +import assemblyai as aai + +pytestmark = pytest.mark.asyncio + +aai.settings.api_key = "test" + +_BEGIN = { + "type": "Begin", + "id": "sess-1", + "expires_at": 1789556741, + "configuration": {"model": "assemblyai-tts-1"}, +} +_TERMINATION = { + "type": "Termination", + "session_duration_seconds": 3, + "total_input_char_length": 20, + "audio_duration_seconds": 1.5, +} + + +class _FakeAsyncServer: + def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): + self.sent: List[dict] = [] + self.inbox: "asyncio.Queue[Any]" = asyncio.Queue() + self.closed = False + self.text = "" + self.flush_id = 0 + self.answer_flush = answer_flush + if begin is not None: + self.push(begin) + + def push(self, item: Any) -> None: + self.inbox.put_nowait(json.dumps(item) if isinstance(item, dict) else item) + + async def send(self, raw: str) -> None: + frame = json.loads(raw) + self.sent.append(frame) + kind = frame["type"] + if kind == "Generate": + self.text += frame["text"] + elif kind == "Flush": + if self.answer_flush: + audio = base64.b64encode(self.text.encode()).decode() + self.push({"type": "Audio", "audio": audio, "flush_id": self.flush_id}) + self.push( + { + "type": "FlushDone", + "flush_id": self.flush_id, + "audio_duration_ms": 80, + } + ) + self.flush_id += 1 + self.text = "" + elif kind == "Cancel": + self.flush_id += 1 + self.text = "" + self.push({"type": "Cancelled", "audio_duration_seconds": 0.1}) + elif kind == "Terminate": + self.push(_TERMINATION) + self.push(ConnectionClosedOK(Close(1000, "Session ended"), None)) + + async def recv(self) -> str: + item = await self.inbox.get() + if isinstance(item, Exception): + raise item + return item + + async def close(self) -> None: + self.closed = True + self.push(ConnectionClosedOK(Close(1000, ""), None)) + + +@pytest.fixture +def server(mocker): + fake = _FakeAsyncServer() + + async def connect(uri, additional_headers): + connect.uri = uri + connect.headers = additional_headers + return fake + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + fake.connect = connect + return fake + + +async def test_async_with_connects_speaks_and_terminates(server): + async with aai.AsyncTTSSession( + config=aai.TTSConfig(voice_id="aai-alba") + ) as session: + assert session.begin is not None and session.begin.id == "sess-1" + chunks = [chunk async for chunk in session.speak("Hello there.")] + + assert b"".join(chunks) == b"Hello there." + assert parse_qs(urlsplit(server.connect.uri).query)["voice_id"] == ["aai-alba"] + assert server.connect.headers["Authorization"] == "test" + assert [f["type"] for f in server.sent] == ["Generate", "Flush", "Terminate"] + assert session.termination is not None + assert server.closed + + +async def test_pipelined_requests_and_cancel(server): + async with aai.AsyncTTSSession() as session: + await session.generate("One. ") + first = await session.flush() + second = await session.flush() + assert b"".join([c async for c in session.audio(first)]) == b"One. " + assert b"".join([c async for c in session.audio(second)]) == b"" + + server.answer_flush = False + third = await session.flush() + await session.cancel() + assert [c async for c in session.audio(third)] == [] + + +async def test_an_error_before_begin_raises_from_async_with(mocker): + fake = _FakeAsyncServer(begin=None) + fake.push( + { + "type": "Error", + "error_code": 3009, + "error": "Too many sessions", + "error_type": "rate_limited", + "retryable": True, + "param": None, + "retry_after_seconds": 17, + } + ) + + async def connect(uri, additional_headers): + return fake + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession(): + pass + assert exc_info.value.error_code == "rate_limited" + assert exc_info.value.retry_after == 17 + assert fake.closed + + +async def test_handlers_may_be_coroutines(server): + received = [] + + async def on_audio(event): + received.append(event.audio) + + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Audio, on_audio) + await session.generate("Pushed.") + flush_id = await session.flush() + assert [c async for c in session.audio(flush_id)] == [] + + assert received == [b"Pushed."] + + +async def test_a_lost_connection_raises_a_retryable_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(None, None)) # abnormal: no close frame + [c async for c in session.audio(flush_id)] + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1006 + + +async def test_a_cancel_drops_every_pipelined_request(server): + server.answer_flush = False + async with aai.AsyncTTSSession() as session: + first = await session.flush() + second = await session.flush() + await session.cancel() + assert [c async for c in session.audio(first)] == [] + assert [c async for c in session.audio(second)] == [] + + +async def test_a_session_object_can_be_reused(mocker): + fakes = [_FakeAsyncServer(), _FakeAsyncServer()] + + async def connect(uri, additional_headers): + return fakes.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + async with session: + assert b"".join([c async for c in session.speak("One.")]) == b"One." + async with session: + assert b"".join([c async for c in session.speak("Two.")]) == b"Two." + + +async def test_missing_key_is_a_clear_error(monkeypatch): + monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) + monkeypatch.setattr(aai.settings, "api_key", None) + with pytest.raises(ValueError, match="api_key= or token="): + aai.AsyncTTSSession() From 32db4edae13596c5e036401d1718ec1939961750 Mon Sep 17 00:00:00 2001 From: Ben Gotthold Date: Wed, 21 Oct 2026 23:32:02 +0000 Subject: [PATCH 02/21] test(tts): build the async fake server inside the event loop (Python 3.9) Co-Authored-By: Claude Fable 5.1 --- tests/unit/test_tts_session_async.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 972513ac..3931f847 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -7,6 +7,7 @@ from urllib.parse import parse_qs, urlsplit import pytest +import pytest_asyncio from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from websockets.frames import Close @@ -82,8 +83,10 @@ async def close(self) -> None: self.push(ConnectionClosedOK(Close(1000, ""), None)) -@pytest.fixture -def server(mocker): +@pytest_asyncio.fixture +async def server(mocker): + # Built inside the running loop: on Python 3.9 an asyncio.Queue binds to + # the loop that exists when it is created. fake = _FakeAsyncServer() async def connect(uri, additional_headers): From 1e822d3345ddc06634de61e74527f80c74b0e8e9 Mon Sep 17 00:00:00 2001 From: Ben Gotthold Date: Thu, 22 Oct 2026 01:06:33 +0000 Subject: [PATCH 03/21] fix(tts): session fixes from three first-use runs - `cancel()` ends the current `audio()` iteration at once and discards the request's in-transit audio, so a barge-in does not play stale speech. - `audio(flush_id)` raises for an id no `flush()` has returned instead of waiting forever; with no id it follows the request being written, so a player thread can start before the flush. - `on()` refuses an unknown event name; pulling while an `Audio` handler is registered raises instead of yielding nothing; a send on a session that is not open says how to open one. - `container`/`bit_rate` in a session's config are refused, not dropped. - Session errors carry the session id as `request_id`. - No warnings under `-W error`: the shared config dump uses `model_dump` when available, and the sync seam passes `legacy=True` on websockets 17. - Docs: raw-audio-only sessions, the text-rate limit, the 120 s default idle timeout, `Begin` and `session.begin`, word-boundary timing, an async code block, the error fields on a session, and a barge-in example that breaks. Co-Authored-By: Claude Fable 5.1 --- CLAUDE.md | 5 ++- README.md | 34 ++++++++++++---- assemblyai/_audio.py | 5 ++- assemblyai/tts/_base_ws.py | 5 +++ assemblyai/tts/v1/async_session.py | 40 ++++++++++++++++--- assemblyai/tts/v1/session.py | 60 +++++++++++++++++++++++----- assemblyai/tts/v1/ws_models.py | 20 +++++++++- tests/unit/test_tts_session.py | 57 +++++++++++++++++++++++++- tests/unit/test_tts_session_async.py | 4 +- 9 files changed, 198 insertions(+), 32 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 4f0209b3..d2a0bc08 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -54,7 +54,7 @@ aai.settings.api_key = "your-key" - `aai.DictationResponse` — Dictation result: `.text`, `.words` (`DictationWord` with `text` + `confidence`), `.confidence`, `.llm_response`, `.llm_error`, `.audio_duration_ms`, `.session_id`, `.request_time_ms`, `.sync_time_ms`, and `.final_text` (the LLM rewrite, falling back to `.text`) - `aai.TextToSpeech` — Text-to-Speech HTTP API: one text per request, audio out. Methods: `synthesize()` (complete audio as a `TTSResponse`), `stream()` (a `TTSStream` context manager yielding bytes as the API generates them), `close()` / `with` for a client built from `api_key=`. Host from `aai.settings.tts_base_url` (`https://tts.us.assemblyai.com`; set `https://tts.eu.assemblyai.com` for the EU zone) - `aai.AsyncTextToSpeech` — Asyncio counterpart of `TextToSpeech`; `synthesize()` is a coroutine and `stream()` returns an `AsyncTTSStream`. Owns an HTTP pool: use `async with` or `await aclose()` -- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin`; `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining; `cancel()` for barge-in; `keep_alive()`; `on(TTSEvents.X, handler)` for `Audio`, `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination`. Refusals raise `TTSError`. `token=` for temporary tokens +- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin` (`session.begin`); `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining (`audio()` with no id follows the request being written, from another thread); `cancel()` for barge-in (ends the iteration at once); `keep_alive()`; `on(TTSEvents.X, handler(event))` for `Begin`, `Audio` (switches audio from pull to push), `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination` (`None` after a failure). Refusals and failures raise `TTSError` (`status_code` None, `request_id` = session id). Raw PCM/G.711 only; 2,000 chars per generate; ~100 chars/s sustained. `token=` for temporary tokens - `aai.AsyncTTSSession` — Asyncio counterpart: `async with`, coroutine sends, `async for chunk in session.speak(text)` - `aai.TTSConfig` — Text-to-Speech options: `voice_id` (a library voice such as `aai-jane`, or a cloned voice by UUID or slug), `model_id` (`assemblyai-tts-1`), `language`, `sample_rate`, `encoding` (`aai.TTSEncoding`: `pcm_s16le` default | `pcm_mulaw` | `pcm_alaw` | `mp3`), `container` (`aai.TTSContainer`: `raw` default | `wav`), `bit_rate`, `extend_biometric_consent` (owner-only consent renewal of a cloned voice). Rejects unknown fields - `aai.TTSResponse` — Text-to-Speech result: `.audio` (bytes) plus the response headers: `.request_id`, `.region`, `.model`, `.voice_id`, `.language`, `.encoding`, `.sample_rate`, `.container`, `.bit_rate`, `.content_type`, `.voice_expires_at`, `.billed_characters`, `.audio_duration_ms`. `TTSStream.info` carries the same fields without `audio`; `audio_duration_ms` is `None` on a stream with audio (`0` when the text had nothing to speak) @@ -484,9 +484,10 @@ tts.synthesize("Hello.", aai.TTSConfig(voice_id="aai-alba")) # A WebSocket session for a voice agent: no handshake per sentence, barge-in with cancel() with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: - for sentence in reply_sentences(): + for sentence in llm_reply_sentences(): for chunk in session.speak(sentence): play(chunk) +# session.termination is the summary (None if the session failed); raw PCM/G.711 only, no mp3/container ``` ## LLM Gateway diff --git a/README.md b/README.md index 594faa44..b60bae90 100644 --- a/README.md +++ b/README.md @@ -1016,7 +1016,9 @@ asyncio.run(main())
Stream text in and audio out over a WebSocket session (voice agents) -`aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary. +`aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary (`None` if the session failed instead). + +A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, `container` and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`. ```python import assemblyai as aai @@ -1028,7 +1030,7 @@ with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: for chunk in session.speak(sentence): # raw 24 kHz 16-bit mono PCM play(chunk) -print(session.termination.audio_duration_seconds, "s spoken") +print("spoken:", session.termination) # the session's summary ``` Streaming the text in as it is generated, with a barge-in: @@ -1041,17 +1043,33 @@ with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: for chunk in session.audio(request): if user_started_talking(): - session.cancel() # stops this request; the iteration ends + session.cancel() # ends this iteration at once; the server drops the rest + break play(chunk) ``` -Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. A `Generate` frame takes at most 2,000 characters; split longer text across calls. The session ends after an hour, or after `inactivity_timeout` seconds without a frame; send `keep_alive()` between sparse requests. +Audio starts before `flush()` once whole sentences have arrived. To play it while the text is still being sent, read `session.audio()` with no id from another thread (or task): it follows the request being written, and ends at that request's `FlushDone`. Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. Every `flush()` and every `cancel()` takes the next `flush_id`, so ids skip one after a cancel. + +Limits: a `Generate` frame takes at most 2,000 characters (split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) without a frame; send `keep_alive()` between sparse requests. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. + +Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. -Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries` (word timings per request, with `word_boundaries=True`), `TTSEvents.Cancelled`, `TTSEvents.Termination`, `TTSEvents.Error`. An `Audio` handler receives the audio frames instead of the pull iterators. A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the audio iteration and reaches the `Error` handler. +A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status) and `request_id` is the session id; `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. -For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`). +For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`, which answers `{"token": ..., "expires_in_seconds": 60}`). -The asyncio twin is `aai.AsyncTTSSession`: `async with`, coroutine `generate()`/`flush()`/`cancel()`, and `async for chunk in session.speak(text)`. +The asyncio twin is `aai.AsyncTTSSession`: + +```python +async with aai.AsyncTTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + async for chunk in session.speak("Hello there."): + await play(chunk) + await session.generate("Streamed ") + await session.generate("in.") + request = await session.flush() + async for chunk in session.audio(request): + await play(chunk) +```
@@ -2144,7 +2162,7 @@ All custom SDK exceptions inherit from `aai.AssemblyAIError`, except the streami | `aai.AssemblyAIError` | Base exception for all AssemblyAI errors. | `status_code: Optional[int]` | | `aai.TranscriptError` | Raised when a prerecorded transcript API call fails at the HTTP level or polling times out. | `status_code: Optional[int]` | | `aai.SyncTranscriptError` | Raised when synchronous transcription fails (`aai.SyncTranscriber`). | `status_code: Optional[int]`, `error_code: Optional[str]`, `retry_after: Optional[int]` | -| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | +| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`) or a session is refused or fails (`aai.TTSSession`, where `status_code` is `None` and `request_id` is the session id). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | | `assemblyai.types.RedactedAudioIncompleteError` | Raised when requesting redacted audio before processing completes. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioExpiredError` | Raised when requesting redacted audio that has expired and is no longer available. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioUnavailableError` | Raised when requesting redacted audio that is unavailable at the given URL. | `status_code: Optional[int]` | diff --git a/assemblyai/_audio.py b/assemblyai/_audio.py index e0bab838..3728c11e 100644 --- a/assemblyai/_audio.py +++ b/assemblyai/_audio.py @@ -164,7 +164,10 @@ def _config_to_json(config: Any, *, exclude: Tuple[str, ...] = ()) -> Optional[d Returns: the JSON-ready dict, or `None` when nothing is set. """ - data = config.dict(exclude_none=True) + if hasattr(config, "model_dump"): + data = config.model_dump(exclude_none=True) + else: + data = config.dict(exclude_none=True) for field in exclude: data.pop(field, None) return data or None diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 9ea66a46..9bb9d99f 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -72,6 +72,11 @@ def build_uri( ) -> Tuple[str, Dict[str, Optional[str]]]: """The session URI and handshake headers from the settings and config.""" + if config.container is not None or config.bit_rate is not None: + raise ValueError( + "container and bit_rate do not apply to a session, which streams raw " + "pcm_s16le, pcm_mulaw or pcm_alaw; leave them unset in its TTSConfig" + ) parts = urlsplit(base_url if "://" in base_url else f"https://{base_url}") scheme = "ws" if parts.scheme in ("http", "ws") else "wss" params: Dict[str, Any] = ( diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index 603053c5..d011baf8 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -96,7 +96,12 @@ def __init__( def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: """Registers `handler(event)`, a function or coroutine function.""" - self._handlers.setdefault(str(event), []).append(handler) + name = str(event) + if name not in TTSEvents.__members__: + raise ValueError( + f"unknown event {name!r}; one of {', '.join(TTSEvents.__members__)}" + ) + self._handlers.setdefault(name, []).append(handler) async def connect(self) -> TTSBeginEvent: """Opens the socket and waits for `Begin`; see `TTSSession.connect`.""" @@ -214,8 +219,10 @@ async def flush(self) -> int: async def cancel(self) -> None: """Barge-in; see `TTSSession.cancel`.""" - self._cancel_ids.append(self._next_flush_id) + cancel_id = self._next_flush_id self._next_flush_id += 1 + self._cancel_ids.append(cancel_id) + self._cancelled_before = max(self._cancelled_before, cancel_id) await self._send({"type": "Cancel"}) async def keep_alive(self) -> None: @@ -229,11 +236,25 @@ async def speak(self, text: str) -> AsyncIterator[bytes]: async for chunk in self.audio(flush_id): yield chunk - async def audio(self, flush_id: int) -> AsyncIterator[bytes]: - """Yields the audio of request `flush_id`; see `TTSSession.audio`.""" + async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: + """ + Yields the audio of request `flush_id`; see `TTSSession.audio`. + + Without `flush_id`, the request being written now, so a task can + play it while another still sends text. + """ if self._queue is None: raise RuntimeError("the session is not open") + if flush_id is None: + flush_id = self._next_flush_id # the request being written now + elif flush_id > self._next_flush_id: + raise ValueError(f"flush_id {flush_id} has not been flushed yet") + if TTSEvents.Audio.value in self._handlers: + raise RuntimeError( + "audio goes to the registered Audio handler; pull with audio() " + "or speak() only without one" + ) while True: if flush_id < self._cancelled_before: # dropped by a Cancel return @@ -284,6 +305,8 @@ def _put(self, item: _Item) -> None: async def _fail(self, error: TTSError) -> None: if self.error is None: + if error.request_id is None and self.begin is not None: + error.request_id = self.begin.id self.error = error await self._emit(TTSEvents.Error, error) self._put(("error", None, self.error)) @@ -294,7 +317,9 @@ async def _send(self, frame: Dict[str, Any]) -> None: if self._ws is None or self._stopped.is_set(): if self.error is not None: raise self.error - raise RuntimeError("the session is not open") + raise RuntimeError( + "the session is not open: use `with` / `async with`, or connect()" + ) await self._ws.send(json.dumps(frame)) async def _read_loop(self) -> None: @@ -361,7 +386,10 @@ async def _on_closed(self, exc: Exception) -> None: code = getattr(exc, "code", None) reason = getattr(exc, "reason", "") or "" if code != 1000 and self.error is None: - self.error = error_from_close(code, reason) + error = error_from_close(code, reason) + if self.begin is not None: + error.request_id = self.begin.id + self.error = error await self._emit(TTSEvents.Error, self.error) self._stopped.set() self._begun.set() diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 66c404a2..d3036521 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -29,8 +29,16 @@ def websocket_connect(uri: str, **kwargs: Any) -> Any: """Opens a sync websocket; the streaming client's seam, imported lazily so `import assemblyai` does not load websockets. Tests patch this name.""" + import inspect + + from websockets.sync.client import connect as _ws_connect + from ...streaming.v3.client import websocket_connect as _connect + # websockets 17 deprecates a bare connect() unless `legacy=True` (older + # releases have no such parameter). + if "legacy" in inspect.signature(_ws_connect).parameters: + kwargs["legacy"] = True return _connect(uri, **kwargs) @@ -116,7 +124,8 @@ def __init__( self.begin: Optional[TTSBeginEvent] = None "The `Begin` frame, once connected." self.termination: Optional[TTSTerminationEvent] = None - "The `Termination` frame, once the session ended cleanly." + """The `Termination` summary once the session ended cleanly; `None` + before that, after a failure, or when the wait for it timed out.""" self.error: Optional[TTSError] = None "The error that ended the session, if one did." self._reset() @@ -141,7 +150,12 @@ def _reset(self) -> None: def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: """Registers `handler(event)` for a server frame type.""" - self._handlers.setdefault(str(event), []).append(handler) + name = str(event) + if name not in TTSEvents.__members__: + raise ValueError( + f"unknown event {name!r}; one of {', '.join(TTSEvents.__members__)}" + ) + self._handlers.setdefault(name, []).append(handler) def connect(self) -> TTSBeginEvent: """ @@ -264,12 +278,15 @@ def cancel(self) -> None: """ Barge-in: stops the request in flight and drops queued requests. - The server answers `Cancelled`; audio already in transit may still - arrive before it. Text sent and not flushed is discarded. An - `audio()` iteration of a cancelled request ends. + An `audio()` iteration of a cancelled request ends at once; audio + of it still in transit is discarded. The server answers `Cancelled`. + Text sent and not flushed is discarded. Like `flush()`, a cancel + takes a `flush_id`, so the next request's id skips one. """ - self._cancel_ids.append(self._next_flush_id) + cancel_id = self._next_flush_id self._next_flush_id += 1 + self._cancel_ids.append(cancel_id) + self._cancelled_before = max(self._cancelled_before, cancel_id) self._send({"type": "Cancel"}) def keep_alive(self) -> None: @@ -290,16 +307,30 @@ def speak(self, text: str) -> Iterator[bytes]: self.generate(text) return self.audio(self.flush()) - def audio(self, flush_id: int) -> Iterator[bytes]: + def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: """ Yields the audio of request `flush_id` until its `FlushDone`. - Requests are served in order, so read them in order. Audio is queued - for pulling only while no `TTSEvents.Audio` handler is registered. + Without `flush_id`, the request being written now: a thread can read + it while `generate()` is still sending text, since audio starts + before `flush()`. Requests are served in order, so read them in + order. Audio is queued for pulling only while no `TTSEvents.Audio` + handler is registered. Raises: TTSError: if the session fails before the request completes. + ValueError: for a `flush_id` no `flush()` has returned. + RuntimeError: when an `Audio` handler receives the audio instead. """ + if flush_id is None: + flush_id = self._next_flush_id # the request being written now + elif flush_id > self._next_flush_id: + raise ValueError(f"flush_id {flush_id} has not been flushed yet") + if TTSEvents.Audio.value in self._handlers: + raise RuntimeError( + "audio goes to the registered Audio handler; pull with audio() " + "or speak() only without one" + ) while True: if flush_id < self._cancelled_before: # dropped by a Cancel return @@ -352,7 +383,9 @@ def _send(self, frame: Dict[str, Any]) -> None: if self._ws is None or self._stopped.is_set(): if self.error is not None: raise self.error - raise RuntimeError("the session is not open") + raise RuntimeError( + "the session is not open: use `with` / `async with`, or connect()" + ) with self._send_lock: self._ws.send(json.dumps(frame)) @@ -390,6 +423,8 @@ def _fail(self, error: TTSError) -> None: """Ends the session on an error: handlers, waiters, state.""" if self.error is None: + if error.request_id is None and self.begin is not None: + error.request_id = self.begin.id self.error = error self._emit(TTSEvents.Error, error) self._put(("error", None, self.error)) @@ -430,7 +465,10 @@ def _on_closed(self, exc: Exception) -> None: code = getattr(exc, "code", None) reason = getattr(exc, "reason", "") or "" if code != 1000 and self.error is None: - self.error = error_from_close(code, reason) + error = error_from_close(code, reason) + if self.begin is not None: + error.request_id = self.begin.id + self.error = error self._emit(TTSEvents.Error, self.error) self._stopped.set() self._begun.set() diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py index 8f273cdd..5c39b903 100644 --- a/assemblyai/tts/v1/ws_models.py +++ b/assemblyai/tts/v1/ws_models.py @@ -50,12 +50,19 @@ class TTSEvents(str, Enum): """The server frames of a session, as `TTSSession.on()` names them.""" Begin = "Begin" + "Handler receives a `TTSBeginEvent`." Audio = "Audio" + "Handler receives a `TTSAudioEvent`; registering one turns off pulling." FlushDone = "FlushDone" + "Handler receives a `TTSFlushDoneEvent`." WordBoundaries = "WordBoundaries" + "Handler receives a `TTSWordBoundariesEvent`, after the request's `FlushDone`." Cancelled = "Cancelled" + "Handler receives a `TTSCancelledEvent`." Termination = "Termination" + "Handler receives a `TTSTerminationEvent`." Error = "Error" + "Handler receives the `TTSError` that ended the session." def __str__(self): return self.value @@ -95,14 +102,23 @@ class TTSFlushDoneEvent(BaseModel): class TTSWord(BaseModel): word: str + "One whitespace-delimited token of the request's text, verbatim." start: int - "Milliseconds on the session's audio timeline (not the request's)." + """Milliseconds on the session's whole audio timeline, not the request's: + subtract the event's `audio_start_ms` for a time within the request.""" end: int confidence: float class TTSWordBoundariesEvent(BaseModel): - """Word timings of a request, best effort, after its `FlushDone`.""" + """ + Word timings of a request: best effort, at most one per request. + + Arrives after the request's `FlushDone`, so usually after its `audio()` + iteration has ended, and possibly after the next request's first audio; + match it by `flush_id`. `start` and `end` count from the start of the + session's audio; `start - audio_start_ms` is the time within the request. + """ flush_id: int audio_start_ms: int diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 13bf0b47..3ad013d5 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -180,6 +180,57 @@ def test_cancel_ends_the_current_request_but_not_a_later_one(server): ] +def test_cancel_ends_the_iteration_before_the_server_answers(server): + server.answer_flush = False + with aai.TTSSession() as session: + flush_id = session.flush() + server.push({"type": "Audio", "audio": "AA==", "flush_id": 0}) + stream = session.audio(flush_id) + assert next(stream) == b"\x00" + session.cancel() + # in-transit audio of the cancelled request is discarded, not yielded + server.push({"type": "Audio", "audio": "AA==", "flush_id": 0}) + assert list(stream) == [] + + +def test_audio_for_an_unflushed_id_raises(server): + with aai.TTSSession() as session: + with pytest.raises(ValueError, match="not been flushed"): + next(session.audio(5)) + + +def test_audio_without_an_id_follows_the_request_being_written(server): + import threading + + with aai.TTSSession() as session: + got = [] + reader = threading.Thread(target=lambda: got.extend(session.audio())) + reader.start() # reading before the flush, like a player thread + session.generate("Early ") + session.generate("audio.") + session.flush() + reader.join(timeout=5) + assert b"".join(got) == b"Early audio." + + +def test_a_misspelled_event_name_is_refused(server): + session = aai.TTSSession() + with pytest.raises(ValueError, match="FlushDone"): + session.on("FlushDon", print) + session.on("FlushDone", print) # a string spelling of a member is fine + + +def test_sending_before_connect_says_how_to_open(server): + with pytest.raises(RuntimeError, match="with"): + aai.TTSSession().generate("Hi.") + + +def test_http_only_options_are_refused_on_a_session(server): + with pytest.raises(ValueError, match="container and bit_rate"): + with aai.TTSSession(config=aai.TTSConfig(container="wav")): + pass + + def test_a_cancel_before_a_request_does_not_end_it(server): with aai.TTSSession() as session: session.cancel() # nothing in flight; Cancelled arrives later @@ -211,6 +262,7 @@ def test_an_error_before_begin_raises_from_with(mocker): assert err.param == "voice_id" assert err.retryable is False assert err.body["error_code"] == 3012 + assert err.request_id is None # no Begin, so no session id yet assert fake.closed @@ -236,6 +288,7 @@ def test_an_error_mid_session_raises_from_the_audio_and_reaches_the_handler(serv assert exc_info.value.error_code == "at_capacity" assert exc_info.value.retry_after == 2 + assert exc_info.value.request_id == "sess-1" # the session id assert errors == [exc_info.value] assert session.error is exc_info.value @@ -320,7 +373,9 @@ def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): session.on(aai.TTSEvents.FlushDone, done.append) session.generate("Pushed.") flush_id = session.flush() - assert list(session.audio(flush_id)) == [] # nothing queued for pulling + with pytest.raises(RuntimeError, match="Audio handler"): + list(session.audio(flush_id)) + session._stopped.wait(0.2) # let the frames land assert [e.audio for e in received] == [b"Pushed."] assert received[0].flush_id == 0 diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 3931f847..38c41245 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -169,7 +169,9 @@ async def on_audio(event): session.on(aai.TTSEvents.Audio, on_audio) await session.generate("Pushed.") flush_id = await session.flush() - assert [c async for c in session.audio(flush_id)] == [] + with pytest.raises(RuntimeError, match="Audio handler"): + await session.audio(flush_id).__anext__() + await asyncio.sleep(0.05) assert received == [b"Pushed."] From c287dffdf8a416cb4a65456ff1e21885d8b0c46b Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:47:11 +0000 Subject: [PATCH 04/21] test(tts): bound every audio wait in the session tests Drain audio() and speak() through a helper with a deadline (a thread joined with a timeout for the sync session, asyncio.wait_for for the async one), so a regression in the flush_id or cancel bookkeeping fails the run instead of hanging it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- tests/unit/test_tts_session.py | 63 ++++++++++++++++++---------- tests/unit/test_tts_session_async.py | 29 ++++++++----- 2 files changed, 61 insertions(+), 31 deletions(-) diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 3ad013d5..d7938328 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -8,8 +8,9 @@ import base64 import json import queue +import threading import time -from typing import Any, List, Optional +from typing import Any, Iterable, List, Optional from urllib.parse import parse_qs, urlsplit import pytest @@ -106,6 +107,28 @@ def server(mocker): return fake +def _drain(items: Iterable[Any], timeout: float = 5.0) -> List[Any]: + """Collects `items` on a thread, failing the test if it does not end.""" + + got: List[Any] = [] + raised: List[BaseException] = [] + + def collect() -> None: + try: + got.extend(items) + except BaseException as exc: + raised.append(exc) + + thread = threading.Thread(target=collect, daemon=True) + thread.start() + thread.join(timeout) + if thread.is_alive(): + pytest.fail(f"the iteration did not end within {timeout} s") + if raised: + raise raised[0] + return got + + def _frames(server: _FakeServer, kind: str) -> List[dict]: return [f for f in server.sent if f["type"] == kind] @@ -138,7 +161,7 @@ def test_with_connects_waits_for_begin_and_terminates_on_exit(server): def test_speak_sends_generate_and_flush_and_yields_the_audio(server): with aai.TTSSession() as session: - chunks = list(session.speak("Hello there.")) + chunks = _drain(session.speak("Hello there.")) assert b"".join(chunks) == b"Hello there." assert server.sent[:2] == [ @@ -154,8 +177,8 @@ def test_requests_are_pipelined_and_read_in_flush_id_order(server): session.generate("Two.") second = session.flush() assert (first, second) == (0, 1) - assert b"".join(session.audio(first)) == b"One. " - assert b"".join(session.audio(second)) == b"Two." + assert b"".join(_drain(session.audio(first))) == b"One. " + assert b"".join(_drain(session.audio(second))) == b"Two." def test_cancel_ends_the_current_request_but_not_a_later_one(server): @@ -165,11 +188,11 @@ def test_cancel_ends_the_current_request_but_not_a_later_one(server): flush_id = session.flush() stream = session.audio(flush_id) session.cancel() # barge-in - assert list(stream) == [] + assert _drain(stream) == [] assert session._next_flush_id == 2 server.answer_flush = True - assert b"".join(session.speak("Next.")) == b"Next." + assert b"".join(_drain(session.speak("Next."))) == b"Next." assert [f["type"] for f in server.sent] == [ "Generate", "Flush", @@ -190,7 +213,7 @@ def test_cancel_ends_the_iteration_before_the_server_answers(server): session.cancel() # in-transit audio of the cancelled request is discarded, not yielded server.push({"type": "Audio", "audio": "AA==", "flush_id": 0}) - assert list(stream) == [] + assert _drain(stream) == [] def test_audio_for_an_unflushed_id_raises(server): @@ -200,8 +223,6 @@ def test_audio_for_an_unflushed_id_raises(server): def test_audio_without_an_id_follows_the_request_being_written(server): - import threading - with aai.TTSSession() as session: got = [] reader = threading.Thread(target=lambda: got.extend(session.audio())) @@ -234,7 +255,7 @@ def test_http_only_options_are_refused_on_a_session(server): def test_a_cancel_before_a_request_does_not_end_it(server): with aai.TTSSession() as session: session.cancel() # nothing in flight; Cancelled arrives later - assert b"".join(session.speak("Hi.")) == b"Hi." + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." def test_an_error_before_begin_raises_from_with(mocker): @@ -284,7 +305,7 @@ def test_an_error_mid_session_raises_from_the_audio_and_reaches_the_handler(serv "retry_after_seconds": 2, } ) - list(session.audio(flush_id)) + _drain(session.audio(flush_id)) assert exc_info.value.error_code == "at_capacity" assert exc_info.value.retry_after == 2 @@ -301,7 +322,7 @@ def test_a_lost_connection_is_a_retryable_error(server): server.answer_flush = False flush_id = session.flush() server.push(ConnectionClosedError(None, None)) # abnormal: no close frame - list(session.audio(flush_id)) + _drain(session.audio(flush_id)) assert exc_info.value.retryable is True assert exc_info.value.body == {"close_code": 1006, "reason": ""} @@ -323,8 +344,8 @@ def test_a_cancel_drops_every_pipelined_request(server): first = session.flush() second = session.flush() session.cancel() - assert list(session.audio(first)) == [] - assert list(session.audio(second)) == [] # also dropped, must not hang + assert _drain(session.audio(first)) == [] + assert _drain(session.audio(second)) == [] # also dropped, must not hang def test_an_unreadable_frame_ends_the_session_with_an_error(server): @@ -333,7 +354,7 @@ def test_an_unreadable_frame_ends_the_session_with_an_error(server): server.answer_flush = False flush_id = session.flush() server.push({"type": "Audio", "flush_id": "zero"}) # no audio, wrong type - list(session.audio(flush_id)) + _drain(session.audio(flush_id)) def test_a_session_object_can_be_reused(mocker): @@ -341,9 +362,9 @@ def test_a_session_object_can_be_reused(mocker): mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=fakes) session = aai.TTSSession() with session: - assert b"".join(session.speak("One.")) == b"One." + assert b"".join(_drain(session.speak("One."))) == b"One." with session: - assert b"".join(session.speak("Two.")) == b"Two." + assert b"".join(_drain(session.speak("Two."))) == b"Two." assert [f["type"] for f in fakes[1].sent] == ["Generate", "Flush", "Terminate"] @@ -360,7 +381,7 @@ def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): assert session._queue.full() session.generate("y") flush_id = session.flush() - assert b"".join(session.audio(flush_id)) == b"y" # still served after the flood + assert b"".join(_drain(session.audio(flush_id))) == b"y" # still served after the flood assert session._queue.maxsize == PULL_QUEUE_FRAMES assert "dropping audio" in caplog.text @@ -374,7 +395,7 @@ def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): session.generate("Pushed.") flush_id = session.flush() with pytest.raises(RuntimeError, match="Audio handler"): - list(session.audio(flush_id)) + _drain(session.audio(flush_id)) session._stopped.wait(0.2) # let the frames land assert [e.audio for e in received] == [b"Pushed."] @@ -388,7 +409,7 @@ def test_word_boundaries_and_cancelled_reach_their_handlers(server): with aai.TTSSession(word_boundaries=True) as session: session.on(aai.TTSEvents.WordBoundaries, words.append) session.on(aai.TTSEvents.Cancelled, cancelled.append) - list(session.speak("Hello there.")) + _drain(session.speak("Hello there.")) server.push( { "type": "WordBoundaries", @@ -401,7 +422,7 @@ def test_word_boundaries_and_cancelled_reach_their_handlers(server): ) session.cancel() session.keep_alive() - list(session.speak("Again.")) + _drain(session.speak("Again.")) assert parse_qs(urlsplit(server.connect.call_args.args[0]).query)[ "word_boundaries" diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 38c41245..9460e544 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -3,7 +3,7 @@ import asyncio import base64 import json -from typing import Any, List, Optional +from typing import Any, AsyncIterable, List, Optional from urllib.parse import parse_qs, urlsplit import pytest @@ -83,6 +83,15 @@ async def close(self) -> None: self.push(ConnectionClosedOK(Close(1000, ""), None)) +async def _drain(items: AsyncIterable[Any], timeout: float = 5.0) -> List[Any]: + """Collects `items`, failing the test if it does not end in time.""" + + async def collect() -> List[Any]: + return [item async for item in items] + + return await asyncio.wait_for(collect(), timeout) + + @pytest_asyncio.fixture async def server(mocker): # Built inside the running loop: on Python 3.9 an asyncio.Queue binds to @@ -106,7 +115,7 @@ async def test_async_with_connects_speaks_and_terminates(server): config=aai.TTSConfig(voice_id="aai-alba") ) as session: assert session.begin is not None and session.begin.id == "sess-1" - chunks = [chunk async for chunk in session.speak("Hello there.")] + chunks = await _drain(session.speak("Hello there.")) assert b"".join(chunks) == b"Hello there." assert parse_qs(urlsplit(server.connect.uri).query)["voice_id"] == ["aai-alba"] @@ -121,13 +130,13 @@ async def test_pipelined_requests_and_cancel(server): await session.generate("One. ") first = await session.flush() second = await session.flush() - assert b"".join([c async for c in session.audio(first)]) == b"One. " - assert b"".join([c async for c in session.audio(second)]) == b"" + assert b"".join(await _drain(session.audio(first))) == b"One. " + assert b"".join(await _drain(session.audio(second))) == b"" server.answer_flush = False third = await session.flush() await session.cancel() - assert [c async for c in session.audio(third)] == [] + assert await _drain(session.audio(third)) == [] async def test_an_error_before_begin_raises_from_async_with(mocker): @@ -182,7 +191,7 @@ async def test_a_lost_connection_raises_a_retryable_error(server): server.answer_flush = False flush_id = await session.flush() server.push(ConnectionClosedError(None, None)) # abnormal: no close frame - [c async for c in session.audio(flush_id)] + await _drain(session.audio(flush_id)) assert exc_info.value.retryable is True assert exc_info.value.body["close_code"] == 1006 @@ -193,8 +202,8 @@ async def test_a_cancel_drops_every_pipelined_request(server): first = await session.flush() second = await session.flush() await session.cancel() - assert [c async for c in session.audio(first)] == [] - assert [c async for c in session.audio(second)] == [] + assert await _drain(session.audio(first)) == [] + assert await _drain(session.audio(second)) == [] async def test_a_session_object_can_be_reused(mocker): @@ -208,9 +217,9 @@ async def connect(uri, additional_headers): ) session = aai.AsyncTTSSession() async with session: - assert b"".join([c async for c in session.speak("One.")]) == b"One." + assert b"".join(await _drain(session.speak("One."))) == b"One." async with session: - assert b"".join([c async for c in session.speak("Two.")]) == b"Two." + assert b"".join(await _drain(session.speak("Two."))) == b"Two." async def test_missing_key_is_a_clear_error(monkeypatch): From 36b4b4a5178048cadfb0488bb1da6fa137910706 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:52:35 +0000 Subject: [PATCH 05/21] fix(tts): read a session's close code from the close frame ConnectionClosed.code and .reason are deprecated since websockets 13.1. Under -W error the warning was raised inside the reader before the session was marked stopped, so audio() waited forever. Read the code and reason from exc.rcvd instead (1006 when no close frame arrived, which also covers a socket error that is not ConnectionClosed), and make sure any failure in the reader still ends the session with a TTSError. A bare close (no Error frame) is now also retryable for 3005, 3008 and 3009, as the API marks those codes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/async_session.py | 17 +++++++-- assemblyai/tts/v1/session.py | 17 +++++++-- assemblyai/tts/v1/ws_models.py | 24 ++++++++++-- tests/unit/test_tts_session.py | 55 +++++++++++++++++++++++++++- tests/unit/test_tts_session_async.py | 52 ++++++++++++++++++++++++++ 5 files changed, 155 insertions(+), 10 deletions(-) diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index d011baf8..a1217342 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -11,6 +11,7 @@ TTSBeginEvent, TTSEvents, TTSTerminationEvent, + close_code_and_reason, error_from_close, error_from_frame, parse_event, @@ -323,6 +324,18 @@ async def _send(self, frame: Dict[str, Any]) -> None: await self._ws.send(json.dumps(frame)) async def _read_loop(self) -> None: + try: + await self._read_frames() + except Exception as exc: # a failure in the reader itself + await self._fail(TTSError(f"The session's reader failed: {exc!r}")) + finally: + # However the reader ended, nothing may wait for it any longer. + if not self._stopped.is_set(): + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + async def _read_frames(self) -> None: from websockets.exceptions import ConnectionClosed ws = self._ws @@ -382,9 +395,7 @@ async def _dispatch(self, data: Dict[str, Any]) -> None: async def _on_closed(self, exc: Exception) -> None: if self._stopped.is_set(): return - # `.code` is 1006 for an abnormal close, where `.rcvd` is None. - code = getattr(exc, "code", None) - reason = getattr(exc, "reason", "") or "" + code, reason = close_code_and_reason(exc) if code != 1000 and self.error is None: error = error_from_close(code, reason) if self.begin is not None: diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index d3036521..63be5e81 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -12,6 +12,7 @@ TTSBeginEvent, TTSEvents, TTSTerminationEvent, + close_code_and_reason, error_from_close, error_from_frame, parse_event, @@ -390,6 +391,18 @@ def _send(self, frame: Dict[str, Any]) -> None: self._ws.send(json.dumps(frame)) def _read_loop(self) -> None: + try: + self._read_frames() + except Exception as exc: # a failure in the reader itself + self._fail(TTSError(f"The session's reader failed: {exc!r}")) + finally: + # However the reader ended, nothing may wait for it any longer. + if not self._stopped.is_set(): + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + def _read_frames(self) -> None: from websockets.exceptions import ConnectionClosed ws = self._ws @@ -461,9 +474,7 @@ def _dispatch(self, data: Dict[str, Any]) -> None: def _on_closed(self, exc: Exception) -> None: if self._stopped.is_set(): return - # `.code` is 1006 for an abnormal close, where `.rcvd` is None. - code = getattr(exc, "code", None) - reason = getattr(exc, "reason", "") or "" + code, reason = close_code_and_reason(exc) if code != 1000 and self.error is None: error = error_from_close(code, reason) if self.begin is not None: diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py index 5c39b903..71e382b8 100644 --- a/assemblyai/tts/v1/ws_models.py +++ b/assemblyai/tts/v1/ws_models.py @@ -7,7 +7,7 @@ import base64 from enum import Enum -from typing import Any, Dict, List, Optional, Type, TypeVar +from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar try: # pydantic v2 import (`model_validator` exists only in v2) @@ -34,8 +34,12 @@ ] # Close codes after which the same session parameters can succeed on a new -# connection: lost connection, a transient check failure, a server restart. -_RETRYABLE_CLOSE_CODES = frozenset({1006, 1011, 1012}) +# connection, for a close that came without an `Error` frame: a lost +# connection, a transient check failure, a server restart, a server-side +# failure, an expired session and the account's session limit. +_RETRYABLE_CLOSE_CODES = frozenset({1006, 1011, 1012, 3005, 3008, 3009}) +# The code of a connection that ended without a close frame. +_ABNORMAL_CLOSURE = 1006 _M = TypeVar("_M", bound=BaseModel) @@ -178,9 +182,23 @@ def error_from_frame(data: Dict[str, Any]) -> TTSError: ) +def close_code_and_reason(exc: BaseException) -> Tuple[int, str]: + """ + The code and reason of the close frame the API sent, from the exception + that ended the connection; `1006` and no reason when none arrived. + """ + + rcvd = getattr(exc, "rcvd", None) + if rcvd is None: + return _ABNORMAL_CLOSURE, "" + return int(rcvd.code), rcvd.reason or "" + + def error_from_close(code: Optional[int], reason: str) -> TTSError: """The `TTSError` for a connection closed without a clean `Termination`.""" + if code is None: + code = _ABNORMAL_CLOSURE text = f"Connection closed (code={code})" if reason: text = f"{text}: {reason}" diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index d7938328..c10ea38e 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -329,6 +329,57 @@ def test_a_lost_connection_is_a_retryable_error(server): assert len(errors) == 1 +@pytest.mark.filterwarnings("error::DeprecationWarning") +@pytest.mark.parametrize( + "code, retryable", + [ + (1011, True), + (1012, True), + (3005, True), + (3008, True), + (3009, True), + (3006, False), + (3010, False), + ], +) +def test_a_close_without_an_error_frame_is_retryable_by_its_code( + server, code, retryable +): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(Close(code, "Upstream failed"), None)) + _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is retryable + assert exc_info.value.body == {"close_code": code, "reason": "Upstream failed"} + + +def test_a_socket_error_is_a_retryable_lost_connection(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(OSError("Connection reset by peer")) + _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + + +def test_a_failure_in_the_reader_still_ends_the_iteration(server, mocker): + mocker.patch( + "assemblyai.tts.v1.session.error_from_close", side_effect=RuntimeError("bug") + ) + with pytest.raises(aai.TTSError, match="reader failed"): + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(None, None)) + _drain(session.audio(flush_id)) + + def test_a_send_after_a_failure_raises_the_stored_error(server): with pytest.raises(aai.TTSError) as exc_info: with aai.TTSSession() as session: @@ -381,7 +432,9 @@ def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): assert session._queue.full() session.generate("y") flush_id = session.flush() - assert b"".join(_drain(session.audio(flush_id))) == b"y" # still served after the flood + assert ( + b"".join(_drain(session.audio(flush_id))) == b"y" + ) # still served after the flood assert session._queue.maxsize == PULL_QUEUE_FRAMES assert "dropping audio" in caplog.text diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 9460e544..35988837 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -196,6 +196,58 @@ async def test_a_lost_connection_raises_a_retryable_error(server): assert exc_info.value.body["close_code"] == 1006 +@pytest.mark.filterwarnings("error::DeprecationWarning") +@pytest.mark.parametrize( + "code, retryable", + [ + (1011, True), + (1012, True), + (3005, True), + (3008, True), + (3009, True), + (3006, False), + (3010, False), + ], +) +async def test_a_close_without_an_error_frame_is_retryable_by_its_code( + server, code, retryable +): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(Close(code, "Upstream failed"), None)) + await _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is retryable + assert exc_info.value.body == {"close_code": code, "reason": "Upstream failed"} + + +async def test_a_socket_error_is_a_retryable_lost_connection(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(OSError("Connection reset by peer")) + await _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + + +async def test_a_failure_in_the_reader_still_ends_the_iteration(server, mocker): + mocker.patch( + "assemblyai.tts.v1.async_session.error_from_close", + side_effect=RuntimeError("bug"), + ) + with pytest.raises(aai.TTSError, match="reader failed"): + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(None, None)) + await _drain(session.audio(flush_id)) + + async def test_a_cancel_drops_every_pipelined_request(server): server.answer_flush = False async with aai.AsyncTTSSession() as session: From 60b8e96ad80402dcfdd0a70b98184bb6ebd136b5 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:54:54 +0000 Subject: [PATCH 06/21] fix(tts): never drop the end of a request from a full pull queue When the pull queue was full, the oldest item of any kind was evicted, including the FlushDone, Cancelled, error and closed markers. Losing the marker that ends the request being read left audio(flush_id) discarding every later request's items and then waiting until the session ended. Only audio frames now count toward the cap and only the oldest audio frame is dropped; markers are always kept. The warning also covers a reader that is reading, just more slowly than the audio arrives. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/_base_ws.py | 54 +++++++++++++++++++++++- assemblyai/tts/v1/async_session.py | 62 +++++++++++++++++----------- assemblyai/tts/v1/session.py | 60 ++++++++++++++++----------- tests/unit/test_tts_session.py | 35 +++++++++++++--- tests/unit/test_tts_session_async.py | 27 ++++++++++++ 5 files changed, 180 insertions(+), 58 deletions(-) diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 9bb9d99f..48c749a4 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -1,7 +1,8 @@ """Helpers shared by the sync and async Text-to-Speech sessions.""" import platform -from typing import Any, Dict, Optional, Tuple +from collections import deque +from typing import Any, Deque, Dict, Optional, Tuple from urllib.parse import urlencode, urlsplit from .. import _audio @@ -14,9 +15,58 @@ # Characters a single `Generate` frame may carry (Unicode code points). MAX_GENERATE_CHARS = 2000 # Audio frames queued for pulling before the oldest is dropped: about 160 s -# at the API's ~80 ms per frame. Bounds memory when nothing reads them. +# at the API's ~80 ms per frame. Bounds memory when audio is not read, or is +# read more slowly than it arrives. PULL_QUEUE_FRAMES = 2000 +# An item `audio()` reads: (kind, flush_id or cancel id, payload). `audio` +# items carry a request's audio; the markers `done` and `cancelled` end a +# request, and `error` and `closed` end the session. +PullItem = Tuple[str, Any, Any] + + +class PullBuffer: + """ + The items `audio()` reads, in arrival order. + + Markers are always kept, since an `audio()` iteration ends only at one. + Past `max_audio` audio frames, the oldest audio frame is dropped. + """ + + def __init__(self, max_audio: int) -> None: + self.max_audio = max_audio + self._items: Deque[PullItem] = deque() + self._audio = 0 + + def __len__(self) -> int: + return len(self._items) + + def append(self, item: PullItem) -> bool: + """Adds `item`; returns whether an older audio frame was dropped.""" + + self._items.append(item) + if item[0] != "audio": + return False + self._audio += 1 + if self._audio <= self.max_audio: + return False + for index, queued in enumerate(self._items): + if queued[0] == "audio": + del self._items[index] + self._audio -= 1 + return True + return False + + def popleft(self) -> Optional[PullItem]: + """Removes and returns the oldest item; `None` when there is none.""" + + if not self._items: + return None + item = self._items.popleft() + if item[0] == "audio": + self._audio -= 1 + return item + def check_text_frame(text: object) -> None: if not isinstance(text, str): diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index a1217342..da27120b 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -1,11 +1,18 @@ import asyncio import json import logging -from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union from ... import _audio from ... import async_client as _async_client -from .._base_ws import PULL_QUEUE_FRAMES, build_uri, check_text_frame, resolve_settings +from .._base_ws import ( + PULL_QUEUE_FRAMES, + PullBuffer, + PullItem, + build_uri, + check_text_frame, + resolve_settings, +) from .models import TTSConfig, TTSError from .ws_models import ( TTSBeginEvent, @@ -22,7 +29,6 @@ logger = logging.getLogger(__name__) _Handler = Callable[[Any], Any] -_Item = Tuple[str, Any, Any] def websocket_connect_async( @@ -86,6 +92,7 @@ def __init__( self._begun: Any = None self._stopped: Any = None self._queue: Any = None + self._queue_ready: Any = None self._next_flush_id = 0 self._cancel_ids: List[int] = [] self._cancelled_before = 0 @@ -111,7 +118,8 @@ async def connect(self) -> TTSBeginEvent: raise RuntimeError("the session is already connected") self._begun = asyncio.Event() self._stopped = asyncio.Event() - self._queue = asyncio.Queue(maxsize=PULL_QUEUE_FRAMES) + self._queue = PullBuffer(PULL_QUEUE_FRAMES) + self._queue_ready = asyncio.Event() self._next_flush_id = 0 self._cancel_ids = [] self._cancelled_before = 0 @@ -259,16 +267,14 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: while True: if flush_id < self._cancelled_before: # dropped by a Cancel return - try: - kind, ident, payload = await asyncio.wait_for( - self._queue.get(), timeout=0.5 - ) - except asyncio.TimeoutError: + item = await self._get(timeout=0.5) + if item is None: if self._stopped.is_set(): if self.error is not None: raise self.error return continue + kind, ident, payload = item if kind == "audio": if ident == flush_id: yield payload @@ -285,24 +291,30 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: raise self.error return - def _put(self, item: _Item) -> None: - """Queues an item for `audio()`, dropping the oldest audio when full.""" + def _put(self, item: PullItem) -> None: + """Queues an item for `audio()`; past the cap the oldest audio drops.""" + + dropped = self._queue.append(item) + self._queue_ready.set() + if dropped and not self._drop_warned: + self._drop_warned = True + logger.warning( + "AsyncTTSSession: over %d audio frames are waiting to be read, so the " + "oldest are being dropped; read audio() as fast as the audio " + "arrives, or register an Audio handler", + self._queue.max_audio, + ) + + async def _get(self, timeout: float) -> Optional[PullItem]: + """The next item for `audio()`; `None` when none came in `timeout` s.""" - while True: + if not self._queue: + self._queue_ready.clear() try: - self._queue.put_nowait(item) - return - except asyncio.QueueFull: - try: - dropped = self._queue.get_nowait() - except asyncio.QueueEmpty: - continue - if dropped[0] == "audio" and not self._drop_warned: - self._drop_warned = True - logger.warning( - "AsyncTTSSession: dropping audio nobody reads; iterate " - "audio() or register an Audio handler" - ) + await asyncio.wait_for(self._queue_ready.wait(), timeout) + except asyncio.TimeoutError: + return None + return self._queue.popleft() async def _fail(self, error: TTSError) -> None: if self.error is None: diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 63be5e81..37423112 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -1,12 +1,18 @@ import json import logging -import queue import threading -from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterator, List, Optional, Union from ... import _audio from ... import client as _client -from .._base_ws import PULL_QUEUE_FRAMES, build_uri, check_text_frame, resolve_settings +from .._base_ws import ( + PULL_QUEUE_FRAMES, + PullBuffer, + PullItem, + build_uri, + check_text_frame, + resolve_settings, +) from .models import TTSConfig, TTSError from .ws_models import ( TTSBeginEvent, @@ -23,7 +29,6 @@ logger = logging.getLogger(__name__) _Handler = Callable[[Any], None] -_Item = Tuple[str, Any, Any] # (kind, flush_id or cancel id, payload) def websocket_connect(uri: str, **kwargs: Any) -> Any: @@ -136,7 +141,8 @@ def _reset(self) -> None: self._begun = threading.Event() self._stopped = threading.Event() - self._queue: "queue.Queue[_Item]" = queue.Queue(maxsize=PULL_QUEUE_FRAMES) + self._queue = PullBuffer(PULL_QUEUE_FRAMES) + self._queue_ready = threading.Condition() self._next_flush_id = 0 self._cancel_ids: List[int] = [] # Requests with a lower flush_id than this were dropped by a Cancel. @@ -335,14 +341,14 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: while True: if flush_id < self._cancelled_before: # dropped by a Cancel return - try: - kind, ident, payload = self._queue.get(timeout=0.5) - except queue.Empty: + item = self._get(timeout=0.5) + if item is None: if self._stopped.is_set(): if self.error is not None: raise self.error return continue + kind, ident, payload = item if kind == "audio": if ident == flush_id: yield payload @@ -361,24 +367,28 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: # -- internals --------------------------------------------------------------- - def _put(self, item: _Item) -> None: - """Queues an item for `audio()`, dropping the oldest audio when full.""" + def _put(self, item: PullItem) -> None: + """Queues an item for `audio()`; past the cap the oldest audio drops.""" + + with self._queue_ready: + dropped = self._queue.append(item) + self._queue_ready.notify_all() + if dropped and not self._drop_warned: + self._drop_warned = True + logger.warning( + "TTSSession: over %d audio frames are waiting to be read, so the " + "oldest are being dropped; read audio() as fast as the audio " + "arrives, or register an Audio handler", + self._queue.max_audio, + ) - while True: - try: - self._queue.put_nowait(item) - return - except queue.Full: - try: - dropped = self._queue.get_nowait() - except queue.Empty: - continue - if dropped[0] == "audio" and not self._drop_warned: - self._drop_warned = True - logger.warning( - "TTSSession: dropping audio nobody reads; iterate audio() " - "or register an Audio handler" - ) + def _get(self, timeout: float) -> Optional[PullItem]: + """The next item for `audio()`; `None` when none came in `timeout` s.""" + + with self._queue_ready: + if not self._queue: + self._queue_ready.wait(timeout) + return self._queue.popleft() def _send(self, frame: Dict[str, Any]) -> None: if self._ws is None or self._stopped.is_set(): diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index c10ea38e..6dc89bac 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -429,14 +429,37 @@ def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): deadline = time.monotonic() + 5 while not server.inbox.empty() and time.monotonic() < deadline: time.sleep(0.01) # let the reader queue the flood before anyone reads - assert session._queue.full() + assert len(session._queue) == PULL_QUEUE_FRAMES session.generate("y") flush_id = session.flush() - assert ( - b"".join(_drain(session.audio(flush_id))) == b"y" - ) # still served after the flood - assert session._queue.maxsize == PULL_QUEUE_FRAMES - assert "dropping audio" in caplog.text + # still served after the flood + assert b"".join(_drain(session.audio(flush_id))) == b"y" + assert session._queue.max_audio == PULL_QUEUE_FRAMES + assert "oldest are being dropped" in caplog.text + + +def test_an_overflow_drops_audio_but_never_the_end_of_a_request(server, mocker, caplog): + mocker.patch("assemblyai.tts.v1.session.PULL_QUEUE_FRAMES", 6) + server.answer_flush = False + done: List[int] = [] + with aai.TTSSession() as session: + session.on(aai.TTSEvents.FlushDone, lambda event: done.append(event.flush_id)) + flush_ids = [session.flush() for _ in range(3)] + for flush_id in flush_ids: # 4 frames each, all before anything is read + for _ in range(4): + server.push({"type": "Audio", "audio": "AA==", "flush_id": flush_id}) + server.push( + {"type": "FlushDone", "flush_id": flush_id, "audio_duration_ms": 320} + ) + deadline = time.monotonic() + 5 + while len(done) < 3 and time.monotonic() < deadline: + time.sleep(0.01) + + got = [_drain(session.audio(flush_id)) for flush_id in flush_ids] + + # 12 frames for room for 6: the oldest 6 went, every request still ended + assert [len(chunks) for chunks in got] == [0, 2, 4] + assert "oldest are being dropped" in caplog.text def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 35988837..3a51cbcc 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -258,6 +258,33 @@ async def test_a_cancel_drops_every_pipelined_request(server): assert await _drain(session.audio(second)) == [] +async def test_an_overflow_drops_audio_but_never_the_end_of_a_request( + server, mocker, caplog +): + mocker.patch("assemblyai.tts.v1.async_session.PULL_QUEUE_FRAMES", 6) + server.answer_flush = False + done: List[int] = [] + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.FlushDone, lambda event: done.append(event.flush_id)) + flush_ids = [await session.flush() for _ in range(3)] + for flush_id in flush_ids: # 4 frames each, all before anything is read + for _ in range(4): + server.push({"type": "Audio", "audio": "AA==", "flush_id": flush_id}) + server.push( + {"type": "FlushDone", "flush_id": flush_id, "audio_duration_ms": 320} + ) + for _ in range(500): + if len(done) == 3: + break + await asyncio.sleep(0.01) + + got = [await _drain(session.audio(flush_id)) for flush_id in flush_ids] + + # 12 frames for room for 6: the oldest 6 went, every request still ended + assert [len(chunks) for chunks in got] == [0, 2, 4] + assert "oldest are being dropped" in caplog.text + + async def test_a_session_object_can_be_reused(mocker): fakes = [_FakeAsyncServer(), _FakeAsyncServer()] From 81bb42f4adc1b5f26e9dec47d9add6e155abd344 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:56:15 +0000 Subject: [PATCH 07/21] fix(tts): number flushes and cancels in wire order across threads flush() and cancel() took their flush_id before taking the send lock, so a cancel() from another thread (a barge-in from a VAD or STT callback while the app thread streams text) could reach the socket in a different order from its id. The API numbers sequences in wire order, so audio() then waited for an id the API never sent, or dropped a request's audio as cancelled. Take the id, record the cancel and write the frame under one reentrant lock. The reader no longer raises the cancel watermark itself: cancel() has already raised it, and a second writer on another thread could lower it again. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/session.py | 24 ++++++++++++++---------- tests/unit/test_tts_session.py | 23 +++++++++++++++++++++++ 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 37423112..a6baded7 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -126,7 +126,9 @@ def __init__( self._handlers: Dict[str, List[_Handler]] = {} self._ws: Any = None self._reader: Optional[threading.Thread] = None - self._send_lock = threading.Lock() + # Held while a frame is written. flush() and cancel() hold it while + # they also take the frame's flush_id, so the ids follow wire order. + self._send_lock = threading.RLock() self.begin: Optional[TTSBeginEvent] = None "The `Begin` frame, once connected." self.termination: Optional[TTSTerminationEvent] = None @@ -276,9 +278,10 @@ def flush(self) -> int: Everything sent so far is spoken, then `FlushDone` arrives for this id. Flushes can be pipelined; the audio comes in order. """ - flush_id = self._next_flush_id - self._next_flush_id += 1 - self._send({"type": "Flush"}) + with self._send_lock: + flush_id = self._next_flush_id + self._next_flush_id += 1 + self._send({"type": "Flush"}) return flush_id def cancel(self) -> None: @@ -290,11 +293,12 @@ def cancel(self) -> None: Text sent and not flushed is discarded. Like `flush()`, a cancel takes a `flush_id`, so the next request's id skips one. """ - cancel_id = self._next_flush_id - self._next_flush_id += 1 - self._cancel_ids.append(cancel_id) - self._cancelled_before = max(self._cancelled_before, cancel_id) - self._send({"type": "Cancel"}) + with self._send_lock: + cancel_id = self._next_flush_id + self._next_flush_id += 1 + self._cancel_ids.append(cancel_id) + self._cancelled_before = max(self._cancelled_before, cancel_id) + self._send({"type": "Cancel"}) def keep_alive(self) -> None: """Resets the inactivity timer between sparse requests.""" @@ -471,8 +475,8 @@ def _dispatch(self, data: Dict[str, Any]) -> None: elif kind == "FlushDone": self._put(("done", event.flush_id, None)) elif kind == "Cancelled": + # cancel() already raised the watermark for this id. cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 - self._cancelled_before = max(self._cancelled_before, cancel_id) self._put(("cancelled", cancel_id, None)) elif kind == "Termination": self.termination = event diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 6dc89bac..108aaf9f 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -51,6 +51,7 @@ def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): self.text = "" self.flush_id = 0 self.answer_flush = answer_flush + self.generate_delay = 0.0 if begin is not None: self.push(begin) @@ -62,6 +63,7 @@ def send(self, raw: str) -> None: self.sent.append(frame) kind = frame["type"] if kind == "Generate": + time.sleep(self.generate_delay) self.text += frame["text"] elif kind == "Flush": if self.answer_flush: @@ -203,6 +205,27 @@ def test_cancel_ends_the_current_request_but_not_a_later_one(server): ] +def test_a_cancel_from_another_thread_keeps_flush_ids_in_wire_order(server): + server.generate_delay = 0.02 # a slow write the cancel lands in + with aai.TTSSession() as session: + for _ in range(5): + flush_ids: List[int] = [] + + def write_reply() -> None: + session.generate("Hello there. ") + flush_ids.append(session.flush()) + + writer = threading.Thread(target=write_reply) + writer.start() + time.sleep(0.01) + session.cancel() # barge-in from a VAD callback, say + writer.join(5) + _drain(session.audio(flush_ids[0])) # ends, cancelled or spoken + + # every Flush and Cancel took the id the API gives it + assert session._next_flush_id == server.flush_id == 10 + + def test_cancel_ends_the_iteration_before_the_server_answers(server): server.answer_flush = False with aai.TTSSession() as session: From 0a0d8ed924bc5686764b327cb09e4ce170fe8812 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:57:14 +0000 Subject: [PATCH 08/21] fix(tts): raise TTSError when a send races the API's close If the API closed the socket (a 3010 input_rate_exceeded while text is streaming in, say) and generate(), flush() or cancel() ran before the reader had handled the Error frame, the call raised websockets' ConnectionClosedError, losing error_type, param and retryable. A send that hits a closed socket now gives the reader a moment to take in what arrived first and raises the session's TTSError, or one built from the close code. The socket is read into a local once, so a close() on another thread can no longer turn the send into an AttributeError. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/async_session.py | 28 +++++++++++++++-- assemblyai/tts/v1/session.py | 35 +++++++++++++++++----- tests/unit/test_tts_session.py | 45 ++++++++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 45 ++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+), 9 deletions(-) diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index da27120b..72e92e35 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -327,13 +327,37 @@ async def _fail(self, error: TTSError) -> None: self._begun.set() async def _send(self, frame: Dict[str, Any]) -> None: - if self._ws is None or self._stopped.is_set(): + from websockets.exceptions import ConnectionClosed + + ws = self._ws + if ws is None or self._stopped.is_set(): if self.error is not None: raise self.error raise RuntimeError( "the session is not open: use `with` / `async with`, or connect()" ) - await self._ws.send(json.dumps(frame)) + try: + await ws.send(json.dumps(frame)) + except ConnectionClosed as exc: + raise await self._closed_error(ws, exc) from exc + + async def _closed_error(self, ws: Any, exc: Exception) -> Exception: + """ + What a send on a socket that closed raises: the session's error, once + the reader has taken in the frames that arrived before the close. + """ + + if asyncio.current_task() is not self._reader: + try: + await asyncio.wait_for(self._stopped.wait(), timeout=1) + except asyncio.TimeoutError: + pass + if self.error is not None: + return self.error + code, reason = close_code_and_reason(exc) + if self._ws is not ws or code == 1000: # closed by close() or cleanly + return RuntimeError("the session is not open") + return error_from_close(code, reason) async def _read_loop(self) -> None: try: diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index a6baded7..79f56f2f 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -395,14 +395,35 @@ def _get(self, timeout: float) -> Optional[PullItem]: return self._queue.popleft() def _send(self, frame: Dict[str, Any]) -> None: - if self._ws is None or self._stopped.is_set(): - if self.error is not None: - raise self.error - raise RuntimeError( - "the session is not open: use `with` / `async with`, or connect()" - ) + from websockets.exceptions import ConnectionClosed + with self._send_lock: - self._ws.send(json.dumps(frame)) + ws = self._ws + if ws is None or self._stopped.is_set(): + if self.error is not None: + raise self.error + raise RuntimeError( + "the session is not open: use `with` / `async with`, or connect()" + ) + try: + ws.send(json.dumps(frame)) + except ConnectionClosed as exc: + raise self._closed_error(ws, exc) from exc + + def _closed_error(self, ws: Any, exc: Exception) -> Exception: + """ + What a send on a socket that closed raises: the session's error, once + the reader has taken in the frames that arrived before the close. + """ + + if threading.current_thread() is not self._reader: + self._stopped.wait(1) + if self.error is not None: + return self.error + code, reason = close_code_and_reason(exc) + if self._ws is not ws or code == 1000: # closed by close() or cleanly + return RuntimeError("the session is not open") + return error_from_close(code, reason) def _read_loop(self) -> None: try: diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 108aaf9f..63c060d9 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -39,6 +39,17 @@ } +_INPUT_RATE_EXCEEDED = { + "type": "Error", + "error_code": 3010, + "error": "Text arrived faster than it can be spoken.", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, +} + + def _closed_ok() -> ConnectionClosedOK: return ConnectionClosedOK(Close(1000, "Session ended"), None) @@ -52,13 +63,26 @@ def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): self.flush_id = 0 self.answer_flush = answer_flush self.generate_delay = 0.0 + self.closing: Optional[List[Any]] = None if begin is not None: self.push(begin) + def close_on_next_send(self, code: int, error: Optional[dict] = None) -> None: + """The API closes the session: the next send raises, and the reader + has yet to take in the frames that came first.""" + self.closing = ([error] if error else []) + [ + ConnectionClosedError(Close(code, "See Error message for details"), None) + ] + def push(self, item: Any) -> None: self.inbox.put(json.dumps(item) if isinstance(item, dict) else item) def send(self, raw: str) -> None: + if self.closing is not None: + closing, self.closing = self.closing, None + for item in closing: + self.push(item) + raise ConnectionClosedError(closing[-1].rcvd, None) frame = json.loads(raw) self.sent.append(frame) kind = frame["type"] @@ -412,6 +436,27 @@ def test_a_send_after_a_failure_raises_the_stored_error(server): assert exc_info.value.body["close_code"] == 1006 +def test_a_send_the_api_closed_raises_the_sessions_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.close_on_next_send(3010, _INPUT_RATE_EXCEEDED) + session.generate("Too much, too fast.") + + assert exc_info.value.error_code == "input_rate_exceeded" + assert exc_info.value.retryable is False + assert exc_info.value is session.error + + +def test_a_send_on_a_bare_close_raises_a_tts_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.close_on_next_send(1012) + session.flush() + + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1012 + + def test_a_cancel_drops_every_pipelined_request(server): server.answer_flush = False with aai.TTSSession() as session: diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 3a51cbcc..6d4b3922 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -31,6 +31,17 @@ } +_INPUT_RATE_EXCEEDED = { + "type": "Error", + "error_code": 3010, + "error": "Text arrived faster than it can be spoken.", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, +} + + class _FakeAsyncServer: def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): self.sent: List[dict] = [] @@ -39,13 +50,26 @@ def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): self.text = "" self.flush_id = 0 self.answer_flush = answer_flush + self.closing: Optional[List[Any]] = None if begin is not None: self.push(begin) + def close_on_next_send(self, code: int, error: Optional[dict] = None) -> None: + """The API closes the session: the next send raises, and the reader + has yet to take in the frames that came first.""" + self.closing = ([error] if error else []) + [ + ConnectionClosedError(Close(code, "See Error message for details"), None) + ] + def push(self, item: Any) -> None: self.inbox.put_nowait(json.dumps(item) if isinstance(item, dict) else item) async def send(self, raw: str) -> None: + if self.closing is not None: + closing, self.closing = self.closing, None + for item in closing: + self.push(item) + raise ConnectionClosedError(closing[-1].rcvd, None) frame = json.loads(raw) self.sent.append(frame) kind = frame["type"] @@ -248,6 +272,27 @@ async def test_a_failure_in_the_reader_still_ends_the_iteration(server, mocker): await _drain(session.audio(flush_id)) +async def test_a_send_the_api_closed_raises_the_sessions_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.close_on_next_send(3010, _INPUT_RATE_EXCEEDED) + await session.generate("Too much, too fast.") + + assert exc_info.value.error_code == "input_rate_exceeded" + assert exc_info.value.retryable is False + assert exc_info.value is session.error + + +async def test_a_send_on_a_bare_close_raises_a_tts_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.close_on_next_send(1012) + await session.flush() + + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1012 + + async def test_a_cancel_drops_every_pipelined_request(server): server.answer_flush = False async with aai.AsyncTTSSession() as session: From 377107f8529cbc9eafae69cf2a21fc13768d781a Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:58:23 +0000 Subject: [PATCH 09/21] fix(tts): make AsyncTTSSession cancellation safe A cancellation (an outer timeout, a TaskGroup, a call hanging up) could leave the session half open or hide itself: - connect() cancelled while waiting for Begin left the socket open and the reader running, and the next connect() said the session was already connected. It now closes before re-raising. - terminate() cancelled during its waits skipped close(). close() now runs in a finally. - close() caught CancelledError from its wait on the reader, swallowing the caller's own cancellation. It now waits with asyncio.wait, cancels a reader still running, and lets the cancellation propagate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/async_session.py | 51 +++++++++++++---------- tests/unit/test_tts_session_async.py | 61 +++++++++++++++++++++++++++- 2 files changed, 90 insertions(+), 22 deletions(-) diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index 72e92e35..f2eec3f1 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -152,6 +152,9 @@ async def connect(self) -> TTSBeginEvent: except asyncio.TimeoutError: await self.close() raise TTSError("The session did not begin in time") from None + except BaseException: # cancelled: leave no socket or reader behind + await self.close() + raise if self.error is not None: await self.close() raise self.error @@ -165,18 +168,20 @@ async def terminate(self) -> Optional[TTSTerminationEvent]: if self._ws is None: return None - if not self._stopped.is_set(): - try: - await self._send({"type": "Terminate"}) - except Exception: - pass - try: - await asyncio.wait_for( - self._stopped.wait(), timeout=self._terminate_timeout - ) - except asyncio.TimeoutError: - pass - await self.close() + try: + if not self._stopped.is_set(): + try: + await self._send({"type": "Terminate"}) + except Exception: + pass + try: + await asyncio.wait_for( + self._stopped.wait(), timeout=self._terminate_timeout + ) + except asyncio.TimeoutError: + pass + finally: + await self.close() return self.termination async def close(self) -> None: @@ -187,18 +192,22 @@ async def close(self) -> None: self._stopped.set() if self._begun is not None: self._begun.set() - if ws is not None: - try: - await ws.close() - except Exception: - pass if self._queue is not None: self._put(("closed", None, None)) reader = self._reader - if reader is not None and reader is not asyncio.current_task(): - try: - await asyncio.wait_for(reader, timeout=self._terminate_timeout) - except (asyncio.TimeoutError, asyncio.CancelledError): + if reader is asyncio.current_task(): + reader = None + try: + if ws is not None: + try: + await ws.close() + except Exception: + pass + if reader is not None and not reader.done(): + await asyncio.wait({reader}, timeout=self._terminate_timeout) + finally: + # Also when the caller is cancelled, which then propagates. + if reader is not None and not reader.done(): reader.cancel() async def __aenter__(self) -> "AsyncTTSSession": diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 6d4b3922..f452cb2e 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -50,6 +50,7 @@ def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): self.text = "" self.flush_id = 0 self.answer_flush = answer_flush + self.answer_terminate = True self.closing: Optional[List[Any]] = None if begin is not None: self.push(begin) @@ -92,7 +93,7 @@ async def send(self, raw: str) -> None: self.flush_id += 1 self.text = "" self.push({"type": "Cancelled", "audio_duration_seconds": 0.1}) - elif kind == "Terminate": + elif kind == "Terminate" and self.answer_terminate: self.push(_TERMINATION) self.push(ConnectionClosedOK(Close(1000, "Session ended"), None)) @@ -346,6 +347,64 @@ async def connect(uri, additional_headers): assert b"".join(await _drain(session.speak("Two."))) == b"Two." +async def test_a_cancelled_connect_leaves_nothing_open(mocker): + slow = _FakeAsyncServer(begin=None) # Begin never comes + fakes = [slow, _FakeAsyncServer()] + + async def connect(uri, additional_headers): + return fakes.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(session.connect(), 0.3) + + assert slow.closed + assert session._reader is not None and session._reader.done() + begin = await asyncio.wait_for(session.connect(), 5) # not "already connected" + assert begin.id == "sess-1" + await session.close() + + +async def test_a_cancelled_terminate_still_closes(server): + server.answer_terminate = False + session = aai.AsyncTTSSession(terminate_timeout=30) + await session.connect() + ending = asyncio.create_task(session.terminate()) + await asyncio.sleep(0.05) + ending.cancel() + await asyncio.wait({ending}, timeout=5) + + assert ending.cancelled() + assert server.closed + assert session._ws is None + + +async def test_a_cancelled_close_propagates_the_cancellation(server): + in_handler = asyncio.Event() + + async def busy_handler(event): + in_handler.set() + await asyncio.sleep(30) + + session = aai.AsyncTTSSession(terminate_timeout=30) + await session.connect() + session.on(aai.TTSEvents.FlushDone, busy_handler) + await session.flush() + await asyncio.wait_for(in_handler.wait(), 5) + + closing = asyncio.create_task(session.close()) + await asyncio.sleep(0.05) # close() now waits for the busy reader + closing.cancel() + await asyncio.wait({closing}, timeout=5) + + assert closing.cancelled() + await asyncio.sleep(0) + assert session._reader is not None and session._reader.done() + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 2992e609dae6c90e101ad76d2ca9ba1257b54167 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:00:28 +0000 Subject: [PATCH 10/21] fix(tts): keep session handlers from waiting on their own reader Handlers run on the session's reader, and nothing stopped them from calling methods that wait on it. A FlushDone handler that iterated speak() for the next sentence deadlocked both itself and the app's audio(), and an Error handler that called terminate() blocked the reader for the whole terminate_timeout. speak(), audio() and terminate() now raise RuntimeError when called from a handler. When the session fails, waiters see the stopped state and the error before the Error handler runs, so a handler that waits on the app no longer holds the app's audio() open. The README says that sending from a handler is fine and waiting is not. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- README.md | 2 +- assemblyai/tts/v1/async_session.py | 28 ++++++++++++++--- assemblyai/tts/v1/session.py | 36 ++++++++++++++++++--- tests/unit/test_tts_session.py | 40 +++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 47 ++++++++++++++++++++++++++++ 5 files changed, 143 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index b60bae90..db4d2359 100644 --- a/README.md +++ b/README.md @@ -1052,7 +1052,7 @@ Audio starts before `flush()` once whole sentences have arrived. To play it whil Limits: a `Generate` frame takes at most 2,000 characters (split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) without a frame; send `keep_alive()` between sparse requests. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. -Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. +Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. A handler may send (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but must not wait on the session: the reader that runs the handler is what `speak()`, `audio()` and `terminate()` wait for, so they raise `RuntimeError` there. Hand the work to another thread (or task) instead. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status) and `request_id` is the session id; `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index f2eec3f1..d83cc5f4 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -50,7 +50,9 @@ class AsyncTTSSession: Open it with `async with`, which connects and waits for `Begin`. `generate()`, `flush()`, `cancel()` and `keep_alive()` are coroutines; `speak()` and `audio()` are async iterators of audio bytes. Handlers - registered with `on()` may be plain functions or coroutines. + registered with `on()` may be plain functions or coroutines; they run on + the session's reader task, so, as with `TTSSession`, they may send but + not wait on the session (`speak()`, `audio()`, `terminate()`). Example: ```python @@ -166,6 +168,7 @@ async def connect(self) -> TTSBeginEvent: async def terminate(self) -> Optional[TTSTerminationEvent]: """Ends the session cleanly; see `TTSSession.terminate`.""" + self._check_not_in_handler("terminate()") if self._ws is None: return None try: @@ -249,6 +252,7 @@ async def keep_alive(self) -> None: async def speak(self, text: str) -> AsyncIterator[bytes]: """Speaks one text and yields its audio; see `TTSSession.speak`.""" + self._check_not_in_handler("speak()") await self.generate(text) flush_id = await self.flush() async for chunk in self.audio(flush_id): @@ -262,6 +266,7 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: play it while another still sends text. """ + self._check_not_in_handler("audio()") if self._queue is None: raise RuntimeError("the session is not open") if flush_id is None: @@ -300,6 +305,16 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: raise self.error return + def _check_not_in_handler(self, method: str) -> None: + """Refuses a method that waits on the reader when the reader, running + a handler, is the caller: it would wait for itself.""" + + if self._reader is not None and asyncio.current_task() is self._reader: + raise RuntimeError( + f"{method} cannot be called from a session handler: it waits " + "for the session's reader, which runs the handlers" + ) + def _put(self, item: PullItem) -> None: """Queues an item for `audio()`; past the cap the oldest audio drops.""" @@ -326,14 +341,17 @@ async def _get(self, timeout: float) -> Optional[PullItem]: return self._queue.popleft() async def _fail(self, error: TTSError) -> None: - if self.error is None: + first = self.error is None + if first: if error.request_id is None and self.begin is not None: error.request_id = self.begin.id self.error = error - await self._emit(TTSEvents.Error, error) + # Waiters learn of the failure before any handler runs. self._put(("error", None, self.error)) self._stopped.set() self._begun.set() + if first: + await self._emit(TTSEvents.Error, error) async def _send(self, frame: Dict[str, Any]) -> None: from websockets.exceptions import ConnectionClosed @@ -441,15 +459,17 @@ async def _on_closed(self, exc: Exception) -> None: if self._stopped.is_set(): return code, reason = close_code_and_reason(exc) + error = None if code != 1000 and self.error is None: error = error_from_close(code, reason) if self.begin is not None: error.request_id = self.begin.id self.error = error - await self._emit(TTSEvents.Error, self.error) self._stopped.set() self._begun.set() self._put(("closed", None, None)) + if error is not None: + await self._emit(TTSEvents.Error, error) async def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: for handler in self._handlers.get(str(event), ()): diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 79f56f2f..68498928 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -63,7 +63,10 @@ class TTSSession: pushed to a handler registered with `on(TTSEvents.Audio, ...)`. Register handlers for the other events (`FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`) as needed; a handler takes the event - and runs on the session's reader thread. + and runs on the session's reader thread. A handler may send + (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but not + wait on the session: `speak()`, `audio()` and `terminate()` raise + `RuntimeError` there. Example: ```python @@ -218,7 +221,11 @@ def terminate(self) -> Optional[TTSTerminationEvent]: Text sent and not flushed is spoken first. Returns the summary, or `None` if the session had already ended or the wait timed out. + + Raises: + RuntimeError: from a session handler. """ + self._check_not_in_handler("terminate()") if self._ws is None: return None if not self._stopped.is_set(): @@ -314,7 +321,9 @@ def speak(self, text: str) -> Iterator[bytes]: Raises: TTSError: from the iteration, if the session fails. + RuntimeError: from a session handler. """ + self._check_not_in_handler("speak()") self.generate(text) return self.audio(self.flush()) @@ -331,8 +340,10 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: Raises: TTSError: if the session fails before the request completes. ValueError: for a `flush_id` no `flush()` has returned. - RuntimeError: when an `Audio` handler receives the audio instead. + RuntimeError: when an `Audio` handler receives the audio instead, + or from a session handler. """ + self._check_not_in_handler("audio()") if flush_id is None: flush_id = self._next_flush_id # the request being written now elif flush_id > self._next_flush_id: @@ -371,6 +382,16 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: # -- internals --------------------------------------------------------------- + def _check_not_in_handler(self, method: str) -> None: + """Refuses a method that waits on the reader when the reader, running + a handler, is the caller: it would wait for itself.""" + + if self._reader is not None and threading.current_thread() is self._reader: + raise RuntimeError( + f"{method} cannot be called from a session handler: it waits " + "for the session's reader, which runs the handlers" + ) + def _put(self, item: PullItem) -> None: """Queues an item for `audio()`; past the cap the oldest audio drops.""" @@ -470,14 +491,17 @@ def _read_frames(self) -> None: def _fail(self, error: TTSError) -> None: """Ends the session on an error: handlers, waiters, state.""" - if self.error is None: + first = self.error is None + if first: if error.request_id is None and self.begin is not None: error.request_id = self.begin.id self.error = error - self._emit(TTSEvents.Error, error) + # Waiters learn of the failure before any handler runs. self._put(("error", None, self.error)) self._stopped.set() self._begun.set() + if first: + self._emit(TTSEvents.Error, error) def _dispatch(self, data: Dict[str, Any]) -> None: kind = data.get("type") @@ -510,15 +534,17 @@ def _on_closed(self, exc: Exception) -> None: if self._stopped.is_set(): return code, reason = close_code_and_reason(exc) + error = None if code != 1000 and self.error is None: error = error_from_close(code, reason) if self.begin is not None: error.request_id = self.begin.id self.error = error - self._emit(TTSEvents.Error, self.error) self._stopped.set() self._begun.set() self._put(("closed", None, None)) + if error is not None: + self._emit(TTSEvents.Error, error) def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: for handler in self._handlers.get(str(event), ()): diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 63c060d9..a701357e 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -547,6 +547,46 @@ def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): assert done[0].audio_duration_ms == 80 +def test_the_error_reaches_waiters_before_the_error_handler(server): + app_saw_error = threading.Event() + handler_waited: List[bool] = [] + + def on_error(error: aai.TTSError) -> None: + handler_waited.append(app_saw_error.wait(5)) # waits on the app + + with pytest.raises(aai.TTSError): + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, on_error) + server.answer_flush = False + flush_id = session.flush() + server.push(_INPUT_RATE_EXCEEDED) + try: + _drain(session.audio(flush_id), timeout=3) + finally: + app_saw_error.set() + + assert handler_waited == [True] + + +def test_waiting_from_a_handler_raises_instead_of_deadlocking(server): + refused: List[str] = [] + + def speak_next(event: aai.TTSFlushDoneEvent) -> None: + for wait in (lambda: list(session.speak("Next.")), session.terminate): + try: + wait() + except RuntimeError as exc: + refused.append(str(exc)) + + with aai.TTSSession(terminate_timeout=3) as session: + session.on(aai.TTSEvents.FlushDone, speak_next) + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." + + assert len(refused) == 2 + assert all("from a session handler" in message for message in refused) + assert _frames(server, "Flush") == [{"type": "Flush"}] # nothing more was sent + + def test_word_boundaries_and_cancelled_reach_their_handlers(server): words = [] cancelled = [] diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index f452cb2e..8a08085a 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -210,6 +210,53 @@ async def on_audio(event): assert received == [b"Pushed."] +async def test_the_error_reaches_waiters_before_the_error_handler(server): + app_saw_error = asyncio.Event() + handler_waited: List[bool] = [] + + async def on_error(error: aai.TTSError) -> None: + try: # waits on the app + await asyncio.wait_for(app_saw_error.wait(), 5) + handler_waited.append(True) + except asyncio.TimeoutError: + handler_waited.append(False) + + with pytest.raises(aai.TTSError): + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Error, on_error) + server.answer_flush = False + flush_id = await session.flush() + server.push(_INPUT_RATE_EXCEEDED) + try: + await _drain(session.audio(flush_id), timeout=3) + finally: + app_saw_error.set() + + assert handler_waited == [True] + + +async def test_waiting_from_a_handler_raises_instead_of_deadlocking(server): + refused: List[str] = [] + + async def speak_next(event: aai.TTSFlushDoneEvent) -> None: + try: + [chunk async for chunk in session.speak("Next.")] + except RuntimeError as exc: + refused.append(str(exc)) + try: + await session.terminate() + except RuntimeError as exc: + refused.append(str(exc)) + + async with aai.AsyncTTSSession(terminate_timeout=3) as session: + session.on(aai.TTSEvents.FlushDone, speak_next) + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + + assert len(refused) == 2 + assert all("from a session handler" in message for message in refused) + assert [f for f in server.sent if f["type"] == "Flush"] == [{"type": "Flush"}] + + async def test_a_lost_connection_raises_a_retryable_error(server): with pytest.raises(aai.TTSError) as exc_info: async with aai.AsyncTTSSession() as session: From 021091208b08249f81bfb33edcb1034926ef15d4 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:02:35 +0000 Subject: [PATCH 11/21] fix(tts): reconnect a session the API has ended After the API closed a session (an inactivity_timeout, an abnormal close), connect() kept raising "the session is already connected" until the caller called close(). It now closes the ended connection and opens a new one; a session that is still live is refused as before. A reader that outlived close() (an Audio handler blocked on a playback write, say) kept dispatching the old socket's frames into the new connection's state and finally stopped it. connect() now waits up to connect_timeout for the previous reader to finish, and refuses to run from a session handler. The README and the token docstring now say that a temporary token is spent when the session begins, so a token= session cannot reconnect: mint a new token and build a new session. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- README.md | 4 +- assemblyai/tts/v1/async_session.py | 13 +++++- assemblyai/tts/v1/session.py | 23 +++++++++- tests/unit/test_tts_session.py | 61 +++++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 68 ++++++++++++++++++++++++++++ 5 files changed, 165 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index db4d2359..cceb5852 100644 --- a/README.md +++ b/README.md @@ -1056,7 +1056,9 @@ Register handlers for the other frames with `session.on(event, handler)`, where A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status) and `request_id` is the session id; `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. -For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`, which answers `{"token": ..., "expires_in_seconds": 60}`). +Once the API has ended a session (an `inactivity_timeout`, say), the same `TTSSession` object opens a new one with `with session:` or `session.connect()`. + +For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`, which answers `{"token": ..., "expires_in_seconds": 60}`). The API spends the token when the session begins, so a token opens one session: to reconnect, mint a new token and build a new `TTSSession` with it. The asyncio twin is `aai.AsyncTTSSession`: diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index d83cc5f4..e077f1f4 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -116,8 +116,19 @@ def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: async def connect(self) -> TTSBeginEvent: """Opens the socket and waits for `Begin`; see `TTSSession.connect`.""" + self._check_not_in_handler("connect()") if self._ws is not None: - raise RuntimeError("the session is already connected") + if not self._stopped.is_set(): + raise RuntimeError("the session is already connected") + await self.close() # the previous connection has ended + reader = self._reader + if reader is not None and not reader.done(): + # Its reader would write into the new connection's state. + await asyncio.wait({reader}, timeout=self._connect_timeout) + if not reader.done(): + raise RuntimeError( + "a handler of the previous connection is still running" + ) self._begun = asyncio.Event() self._stopped = asyncio.Event() self._queue = PullBuffer(PULL_QUEUE_FRAMES) diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 68498928..ba7bfaf7 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -105,7 +105,9 @@ def __init__( api_key: The API key to authenticate with; takes precedence over `client` and the global settings. token: A temporary token instead of an API key, for a session - opened on a user's device. + opened on a user's device. The API spends the token when the + session begins, so such a session connects once: to + reconnect, mint a new token and build a new session. word_boundaries: Whether the server sends `WordBoundaries` frames. inactivity_timeout: Seconds (5 to 120) after which an idle session is closed; `None` for the account's default. @@ -173,12 +175,29 @@ def connect(self) -> TTSBeginEvent: """ Opens the socket and waits for `Begin`. + A session whose previous connection has ended (cleanly, by `close()` + or by the API, such as after an `inactivity_timeout`) connects again; + one built with `token=` cannot, since its token is already spent. + Raises: TTSError: if the API refuses the session (its `Error` frame), the handshake is rejected, or `Begin` does not arrive in time. + RuntimeError: if the session is still connected, from a session + handler, or while a handler of the previous connection still + runs after `connect_timeout`. """ + self._check_not_in_handler("connect()") if self._ws is not None: - raise RuntimeError("the session is already connected") + if not self._stopped.is_set(): + raise RuntimeError("the session is already connected") + self.close() # the previous connection has ended + if self._reader is not None and self._reader.is_alive(): + # Its reader would write into the new connection's state. + self._reader.join(self._connect_timeout) + if self._reader.is_alive(): + raise RuntimeError( + "a handler of the previous connection is still running" + ) self._reset() uri, headers = build_uri( diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index a701357e..6d60da8f 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -39,6 +39,15 @@ } +_INACTIVITY_TIMEOUT = { + "type": "Error", + "error_code": 3006, + "error": "No message received for 120 seconds; the idle session was closed.", + "error_type": "inactivity_timeout", + "retryable": True, + "param": None, + "retry_after_seconds": None, +} _INPUT_RATE_EXCEEDED = { "type": "Error", "error_code": 3010, @@ -487,6 +496,58 @@ def test_a_session_object_can_be_reused(mocker): assert [f["type"] for f in fakes[1].sent] == ["Generate", "Flush", "Terminate"] +def test_connect_after_the_api_ended_the_session_opens_a_new_one(mocker): + first, second = _FakeServer(), _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", side_effect=[first, second] + ) + session = aai.TTSSession() + session.connect() + first.push(_INACTIVITY_TIMEOUT) + first.push( + ConnectionClosedError(Close(3006, "See Error message for details"), None) + ) + assert session._stopped.wait(5) + assert session.error is not None and session.error.retryable is True + + session.connect() # no close() needed first + assert first.closed + assert session.error is None + assert b"".join(_drain(session.speak("Hello again."))) == b"Hello again." + session.terminate() + + +def test_connect_on_a_live_session_is_refused(server): + with aai.TTSSession() as session: + with pytest.raises(RuntimeError, match="already connected"): + session.connect() + + +def test_a_new_connection_waits_for_the_previous_connections_handler(mocker): + first, second = _FakeServer(), _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", side_effect=[first, second] + ) + in_handler = threading.Event() + + def blocking_player(event: aai.TTSAudioEvent) -> None: + in_handler.set() + time.sleep(1) # a playback write that blocks + + session = aai.TTSSession(terminate_timeout=0.2) + session.on(aai.TTSEvents.Audio, blocking_player) + with session: + session.generate("Hi.") + session.flush() + assert in_handler.wait(5) + # the previous reader is still in the handler here + with session: + time.sleep(1.5) # the previous handler has returned by now + session.generate("Still open.") + assert not session._stopped.is_set() + assert session.termination is None + + def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): from assemblyai.tts._base_ws import PULL_QUEUE_FRAMES diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 8a08085a..c43ef24b 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -31,6 +31,15 @@ } +_INACTIVITY_TIMEOUT = { + "type": "Error", + "error_code": 3006, + "error": "No message received for 120 seconds; the idle session was closed.", + "error_type": "inactivity_timeout", + "retryable": True, + "param": None, + "retry_after_seconds": None, +} _INPUT_RATE_EXCEEDED = { "type": "Error", "error_code": 3010, @@ -452,6 +461,65 @@ async def busy_handler(event): assert session._reader is not None and session._reader.done() +def _connect_to(mocker, *fakes: "_FakeAsyncServer") -> None: + queued = list(fakes) + + async def connect(uri, additional_headers): + return queued.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + + +async def test_connect_after_the_api_ended_the_session_opens_a_new_one(mocker): + first, second = _FakeAsyncServer(), _FakeAsyncServer() + _connect_to(mocker, first, second) + session = aai.AsyncTTSSession() + await session.connect() + first.push(_INACTIVITY_TIMEOUT) + first.push( + ConnectionClosedError(Close(3006, "See Error message for details"), None) + ) + await asyncio.wait_for(session._stopped.wait(), 5) + assert session.error is not None and session.error.retryable is True + + await session.connect() # no close() needed first + assert first.closed + assert session.error is None + assert b"".join(await _drain(session.speak("Hello again."))) == b"Hello again." + await session.terminate() + + +async def test_connect_on_a_live_session_is_refused(server): + async with aai.AsyncTTSSession() as session: + with pytest.raises(RuntimeError, match="already connected"): + await session.connect() + + +async def test_connect_right_after_close_gave_up_on_a_handler(mocker): + first, second = _FakeAsyncServer(), _FakeAsyncServer() + _connect_to(mocker, first, second) + in_handler = asyncio.Event() + + async def busy_handler(event): + in_handler.set() + await asyncio.sleep(30) + + session = aai.AsyncTTSSession(terminate_timeout=0.1) + await session.connect() + session.on(aai.TTSEvents.Cancelled, busy_handler) + await session.cancel() + await asyncio.wait_for(in_handler.wait(), 5) + await session.close() # gives up on the busy reader and cancels it + + await asyncio.wait_for(session.connect(), 5) # at once + await asyncio.sleep(0.05) + assert not session._stopped.is_set() + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + await session.terminate() + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 418feae5660d438048faf6743555c22a060a1253 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:02:50 +0000 Subject: [PATCH 12/21] docs(tts): describe the session idle timeout as the API applies it inactivity_timeout has no account default: every session closes after 120 s with no frame either way, and the argument can only shorten it. The README now also says that playing audio already received, or waiting for the user, counts as idle time for the API, so keep_alive() is needed during both. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- README.md | 2 +- assemblyai/tts/v1/session.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index cceb5852..d7f4ea33 100644 --- a/README.md +++ b/README.md @@ -1050,7 +1050,7 @@ with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: Audio starts before `flush()` once whole sentences have arrived. To play it while the text is still being sent, read `session.audio()` with no id from another thread (or task): it follows the request being written, and ends at that request's `FlushDone`. Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. Every `flush()` and every `cancel()` takes the next `flush_id`, so ids skip one after a cancel. -Limits: a `Generate` frame takes at most 2,000 characters (split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) without a frame; send `keep_alive()` between sparse requests. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. +Limits: a `Generate` frame takes at most 2,000 characters (split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) with no frame either way; send `keep_alive()` between sparse requests. Audio the SDK has already received does not hold the session open: the API counts idle time from the last frame it sent or received, so playing buffered audio, or waiting for the user, counts as idle. Send `keep_alive()` at least once per idle window during both. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. A handler may send (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but must not wait on the session: the reader that runs the handler is what `speak()`, `audio()` and `terminate()` wait for, so they raise `RuntimeError` there. Hand the work to another thread (or task) instead. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index ba7bfaf7..5392c848 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -109,8 +109,9 @@ def __init__( session begins, so such a session connects once: to reconnect, mint a new token and build a new session. word_boundaries: Whether the server sends `WordBoundaries` frames. - inactivity_timeout: Seconds (5 to 120) after which an idle session - is closed; `None` for the account's default. + inactivity_timeout: Seconds (5 to 120) with no frame either way + after which the API closes the session; `None` keeps the + API's 120 s. connect_timeout: Seconds to wait for the handshake and `Begin`. terminate_timeout: Seconds to wait for `Termination` on exit. @@ -327,7 +328,11 @@ def cancel(self) -> None: self._send({"type": "Cancel"}) def keep_alive(self) -> None: - """Resets the inactivity timer between sparse requests.""" + """ + Restarts the API's idle timer. Send one at least once per + `inactivity_timeout` while the session waits, which includes the + time already received audio takes to play. + """ self._send({"type": "KeepAlive"}) From d36dcbe9b5938f47fe52526fdf3c7761ea2a63c7 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:04:07 +0000 Subject: [PATCH 13/21] fix(tts): keep a failed session handshake's status and retryability The status of a rejected upgrade was read only from exc.response, which the legacy asyncio client of websockets 11 and 12 does not have (its InvalidStatusCode carries status_code), so a 429 or 5xx reported status_code None. Neither session set retryable for a handshake failure. Both sessions now build the error in one place: the status from either exception shape, retryable for the statuses the HTTP client retries (408, 429, 500, 502, 503, 504), and with no HTTP answer, retryable for a refused or dropped connection and a timeout. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/_base_ws.py | 30 ++++++++++++++++- assemblyai/tts/v1/async_session.py | 8 ++--- assemblyai/tts/v1/session.py | 6 ++-- tests/unit/test_tts_session.py | 44 ++++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 50 ++++++++++++++++++++++++++++ 5 files changed, 128 insertions(+), 10 deletions(-) diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 48c749a4..04b6b4d0 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -1,6 +1,8 @@ """Helpers shared by the sync and async Text-to-Speech sessions.""" +import asyncio import platform +import ssl from collections import deque from typing import Any, Deque, Dict, Optional, Tuple from urllib.parse import urlencode, urlsplit @@ -8,7 +10,8 @@ from .. import _audio from ..__version__ import __version__ from ..types import Settings -from .v1.models import TTSConfig +from .v1.api import _RETRYABLE_STATUSES +from .v1.models import TTSConfig, TTSError # The route of the WebSocket API; the host comes from `settings.tts_base_url`. WS_PATH = "/v1/ws" @@ -104,6 +107,31 @@ def resolve_settings( return settings +def handshake_error(exc: BaseException) -> TTSError: + """ + The `TTSError` for a WebSocket handshake that failed with `exc`. + + An HTTP answer is retryable by its status, as on the HTTP API; with no + answer, a refused or dropped connection and a timeout are retryable. + """ + + # The legacy client's InvalidStatusCode carries `status_code`; the + # current clients' InvalidStatus carries the `response`. + status = getattr(exc, "status_code", None) + if status is None: + status = getattr(getattr(exc, "response", None), "status_code", None) + if isinstance(status, int): + retryable = status in _RETRYABLE_STATUSES + else: + status = None + retryable = isinstance(exc, (OSError, asyncio.TimeoutError)) and not isinstance( + exc, ssl.CertificateError + ) + return TTSError( + f"WebSocket handshake failed: {exc}", status_code=status, retryable=retryable + ) + + def _user_agent() -> str: return ( f"AssemblyAI/1.0 (sdk=Python/{__version__} " diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index e077f1f4..44ba8896 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -11,6 +11,7 @@ PullItem, build_uri, check_text_frame, + handshake_error, resolve_settings, ) from .models import TTSConfig, TTSError @@ -153,11 +154,8 @@ async def connect(self) -> TTSBeginEvent: self._ws = await asyncio.wait_for( websocket_connect_async(uri, headers), timeout=self._connect_timeout ) - except Exception as exc: - status = getattr(getattr(exc, "response", None), "status_code", None) - raise TTSError( - f"WebSocket handshake failed: {exc}", status_code=status - ) from exc + except Exception as exc: # the handshake failed + raise handshake_error(exc) from exc self._reader = asyncio.create_task(self._read_loop()) try: diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 5392c848..04e8d983 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -11,6 +11,7 @@ PullItem, build_uri, check_text_frame, + handshake_error, resolve_settings, ) from .models import TTSConfig, TTSError @@ -214,10 +215,7 @@ def connect(self) -> TTSBeginEvent: uri, additional_headers=headers, open_timeout=self._connect_timeout ) except Exception as exc: # the handshake failed - status = getattr(getattr(exc, "response", None), "status_code", None) - raise TTSError( - f"WebSocket handshake failed: {exc}", status_code=status - ) from exc + raise handshake_error(exc) from exc self._reader = threading.Thread( target=self._read_loop, name="assemblyai-tts-session", daemon=True diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 6d60da8f..339a7775 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -10,6 +10,7 @@ import queue import threading import time +from types import SimpleNamespace from typing import Any, Iterable, List, Optional from urllib.parse import parse_qs, urlsplit @@ -59,6 +60,36 @@ } +class _Rejected(Exception): + """A handshake answered with an HTTP status, as the current clients raise.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = SimpleNamespace(status_code=status) + + +class _LegacyRejected(Exception): + """The same, as the legacy client of websockets 11 and 12 raises it.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.status_code = status + + +_HANDSHAKE_FAILURES = [ + pytest.param(_Rejected(503), 503, True, id="503"), + pytest.param(_Rejected(502), 502, True, id="502"), + pytest.param(_Rejected(401), 401, False, id="401"), + pytest.param(_LegacyRejected(429), 429, True, id="legacy-429"), + pytest.param(_LegacyRejected(403), 403, False, id="legacy-403"), + pytest.param( + ConnectionRefusedError(111, "Connection refused"), None, True, id="refused" + ), + pytest.param(TimeoutError("timed out"), None, True, id="timeout"), + pytest.param(ValueError("not a WebSocket URI"), None, False, id="bad-uri"), +] + + def _closed_ok() -> ConnectionClosedOK: return ConnectionClosedOK(Close(1000, "Session ended"), None) @@ -343,6 +374,19 @@ def test_an_error_before_begin_raises_from_with(mocker): assert fake.closed +@pytest.mark.parametrize("failure, status, retryable", _HANDSHAKE_FAILURES) +def test_a_failed_handshake_keeps_its_status_and_says_if_retryable( + mocker, failure, status, retryable +): + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=failure) + with pytest.raises(aai.TTSError, match="handshake failed") as exc_info: + with aai.TTSSession(): + pass + assert exc_info.value.status_code == status + assert exc_info.value.retryable is retryable + assert exc_info.value.__cause__ is failure + + def test_an_error_mid_session_raises_from_the_audio_and_reaches_the_handler(server): errors = [] with pytest.raises(aai.TTSError) as exc_info: diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index c43ef24b..a0cc1a31 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -3,6 +3,7 @@ import asyncio import base64 import json +from types import SimpleNamespace from typing import Any, AsyncIterable, List, Optional from urllib.parse import parse_qs, urlsplit @@ -31,6 +32,37 @@ } +class _Rejected(Exception): + """A handshake answered with an HTTP status, as the current clients raise.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = SimpleNamespace(status_code=status) + + +class _LegacyRejected(Exception): + """The same, as the legacy client of websockets 11 and 12 raises it.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.status_code = status + + +_HANDSHAKE_FAILURES = [ + pytest.param(_Rejected(503), 503, True, id="503"), + pytest.param(_Rejected(502), 502, True, id="502"), + pytest.param(_Rejected(401), 401, False, id="401"), + pytest.param(_LegacyRejected(429), 429, True, id="legacy-429"), + pytest.param(_LegacyRejected(403), 403, False, id="legacy-403"), + pytest.param( + ConnectionRefusedError(111, "Connection refused"), None, True, id="refused" + ), + pytest.param(TimeoutError("timed out"), None, True, id="timeout"), + pytest.param(asyncio.TimeoutError(), None, True, id="asyncio-timeout"), + pytest.param(ValueError("not a WebSocket URI"), None, False, id="bad-uri"), +] + + _INACTIVITY_TIMEOUT = { "type": "Error", "error_code": 3006, @@ -202,6 +234,24 @@ async def connect(uri, additional_headers): assert fake.closed +@pytest.mark.parametrize("failure, status, retryable", _HANDSHAKE_FAILURES) +async def test_a_failed_handshake_keeps_its_status_and_says_if_retryable( + mocker, failure, status, retryable +): + async def connect(uri, additional_headers): + raise failure + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + with pytest.raises(aai.TTSError, match="handshake failed") as exc_info: + async with aai.AsyncTTSSession(): + pass + assert exc_info.value.status_code == status + assert exc_info.value.retryable is retryable + assert exc_info.value.__cause__ is failure + + async def test_handlers_may_be_coroutines(server): received = [] From 8da63b12bc6a547c1e27010053e37e2dd6fde484 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:04:44 +0000 Subject: [PATCH 14/21] fix(tts): one key for a session error's close code An Error frame's TTSError kept the close code as body["error_code"] (the body is the frame), and a close without a frame as body["close_code"], so retry logic keyed on close codes had to check both. The frame's error now also carries body["close_code"]. The TTSError docstring, CLAUDE.md and the README now describe session errors: no HTTP status, the close code in body["close_code"], retryable from the API's Error frame, and biometric_terms_not_accepted and auth_unavailable arriving as closes 3006 and 1011. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- CLAUDE.md | 2 +- README.md | 4 ++-- assemblyai/tts/v1/models.py | 15 +++++++++++---- assemblyai/tts/v1/ws_models.py | 10 ++++++++-- tests/unit/test_tts_session.py | 1 + tests/unit/test_tts_session_async.py | 1 + 6 files changed, 24 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index d2a0bc08..6ab78f4c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -58,7 +58,7 @@ aai.settings.api_key = "your-key" - `aai.AsyncTTSSession` — Asyncio counterpart: `async with`, coroutine sends, `async for chunk in session.speak(text)` - `aai.TTSConfig` — Text-to-Speech options: `voice_id` (a library voice such as `aai-jane`, or a cloned voice by UUID or slug), `model_id` (`assemblyai-tts-1`), `language`, `sample_rate`, `encoding` (`aai.TTSEncoding`: `pcm_s16le` default | `pcm_mulaw` | `pcm_alaw` | `mp3`), `container` (`aai.TTSContainer`: `raw` default | `wav`), `bit_rate`, `extend_biometric_consent` (owner-only consent renewal of a cloned voice). Rejects unknown fields - `aai.TTSResponse` — Text-to-Speech result: `.audio` (bytes) plus the response headers: `.request_id`, `.region`, `.model`, `.voice_id`, `.language`, `.encoding`, `.sample_rate`, `.container`, `.bit_rate`, `.content_type`, `.voice_expires_at`, `.billed_characters`, `.audio_duration_ms`. `TTSStream.info` carries the same fields without `audio`; `audio_duration_ms` is `None` on a stream with audio (`0` when the text had nothing to speak) -- `aai.TTSError` — Raised on any Text-to-Speech refusal: `.status_code`, `.error_code` (stable API code; `biometric_terms_not_accepted` 403 until the account accepts the biometric data terms in the dashboard, `auth_unavailable` 503 retryable), `.param`, `.request_id`, `.retryable`, `.retry_after`, `.errors`, `.body` (the whole envelope) +- `aai.TTSError` — Raised on any Text-to-Speech refusal: `.status_code`, `.error_code` (stable API code; `biometric_terms_not_accepted` 403 until the account accepts the biometric data terms in the dashboard, `auth_unavailable` 503 retryable; on a session these arrive as close 3006 and 1011, with `status_code` None, the close code in `.body["close_code"]` and `.retryable` from the API's Error frame), `.param`, `.request_id`, `.retryable`, `.retry_after`, `.errors`, `.body` (the whole envelope) - `assemblyai.streaming.v3.RealTimeTranscriber` — Real-time streaming with event-based API (threaded) - `assemblyai.streaming.v3.AsyncRealTimeTranscriber` — Asyncio-native counterpart; same options/events - `aai.LLMGateway` — LLM Gateway client. Resources: `models.list()`, `chat.completions.create()` (incl. `stream=True`), `understanding.create()`/`.validate()` diff --git a/README.md b/README.md index d7f4ea33..4975df9b 100644 --- a/README.md +++ b/README.md @@ -1054,7 +1054,7 @@ Limits: a `Generate` frame takes at most 2,000 characters (split longer text acr Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. A handler may send (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but must not wait on the session: the reader that runs the handler is what `speak()`, `audio()` and `terminate()` wait for, so they raise `RuntimeError` there. Hand the work to another thread (or task) instead. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. -A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status) and `request_id` is the session id; `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. +A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status), `request_id` is the session id and `body["close_code"]` is the WebSocket close code (`biometric_terms_not_accepted` is a `3006`, `auth_unavailable` a retryable `1011`); `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. Once the API has ended a session (an `inactivity_timeout`, say), the same `TTSSession` object opens a new one with `with session:` or `session.connect()`. @@ -2164,7 +2164,7 @@ All custom SDK exceptions inherit from `aai.AssemblyAIError`, except the streami | `aai.AssemblyAIError` | Base exception for all AssemblyAI errors. | `status_code: Optional[int]` | | `aai.TranscriptError` | Raised when a prerecorded transcript API call fails at the HTTP level or polling times out. | `status_code: Optional[int]` | | `aai.SyncTranscriptError` | Raised when synchronous transcription fails (`aai.SyncTranscriber`). | `status_code: Optional[int]`, `error_code: Optional[str]`, `retry_after: Optional[int]` | -| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`) or a session is refused or fails (`aai.TTSSession`, where `status_code` is `None` and `request_id` is the session id). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | +| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`) or a session is refused or fails (`aai.TTSSession`, where `status_code` is `None`, `request_id` is the session id and `body["close_code"]` is the WebSocket close code). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | | `assemblyai.types.RedactedAudioIncompleteError` | Raised when requesting redacted audio before processing completes. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioExpiredError` | Raised when requesting redacted audio that has expired and is no longer available. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioUnavailableError` | Raised when requesting redacted audio that is unavailable at the given URL. | `status_code: Optional[int]` | diff --git a/assemblyai/tts/v1/models.py b/assemblyai/tts/v1/models.py index bdebdbd4..d56afa59 100644 --- a/assemblyai/tts/v1/models.py +++ b/assemblyai/tts/v1/models.py @@ -50,10 +50,17 @@ class TTSError(AssemblyAIError): body over its limit adds `limit`). `biometric_terms_not_accepted` (`403`, not retryable) refuses one of the - account's own cloned voices, or `extend_biometric_consent`, until the - account accepts the biometric data terms in the AssemblyAI dashboard. - `auth_unavailable` (`503`, retryable) means the credential, or those - terms, could not be checked just now. + account's own cloned voices, an inline voice clip or + `extend_biometric_consent`, until the account accepts the biometric data + terms in the AssemblyAI dashboard. `auth_unavailable` (`503`, retryable) + means the credential, or those terms, could not be checked just now. + + On a WebSocket session (`TTSSession`) there is no HTTP status: + `status_code` is `None`, `request_id` is the session id, `retryable` is + the API's (from its `Error` frame, else from the close code), and + `body["close_code"]` is the WebSocket close code. There + `biometric_terms_not_accepted` arrives as close `3006` and + `auth_unavailable` as close `1011`. A streamed response cut after its first audio raises a `TTSError` with `status_code` `200`, no `error_code`, `retryable` `True` and the diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py index 71e382b8..63d92b95 100644 --- a/assemblyai/tts/v1/ws_models.py +++ b/assemblyai/tts/v1/ws_models.py @@ -169,16 +169,22 @@ def parse_event(data: Dict[str, Any]) -> Any: def error_from_frame(data: Dict[str, Any]) -> TTSError: - """The `TTSError` for an `Error` frame; `body` keeps the frame.""" + """ + The `TTSError` for an `Error` frame. `body` keeps the frame, plus + `close_code`: the code of the close that follows it (the frame's + `error_code`), under the same key as for a close without a frame. + """ retry_after = data.get("retry_after_seconds") + body = dict(data) + body.setdefault("close_code", data.get("error_code")) return TTSError( data.get("error") or "text-to-speech session failed", error_code=data.get("error_type"), param=data.get("param"), retryable=data.get("retryable"), retry_after=int(retry_after) if isinstance(retry_after, (int, float)) else None, - body=data, + body=body, ) diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 339a7775..b6701b04 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -370,6 +370,7 @@ def test_an_error_before_begin_raises_from_with(mocker): assert err.param == "voice_id" assert err.retryable is False assert err.body["error_code"] == 3012 + assert err.body["close_code"] == 3012 # the key a bare close uses too assert err.request_id is None # no Begin, so no session id yet assert fake.closed diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index a0cc1a31..d007ee6d 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -231,6 +231,7 @@ async def connect(uri, additional_headers): pass assert exc_info.value.error_code == "rate_limited" assert exc_info.value.retry_after == 17 + assert exc_info.value.body["close_code"] == 3009 assert fake.closed From 9a19c690cb326128888fe9d12286b2781d07e35e Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:05:41 +0000 Subject: [PATCH 15/21] fix(tts): check Generate text against the limit Begin reports The 2,000-character limit of one Generate frame was hard-coded, while the API reports the limit in force in Begin.configuration.limits.max_generate_text_length and says it may change. A raised limit made the SDK refuse valid text; a lowered one turned the SDK's clear ValueError into a 3006 close. generate() now checks against the reported limit (0 means none) and falls back to 2,000. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- CLAUDE.md | 2 +- README.md | 2 +- assemblyai/tts/_base_ws.py | 26 ++++++++++++++++++++++---- assemblyai/tts/v1/async_session.py | 3 ++- assemblyai/tts/v1/session.py | 6 ++++-- tests/unit/test_tts_session.py | 15 +++++++++++++++ tests/unit/test_tts_session_async.py | 14 ++++++++++++++ 7 files changed, 59 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 6ab78f4c..db0f58e5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -54,7 +54,7 @@ aai.settings.api_key = "your-key" - `aai.DictationResponse` — Dictation result: `.text`, `.words` (`DictationWord` with `text` + `confidence`), `.confidence`, `.llm_response`, `.llm_error`, `.audio_duration_ms`, `.session_id`, `.request_time_ms`, `.sync_time_ms`, and `.final_text` (the LLM rewrite, falling back to `.text`) - `aai.TextToSpeech` — Text-to-Speech HTTP API: one text per request, audio out. Methods: `synthesize()` (complete audio as a `TTSResponse`), `stream()` (a `TTSStream` context manager yielding bytes as the API generates them), `close()` / `with` for a client built from `api_key=`. Host from `aai.settings.tts_base_url` (`https://tts.us.assemblyai.com`; set `https://tts.eu.assemblyai.com` for the EU zone) - `aai.AsyncTextToSpeech` — Asyncio counterpart of `TextToSpeech`; `synthesize()` is a coroutine and `stream()` returns an `AsyncTTSStream`. Owns an HTTP pool: use `async with` or `await aclose()` -- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin` (`session.begin`); `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining (`audio()` with no id follows the request being written, from another thread); `cancel()` for barge-in (ends the iteration at once); `keep_alive()`; `on(TTSEvents.X, handler(event))` for `Begin`, `Audio` (switches audio from pull to push), `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination` (`None` after a failure). Refusals and failures raise `TTSError` (`status_code` None, `request_id` = session id). Raw PCM/G.711 only; 2,000 chars per generate; ~100 chars/s sustained. `token=` for temporary tokens +- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin` (`session.begin`); `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining (`audio()` with no id follows the request being written, from another thread); `cancel()` for barge-in (ends the iteration at once); `keep_alive()`; `on(TTSEvents.X, handler(event))` for `Begin`, `Audio` (switches audio from pull to push), `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination` (`None` after a failure). Refusals and failures raise `TTSError` (`status_code` None, `request_id` = session id). Raw PCM/G.711 only; 2,000 chars per generate unless `Begin` reports another `limits.max_generate_text_length`; ~100 chars/s sustained. `token=` for temporary tokens - `aai.AsyncTTSSession` — Asyncio counterpart: `async with`, coroutine sends, `async for chunk in session.speak(text)` - `aai.TTSConfig` — Text-to-Speech options: `voice_id` (a library voice such as `aai-jane`, or a cloned voice by UUID or slug), `model_id` (`assemblyai-tts-1`), `language`, `sample_rate`, `encoding` (`aai.TTSEncoding`: `pcm_s16le` default | `pcm_mulaw` | `pcm_alaw` | `mp3`), `container` (`aai.TTSContainer`: `raw` default | `wav`), `bit_rate`, `extend_biometric_consent` (owner-only consent renewal of a cloned voice). Rejects unknown fields - `aai.TTSResponse` — Text-to-Speech result: `.audio` (bytes) plus the response headers: `.request_id`, `.region`, `.model`, `.voice_id`, `.language`, `.encoding`, `.sample_rate`, `.container`, `.bit_rate`, `.content_type`, `.voice_expires_at`, `.billed_characters`, `.audio_duration_ms`. `TTSStream.info` carries the same fields without `audio`; `audio_duration_ms` is `None` on a stream with audio (`0` when the text had nothing to speak) diff --git a/README.md b/README.md index 4975df9b..528dbbe3 100644 --- a/README.md +++ b/README.md @@ -1050,7 +1050,7 @@ with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: Audio starts before `flush()` once whole sentences have arrived. To play it while the text is still being sent, read `session.audio()` with no id from another thread (or task): it follows the request being written, and ends at that request's `FlushDone`. Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. Every `flush()` and every `cancel()` takes the next `flush_id`, so ids skip one after a cancel. -Limits: a `Generate` frame takes at most 2,000 characters (split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) with no frame either way; send `keep_alive()` between sparse requests. Audio the SDK has already received does not hold the session open: the API counts idle time from the last frame it sent or received, so playing buffered audio, or waiting for the user, counts as idle. Send `keep_alive()` at least once per idle window during both. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. +Limits: a `Generate` frame takes at most 2,000 characters by default (`session.begin.configuration["limits"]["max_generate_text_length"]` has the limit in force; split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) with no frame either way; send `keep_alive()` between sparse requests. Audio the SDK has already received does not hold the session open: the API counts idle time from the last frame it sent or received, so playing buffered audio, or waiting for the user, counts as idle. Send `keep_alive()` at least once per idle window during both. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. A handler may send (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but must not wait on the session: the reader that runs the handler is what `speak()`, `audio()` and `terminate()` wait for, so they raise `RuntimeError` there. Hand the work to another thread (or task) instead. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 04b6b4d0..149d243d 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -15,7 +15,8 @@ # The route of the WebSocket API; the host comes from `settings.tts_base_url`. WS_PATH = "/v1/ws" -# Characters a single `Generate` frame may carry (Unicode code points). +# Characters a single `Generate` frame may carry (Unicode code points) when +# `Begin` does not report the limit in force. MAX_GENERATE_CHARS = 2000 # Audio frames queued for pulling before the oldest is dropped: about 160 s # at the API's ~80 ms per frame. Bounds memory when audio is not read, or is @@ -71,12 +72,29 @@ def popleft(self) -> Optional[PullItem]: return item -def check_text_frame(text: object) -> None: +def max_generate_chars(begin: Any) -> int: + """ + The characters one `Generate` may carry: the limit the session's + `Begin` reports (`limits.max_generate_text_length`, where `0` means + none), else `MAX_GENERATE_CHARS`. + """ + + configuration = getattr(begin, "configuration", None) or {} + limits = configuration.get("limits") + limit = limits.get("max_generate_text_length") if isinstance(limits, dict) else None + if isinstance(limit, int) and not isinstance(limit, bool) and limit >= 0: + return limit + return MAX_GENERATE_CHARS + + +def check_text_frame(text: object, limit: int = MAX_GENERATE_CHARS) -> None: + """Refuses what one `Generate` cannot carry; a `limit` of `0` is none.""" + if not isinstance(text, str): raise TypeError(f"text must be a str, not {type(text).__name__}") - if len(text) > MAX_GENERATE_CHARS: + if limit and len(text) > limit: raise ValueError( - f"text is {len(text)} characters, over the {MAX_GENERATE_CHARS}-character " + f"text is {len(text)} characters, over the {limit}-character " "limit of one Generate frame; split it across several generate() calls" ) diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index 44ba8896..5741af50 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -12,6 +12,7 @@ build_uri, check_text_frame, handshake_error, + max_generate_chars, resolve_settings, ) from .models import TTSConfig, TTSError @@ -235,7 +236,7 @@ async def __aexit__(self, exc_type: Any, *exc: object) -> None: async def generate(self, text: str) -> None: """Appends `text` to the current request; see `TTSSession.generate`.""" - check_text_frame(text) + check_text_frame(text, max_generate_chars(self.begin)) await self._send({"type": "Generate", "text": text}) async def flush(self) -> int: diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 04e8d983..56e7c11b 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -12,6 +12,7 @@ build_uri, check_text_frame, handshake_error, + max_generate_chars, resolve_settings, ) from .models import TTSConfig, TTSError @@ -291,9 +292,10 @@ def generate(self, text: str) -> None: arrived; `flush()` is what concludes the request. Raises: - ValueError: for a text over 2,000 characters (split it). + ValueError: for a text over the limit of one `Generate` frame + (split it): the limit `Begin` reports, else 2,000 characters. """ - check_text_frame(text) + check_text_frame(text, max_generate_chars(self.begin)) self._send({"type": "Generate", "text": text}) def flush(self) -> int: diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index b6701b04..9d21ccd3 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -743,6 +743,21 @@ def test_generate_refuses_text_over_one_frame(server): session.generate(b"bytes") # type: ignore[arg-type] +@pytest.mark.parametrize( + "limit, allowed, refused", [(4000, 4000, 4001), (500, 500, 501)] +) +def test_generate_follows_the_limit_begin_reports(mocker, limit, allowed, refused): + begin = dict(_BEGIN, configuration={"limits": {"max_generate_text_length": limit}}) + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + return_value=_FakeServer(begin=begin), + ) + with aai.TTSSession() as session: + session.generate("x" * allowed) + with pytest.raises(ValueError, match=f"{limit}-character"): + session.generate("x" * refused) + + def test_the_eu_host_comes_from_the_settings(server): client = aai.Client(api_key="test") client.settings.tts_base_url = "https://tts.eu.assemblyai.com" diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index d007ee6d..79ab88e6 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -571,6 +571,20 @@ async def busy_handler(event): await session.terminate() +@pytest.mark.parametrize( + "limit, allowed, refused", [(4000, 4000, 4001), (500, 500, 501)] +) +async def test_generate_follows_the_limit_begin_reports( + mocker, limit, allowed, refused +): + begin = dict(_BEGIN, configuration={"limits": {"max_generate_text_length": limit}}) + _connect_to(mocker, _FakeAsyncServer(begin=begin)) + async with aai.AsyncTTSSession() as session: + await session.generate("x" * allowed) + with pytest.raises(ValueError, match=f"{limit}-character"): + await session.generate("x" * refused) + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 27b13b34f6d2747a911b53729418f2214ff11e42 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:06:09 +0000 Subject: [PATCH 16/21] fix(tts): open a session whose Begin carries no configuration TTSBeginEvent.configuration was required, while the API declares it optional (default null). A Begin without it, or with null, ended the session as an unreadable frame instead of opening it. The field now defaults to an empty dict, and a null is read as one. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/ws_models.py | 9 ++++++--- tests/unit/test_tts_session.py | 14 ++++++++++++++ tests/unit/test_tts_session_async.py | 11 +++++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py index 63d92b95..27c3ebc0 100644 --- a/assemblyai/tts/v1/ws_models.py +++ b/assemblyai/tts/v1/ws_models.py @@ -81,9 +81,10 @@ class TTSBeginEvent(BaseModel): expires_at: int "Unix time in seconds at which the session's maximum duration ends." - configuration: Dict[str, Any] + configuration: Dict[str, Any] = Field(default_factory=dict) """The session as created: the effective connect parameters, `model`, - `voice_id`, `session_ttl_sec`, `limits` and more. An open set.""" + `voice_id`, `session_ttl_sec`, `limits` and more. An open set; empty + when the API sends none.""" class TTSAudioEvent(BaseModel): @@ -156,7 +157,9 @@ def parse_event(data: Dict[str, Any]) -> Any: flush_id=data["flush_id"], audio=base64.b64decode(data["audio"]) ) if kind == "Begin": - return _parse(TTSBeginEvent, data) + return _parse( + TTSBeginEvent, {**data, "configuration": data.get("configuration") or {}} + ) if kind == "FlushDone": return _parse(TTSFlushDoneEvent, data) if kind == "WordBoundaries": diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 9d21ccd3..5ebd9680 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -786,6 +786,20 @@ def test_api_key_argument_is_used_without_touching_the_shared_settings(server): assert aai.settings.api_key == "test" +@pytest.mark.parametrize("configuration", [None, "missing"]) +def test_a_begin_without_configuration_opens_the_session(mocker, configuration): + begin = {k: v for k, v in _BEGIN.items() if k != "configuration"} + if configuration != "missing": + begin["configuration"] = configuration + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + return_value=_FakeServer(begin=begin), + ) + with aai.TTSSession() as session: + assert session.begin is not None and session.begin.configuration == {} + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." + + def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 79ab88e6..626b159e 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -585,6 +585,17 @@ async def test_generate_follows_the_limit_begin_reports( await session.generate("x" * refused) +@pytest.mark.parametrize("configuration", [None, "missing"]) +async def test_a_begin_without_configuration_opens_the_session(mocker, configuration): + begin = {k: v for k, v in _BEGIN.items() if k != "configuration"} + if configuration != "missing": + begin["configuration"] = configuration + _connect_to(mocker, _FakeAsyncServer(begin=begin)) + async with aai.AsyncTTSSession() as session: + assert session.begin is not None and session.begin.configuration == {} + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 6c8320497647e5007fc34bb5c4ca8577571e5ece Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:06:31 +0000 Subject: [PATCH 17/21] fix(tts): raise from TTSSession.audio() before the first connect audio() on a TTSSession that never connected polled forever, since its queue exists from __init__ and nothing would ever stop it. It now raises RuntimeError('the session is not open'), as AsyncTTSSession.audio() already did; the async test pins that. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/v1/session.py | 4 +++- tests/unit/test_tts_session.py | 5 +++++ tests/unit/test_tts_session_async.py | 5 +++++ 3 files changed, 13 insertions(+), 1 deletion(-) diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 56e7c11b..17f985e1 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -365,9 +365,11 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: TTSError: if the session fails before the request completes. ValueError: for a `flush_id` no `flush()` has returned. RuntimeError: when an `Audio` handler receives the audio instead, - or from a session handler. + before the session first connects, or from a session handler. """ self._check_not_in_handler("audio()") + if self._reader is None: + raise RuntimeError("the session is not open") if flush_id is None: flush_id = self._next_flush_id # the request being written now elif flush_id > self._next_flush_id: diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 5ebd9680..e4b78242 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -333,6 +333,11 @@ def test_sending_before_connect_says_how_to_open(server): aai.TTSSession().generate("Hi.") +def test_audio_before_connect_says_the_session_is_not_open(server): + with pytest.raises(RuntimeError, match="not open"): + _drain(aai.TTSSession().audio()) + + def test_http_only_options_are_refused_on_a_session(server): with pytest.raises(ValueError, match="container and bit_rate"): with aai.TTSSession(config=aai.TTSConfig(container="wav")): diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 626b159e..a344f402 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -596,6 +596,11 @@ async def test_a_begin_without_configuration_opens_the_session(mocker, configura assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." +async def test_audio_before_connect_says_the_session_is_not_open(server): + with pytest.raises(RuntimeError, match="not open"): + await _drain(aai.AsyncTTSSession().audio()) + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 3775aeb29ae3233dc1a015248a2b3b7e51a77749 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:06:54 +0000 Subject: [PATCH 18/21] fix(tts): keep the path prefix of tts_base_url for a session build_uri dropped any path in tts_base_url, so https://proxy.example.com/tts became wss://proxy.example.com/v1/ws while the HTTP client kept the prefix. A deployment behind a path-prefixed proxy then worked over HTTP and failed over WebSocket. The prefix now goes before /v1/ws. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- assemblyai/tts/_base_ws.py | 6 ++++-- tests/unit/test_tts_session.py | 16 ++++++++++++++++ tests/unit/test_tts_session_async.py | 7 +++++++ 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 149d243d..e58224a6 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -13,7 +13,8 @@ from .v1.api import _RETRYABLE_STATUSES from .v1.models import TTSConfig, TTSError -# The route of the WebSocket API; the host comes from `settings.tts_base_url`. +# The route of the WebSocket API, after the host and any path prefix of +# `settings.tts_base_url`. WS_PATH = "/v1/ws" # Characters a single `Generate` frame may carry (Unicode code points) when # `Begin` does not report the limit in force. @@ -190,4 +191,5 @@ def build_uri( headers: Dict[str, Optional[str]] = {"User-Agent": _user_agent()} if token is None and api_key: headers["Authorization"] = api_key - return f"{scheme}://{parts.netloc}{WS_PATH}?{query}", headers + prefix = parts.path.rstrip("/") + return f"{scheme}://{parts.netloc}{prefix}{WS_PATH}?{query}", headers diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index e4b78242..972478df 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -779,6 +779,22 @@ def test_a_plain_http_host_opens_a_plain_websocket(server): assert server.connect.call_args.args[0].startswith("ws://localhost:8080/v1/ws?") +@pytest.mark.parametrize( + "base_url, expected", + [ + ("https://proxy.example.com/tts", "wss://proxy.example.com/tts/v1/ws?"), + ("https://proxy.example.com/tts/", "wss://proxy.example.com/tts/v1/ws?"), + ("https://tts.us.assemblyai.com/", "wss://tts.us.assemblyai.com/v1/ws?"), + ], +) +def test_a_path_prefix_in_the_host_setting_is_kept(server, base_url, expected): + client = aai.Client(api_key="test") + client.settings.tts_base_url = base_url + with aai.TTSSession(client=client): + pass + assert server.connect.call_args.args[0].startswith(expected) + + def test_api_key_argument_is_used_without_touching_the_shared_settings(server): with aai.TTSSession( api_key="other", config=aai.TTSConfig(extend_biometric_consent=True) diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index a344f402..6cb3cdbc 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -601,6 +601,13 @@ async def test_audio_before_connect_says_the_session_is_not_open(server): await _drain(aai.AsyncTTSSession().audio()) +async def test_a_path_prefix_in_the_host_setting_is_kept(server, monkeypatch): + monkeypatch.setattr(aai.settings, "tts_base_url", "https://proxy.example.com/tts/") + async with aai.AsyncTTSSession(): + pass + assert server.connect.uri.startswith("wss://proxy.example.com/tts/v1/ws?") + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From c6fbbbb2959f72c60a60be3e4c4d77358624f8e4 Mon Sep 17 00:00:00 2001 From: Aleksandar Mitov <140423361+aleks-mitov@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:07:29 +0000 Subject: [PATCH 19/21] fix(tts): accept container=raw on a session and check at construction The session refused every container value, including raw, which is what it streams, so a TTSConfig shared with the HTTP client that set raw explicitly failed. The refusal also surfaced only at connect. Only a wav container and bit_rate are refused now, and both sessions check when they are built. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KcyY6MAEUVjqBt5pMpetvD --- CLAUDE.md | 2 +- README.md | 2 +- assemblyai/tts/_base_ws.py | 18 ++++++++++++------ assemblyai/tts/v1/async_session.py | 2 ++ assemblyai/tts/v1/session.py | 8 ++++++-- tests/unit/test_tts_session.py | 22 ++++++++++++++++++---- tests/unit/test_tts_session_async.py | 16 ++++++++++++++++ 7 files changed, 56 insertions(+), 14 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index db0f58e5..ede6f2ae 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -487,7 +487,7 @@ with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: for sentence in llm_reply_sentences(): for chunk in session.speak(sentence): play(chunk) -# session.termination is the summary (None if the session failed); raw PCM/G.711 only, no mp3/container +# session.termination is the summary (None if the session failed); raw PCM/G.711 only, no mp3, wav or bit_rate ``` ## LLM Gateway diff --git a/README.md b/README.md index 528dbbe3..d7d32444 100644 --- a/README.md +++ b/README.md @@ -1018,7 +1018,7 @@ asyncio.run(main()) `aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary (`None` if the session failed instead). -A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, `container` and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`. +A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, a `wav` container and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`. ```python import assemblyai as aai diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index e58224a6..a5e9894b 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -11,7 +11,7 @@ from ..__version__ import __version__ from ..types import Settings from .v1.api import _RETRYABLE_STATUSES -from .v1.models import TTSConfig, TTSError +from .v1.models import TTSConfig, TTSContainer, TTSError # The route of the WebSocket API, after the host and any path prefix of # `settings.tts_base_url`. @@ -158,6 +158,16 @@ def _user_agent() -> str: ) +def check_session_config(config: TTSConfig) -> None: + """Refuses the options a session cannot honour: it streams raw audio.""" + + if config.container not in (None, TTSContainer.raw) or config.bit_rate is not None: + raise ValueError( + "a wav container and bit_rate do not apply to a session, which streams " + "raw pcm_s16le, pcm_mulaw or pcm_alaw; leave them unset in its TTSConfig" + ) + + def build_uri( base_url: str, config: TTSConfig, @@ -169,11 +179,7 @@ def build_uri( ) -> Tuple[str, Dict[str, Optional[str]]]: """The session URI and handshake headers from the settings and config.""" - if config.container is not None or config.bit_rate is not None: - raise ValueError( - "container and bit_rate do not apply to a session, which streams raw " - "pcm_s16le, pcm_mulaw or pcm_alaw; leave them unset in its TTSConfig" - ) + check_session_config(config) parts = urlsplit(base_url if "://" in base_url else f"https://{base_url}") scheme = "ws" if parts.scheme in ("http", "ws") else "wss" params: Dict[str, Any] = ( diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index 5741af50..13951981 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -10,6 +10,7 @@ PullBuffer, PullItem, build_uri, + check_session_config, check_text_frame, handshake_error, max_generate_chars, @@ -80,6 +81,7 @@ def __init__( _audio.check_config(type(self).__name__, config, TTSConfig) self.config = config or TTSConfig() + check_session_config(self.config) self._settings = resolve_settings( type(self).__name__, client.settings if client else None, api_key, token ) diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 17f985e1..89970ee2 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -10,6 +10,7 @@ PullBuffer, PullItem, build_uri, + check_session_config, check_text_frame, handshake_error, max_generate_chars, @@ -100,8 +101,9 @@ def __init__( Args: config: The voice and audio options (`voice_id`, `model_id`, - `language`, `sample_rate`, `encoding`). `container` and - `bit_rate` do not apply: a session streams raw PCM or G.711. + `language`, `sample_rate`, `encoding`). A session streams raw + PCM or G.711: `container` may only be `raw`, and `bit_rate` + does not apply. client: A `Client` whose settings (API key, `tts_base_url`) to use. Defaults to the shared default client. api_key: The API key to authenticate with; takes precedence over @@ -119,9 +121,11 @@ def __init__( Raises: TypeError: if `config` is not a `TTSConfig`. + ValueError: if `config` sets a `wav` container or `bit_rate`. """ _audio.check_config(type(self).__name__, config, TTSConfig) self.config = config or TTSConfig() + check_session_config(self.config) self._settings = resolve_settings( type(self).__name__, client.settings if client else None, api_key, token ) diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index 972478df..c946e0a9 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -338,10 +338,24 @@ def test_audio_before_connect_says_the_session_is_not_open(server): _drain(aai.TTSSession().audio()) -def test_http_only_options_are_refused_on_a_session(server): - with pytest.raises(ValueError, match="container and bit_rate"): - with aai.TTSSession(config=aai.TTSConfig(container="wav")): - pass +@pytest.mark.parametrize( + "config", + [aai.TTSConfig(container="wav"), aai.TTSConfig(bit_rate=64000)], + ids=["wav", "bit_rate"], +) +def test_http_only_options_are_refused_when_the_session_is_built(server, config): + with pytest.raises(ValueError, match="wav container and bit_rate"): + aai.TTSSession(config=config) + assert not server.connect.called + + +@pytest.mark.parametrize( + "container", ["raw", aai.TTSContainer.raw], ids=["str", "enum"] +) +def test_a_raw_container_is_what_a_session_streams(server, container): + with aai.TTSSession(config=aai.TTSConfig(container=container)): + pass + assert "container" not in parse_qs(urlsplit(server.connect.call_args.args[0]).query) def test_a_cancel_before_a_request_does_not_end_it(server): diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 6cb3cdbc..256292b7 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -608,6 +608,22 @@ async def test_a_path_prefix_in_the_host_setting_is_kept(server, monkeypatch): assert server.connect.uri.startswith("wss://proxy.example.com/tts/v1/ws?") +@pytest.mark.parametrize( + "config", + [aai.TTSConfig(container="wav"), aai.TTSConfig(bit_rate=64000)], + ids=["wav", "bit_rate"], +) +async def test_http_only_options_are_refused_when_the_session_is_built(config): + with pytest.raises(ValueError, match="wav container and bit_rate"): + aai.AsyncTTSSession(config=config) + + +async def test_a_raw_container_is_what_a_session_streams(server): + async with aai.AsyncTTSSession(config=aai.TTSConfig(container="raw")): + pass + assert "container" not in parse_qs(urlsplit(server.connect.uri).query) + + async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) monkeypatch.setattr(aai.settings, "api_key", None) From 1946d4e80956f7506f8eebd62da2d7394a8c6e94 Mon Sep 17 00:00:00 2001 From: Ben Gotthold Date: Thu, 22 Oct 2026 21:04:13 +0000 Subject: [PATCH 20/21] fix(tts): review follow-ups on the session fixes Design calls from Aleks's review: - a session failure nothing in the `with` block observed (push mode, or text flushed and never read) is raised when the block ends; - `speak()` splits text longer than one Generate frame across frames; - `word_boundaries` is documented as off by default here while the API's default is on; the async `speak()` is documented as sending on first iteration. Found reviewing the fixes: - a stale `audio()` iterator ends when the session reconnects instead of reading the new connection's queue; - a send that meets the API's close waits for the reader outside the send lock, so a handler that sends cannot block the reader; - a handshake that fails on a reconnect leaves no fresh state for `audio()` to wait on; - a connection dropped during the handshake (`InvalidMessage`) is retryable; - the async dispatch no longer moves the cancel watermark twice. Co-Authored-By: Claude Fable 5.1 --- README.md | 2 +- assemblyai/tts/_base_ws.py | 8 +- assemblyai/tts/v1/async_session.py | 44 ++++++--- assemblyai/tts/v1/session.py | 52 +++++++--- tests/unit/test_tts_session.py | 142 +++++++++++++++++++++++++++ tests/unit/test_tts_session_async.py | 53 ++++++++++ 6 files changed, 274 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index d7d32444..c90f9cff 100644 --- a/README.md +++ b/README.md @@ -1018,7 +1018,7 @@ asyncio.run(main()) `aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary (`None` if the session failed instead). -A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, a `wav` container and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`. +A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, a `wav` container and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`; word boundaries are off unless asked for. A session that fails raises `TTSError` wherever the app meets it first: the `audio()` iteration, the next send, or, for an app that pushes audio to a handler and never reads, the end of the `with` block. ```python import assemblyai as aai diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index a5e9894b..83e34fb0 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -142,10 +142,12 @@ def handshake_error(exc: BaseException) -> TTSError: if isinstance(status, int): retryable = status in _RETRYABLE_STATUSES else: + from websockets.exceptions import InvalidMessage # a dropped handshake + status = None - retryable = isinstance(exc, (OSError, asyncio.TimeoutError)) and not isinstance( - exc, ssl.CertificateError - ) + retryable = isinstance( + exc, (OSError, asyncio.TimeoutError, InvalidMessage) + ) and not isinstance(exc, ssl.CertificateError) return TTSError( f"WebSocket handshake failed: {exc}", status_code=status, retryable=retryable ) diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index 13951981..e241f57d 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -144,6 +144,7 @@ async def connect(self) -> TTSBeginEvent: self.begin = None self.termination = None self.error = None + self._error_raised = False uri, headers = build_uri( self._settings.tts_base_url, @@ -158,6 +159,7 @@ async def connect(self) -> TTSBeginEvent: websocket_connect_async(uri, headers), timeout=self._connect_timeout ) except Exception as exc: # the handshake failed + await self.close() # ends any audio() waiting on the fresh state raise handshake_error(exc) from exc self._reader = asyncio.create_task(self._read_loop()) @@ -171,7 +173,7 @@ async def connect(self) -> TTSBeginEvent: raise if self.error is not None: await self.close() - raise self.error + raise self._surfaced() if self.begin is None: await self.close() raise TTSError("The session closed before it began") @@ -230,10 +232,12 @@ async def __aenter__(self) -> "AsyncTTSSession": return self async def __aexit__(self, exc_type: Any, *exc: object) -> None: - if exc_type is None: - await self.terminate() - else: + if exc_type is not None: await self.close() + return + await self.terminate() + if self.error is not None and not self._error_raised: + raise self._surfaced() async def generate(self, text: str) -> None: """Appends `text` to the current request; see `TTSSession.generate`.""" @@ -262,10 +266,18 @@ async def keep_alive(self) -> None: await self._send({"type": "KeepAlive"}) async def speak(self, text: str) -> AsyncIterator[bytes]: - """Speaks one text and yields its audio; see `TTSSession.speak`.""" + """ + Speaks one text and yields its audio; see `TTSSession.speak`. + + Unlike `TTSSession.speak()`, nothing is sent until iteration starts, + so pipeline requests with `generate()`, `flush()` and `audio()`. + """ self._check_not_in_handler("speak()") - await self.generate(text) + check_text_frame(text, 0) + limit = max_generate_chars(self.begin) or max(len(text), 1) + for start in range(0, len(text), limit): + await self.generate(text[start : start + limit]) flush_id = await self.flush() async for chunk in self.audio(flush_id): yield chunk @@ -290,14 +302,17 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: "audio goes to the registered Audio handler; pull with audio() " "or speak() only without one" ) + connection = self._queue # this connection's; a reconnect makes a new one while True: + if self._queue is not connection: + return if flush_id < self._cancelled_before: # dropped by a Cancel return item = await self._get(timeout=0.5) if item is None: if self._stopped.is_set(): if self.error is not None: - raise self.error + raise self._surfaced() return continue kind, ident, payload = item @@ -311,12 +326,20 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: if ident > flush_id: return elif kind == "error": - raise payload + raise self._surfaced() elif kind == "closed": if self.error is not None: - raise self.error + raise self._surfaced() return + def _surfaced(self) -> TTSError: + """The session's error, marked as raised to the caller.""" + + if self.error is None: + raise RuntimeError("the session has no error") + self._error_raised = True + return self.error + def _check_not_in_handler(self, method: str) -> None: """Refuses a method that waits on the reader when the reader, running a handler, is the caller: it would wait for itself.""" @@ -371,7 +394,7 @@ async def _send(self, frame: Dict[str, Any]) -> None: ws = self._ws if ws is None or self._stopped.is_set(): if self.error is not None: - raise self.error + raise self._surfaced() raise RuntimeError( "the session is not open: use `with` / `async with`, or connect()" ) @@ -458,7 +481,6 @@ async def _dispatch(self, data: Dict[str, Any]) -> None: self._put(("done", event.flush_id, None)) elif kind == "Cancelled": cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 - self._cancelled_before = max(self._cancelled_before, cancel_id) self._put(("cancelled", cancel_id, None)) elif kind == "Termination": self.termination = event diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 89970ee2..1eabaa6c 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -113,6 +113,8 @@ def __init__( session begins, so such a session connects once: to reconnect, mint a new token and build a new session. word_boundaries: Whether the server sends `WordBoundaries` frames. + Off by default here (the API's own default is on); pass + `True` to receive them. inactivity_timeout: Seconds (5 to 120) with no frame either way after which the API closes the session; `None` keeps the API's 120 s. @@ -165,6 +167,7 @@ def _reset(self) -> None: self.begin = None self.termination = None self.error = None + self._error_raised = False # -- lifecycle ----------------------------------------------------------- @@ -220,6 +223,7 @@ def connect(self) -> TTSBeginEvent: uri, additional_headers=headers, open_timeout=self._connect_timeout ) except Exception as exc: # the handshake failed + self.close() # ends any audio() waiting on the fresh state raise handshake_error(exc) from exc self._reader = threading.Thread( @@ -232,7 +236,7 @@ def connect(self) -> TTSBeginEvent: raise TTSError("The session did not begin in time") if self.error is not None: self.close() - raise self.error + raise self._surfaced() if self.begin is None: self.close() raise TTSError("The session closed before it began") @@ -280,10 +284,14 @@ def __enter__(self) -> "TTSSession": return self def __exit__(self, exc_type: Any, *exc: object) -> None: - if exc_type is None: - self.terminate() - else: + if exc_type is not None: self.close() + return + self.terminate() + # A failure nothing in the block observed (push mode, or text sent + # and never read) must not pass silently. + if self.error is not None and not self._error_raised: + raise self._surfaced() # -- sending --------------------------------------------------------------- @@ -344,15 +352,21 @@ def speak(self, text: str) -> Iterator[bytes]: """ Speaks one text: `generate()`, `flush()`, then its audio. - Returns an iterator of the request's audio bytes that ends at its - `FlushDone`, or early if the request is cancelled. + Text longer than one `Generate` frame is split across frames; very + long text (several thousand characters) should be paced instead, as + the API limits unspoken text. Returns an iterator of the request's + audio bytes that ends at its `FlushDone`, or early if the request + is cancelled. Raises: TTSError: from the iteration, if the session fails. RuntimeError: from a session handler. """ self._check_not_in_handler("speak()") - self.generate(text) + check_text_frame(text, 0) + limit = max_generate_chars(self.begin) or max(len(text), 1) + for start in range(0, len(text), limit): + self.generate(text[start : start + limit]) return self.audio(self.flush()) def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: @@ -383,14 +397,17 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: "audio goes to the registered Audio handler; pull with audio() " "or speak() only without one" ) + connection = self._queue # this connection's; a reconnect makes a new one while True: + if self._queue is not connection: + return if flush_id < self._cancelled_before: # dropped by a Cancel return item = self._get(timeout=0.5) if item is None: if self._stopped.is_set(): if self.error is not None: - raise self.error + raise self._surfaced() return continue kind, ident, payload = item @@ -404,14 +421,22 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: if ident > flush_id: # a Cancel sent after this request return elif kind == "error": - raise payload + raise self._surfaced() elif kind == "closed": if self.error is not None: - raise self.error + raise self._surfaced() return # -- internals --------------------------------------------------------------- + def _surfaced(self) -> TTSError: + """The session's error, marked as raised to the caller.""" + + if self.error is None: + raise RuntimeError("the session has no error") + self._error_raised = True + return self.error + def _check_not_in_handler(self, method: str) -> None: """Refuses a method that waits on the reader when the reader, running a handler, is the caller: it would wait for itself.""" @@ -452,14 +477,17 @@ def _send(self, frame: Dict[str, Any]) -> None: ws = self._ws if ws is None or self._stopped.is_set(): if self.error is not None: - raise self.error + raise self._surfaced() raise RuntimeError( "the session is not open: use `with` / `async with`, or connect()" ) try: ws.send(json.dumps(frame)) + return except ConnectionClosed as exc: - raise self._closed_error(ws, exc) from exc + closed = exc + # Outside the lock: a handler that sends must not block the reader. + raise self._closed_error(ws, closed) from closed def _closed_error(self, ws: Any, exc: Exception) -> Exception: """ diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index c946e0a9..b96608a3 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -853,3 +853,145 @@ def test_an_exception_inside_with_closes_without_terminate(server): def test_wrong_config_type_is_a_type_error(): with pytest.raises(TypeError): aai.TTSSession(config=aai.DictationConfig()) # type: ignore[arg-type] + + +def test_a_failure_nobody_read_is_raised_when_the_block_ends(server): + # push mode: an Audio handler, text flushed, then the API fails + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Audio, lambda e: None) + server.answer_flush = False + session.generate("Never read.") + session.flush() + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + session._stopped.wait(2) + assert exc_info.value.error_code == "input_rate_exceeded" + + +def test_an_error_the_block_already_saw_is_not_raised_twice(server): + seen = [] + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push( + { + "type": "Error", + "error_code": 3005, + "error": "No room", + "error_type": "at_capacity", + "retryable": True, + "param": None, + "retry_after_seconds": 2, + } + ) + try: + list(session.audio(flush_id)) + except aai.TTSError as err: + seen.append(err) + assert len(seen) == 1 # the with block exited quietly + + +def test_speak_splits_long_text_across_generate_frames(server): + with aai.TTSSession() as session: + text = "a" * 2000 + "b" * 2000 + "c" * 1000 + assert b"".join(session.speak(text)) == text.encode() + sent = [f["text"] for f in server.sent if f["type"] == "Generate"] + assert [len(t) for t in sent] == [2000, 2000, 1000] + + +def test_a_stale_audio_iterator_ends_when_the_session_reconnects(mocker): + fakes = [_FakeServer(answer_flush=False), _FakeServer()] + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=fakes) + session = aai.TTSSession() + session.connect() + stale = session.audio(session.flush()) # a player still holding request 0 + fakes[0].push(ConnectionClosedError(Close(3006, "idle"), None)) + session._stopped.wait(2) + with pytest.raises(aai.TTSError): + next(stale) # the old connection's failure, raised once + + session.connect() # same object, new connection, ids start again at 0 + assert list(stale) == [] # the old iterator does not read the new queue + assert b"".join(session.speak("NEW")) == b"NEW" + session.terminate() + + +def test_a_failed_reconnect_does_not_leave_audio_waiting(mocker): + import time + + first = _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + side_effect=[first, OSError("refused")], + ) + session = aai.TTSSession() + with session: + pass + with pytest.raises(aai.TTSError) as exc_info: + session.connect() + assert exc_info.value.retryable is True + started = time.monotonic() + assert list(session.audio(0)) == [] # ends at once instead of polling forever + assert time.monotonic() - started < 2 + + +def test_a_dropped_handshake_is_retryable(): + from websockets.exceptions import InvalidMessage + + from assemblyai.tts._base_ws import handshake_error + + assert handshake_error( + InvalidMessage("did not receive a valid HTTP response") + ).retryable + assert handshake_error(ConnectionRefusedError()).retryable + assert not handshake_error(ValueError("bad uri")).retryable + + +def test_a_send_that_meets_a_close_does_not_hold_the_lock_while_waiting(server): + import time + + server.answer_flush = False + blocked_for = {} + + def on_flush_done(event): + t0 = time.monotonic() + try: + session.keep_alive() + except Exception: + pass + blocked_for["s"] = time.monotonic() - t0 + + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.FlushDone, on_flush_done) + session.flush() + # the API closes with an Error frame; the main thread's send races it + server.inbox.put( + json.dumps({"type": "FlushDone", "flush_id": 0, "audio_duration_ms": 1}) + ) + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + server.push(ConnectionClosedError(Close(3010, ""), None)) + session._stopped.wait(2) + session.generate("late") + assert exc_info.value.error_code == "input_rate_exceeded" + assert blocked_for.get("s", 0) < 0.5 diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 256292b7..f6270447 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -629,3 +629,56 @@ async def test_missing_key_is_a_clear_error(monkeypatch): monkeypatch.setattr(aai.settings, "api_key", None) with pytest.raises(ValueError, match="api_key= or token="): aai.AsyncTTSSession() + + +async def test_a_failure_nobody_read_is_raised_when_the_block_ends(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Audio, lambda e: None) + server.answer_flush = False + await session.generate("Never read.") + await session.flush() + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + await asyncio.wait_for(session._stopped.wait(), 2) + assert exc_info.value.error_code == "input_rate_exceeded" + + +async def test_speak_splits_long_text_across_generate_frames(server): + async with aai.AsyncTTSSession() as session: + text = "a" * 2000 + "b" * 500 + assert b"".join([c async for c in session.speak(text)]) == text.encode() + sent = [f["text"] for f in server.sent if f["type"] == "Generate"] + assert [len(t) for t in sent] == [2000, 500] + + +async def test_a_stale_audio_iterator_ends_when_the_session_reconnects(mocker): + fakes = [_FakeAsyncServer(answer_flush=False), _FakeAsyncServer()] + + async def connect(uri, additional_headers): + return fakes.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + await session.connect() + stale = session.audio(await session.flush()) + fakes[0].push(ConnectionClosedError(Close(3006, "idle"), None)) + await asyncio.wait_for(session._stopped.wait(), 2) + with pytest.raises(aai.TTSError): + await stale.__anext__() + + await session.connect() + assert [c async for c in stale] == [] + assert b"".join([c async for c in session.speak("NEW")]) == b"NEW" + await session.terminate() From dd30e412d0c65026ed9cfd40feebe3275c9a23ac Mon Sep 17 00:00:00 2001 From: Ben Gotthold Date: Thu, 22 Oct 2026 21:07:02 +0000 Subject: [PATCH 21/21] test(tts): keep a handle on the first fake server in the async reconnect test Co-Authored-By: Claude Fable 5.1 --- tests/unit/test_tts_session_async.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index f6270447..4c4c5f3f 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -662,10 +662,11 @@ async def test_speak_splits_long_text_across_generate_frames(server): async def test_a_stale_audio_iterator_ends_when_the_session_reconnects(mocker): - fakes = [_FakeAsyncServer(answer_flush=False), _FakeAsyncServer()] + first, second = _FakeAsyncServer(answer_flush=False), _FakeAsyncServer() + servers = [first, second] async def connect(uri, additional_headers): - return fakes.pop(0) + return servers.pop(0) mocker.patch( "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect @@ -673,7 +674,7 @@ async def connect(uri, additional_headers): session = aai.AsyncTTSSession() await session.connect() stale = session.audio(await session.flush()) - fakes[0].push(ConnectionClosedError(Close(3006, "idle"), None)) + first.push(ConnectionClosedError(Close(3006, "idle"), None)) await asyncio.wait_for(session._stopped.wait(), 2) with pytest.raises(aai.TTSError): await stale.__anext__()