diff --git a/.changeset/fishaudio-pool-prebuffer.md b/.changeset/fishaudio-pool-prebuffer.md new file mode 100644 index 000000000..733f47944 --- /dev/null +++ b/.changeset/fishaudio-pool-prebuffer.md @@ -0,0 +1,5 @@ +--- +'@livekit/agents-plugin-fishaudio': patch +--- + +Reuse and prewarm Fish Audio streaming TTS websockets to avoid repeated connection setup. diff --git a/plugins/fishaudio/src/tts.test.ts b/plugins/fishaudio/src/tts.test.ts index 126d19972..179964cd1 100644 --- a/plugins/fishaudio/src/tts.test.ts +++ b/plugins/fishaudio/src/tts.test.ts @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: 2026 LiveKit, Inc. // // SPDX-License-Identifier: Apache-2.0 -import { tts } from '@livekit/agents'; +import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS, tts } from '@livekit/agents'; import { STT } from '@livekit/agents-plugin-openai'; import { tts as testTts } from '@livekit/agents-plugins-test'; import { decode, encode } from '@msgpack/msgpack'; @@ -60,6 +60,26 @@ async function startSynthesis(wss: WebSocketServer) { return { fishAudio, stream, ws: await waitFor(stopReceived) }; } +async function synthesizeTurn( + fishAudio: TTS, + text: string, + connOptions?: APIConnectOptions, +): Promise { + const stream = fishAudio.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(); + } +} + const hasFishAudioConfig = Boolean(process.env.FISH_API_KEY && process.env.OPENAI_API_KEY); if (hasFishAudioConfig) { @@ -73,6 +93,242 @@ if (hasFishAudioConfig) { } describe('FishAudio streaming', () => { + it('reuses one websocket across sequential turns', async () => { + const { wss, baseURL } = await startWebSocketServer(); + let connectionCount = 0; + + wss.on('connection', (ws) => { + connectionCount++; + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event === 'stop') { + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + } + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL }); + try { + expect(await synthesizeTurn(fishAudio, 'first turn.')).not.toHaveLength(0); + expect(await synthesizeTurn(fishAudio, 'second turn.')).not.toHaveLength(0); + expect(connectionCount).toBe(1); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('prewarms and reuses the ready websocket', async () => { + const { wss, baseURL } = await startWebSocketServer(); + let connectionCount = 0; + const connected = new Promise((resolve) => { + wss.on('connection', (ws) => { + connectionCount++; + resolve(); + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event === 'stop') { + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + } + }); + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL }); + try { + fishAudio.prewarm(); + await waitFor(connected); + expect(await synthesizeTurn(fishAudio, 'prewarmed turn.')).not.toHaveLength(0); + expect(connectionCount).toBe(1); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('opens a websocket with the updated model after model changes', async () => { + const { wss, baseURL } = await startWebSocketServer(); + const models: string[] = []; + + wss.on('connection', (ws, request) => { + const model = request.headers.model; + if (typeof model === 'string') models.push(model); + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event === 'stop') { + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + } + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL, model: 's2-pro' }); + try { + await synthesizeTurn(fishAudio, 'old model.'); + fishAudio.updateOptions({ model: 's2.1-pro' }); + await synthesizeTurn(fishAudio, 'new model.'); + expect(models).toEqual(['s2-pro', 's2.1-pro']); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('reconnects with the new model when an in-flight prewarm becomes stale', async () => { + let approveFirstUpgrade: (() => void) | undefined; + let firstUpgradeStarted: (() => void) | undefined; + const firstUpgrade = new Promise((resolve) => { + firstUpgradeStarted = resolve; + }); + let upgradeCount = 0; + const wss = new WebSocketServer({ + host: '127.0.0.1', + port: 0, + verifyClient: (_info, approve) => { + upgradeCount++; + if (upgradeCount === 1) { + approveFirstUpgrade = () => approve(true); + firstUpgradeStarted?.(); + return; + } + approve(true); + }, + }); + await once(wss, 'listening'); + const address = wss.address(); + if (address === null || typeof address === 'string') { + throw new Error('expected websocket server to listen on a TCP port'); + } + const models: string[] = []; + wss.on('connection', (ws, request) => { + const model = request.headers.model; + if (typeof model === 'string') models.push(model); + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event === 'stop') { + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + } + }); + }); + + const fishAudio = new TTS({ + apiKey: 'test-key', + baseURL: `http://127.0.0.1:${address.port}`, + model: 's2-pro', + }); + try { + fishAudio.prewarm(); + await waitFor(firstUpgrade); + fishAudio.updateOptions({ model: 's2.1-pro' }); + approveFirstUpgrade?.(); + expect(await synthesizeTurn(fishAudio, 'new model.')).not.toHaveLength(0); + expect(models).toEqual(['s2-pro', 's2.1-pro']); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('discards a failed websocket before the next turn', async () => { + const { wss, baseURL } = await startWebSocketServer(); + let connectionCount = 0; + + wss.on('connection', (ws) => { + connectionCount++; + const connectionNumber = connectionCount; + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event !== 'stop') return; + if (connectionNumber === 1) { + ws.close(1011, 'provider failure'); + return; + } + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL }); + try { + expect( + await synthesizeTurn(fishAudio, 'failing turn.', { + ...DEFAULT_API_CONNECT_OPTIONS, + maxRetry: 0, + }), + ).toHaveLength(0); + expect(await synthesizeTurn(fishAudio, 'recovery turn.')).not.toHaveLength(0); + expect(connectionCount).toBe(2); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('replaces a websocket that closed while idle', async () => { + const { wss, baseURL } = await startWebSocketServer(); + let connectionCount = 0; + let firstConnectionClosed: (() => void) | undefined; + const firstClosed = new Promise((resolve) => { + firstConnectionClosed = resolve; + }); + + wss.on('connection', (ws) => { + connectionCount++; + const connectionNumber = connectionCount; + ws.on('close', () => { + if (connectionNumber === 1) firstConnectionClosed?.(); + }); + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event !== 'stop') return; + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + if (connectionNumber === 1) setTimeout(() => ws.close(), 10); + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL }); + try { + expect(await synthesizeTurn(fishAudio, 'first turn.')).not.toHaveLength(0); + await waitFor(firstClosed); + expect(await waitFor(synthesizeTurn(fishAudio, 'second turn.'))).not.toHaveLength(0); + expect(connectionCount).toBe(2); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + + it('closes a pooled websocket when the TTS closes', async () => { + const { wss, baseURL } = await startWebSocketServer(); + const closed = new Promise((resolve) => { + wss.on('connection', (ws) => { + ws.on('close', () => resolve()); + ws.on('message', (raw) => { + const message = decode(Buffer.from(raw as ArrayBuffer)) as Record; + if (message.event === 'stop') { + ws.send(encode({ event: 'audio', audio: Buffer.alloc(4800) })); + ws.send(encode({ event: 'finish', reason: 'stop' })); + } + }); + }); + }); + + const fishAudio = new TTS({ apiKey: 'test-key', baseURL }); + try { + await synthesizeTurn(fishAudio, 'closing turn.'); + await fishAudio.close(); + await waitFor(closed); + expect(wss.clients.size).toBe(0); + } finally { + await fishAudio.close(); + await closeWebSocketServer(wss); + } + }); + it('emits the first complete frame before the next provider event', async () => { const { wss } = await startWebSocketServer(); const { fishAudio, stream, ws } = await startSynthesis(wss); diff --git a/plugins/fishaudio/src/tts.ts b/plugins/fishaudio/src/tts.ts index 1b47914bf..ec00798d5 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, @@ -24,6 +25,8 @@ const NUM_CHANNELS = 1; // Fish Audio's default sample rate for raw PCM output. const DEFAULT_SAMPLE_RATE = 24000; +const connectionPools = new WeakMap>(); + export interface TTSOptions { apiKey?: string; model?: TTSModels | string; @@ -113,6 +116,8 @@ const buildTtsRequest = (opts: ResolvedTTSOptions, text: string = ''): Record; + #closed = false; label = 'fishaudio.TTS'; constructor(opts: TTSOptions = {}) { @@ -148,6 +153,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, + }); + connectionPools.set(this, this.#pool); } get model(): string { @@ -166,7 +179,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) { @@ -188,6 +206,38 @@ 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 { + this.#closed = true; + await this.#pool.close(); + await super.close(); + } + + async #connectWebSocket(timeoutMs: number): Promise { + const wsUrl = `${this.#opts.baseURL.replace(/^http/, 'ws')}/v1/tts/live`; + const model = this.#opts.model; + const ws = await connectWebSocket({ + url: wsUrl, + headers: { + Authorization: `Bearer ${this.#opts.apiKey}`, + model, + }, + timeoutMs, + }); + if (this.#closed) { + closeWebSocket(ws); + throw new APIConnectionError({ message: 'Fish Audio TTS is closed' }); + } + if (model !== this.#opts.model) { + closeWebSocket(ws); + return await this.#connectWebSocket(timeoutMs); + } + return ws; + } } export class ChunkedStream extends tts.ChunkedStream { @@ -326,10 +376,14 @@ export class ChunkedStream extends tts.ChunkedStream { export class SynthesizeStream extends tts.SynthesizeStream { label = 'fishaudio.SynthesizeStream'; #logger = log(); + #pool: ConnectionPool; #opts: ResolvedTTSOptions; constructor(tts: TTS, opts: ResolvedTTSOptions, connOptions?: APIConnectOptions) { super(tts, connOptions); + const pool = connectionPools.get(tts); + if (!pool) throw new Error('Fish Audio connection pool is not initialized'); + this.#pool = pool; this.#opts = opts; } @@ -344,24 +398,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 () => { @@ -380,25 +416,25 @@ 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' }))); } }; - const recvTask = async () => { + 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) => { @@ -474,21 +510,31 @@ 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 this.#pool.withConnection( + async (ws) => { + if (ws.readyState !== WebSocket.OPEN) { + throw new APIConnectionError({ + message: 'Fish Audio pooled websocket is not open', + }); + } + 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) { @@ -499,13 +545,6 @@ export class SynthesizeStream extends tts.SynthesizeStream { }); } finally { if (!sentStream.closed) sentStream.close(); - if (ws && ws.readyState !== WebSocket.CLOSED) { - try { - ws.close(); - } catch { - // ignore - } - } } } } @@ -519,7 +558,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(); @@ -530,7 +569,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(); @@ -544,7 +583,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); @@ -555,12 +594,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 } @@ -569,3 +603,16 @@ const connectWebSocket = async ({ cleanup(); } }; + +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 + } +};