Add connection pooling for xAI TTS#6317
Conversation
|
|
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| async def _connect_pooled_ws(self, timeout: float) -> aiohttp.ClientWebSocketResponse: | ||
| return await self._connect_ws(timeout, self._opts) |
There was a problem hiding this comment.
🚩 Pooled connections now use TTS-level options instead of stream-level snapshot
The old code at the removed line ws = await self._tts._connect_ws(self._conn_options.timeout, self._opts) used self._opts which is the stream's snapshot copy (set at tts.py:256 via self._opts = replace(tts._opts)). The new _connect_pooled_ws at tts.py:163-164 uses self._opts referring to TTS._opts — the TTS instance's current (possibly mutated) options. This means if update_options() is called after a stream is created but before its _run_ws executes, the WebSocket connection parameters (voice, language, speed, etc.) will reflect the new options rather than the options the stream was created with. This is consistent with how Deepgram and Cartesia plugins handle pooled connections, so it appears intentional, but it is a behavioral change from the previous implementation.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Intentional. With pooled websockets, the connection parameters are scoped to the pool/TTS options rather than an individual stream snapshot; this matches the existing Deepgram/Cartesia pooling pattern. update_options() invalidates the pool when URL-affecting options change, so later requests connect with the updated params instead of reusing stale sockets. The stream snapshot remains useful for per-stream tokenizer state.
| async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws: | ||
| self._acquire_time = self._tts._pool.last_acquire_time | ||
| self._connection_reused = self._tts._pool.last_connection_reused | ||
| tasks = [ | ||
| asyncio.create_task(_send_task(ws)), | ||
| asyncio.create_task(_recv_task(ws)), | ||
| ] | ||
| try: | ||
| await asyncio.gather(*tasks) | ||
| finally: | ||
| await utils.aio.gracefully_cancel(*tasks) |
There was a problem hiding this comment.
🚩 Connection reuse assumes xAI WebSocket supports sequential request multiplexing
The connection pool returns the WebSocket to the available set after a successful _run_ws call (via the connection() context manager's put() at connection_pool.py:94). The next _run_ws call then reuses the same WebSocket, sending new text.delta messages. This assumes the xAI server supports receiving a new sequence of text.delta → text.done on the same WebSocket after a previous audio.done has been received. If the xAI server instead expects one request per WebSocket connection, the second segment would fail or produce incorrect audio. The test at test_plugin_xai_tts.py:110-116 validates this with a fake WebSocket but cannot confirm server-side behavior.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Verified against xAI docs: the streaming TTS websocket supports multi-utterance on the same connection. After audio.done, clients can send another text.delta / text.done sequence, and xAI documents that the connection remains open for this pattern: https://docs.x.ai/developers/rest-api-reference/inference/voice. That is the provider behavior this pool relies on.
There was a problem hiding this comment.
🚩 Pre-existing: last input stream not finalized when channel closes without flush
In _tokenize_input at lines 270-283, if the _input_ch closes (the async for loop ends) while an input_stream is active (i.e., text was pushed but no FlushSentinel was received), end_input() is never called on that stream. The segments_ch is closed, but the last input_stream remains open. This means _run_ws would iterate async for word in input_stream and never finish, causing the pooled WebSocket connection to be held indefinitely.
This is pre-existing behavior (not introduced by this PR), but the impact is amplified with connection pooling since the held connection blocks the pool rather than just leaking a single WebSocket. Other plugins like Deepgram have the same pattern (livekit-plugins/livekit-plugins-deepgram/livekit/plugins/deepgram/tts.py:300-314), so this appears to be a framework-level convention where the caller is expected to always send a flush before closing.
(Refers to lines 270-283)
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
This is pre-existing and outside this PR scope. The normal public path is SynthesizeStream.end_input(), which calls flush() before closing _input_ch, so the active segment is finalized. Directly closing _input_ch without a flush would affect the same tokenizer pattern in other plugins such as Deepgram as well; this PR keeps the change scoped to xAI websocket connection reuse.
3 participants
Summary
ConnectionPoolto the xAI TTS plugin so streamed synthesis segments can reuse websocket connections.prewarm()through the pool and close pooled connections fromaclose().Fixes #6078
Root cause
The xAI TTS stream opened a websocket with
_connect_ws()for each synthesized segment and always closed it in_run_ws()cleanup. Unlike other websocket-backed TTS plugins, it did not use the sharedutils.ConnectionPool, so repeated streamed segments could not reuse an existing connection or be prewarmed.Notes
xAI documents the streaming TTS websocket as multi-utterance: after
audio.done, clients can send anothertext.delta/text.donesequence on the same connection.Test plan
uv run pytest tests/test_plugin_xai_tts.py --plugin xai -quv run pytest tests/test_realtime/test_xai_realtime_model.py --unit -quv run ruff check tests/test_plugin_xai_tts.py livekit-plugins/livekit-plugins-xai/livekit/plugins/xai/tts.pyuv run ruff format --check tests/test_plugin_xai_tts.py livekit-plugins/livekit-plugins-xai/livekit/plugins/xai/tts.pyuv run --group typing mypy --untyped-calls-exclude=smithy_aws_core -p livekit.plugins.xai