diff --git a/.changeset/fishaudio-pool-prebuffer.md b/.changeset/fishaudio-pool-prebuffer.md new file mode 100644 index 000000000..2648d669a --- /dev/null +++ b/.changeset/fishaudio-pool-prebuffer.md @@ -0,0 +1,5 @@ +--- +'@livekit/agents-plugin-fishaudio': minor +--- + +Reuse streaming TTS websockets and prebuffer initial FishAudio chunks to avoid cold-start underruns. diff --git a/plugins/fishaudio/src/tts.ts b/plugins/fishaudio/src/tts.ts index 6bae25dd8..fad2198e4 100644 --- a/plugins/fishaudio/src/tts.ts +++ b/plugins/fishaudio/src/tts.ts @@ -6,6 +6,7 @@ import { APIConnectionError, APIStatusError, AudioByteStream, + ConnectionPool, Future, log, shortuuid, @@ -25,6 +26,12 @@ const NUM_CHANNELS = 1; // Fish Audio's default sample rate for raw PCM output. const DEFAULT_SAMPLE_RATE = 24000; +// Hold the first N audio chunks before releasing any audio, so playout starts with +// enough buffered to ride out the gap after Fish's small first chunk. Values <= 1 disable it. +const PREBUFFER_CHUNKS = 2; + +const ttsConnectionPools = new WeakMap>(); + export interface TTSOptions { apiKey?: string; model?: TTSModels | string; @@ -114,6 +121,7 @@ const buildTtsRequest = (opts: ResolvedTTSOptions, text: string = ''): Record; label = 'fishaudio.TTS'; constructor(opts: TTSOptions = {}) { @@ -149,6 +157,14 @@ export class TTS extends tts.TTS { volume: opts.volume, tokenizer, }; + + this.#pool = new ConnectionPool({ + connectCb: (timeoutMs) => this.#connectWebSocket(timeoutMs), + closeCb: async (ws) => closeWebSocket(ws), + maxSessionDuration: 300_000, + markRefreshedOnGet: true, + }); + ttsConnectionPools.set(this, this.#pool); } get model(): string { @@ -167,7 +183,12 @@ export class TTS extends tts.TTS { speed?: number; volume?: number; }): void { - if (opts.model !== undefined) this.#opts.model = opts.model; + if (opts.model !== undefined && opts.model !== this.#opts.model) { + this.#opts.model = opts.model; + // The model is sent as a connection header at ws-handshake time, not in the + // per-request body, so a pooled socket keeps the old model. + this.#pool.invalidate(); + } if (opts.voiceId !== undefined) this.#opts.voiceId = opts.voiceId; if (opts.latencyMode !== undefined) this.#opts.latencyMode = opts.latencyMode; if (opts.chunkLength !== undefined) { @@ -189,6 +210,27 @@ export class TTS extends tts.TTS { stream(options?: { connOptions?: APIConnectOptions }): tts.SynthesizeStream { return new SynthesizeStream(this, this.#opts, options?.connOptions); } + + prewarm(): void { + this.#pool.prewarm(); + } + + override async close(): Promise { + await this.#pool.close(); + await super.close(); + } + + async #connectWebSocket(timeoutMs: number): Promise { + const wsUrl = `${this.#opts.baseURL.replace(/^http/, 'ws')}/v1/tts/live`; + return await connectWebSocket({ + url: wsUrl, + headers: { + Authorization: `Bearer ${this.#opts.apiKey}`, + model: this.#opts.model, + }, + timeoutMs, + }); + } } export class ChunkedStream extends tts.ChunkedStream { @@ -327,10 +369,12 @@ export class ChunkedStream extends tts.ChunkedStream { export class SynthesizeStream extends tts.SynthesizeStream { label = 'fishaudio.SynthesizeStream'; #logger = log(); + #tts: TTS; #opts: ResolvedTTSOptions; constructor(tts: TTS, opts: ResolvedTTSOptions, connOptions?: APIConnectOptions) { super(tts, connOptions); + this.#tts = tts; this.#opts = opts; } @@ -345,24 +389,6 @@ export class SynthesizeStream extends tts.SynthesizeStream { // rather than mid-clause. const sentStream = this.#opts.tokenizer.stream(); - const wsUrl = `${this.#opts.baseURL.replace(/^http/, 'ws')}/v1/tts/live`; - let ws: WebSocket | undefined; - try { - ws = await connectWebSocket({ - url: wsUrl, - headers: { - Authorization: `Bearer ${this.#opts.apiKey}`, - model: this.#opts.model, - }, - timeoutMs: this.connOptions.timeoutMs, - abortSignal: this.abortSignal, - }); - } catch (e) { - throw new APIConnectionError({ - message: `Fish Audio websocket connect failed: ${(e as Error).message ?? 'unknown error'}`, - }); - } - const finished = new Future(); const inputTask = async () => { @@ -381,21 +407,21 @@ export class SynthesizeStream extends tts.SynthesizeStream { } }; - const sendTask = async () => { + const sendTask = async (ws: WebSocket) => { const startMsg = { event: 'start', request: buildTtsRequest(this.#opts) }; - ws!.send(Buffer.from(encode(startMsg))); + ws.send(Buffer.from(encode(startMsg))); for await (const ev of sentStream) { if (this.abortController.signal.aborted) break; const sentence = ev.token; if (!sentence) continue; this.markStarted(); - ws!.send(Buffer.from(encode({ event: 'text', text: sentence + ' ' }))); - ws!.send(Buffer.from(encode({ event: 'flush' }))); + ws.send(Buffer.from(encode({ event: 'text', text: sentence + ' ' }))); + ws.send(Buffer.from(encode({ event: 'flush' }))); } if (!this.abortController.signal.aborted) { - ws!.send(Buffer.from(encode({ event: 'stop' }))); + ws.send(Buffer.from(encode({ event: 'stop' }))); } }; @@ -407,7 +433,40 @@ export class SynthesizeStream extends tts.SynthesizeStream { } }; - const recvTask = async () => { + let prebuffering = PREBUFFER_CHUNKS > 1; + let chunksSeen = 0; + const prebuffer: Buffer[] = []; + + const writeAudio = (audio: Uint8Array) => { + for (const f of bstream.write(audio)) { + sendLastFrame(false); + lastFrame = f; + } + }; + + const pushAudio = (audio: Uint8Array) => { + if (prebuffering) { + prebuffer.push(Buffer.from(audio)); + chunksSeen += 1; + if (chunksSeen < PREBUFFER_CHUNKS) return; + writeAudio(Buffer.concat(prebuffer)); + prebuffer.length = 0; + prebuffering = false; + return; + } + writeAudio(audio); + }; + + const flushPrebuffer = () => { + // Stream ended before PREBUFFER_CHUNKS (short utterance); release held audio. + if (prebuffering && prebuffer.length > 0) { + writeAudio(Buffer.concat(prebuffer)); + prebuffer.length = 0; + } + prebuffering = false; + }; + + const recvTask = async (ws: WebSocket) => { // No per-receive timeout: Fish has natural inter-sentence gaps that can // exceed connOptions.timeoutMs when the LLM is slow. const onMessage = (raw: RawData) => { @@ -432,10 +491,7 @@ export class SynthesizeStream extends tts.SynthesizeStream { if (event === 'audio') { const audio = parsed.audio as Uint8Array | undefined; if (audio && audio.byteLength > 0) { - for (const f of bstream.write(audio)) { - sendLastFrame(false); - lastFrame = f; - } + pushAudio(audio); } } else if (event === 'finish') { const reason = parsed.reason as string | undefined; @@ -448,6 +504,7 @@ export class SynthesizeStream extends tts.SynthesizeStream { ); return; } + flushPrebuffer(); for (const f of bstream.flush()) { sendLastFrame(false); lastFrame = f; @@ -480,21 +537,27 @@ export class SynthesizeStream extends tts.SynthesizeStream { if (!finished.done) finished.reject(err); }; - ws!.on('message', onMessage); - ws!.on('close', onClose); - ws!.on('error', onError); + ws.on('message', onMessage); + ws.on('close', onClose); + ws.on('error', onError); try { await finished.await; } finally { - ws!.off('message', onMessage); - ws!.off('close', onClose); - ws!.off('error', onError); + ws.off('message', onMessage); + ws.off('close', onClose); + ws.off('error', onError); } }; try { - await Promise.all([inputTask(), sendTask(), recvTask()]); + await withPooledConnection( + this.#tts, + async (ws) => { + await Promise.all([inputTask(), sendTask(ws), recvTask(ws)]); + }, + { timeout: this.connOptions.timeoutMs, signal: this.abortSignal }, + ); } catch (e) { if (this.abortSignal.aborted) return; if (e instanceof APIStatusError || e instanceof APIConnectionError) { @@ -505,13 +568,6 @@ export class SynthesizeStream extends tts.SynthesizeStream { }); } finally { if (!sentStream.closed) sentStream.close(); - if (ws && ws.readyState !== WebSocket.CLOSED) { - try { - ws.close(); - } catch { - // ignore - } - } } } } @@ -525,7 +581,7 @@ const connectWebSocket = async ({ url: string; headers: Record; timeoutMs: number; - abortSignal: AbortSignal; + abortSignal?: AbortSignal; }): Promise => { const ws = new WebSocket(url, { headers, handshakeTimeout: timeoutMs }); const fut = new Future(); @@ -536,7 +592,7 @@ const connectWebSocket = 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(); @@ -550,7 +606,7 @@ const connectWebSocket = 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); @@ -561,12 +617,7 @@ const connectWebSocket = async ({ return ws; } catch (e) { try { - ws.on('error', () => {}); - if (ws.readyState === WebSocket.CONNECTING) { - ws.close(); - } else { - ws.terminate(); - } + closeWebSocket(ws); } catch { // ignore } @@ -575,3 +626,60 @@ const connectWebSocket = async ({ cleanup(); } }; + +const withPooledConnection = async ( + ttsInstance: TTS, + fn: (ws: WebSocket) => Promise, + options?: { timeout?: number; signal?: AbortSignal }, +): Promise => { + if (options?.signal?.aborted) throw abortError(); + + const pool = ttsConnectionPools.get(ttsInstance); + if (!pool) throw new Error('Fish Audio connection pool is not initialized'); + + const ws = await pool.get(options?.timeout); + let onAbort: (() => void) | undefined; + + try { + const result = options?.signal + ? await Promise.race([ + fn(ws), + new Promise((_, reject) => { + onAbort = () => reject(abortError()); + options.signal?.addEventListener('abort', onAbort, { once: true }); + }), + ]) + : await fn(ws); + + if (ws.readyState === WebSocket.OPEN) { + pool.put(ws); + } else { + pool.remove(ws); + } + return result; + } catch (error) { + pool.remove(ws); + throw error; + } finally { + if (onAbort) options?.signal?.removeEventListener('abort', onAbort); + } +}; + +const closeWebSocket = (ws: WebSocket) => { + try { + ws.on('error', () => {}); + 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 abortError = () => { + const error = new Error('The operation was aborted.'); + error.name = 'AbortError'; + return error; +};