fix(inference): treat a dropped TTS gateway session as a failed attempt

.changeset/tts-gateway-session-drop.md

@@ -0,0 +1,19 @@
+---
+'@livekit/agents': patch
+---
+
+fix(inference): treat a dropped TTS gateway session as a failed attempt
+
+When the inference gateway ends a TTS session with `session.closed` part-way through a
+reply, the JS client treated it as a successful completion. That single mistake had two
+consequences. The rest of the reply's text was discarded with no error, no warning and no
+retry — in the trace this came from, a ~9000-character reply had only 2017 characters
+submitted before the drop. And the gateway websocket went back into the `ConnectionPool`
+while the gateway was still mid-synthesis, so the next reply picked it up and read the
+previous reply's outstanding audio as its own: after one barge-in the following reply spoke
+53.8s of the previous answer while the transcript showed the new one.
+
+The dropped session now flushes the audio it did produce, marks that segment's last frame
+final, and rejects with a retryable `APIStatusError`, so the socket is evicted and the retry
+finishes the reply. Socket reuse is additionally gated on having observed the gateway's
+`done`.

agents/src/inference/tts.ts

@@ -581,6 +581,14 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
   protected async run(): Promise<void> {
     let closing = false;
     let lastFrame: AudioFrame | undefined;
+    // Only a `done` from the gateway proves the session owes us no more audio, and a socket
+    // recycled before that hands the leftover audio to whichever SynthesizeStream picks it
+    // up next. `session.closed` is the exit that reaches the pool: it returns from this run
+    // normally, so nothing else evicts the socket. The remaining non-`done` exits — a closed
+    // event channel, a swallowed abort — are only ever reached after `onClose` / `onAbort`
+    // has already removed the socket, so gating reuse on `done` is what keeps reuse tied to
+    // the one event that proves the session is drained rather than to each exit remembering.
+    let sessionDrained = false;
     // Timestamps are delivered in their own WS message; buffer them and attach
     // to the next audio frame that we forward to the output emitter. This
     // mirrors the semantics of `output_emitter.push_timed_transcript` on the
@@ -631,13 +639,21 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
     };
 
     const createInputTask = async (signal: AbortSignal) => {
-      for await (const data of this.input) {
-        if (signal.aborted || closing) break;
-        if (data === SynthesizeStream.FLUSH_SENTINEL) {
+      while (!signal.aborted && !closing) {
+        // Read with the signal so a cancelled attempt stops waiting instead of taking —
+        // and dropping — the next chunk of text, which belongs to the retry.
+        const { done, value } = await this.input.next({ signal });
+        // `resourceCleanup` can land while that read is parked; it closes
+        // `sendTokenizerStream` synchronously, so pushing the chunk that just arrived would
+        // throw a plain `Error('Stream is closed')`. That is not an `APIError`, so
+        // `SynthesizeStream` would report the whole request as unrecoverable instead of
+        // retrying the attempt.
+        if (done || closing || signal.aborted) break;
+        if (value === SynthesizeStream.FLUSH_SENTINEL) {
           sendTokenizerStream.flush();
           continue;
         }
-        sendTokenizerStream.pushText(data);
+        sendTokenizerStream.pushText(value);
       }
       // Only call endInput if the stream hasn't been closed by cleanup
       if (!closing) {
@@ -830,12 +846,29 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
               }
               sendLastFrame(currentSessionId!, true);
               this.queue.put(SynthesizeStream.END_OF_STREAM);
+              sessionDrained = true;
               await resourceCleanup();
               completionFuture.resolve();
               return;
             case 'session.closed':
+              // The gateway dropped the session before it finished the reply. Hand over
+              // the audio it did produce, then fail the attempt: Python has no
+              // `session.closed` branch at all, so the dropped session surfaces there as
+              // a read timeout or a closed socket, i.e. an error that evicts the socket
+              // and lets the retry machinery resynthesize what is left. Resolving here
+              // instead reports a truncated reply as a completed one.
+              for (const frame of bstream.flush()) {
+                sendLastFrame(currentSessionId!, false);
+                lastFrame = frame;
+              }
+              sendLastFrame(currentSessionId!, true);
               await resourceCleanup();
-              completionFuture.resolve();
+              completionFuture.reject(
+                new APIStatusError({
+                  message: 'Gateway closed the TTS session before synthesis completed',
+                  options: { requestId },
+                }),
+              );
               return;
             case 'error':
               this.#logger.error(
@@ -920,6 +953,9 @@ export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeSt
               await resourceCleanup();
               await cancelAndWait(tasks, 5000);
               this.abortController.signal.removeEventListener('abort', onStreamAbort);
+              if (!sessionDrained) {
+                this.tts.pool.remove(ws);
+              }
             }
           } catch (e) {
             // If aborted, don't throw - let cleanup handle it

agents/src/inference/tts_pool_reuse.test.ts

@@ -0,0 +1,168 @@
+// SPDX-FileCopyrightText: 2026 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { AddressInfo } from 'node:net';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { WebSocketServer } from 'ws';
+import type { WebSocket as WsSocket } from 'ws';
+import { initializeLogger } from '../log.js';
+import { TTS } from './tts.js';
+
+/**
+ * Pins the user-visible invariant: a reply must never inherit the audio a previous, dropped
+ * session never finished delivering.
+ *
+ * `session.closed` is the only route that reaches that leak. Once the `session.closed`
+ * handler rejects the attempt instead of resolving it, the exception path in
+ * `ConnectionPool.withConnection` evicts the socket by itself and this test passes with or
+ * without the `sessionDrained` eviction in `SynthesizeStream.run`. Read it as coverage of
+ * the behaviour, not of that eviction.
+ */
+
+initializeLogger({ pretty: false });
+
+const SAMPLE_RATE = 16000;
+const FRAME_MS = 20;
+const SAMPLES_PER_FRAME = (SAMPLE_RATE * FRAME_MS) / 1000;
+
+/** Audio from the session that gets dropped and audio from a healthy session carry
+ *  distinct constant samples so the test can tell, per frame, which session's synthesis a
+ *  frame actually came from. */
+const DROPPED_SESSION_SAMPLE = 1000;
+const HEALTHY_SESSION_SAMPLE = 2000;
+
+/** ~6s of already-synthesized audio the gateway still owes after it drops the session. */
+const BACKLOG_FRAMES = 300;
+const OWN_FRAMES = 25;
+
+function audioEvent(sessionId: string, sample: number): string {
+  const pcm = Buffer.alloc(SAMPLES_PER_FRAME * 2);
+  for (let i = 0; i < SAMPLES_PER_FRAME; i++) {
+    pcm.writeInt16LE(sample, i * 2);
+  }
+  return JSON.stringify({
+    type: 'output_audio',
+    session_id: sessionId,
+    audio: pcm.toString('base64'),
+  });
+}
+
+/**
+ * Gateway stand-in for the production trace in which the first reply's session was dropped
+ * with `session.closed` while ~90s of synthesis was still outstanding. On the first
+ * connection it streams a short prefix, drops the session without a `done`, then keeps
+ * flushing the rest of the backlog onto the same socket. Any later connection behaves
+ * normally.
+ */
+async function startFakeGateway() {
+  const wss = new WebSocketServer({ port: 0, host: '127.0.0.1' });
+  await new Promise<void>((resolve) => wss.once('listening', () => resolve()));
+  let connections = 0;
+  const sockets: WsSocket[] = [];
+
+  wss.on('connection', (ws: WsSocket) => {
+    sockets.push(ws);
+    const index = ++connections;
+    const sessionId = `session-${index}`;
+    const sample = index === 1 ? DROPPED_SESSION_SAMPLE : HEALTHY_SESSION_SAMPLE;
+    let dropped = false;
+
+    const send = (payload: string) => {
+      if (ws.readyState === ws.OPEN) ws.send(payload);
+    };
+
+    ws.on('message', async (raw: Buffer) => {
+      const event = JSON.parse(raw.toString()) as { type: string };
+      if (event.type === 'session.create') {
+        send(JSON.stringify({ type: 'session.created', session_id: sessionId }));
+        return;
+      }
+      if (event.type !== 'session.flush') return;
+
+      // A healthy session serves every reply it is asked for, so a pooled socket can be
+      // reused across replies.
+      if (index > 1) {
+        for (let i = 0; i < OWN_FRAMES; i++) send(audioEvent(sessionId, sample));
+        send(JSON.stringify({ type: 'done', session_id: sessionId }));
+        return;
+      }
+
+      if (dropped) return;
+      dropped = true;
+
+      // Reply 1: hand over a short prefix, then drop the session mid-synthesis.
+      for (let i = 0; i < OWN_FRAMES; i++) send(audioEvent(sessionId, sample));
+      send(JSON.stringify({ type: 'session.closed', session_id: sessionId }));
+      // The synthesis that was already in flight keeps arriving on this socket.
+      await new Promise((resolve) => setTimeout(resolve, 20));
+      for (let i = 0; i < BACKLOG_FRAMES; i++) send(audioEvent(sessionId, sample));
+      send(JSON.stringify({ type: 'done', session_id: sessionId }));
+    });
+  });
+
+  const { port } = wss.address() as AddressInfo;
+  return {
+    baseURL: `http://127.0.0.1:${port}/v1`,
+    get connections() {
+      return connections;
+    },
+    close: () => {
+      for (const socket of sockets) socket.terminate();
+      return new Promise<void>((resolve) => wss.close(() => resolve()));
+    },
+  };
+}
+
+async function synthesize(tts: TTS<string>, text: string) {
+  const stream = tts.stream();
+  stream.pushText(text);
+  stream.endInput();
+
+  const samples = new Set<number>();
+  let frames = 0;
+  for await (const event of stream) {
+    if (typeof event === 'symbol' || event.frame.samplesPerChannel === 0) continue;
+    frames++;
+    samples.add(event.frame.data[0]!);
+  }
+  await stream.close();
+  return { frames, samples };
+}
+
+describe('inference TTS pooled socket reuse', () => {
+  let gateway: Awaited<ReturnType<typeof startFakeGateway>>;
+
+  beforeEach(async () => {
+    gateway = await startFakeGateway();
+  });
+
+  afterEach(async () => {
+    await gateway.close();
+  });
+
+  it('does not hand a dropped session\u2019s outstanding audio to the next reply', async () => {
+    const tts = new TTS({
+      model: 'inworld/inworld-tts-2',
+      voice: 'Sarah',
+      sampleRate: SAMPLE_RATE,
+      baseURL: gateway.baseURL,
+      apiKey: 'devkey',
+      apiSecret: 'secret'.padEnd(32, 'x'),
+    });
+
+    // The dropped session fails the attempt, so the first reply is the prefix it did
+    // deliver plus the audio from the retry that finishes it on a healthy session.
+    const first = await synthesize(tts, 'Tell me a long story about the lighthouse.');
+    expect(first.samples).toEqual(new Set([DROPPED_SESSION_SAMPLE, HEALTHY_SESSION_SAMPLE]));
+
+    const second = await synthesize(tts, 'Tell me a long joke about skeletons.');
+
+    // The second reply must speak only its own synthesis, and must not inherit the
+    // seconds of audio the first session never finished delivering.
+    expect(second.samples).toEqual(new Set([HEALTHY_SESSION_SAMPLE]));
+    expect(second.frames).toBeLessThan(BACKLOG_FRAMES);
+    expect(gateway.connections).toBe(2);
+
+    await tts.close();
+  }, 20_000);
+});