From 6c6c78c14005ec8a58035ba20eac9e5c1c19a111 Mon Sep 17 00:00:00 2001 From: Chase Fagen Date: Sat, 5 Sep 2026 05:17:44 -0700 Subject: [PATCH] perf(rime): stream sentence batches and reuse WebSocket connections --- .github/workflows/test.yml | 3 + plugins/rime/README.md | 37 +- plugins/rime/etc/agents-plugin-rime.api.md | 8 +- plugins/rime/src/tts.ts | 432 +++++++++++++++++- plugins/rime/src/tts.websocket.test.ts | 485 +++++++++++++++++++++ 5 files changed, 961 insertions(+), 4 deletions(-) create mode 100644 plugins/rime/src/tts.websocket.test.ts diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 81e57b35e3..d39c677e1f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -16,6 +16,7 @@ on: - 'agents/**' - 'package.json' - 'plugins/test/**' + - 'plugins/rime/**' - 'pnpm-lock.yaml' - 'pnpm-workspace.yaml' - 'turbo.json' @@ -48,6 +49,8 @@ jobs: run: pnpm build - name: Test agents run: pnpm test agents --silent + - name: Test Rime plugin + run: pnpm test plugins/rime --silent - name: Test LiveKit inference env: LIVEKIT_INFERENCE_TESTS: '1' diff --git a/plugins/rime/README.md b/plugins/rime/README.md index e9f7445477..330251f41f 100644 --- a/plugins/rime/README.md +++ b/plugins/rime/README.md @@ -3,6 +3,7 @@ SPDX-FileCopyrightText: 2024 LiveKit, Inc. SPDX-License-Identifier: Apache-2.0 --> + # Rime plugin for LiveKit Agents The Agents Framework is designed for building realtime, programmable @@ -13,4 +14,38 @@ This package contains the Rime plugin, which provides high-quality text-to-speec [documentation](https://docs.livekit.io/agents/overview/) for information on how to use it, or browse the [API reference](https://docs.livekit.io/agents-js/modules/plugins_agents_plugin_rime.html). See the [repository](https://github.com/livekit/agents-js) for more information -about the framework as a whole. \ No newline at end of file +about the framework as a whole. + +## WebSocket sentence streaming and reuse + +For `segment: 'never'`, sentence flushing and connection reuse are opt-in: + +```ts +const speech = new rime.TTS({ + modelId: 'coda', + speaker: 'lyra', + useWebsocket: true, + segment: 'never', + flushSentences: true, + reuseWebsocket: true, +}); +``` + +`flushSentences` sends each sentence emitted by the configured tokenizer and +waits for that synthesis batch's `done` event before sending the next sentence. +This prevents overlapping flushes from being combined by Rime. Explicit SDK +`stream.flush()` calls still finish separate audio segments and metrics; +whitespace-only segments are ignored. The +tokenizer retains control over when text is ready; this option does not bypass +its buffering or alter its sentence boundaries. + +`reuseWebsocket` retains at most one successfully completed connection per TTS +instance for up to 30 idle seconds. Simultaneous streams use separate sockets. +Interrupted or failed streams discard their connections; voice, language, or +other connection-option changes prevent reuse of the old connection. Call +`await speech.close()` during application shutdown to release active and idle +connections and cancel pending connections. + +Both options default to `false`; existing HTTP and WebSocket behavior is unchanged +unless enabled. See the [Rime WebSocket segmentation contract](https://docs.rime.ai/docs/websockets-segment) +for provider flush and completion semantics. diff --git a/plugins/rime/etc/agents-plugin-rime.api.md b/plugins/rime/etc/agents-plugin-rime.api.md index a6464641d5..67d2123d0c 100644 --- a/plugins/rime/etc/agents-plugin-rime.api.md +++ b/plugins/rime/etc/agents-plugin-rime.api.md @@ -32,8 +32,11 @@ export type DefaultLanguages = 'eng' | 'spa' | 'fra' | 'ger'; // @public (undocumented) export class SynthesizeStream extends tts.SynthesizeStream { constructor(tts: TTS, opts: TTSOptions, connOptions?: APIConnectOptions); + flush(): void; // (undocumented) label: string; + // @deprecated + pushText(text: string): void; // (undocumented) protected run(): Promise; } @@ -43,6 +46,7 @@ export class SynthesizeStream extends tts.SynthesizeStream { // @public (undocumented) export class TTS extends tts.TTS { constructor(opts?: Partial); + close(): Promise; // (undocumented) label: string; // (undocumented) @@ -60,7 +64,7 @@ export class TTS extends tts.TTS { // Warning: (ae-missing-release-tag) "TTSModels" is part of the package's API, but it is missing a release tag (@alpha, @beta, @public, or @internal) // // @public -export type TTSModels = 'arcana' | 'coda' | 'mistv2' | 'mistv3'; +export type TTSModels = 'coda' | 'mistv2' | 'mistv3'; // Warning: (ae-missing-release-tag) "TTSOptions" is part of the package's API, but it is missing a release tag (@alpha, @beta, @public, or @internal) // @@ -71,6 +75,7 @@ export interface TTSOptions { apiKey?: string; // (undocumented) baseURL?: string; + flushSentences?: boolean; // (undocumented) inlineSpeedAlpha?: string; // (undocumented) @@ -89,6 +94,7 @@ export interface TTSOptions { reduceLatency?: boolean; // (undocumented) repetition_penalty?: number; + reuseWebsocket?: boolean; // (undocumented) samplingRate?: number; // (undocumented) diff --git a/plugins/rime/src/tts.ts b/plugins/rime/src/tts.ts index 5fcb7e3e8f..14281adbc3 100644 --- a/plugins/rime/src/tts.ts +++ b/plugins/rime/src/tts.ts @@ -7,6 +7,7 @@ import { APIError, APIStatusError, APITimeoutError, + AsyncIterableQueue, AudioByteStream, Future, type TimedString, @@ -18,7 +19,7 @@ import { tokenize, tts, } from '@livekit/agents'; -import type { AudioFrame } from '@livekit/rtc-node'; +import { AudioFrame } from '@livekit/rtc-node'; import { type RawData, WebSocket } from 'ws'; import type { DefaultLanguages, TTSModels } from './models.js'; @@ -57,6 +58,22 @@ export interface TTSOptions { apiKey?: string; useWebsocket?: boolean; segment?: string; + /** + * Reuse a successfully completed WebSocket for later speech streams. + * Requires `useWebsocket: true` and `segment: 'never'`. Concurrent streams + * have separate connections; at most one idle connection is kept for 30 seconds. + * Call {@link TTS.close} when the provider is no longer needed. + * @defaultValue false + */ + reuseWebsocket?: boolean; + /** + * Flush each sentence emitted by the configured tokenizer to Rime without + * waiting for input to end. Requires `useWebsocket: true` and `segment: 'never'`. + * Synthesis batches run sequentially so Rime cannot coalesce overlapping flushes. + * Tokenizer buffering still determines when a sentence is available. + * @defaultValue false + */ + flushSentences?: boolean; tokenizer?: tokenize.SentenceTokenizer; lang?: DefaultLanguages | string; repetition_penalty?: number; @@ -139,6 +156,8 @@ function fetchPayload(opts: TTSOptions, text: string): Record { 'useWebsocket', 'segment', 'tokenizer', + 'reuseWebsocket', + 'flushSentences', 'speaker', 'modelId', 'lang', @@ -200,12 +219,28 @@ function resolveOptions(opts: Partial): TTSOptions { throw new Error('timeScaleFactor is not supported by the mistv2 model; use mistv3 or coda.'); } + if ( + (resolved.reuseWebsocket || resolved.flushSentences) && + (!resolved.useWebsocket || resolved.segment !== 'never') + ) { + throw new Error('Rime sentence flushing and connection reuse require WebSocket segment=never'); + } + return resolved; } +const connectionPools = new WeakMap(); + export class TTS extends tts.TTS { private opts: TTSOptions; label = 'rime.TTS'; + private connectionPool = new RimeConnectionPool(); + + /** Close active and idle WebSockets and cancel pending connections. */ + async close(): Promise { + this.connectionPool.close(); + await super.close(); + } /** * Create a new instance of Rime TTS. @@ -229,6 +264,7 @@ export class TTS extends tts.TTS { if (this.opts.apiKey === undefined) { throw new Error('RIME API key is required, whether as an argument or as $RIME_API_KEY'); } + connectionPools.set(this, this.connectionPool); warnIfArcana(opts.modelId); } @@ -247,7 +283,11 @@ export class TTS extends tts.TTS { */ updateOptions(opts: Partial) { warnIfArcana(opts.modelId); - this.opts = resolveOptions({ ...this.opts, ...opts }); + const updated = resolveOptions({ ...this.opts, ...opts }); + if (wsUrl(updated) !== wsUrl(this.opts) || updated.apiKey !== this.opts.apiKey) { + this.connectionPool.invalidate(); + } + this.opts = updated; } /** @@ -367,14 +407,49 @@ export class SynthesizeStream extends tts.SynthesizeStream { #opts: TTSOptions; #logger = log(); #tokenizer: tokenize.SentenceStream; + #provider: TTS; + #leadingWhitespace = ''; + #segmentHasText = false; + private connectionPool: RimeConnectionPool; constructor(tts: TTS, opts: TTSOptions, connOptions?: APIConnectOptions) { super(tts, connOptions); + this.#provider = tts; + const pool = connectionPools.get(tts); + if (!pool) throw new Error('Rime TTS connection pool is unavailable'); + this.connectionPool = pool; this.#opts = opts; this.#tokenizer = (opts.tokenizer ?? new tokenize.basic.SentenceTokenizer()).stream(); } + /** + * Buffer leading whitespace so empty SDK segments do not consume speech metrics. + * @deprecated Use `updateInputStream` instead. + */ + override pushText(text: string): void { + if (this.#opts.flushSentences || this.#opts.reuseWebsocket) { + if (!this.#segmentHasText) { + this.#leadingWhitespace += text; + if (!this.#leadingWhitespace.trim()) return; + text = this.#leadingWhitespace; + this.#leadingWhitespace = ''; + this.#segmentHasText = true; + } + } + super.pushText(text); + } + + /** Finish the current SDK segment, discarding whitespace-only input. */ + override flush(): void { + this.#leadingWhitespace = ''; + this.#segmentHasText = false; + super.flush(); + } + protected async run() { + if (this.#opts.flushSentences || this.#opts.reuseWebsocket) { + return this.runSentences(); + } const requestId = shortuuid(); const contextId = shortuuid(); const bstream = new AudioByteStream(getSampleRate(this.#opts), RIME_TTS_CHANNELS); @@ -549,6 +624,255 @@ export class SynthesizeStream extends tts.SynthesizeStream { } } } + // /ws3 can coalesce overlapping flushes. Wait for each matching done before + // sending the next sentence, and rebase synthesis-local timestamps onto PCM time. + private async runSentences() { + const requestId = shortuuid(); + let segmentId = ''; + const sampleRate = getSampleRate(this.#opts); + const segments = new AsyncIterableQueue(); + const tokenizers = new Set([this.#tokenizer]); + let inputTokenizer: tokenize.SentenceStream | undefined = this.#tokenizer; + segments.put(this.#tokenizer); + const messages = stream.createStreamChannel>(); + const reader = messages.stream().getReader(); + const failure = new Future(); + const fail = (message: string) => { + if (!failure.done) + failure.resolve(new APIConnectionError({ message, options: { retryable: false } })); + }; + let lease: Awaited> | undefined; + let complete = false; + let activeContext: string | undefined; + let totalSamples = 0; + let lastFrame: AudioFrame | undefined; + let pendingTranscripts: TimedString[] = []; + const emitFrame = (final: boolean) => { + if (!lastFrame || this.closed || this.abortSignal.aborted || this.queue.closed) return; + this.queue.put({ + requestId, + segmentId, + frame: lastFrame, + final, + timedTranscripts: pendingTranscripts.length ? pendingTranscripts : undefined, + }); + lastFrame = undefined; + pendingTranscripts = []; + }; + const emitAvailable = (frame: AudioFrame) => { + emitFrame(false); + const samples = frame.samplesPerChannel; + if (samples > 1) { + lastFrame = new AudioFrame( + frame.data.slice(0, -1), + sampleRate, + RIME_TTS_CHANNELS, + samples - 1, + ); + emitFrame(false); + } + // Keep one real sample for the SDK final frame, never a whole audio chunk. + lastFrame = new AudioFrame(frame.data.slice(-1), sampleRate, RIME_TTS_CHANNELS, 1); + }; + const onMessage = (raw: RawData) => { + try { + const data = JSON.parse(raw.toString()); + if (data.type === 'error') { + // Provider error text can echo input or credentials. Keep it out of logs. + fail('Rime WebSocket synthesis failed'); + } else if (activeContext && data.contextId === activeContext) { + void messages.write(data).catch(() => {}); + } + } catch { + fail('Rime WebSocket returned invalid JSON'); + } + }; + const onClose = () => fail('Rime WebSocket closed before synthesis completed'); + const onError = () => fail('Rime WebSocket transport failed'); + const onAbort = () => { + fail('Rime WebSocket synthesis cancelled'); + lease?.socket.terminate(); + }; + const inputTask = async () => { + for await (const data of this.input) { + if (this.abortSignal.aborted) break; + if (data === SynthesizeStream.FLUSH_SENTINEL) { + inputTokenizer?.endInput(); + inputTokenizer = undefined; + } else if (data) { + if (!inputTokenizer) { + inputTokenizer = ( + this.#opts.tokenizer ?? new tokenize.basic.SentenceTokenizer() + ).stream(); + tokenizers.add(inputTokenizer); + segments.put(inputTokenizer); + } + inputTokenizer.pushText(data); + } + } + inputTokenizer?.endInput(); + segments.close(); + }; + const synthesize = async (text: string) => { + if (!text.trim()) return; + if (this.abortSignal.aborted || lease!.socket.readyState !== WebSocket.OPEN) { + throw new APIConnectionError({ + message: 'Rime WebSocket connection is closed', + options: { retryable: false }, + }); + } + activeContext = shortuuid(); + const contextId = activeContext; + segmentId ||= contextId; + this.noteProviderRequestId(contextId); + const offset = totalSamples / sampleRate; + const samplesBeforeSynthesis = totalSamples; + const bytes = new AudioByteStream(sampleRate, RIME_TTS_CHANNELS); + this.markStarted(); + lease!.socket.send(JSON.stringify({ text, contextId })); + lease!.socket.send(JSON.stringify({ operation: 'flush', contextId })); + let timer: NodeJS.Timeout | undefined; + const armTimeout = () => { + if (timer) clearTimeout(timer); + if (this.connOptions.timeoutMs > 0) + timer = setTimeout( + () => fail('Rime WebSocket synthesis timed out'), + this.connOptions.timeoutMs, + ); + }; + armTimeout(); + try { + while (true) { + const event = await reader.read(); + if (event.done) + throw new APIConnectionError({ + message: 'Rime WebSocket stream ended early', + options: { retryable: false }, + }); + const data = event.value; + if (data.contextId !== contextId) continue; + armTimeout(); + if (data.type === 'chunk') { + const audio = Buffer.from(data.data as string, 'base64'); + totalSamples += audio.byteLength / 2; + for (const frame of bytes.write(audio)) emitAvailable(frame); + } else if (data.type === 'timestamps') { + const timestamps = data.word_timestamps as + | { words?: string[]; start?: number[]; end?: number[] } + | undefined; + if (timestamps?.words && timestamps.start && timestamps.end) { + const count = Math.min( + timestamps.words.length, + timestamps.start.length, + timestamps.end.length, + ); + for (let i = 0; i < count; i++) + pendingTranscripts.push( + createTimedString({ + text: `${timestamps.words[i]} `, + startTime: offset + timestamps.start[i]!, + endTime: offset + timestamps.end[i]!, + }), + ); + } + } else if (data.type === 'done') { + if (totalSamples === samplesBeforeSynthesis) { + throw new APIConnectionError({ + message: 'Rime WebSocket synthesis completed without audio', + options: { retryable: false }, + }); + } + for (const frame of bytes.flush()) { + if (frame.samplesPerChannel > 0) emitAvailable(frame); + } + activeContext = undefined; + return; + } + } + } finally { + if (timer) clearTimeout(timer); + } + }; + const finishSegment = async () => { + if (!lastFrame) return; + const emitted = new Future(); + const onMetrics = (event: { requestId: string }) => { + if (event.requestId === requestId && !emitted.done) emitted.resolve(); + }; + this.#provider.on('metrics_collected', onMetrics); + try { + emitFrame(true); + // markStarted is reset by the SDK metrics consumer, not queue.put. + // Do not let a queued segment inherit the previous segment's timing. + await Promise.race([ + emitted.await, + failure.await.then((error) => { + throw error; + }), + ]); + } finally { + this.#provider.off('metrics_collected', onMetrics); + } + }; + const synthesisTask = async () => { + for await (const tokenizer of segments) { + segmentId = ''; + totalSamples = 0; + let bufferedText = ''; + for await (const event of tokenizer) { + if (this.#opts.flushSentences) await synthesize(`${event.token} `); + else bufferedText += `${event.token} `; + } + if (!this.#opts.flushSentences) await synthesize(bufferedText); + await finishSegment(); + tokenizer.close(); + tokenizers.delete(tokenizer); + } + complete = true; + }; + this.abortSignal.addEventListener('abort', onAbort, { once: true }); + try { + lease = await this.connectionPool.acquire( + this.#opts, + this.connOptions.timeoutMs, + this.abortSignal, + ); + lease.socket.on('message', onMessage); + lease.socket.on('error', onError); + lease.socket.on('close', onClose); + await Promise.race([ + Promise.all([inputTask(), synthesisTask()]), + failure.await.then((error) => { + throw error; + }), + ]); + } catch (error) { + if (!this.abortSignal.aborted) { + if (error instanceof APIConnectionError && !error.retryable) throw error; + // Do not replay partial audio through the SDK retry loop. + throw new APIConnectionError({ + message: 'Rime WebSocket synthesis failed', + options: { retryable: false }, + }); + } + } finally { + this.abortSignal.removeEventListener('abort', onAbort); + for (const tokenizer of tokenizers) tokenizer.close(); + segments.close(); + if (!this.input.closed) this.input.close(); + await reader.cancel(); + reader.releaseLock(); + if (lease) { + lease.socket.off('message', onMessage); + lease.socket.off('error', onError); + lease.socket.off('close', onClose); + this.connectionPool.release( + lease, + Boolean(this.#opts.reuseWebsocket && complete && !this.abortSignal.aborted), + ); + } + } + } } async function connectRimeWebSocket({ @@ -616,3 +940,107 @@ function closeRimeWebSocket(ws: WebSocket) { // best-effort close } } + +// A socket is never shared by simultaneous contexts. Only fully drained leases return idle. +class RimeConnectionPool { + private sockets = new Set(); + private pending = new Set(); + private idle?: { socket: WebSocket; url: string; apiKey: string; timer: NodeJS.Timeout }; + private generation = 0; + private closed = false; + + async acquire(opts: TTSOptions, timeoutMs: number, signal: AbortSignal) { + if (this.closed || signal.aborted) + throw new APIConnectionError({ + message: 'Rime connection closed', + options: { retryable: false }, + }); + const url = wsUrl(opts); + const generation = this.generation; + if (this.idle) { + const idle = this.idle; + this.idle = undefined; + clearTimeout(idle.timer); + if ( + opts.reuseWebsocket && + idle.url === url && + idle.apiKey === opts.apiKey && + idle.socket.readyState === WebSocket.OPEN + ) { + return { socket: idle.socket, generation, url, apiKey: opts.apiKey! }; + } + idle.socket.terminate(); + } + const controller = new AbortController(); + const abort = () => controller.abort(); + signal.addEventListener('abort', abort, { once: true }); + this.pending.add(controller); + try { + const socket = await connectRimeWebSocket({ + url, + apiKey: opts.apiKey!, + timeoutMs, + abortSignal: controller.signal, + }); + // Own error events while idle and while a lease changes listeners. + socket.on('error', () => {}); + socket.once('close', () => { + this.sockets.delete(socket); + if (this.idle?.socket === socket) { + clearTimeout(this.idle.timer); + this.idle = undefined; + } + }); + this.sockets.add(socket); + if (this.closed || signal.aborted) { + socket.terminate(); + throw new APIConnectionError({ + message: 'Rime connection closed', + options: { retryable: false }, + }); + } + return { socket, generation, url, apiKey: opts.apiKey! }; + } finally { + this.pending.delete(controller); + signal.removeEventListener('abort', abort); + } + } + + release( + lease: { socket: WebSocket; generation: number; url: string; apiKey: string }, + reusable: boolean, + ) { + if ( + !reusable || + this.closed || + lease.generation !== this.generation || + this.idle || + lease.socket.readyState !== WebSocket.OPEN + ) { + lease.socket.terminate(); + return; + } + const timer = setTimeout(() => { + if (this.idle?.socket === lease.socket) this.idle = undefined; + lease.socket.terminate(); + }, 30_000); + timer.unref(); + this.idle = { ...lease, timer }; + } + + invalidate() { + this.generation += 1; + if (this.idle) { + clearTimeout(this.idle.timer); + this.idle.socket.terminate(); + this.idle = undefined; + } + } + + close() { + this.closed = true; + this.invalidate(); + for (const controller of this.pending) controller.abort(); + for (const socket of this.sockets) socket.terminate(); + } +} diff --git a/plugins/rime/src/tts.websocket.test.ts b/plugins/rime/src/tts.websocket.test.ts new file mode 100644 index 0000000000..087393bd5e --- /dev/null +++ b/plugins/rime/src/tts.websocket.test.ts @@ -0,0 +1,485 @@ +// SPDX-FileCopyrightText: 2026 LiveKit, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +import { AgentSession, initializeLogger, tts as livekitTts, voice } from '@livekit/agents'; +import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest'; +import { TTS } from './tts.js'; + +interface FakeSocket { + url: string; + readyState: number; + sent: Record[]; + terminate(): void; + receive(data: Record): void; +} + +const transport = vi.hoisted(() => ({ sockets: [] as FakeSocket[] })); + +vi.mock('ws', async () => { + const { EventEmitter } = await import('node:events'); + class FakeWebSocket extends EventEmitter { + static OPEN = 1; + static CLOSED = 3; + readyState = 0; + sent: Record[] = []; + constructor(readonly url: string) { + super(); + transport.sockets.push(this); + queueMicrotask(() => { + if (this.readyState !== 0) return; + this.readyState = 1; + this.emit('open'); + }); + } + send(value: string) { + this.sent.push(JSON.parse(value)); + } + terminate() { + this.readyState = 3; + this.emit('close', 1000, Buffer.alloc(0)); + } + close() { + this.terminate(); + } + receive(data: Record) { + this.emit('message', Buffer.from(JSON.stringify(data))); + } + } + return { WebSocket: FakeWebSocket }; +}); + +const providers: TTS[] = []; +beforeAll(() => initializeLogger({ level: 'silent', pretty: false })); +afterEach(async () => { + await Promise.all(providers.splice(0).map((provider) => provider.close())); + transport.sockets = []; + vi.useRealTimers(); +}); +function provider() { + const result = new TTS({ + modelId: 'coda', + speaker: 'lyra', + lang: 'eng', + samplingRate: 16000, + segment: 'never', + apiKey: 'fake-test-credential', + baseURL: 'ws://127.0.0.1:1', + reuseWebsocket: true, + flushSentences: true, + }); + providers.push(result); + result.on('error', () => {}); + return result; +} +function start(tts: TTS, text = 'Hello. ', end = true) { + const stream = tts.stream({ + connOptions: { timeoutMs: 1000, maxRetry: 2, retryIntervalMs: 0 }, + }); + const frames: livekitTts.SynthesizedAudio[] = []; + const done = (async () => { + for await (const frame of stream) { + if (frame !== livekitTts.SynthesizeStream.END_OF_STREAM) frames.push(frame); + } + })(); + stream.pushText(text); + if (end) stream.endInput(); + return { stream, frames, done }; +} +async function flushed(index = 0, count = 1) { + await vi.waitFor(() => { + expect(transport.sockets[index]?.sent.filter((m) => m.operation === 'flush')).toHaveLength( + count, + ); + }); + return transport.sockets[index]!; +} +function complete(socket: FakeSocket, sample = 10, samples = 320) { + const contextId = [...socket.sent].reverse().find((m) => m.text)?.contextId; + socket.receive({ + type: 'timestamps', + contextId, + word_timestamps: { words: ['Hello'], start: [0], end: [0.01] }, + }); + socket.receive({ + type: 'chunk', + contextId, + data: Buffer.from(new Int16Array(samples).fill(sample).buffer).toString('base64'), + }); + socket.receive({ type: 'done', contextId }); +} + +describe('Rime WebSocket transport', () => { + it('delivers first sentence audio before model EOF and rebases timestamps', async () => { + const tts = provider(); + const metrics: { charactersCount: number; audioDurationMs: number }[] = []; + tts.on('metrics_collected', (event) => metrics.push(event)); + const run = start(tts, 'I can help with that. Next', false); + const socket = await flushed(); + run.stream.pushText(' sentence. '); + expect(socket.sent.filter((m) => m.operation === 'flush')).toHaveLength(1); + complete(socket); + await vi.waitFor(() => expect(run.frames.length).toBeGreaterThan(0)); + run.stream.endInput(); + await flushed(0, 2); + complete(socket, 20); + await run.done; + expect( + run.frames.flatMap((frame) => frame.timedTranscripts ?? []).map((entry) => entry.startTime), + ).toEqual([0, 0.02]); + expect(new Set(run.frames.map((frame) => frame.segmentId)).size).toBe(1); + expect(run.frames.filter((frame) => frame.final)).toHaveLength(1); + expect(run.frames.at(-1)?.final).toBe(true); + expect(run.frames.flatMap((frame) => Array.from(frame.frame.data))).toEqual([ + ...Array(320).fill(10), + ...Array(320).fill(20), + ]); + expect(metrics).toHaveLength(1); + expect(metrics[0]?.audioDurationMs).toBeCloseTo(40); + expect(metrics[0]?.charactersCount).toBe('I can help with that. Next sentence. '.length); + }); + + it('serializes multiple ready sentences until each matching context completes', async () => { + const tts = provider(); + const sentences = [ + 'This is the first complete sentence.', + 'This is the second complete sentence.', + 'This is the third complete sentence.', + ]; + // EOF makes all three sentences available independently of tokenizer lookahead. + const run = start(tts, sentences.join(' ')); + const socket = await flushed(); + await new Promise((resolve) => setImmediate(resolve)); + expect(socket.sent.filter((message) => message.text).map((message) => message.text)).toEqual([ + `${sentences[0]} `, + ]); + const firstContext = socket.sent[0]!.contextId; + complete(socket, 10); + await flushed(0, 2); + await new Promise((resolve) => setImmediate(resolve)); + expect(socket.sent.filter((message) => message.text).map((message) => message.text)).toEqual([ + `${sentences[0]} `, + `${sentences[1]} `, + ]); + socket.receive({ type: 'done', contextId: firstContext }); + await new Promise((resolve) => setImmediate(resolve)); + expect(socket.sent.filter((message) => message.operation === 'flush')).toHaveLength(2); + complete(socket, 20); + await flushed(0, 3); + complete(socket, 30); + await run.done; + expect(socket.sent.filter((message) => message.text).map((message) => message.text)).toEqual( + sentences.map((sentence) => `${sentence} `), + ); + expect( + new Set(socket.sent.filter((message) => message.text).map((message) => message.contextId)) + .size, + ).toBe(3); + expect(run.frames.flatMap((frame) => Array.from(frame.frame.data))).toEqual([ + ...Array(320).fill(10), + ...Array(320).fill(20), + ...Array(320).fill(30), + ]); + expect(run.frames.filter((frame) => frame.final)).toHaveLength(1); + }); + + it('honors explicit SDK flush boundaries with separate final segments and metrics', async () => { + const tts = provider(); + const metrics: { charactersCount: number; audioDurationMs: number }[] = []; + tts.on('metrics_collected', (event) => metrics.push(event)); + const run = start(tts, 'First segment. ', false); + run.stream.flush(); + const socket = await flushed(); + complete(socket, 10); + await vi.waitFor(() => expect(run.frames.filter((frame) => frame.final)).toHaveLength(1)); + expect(metrics).toHaveLength(1); + run.stream.pushText('Second segment. '); + run.stream.endInput(); + await flushed(0, 2); + complete(socket, 20); + await run.done; + expect(new Set(run.frames.map((frame) => frame.segmentId)).size).toBe(2); + expect(run.frames.filter((frame) => frame.final)).toHaveLength(2); + expect(metrics.map((event) => event.charactersCount)).toEqual([ + 'First segment. '.length, + 'Second segment. '.length, + ]); + expect(metrics.map((event) => event.audioDurationMs)).toEqual([20, 20]); + }); + + it('records metrics for SDK segments queued before previous audio has drained', async () => { + const tts = provider(); + const metrics: { charactersCount: number; audioDurationMs: number }[] = []; + tts.on('metrics_collected', (event) => metrics.push(event)); + const run = start(tts, 'First segment. ', false); + run.stream.flush(); + run.stream.pushText('Second segment. '); + run.stream.endInput(); + const socket = await flushed(); + complete(socket, 10, 6400); + await flushed(0, 2); + complete(socket, 20, 6400); + await run.done; + expect(metrics.map((event) => event.charactersCount)).toEqual([ + 'First segment. '.length, + 'Second segment. '.length, + ]); + expect(metrics.map((event) => event.audioDurationMs)).toEqual([400, 400]); + }); + + it('does not attribute a spoken segment to earlier whitespace-only input', async () => { + const tts = provider(); + const metrics: { charactersCount: number; audioDurationMs: number }[] = []; + tts.on('metrics_collected', (event) => metrics.push(event)); + const run = start(tts, ' ', false); + run.stream.flush(); + run.stream.pushText('Hello. '); + run.stream.endInput(); + complete(await flushed()); + await run.done; + expect(metrics).toHaveLength(1); + expect(metrics[0]).toMatchObject({ charactersCount: 'Hello. '.length, audioDurationMs: 20 }); + expect(run.frames.filter((frame) => frame.final)).toHaveLength(1); + }); + + it('fails visibly when a nonempty synthesis completes without audio', async () => { + const tts = provider(); + const errors: Error[] = []; + tts.on('error', (event) => errors.push(event.error)); + const run = start(tts); + const socket = await flushed(); + socket.receive({ type: 'done', contextId: socket.sent[0]!.contextId }); + await run.done; + expect(run.frames).toHaveLength(0); + expect(errors).toHaveLength(1); + expect(errors[0]?.message).toBe('Rime WebSocket synthesis completed without audio'); + expect(socket.readyState).toBe(3); + }); + + it('records the actual Rime context IDs for provider trace correlation', async () => { + const run = start(provider(), 'This is the first sentence. This is the second sentence.'); + const noted = vi.spyOn( + run.stream as unknown as { noteProviderRequestId(id: string): void }, + 'noteProviderRequestId', + ); + const socket = await flushed(); + const firstContext = socket.sent[0]!.contextId; + complete(socket); + await flushed(0, 2); + const secondContext = [...socket.sent].reverse().find((message) => message.text)?.contextId; + complete(socket); + await run.done; + expect(firstContext).not.toBe(secondContext); + expect(run.frames[0]?.segmentId).toBe(firstContext); + expect(noted).toHaveBeenCalledWith(firstContext); + expect(noted).toHaveBeenCalledWith(secondContext); + }); + + it('reuses only a completed socket and ignores audio/done from an old context', async () => { + const tts = provider(); + const first = start(tts); + const socket = await flushed(); + const staleContext = socket.sent[0]!.contextId; + complete(socket); + await first.done; + const second = start(tts); + await flushed(0, 2); + socket.receive({ + type: 'chunk', + contextId: staleContext, + data: Buffer.from(new Int16Array(320).fill(99).buffer).toString('base64'), + }); + socket.receive({ type: 'done', contextId: staleContext }); + complete(socket, 20); + await second.done; + expect(transport.sockets).toHaveLength(1); + expect(second.frames.flatMap((frame) => Array.from(frame.frame.data))).not.toContain(99); + }); + + it('gives concurrent streams separate sockets and retains only one idle connection', async () => { + const tts = provider(); + const first = start(tts); + const second = start(tts); + const a = await flushed(0); + const b = await flushed(1); + complete(a, 10); + complete(b, 20); + await Promise.all([first.done, second.done]); + expect(transport.sockets.filter((socket) => socket.readyState === 1)).toHaveLength(1); + expect(first.frames[0]!.frame.data[0]).toBe(10); + expect(second.frames[0]!.frame.data[0]).toBe(20); + }); + + it('terminates an interrupted synthesis and cannot reuse or emit its late audio', async () => { + const tts = provider(); + const first = start(tts, 'This sentence is interrupted. Next', false); + const old = await flushed(); + first.stream.close(); + await first.done; + await vi.waitFor(() => expect(old.readyState).toBe(3)); + const second = start(tts); + const fresh = await flushed(1); + complete(old, 99); + complete(fresh, 20); + await second.done; + expect(first.frames).toHaveLength(0); + expect(second.frames.flatMap((frame) => Array.from(frame.frame.data))).not.toContain(99); + }); + + it('reconnects after an idle socket closes and after voice/language changes', async () => { + const tts = provider(); + const first = start(tts); + complete(await flushed()); + await first.done; + transport.sockets[0]!.terminate(); + const second = start(tts); + complete(await flushed(1)); + await second.done; + tts.updateOptions({ lang: 'spa', speaker: 'luz' }); + expect(transport.sockets[1]!.readyState).toBe(3); + const third = start(tts); + const socket = await flushed(2); + expect(new URL(socket.url).searchParams.get('lang')).toBe('spa'); + expect(new URL(socket.url).searchParams.get('speaker')).toBe('luz'); + expect(socket.url).not.toContain('fake-test-credential'); + complete(socket); + await third.done; + }); + + it('reports socket failure without replaying partial speech or exposing provider error text', async () => { + const tts = provider(); + const errors: Error[] = []; + tts.on('error', (event) => errors.push(event.error)); + const run = start(tts); + const socket = await flushed(); + socket.receive({ + type: 'error', + message: 'fake-test-credential raw patient input', + }); + await run.done; + expect(errors).toHaveLength(1); + expect(errors[0]?.message).toBe('Rime WebSocket synthesis failed'); + expect(transport.sockets).toHaveLength(1); + expect(socket.readyState).toBe(3); + const next = start(tts); + complete(await flushed(1)); + await next.done; + }); + + it('handles empty input and sub-frame PCM without waiting for another sentence', async () => { + const tts = provider(); + const empty = start(tts, ' '); + await empty.done; + expect(empty.frames).toHaveLength(0); + const short = start(tts); + complete(await flushed(), 7, 4); + await short.done; + expect(short.frames).toHaveLength(2); + expect(short.frames.reduce((sum, frame) => sum + frame.frame.samplesPerChannel, 0)).toBe(4); + expect(short.frames.at(-1)?.final).toBe(true); + expect(short.frames[0]?.frame.data[0]).toBe(7); + }); + + it('close releases active and idle sockets and prevents reopening', async () => { + const tts = provider(); + const idle = start(tts); + complete(await flushed()); + await idle.done; + const active = start(tts, 'This sentence is currently active. Next', false); + await flushed(0, 2); + await tts.close(); + await active.done; + expect(transport.sockets.every((socket) => socket.readyState === 3)).toBe(true); + const later = start(tts); + await later.done; + expect(transport.sockets).toHaveLength(1); + }); + it('the unchanged upstream path delays synthesis until EOF and opens another socket next turn', async () => { + const tts = provider(); + tts.updateOptions({ flushSentences: false, reuseWebsocket: false }); + const run = start(tts, 'This is the first sentence. Next', false); + await vi.waitFor(() => + expect(transport.sockets[0]?.sent.some((message) => message.text)).toBe(true), + ); + const socket = transport.sockets[0]!; + expect(socket.sent.filter((message) => message.operation === 'flush')).toHaveLength(0); + expect(run.frames).toHaveLength(0); + run.stream.endInput(); + await flushed(); + complete(socket); + await run.done; + expect(socket.readyState).toBe(3); + const next = start(tts); + complete(await flushed(1)); + await next.done; + expect(transport.sockets).toHaveLength(2); + }); + + it('expires a warm idle socket after 30 seconds', async () => { + const tts = provider(); + const run = start(tts); + const socket = await flushed(); + vi.useFakeTimers(); + complete(socket); + await run.done; + await vi.advanceTimersByTimeAsync(30_000); + expect(socket.readyState).toBe(3); + }); + + it('times out a stalled synthesis with a useful sanitized error and no retry', async () => { + const tts = provider(); + const errors: Error[] = []; + tts.on('error', (event) => errors.push(event.error)); + const run = start(tts); + const socket = await flushed(); + await run.done; + expect(errors[0]?.message).toBe('Rime WebSocket synthesis timed out'); + expect(socket.readyState).toBe(3); + expect(transport.sockets).toHaveLength(1); + }); + + it('does not retain an old active voice after options change', async () => { + const tts = provider(); + const old = start(tts); + const socket = await flushed(); + tts.updateOptions({ lang: 'spa', speaker: 'luz' }); + complete(socket); + await old.done; + expect(socket.readyState).toBe(3); + const next = start(tts); + complete(await flushed(1)); + await next.done; + }); + it('delivers PCM through the actual LiveKit default ttsNode before model EOF', async () => { + const tts = provider(); + const agent = new voice.Agent({ instructions: 'Test voice output.' }); + const session = new AgentSession({ tts }); + await session.start({ agent }); + let source!: ReadableStreamDefaultController; + const input = new ReadableStream({ + start(controller) { + source = controller; + }, + }); + const output = await voice.Agent.default.ttsNode(agent, input, {}); + const reader = output!.getReader(); + try { + source.enqueue('I can help with that. Next'); + const socket = await flushed(); + complete(socket, 7, 4); + const first = await reader.read(); + expect(first.done).toBe(false); + expect(first.value?.samplesPerChannel).toBe(3); + source.close(); + await flushed(0, 2); + complete(socket, 8, 4); + const last = await reader.read(); + expect(last.value?.samplesPerChannel).toBe(1); + expect((await reader.read()).value?.samplesPerChannel).toBe(3); + expect((await reader.read()).value?.samplesPerChannel).toBe(1); + expect((await reader.read()).done).toBe(true); + } finally { + reader.releaseLock(); + await session.close(); + } + }); +});