pi-asr/src/client.ts

/**
 * NVIDIA Speech NIM realtime ASR client (Riva realtime WebSocket API).
 * Zero dependencies, Node >= 22 (global WebSocket).
 *
 * Docs: https://docs.nvidia.com/nim/speech/latest/reference/api-references/asr/realtime-asr.html
 *
 * Verified against the live deployment (behind Traefik `replacePathRegex`
 * ^/openai/streaming-asr/(.*) -> /v1/$1 plus the `qwen-audio-api-key` middleware):
 *
 *   1. GET  {base}/health/live | /health/ready            Bearer key   -> {status:'live'|'ready'}
 *   2. POST {base}/realtime/transcription_sessions        Bearer key   -> session + client_secret.value
 *   3. WS   {base}/realtime?intent=transcription
 *           Sec-WebSocket-Protocol: "realtime, realtime-token.<value>"   (never a URL query param)
 *   4. transcription_session.update  -> transcription_session.updated
 *   5. input_audio_buffer.append (base64 pcm16) * N
 *   6. input_audio_buffer.commit     -> input_audio_buffer.committed   <-- REQUIRED, see below
 *   7. server emits conversation.item.input_audio_transcription.delta|.completed|.failed
 *
 * Three behaviours differ from the published docs and are handled here:
 *   a) `input_audio_buffer.done` alone never flushes the decoder -> zero results. An explicit
 *      `input_audio_buffer.commit` is what triggers recognition.
 *   b) `delta` is CUMULATIVE (each frame repeats the whole utterance so far), not incremental.
 *   c) `is_last_result` never arrives as true, so completion is detected by event quiescence.
 *      And the server sends two `.completed` frames per segment: the first carries `words_info`,
 *      the second has empty `words_info` -> the word-bearing one must be kept.
 *   d) Auth failures arrive as an HTTP **403 during the upgrade handshake**, not as WebSocket
 *      close code 1008 like the docs claim. A tokenless handshake returns 101 with NO echoed
 *      subprotocol, which Node's global WebSocket (undici) treats as a fatal error
 *      ("Server did not respond with sent protocols") -> the minted token is mandatory in Node.
 */

import { spawn } from "node:child_process";
import { readFileSync } from "node:fs";
import { readFile } from "node:fs/promises";
import { homedir } from "node:os";
import { extname } from "node:path";

export const EV = {
	conversationCreated: "conversation.created",
	sessionUpdated: "transcription_session.updated",
	bufferCommitted: "input_audio_buffer.committed",
	bufferCleared: "input_audio_buffer.cleared",
	delta: "conversation.item.input_audio_transcription.delta",
	completed: "conversation.item.input_audio_transcription.completed",
	failed: "conversation.item.input_audio_transcription.failed",
	error: "error",
} as const;

export interface AsrConfig {
	baseUrl: string;
	apiKey: string;
	language: string;
	model: string | null;
	sampleRate: number;
	channels: number;
	punctuation: boolean;
	wordOffsets: boolean;
}

export const DEFAULT_BASE_URL = "https://ai.iaminsane.online/openai/streaming-asr";

export interface AsrFileConfig { baseUrl?: string; apiKey?: string }

/** Optional on-disk config so the key never has to live in source: ~/.pi/agent/asr.json (or $ASR_CONFIG). */
export function configFilePath(): string {
	return process.env.ASR_CONFIG ?? [homedir(), ".pi", "agent", "asr.json"].join("/");
}

let cachedFileConfig: AsrFileConfig | null = null;
export function readConfigFile(): AsrFileConfig {
	if (cachedFileConfig) return cachedFileConfig;
	try {
		cachedFileConfig = JSON.parse(readFileSync(configFilePath(), "utf-8")) as AsrFileConfig;
	} catch {
		cachedFileConfig = {};
	}
	return cachedFileConfig;
}

/** Env wins, so the key can live outside the file: ASR_API_KEY / ASR_NIM_URL. */
export function defaultConfig(overrides: Partial<AsrConfig> = {}): AsrConfig {
	return {
		baseUrl: (overrides.baseUrl ?? process.env.ASR_NIM_URL ?? readConfigFile().baseUrl ?? DEFAULT_BASE_URL).replace(/\/+$/, ""),
		apiKey: overrides.apiKey ?? process.env.ASR_API_KEY ?? readConfigFile().apiKey ?? "",
		language: overrides.language ?? "en-US",
		model: overrides.model ?? null, // null => whatever the NIM advertises on mint
		sampleRate: overrides.sampleRate ?? 16000,
		channels: overrides.channels ?? 1,
		punctuation: overrides.punctuation ?? true,
		wordOffsets: overrides.wordOffsets ?? false,
	};
}

const uid = (p: string = "event"): string => `${p}_${crypto.randomUUID()}`;
const sleep = (ms: number): Promise<void> => new Promise(r => setTimeout(r, ms));

export interface WordInfo {
	word: string;
	start_time: number;
	end_time: number;
	confidence?: number;
	speaker_tag?: number;
}

export interface TranscriptResult {
	text: string;
	segments: { transcript: string; isLast: boolean }[];
	words: WordInfo[];
	failed: unknown[];
	seconds: number;
	model: string | null;
	sessionId: string;
	protocol: string | null;
}

type Listener = (ev: any) => void;

export class RealtimeAsrClient {
	cfg: AsrConfig;
	verbose: boolean;
	private listeners: Map<string, Set<Listener>> = new Map();

	constructor(opts: Partial<AsrConfig> & { verbose?: boolean } = {}) {
		this.cfg = defaultConfig(opts);
		this.verbose = opts.verbose ?? false;
	}

	log(...a: unknown[]): void {
		if (this.verbose) console.error("[asr]", ...a);
	}

	on(type: string, fn: Listener): () => void {
		if (!this.listeners.has(type)) this.listeners.set(type, new Set());
		this.listeners.get(type)!.add(fn);
		return () => this.listeners.get(type)?.delete(fn);
	}

	private emit(ev: any): void {
		for (const fn of this.listeners.get(ev?.type) ?? []) { try { fn(ev); } catch {} }
		for (const fn of this.listeners.get("*") ?? []) { try { fn(ev); } catch {} }
	}

	private headers(extra: Record<string, string> = {}): Record<string, string> {
		const h: Record<string, string> = { ...extra };
		// Gateway middleware `qwen-audio-api-key` -> 403 {"Invalid or missing API Key"} without it.
		if (this.cfg.apiKey) h.Authorization = `Bearer ${this.cfg.apiKey}`;
		return h;
	}

	get wsUrl(): string {
		return `${this.cfg.baseUrl.replace(/^http/, "ws")}/realtime?intent=transcription`;
	}

	/** The gateway answers 403 {Invalid or missing API Key} without a key, which is cryptic; fail with instructions instead. */
	private requireKey(): void {
		if (this.cfg.apiKey) return;
		throw new Error(
			"No ASR API key configured. Set ASR_API_KEY, or write { \"apiKey\": \"sk-...\" } to " + configFilePath() +
			" (or point ASR_CONFIG at your file). Base URL: " + this.cfg.baseUrl,
		);
	}

	/** GET {base}/health/live | /health/ready */
	async health(which: "live" | "ready" = "live", timeout: number = 15000): Promise<any> {
		this.requireKey();
		const r = await fetch(`${this.cfg.baseUrl}/health/${which}`, {
			headers: this.headers(),
			signal: AbortSignal.timeout?.(timeout),
		});
		const t = await r.text();
		if (!r.ok) throw new Error(`health/${which} -> HTTP ${r.status}: ${t.slice(0, 200)}`);
		try { return JSON.parse(t); } catch { return { raw: t }; }
	}

	/** POST {base}/realtime/transcription_sessions — single-use token, 60s TTL, so mint right before connecting. */
	async mint(timeout: number = 15000): Promise<any> {
		this.requireKey();
		const r = await fetch(`${this.cfg.baseUrl}/realtime/transcription_sessions`, {
			method: "POST",
			headers: this.headers({ "Content-Type": "application/json" }),
			signal: AbortSignal.timeout?.(timeout),
		});
		const t = await r.text();
		if (!r.ok) {
			const why = r.status === 403 || r.status === 401 ? "gateway rejected the API key" : `HTTP ${r.status}`;
			throw new Error(`mint session failed: ${why}: ${t.slice(0, 200)}`);
		}
		return JSON.parse(t);
	}

	/**
	 * Payload for transcription_session.update. The server substitutes *defaults* for anything
	 * omitted (model silently becomes "conformer", which is not loaded here and yields zero
	 * results), so echo back the minted session and apply overrides on top.
	 */
	sessionConfig(minted: any = {}, overrides: Record<string, any> = {}): Record<string, any> {
		const base: Record<string, any> = {
			modalities: minted.modalities ?? ["text"],
			input_audio_format: "pcm16",
			input_audio_transcription: {
				language: this.cfg.language,
				model: this.cfg.model ?? minted.input_audio_transcription?.model,
				prompt: minted.input_audio_transcription?.prompt ?? "",
			},
			input_audio_params: {
				sample_rate_hz: this.cfg.sampleRate,
				num_channels: this.cfg.channels,
			},
			recognition_config: {
				...(minted.recognition_config ?? {}),
				max_alternatives: 1,
				enable_automatic_punctuation: !!this.cfg.punctuation,
				enable_word_time_offsets: !!this.cfg.wordOffsets,
				enable_profanity_filter: false,
				enable_verbatim_transcripts: false,
			},
			speaker_diarization: minted.speaker_diarization ?? { enable_speaker_diarization: false, max_speaker_count: 8 },
			word_boosting: minted.word_boosting ?? { enable_word_boosting: false, word_boosting_list: [] },
			endpointing_config: minted.endpointing_config ?? {
				start_history: 0, start_threshold: 0, stop_history: 0, stop_threshold: 0,
				stop_history_eou: 0, stop_threshold_eou: 0,
			},
		};
		return { ...base, ...overrides };
	}

	/** Open the socket. `onDelta` gets every cumulative partial, `onFinal` each completed segment. */
	async open(opts: { token?: string; timeout?: number; onDelta?: (text: string) => void; onFinal?: (ev: any) => void } = {}): Promise<SocketHandle> {
		const token = opts.token ?? null;
		const protocols = token ? ["realtime", `realtime-token.${token}`] : ["realtime"];
		const ws = new WebSocket(this.wsUrl, { protocols, headers: this.headers() });
		const st: SocketHandle = { events: [], closed: null, protocol: null };

		await new Promise<void>((res, rej) => {
			const t = setTimeout(() => { try { ws.close(); } catch {} rej(new Error("WebSocket handshake timed out")); }, opts.timeout ?? 20000);
			ws.onopen = () => { clearTimeout(t); st.protocol = ws.protocol; res(); };
			ws.onerror = (e: any) => {
				clearTimeout(t);
				const msg = e?.message || e?.error?.message || e?.type || "unknown";
				rej(new Error(msg.includes("did not respond with sent protocols")
					? `${msg} — the server accepted a tokenless handshake but echoed no subprotocol; always mint a client_secret first`
					: `WebSocket error: ${msg}`));
			};
			ws.onclose = (e: any) => { st.closed = { code: e?.code, reason: e?.reason || "" }; this.log("closed", st.closed.code, st.closed.reason); };
			ws.onmessage = (m: any) => {
				const raw = m?.data ?? m;
				let ev: any;
				try { ev = JSON.parse(typeof raw === "string" ? raw : Buffer.from(raw as any).toString("utf-8")); } catch { return; }
				st.events.push(ev);
				this.log("<-", ev.type);
				if (ev.type === EV.delta && opts.onDelta) opts.onDelta(ev.delta ?? "");
				if (ev.type === EV.completed && opts.onFinal) opts.onFinal(ev);
				this.emit(ev);
			};
		});

		st.send = (obj: Record<string, any>) => ws.send(JSON.stringify({ event_id: uid(), ...obj }));
		st.expect = (type: string, ms: number = 20000) => {
			const hit = st.events.find(e => e.type === type);
			if (hit) return Promise.resolve(hit);
			return new Promise<any>((res, rej) => {
				const t = setTimeout(() => rej(new Error(`timeout waiting for "${type}" (saw: ${st.events.map(e => e.type).join(",") || "nothing"})`)), ms);
				const off = this.on(type, ev => { clearTimeout(t); res(ev); });
				if (st.closed) { clearTimeout(t); off(); rej(new Error(`socket closed (code ${st.closed.code}) while waiting for "${type}"`)); }
			});
		};
		st.close = () => { try { ws.close(1000, "bye"); } catch {} };
		return st;
	}

	/** Convenience wrapper: mint -> open -> configure (caller keeps the handle for live/mic streaming). */
	async openSession(opts: { session?: Record<string, any>; onDelta?: (t: string) => void; onFinal?: (ev: any) => void } = {}): Promise<{ handle: SocketHandle; minted: any }> {
		const minted = await this.mint();
		const handle = await this.open({ token: minted.client_secret?.value ?? null, onDelta: opts.onDelta, onFinal: opts.onFinal });
		handle.send({ type: "transcription_session.update", session: this.sessionConfig(minted, opts.session ?? {}) });
		await handle.expect(EV.sessionUpdated, 20000);
		return { handle, minted };
	}

	/**
	 * One-shot transcribe of raw s16le pcm16 audio.
	 * `onDelta` receives cumulative partials (great for a live UI), return value has final text.
	 */
	async transcribe(pcm16: Buffer, opts: {
		frameMs?: number; pace?: boolean; paceMs?: number; session?: Record<string, any>;
		quietMs?: number; emptyGraceMs?: number; finalTimeout?: number; tailMs?: number;
		onDelta?: (t: string) => void; signal?: AbortSignal;
	} = {}): Promise<TranscriptResult> {
		const buf = toBuffer(pcm16);
		const frameMs = opts.frameMs ?? 100;
		const frameBytes = Math.max(2, Math.round((this.cfg.sampleRate * frameMs) / 1000) * 2 * this.cfg.channels);
		const secs = buf.length / (this.cfg.sampleRate * 2 * this.cfg.channels);
		const finalTimeout = opts.finalTimeout ?? Math.max(30000, Math.round(secs * 1000) + 30000);

		const minted = await this.mint();
		if (opts.signal?.aborted) throw new Error("cancelled");
		const handle = await this.open({ token: minted.client_secret?.value ?? null });

		const deltas: string[] = [];
		const finals: Map<string, any> = new Map();
		const failures: unknown[] = [];
		const offs = [
			this.on(EV.delta, e => deltas.push(e.delta ?? "")),
			this.on(EV.completed, (e: any) => {
				// keep the word-bearing copy: the server sends a second, words-less duplicate per segment
				const key = e.item_id ?? "_";
				const prev = finals.get(key);
				const nw = e.words_info?.words?.length ?? 0;
				const pw = prev?.words_info?.words?.length ?? 0;
				finals.set(key, prev && pw > nw && prev.transcript ? { ...e, words_info: prev.words_info } : e);
			}),
			this.on(EV.failed, (e: any) => failures.push(e)),
		];

		try {
			handle.send({ type: "transcription_session.update", session: this.sessionConfig(minted, opts.session ?? {}) });
			await handle.expect(EV.sessionUpdated, 20000);

			for (let i = 0; i < buf.length; i += frameBytes) {
				if (opts.signal?.aborted) throw new Error("cancelled");
				handle.send({
					type: "input_audio_buffer.append",
					audio: buf.subarray(i, Math.min(i + frameBytes, buf.length)).toString("base64"),
				});
				if (opts.pace !== false) await sleep(opts.paceMs ?? frameMs);
			}

			handle.send({ type: "input_audio_buffer.commit" });     // <-- the actual flush
			const t0 = Date.now(), deadline = t0 + finalTimeout;
			const quietMs = opts.quietMs ?? 1500, emptyGraceMs = opts.emptyGraceMs ?? 8000;
			let lastCount = 0, lastChange = t0;
			while (Date.now() < deadline) {
				if (opts.signal?.aborted) break;
				const n = finals.size;
				if (n !== lastCount) { lastCount = n; lastChange = Date.now(); }
				if (failures.length) break;
				if (n > 0 && Date.now() - lastChange >= quietMs) break;
				if (n === 0 && Date.now() - t0 >= emptyGraceMs) break;
				await sleep(120);
			}

			handle.send({ type: "input_audio_buffer.done" });
			await sleep(opts.tailMs ?? 600);

			const list = Array.from(finals.values());
			const text = list.length
				? list.map(f => f.transcript ?? "").filter(Boolean).join(" ").replace(/\s+/g, " ").trim()
				: (deltas.at(-1) ?? "").replace(/\s+/g, " ").trim();

			return {
				text,
				segments: list.map(f => ({ transcript: f.transcript ?? "", isLast: f.is_last_result ?? false })),
				words: list.flatMap(f => f.words_info?.words ?? []),
				failed: failures,
				seconds: +secs.toFixed(3),
				model: minted.input_audio_transcription?.model ?? null,
				sessionId: minted.id ?? "",
				protocol: handle.protocol,
			};
		} finally {
			offs.forEach(f => f());
			try { handle.send?.({ type: "input_audio_buffer.clear" }); } catch {}
			handle.close();
		}
	}
}

export interface SocketHandle {
	events: any[];
	closed: { code: number; reason: string } | null;
	protocol: string | null;
	send?: (obj: Record<string, any>) => void;
	expect?: (type: string, ms?: number) => Promise<any>;
	close?: () => void;
}

function toBuffer(x: any): Buffer {
	if (Buffer.isBuffer(x)) return x;
	if (x instanceof Uint8Array) return Buffer.from(x);
	if (x instanceof ArrayBuffer) return Buffer.from(new Uint8Array(x));
	if (x?.buffer instanceof ArrayBuffer) return Buffer.from(new Uint8Array(x.buffer));
	throw new Error("expected pcm16 audio as Buffer/Uint8Array/ArrayBuffer");
}

/* ------------------------------------------------------------------ audio -- */

/** Decode wav/mp3/m4a/aac/ogg/flac/mp4/video... to interleaved s16le at the target rate. Needs ffmpeg. */
export async function decodeToPcm16(
	path: string,
	{ sampleRate = 16000, channels = 1, start = 0, speed = 1 }: { sampleRate?: number; channels?: number; start?: number; speed?: number } = {},
): Promise<Buffer> {
	const args: string[] = ["-v", "error"];
	if (start > 0) args.push("-ss", String(start));
	args.push("-i", path);
	if (speed !== 1 && speed > 0) {
		// atempo only accepts 0.5..100; chain it for extremes
		let rate = Math.max(0.25, Math.min(4, speed));
		const filters: string[] = [];
		while (rate < 0.5) { filters.push("atempo=0.5"); rate /= 0.5; }
		while (rate > 2) { filters.push("atempo=2.0"); rate /= 2.0; }
		if (rate !== 1) filters.push(`atempo=${rate.toFixed(6)}`);
		if (filters.length) args.push("-af", filters.join(","));
	}
	args.push("-ac", String(channels), "-ar", String(sampleRate), "-f", "s16le", "-");
	const chunks: Buffer[] = [];
	const { code, stderr } = await run("ffmpeg", args, b => chunks.push(b));
	const out = Buffer.concat(chunks);
	if (!out.length) throw new Error(`ffmpeg produced no audio for ${path} (exit ${code}) ${stderr.slice(0, 300)}`.trim());
	return out;
}

/** Native RIFF/EXTENSIBLE parser so plain 16k-mono wav files work without ffmpeg on PATH. */
export async function readWavAsPcm16(path: string, sampleRate: number = 16000): Promise<Buffer> {
	const buf = await readFile(path);
	if (buf.toString("ascii", 0, 4) !== "RIFF" || buf.toString("ascii", 8, 12) !== "WAVE") throw new Error("not a wav file");
	let pos = 12, fmt: any = null, data: Buffer | null = null;
	while (pos + 8 <= buf.length) {
		const id = buf.toString("ascii", pos, pos + 4);
		const size = buf.readUInt32LE(pos + 4);
		const body = buf.subarray(pos + 8, Math.min(pos + 8 + size, buf.length));
		if (id === "fmt ") fmt = { format: body.readUInt16LE(0), channels: body.readUInt16LE(2), rate: body.readUInt32LE(4), bits: body.readUInt16LE(14) };
		else if (id === "data") data = body;
		else if (id === "LIST" || id === "INFO" || id === "fact") { /* skip */ }
		pos += 8 + size + (size % 2);
	}
	if (!fmt || !data) throw new Error("wav is missing fmt /data chunks");
	if (fmt.bits !== 16 || fmt.format !== 1) throw new Error(`wav is ${fmt.bits}bit/format${fmt.format}, need 16bit PCM (use ffmpeg) — falling back`);
	if (fmt.channels !== 1) throw new Error(`wav has ${fmt.channels} channels, need mono (use ffmpeg)`);
	if (fmt.rate !== sampleRate) throw new Error(`wav is ${fmt.rate}Hz, need ${sampleRate}Hz (use ffmpeg)`);
	return data;
}

/** PCM16 with ffmpeg-first strategy and a native wav fast path. */
export async function loadPcm16(path: string, opts: { sampleRate?: number; channels?: number; start?: number; speed?: number } = {}): Promise<Buffer> {
	const sampleRate = opts.sampleRate ?? 16000;
	const channels = opts.channels ?? 1;
	if (!opts.start && !opts.speed && extname(path).toLowerCase() === ".wav") {
		try { return await readWavAsPcm16(path, sampleRate); } catch { /* fall through to ffmpeg */ }
	}
	return decodeToPcm16(path, { sampleRate, channels, start: opts.start ?? 0, speed: opts.speed ?? 1 });
}

export async function hasFfmpeg(): Promise<boolean> {
	try { await run("ffmpeg", ["-version"], () => {}); return true; } catch { return false; }
}

function run(cmd: string, args: string[], onData: (b: Buffer) => void): Promise<{ code: number; stderr: string }> {
	return new Promise((res, rej) => {
		const child = spawn(cmd, args, { stdio: ["ignore", "pipe", "pipe"] });
		let stderr = "";
		child.stdout.on("data", onData);
		child.stderr.on("data", b => { stderr += b.toString(); });
		child.on("error", e => rej(new Error((e as any).code === "ENOENT" ? `${cmd} is not installed / not on PATH` : e.message)));
		child.on("close", code => res({ code: code ?? 0, stderr }));
	});
}