xai TTS: input_ended race in _run_ws can drop audio.done and hang the synthesis

[moamen-bytesai]

Plugin / version

livekit-plugins-xai 1.6.0, livekit-agents 1.6.0, Python 3.13

Verified on the latest release: the 1.6.0 wheel from PyPI (uploaded 2026-06-11) contains the same pattern in SynthesizeStream._run_ws (tts.py lines 247–281), as does the current main.

Problem

In SynthesizeStream._run_ws, input_ended is a plain bool set by _send_task after the await ws.send_str(...) for text.done returns:

async def _send_task(ws):
    async for word in input_stream:
        self._mark_started()
        await ws.send_str(json.dumps({"type": "text.delta", "delta": word.token}))
    await ws.send_str(json.dumps({"type": "text.done"}))
    input_ended = True   # set only when _send_task resumes

_recv_task only terminates on audio.done when the flag is already set:

elif msg_type == "audio.done":
    if input_ended:
        output_emitter.end_segment()
        break

While _send_task is suspended in the final await ws.send_str(...), the event loop can run _recv_task. If the server's audio.done arrives and is processed in that window — i.e. before the loop resumes _send_task's continuation that sets the flag — the audio.done is silently ignored (no break, no end_segment). _recv_task then keeps waiting on ws.receive() for messages that will never come, and the synthesis hangs until the connection closes or an outer timeout fires.

The window is small (it requires the server's final message to be processed between the send completing and the sender task resuming), but it is real — task resumption order after an await is not guaranteed, and the likelihood grows under event-loop load, with short inputs, and with low-latency synthesis settings.

Expected behavior

audio.done handling should not depend on scheduler ordering — e.g. an asyncio.Event that _recv_task awaits on audio.done (yielding immediately if already set) instead of a plain bool check, or any equivalent synchronization that guarantees the final message is never dropped.

[longcw]

  await ws.send_str(json.dumps({"type": "text.done"}))   # line 255
  input_ended = True                                       # line 256

ws.send_str returns when the data is drained from the local app buffer — which is causally before the server even receives text.done, input_ended = True runs synchronously the instant _send_task resumes. I don't think it makes sense audio.done can be received before input ended.

can you share a reproducible example, probably it's a different issue that caused the synthesis hang.