fishaudio: reuse streaming websocket and prebuffer initial chunks

.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.

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<tts.SynthesizedAudio[]> {
+  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<string, unknown>;
+        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<void>((resolve) => {
+      wss.on('connection', (ws) => {
+        connectionCount++;
+        resolve();
+        ws.on('message', (raw) => {
+          const message = decode(Buffer.from(raw as ArrayBuffer)) as Record<string, unknown>;
+          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<string, unknown>;
+        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<void>((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<string, unknown>;
+        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<string, unknown>;
+        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<void>((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<string, unknown>;
+        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<void>((resolve) => {
+      wss.on('connection', (ws) => {
+        ws.on('close', () => resolve());
+        ws.on('message', (raw) => {
+          const message = decode(Buffer.from(raw as ArrayBuffer)) as Record<string, unknown>;
+          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);