kugelaudio 0.8.0 → 0.9.0

package/src/livekit/tts.functional.test.ts

@@ -0,0 +1,348 @@
+/**
+ * Functional tests for the LiveKit plugin against a mock `/ws/tts/multi`
+ * server that speaks the real ingress wire protocol (see
+ * `backend/ingress/src/ingress/routes/ws_multi.py`):
+ *
+ *  - first message of a context carries `model_id` / `sample_rate` / config
+ *  - `flush: true` triggers audio frames + `chunk_complete`
+ *  - `{close_context}` is answered with `{context_closed}` (terminal)
+ *  - `immediate: true` close cancels without draining
+ *
+ * These tests exercise the actual WebSocket state machine end to end:
+ * chunked synthesis, streaming synthesis, error frames, barge-in, and
+ * connection reuse.
+ */
+
+import { createServer, type Server } from 'node:http';
+import { AddressInfo } from 'node:net';
+import { afterEach, beforeAll, describe, expect, it } from 'vitest';
+import { WebSocketServer, type WebSocket as ServerWebSocket } from 'ws';
+
+import { initializeLogger, type APIError } from '@livekit/agents';
+
+import { SynthesizeStream, TTS } from './tts';
+
+/** 100ms of 16-bit mono PCM at 24kHz. */
+const FRAME_BYTES = 4800;
+
+interface MockContextState {
+  text: string;
+  configMessages: Record<string, unknown>[];
+  closed: boolean;
+  immediateClose: boolean;
+}
+
+/**
+ * Minimal in-process implementation of the `/ws/tts/multi` protocol. Each
+ * flush emits `audioFramesPerFlush` PCM frames followed by `chunk_complete`;
+ * `close_context` drains (unless immediate) and answers `context_closed`.
+ */
+class MockMultiServer {
+  server: Server;
+  wss: WebSocketServer;
+  connectionCount = 0;
+  contexts = new Map<string, MockContextState>();
+  messages: Record<string, unknown>[] = [];
+  audioFramesPerFlush = 2;
+  /** When set, the first text message of a context is answered with this error frame. */
+  errorFrame: Record<string, unknown> | null = null;
+  /** When true, never respond to anything (for idle-timeout tests). */
+  silent = false;
+
+  constructor() {
+    this.server = createServer();
+    this.wss = new WebSocketServer({ server: this.server, path: '/ws/tts/multi' });
+    this.wss.on('connection', (ws) => {
+      this.connectionCount += 1;
+      ws.on('message', (raw) => this.onMessage(ws, JSON.parse(raw.toString())));
+    });
+  }
+
+  async listen(): Promise<string> {
+    await new Promise<void>((resolve) => this.server.listen(0, '127.0.0.1', resolve));
+    const { port } = this.server.address() as AddressInfo;
+    return `http://127.0.0.1:${port}`;
+  }
+
+  async close(): Promise<void> {
+    for (const client of this.wss.clients) client.terminate();
+    await new Promise<void>((resolve) => this.wss.close(() => resolve()));
+    await new Promise<void>((resolve) => this.server.close(() => resolve()));
+  }
+
+  private onMessage(ws: ServerWebSocket, msg: Record<string, unknown>): void {
+    this.messages.push(msg);
+    if (this.silent) return;
+
+    if (msg.close_socket) {
+      ws.send(JSON.stringify({ session_closed: true }));
+      ws.close();
+      return;
+    }
+
+    const contextId = msg.context_id as string;
+
+    if (msg.close_context) {
+      const ctx = this.contexts.get(contextId);
+      if (ctx) {
+        ctx.closed = true;
+        ctx.immediateClose = Boolean(msg.immediate);
+      }
+      ws.send(JSON.stringify({ context_closed: true, context_id: contextId }));
+      return;
+    }
+
+    let ctx = this.contexts.get(contextId);
+    if (!ctx) {
+      ctx = { text: '', configMessages: [], closed: false, immediateClose: false };
+      this.contexts.set(contextId, ctx);
+      ws.send(JSON.stringify({ context_created: true, context_id: contextId }));
+    }
+    if (msg.model_id !== undefined) ctx.configMessages.push(msg);
+    if (typeof msg.text === 'string') ctx.text += msg.text;
+
+    if (msg.flush) {
+      if (this.errorFrame) {
+        ws.send(JSON.stringify({ ...this.errorFrame, context_id: contextId }));
+        this.contexts.delete(contextId);
+        return;
+      }
+      if (ctx.text.length > 0) {
+        for (let i = 0; i < this.audioFramesPerFlush; i++) {
+          const pcm = Buffer.alloc(FRAME_BYTES, i + 1);
+          ws.send(
+            JSON.stringify({ audio: pcm.toString('base64'), context_id: contextId }),
+          );
+        }
+        ws.send(
+          JSON.stringify({
+            word_timestamps: [{ word: 'hello', start_ms: 0, end_ms: 500 }],
+            context_id: contextId,
+          }),
+        );
+        ws.send(JSON.stringify({ chunk_complete: true, context_id: contextId }));
+        ctx.text = '';
+      }
+    }
+  }
+}
+
+let server: MockMultiServer;
+let baseURL: string;
+const openTTS: TTS[] = [];
+
+beforeAll(() => {
+  initializeLogger({ pretty: false, level: 'silent' });
+});
+
+afterEach(async () => {
+  for (const t of openTTS.splice(0)) await t.close();
+  await server?.close();
+});
+
+async function setup(): Promise<TTS> {
+  server = new MockMultiServer();
+  baseURL = await server.listen();
+  const instance = new TTS({ apiKey: 'test-key', baseURL, language: 'en' });
+  openTTS.push(instance);
+  return instance;
+}
+
+describe('LiveKit plugin against mock /ws/tts/multi server', () => {
+  it('chunked synthesize returns all audio with the final flag on the last frame', async () => {
+    const instance = await setup();
+    const stream = instance.synthesize('Hallo Welt');
+
+    const events = [];
+    for await (const event of stream) events.push(event);
+
+    expect(events.length).toBeGreaterThan(0);
+    const totalSamples = events.reduce((sum, e) => sum + e.frame.samplesPerChannel, 0);
+    expect(totalSamples).toBe((2 * FRAME_BYTES) / 2);
+    expect(events[events.length - 1]!.final).toBe(true);
+    expect(events.slice(0, -1).every((e) => !e.final)).toBe(true);
+    expect(events.every((e) => e.frame.sampleRate === 24000)).toBe(true);
+
+    // Word timestamps ride on the frame emitted after they arrived.
+    const timed = events.flatMap((e) => e.timedTranscripts ?? []);
+    expect(timed).toHaveLength(1);
+    expect(timed[0]!.text).toBe('hello');
+    expect(timed[0]!.startTime).toBe(0);
+    expect(timed[0]!.endTime).toBe(0.5);
+
+    // First message of the context carried the session config exactly once.
+    const [ctx] = [...server.contexts.values()];
+    expect(ctx!.configMessages).toHaveLength(1);
+    expect(ctx!.configMessages[0]).toMatchObject({
+      model_id: 'kugel-3',
+      sample_rate: 24000,
+      language: 'en',
+      normalize: true,
+      word_timestamps: false,
+    });
+    expect(ctx!.closed).toBe(true);
+    expect(ctx!.immediateClose).toBe(false);
+  });
+
+  it('streaming synthesize forwards tokens and drains the tail gracefully', async () => {
+    const instance = await setup();
+    const stream = instance.stream();
+
+    stream.pushText('Guten ');
+    stream.pushText('Tag');
+    stream.endInput();
+
+    const events = [];
+    for await (const event of stream) {
+      if (event === SynthesizeStream.END_OF_STREAM) break;
+      events.push(event);
+    }
+
+    expect(events.length).toBeGreaterThan(0);
+    const totalSamples = events.reduce((sum, e) => sum + e.frame.samplesPerChannel, 0);
+    expect(totalSamples).toBe((2 * FRAME_BYTES) / 2);
+    expect(events[events.length - 1]!.final).toBe(true);
+
+    // The server received the concatenated text and a graceful close.
+    const [ctx] = [...server.contexts.values()];
+    expect(ctx!.closed).toBe(true);
+    expect(ctx!.immediateClose).toBe(false);
+    const textMessages = server.messages
+      .filter((m) => typeof m.text === 'string')
+      .map((m) => m.text);
+    expect(textMessages.join('')).toBe('Guten Tag');
+  });
+
+  it('reuses one WebSocket connection across sequential requests', async () => {
+    const instance = await setup();
+
+    for (const text of ['erste', 'zweite']) {
+      const stream = instance.synthesize(text);
+      for await (const _ of stream) {
+        // drain
+      }
+    }
+
+    expect(server.connectionCount).toBe(1);
+    // Each context got its own config-bearing first message.
+    expect(server.contexts.size).toBe(2);
+    for (const ctx of server.contexts.values()) {
+      expect(ctx.configMessages).toHaveLength(1);
+    }
+  });
+
+  it('maps server error frames to a non-retryable APIStatusError', async () => {
+    const instance = await setup();
+    server.errorFrame = {
+      error: 'text validation failed',
+      error_code: 'VALIDATION_ERROR',
+      code: 400,
+    };
+
+    const errors: Error[] = [];
+    instance.on('error', (event) => errors.push(event.error));
+
+    const stream = instance.stream();
+    stream.pushText('kaputt');
+    stream.endInput();
+
+    const events = [];
+    for await (const event of stream) {
+      if (event === SynthesizeStream.END_OF_STREAM) break;
+      events.push(event);
+    }
+
+    expect(events).toHaveLength(0);
+    expect(errors).toHaveLength(1);
+    expect(errors[0]!.message).toContain('text validation failed');
+    // 400-class errors must not be retried by the framework.
+    expect((errors[0] as APIError).retryable).toBe(false);
+  });
+
+  it('barge-in (stream.close) sends an immediate close_context', async () => {
+    const instance = await setup();
+    // Many frames so generation is "in flight" when we barge in.
+    server.audioFramesPerFlush = 50;
+
+    const stream = instance.stream();
+    stream.pushText('Eine sehr lange Ansage');
+    stream.flush();
+
+    // Consume a couple of frames, then barge in.
+    let received = 0;
+    for await (const event of stream) {
+      if (event === SynthesizeStream.END_OF_STREAM) break;
+      received += 1;
+      if (received >= 2) {
+        stream.close();
+        break;
+      }
+    }
+    expect(received).toBe(2);
+
+    // The immediate close reaches the server asynchronously.
+    await expect
+      .poll(() =>
+        server.messages.some((m) => m.close_context === true && m.immediate === true),
+      )
+      .toBe(true);
+  });
+
+  it('fails promptly on a server error even while the input channel stays open', async () => {
+    const instance = await setup();
+    server.errorFrame = {
+      error: 'model exploded',
+      error_code: 'INTERNAL_ERROR',
+      code: 500,
+    };
+
+    const errors: Error[] = [];
+    instance.on('error', (event) => errors.push(event.error));
+
+    const stream = instance.stream({
+      connOptions: { maxRetry: 0, retryIntervalMs: 10, timeoutMs: 5000 },
+    });
+    // Flush triggers the error frame, but the input channel is deliberately
+    // never ended — simulating an LLM that is still producing tokens. Without
+    // input-loop cancellation this would hang until the input closes.
+    stream.pushText('boom');
+    stream.flush();
+
+    const before = Date.now();
+    const events = [];
+    for await (const event of stream) {
+      if (event === SynthesizeStream.END_OF_STREAM) break;
+      events.push(event);
+    }
+
+    expect(Date.now() - before).toBeLessThan(2000);
+    expect(events).toHaveLength(0);
+    expect(errors).toHaveLength(1);
+    expect(errors[0]!.message).toContain('model exploded');
+    stream.close();
+  }, 10_000);
+
+  it('rejects the run with APITimeoutError when the server goes silent', async () => {
+    const instance = await setup();
+    server.silent = true;
+
+    const errors: Error[] = [];
+    instance.on('error', (event) => errors.push(event.error));
+
+    const stream = instance.stream({
+      connOptions: { maxRetry: 0, retryIntervalMs: 10, timeoutMs: 500 },
+    });
+    stream.pushText('niemand antwortet');
+    stream.endInput();
+
+    const events = [];
+    for await (const event of stream) {
+      if (event === SynthesizeStream.END_OF_STREAM) break;
+      events.push(event);
+    }
+
+    expect(events).toHaveLength(0);
+    expect(errors).toHaveLength(1);
+    expect(errors[0]!.name).toBe('APITimeoutError');
+  }, 10_000);
+});