diff --git a/.changeset/cartesia-persistent-tts-websocket.md b/.changeset/cartesia-persistent-tts-websocket.md new file mode 100644 index 000000000..34e60bf8c --- /dev/null +++ b/.changeset/cartesia-persistent-tts-websocket.md @@ -0,0 +1,5 @@ +--- +'@livekit/agents-plugin-cartesia': patch +--- + +Reuse a single Cartesia TTS WebSocket across generations instead of opening and closing one per synthesis. The plugin now holds a `ConnectionPool` on the `TTS` instance (matching the Python plugin and the fishaudio/inworld/xai plugins), so only the first turn pays the connect and later turns skip the TCP/TLS and WebSocket handshake. Adds `TTS.prewarm()` to open the socket before the first turn and a `TTS.close()` that drains the pool. diff --git a/plugins/cartesia/src/tts.test.ts b/plugins/cartesia/src/tts.test.ts index c2ccead8b..4b64ce9bf 100644 --- a/plugins/cartesia/src/tts.test.ts +++ b/plugins/cartesia/src/tts.test.ts @@ -1,19 +1,268 @@ // SPDX-FileCopyrightText: 2024 LiveKit, Inc. // // SPDX-License-Identifier: Apache-2.0 +import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS, tts } from '@livekit/agents'; import { STT } from '@livekit/agents-plugin-openai'; -import { tts } from '@livekit/agents-plugins-test'; -import { describe, it } from 'vitest'; +import { tts as testTts } from '@livekit/agents-plugins-test'; +import { once } from 'node:events'; +import type { AddressInfo } from 'node:net'; +import { describe, expect, it } from 'vitest'; +import { type WebSocket, WebSocketServer } from 'ws'; import { TTS } from './tts.js'; const hasCartesiaConfig = Boolean(process.env.CARTESIA_API_KEY && process.env.OPENAI_API_KEY); if (hasCartesiaConfig) { describe('Cartesia', async () => { - await tts(new TTS(), new STT()); + await testTts(new TTS(), new STT()); }); } else { describe('Cartesia', () => { it.skip('requires CARTESIA_API_KEY and OPENAI_API_KEY', () => {}); }); } + +// A single 24 kHz mono s16le frame's worth of silence, base64-encoded the way +// Cartesia sends audio chunks. +const CHUNK_BASE64 = Buffer.alloc(4800).toString('base64'); + +async function startWebSocketServer() { + const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 }); + await once(wss, 'listening'); + const address = wss.address() as AddressInfo; + return { wss, baseURL: `http://127.0.0.1:${address.port}` }; +} + +async function closeWebSocketServer(wss: WebSocketServer): Promise { + for (const client of wss.clients) { + client.close(); + } + await new Promise((resolve) => wss.close(() => resolve())); +} + +async function waitFor(promise: Promise, timeoutMs = 1000): Promise { + let timeout: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timeout = setTimeout(() => reject(new Error('timed out waiting for promise')), timeoutMs); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } +} + +// A minimal Cartesia TTS WebSocket server: for every generation it replies with +// one audio chunk and a done message, echoing the caller's context_id. `onStop` +// lets a test override the reply (e.g. to simulate a provider failure); return +// false to suppress the normal chunk/done reply. +function serveCartesia( + wss: WebSocketServer, + onStop?: (ws: WebSocket, contextId: string, connectionNumber: number) => boolean, +): { connectionCount: () => number } { + let connectionCount = 0; + wss.on('connection', (ws) => { + connectionCount++; + const connectionNumber = connectionCount; + ws.on('message', (raw) => { + const message = JSON.parse(raw.toString()) as { context_id: string; continue?: boolean }; + if (message.continue !== false) return; // only reply once the turn is closed + const contextId = message.context_id; + if (onStop && !onStop(ws, contextId, connectionNumber)) return; + ws.send( + JSON.stringify({ + type: 'chunk', + data: CHUNK_BASE64, + done: false, + status_code: 200, + step_time: 0, + context_id: contextId, + }), + ); + ws.send( + JSON.stringify({ type: 'done', done: true, status_code: 200, context_id: contextId }), + ); + }); + }); + return { connectionCount: () => connectionCount }; +} + +async function synthesizeTurn( + cartesia: TTS, + text: string, + connOptions?: APIConnectOptions, +): Promise { + const stream = cartesia.stream({ connOptions }); + stream.pushText(text); + stream.endInput(); + + try { + const events: tts.SynthesizedAudio[] = []; + for await (const event of stream) { + if (event !== tts.SynthesizeStream.END_OF_STREAM) events.push(event); + } + return events; + } finally { + stream.close(); + } +} + +describe('Cartesia streaming pool', () => { + it('reuses one websocket across sequential turns', async () => { + const { wss, baseURL } = await startWebSocketServer(); + const server = serveCartesia(wss); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + expect(await synthesizeTurn(cartesia, 'first turn.')).not.toHaveLength(0); + expect(await synthesizeTurn(cartesia, 'second turn.')).not.toHaveLength(0); + expect(server.connectionCount()).toBe(1); + } finally { + await cartesia.close(); + await closeWebSocketServer(wss); + } + }); + + it('prewarms and reuses the ready websocket', async () => { + const { wss, baseURL } = await startWebSocketServer(); + const server = serveCartesia(wss); + const connected = once(wss, 'connection'); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + cartesia.prewarm(); + await waitFor(connected); + expect(await synthesizeTurn(cartesia, 'prewarmed turn.')).not.toHaveLength(0); + expect(server.connectionCount()).toBe(1); + } finally { + await cartesia.close(); + await closeWebSocketServer(wss); + } + }); + + it('discards a poisoned websocket after a failure', async () => { + const { wss, baseURL } = await startWebSocketServer(); + // The first connection drops the turn; the second serves it normally. + const server = serveCartesia(wss, (ws, _contextId, connectionNumber) => { + if (connectionNumber === 1) { + ws.close(1011, 'provider failure'); + return false; + } + return true; + }); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + expect( + await synthesizeTurn(cartesia, 'failing turn.', { + ...DEFAULT_API_CONNECT_OPTIONS, + maxRetry: 0, + }), + ).toHaveLength(0); + expect(await synthesizeTurn(cartesia, 'recovery turn.')).not.toHaveLength(0); + expect(server.connectionCount()).toBe(2); + } finally { + await cartesia.close(); + await closeWebSocketServer(wss); + } + }); + + it('fails over when the socket drops mid-generation instead of ending silently', async () => { + const { wss, baseURL } = await startWebSocketServer(); + // Connection 1 emits one audio chunk, then drops WITHOUT a done message, + // i.e. mid-speech. Connection 2 serves the recovery turn normally. + const server = serveCartesia(wss, (ws, contextId, connectionNumber) => { + if (connectionNumber === 1) { + ws.send( + JSON.stringify({ + type: 'chunk', + data: CHUNK_BASE64, + done: false, + status_code: 200, + step_time: 0, + context_id: contextId, + }), + ); + setTimeout(() => ws.close(1011, 'mid-speech drop'), 5); + return false; // suppress the normal chunk/done reply + } + return true; + }); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + // The dropped turn does not complete successfully (it fails over rather + // than silently ending); at maxRetry: 0 that surfaces as no audio. + expect( + await synthesizeTurn(cartesia, 'dropping turn.', { + ...DEFAULT_API_CONNECT_OPTIONS, + maxRetry: 0, + }), + ).toHaveLength(0); + // The dead socket is discarded, so the next turn opens a fresh one. + expect(await synthesizeTurn(cartesia, 'recovery turn.')).not.toHaveLength(0); + expect(server.connectionCount()).toBe(2); + } finally { + await cartesia.close(); + await closeWebSocketServer(wss); + } + }); + + it('replaces a websocket that closed while idle', async () => { + const { wss, baseURL } = await startWebSocketServer(); + let firstConnectionClosed: (() => void) | undefined; + const firstClosed = new Promise((resolve) => { + firstConnectionClosed = resolve; + }); + const server = serveCartesia(wss, (ws, _contextId, connectionNumber) => { + if (connectionNumber === 1) { + ws.on('close', () => firstConnectionClosed?.()); + // Serve the turn, then drop the idle socket so the next turn reconnects. + setTimeout(() => ws.close(), 10); + } + return true; + }); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + expect(await synthesizeTurn(cartesia, 'first turn.')).not.toHaveLength(0); + await waitFor(firstClosed); + // Let the client observe the close so the idle handler removes the socket + // before the next checkout, making the maxRetry: 0 assertion deterministic. + await new Promise((resolve) => setTimeout(resolve, 100)); + // maxRetry: 0 proves the idle-closed socket was dropped from the pool, not + // handed back to burn the turn's only attempt. + expect( + await waitFor( + synthesizeTurn(cartesia, 'second turn.', { ...DEFAULT_API_CONNECT_OPTIONS, maxRetry: 0 }), + ), + ).not.toHaveLength(0); + expect(server.connectionCount()).toBe(2); + } finally { + await cartesia.close(); + await closeWebSocketServer(wss); + } + }); + + it('closes the pooled websocket when the TTS closes', async () => { + const { wss, baseURL } = await startWebSocketServer(); + serveCartesia(wss); + + const cartesia = new TTS({ apiKey: 'test-key', baseUrl: baseURL }); + try { + await synthesizeTurn(cartesia, 'closing turn.'); + await cartesia.close(); + // close() drains the pooled socket; give the close frame a beat to land. + await waitFor( + (async () => { + while (wss.clients.size > 0) await new Promise((r) => setTimeout(r, 5)); + })(), + ); + expect(wss.clients.size).toBe(0); + } finally { + await closeWebSocketServer(wss); + } + }); +}); diff --git a/plugins/cartesia/src/tts.ts b/plugins/cartesia/src/tts.ts index cc622f36d..218e8f925 100644 --- a/plugins/cartesia/src/tts.ts +++ b/plugins/cartesia/src/tts.ts @@ -8,6 +8,7 @@ import { APIStatusError, APITimeoutError, AudioByteStream, + ConnectionPool, Future, type TimedString, asError, @@ -48,6 +49,13 @@ const API_VERSION_WITH_EXPERIMENTAL_CONTROLS = '2024-11-13'; const MODEL_WITH_EXPERIMENTAL_CONTROLS = 'sonic-2-2025-03-07'; const NUM_CHANNELS = 1; const BUFFERED_WORDS_COUNT = 8; +// Cartesia refreshes a pooled socket after this long so a very long call cannot +// keep one connection open indefinitely. Matches the Python plugin's 300s. +const MAX_SESSION_DURATION_MS = 300_000; + +// Lets each SynthesizeStream reach the pool owned by the TTS that created it, +// without widening the constructor signature the base class fixes. +const connectionPools = new WeakMap>(); export interface TTSOptions { model: TTSModels | string; @@ -128,6 +136,8 @@ const checkGenerationConfig = (opts: TTSOptions) => { export class TTS extends tts.TTS { #opts: TTSOptions; + #pool: ConnectionPool; + #closed = false; label = 'cartesia.TTS'; get model(): string { @@ -166,9 +176,31 @@ export class TTS extends tts.TTS { ) { checkGenerationConfig(this.#opts); } + + // One socket, reused across generations. Cartesia recommends a single + // preconnected WebSocket for many generations because a fresh connection + // repays TCP/TLS setup on every turn: + // https://docs.cartesia.ai/use-the-api/compare-tts-endpoints + this.#pool = new ConnectionPool({ + connectCb: (timeoutMs) => this.#connectWebSocket(timeoutMs), + closeCb: async (ws) => safeCloseWebSocket(ws), + maxSessionDuration: MAX_SESSION_DURATION_MS, + markRefreshedOnGet: true, + }); + connectionPools.set(this, this.#pool); } updateOptions(opts: Partial) { + // Only these three fields reach Cartesia at WebSocket-handshake time (auth + // header, version header, host). Everything else (model, voice, encoding, + // sample rate, speed, emotion, volume, language) is sent in-band on each + // generation, so a pooled socket serves the new value without reconnecting. + // Reconnect only when one of the handshake inputs actually changes. + const handshakeChanged = + (opts.apiKey !== undefined && opts.apiKey !== this.#opts.apiKey) || + (opts.apiVersion !== undefined && opts.apiVersion !== this.#opts.apiVersion) || + (opts.baseUrl !== undefined && opts.baseUrl !== this.#opts.baseUrl); + this.#opts = { ...this.#opts, ...opts }; if (opts.language !== undefined) { this.#opts.language = normalizeLanguage(opts.language); @@ -182,6 +214,10 @@ export class TTS extends tts.TTS { ) { checkGenerationConfig(this.#opts); } + + if (handshakeChanged) { + this.#pool.invalidate(); + } } synthesize( @@ -189,11 +225,67 @@ export class TTS extends tts.TTS { connOptions?: APIConnectOptions, abortSignal?: AbortSignal, ): tts.ChunkedStream { - return new ChunkedStream(this, text, this.#opts, connOptions, abortSignal); + return new ChunkedStream(this, text, { ...this.#opts }, connOptions, abortSignal); } stream(options?: { connOptions?: APIConnectOptions }): SynthesizeStream { - return new SynthesizeStream(this, this.#opts, options?.connOptions); + return new SynthesizeStream(this, { ...this.#opts }, options?.connOptions); + } + + /** + * Open the pooled WebSocket ahead of the first generation so the first turn + * does not pay the connect. Safe to call more than once; it is a no-op when a + * connection already exists. + */ + prewarm(): void { + this.#pool.prewarm(); + } + + override async close(): Promise { + this.#closed = true; + await this.#pool.close(); + await super.close(); + } + + async #connectWebSocket(timeoutMs: number): Promise { + // Snapshot the handshake inputs. If a concurrent updateOptions() changes one + // of them while this connect is in flight, reconnect on the new value rather + // than pooling a socket built on stale credentials (mirrors the fishaudio + // plugin's model re-check). + const apiKey = this.#opts.apiKey!; + const apiVersion = this.#opts.apiVersion; + const baseUrl = this.#opts.baseUrl; + const url = `${baseUrl.replace(/^http/, 'ws')}/tts/websocket`; + const ws = await connectCartesiaWebSocket({ + url, + headers: { + [AUTHORIZATION_HEADER]: apiKey, + [VERSION_HEADER]: apiVersion, + }, + timeoutMs, + }); + if (this.#closed) { + safeCloseWebSocket(ws); + throw new APIConnectionError({ message: 'Cartesia TTS is closed' }); + } + if ( + apiKey !== this.#opts.apiKey || + apiVersion !== this.#opts.apiVersion || + baseUrl !== this.#opts.baseUrl + ) { + safeCloseWebSocket(ws); + return await this.#connectWebSocket(timeoutMs); + } + // Drop a socket that closes (or errors) while idle in the pool. Between turns + // no generation listeners are attached, so without this the pool keeps a dead + // socket in `available` and the next turn spends a retry to discard it, or + // fails outright at maxRetry:0. A generation attaches its own listeners on top + // of these; the no-op error listener also stops an idle 'error' from crashing + // the process. Remove is a no-op once the socket is no longer pooled, so this + // is safe during an active generation and during close(). + ws.on('error', () => {}); + ws.on('close', () => this.#pool.remove(ws)); + return ws; } } @@ -290,6 +382,7 @@ export class ChunkedStream extends tts.ChunkedStream { export class SynthesizeStream extends tts.SynthesizeStream { #opts: TTSOptions; + #pool: ConnectionPool; #logger = log(); #tokenizer = new tokenize.basic.SentenceTokenizer({ minSentenceLength: BUFFERED_WORDS_COUNT, @@ -298,6 +391,9 @@ export class SynthesizeStream extends tts.SynthesizeStream { constructor(tts: TTS, opts: TTSOptions, connOptions?: APIConnectOptions) { super(tts, connOptions); + const pool = connectionPools.get(tts); + if (!pool) throw new Error('Cartesia connection pool is not initialized'); + this.#pool = pool; this.#opts = opts; } @@ -316,8 +412,7 @@ export class SynthesizeStream extends tts.SynthesizeStream { protected async run() { const requestId = shortuuid(); - let closing = false; - // Only close WebSocket when both: 1) Cartesia returns done, AND 2) all sentences have been sent + // Only finish the generation once both: 1) Cartesia returns done, AND 2) all sentences have been sent let sentenceStreamClosed = false; const sentenceStreamTask = async (ws: WebSocket) => { @@ -384,6 +479,14 @@ export class SynthesizeStream extends tts.SynthesizeStream { }; let timeout: NodeJS.Timeout | null = null; + // Set when the chunk watchdog fires: the socket is discarded, not pooled. + let timedOut = false; + // Set once this generation's `done` has been handled. Until then, a socket + // close or error is a mid-generation drop, not a normal end. + let completed = false; + // A socket close/error before completion. Thrown after the loop so the turn + // fails over (and the dead socket is discarded) instead of ending silently. + let streamError: Error | undefined; const clearTTSChunkTimeout = () => { if (timeout) { @@ -400,15 +503,25 @@ export class SynthesizeStream extends tts.SynthesizeStream { }; const onClose = (code: number, reason: Buffer) => { - if (!closing) { - this.#logger.debug(`WebSocket closed with code ${code}: ${reason.toString()}`); - } + // A close during an active generation is unexpected: the pool owns the + // socket lifecycle and does not close it between turns. If it happens + // before `done`, surface it so the turn retries rather than ending mid + // speech, and so withConnection discards the dead socket. + this.#logger.debug(`WebSocket closed with code ${code}: ${reason.toString()}`); clearTTSChunkTimeout(); + if (!completed && !timedOut && !streamError) { + streamError = new APIConnectionError({ + message: `Cartesia WebSocket closed mid-generation (code=${code})`, + }); + } void eventChannel.close(); }; const onError = (err: Error) => { this.#logger.error({ err }, 'Cartesia WebSocket error'); + if (!completed && !timedOut && !streamError) { + streamError = err instanceof APIError ? err : toRetryableConnectionError(err); + } void eventChannel.close(); }; @@ -476,7 +589,12 @@ export class SynthesizeStream extends tts.SynthesizeStream { this.#logger.debug( `Cartesia WebSocket TTS chunk stream timeout after ${this.#opts.chunkTimeout}ms`, ); - ws.close(); + // The socket is stuck mid-generation, so it must not return to the + // pool. Poison it and unblock the reader; the post-loop check turns + // this into a retryable error so withConnection discards the socket. + timedOut = true; + safeCloseWebSocket(ws); + void eventChannel.close(); }, this.#opts.chunkTimeout); } else if (this.#opts.wordTimestamps !== false && hasWordTimestamps(serverMsg)) { const wordTimestamps = serverMsg.word_timestamps; @@ -507,9 +625,9 @@ export class SynthesizeStream extends tts.SynthesizeStream { } if (segmentId === requestId) { - closing = true; clearTTSChunkTimeout(); - ws.close(); + completed = true; + // Leave the socket open so the pool reuses it on the next turn. break; // Exit the loop } } @@ -520,6 +638,15 @@ export class SynthesizeStream extends tts.SynthesizeStream { this.#logger.warn({ message: serverMsg }, 'Unknown Cartesia message'); } } + + if (timedOut) { + throw new APITimeoutError({ + message: `Cartesia TTS chunk stream timed out after ${this.#opts.chunkTimeout}ms`, + }); + } + if (streamError) { + throw streamError; + } } catch (err) { // Always propagate API errors so the base SynthesizeStream can retry // and emit tts_error once retries are exhausted. @@ -547,32 +674,25 @@ export class SynthesizeStream extends tts.SynthesizeStream { } }; - const wsUrl = this.#opts.baseUrl.replace(/^http/, 'ws'); - const url = `${wsUrl}/tts/websocket`; - - let ws: WebSocket | undefined; try { - ws = await connectCartesiaWebSocket({ - url, - headers: { - [AUTHORIZATION_HEADER]: this.#opts.apiKey!, - [VERSION_HEADER]: this.#opts.apiVersion, + // The pool hands back one live socket per call and reclaims it on success + // (put) or discards it on any thrown error (remove). A generation never + // closes the socket itself, so the next turn skips the handshake. + await this.#pool.withConnection( + async (ws) => { + if (ws.readyState !== WebSocket.OPEN) { + throw new APIConnectionError({ message: 'Cartesia pooled websocket is not open' }); + } + await Promise.all([inputTask(), sentenceStreamTask(ws), recvTask(ws)]); }, - timeoutMs: this.connOptions.timeoutMs, - abortSignal: this.abortSignal, - }); - await Promise.all([inputTask(), sentenceStreamTask(ws), recvTask(ws)]); + { timeout: this.connOptions.timeoutMs, signal: this.abortSignal }, + ); } catch (e) { if (this.abortSignal.aborted) { return; } if (e instanceof APIError) throw e; throw toRetryableConnectionError(e); - } finally { - // Ensure we don't leak sockets/tasks across retry attempts. - if (ws && ws.readyState !== WebSocket.CLOSED) { - safeTerminateWebSocket(ws); - } } } } @@ -631,9 +751,9 @@ const waitForWsOpen = async ({ }: { ws: WebSocket; timeoutMs: number; - abortSignal: AbortSignal; + abortSignal?: AbortSignal; }) => { - if (abortSignal.aborted) { + if (abortSignal?.aborted) { throw new Error('aborted'); } @@ -645,7 +765,7 @@ const waitForWsOpen = async ({ ws.off('open', onOpen); ws.off('error', onError); ws.off('close', onClose); - abortSignal.removeEventListener('abort', onAbort); + abortSignal?.removeEventListener('abort', onAbort); }; const onOpen = () => fut.resolve(); @@ -659,7 +779,7 @@ const waitForWsOpen = async ({ ws.on('open', onOpen); ws.on('error', onError); ws.on('close', onClose); - abortSignal.addEventListener('abort', onAbort, { once: true }); + abortSignal?.addEventListener('abort', onAbort, { once: true }); if (timeoutMs > 0) { timeout = setTimeout(() => fut.reject(new Error('connect timeout')), timeoutMs); @@ -693,6 +813,30 @@ const safeTerminateWebSocket = (ws: WebSocket) => { } }; +// Graceful close used by the connection pool. A pooled socket is healthy when it +// is retired (session age, option change, or TTS close), so a clean close frame +// is preferable to an abrupt terminate; terminate remains the fallback for a +// socket caught mid-handshake. +const safeCloseWebSocket = (ws: WebSocket) => { + try { + // `ws` can emit 'error' during teardown; without a listener Node treats it as + // unhandled and crashes the process. + ws.on('error', () => {}); + } catch { + // ignore + } + + try { + if (ws.readyState === WebSocket.CONNECTING || ws.readyState === WebSocket.OPEN) { + ws.close(); + } else if (ws.readyState !== WebSocket.CLOSED && ws.readyState !== WebSocket.CLOSING) { + ws.terminate(); + } + } catch { + // ignore + } +}; + const connectCartesiaWebSocket = async ({ url, headers, @@ -702,7 +846,7 @@ const connectCartesiaWebSocket = async ({ url: string; headers: Record; timeoutMs: number; - abortSignal: AbortSignal; + abortSignal?: AbortSignal; }): Promise => { const connectOnce = async (family?: number): Promise => { const ws = new WebSocket(url, { handshakeTimeout: timeoutMs, family, headers });