@livekit/agents 1.1.0 → 1.2.1

package/src/inference/tts.ts

@@ -5,96 +5,208 @@ import type { AudioFrame } from '@livekit/rtc-node';
 import { WebSocket } from 'ws';
 import { APIError, APIStatusError } from '../_exceptions.js';
 import { AudioByteStream } from '../audio.js';
+import { ConnectionPool } from '../connection_pool.js';
+import { type LanguageCode, normalizeLanguage } from '../language.js';
 import { log } from '../log.js';
 import { createStreamChannel } from '../stream/stream_channel.js';
 import { basic as tokenizeBasic } from '../tokenize/index.js';
-import {
-  SynthesizeStream as BaseSynthesizeStream,
-  TTS as BaseTTS,
-  ChunkedStream,
-} from '../tts/index.js';
+import type { ChunkedStream } from '../tts/index.js';
+import { SynthesizeStream as BaseSynthesizeStream, TTS as BaseTTS } from '../tts/index.js';
 import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js';
-import { shortuuid } from '../utils.js';
+import { Event, Future, Task, cancelAndWait, combineSignals, shortuuid } from '../utils.js';
 import {
   type TtsClientEvent,
   type TtsServerEvent,
-  type TtsSessionCreateEvent,
   ttsClientEventSchema,
   ttsServerEventSchema,
 } from './api_protos.js';
-import { type AnyModels, connectWs, createAccessToken } from './utils.js';
+import { type AnyString, connectWs, createAccessToken, getDefaultInferenceUrl } from './utils.js';
 
-type _CartesiaModels = 'cartesia' | 'cartesia/sonic' | 'cartesia/sonic-2' | 'cartesia/sonic-turbo';
+export type CartesiaModels =
+  | 'cartesia/sonic-3'
+  | 'cartesia/sonic-2'
+  | 'cartesia/sonic-turbo'
+  | 'cartesia/sonic';
 
-export type CartesiaModels = _CartesiaModels | `${_CartesiaModels}:${string}`;
+export type DeepgramTTSModels = 'deepgram/aura' | 'deepgram/aura-2';
 
-type _ElevenlabsModels =
-  | 'elevenlabs'
+export type ElevenlabsModels =
   | 'elevenlabs/eleven_flash_v2'
   | 'elevenlabs/eleven_flash_v2_5'
   | 'elevenlabs/eleven_turbo_v2'
   | 'elevenlabs/eleven_turbo_v2_5'
   | 'elevenlabs/eleven_multilingual_v2';
 
-export type ElevenlabsModels = _ElevenlabsModels | `${_ElevenlabsModels}:${string}`;
+export type InworldModels =
+  | 'inworld/inworld-tts-1.5-max'
+  | 'inworld/inworld-tts-1.5-mini'
+  | 'inworld/inworld-tts-1-max'
+  | 'inworld/inworld-tts-1';
 
-export type _RimeModels = 'rime' | 'rime/mist' | 'rime/mistv2' | 'rime/arcana';
+export type RimeModels = 'rime/arcana' | 'rime/mistv2';
 
-export type RimeModels = _RimeModels | `${_RimeModels}:${string}`;
+export interface CartesiaOptions {
+  emotion?: string;
+  /** Maximum duration of audio in seconds. */
+  duration?: number;
+  /** Speech speed. Default: not specified. */
+  speed?: 'slow' | 'normal' | 'fast' | number;
+  volume?: number;
+  max_buffer_delay_ms?: number;
+  add_timestamps?: boolean;
+  add_phoneme_timestamps?: boolean;
+  use_normalized_timestamps?: boolean;
+}
 
-export type _InworldModels = 'inworld' | 'inworld/inworld-tts-1';
+export interface ElevenlabsOptions {
+  /** Inactivity timeout in seconds. Default: 60. */
+  inactivity_timeout?: number;
+  /** Text normalization mode. Default: "auto". */
+  apply_text_normalization?: 'auto' | 'off' | 'on';
+  auto_mode?: boolean;
+  enable_logging?: boolean;
+  enable_ssml_parsing?: boolean;
+  sync_alignment?: boolean;
+  language_code?: string;
+  /** Voice stability tuning, typically in the range [0, 1]. */
+  stability?: number;
+  /** Voice similarity tuning, typically in the range [0, 1]. */
+  similarity_boost?: number;
+  /** Style exaggeration tuning, typically in the range [0, 1]. */
+  style?: number;
+  /** Speech speed multiplier. */
+  speed?: number;
+  use_speaker_boost?: boolean;
+  chunk_length_schedule?: number[];
+  preferred_alignment?: string;
+}
 
-export type InworldModels = _InworldModels | `${_InworldModels}:${string}`;
+export interface DeepgramTTSOptions {
+  /** Default: false. */
+  mip_opt_out?: boolean;
+}
 
-export interface CartesiaOptions {
-  duration?: number; // max duration of audio in seconds
-  speed?: 'slow' | 'normal' | 'fast'; // default: not specified
+export interface RimeOptions {
+  /** Default 1.0, <1 = faster, >1 = slower. */
+  speed_alpha?: number;
+  /** Default false. */
+  pause_between_brackets?: boolean;
+  /** Default false. */
+  phonemize_between_brackets?: boolean;
+  /** Comma-separated speed factors for [bracketed] words. */
+  inline_speed_alpha?: string;
+  /** Default false. */
+  no_text_normalization?: boolean;
 }
 
-export interface ElevenlabsOptions {
-  inactivity_timeout?: number; // default: 60
-  apply_text_normalization?: 'auto' | 'off' | 'on'; // default: "auto"
+export interface InworldOptions {
+  /** Range >0.5, <=1.5. */
+  speaking_rate?: number;
+  /** Range 0-2. */
+  temperature?: number;
+  timestamp_type?: 'TIMESTAMP_TYPE_UNSPECIFIED' | 'WORD' | 'CHARACTER';
+  apply_text_normalization?: 'APPLY_TEXT_NORMALIZATION_UNSPECIFIED' | 'ON' | 'OFF';
+  /** @deprecated Backward-compatible alias. Use `apply_text_normalization`. */
+  text_normalization?: 'ON' | 'OFF';
 }
 
-export interface RimeOptions {}
+type _TTSModels =
+  | CartesiaModels
+  | DeepgramTTSModels
+  | ElevenlabsModels
+  | RimeModels
+  | InworldModels;
 
-export interface InworldOptions {}
+export type TTSModels =
+  | CartesiaModels
+  | DeepgramTTSModels
+  | ElevenlabsModels
+  | RimeModels
+  | InworldModels
+  | AnyString;
 
-export type TTSModels = CartesiaModels | ElevenlabsModels | RimeModels | InworldModels | AnyModels;
+export type ModelWithVoice = `${_TTSModels}:${string}` | TTSModels;
 
 export type TTSOptions<TModel extends TTSModels> = TModel extends CartesiaModels
   ? CartesiaOptions
-  : TModel extends ElevenlabsModels
-    ? ElevenlabsOptions
-    : TModel extends RimeOptions
-      ? RimeOptions
-      : TModel extends InworldOptions
-        ? InworldOptions
-        : Record<string, unknown>;
+  : TModel extends DeepgramTTSModels
+    ? DeepgramTTSOptions
+    : TModel extends ElevenlabsModels
+      ? ElevenlabsOptions
+      : TModel extends RimeModels
+        ? RimeOptions
+        : TModel extends InworldModels
+          ? InworldOptions
+          : Record<string, unknown>;
+
+/** Parse a model string into [model, voice]. Voice is undefined if not specified. */
+export function parseTTSModelString(model: string): [string, string | undefined] {
+  const idx = model.lastIndexOf(':');
+  if (idx !== -1) {
+    return [model.slice(0, idx), model.slice(idx + 1)];
+  }
+  return [model, undefined];
+}
+
+/** A fallback model with optional extra configuration. Extra fields are passed through to the provider. */
+export interface TTSFallbackModel {
+  /** Model name (e.g. "cartesia/sonic", "elevenlabs/eleven_flash_v2", "rime/arcana"). */
+  model: string;
+  /** Voice to use for the model. */
+  voice: string;
+  /** Extra configuration for the model. */
+  extraKwargs?: Record<string, unknown>;
+}
+
+export type TTSFallbackModelType = TTSFallbackModel | string;
+
+/** Normalize a single or list of FallbackModelType into TTSFallbackModel[]. */
+export function normalizeTTSFallback(
+  fallback: TTSFallbackModelType | TTSFallbackModelType[],
+): TTSFallbackModel[] {
+  const makeFallback = (model: TTSFallbackModelType): TTSFallbackModel => {
+    if (typeof model === 'string') {
+      const [name, voice] = parseTTSModelString(model);
+      return { model: name, voice: voice ?? '' };
+    }
+    return model;
+  };
+
+  if (Array.isArray(fallback)) {
+    return fallback.map(makeFallback);
+  }
+  return [makeFallback(fallback)];
+}
 
 type TTSEncoding = 'pcm_s16le';
 
 const DEFAULT_ENCODING: TTSEncoding = 'pcm_s16le';
 const DEFAULT_SAMPLE_RATE = 16000;
-const DEFAULT_BASE_URL = 'https://agent-gateway.livekit.cloud/v1';
 const NUM_CHANNELS = 1;
 const DEFAULT_LANGUAGE = 'en';
 
 export interface InferenceTTSOptions<TModel extends TTSModels> {
   model?: TModel;
   voice?: string;
-  language?: string;
+  language?: LanguageCode;
   encoding: TTSEncoding;
   sampleRate: number;
   baseURL: string;
   apiKey: string;
   apiSecret: string;
-  extraKwargs: TTSOptions<TModel>;
+  /** Flat provider-specific inference options forwarded as the `extra` payload field. */
+  modelOptions: TTSOptions<TModel>;
+  fallback?: TTSFallbackModel[];
+  connOptions?: APIConnectOptions;
 }
 
+/**
+ * Livekit Cloud Inference TTS
+ */
 export class TTS<TModel extends TTSModels> extends BaseTTS {
   private opts: InferenceTTSOptions<TModel>;
   private streams: Set<SynthesizeStream<TModel>> = new Set();
+  pool: ConnectionPool<WebSocket>;
 
   #logger = log();
 
@@ -107,7 +219,10 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
     sampleRate?: number;
     apiKey?: string;
     apiSecret?: string;
-    extraKwargs?: TTSOptions<TModel>;
+    /** Flat provider-specific inference options forwarded as the `extra` payload field. */
+    modelOptions?: TTSOptions<TModel>;
+    fallback?: TTSFallbackModelType | TTSFallbackModelType[];
+    connOptions?: APIConnectOptions;
   }) {
     const sampleRate = opts?.sampleRate ?? DEFAULT_SAMPLE_RATE;
     super(sampleRate, 1, { streaming: true });
@@ -120,10 +235,12 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
       encoding = DEFAULT_ENCODING,
       apiKey,
       apiSecret,
-      extraKwargs = {} as TTSOptions<TModel>,
+      modelOptions = {} as TTSOptions<TModel>,
+      fallback,
+      connOptions,
     } = opts || {};
 
-    const lkBaseURL = baseURL || process.env.LIVEKIT_INFERENCE_URL || DEFAULT_BASE_URL;
+    const lkBaseURL = baseURL || getDefaultInferenceUrl();
     const lkApiKey = apiKey || process.env.LIVEKIT_INFERENCE_API_KEY || process.env.LIVEKIT_API_KEY;
     if (!lkApiKey) {
       throw new Error('apiKey is required: pass apiKey or set LIVEKIT_API_KEY');
@@ -154,27 +271,57 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
       }
     }
 
+    const normalizedFallback = fallback ? normalizeTTSFallback(fallback) : undefined;
+
     this.opts = {
       model: nextModel,
       voice: nextVoice,
-      language,
+      language: normalizeLanguage(language),
       encoding,
       sampleRate,
       baseURL: lkBaseURL,
       apiKey: lkApiKey,
       apiSecret: lkApiSecret,
-      extraKwargs,
+      modelOptions,
+      fallback: normalizedFallback,
+      connOptions: connOptions ?? DEFAULT_API_CONNECT_OPTIONS,
     };
+
+    // Initialize connection pool
+    this.pool = new ConnectionPool<WebSocket>({
+      connectCb: (timeout) => this.connectWs(timeout),
+      closeCb: (ws) => this.closeWs(ws),
+      maxSessionDuration: 300_000,
+      markRefreshedOnGet: true,
+      connectTimeout: 10_000, // 10 seconds default
+    });
   }
 
   get label() {
     return 'inference.TTS';
   }
 
+  get model(): string {
+    return this.opts.model ?? 'unknown';
+  }
+
+  get provider(): string {
+    return 'livekit';
+  }
+
+  static fromModelString(modelString: string): TTS<AnyString> {
+    const [model, voice] = parseTTSModelString(modelString);
+    return new TTS({ model, voice: voice || undefined });
+  }
+
   updateOptions(opts: Partial<Pick<InferenceTTSOptions<TModel>, 'model' | 'voice' | 'language'>>) {
-    this.opts = { ...this.opts, ...opts };
+    this.opts = {
+      ...this.opts,
+      ...opts,
+      language: opts.language !== undefined ? normalizeLanguage(opts.language) : this.opts.language,
+    };
     for (const stream of this.streams) {
-      stream.updateOptions(opts);
+      stream.updateOptions(this.opts);
     }
   }
 
@@ -183,7 +330,7 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
   }
 
   stream(options?: { connOptions?: APIConnectOptions }): SynthesizeStream<TModel> {
-    const { connOptions = DEFAULT_API_CONNECT_OPTIONS } = options || {};
+    const { connOptions = this.opts.connOptions ?? DEFAULT_API_CONNECT_OPTIONS } = options || {};
     const stream = new SynthesizeStream(this, { ...this.opts }, connOptions);
     this.streams.add(stream);
     return stream;
@@ -203,13 +350,31 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
       type: 'session.create',
       sample_rate: String(this.opts.sampleRate),
       encoding: this.opts.encoding,
-      extra: this.opts.extraKwargs,
-    } as TtsSessionCreateEvent;
+      extra: this.opts.modelOptions,
+    } as Record<string, unknown>;
+
+    if (this.opts.voice) (params as Record<string, unknown>).voice = this.opts.voice;
+    if (this.opts.model) (params as Record<string, unknown>).model = this.opts.model;
+    if (this.opts.language) (params as Record<string, unknown>).language = this.opts.language;
+
+    if (this.opts.fallback?.length) {
+      params.fallback = {
+        models: this.opts.fallback.map((m) => ({
+          model: m.model,
+          voice: m.voice,
+          extra: m.extraKwargs ?? {},
+        })),
+      };
+    }
 
-    if (this.opts.voice) params.voice = this.opts.voice;
-    if (this.opts.model) params.model = this.opts.model;
-    if (this.opts.language) params.language = this.opts.language;
+    if (this.opts.connOptions) {
+      params.connection = {
+        timeout: this.opts.connOptions.timeoutMs / 1000,
+        retries: this.opts.connOptions.maxRetry,
+      };
+    }
 
+    this.#logger.debug({ url }, 'inference.TTS creating new websocket connection (pool miss)');
     const socket = await connectWs(url, headers, timeout);
     socket.send(JSON.stringify(params));
     return socket;
@@ -219,18 +384,22 @@ export class TTS<TModel extends TTSModels> extends BaseTTS {
     await ws.close();
   }
 
+  prewarm(): void {
+    this.pool.prewarm();
+  }
+
   async close() {
     for (const stream of this.streams) {
       await stream.close();
     }
     this.streams.clear();
+    await this.pool.close();
   }
 }
 
 export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeStream {
   private opts: InferenceTTSOptions<TModel>;
   private tts: TTS<TModel>;
-  private connOptions: APIConnectOptions;
 
   #logger = log();
 
@@ -238,7 +407,6 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
     super(tts, connOptions);
     this.opts = opts;
     this.tts = tts;
-    this.connOptions = connOptions;
   }
 
   get label() {
@@ -246,31 +414,39 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
   }
 
   updateOptions(opts: Partial<Pick<InferenceTTSOptions<TModel>, 'model' | 'voice' | 'language'>>) {
-    this.opts = { ...this.opts, ...opts };
+    this.opts = {
+      ...this.opts,
+      ...opts,
+      language: opts.language !== undefined ? normalizeLanguage(opts.language) : this.opts.language,
+    };
   }
 
   protected async run(): Promise<void> {
-    let ws: WebSocket | null = null;
     let closing = false;
-    let finalReceived = false;
     let lastFrame: AudioFrame | undefined;
 
     const sendTokenizerStream = new tokenizeBasic.SentenceTokenizer().stream();
     const eventChannel = createStreamChannel<TtsServerEvent>();
     const requestId = shortuuid('tts_request_');
+    const inputSentEvent = new Event();
 
-    const resourceCleanup = () => {
+    // Signal for protocol-driven completion (when 'done' message is received)
+    const completionFuture = new Future<void>();
+
+    const resourceCleanup = async () => {
       if (closing) return;
       closing = true;
       sendTokenizerStream.close();
-      eventChannel.close();
-      ws?.removeAllListeners();
-      ws?.close();
+      // close() returns a promise; don't leak it
+      await eventChannel.close();
     };
 
-    const sendClientEvent = async (event: TtsClientEvent) => {
+    const sendClientEvent = async (event: TtsClientEvent, ws: WebSocket, signal: AbortSignal) => {
+      // Don't send events to a closed WebSocket or aborted controller
+      if (signal.aborted || closing) return;
+
       const validatedEvent = await ttsClientEventSchema.parseAsync(event);
-      if (!ws || ws.readyState !== WebSocket.OPEN) {
+      if (ws.readyState !== WebSocket.OPEN) {
         this.#logger.warn('Trying to send client TTS event to a closed WebSocket');
         return;
       }
@@ -284,67 +460,123 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
       }
     };
 
-    const createInputTask = async () => {
+    const createInputTask = async (signal: AbortSignal) => {
       for await (const data of this.input) {
-        if (this.abortController.signal.aborted) break;
+        if (signal.aborted || closing) break;
         if (data === SynthesizeStream.FLUSH_SENTINEL) {
           sendTokenizerStream.flush();
           continue;
         }
         sendTokenizerStream.pushText(data);
       }
-      sendTokenizerStream.endInput();
+      // Only call endInput if the stream hasn't been closed by cleanup
+      if (!closing) {
+        sendTokenizerStream.endInput();
+      }
     };
 
-    const createSentenceStreamTask = async () => {
+    const createSentenceStreamTask = async (ws: WebSocket, signal: AbortSignal) => {
       for await (const ev of sendTokenizerStream) {
-        if (this.abortController.signal.aborted) break;
-
-        sendClientEvent({
-          type: 'input_transcript',
-          transcript: ev.token + ' ',
-        });
+        if (signal.aborted || closing) break;
+
+        await sendClientEvent(
+          {
+            type: 'input_transcript',
+            transcript: ev.token + ' ',
+          },
+          ws,
+          signal,
+        );
+        inputSentEvent.set();
       }
 
-      sendClientEvent({ type: 'session.flush' });
+      await sendClientEvent({ type: 'session.flush' }, ws, signal);
+      // needed in case empty input is sent
+      inputSentEvent.set();
     };
 
-    const createWsListenerTask = async (ws: WebSocket) => {
-      return new Promise<void>((resolve, reject) => {
-        this.abortController.signal.addEventListener('abort', () => {
-          resourceCleanup();
-          reject(new Error('WebSocket connection aborted'));
-        });
-
-        ws.on('message', async (data) => {
+    // Handles WebSocket message routing and error handling
+    // Completes based on protocol messages, NOT on ws.close()
+    const createWsListenerTask = async (ws: WebSocket, signal: AbortSignal) => {
+      const onMessage = (data: Buffer) => {
+        try {
           const eventJson = JSON.parse(data.toString()) as Record<string, unknown>;
           const validatedEvent = ttsServerEventSchema.parse(eventJson);
-          eventChannel.write(validatedEvent);
-        });
-
-        ws.on('error', (e) => {
-          this.#logger.error({ error: e }, 'WebSocket error');
-          resourceCleanup();
-          reject(e);
-        });
-
-        ws.on('close', () => {
-          resourceCleanup();
-
-          if (!closing) return this.#logger.error('WebSocket closed unexpectedly');
-          if (finalReceived) return resolve();
+          // writer.write returns a promise; avoid unhandled rejections if stream is closed
+          void eventChannel.write(validatedEvent).catch((error) => {
+            this.#logger.debug(
+              { error },
+              'Failed writing TTS event to stream channel (likely closed)',
+            );
+          });
+        } catch (e) {
+          this.#logger.error({ error: e }, 'Error parsing WebSocket message');
+        }
+      };
 
-          reject(
+      const onError = (e: Error) => {
+        this.#logger.error({ error: e }, 'WebSocket error');
+        void resourceCleanup();
+        try {
+          // If the ws is misbehaving, hard-stop it immediately to avoid buffering.
+          ws.terminate?.();
+        } catch {
+          // ignore
+        }
+        // Ensure this ws is not reused
+        this.tts.pool.remove(ws);
+        completionFuture.reject(e);
+      };
+
+      const onClose = () => {
+        // WebSocket closed unexpectedly (not by us)
+        if (!closing) {
+          this.#logger.error('WebSocket closed unexpectedly');
+          void resourceCleanup();
+          // Ensure this ws is not reused
+          this.tts.pool.remove(ws);
+          completionFuture.reject(
             new APIStatusError({
               message: 'Gateway connection closed unexpectedly',
               options: { requestId },
             }),
           );
-        });
-      });
+        }
+      };
+
+      const onAbort = () => {
+        void resourceCleanup();
+        try {
+          // On interruption/abort, close the websocket immediately so the server stops streaming
+          // and the ws library doesn't buffer unread frames in memory.
+          ws.terminate?.();
+        } catch {
+          // ignore
+        }
+        this.tts.pool.remove(ws);
+        inputSentEvent.set();
+        completionFuture.resolve();
+      };
+
+      // Attach listeners
+      ws.on('message', onMessage);
+      ws.on('error', onError);
+      ws.on('close', onClose);
+      signal.addEventListener('abort', onAbort);
+
+      try {
+        // Wait for protocol-driven completion or error
+        await completionFuture.await;
+      } finally {
+        // IMPORTANT: Remove listeners so connection can be reused
+        ws.off('message', onMessage);
+        ws.off('error', onError);
+        ws.off('close', onClose);
+        signal.removeEventListener('abort', onAbort);
+      }
     };
 
-    const createRecvTask = async () => {
+    const createRecvTask = async (signal: AbortSignal) => {
       let currentSessionId: string | null = null;
 
       const bstream = new AudioByteStream(this.opts.sampleRate, NUM_CHANNELS);
@@ -352,9 +584,11 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
       const reader = serverEventStream.getReader();
 
       try {
-        while (!this.closed && !this.abortController.signal.aborted) {
+        await inputSentEvent.wait();
+
+        while (!this.closed && !signal.aborted) {
           const result = await reader.read();
-          if (this.abortController.signal.aborted) return;
+          if (signal.aborted) return;
           if (result.done) return;
 
           const serverEvent = result.value;
@@ -370,24 +604,29 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
               }
               break;
             case 'done':
-              finalReceived = true;
               for (const frame of bstream.flush()) {
                 sendLastFrame(currentSessionId!, false);
                 lastFrame = frame;
               }
               sendLastFrame(currentSessionId!, true);
               this.queue.put(SynthesizeStream.END_OF_STREAM);
-              break;
+              await resourceCleanup();
+              completionFuture.resolve();
+              return;
             case 'session.closed':
-              resourceCleanup();
-              break;
+              await resourceCleanup();
+              completionFuture.resolve();
+              return;
             case 'error':
               this.#logger.error(
                 { serverEvent },
                 'Received error message from LiveKit TTS WebSocket',
               );
-              resourceCleanup();
-              throw new APIError(`LiveKit TTS returned error: ${serverEvent.message}`);
+              await resourceCleanup();
+              completionFuture.reject(
+                new APIError(`LiveKit TTS returned error: ${serverEvent.message}`),
+              );
+              return;
             default:
               this.#logger.warn('Unexpected message %s', serverEvent);
               break;
@@ -404,18 +643,81 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
     };
 
     try {
-      ws = await this.tts.connectWs(this.connOptions.timeoutMs);
-
-      await Promise.all([
-        createInputTask(),
-        createSentenceStreamTask(),
-        createWsListenerTask(ws),
-        createRecvTask(),
-      ]);
+      await this.tts.pool.withConnection(
+        async (ws: WebSocket) => {
+          try {
+            // IMPORTANT: don't cancel the stream's controller on normal completion,
+            // otherwise the pool will remove+close the ws and every run becomes a pool miss.
+            const runController = new AbortController();
+            const onStreamAbort = () => runController.abort(this.abortController.signal.reason);
+            this.abortController.signal.addEventListener('abort', onStreamAbort, { once: true });
+
+            const tasks = [
+              Task.from(
+                async (controller) => {
+                  const combined = combineSignals(runController.signal, controller.signal);
+                  await createInputTask(combined);
+                },
+                undefined,
+                'inference-tts-input',
+              ),
+              Task.from(
+                async (controller) => {
+                  const combined = combineSignals(runController.signal, controller.signal);
+                  await createSentenceStreamTask(ws, combined);
+                },
+                undefined,
+                'inference-tts-sentence',
+              ),
+              Task.from(
+                async (controller) => {
+                  const combined = combineSignals(runController.signal, controller.signal);
+                  await createWsListenerTask(ws, combined);
+                },
+                undefined,
+                'inference-tts-ws-listener',
+              ),
+              Task.from(
+                async (controller) => {
+                  const combined = combineSignals(runController.signal, controller.signal);
+                  await createRecvTask(combined);
+                },
+                undefined,
+                'inference-tts-recv',
+              ),
+            ];
+
+            try {
+              await Promise.all(tasks.map((t) => t.result));
+            } finally {
+              // Mirror python finally: unblock recv and cancel all tasks.
+              inputSentEvent.set();
+              await resourceCleanup();
+              await cancelAndWait(tasks, 5000);
+              this.abortController.signal.removeEventListener('abort', onStreamAbort);
+            }
+          } catch (e) {
+            // If aborted, don't throw - let cleanup handle it
+            if (e instanceof Error && e.name === 'AbortError') {
+              return;
+            }
+            throw e;
+          }
+        },
+        {
+          timeout: this.connOptions.timeoutMs,
+        },
+      );
     } catch (e) {
-      this.#logger.error('Error in SynthesizeStream', { error: e });
+      // Handle connection errors
+      if (e instanceof Error && e.name === 'AbortError') {
+        // Abort is expected during normal shutdown
+        return;
+      }
+      throw e;
     } finally {
-      resourceCleanup();
+      // Ensure cleanup always runs (and don't leak the promise)
+      await resourceCleanup();
     }
   }
 }