feat: unified speech setup and introduced local STT/TTS capabilities

- Added unified `omp setup speech` flow with JSON/check modes and model picker.
- Added local STT pipeline with sherpa workers, recorder/download flow, and streaming inference.
- Added local TTS pipeline with `omp say`, backend selection, and streaming vocalization.
- Replaced legacy speech settings with unified `speech`/`speechgen` configuration keys.
This commit is contained in:
can1357
2026-06-14 16:07:44 +02:00
parent 46ea12f1a1
commit b830f7912b
55 changed files with 6702 additions and 505 deletions
@@ -29,6 +29,7 @@ export const commands: CommandEntry[] = [
{ name: "join", load: () => import("./commands/join").then(m => m.default) },
{ name: "models", load: () => import("./commands/models").then(m => m.default) },
{ name: "plugin", load: () => import("./commands/plugin").then(m => m.default) },
{ name: "say", load: () => import("./commands/say").then(m => m.default) },
{ name: "setup", load: () => import("./commands/setup").then(m => m.default) },
{ name: "shell", load: () => import("./commands/shell").then(m => m.default) },
{ name: "read", load: () => import("./commands/read").then(m => m.default) },
+45 -13
View File
@@ -56,6 +56,8 @@ async function showHelp(config: CliConfig): Promise<void> {
async function runSmokeTest(): Promise<void> {
const { smokeTestSyncWorker, startServer } = await import("@oh-my-pi/omp-stats");
const { smokeTestTinyTitleWorker } = await import("./tiny/title-client");
const { smokeTestSttWorker } = await import("./stt/asr-client");
const { smokeTestTtsWorker } = await import("./tts/tts-client");
await smokeTestSyncWorker();
const statsServer = await startServer(0);
@@ -71,6 +73,8 @@ async function runSmokeTest(): Promise<void> {
}
await smokeTestTinyTitleWorker();
await smokeTestSttWorker();
await smokeTestTtsWorker();
process.stdout.write("smoke-test: ok\n");
}
@@ -78,6 +82,8 @@ const TINY_WORKER_ARGS = new Set(["--tiny-worker", "__tiny_worker"]);
const STATS_SYNC_WORKER_ARG = "__omp_stats_sync_worker";
const TAB_WORKER_ARG = "__omp_tab_worker";
const JS_EVAL_WORKER_ARG = "__omp_js_eval_worker";
const STT_WORKER_ARG = "__omp_stt_worker";
const TTS_WORKER_ARG = "__omp_tts_worker";
async function runWorkerEntrypoint(arg: string | undefined): Promise<boolean> {
if (arg === STATS_SYNC_WORKER_ARG) {
@@ -110,21 +116,34 @@ async function runWorkerEntrypoint(arg: string | undefined): Promise<boolean> {
await import("./eval/js/worker-entry");
return true;
}
if (arg === STT_WORKER_ARG) {
const { startSttWorker } = await import("./stt/asr-worker");
await runIpcSubprocessWorker(startSttWorker);
return true;
}
if (arg === TTS_WORKER_ARG) {
const { startTtsWorker } = await import("./tts/tts-worker");
await runIpcSubprocessWorker(startTtsWorker);
return true;
}
return false;
}
/**
* Hidden subcommand that boots the tiny-model worker inside this process
* over the parent's IPC channel. The agent's main process spawns the same
* binary with this flag so `onnxruntime-node` (loaded transitively by
* `@huggingface/transformers`) lives in a child address space. The parent
* `SIGKILL`s the child on shutdown so the NAPI finalizer never runs in
* either process — that finalizer segfaults Bun on Windows (issue #1606).
* Boot a subprocess-isolated transformers.js worker over the parent's IPC
* channel and block until the parent disconnects. The tiny-model, STT, and TTS
* workers each run `onnxruntime-node` (loaded transitively by
* `@huggingface/transformers`) in a child address space because its NAPI
* finalizer segfaults Bun on shutdown (issue #1606); the parent `SIGKILL`s the
* child so that finalizer never runs in either process. This wires `process`
* IPC to the worker's typed transport, keeps the event loop alive while the
* worker is idle, and hard-kills the process on parent `disconnect`.
*/
async function runTinyWorker(): Promise<void> {
const { startTinyTitleWorker } = await import("./tiny/worker");
async function runIpcSubprocessWorker<In, Out>(
start: (transport: { send(message: Out): void; onMessage(handler: (message: In) => void): () => void }) => void,
): Promise<void> {
const { promise: shuttingDown, resolve: shutdown } = Promise.withResolvers<void>();
const send = (message: unknown): void => {
const send = (message: Out): void => {
// `process.send` only exists when spawned with an IPC channel; the
// parent always spawns us that way. If it's missing, the parent
// vanished and there's no one to talk to.
@@ -139,10 +158,10 @@ async function runTinyWorker(): Promise<void> {
shutdown();
}
};
startTinyTitleWorker({
start({
send,
onMessage(handler) {
const wrap = (data: unknown): void => handler(data as never);
const wrap = (data: unknown): void => handler(data as In);
process.on("message", wrap);
return () => {
process.off("message", wrap);
@@ -151,8 +170,8 @@ async function runTinyWorker(): Promise<void> {
});
const keepalive = setInterval(() => {}, 2 ** 30);
// Parent went away (crashed, SIGKILL, etc.) — commit suicide so we don't
// linger as an orphan. SIGKILL via `process.kill` keeps us symmetrical
// with the parent's hard-kill on shutdown: skip every JS/native finalizer.
// linger as an orphan. SIGKILL via `process.kill` keeps us symmetrical with
// the parent's hard-kill on shutdown: skip every JS/native finalizer.
process.on("disconnect", () => shutdown());
try {
await shuttingDown;
@@ -162,6 +181,19 @@ async function runTinyWorker(): Promise<void> {
process.kill(process.pid, "SIGKILL");
}
/**
* Hidden subcommand that boots the tiny-model worker inside this process over
* the parent's IPC channel. The agent's main process spawns the same binary
* with this flag so `onnxruntime-node` (loaded transitively by
* `@huggingface/transformers`) lives in a child address space. The parent
* `SIGKILL`s the child on shutdown so the NAPI finalizer never runs in either
* process — that finalizer segfaults Bun on Windows (issue #1606).
*/
async function runTinyWorker(): Promise<void> {
const { startTinyTitleWorker } = await import("./tiny/worker");
await runIpcSubprocessWorker(startTinyTitleWorker);
}
/** Run the CLI with the given argv (no `process.argv` prefix). */
export async function runCli(argv: string[]): Promise<void> {
if (argv[0] === "--smoke-test") {
+149 -48
View File
@@ -4,12 +4,18 @@
* Handles `omp setup` for onboarding and `omp setup <component>` for optional dependencies.
*/
import * as path from "node:path";
import { $which, APP_NAME, getPythonEnvDir } from "@oh-my-pi/pi-utils";
import { $which, APP_NAME, getProjectDir, getPythonEnvDir } from "@oh-my-pi/pi-utils";
import { $ } from "bun";
import chalk from "chalk";
import { Settings, settings } from "../config/settings";
import { theme } from "../modes/theme/theme";
import { downloadSttModel, isSttModelCached } from "../stt/downloader";
import { isSttModelKey, STT_MODEL_OPTIONS } from "../stt/models";
import { detectRecorder, ensureRecorder } from "../stt/recorder";
import { downloadTtsModel, isTtsLocalModelKey, isTtsModelCached, TTS_LOCAL_MODEL_OPTIONS } from "../tts";
import { selectSetupModel } from "./setup-model-picker";
export type SetupComponent = "python" | "stt";
export type SetupComponent = "python" | "speech";
export interface SetupCommandArgs {
component: SetupComponent;
@@ -19,7 +25,7 @@ export interface SetupCommandArgs {
};
}
const VALID_COMPONENTS: SetupComponent[] = ["python", "stt"];
const VALID_COMPONENTS: SetupComponent[] = ["python", "speech"];
const MANAGED_PYTHON_ENV = getPythonEnvDir();
@@ -114,8 +120,8 @@ export async function runSetupCommand(cmd: SetupCommandArgs): Promise<void> {
case "python":
await handlePythonSetup(cmd.flags);
break;
case "stt":
await handleSttSetup(cmd.flags);
case "speech":
await handleSpeechSetup(cmd.flags);
break;
}
}
@@ -149,58 +155,153 @@ async function handlePythonSetup(flags: { json?: boolean; check?: boolean }): Pr
process.exit(1);
}
async function handleSttSetup(flags: { json?: boolean; check?: boolean }): Promise<void> {
const { checkDependencies, formatDependencyStatus } = await import("../stt/setup");
const status = await checkDependencies();
/**
* One installable speech dependency. `isReady`/`status` are read-only probes;
* `pick` (optional) lets an interactive user choose + persist a model; `ensure`
* performs the download, streaming a normalized progress event.
*/
interface SpeechComponent {
name: string;
isReady(): Promise<boolean>;
status(): Promise<string>;
pick?(): Promise<boolean>;
ensure(onProgress: (progress: { stage: string; percent?: number }) => void): Promise<void>;
}
function buildSpeechComponents(): SpeechComponent[] {
return [
{
name: "Recorder",
isReady: async () => detectRecorder() !== null,
status: async () => {
const recorder = detectRecorder();
return recorder ? `${recorder.tool} (${recorder.bin})` : "none — ffmpeg will be downloaded";
},
ensure: async onProgress => {
await ensureRecorder(onProgress);
},
},
{
name: "Speech-to-Text model",
isReady: () => isSttModelCached(settings.get("stt.modelName")),
status: async () => {
const key = settings.get("stt.modelName");
return (await isSttModelCached(key)) ? key : `${key} — not downloaded`;
},
pick: async () => {
const chosen = await selectSetupModel(
"Speech-to-Text model",
[...STT_MODEL_OPTIONS],
settings.get("stt.modelName"),
);
if (chosen === null) return false;
if (isSttModelKey(chosen)) {
settings.set("stt.modelName", chosen);
await settings.flush();
}
return true;
},
ensure: onProgress =>
downloadSttModel(settings.get("stt.modelName"), progress =>
onProgress({ stage: `Downloading ${progress.label} model`, percent: progress.percent }),
),
},
{
name: "Text-to-Speech model",
isReady: () => isTtsModelCached(settings.get("tts.localModel")),
status: async () => {
const key = settings.get("tts.localModel");
return (await isTtsModelCached(key)) ? key : `${key} — model/runtime not installed`;
},
pick: async () => {
const chosen = await selectSetupModel(
"Text-to-Speech model",
[...TTS_LOCAL_MODEL_OPTIONS],
settings.get("tts.localModel"),
);
if (chosen === null) return false;
if (isTtsLocalModelKey(chosen)) {
settings.set("tts.localModel", chosen);
await settings.flush();
}
return true;
},
ensure: async onProgress => {
const ok = await downloadTtsModel(settings.get("tts.localModel"), progress =>
onProgress({ stage: progress.stage, percent: progress.percent }),
);
if (!ok) throw new Error("Failed to download the local text-to-speech model.");
},
},
];
}
/**
* Unified `omp setup speech` flow. Drives every {@link SpeechComponent} through
* one path: report (`--json`/`--check`) or install (interactive pick + ensure
* with single-line progress; non-TTY skips pickers and installs configured
* values).
*/
async function handleSpeechSetup(flags: { json?: boolean; check?: boolean }): Promise<void> {
await Settings.init({ cwd: getProjectDir() });
const components = buildSpeechComponents();
if (flags.json) {
console.log(JSON.stringify(status, null, 2));
if (!status.recorder.available || !status.python.available || !status.whisper.available) process.exit(1);
return;
}
console.log(formatDependencyStatus(status));
if (status.recorder.available && status.python.available && status.whisper.available) {
console.log(chalk.green(`\n${theme.status.success} Speech-to-text is ready`));
const report: Record<string, { ready: boolean; status: string }> = {};
let allReady = true;
for (const component of components) {
const ready = await component.isReady();
if (!ready) allReady = false;
report[component.name] = { ready, status: await component.status() };
}
console.log(JSON.stringify(report, null, 2));
if (!allReady) process.exit(1);
return;
}
if (flags.check) {
process.exit(1);
console.log(chalk.bold("Speech dependencies:"));
let allReady = true;
for (const component of components) {
const ready = await component.isReady();
if (!ready) allReady = false;
const mark = ready ? chalk.green("[ok]") : chalk.yellow("[missing]");
console.log(` ${mark} ${component.name}: ${await component.status()}`);
}
if (!allReady) process.exit(1);
return;
}
if (!status.python.available) {
console.error(chalk.red(`\n${theme.status.error} Python not found`));
console.error(chalk.dim("Install Python 3.8+ and ensure it's in your PATH"));
process.exit(1);
}
if (!status.recorder.available) {
console.error(chalk.yellow(`\n${theme.status.warning} No recording tool found`));
console.error(chalk.dim(status.recorder.installHint));
}
if (!status.whisper.available) {
console.log(chalk.dim(`\nInstalling openai-whisper...`));
const { resolvePython } = await import("../stt/transcriber");
const pythonCmd = resolvePython()!;
const result = await $`${pythonCmd} -m pip install -q openai-whisper`.nothrow();
if (result.exitCode !== 0) {
console.error(chalk.red(`\n${theme.status.error} Failed to install openai-whisper`));
console.error(chalk.dim("Try manually: pip install openai-whisper"));
const interactive = Boolean(process.stdout.isTTY);
for (const component of components) {
if (interactive && component.pick) {
await component.pick();
}
if (await component.isReady()) {
console.log(chalk.green(`${theme.status.success} ${component.name} ready`));
continue;
}
console.log(chalk.dim(`Preparing ${component.name}...`));
try {
await component.ensure(progress => {
const percent = typeof progress.percent === "number" ? ` (${progress.percent}%)` : "";
process.stdout.write(`\r${chalk.dim(`${progress.stage}${percent}`)}\x1b[K`);
});
process.stdout.write("\n");
} catch (err) {
process.stdout.write("\n");
const msg = err instanceof Error ? err.message : `Failed to set up ${component.name}`;
console.error(chalk.red(`${theme.status.error} ${msg}`));
process.exit(1);
}
}
const recheck = await checkDependencies();
if (recheck.recorder.available && recheck.python.available && recheck.whisper.available) {
console.log(chalk.green(`\n${theme.status.success} Speech-to-text is ready`));
} else {
console.error(chalk.red(`\n${theme.status.error} Setup incomplete`));
console.log(formatDependencyStatus(recheck));
process.exit(1);
}
console.log(chalk.green(`\n${theme.status.success} Speech is ready`));
console.log(
chalk.dim(
"Enable speech-to-text via stt.enabled, then hold Space to talk (or bind app.stt.toggle); enable the speech-generation tool via speechgen.enabled; speak replies aloud via speech.enabled.",
),
);
}
/**
@@ -215,7 +316,7 @@ ${chalk.bold("Usage:")}
${chalk.bold("Components:")}
python Verify a Python 3 interpreter is reachable for code execution
stt Install speech-to-text dependencies (openai-whisper, recording tools)
speech Pick + download the speech-to-text and text-to-speech models and an audio recorder
${chalk.bold("Options:")}
-c, --check Check if dependencies are installed without installing
@@ -224,8 +325,8 @@ ${chalk.bold("Options:")}
${chalk.bold("Examples:")}
${APP_NAME} setup Run the onboarding wizard
${APP_NAME} setup python Check Python execution dependencies
${APP_NAME} setup stt Install speech-to-text dependencies
${APP_NAME} setup stt --check Check if STT dependencies are available
${APP_NAME} setup speech Set up speech (pick STT + TTS models, install a recorder)
${APP_NAME} setup speech --check Check if speech dependencies are available
${APP_NAME} setup python --check Check if Python execution is available
`);
}
@@ -0,0 +1,43 @@
/**
* Standalone TUI model picker used by `omp setup speech`.
*
* Mirrors {@link ./session-picker.ts} for the standalone-TUI lifecycle: spin up
* a one-shot {@link TUI} over a {@link SelectList}, resolve on select/cancel, and
* tear the UI down. The standalone TUI auto-renders on input, so no manual
* render wiring is needed beyond `addChild`/`setFocus`/`start`.
*/
import { ProcessTerminal, type SelectItem, SelectList, TUI } from "@oh-my-pi/pi-tui";
import { getSelectListTheme } from "../modes/theme/theme";
/**
* Show a single-column model picker and resolve with the chosen item's value,
* or `null` if the user cancelled. `currentValue` pre-selects the matching row.
*/
export async function selectSetupModel(
title: string,
items: SelectItem[],
currentValue: string,
): Promise<string | null> {
const { promise, resolve } = Promise.withResolvers<string | null>();
const ui = new TUI(new ProcessTerminal());
let resolved = false;
const finish = (value: string | null): void => {
if (resolved) return;
resolved = true;
ui.stop();
resolve(value);
};
const list = new SelectList(items, Math.min(items.length, 10), getSelectListTheme());
const currentIndex = items.findIndex(item => item.value === currentValue);
if (currentIndex >= 0) list.setSelectedIndex(currentIndex);
list.onSelect = item => finish(item.value);
list.onCancel = () => finish(null);
process.stdout.write(`${title}\n`);
ui.addChild(list);
ui.setFocus(list);
ui.start();
return promise;
}
+102
View File
@@ -0,0 +1,102 @@
/**
* Synthesize text with the local TTS engine and play it (or save it with --out).
*
* Demonstrates the on-device speech stack end to end: the first run downloads
* the configured local model, synthesis happens in the TTS worker subprocess,
* and the resulting WAV is either played through the speakers or written to disk.
*/
import * as os from "node:os";
import * as path from "node:path";
import { getProjectDir, Snowflake } from "@oh-my-pi/pi-utils";
import { Args, Command, Flags } from "@oh-my-pi/pi-utils/cli";
import chalk from "chalk";
import { Settings, settings } from "../config/settings";
import { playAudioFile, removeTempFile } from "../tts/player";
import { shutdownTtsClient, ttsClient } from "../tts/tts-client";
import { encodeWav } from "../tts/wav";
export default class Say extends Command {
static description = "Synthesize text with the local TTS engine and play it through the speakers";
static args = {
text: Args.string({ required: true, description: "Text to speak" }),
};
static flags = {
voice: Flags.string({ description: "Voice id" }),
model: Flags.string({ description: "Local TTS model key" }),
out: Flags.string({ char: "o", description: "Write WAV to this path instead of playing" }),
};
static examples = [
'omp say "hello world"',
'omp say "hello world" --out /tmp/hello.wav',
'omp say "bonjour" --voice af_heart --model kokoro',
];
async run(): Promise<void> {
const { args, flags } = await this.parse(Say);
const text = args.text ?? "";
await Settings.init({ cwd: getProjectDir() });
const model = flags.model ?? settings.get("tts.localModel");
const voice = flags.voice ?? settings.get("tts.localVoice");
let exitCode = 0;
const unsubscribe = ttsClient.onProgress(event => {
if (event.status === "progress" && typeof event.progress === "number") {
process.stderr.write(
`\r${chalk.dim(`downloading ${event.file ?? model}: ${Math.round(event.progress)}%`)}`,
);
} else if (event.status === "done" || event.status === "ready") {
// Clear the progress line once the download finishes.
process.stderr.write("\r\x1b[K");
}
});
try {
const audio = await ttsClient.synthesize(model, text, { voice });
if (!audio) {
process.stderr.write(
chalk.red(
`error: could not synthesize with local TTS model "${model}". ` +
"Run `omp setup speech` to install it.\n",
),
);
exitCode = 1;
return;
}
const wav = encodeWav(audio.pcm, audio.sampleRate);
const durationSec = audio.pcm.length / audio.sampleRate;
if (flags.out) {
await Bun.write(flags.out, wav);
process.stdout.write(
`${chalk.green("saved")} ${flags.out} ` +
`${chalk.dim(`(${voice}, ${model}, ${durationSec.toFixed(1)}s, ${wav.byteLength} bytes)`)}\n`,
);
return;
}
const tmp = path.join(os.tmpdir(), `omp-say-${Snowflake.next()}.wav`);
await Bun.write(tmp, wav);
try {
await playAudioFile(tmp);
process.stdout.write(
`${chalk.green("spoke")} ${chalk.dim(`(${voice}, ${model}, ${durationSec.toFixed(1)}s)`)}\n`,
);
} finally {
await removeTempFile(tmp);
}
} catch (err) {
process.stderr.write(chalk.red(`error: ${err instanceof Error ? err.message : String(err)}\n`));
exitCode = 1;
} finally {
unsubscribe();
await shutdownTtsClient();
}
if (exitCode !== 0) process.exit(exitCode);
}
}
+1 -1
View File
@@ -7,7 +7,7 @@ import { runSetupCommand, type SetupCommandArgs, type SetupComponent } from "../
import { runRootCommand } from "../main";
import { initTheme } from "../modes/theme/theme";
const COMPONENTS: SetupComponent[] = ["python", "stt"];
const COMPONENTS: SetupComponent[] = ["python", "speech"];
export interface OnboardingSetupDependencies {
runRoot?: typeof runRootCommand;
@@ -2,6 +2,7 @@ import { THINKING_EFFORTS } from "@oh-my-pi/pi-ai";
import { DEFAULT_SHARE_URL } from "@oh-my-pi/pi-wire";
import { SHAPE_VARIANT_NAMES } from "@oh-my-pi/snapcompact";
import { DEFAULT_RELAY_URL } from "../collab/protocol";
import { DEFAULT_STT_MODEL_KEY, STT_MODEL_OPTIONS, STT_MODEL_VALUES } from "../stt/models";
import { AUTO_THINKING, getConfiguredThinkingLevelMetadata, getThinkingLevelMetadata } from "../thinking";
import {
TINY_MODEL_DEVICE_DEFAULT,
@@ -24,6 +25,14 @@ import {
TINY_TITLE_MODEL_OPTIONS,
TINY_TITLE_MODEL_VALUES,
} from "../tiny/models";
import {
DEFAULT_TTS_LOCAL_MODEL_KEY,
DEFAULT_TTS_VOICE,
TTS_LOCAL_MODEL_OPTIONS,
TTS_LOCAL_MODEL_VALUES,
TTS_LOCAL_VOICE_OPTIONS,
TTS_LOCAL_VOICE_VALUES,
} from "../tts/models";
import { EDIT_MODES } from "../utils/edit-mode";
import { SEARCH_PROVIDER_OPTIONS, SEARCH_PROVIDER_PREFERENCES } from "../web/search/types";
@@ -1441,24 +1450,15 @@ export const SETTINGS_SCHEMA = {
"stt.modelName": {
type: "enum",
values: ["tiny", "tiny.en", "base", "base.en", "small", "small.en", "medium", "medium.en", "large"] as const,
default: "base.en",
values: STT_MODEL_VALUES,
default: DEFAULT_STT_MODEL_KEY,
ui: {
tab: "interaction",
group: "Speech",
label: "Speech Model",
description: "Whisper model size (larger = more accurate but slower)",
options: [
{ value: "tiny", label: "tiny", description: "Multilingual; fastest, lowest accuracy" },
{ value: "tiny.en", label: "tiny.en", description: "English-only; fastest" },
{ value: "base", label: "base", description: "Multilingual; small and fast" },
{ value: "base.en", label: "base.en", description: "English-only; default" },
{ value: "small", label: "small", description: "Multilingual; balanced" },
{ value: "small.en", label: "small.en", description: "English-only; balanced" },
{ value: "medium", label: "medium", description: "Multilingual; accurate but slower" },
{ value: "medium.en", label: "medium.en", description: "English-only; accurate but slower" },
{ value: "large", label: "large", description: "Multilingual; most accurate" },
],
description:
"Local on-device speech model. Parakeet TDT v3 (sherpa-onnx) is the SoTA default; Whisper base/small/large-v3-turbo tiers (transformers.js) trade size for multilingual coverage. Downloaded on first use.",
options: STT_MODEL_OPTIONS,
},
},
@@ -2971,14 +2971,14 @@ export const SETTINGS_SCHEMA = {
},
},
"tts.enabled": {
"speechgen.enabled": {
type: "boolean",
default: false,
ui: {
tab: "tools",
group: "Available Tools",
label: "Text-to-Speech",
description: "Enable the tts tool for xAI Grok Voice speech synthesis",
label: "Speech Generation",
description: "Enable the tts tool for on-device (Kokoro) or xAI Grok Voice speech-file synthesis",
},
},
@@ -3739,6 +3739,93 @@ export const SETTINGS_SCHEMA = {
],
},
},
"providers.tts": {
type: "enum",
values: ["auto", "local", "xai"] as const,
default: "auto",
ui: {
tab: "providers",
group: "Services",
label: "Text-to-Speech Provider",
description: "Backend for the tts tool: local on-device neural TTS (Kokoro-82M) or xAI Grok Voice",
options: [
{
value: "auto",
label: "Auto",
description: "Prefer local on-device TTS; route .mp3 output to xAI when credentials exist",
},
{ value: "local", label: "Local", description: "On-device neural TTS (Kokoro-82M); output is WAV/PCM16" },
{
value: "xai",
label: "xAI Grok Voice",
description: "Requires xAI Grok OAuth or XAI_API_KEY; MP3 or WAV",
},
],
},
},
"tts.localModel": {
type: "enum",
values: TTS_LOCAL_MODEL_VALUES,
default: DEFAULT_TTS_LOCAL_MODEL_KEY,
ui: {
tab: "providers",
group: "Services",
label: "Local TTS Model",
description: "On-device neural TTS model (Kokoro-82M) used by the local TTS backend",
options: TTS_LOCAL_MODEL_OPTIONS,
},
},
"tts.localVoice": {
type: "enum",
values: TTS_LOCAL_VOICE_VALUES,
default: DEFAULT_TTS_VOICE,
ui: {
tab: "providers",
group: "Services",
label: "Local TTS Voice",
description: "Kokoro voice used by the local TTS backend (American/British, female/male)",
options: TTS_LOCAL_VOICE_OPTIONS,
},
},
"speech.enabled": {
type: "boolean",
default: false,
ui: {
tab: "providers",
group: "Services",
label: "Speech Vocalization",
description: "Speak the assistant's output aloud through the speakers as it streams",
},
},
"speech.mode": {
type: "enum",
values: ["all", "assistant", "yield"] as const,
default: "assistant",
ui: {
tab: "providers",
group: "Services",
label: "Speech Vocalization Mode",
description:
"What to speak: all = assistant messages + thinking; assistant = messages only; yield = only the final message at turn end",
options: [
{ value: "all", label: "All (messages + thinking)" },
{ value: "assistant", label: "Assistant messages" },
{ value: "yield", label: "Final message only" },
],
},
},
"speech.voice": {
type: "enum",
values: TTS_LOCAL_VOICE_VALUES,
default: DEFAULT_TTS_VOICE,
ui: {
tab: "providers",
group: "Services",
label: "Speech Vocalization Voice",
description: "Kokoro voice used when speaking the assistant's output aloud",
options: TTS_LOCAL_VOICE_OPTIONS,
},
},
"providers.tinyModel": {
type: "enum",
values: TINY_TITLE_MODEL_VALUES,
@@ -4296,8 +4383,7 @@ export interface SttSettings {
enabled: boolean;
language: string | undefined;
modelName: string;
whisperPath: string | undefined;
modelPath: string | undefined;
streaming: boolean;
}
export interface BashInterceptorRule {
@@ -2,6 +2,7 @@ import { INTENT_FIELD } from "@oh-my-pi/pi-agent-core";
import { calculatePromptTokens } from "@oh-my-pi/pi-agent-core/compaction/compaction";
import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai";
import { type Component, Loader, TERMINAL } from "@oh-my-pi/pi-tui";
import { extractTextContent } from "../../commit/utils";
import { settings } from "../../config/settings";
import { getFileSnapshotStore } from "../../edit/file-snapshot-store";
import { AssistantMessageComponent } from "../../modes/components/assistant-message";
@@ -20,6 +21,7 @@ import type { PlanApprovalDetails } from "../../plan-mode/approved-plan";
import type { AgentSessionEvent } from "../../session/agent-session";
import { isSilentAbort, readQueueChipText, resolveAbortLabel } from "../../session/messages";
import type { ResolveToolDetails } from "../../tools/resolve";
import { vocalizer } from "../../tts/vocalizer";
import { hasVisibleThinking } from "../../utils/thinking-display";
import { interruptHint } from "../shared";
import { StreamingRevealController } from "./streaming-reveal";
@@ -93,8 +95,8 @@ export class EventController {
this.#handlers = {
agent_start: e => this.#handleAgentStart(e),
agent_end: e => this.#handleAgentEnd(e),
turn_start: async () => {},
turn_end: async () => {},
turn_start: async () => this.#handleTurnStart(),
turn_end: async e => this.#handleTurnEnd(e),
message_start: e => this.#handleMessageStart(e),
message_update: e => this.#handleMessageUpdate(e),
message_end: e => this.#handleMessageEnd(e),
@@ -431,7 +433,47 @@ export class EventController {
}
}
/** A new turn interrupts any speech still queued/playing from the previous one. */
#handleTurnStart(): void {
vocalizer.clear();
}
/**
* Speak streamed assistant output as a side effect of the turn. The mode
* decides which deltas feed the vocalizer (the vocalizer re-checks enabled):
* assistant|all speak text; all also speaks thinking; yield speaks nothing
* live (the final message is spoken at turn end).
*/
#vocalizeDelta(event: Extract<AgentSessionEvent, { type: "message_update" }>): void {
if (!settings.get("speech.enabled")) return;
const mode = settings.get("speech.mode");
const delta = event.assistantMessageEvent;
if (delta.type === "text_delta" && (mode === "assistant" || mode === "all")) {
vocalizer.pushDelta(delta.delta);
} else if (delta.type === "thinking_delta" && mode === "all") {
vocalizer.pushDelta(delta.delta);
}
}
/**
* End-of-turn vocalization: yield mode speaks the final assistant message in
* one shot here (the only mode that is post-hoc); every other mode just makes
* sure the live buffer's trailing partial gets flushed.
*/
#handleTurnEnd(event: Extract<AgentSessionEvent, { type: "turn_end" }>): void {
if (!settings.get("speech.enabled")) return;
if (settings.get("speech.mode") !== "yield") {
vocalizer.flush();
return;
}
if (event.message.role !== "assistant") return;
if (event.message.stopReason === "aborted") return; // interrupted: never speak the aborted partial
const text = extractTextContent(event.message);
if (text) vocalizer.speak(text);
}
async #handleMessageUpdate(event: Extract<AgentSessionEvent, { type: "message_update" }>): Promise<void> {
this.#vocalizeDelta(event);
if (this.ctx.streamingComponent && event.message.role === "assistant") {
this.ctx.streamingMessage = event.message;
this.#streamingReveal.setTarget(this.ctx.streamingMessage);
@@ -560,6 +602,17 @@ export class EventController {
async #handleMessageEnd(event: Extract<AgentSessionEvent, { type: "message_end" }>): Promise<void> {
if (event.message.role === "user") return;
if (event.message.role === "assistant" && settings.get("speech.enabled")) {
if (event.message.stopReason === "aborted") {
// Esc / Ctrl+C / interrupt: stop speaking now and drop the trailing partial.
vocalizer.clear();
} else {
const mode = settings.get("speech.mode");
// Speak the last partial sentence of a completed message; yield mode
// instead speaks the whole final message at turn end.
if (mode === "assistant" || mode === "all") vocalizer.flush();
}
}
if (this.ctx.streamingComponent && event.message.role === "assistant") {
this.ctx.streamingMessage = event.message;
this.#streamingReveal.stop();
@@ -97,6 +97,7 @@ import { setAutoQaConsentHandler } from "../tools/report-tool-issue";
import { type ResolveToolDetails, runResolveInvocation } from "../tools/resolve";
import { formatPhaseDisplayName, selectStickyTodoWindow, todoMatchesAnyDescription } from "../tools/todo";
import { ToolError } from "../tools/tool-errors";
import { vocalizer } from "../tts/vocalizer";
import type { EventBus } from "../utils/event-bus";
import { getEditorCommand, openInEditor } from "../utils/external-editor";
import { getSessionAccentAnsi, getSessionAccentHex } from "../utils/session-color";
@@ -3395,7 +3396,11 @@ export class InteractiveMode implements InteractiveModeContext {
await this.#sttController.toggle(this.editor, {
showWarning: (msg: string) => this.showWarning(msg),
showStatus: (msg: string) => this.showStatus(msg),
requestRender: () => this.ui.requestRender(),
onStateChange: (state: SttState) => {
// Duck assistant speech while the user is talking (push-to-talk); restore after.
if (state === "recording") vocalizer.duck();
else vocalizer.unduck();
if (state === "recording") {
this.#voicePreviousShowHardwareCursor = this.ui.getShowHardwareCursor();
this.#voicePreviousUseTerminalCursor = this.editor.getUseTerminalCursor();
+1 -1
View File
@@ -1728,7 +1728,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
customTools.push(...(imageGenTools as unknown as CustomTool[]));
}
if (settings.get("tts.enabled")) {
if (settings.get("speechgen.enabled")) {
customTools.push(ttsTool as unknown as CustomTool);
}
+520
View File
@@ -0,0 +1,520 @@
import * as path from "node:path";
import { $env, isBunTestRuntime, isCompiledBinary, logger, workerHostEntry } from "@oh-my-pi/pi-utils";
import type { Subprocess } from "bun";
import { settings } from "../config/settings";
import { tinyWorkerEnvOverlay } from "../tiny/title-client";
import type { SttProgressEvent, SttWorkerInbound, SttWorkerOutbound } from "./asr-protocol";
import type { SttModelKey } from "./models";
/**
* Abstraction over the speech-recognition subprocess. Modelled as a worker
* interface so the parent composes lifecycle, ping/pong, and request/response
* correlation uniformly; the runtime implementation is a Bun child process so
* `onnxruntime-node`'s NAPI finalizer never runs inside the main agent address
* space — that destructor segfaults Bun on shutdown (issue #1606).
*/
interface WorkerHandle {
send(message: SttWorkerInbound): void;
onMessage(handler: (message: SttWorkerOutbound) => void): () => void;
onError(handler: (error: Error) => void): () => void;
terminate(): Promise<void>;
}
type PendingRequest =
| { kind: "transcribe"; modelKey: SttModelKey; resolve: (text: string) => void; reject: (error: Error) => void }
| { kind: "download"; modelKey: SttModelKey; resolve: (ok: boolean) => void };
export interface SttTranscribeOptions {
language?: string;
signal?: AbortSignal;
}
export interface SttDownloadOptions {
signal?: AbortSignal;
onProgress?: (event: SttProgressEvent) => void;
}
/** Live streaming session handle returned by {@link SttClient.startStream}. */
export interface SttStreamHandle {
/** Feed 16 kHz mono float samples as the recorder produces them. */
pushAudio(audio: Float32Array): void;
/** Flush the trailing segment and resolve with the full joined transcript. */
stop(): Promise<string>;
/** Tear the session down without a final flush (resolves `stop()` with ""). */
cancel(): void;
}
export interface SttStreamOptions {
language?: string;
signal?: AbortSignal;
/** Volatile transcript of the in-progress segment, refreshed as audio arrives. */
onPartial?: (text: string) => void;
/** A finalized segment, emitted once when the endpointer commits it. */
onSegment?: (text: string, index: number) => void;
}
interface StreamState {
modelKey: SttModelKey;
onPartial: ((text: string) => void) | undefined;
onSegment: ((text: string, index: number) => void) | undefined;
resolve: (text: string) => void;
reject: (error: Error) => void;
/** Run `apply` (resolve/reject) once, then unregister the stream. */
finish: (apply: () => void) => void;
}
// Cold-starting the worker subprocess from a compiled binary (decompress +
// module graph load) is slow on contended CI runners; the probe only needs to
// prove the worker spawns and ponges, so a generous bound removes the flake.
const SMOKE_TEST_TIMEOUT_MS = 30_000;
/**
* Hidden subcommand on the main CLI that boots the speech-recognition worker in
* the spawned subprocess. Kept in sync with the dispatch in `cli.ts`.
*/
export const STT_WORKER_ARG = "__omp_stt_worker";
function readTinyModelSetting(key: "providers.tinyModelDevice" | "providers.tinyModelDtype"): string | undefined {
try {
const value = settings.get(key);
return typeof value === "string" ? value : undefined;
} catch {
// Settings may be uninitialized (e.g. `omp --smoke-test`); fall back to env/default.
return undefined;
}
}
/**
* Env handed to the speech subprocess. The `PI_TINY_DEVICE` / `PI_TINY_DTYPE`
* env vars win; otherwise the persisted `providers.tinyModelDevice` /
* `providers.tinyModelDtype` settings are mapped onto those vars so the
* subprocess's env-based resolution picks them up (shared with tiny models).
*/
function sttWorkerEnv(): Record<string, string> {
const overlay = tinyWorkerEnvOverlay(
$env,
readTinyModelSetting("providers.tinyModelDevice"),
readTinyModelSetting("providers.tinyModelDtype"),
);
const base = $env as Record<string, string | undefined>;
const merged: Record<string, string> = {};
for (const key in base) {
const value = base[key];
if (typeof value === "string") merged[key] = value;
}
for (const key in overlay) merged[key] = overlay[key];
return merged;
}
interface SttWorkerSpawnCommand {
cmd: string[];
cwd?: string;
}
/**
* Resolve the command used to relaunch the agent CLI into stt-worker mode. In a
* compiled binary the entry point is the binary itself; otherwise re-enter the
* declared worker-host entry with a cwd-relative script path (Bun's subprocess
* IPC is more reliable that way under `bun test`), falling back to this
* package's own `src/cli.ts` when no host entry is declared.
*/
function sttWorkerSpawnCmd(): SttWorkerSpawnCommand {
if (isCompiledBinary()) return { cmd: [process.execPath, STT_WORKER_ARG] };
const hostEntry = workerHostEntry();
if (hostEntry) {
return { cmd: [process.execPath, path.basename(hostEntry), STT_WORKER_ARG], cwd: path.dirname(hostEntry) };
}
const packageRoot = path.resolve(import.meta.dir, "..", "..");
return { cmd: [process.execPath, "src/cli.ts", STT_WORKER_ARG], cwd: packageRoot };
}
interface SpawnedSubprocess {
proc: Subprocess<"ignore", "ignore", "ignore">;
inbound: Set<(message: SttWorkerOutbound) => void>;
errors: Set<(error: Error) => void>;
/**
* Flipped to `true` right before the parent SIGKILLs the child so `onExit`
* can distinguish the expected hard-kill from a crash/OOM/external signal.
*/
intentionalExit: { value: boolean };
}
/**
* Spawn the speech worker as a subprocess. Exported for tests and the smoke
* probe; production callers go through {@link spawnSttWorker}.
*/
export function createSttSubprocess(): SpawnedSubprocess {
const inbound = new Set<(message: SttWorkerOutbound) => void>();
const errors = new Set<(error: Error) => void>();
const intentionalExit = { value: false };
const spawnCommand = sttWorkerSpawnCmd();
const proc = Bun.spawn({
cmd: spawnCommand.cmd,
cwd: spawnCommand.cwd,
env: sttWorkerEnv(),
stdin: "ignore",
stdout: "ignore",
stderr: "ignore",
serialization: "advanced",
windowsHide: true,
ipc(message) {
for (const handler of inbound) handler(message as SttWorkerOutbound);
},
onExit(_proc, exitCode, signalCode) {
if (exitCode === 0) return;
// Swallow only the expected SIGKILL from `terminate()`; every other
// signal exit is a real worker death that must fault in-flight
// requests so callers don't await forever.
if (exitCode === null && intentionalExit.value) return;
const reason = exitCode !== null ? `code ${exitCode}` : `signal ${signalCode ?? "unknown"}`;
const err = new Error(`stt subprocess exited with ${reason}`);
for (const handler of errors) handler(err);
},
});
// Don't keep the parent event loop alive on an idle worker; dispose calls
// `terminate()` explicitly. Bun's test runner can starve IPC delivery for
// unref'd subprocesses, so keep it referenced under tests.
if (!isBunTestRuntime()) proc.unref();
return { proc, inbound, errors, intentionalExit };
}
function wrapSubprocess({ proc, inbound, errors, intentionalExit }: SpawnedSubprocess): WorkerHandle {
return {
send(message) {
try {
proc.send(message);
} catch (error) {
logger.debug("stt: send to subprocess failed", {
error: error instanceof Error ? error.message : String(error),
});
}
},
onMessage(handler) {
inbound.add(handler);
return () => inbound.delete(handler);
},
onError(handler) {
errors.add(handler);
return () => errors.delete(handler);
},
async terminate() {
// SIGKILL: the whole point of subprocess isolation is that the parent
// never runs `onnxruntime-node`'s NAPI finalizer. Hard-kill instead —
// the model lives in process memory and the OS reclaims everything.
intentionalExit.value = true;
try {
proc.kill("SIGKILL");
} catch {
// Already gone.
}
},
};
}
function spawnInlineUnavailableWorker(error: unknown): WorkerHandle {
const listeners = new Set<(message: SttWorkerOutbound) => void>();
const errorMessage = error instanceof Error ? error.message : String(error);
const emit = (message: SttWorkerOutbound): void => {
for (const listener of listeners) listener(message);
};
return {
send(message) {
queueMicrotask(() => {
if (message.type === "ping") {
emit({ type: "pong", id: message.id });
return;
}
emit({ type: "error", id: message.id, error: errorMessage });
});
},
onMessage(handler) {
listeners.add(handler);
return () => listeners.delete(handler);
},
onError() {
return () => {};
},
async terminate() {
listeners.clear();
},
};
}
function spawnSttWorker(): WorkerHandle {
try {
return wrapSubprocess(createSttSubprocess());
} catch (error) {
logger.warn("stt worker spawn failed; speech-to-text disabled", {
error: error instanceof Error ? error.message : String(error),
});
return spawnInlineUnavailableWorker(error);
}
}
function logWorkerMessage(message: Extract<SttWorkerOutbound, { type: "log" }>): void {
if (message.level === "debug") logger.debug(message.msg, message.meta);
else if (message.level === "warn") logger.warn(message.msg, message.meta);
else logger.error(message.msg, message.meta);
}
export class SttClient {
#worker: WorkerHandle | null = null;
#unsubscribeMessage: (() => void) | null = null;
#unsubscribeError: (() => void) | null = null;
#pending = new Map<string, PendingRequest>();
#streams = new Map<string, StreamState>();
#progressListeners = new Set<(event: SttProgressEvent) => void>();
#nextRequestId = 0;
#spawnWorker: () => WorkerHandle;
constructor(spawnWorker: () => WorkerHandle = spawnSttWorker) {
this.#spawnWorker = spawnWorker;
}
onProgress(listener: (event: SttProgressEvent) => void): () => void {
this.#progressListeners.add(listener);
return () => this.#progressListeners.delete(listener);
}
/**
* Transcribe 16 kHz mono audio on the warm worker. Rejects with the worker
* error on failure and with an `AbortError` when the signal fires (the warm
* worker keeps the model loaded across calls — the model is never reloaded).
*/
async transcribe(modelKey: SttModelKey, audio: Float32Array, options: SttTranscribeOptions = {}): Promise<string> {
options.signal?.throwIfAborted();
const worker = this.#ensureWorker();
const id = String(++this.#nextRequestId);
const { promise, resolve, reject } = Promise.withResolvers<string>();
this.#pending.set(id, { kind: "transcribe", modelKey, resolve, reject });
const abort = (): void => {
const pending = this.#pending.get(id);
if (pending?.kind !== "transcribe") return;
this.#pending.delete(id);
pending.reject(new DOMException("The operation was aborted.", "AbortError"));
};
options.signal?.addEventListener("abort", abort, { once: true });
try {
worker.send({ type: "transcribe", id, modelKey, audio, language: options.language });
return await promise;
} finally {
options.signal?.removeEventListener("abort", abort);
this.#pending.delete(id);
}
}
/**
* Open a live streaming session on the warm worker. Audio fed through the
* returned handle is segmented by the worker's endpointer: `onSegment` fires
* once per committed segment and `onPartial` for the volatile in-progress
* preview. `stop()` resolves with the full joined transcript; `cancel()` (or
* an aborted signal) tears the session down and resolves `stop()` with "".
*/
startStream(modelKey: SttModelKey, options: SttStreamOptions = {}): SttStreamHandle {
const worker = this.#ensureWorker();
const id = String(++this.#nextRequestId);
const { promise, resolve, reject } = Promise.withResolvers<string>();
const signal = options.signal;
let settled = false;
const onAbort = (): void => handle.cancel();
const finish = (apply: () => void): void => {
if (settled) return;
settled = true;
this.#streams.delete(id);
signal?.removeEventListener("abort", onAbort);
apply();
};
this.#streams.set(id, {
modelKey,
onPartial: options.onPartial,
onSegment: options.onSegment,
resolve,
reject,
finish,
});
worker.send({ type: "stream_start", id, modelKey, language: options.language });
const handle: SttStreamHandle = {
pushAudio: audio => {
if (!settled) worker.send({ type: "stream_audio", id, audio });
},
stop: () => {
if (!settled) worker.send({ type: "stream_stop", id });
return promise;
},
cancel: () => {
if (settled) return;
worker.send({ type: "stream_cancel", id });
finish(() => resolve(""));
},
};
if (signal?.aborted) handle.cancel();
else signal?.addEventListener("abort", onAbort, { once: true });
return handle;
}
async downloadModel(modelKey: SttModelKey, options: SttDownloadOptions = {}): Promise<boolean> {
if (options.signal?.aborted) return false;
const unsubscribe = options.onProgress ? this.onProgress(options.onProgress) : undefined;
try {
const worker = this.#ensureWorker();
const id = String(++this.#nextRequestId);
const { promise, resolve } = Promise.withResolvers<boolean>();
this.#pending.set(id, { kind: "download", modelKey, resolve });
const abort = (): void => {
const pending = this.#pending.get(id);
if (pending?.kind !== "download") return;
this.#pending.delete(id);
pending.resolve(false);
};
options.signal?.addEventListener("abort", abort, { once: true });
try {
worker.send({ type: "download", id, modelKey });
return await promise;
} finally {
options.signal?.removeEventListener("abort", abort);
this.#pending.delete(id);
}
} catch (error) {
logger.debug("stt: local model download failed", {
modelKey,
error: error instanceof Error ? error.message : String(error),
});
return false;
} finally {
unsubscribe?.();
}
}
async terminate(): Promise<void> {
const worker = this.#worker;
this.#worker = null;
this.#unsubscribeMessage?.();
this.#unsubscribeMessage = null;
this.#unsubscribeError?.();
this.#unsubscribeError = null;
for (const pending of this.#pending.values()) {
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "transcribe") pending.reject(new Error("stt worker terminated"));
else pending.resolve(false);
}
this.#pending.clear();
this.#failStreams(new Error("stt worker terminated"));
try {
await worker?.terminate();
} catch {
// Already gone.
}
}
#ensureWorker(): WorkerHandle {
if (this.#worker) return this.#worker;
const worker = this.#spawnWorker();
this.#worker = worker;
this.#unsubscribeMessage = worker.onMessage(message => this.#handleMessage(message));
this.#unsubscribeError = worker.onError(error => this.#handleWorkerError(error));
return worker;
}
#handleMessage(message: SttWorkerOutbound): void {
if (message.type === "log") {
logWorkerMessage(message);
return;
}
if (message.type === "progress") {
this.#emitProgress(message.event);
return;
}
if (message.type === "pong") return;
if (message.type === "partial" || message.type === "segment" || message.type === "stream_done") {
const stream = this.#streams.get(message.id);
if (!stream) return;
if (message.type === "partial") stream.onPartial?.(message.text);
else if (message.type === "segment") stream.onSegment?.(message.text, message.index);
else stream.finish(() => stream.resolve(message.text));
return;
}
const pending = this.#pending.get(message.id);
if (!pending) {
if (message.type === "error") {
const stream = this.#streams.get(message.id);
if (stream) {
this.#emitProgress({ modelKey: stream.modelKey, status: "error" });
stream.finish(() => stream.reject(new Error(message.error)));
}
}
return;
}
this.#pending.delete(message.id);
if (message.type === "transcription") {
if (pending.kind === "transcribe") pending.resolve(message.text);
return;
}
if (message.type === "downloaded") {
if (pending.kind === "download") pending.resolve(true);
return;
}
// message.type === "error"
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "transcribe") pending.reject(new Error(message.error));
else pending.resolve(false);
}
#emitProgress(event: SttProgressEvent): void {
for (const listener of this.#progressListeners) listener(event);
}
#failStreams(error: Error): void {
for (const stream of [...this.#streams.values()]) {
this.#emitProgress({ modelKey: stream.modelKey, status: "error" });
stream.finish(() => stream.reject(error));
}
}
#handleWorkerError(error: Error): void {
logger.warn("stt: worker error", { error: error.message });
for (const pending of this.#pending.values()) {
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "transcribe") pending.reject(error);
else pending.resolve(false);
}
this.#pending.clear();
this.#failStreams(error);
void this.terminate();
}
}
export const sttClient = new SttClient();
export async function shutdownSttClient(): Promise<void> {
await sttClient.terminate();
}
export async function smokeTestSttWorker({
timeoutMs = SMOKE_TEST_TIMEOUT_MS,
}: {
timeoutMs?: number;
} = {}): Promise<void> {
const handle = wrapSubprocess(createSttSubprocess());
const { promise, resolve, reject } = Promise.withResolvers<void>();
const timer = setTimeout(() => reject(new Error(`stt worker did not pong within ${timeoutMs}ms`)), timeoutMs);
const unsubscribeMessage = handle.onMessage(message => {
if (message.type === "pong") {
resolve();
return;
}
if (message.type === "log") return;
reject(new Error(`stt worker: expected pong, got ${JSON.stringify(message)}`));
});
const unsubscribeError = handle.onError(reject);
try {
handle.send({ type: "ping", id: "smoke" } satisfies SttWorkerInbound);
await promise;
} finally {
clearTimeout(timer);
unsubscribeMessage();
unsubscribeError();
await handle.terminate();
}
}
@@ -0,0 +1,65 @@
import type { SttModelKey } from "./models";
export type SttProgressStatus = "initiate" | "download" | "progress" | "progress_total" | "done" | "ready" | "error";
export interface SttProgressFileState {
loaded: number;
total: number;
}
export interface SttProgressEvent {
modelKey: SttModelKey;
status: SttProgressStatus;
name?: string;
file?: string;
progress?: number;
loaded?: number;
total?: number;
files?: Record<string, SttProgressFileState>;
task?: string;
model?: string;
}
export type SttWorkerInbound =
| { type: "ping"; id: string }
| { type: "transcribe"; id: string; modelKey: SttModelKey; audio: Float32Array; language?: string }
| { type: "download"; id: string; modelKey: SttModelKey }
// ── Live streaming session ──
// `stream_start` warms the model and opens a session; `stream_audio` feeds
// 16 kHz mono float frames as they arrive from the recorder; `stream_stop`
// flushes the trailing speech segment and ends the session; `stream_cancel`
// tears it down without a final flush. All carry the same `id`.
| { type: "stream_start"; id: string; modelKey: SttModelKey; language?: string }
| { type: "stream_audio"; id: string; audio: Float32Array }
| { type: "stream_stop"; id: string }
| { type: "stream_cancel"; id: string };
export type SttWorkerOutbound =
| { type: "pong"; id: string }
| { type: "transcription"; id: string; text: string }
| { type: "downloaded"; id: string }
| { type: "error"; id: string; error: string }
| { type: "progress"; id: string; event: SttProgressEvent }
| { type: "log"; level: "debug" | "warn" | "error"; msg: string; meta?: Record<string, unknown> }
// ── Live streaming session ──
// `partial` is the volatile transcript of the in-progress speech segment
// (refreshed as more audio arrives, never appended verbatim); `segment` is a
// finalized segment committed once at an endpoint; `stream_done` carries the
// full transcript (all committed segments joined) when the session ends.
| { type: "partial"; id: string; text: string }
| { type: "segment"; id: string; index: number; text: string }
| { type: "stream_done"; id: string; text: string };
/**
* Wire transport between the parent (`SttClient`) and the speech-recognition
* subprocess. The parent owns the subprocess lifecycle (graceful work, hard
* SIGKILL on shutdown); the protocol therefore carries no explicit close
* handshake — once the parent decides to terminate, it signals the OS to reap
* the child so `onnxruntime-node`'s NAPI finalizer never runs in any shared
* address space (the destructor segfaults Bun on shutdown; issue #1606). See
* `asr-client.ts` for the spawn/kill glue.
*/
export interface SttTransport {
send(message: SttWorkerOutbound): void;
onMessage(handler: (message: SttWorkerInbound) => void): () => void;
}
+790
View File
@@ -0,0 +1,790 @@
import * as fs from "node:fs/promises";
import { createRequire } from "node:module";
import * as os from "node:os";
import * as path from "node:path";
import type {
AutomaticSpeechRecognitionOutput,
AutomaticSpeechRecognitionPipeline,
ProgressInfo,
} from "@huggingface/transformers";
import {
ensureRuntimeInstalled,
getTinyModelsCacheDir,
installRuntimeModuleResolver,
isCompiledBinary,
resolveRuntimeModule,
} from "@oh-my-pi/pi-utils";
import packageJson from "../../package.json" with { type: "json" };
import { resolveTinyModelDevicePreference, type TinyModelDevice, tinyModelDeviceLoadOrder } from "../tiny/device";
import { resolveTinyModelDtypeOverride, type TinyModelDtype } from "../tiny/dtype";
import type { SttProgressEvent, SttTransport, SttWorkerInbound } from "./asr-protocol";
import { type EndpointerEvent, StreamEndpointer } from "./endpointer";
import {
getSttModelSpec,
type SherpaSttModelSpec,
type SttModel,
type SttModelKey,
type TransformersSttModelSpec,
} from "./models";
const ASR_TASK = "automatic-speech-recognition";
const TRANSFORMERS_PACKAGE = "@huggingface/transformers";
const SHERPA_PACKAGE = "sherpa-onnx-node";
const COMPILED_TRANSFORMERS_VERSION = process.env.PI_TINY_TRANSFORMERS_VERSION;
// Whisper long-form decoding: split into 30s windows with 5s overlap so audio of
// any length transcribes without exceeding the 30s receptive field.
const CHUNK_LENGTH_S = 30;
const STRIDE_LENGTH_S = 5;
// The client always resamples to 16 kHz mono float32 before sending; sherpa-onnx
// is told the true input rate (it resamples internally to its feature config).
const ASR_SAMPLE_RATE = 16_000;
// Hub origin for raw sherpa-onnx model files (encoder/decoder/joiner/tokens).
const HF_RESOLVE_BASE = "https://huggingface.co";
// Coalesce download progress so streaming a multi-hundred-MB model file doesn't
// flood the IPC channel with one event per chunk.
const PROGRESS_EMIT_BYTES = 4_000_000;
const sourceRequire = createRequire(import.meta.url);
const sttModelDevicePreference = resolveTinyModelDevicePreference();
const sttModelDtypeOverride = resolveTinyModelDtypeOverride();
/**
* Subset of the transformers.js ASR call options we set. The index signature
* mirrors `GenerationFunctionParameters` so this is assignable to the pipeline's
* `Partial<AutomaticSpeechRecognitionConfig>` param (not re-exported from the
* package root, so we model only what we pass).
*/
interface AsrCallOptions {
chunk_length_s: number;
stride_length_s: number;
return_timestamps: boolean;
task?: string;
language?: string;
[key: string]: unknown;
}
interface TransformersRuntime {
env: {
cacheDir?: string;
allowLocalModels?: boolean;
logLevel?: unknown;
};
LogLevel: {
ERROR: unknown;
};
pipeline: (
task: typeof ASR_TASK,
model: string,
options: {
device: TinyModelDevice;
dtype: TinyModelDtype;
progress_callback: (info: ProgressInfo) => void;
},
) => Promise<AutomaticSpeechRecognitionPipeline>;
}
/** Recognition result returned by `sherpa-onnx-node`'s offline recognizer. */
interface SherpaOfflineResult {
text?: string;
}
/** A sherpa-onnx offline stream that accepts a single waveform before decoding. */
interface SherpaOfflineStream {
acceptWaveform(audio: { samples: Float32Array; sampleRate: number }): void;
}
interface SherpaOfflineRecognizer {
createStream(): SherpaOfflineStream;
decodeAsync(stream: SherpaOfflineStream): Promise<SherpaOfflineResult>;
}
/** Offline recognizer config passed to `sherpa-onnx-node` (transducer family). */
interface SherpaOfflineConfig {
modelConfig: {
transducer: { encoder: string; decoder: string; joiner: string };
tokens: string;
modelType: string;
numThreads: number;
provider: string;
debug: number;
};
decodingMethod: string;
}
/** Subset of the native `sherpa-onnx-node` module surface we use. */
interface SherpaRuntime {
OfflineRecognizer: {
createAsync(config: SherpaOfflineConfig): Promise<SherpaOfflineRecognizer>;
};
}
/** A warm model plus the engine that loaded it; cached per tier key. */
type LoadedModel =
| { engine: "transformers"; pipeline: AutomaticSpeechRecognitionPipeline }
| { engine: "sherpa"; recognizer: SherpaOfflineRecognizer };
const models = new Map<SttModelKey, Promise<LoadedModel>>();
// Serialize all model inference on a single chain: the recognizers are not
// guaranteed reentrant and there is one CPU-bound model per tier. Batch
// transcribes and live-stream segment/partial decodes share this lock.
let modelLock = Promise.resolve();
function runOnModel<T>(work: () => Promise<T>): Promise<T> {
const run = modelLock.then(work, work);
modelLock = run.then(
() => undefined,
() => undefined,
);
return run;
}
let transformersRuntime: Promise<TransformersRuntime> | null = null;
let sherpaRuntime: Promise<SherpaRuntime> | null = null;
let cachedTransformersVersionSpec: string | undefined;
function resolveTransformersVersionSpec(): string {
const manifest = packageJson as {
optionalDependencies?: Record<string, string>;
dependencies?: Record<string, string>;
};
const versionSpec =
manifest.optionalDependencies?.[TRANSFORMERS_PACKAGE] ?? manifest.dependencies?.[TRANSFORMERS_PACKAGE];
if (!versionSpec) throw new Error(`${TRANSFORMERS_PACKAGE} is missing from package.json optionalDependencies`);
if (!versionSpec.startsWith("catalog:")) return versionSpec;
if (COMPILED_TRANSFORMERS_VERSION) return COMPILED_TRANSFORMERS_VERSION;
const installed = sourceRequire(`${TRANSFORMERS_PACKAGE}/package.json`) as { version: string };
return installed.version;
}
/**
* Lazily resolve (and memoize) the transformers version spec. In the `catalog:`
* case this `require`s the installed package manifest, so defer it to the
* compiled-binary runtime-install path (only reached on a real transcribe /
* download) — loading this worker for a smoke ping never triggers the resolve.
*/
function getTransformersVersionSpec(): string {
cachedTransformersVersionSpec ??= resolveTransformersVersionSpec();
return cachedTransformersVersionSpec;
}
let cachedSherpaVersionSpec: string | undefined;
function resolveSherpaVersionSpec(): string {
const manifest = packageJson as {
optionalDependencies?: Record<string, string>;
dependencies?: Record<string, string>;
};
const versionSpec = manifest.optionalDependencies?.[SHERPA_PACKAGE] ?? manifest.dependencies?.[SHERPA_PACKAGE];
if (!versionSpec) throw new Error(`${SHERPA_PACKAGE} is missing from package.json optionalDependencies`);
return versionSpec;
}
function getSherpaVersionSpec(): string {
cachedSherpaVersionSpec ??= resolveSherpaVersionSpec();
return cachedSherpaVersionSpec;
}
function errorText(error: unknown): string {
return error instanceof Error ? (error.stack ?? error.message) : String(error);
}
function errorMessage(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function sendLog(
transport: SttTransport,
level: "debug" | "warn" | "error",
msg: string,
meta?: Record<string, unknown>,
): void {
transport.send({ type: "log", level, msg, meta });
}
function getSttRuntimeDir(): string {
const key = getTransformersVersionSpec().replace(/[^A-Za-z0-9._-]/g, "_");
return path.join(path.dirname(getTinyModelsCacheDir()), "stt-runtime", `transformers-${key}`);
}
function getSherpaRuntimeDir(): string {
const key = getSherpaVersionSpec().replace(/[^A-Za-z0-9._-]/g, "_");
return path.join(path.dirname(getTinyModelsCacheDir()), "stt-runtime", `sherpa-${key}`);
}
function sendRuntimeInstallProgress(
transport: SttTransport,
requestId: string,
modelKey: SttModelKey,
status: "initiate" | "download" | "done",
name: string,
): void {
transport.send({ type: "progress", id: requestId, event: { modelKey, status, name } });
}
/**
* Prepare the freshly-installed compiled runtime for loading: stub `sharp` (the
* speech pipeline is audio-only, so the native image codec is dead weight) and
* patch the module resolver so Transformers.js's bare requires resolve against
* the cache. Returns the absolute Transformers.js entrypoint to `require`.
*/
async function prepareCompiledRuntime(runtimeDir: string): Promise<string> {
const nodeModules = path.join(runtimeDir, "node_modules");
const sharpStub = path.join(runtimeDir, "omp-sharp-stub.cjs");
await Bun.write(sharpStub, "module.exports = {};\n");
installRuntimeModuleResolver({ runtimeNodeModules: nodeModules, stubs: { sharp: sharpStub } });
const entry = resolveRuntimeModule(nodeModules, TRANSFORMERS_PACKAGE);
if (!entry) throw new Error(`Unable to resolve ${TRANSFORMERS_PACKAGE} in compiled runtime at ${nodeModules}`);
return entry;
}
function configureTransformers(transformers: TransformersRuntime): TransformersRuntime {
transformers.env.cacheDir = getTinyModelsCacheDir();
transformers.env.allowLocalModels = false;
transformers.env.logLevel = transformers.LogLevel.ERROR;
return transformers;
}
async function loadTransformers(
transport: SttTransport,
requestId: string,
modelKey: SttModelKey,
): Promise<TransformersRuntime> {
if (transformersRuntime) return transformersRuntime;
transformersRuntime = (async () => {
if (!isCompiledBinary()) return configureTransformers(sourceRequire(TRANSFORMERS_PACKAGE) as TransformersRuntime);
const runtimeDir = await ensureRuntimeInstalled({
runtimeDir: getSttRuntimeDir(),
install: {
dependencies: { [TRANSFORMERS_PACKAGE]: getTransformersVersionSpec() },
trustedDependencies: ["onnxruntime-node"],
},
probePackage: TRANSFORMERS_PACKAGE,
onPhase: phase =>
sendRuntimeInstallProgress(
transport,
requestId,
modelKey,
phase,
`${TRANSFORMERS_PACKAGE}@${getTransformersVersionSpec()}`,
),
});
const entry = await prepareCompiledRuntime(runtimeDir);
const require_ = createRequire(entry);
return configureTransformers(require_(entry) as TransformersRuntime);
})().catch(error => {
transformersRuntime = null;
throw error;
});
return transformersRuntime;
}
/**
* Resolve the native `sherpa-onnx-node` module. In a compiled binary the addon
* (plus its per-platform prebuilt `sherpa-onnx.node` + bundled onnxruntime
* dylibs) is installed into a side runtime dir; the addon resolves its native
* library relative to its own location, so a plain `createRequire` of the entry
* is enough — no module-resolver patch or bare-require stubbing is needed.
* Memoized so the runtime loads once per process.
*/
async function loadSherpaRuntime(
transport: SttTransport,
requestId: string,
modelKey: SttModelKey,
): Promise<SherpaRuntime> {
if (sherpaRuntime) return sherpaRuntime;
sherpaRuntime = (async () => {
if (!isCompiledBinary()) return sourceRequire(SHERPA_PACKAGE) as SherpaRuntime;
const runtimeDir = await ensureRuntimeInstalled({
runtimeDir: getSherpaRuntimeDir(),
install: { dependencies: { [SHERPA_PACKAGE]: getSherpaVersionSpec() } },
probePackage: SHERPA_PACKAGE,
onPhase: phase =>
sendRuntimeInstallProgress(
transport,
requestId,
modelKey,
phase,
`${SHERPA_PACKAGE}@${getSherpaVersionSpec()}`,
),
});
const nodeModules = path.join(runtimeDir, "node_modules");
const entry = resolveRuntimeModule(nodeModules, SHERPA_PACKAGE);
if (!entry) throw new Error(`Unable to resolve ${SHERPA_PACKAGE} in compiled runtime at ${nodeModules}`);
return createRequire(entry)(entry) as SherpaRuntime;
})().catch(error => {
sherpaRuntime = null;
throw error;
});
return sherpaRuntime;
}
function toProgressEvent(modelKey: SttModelKey, info: ProgressInfo): SttProgressEvent {
if (info.status === "ready") {
return { modelKey, status: info.status, task: info.task, model: info.model };
}
if (info.status === "progress_total") {
return {
modelKey,
status: info.status,
name: info.name,
progress: info.progress,
loaded: info.loaded,
total: info.total,
files: info.files,
};
}
if (info.status === "progress") {
return {
modelKey,
status: info.status,
name: info.name,
file: info.file,
progress: info.progress,
loaded: info.loaded,
total: info.total,
};
}
return { modelKey, status: info.status, name: info.name, file: info.file };
}
function sendProgress(transport: SttTransport, id: string, modelKey: SttModelKey, info: ProgressInfo): void {
transport.send({ type: "progress", id, event: toProgressEvent(modelKey, info) });
}
async function loadPipelineOnDevice(
transformers: TransformersRuntime,
spec: TransformersSttModelSpec,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
device: TinyModelDevice,
): Promise<AutomaticSpeechRecognitionPipeline> {
return transformers.pipeline(ASR_TASK, spec.repo, {
device,
dtype: sttModelDtypeOverride ?? spec.dtype,
progress_callback: info => sendProgress(transport, requestId, modelKey, info),
});
}
async function loadPipelineWithDeviceFallback(
transformers: TransformersRuntime,
spec: TransformersSttModelSpec,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
): Promise<{ pipeline: AutomaticSpeechRecognitionPipeline; device: TinyModelDevice }> {
const devices = tinyModelDeviceLoadOrder(sttModelDevicePreference);
if (devices[0] !== sttModelDevicePreference.device) {
sendLog(transport, "warn", "stt: requested device is unsafe in the worker; using CPU", {
modelKey,
repo: spec.repo,
requestedDevice: sttModelDevicePreference.device,
device: devices[0],
});
}
for (let i = 0; i < devices.length; i += 1) {
const device = devices[i]!;
try {
return {
pipeline: await loadPipelineOnDevice(transformers, spec, modelKey, transport, requestId, device),
device,
};
} catch (error) {
if (i === devices.length - 1) throw error;
const fallbackDevice = devices[i + 1]!;
sendLog(transport, "warn", "stt: accelerated device failed; falling back", {
modelKey,
repo: spec.repo,
device,
fallbackDevice,
error: errorMessage(error),
});
}
}
throw new Error("No stt model devices configured");
}
async function loadTransformersModel(
spec: TransformersSttModelSpec,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
): Promise<LoadedModel> {
const transformers = await loadTransformers(transport, requestId, modelKey);
const startedAt = performance.now();
const { pipeline, device } = await loadPipelineWithDeviceFallback(
transformers,
spec,
modelKey,
transport,
requestId,
);
sendLog(transport, "debug", "stt: local model loaded", {
modelKey,
repo: spec.repo,
engine: "transformers",
device,
requestedDevice: sttModelDevicePreference.device,
dtype: sttModelDtypeOverride ?? spec.dtype,
elapsedMs: Math.round(performance.now() - startedAt),
});
return { engine: "transformers", pipeline };
}
/**
* Stream a single sherpa-onnx model file from the Hub into the cache, writing to
* a `.part` sidecar and renaming on completion so an interrupted fetch never
* reads as cached. Emits coalesced per-file progress for the aggregating client.
*/
async function downloadSherpaFile(
repo: string,
filename: string,
dest: string,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
): Promise<void> {
const url = `${HF_RESOLVE_BASE}/${repo}/resolve/main/${filename}`;
const response = await fetch(url, { redirect: "follow" });
if (!response.ok || !response.body) {
throw new Error(`Failed to download ${filename} (${repo}): HTTP ${response.status}`);
}
const total = Number(response.headers.get("content-length") ?? 0);
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status: "download", name: `${repo}/${filename}`, file: filename },
});
const part = `${dest}.part`;
const handle = await fs.open(part, "w");
let loaded = 0;
let lastEmitted = 0;
const reader = response.body.getReader();
try {
for (;;) {
const { done, value } = await reader.read();
if (done) break;
if (!value) continue;
await handle.write(value);
loaded += value.byteLength;
if (loaded - lastEmitted >= PROGRESS_EMIT_BYTES || (total > 0 && loaded >= total)) {
lastEmitted = loaded;
transport.send({
type: "progress",
id: requestId,
event: {
modelKey,
status: "progress",
name: `${repo}/${filename}`,
file: filename,
loaded,
total: total || loaded,
},
});
}
}
} finally {
await handle.close();
}
await fs.rename(part, dest);
}
/**
* Ensure all sherpa-onnx model files for a tier are present in the cache,
* downloading any that are missing, and return their absolute paths.
*/
async function ensureSherpaModelFiles(
spec: SherpaSttModelSpec,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
): Promise<{ encoder: string; decoder: string; joiner: string; tokens: string }> {
const dir = path.join(getTinyModelsCacheDir(), spec.repo);
await fs.mkdir(dir, { recursive: true });
const resolved = {} as { encoder: string; decoder: string; joiner: string; tokens: string };
for (const role in spec.files) {
const key = role as keyof typeof spec.files;
const filename = spec.files[key];
const dest = path.join(dir, filename);
const present = await fs
.stat(dest)
.then(stats => stats.size > 0)
.catch(() => false);
if (!present) await downloadSherpaFile(spec.repo, filename, dest, modelKey, transport, requestId);
resolved[key] = dest;
}
return resolved;
}
async function loadSherpaModel(
spec: SherpaSttModelSpec,
modelKey: SttModelKey,
transport: SttTransport,
requestId: string,
): Promise<LoadedModel> {
const runtime = await loadSherpaRuntime(transport, requestId, modelKey);
const files = await ensureSherpaModelFiles(spec, modelKey, transport, requestId);
const startedAt = performance.now();
const numThreads = Math.max(1, Math.min(4, os.availableParallelism()));
const recognizer = await runtime.OfflineRecognizer.createAsync({
modelConfig: {
transducer: { encoder: files.encoder, decoder: files.decoder, joiner: files.joiner },
tokens: files.tokens,
modelType: spec.modelType,
numThreads,
provider: "cpu",
debug: 0,
},
decodingMethod: "greedy_search",
});
sendLog(transport, "debug", "stt: local model loaded", {
modelKey,
repo: spec.repo,
engine: "sherpa",
provider: "cpu",
numThreads,
elapsedMs: Math.round(performance.now() - startedAt),
});
return { engine: "sherpa", recognizer };
}
async function loadModel(modelKey: SttModelKey, transport: SttTransport, requestId: string): Promise<LoadedModel> {
const spec = getSttModelSpec(modelKey);
if (!spec) throw new Error(`Unknown stt model: ${modelKey}`);
const cached = models.get(modelKey);
if (cached) {
void cached
.then(() => {
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status: "ready", task: ASR_TASK, model: spec.repo },
});
})
.catch(() => undefined);
return cached;
}
const loading =
spec.engine === "sherpa"
? loadSherpaModel(spec, modelKey, transport, requestId)
: loadTransformersModel(spec, modelKey, transport, requestId);
const loaded = loading.then(
model => {
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status: "ready", task: ASR_TASK, model: spec.repo },
});
return model;
},
error => {
models.delete(modelKey);
throw error;
},
);
models.set(modelKey, loaded);
return loaded;
}
async function decodeSegment(
model: LoadedModel,
spec: SttModel,
audio: Float32Array,
language: string | undefined,
): Promise<string> {
if (model.engine === "sherpa") {
const stream = model.recognizer.createStream();
stream.acceptWaveform({ samples: audio, sampleRate: ASR_SAMPLE_RATE });
const result = await model.recognizer.decodeAsync(stream);
return (result.text ?? "").trim();
}
const options: AsrCallOptions = {
chunk_length_s: CHUNK_LENGTH_S,
stride_length_s: STRIDE_LENGTH_S,
return_timestamps: false,
};
// English-only Whisper checkpoints reject `language`/`task`; multilingual ones
// take the configured source language (auto-detected when omitted).
if (!spec.englishOnly) {
options.task = "transcribe";
if (language) options.language = language;
}
const output = (await model.pipeline(audio, options)) as AutomaticSpeechRecognitionOutput;
return (output.text ?? "").trim();
}
async function transcribeAudio(
transport: SttTransport,
requestId: string,
modelKey: SttModelKey,
audio: Float32Array,
language: string | undefined,
): Promise<string> {
const spec = getSttModelSpec(modelKey);
if (!spec) throw new Error(`Unknown stt model: ${modelKey}`);
const model = await loadModel(modelKey, transport, requestId);
return runOnModel(() => decodeSegment(model, spec, audio, language));
}
async function handleBatchRequest(
transport: SttTransport,
request: Extract<SttWorkerInbound, { type: "transcribe" | "download" }>,
): Promise<void> {
try {
if (request.type === "download") {
await loadModel(request.modelKey, transport, request.id);
transport.send({ type: "downloaded", id: request.id });
return;
}
const text = await transcribeAudio(transport, request.id, request.modelKey, request.audio, request.language);
transport.send({ type: "transcription", id: request.id, text });
} catch (error) {
transport.send({ type: "error", id: request.id, error: errorText(error) });
}
}
// ── Live streaming sessions ─────────────────────────────────────────
/** State for one in-flight {@link StreamEndpointer}-driven streaming session. */
interface StreamingSession {
id: string;
spec: SttModel;
language: string | undefined;
model: Promise<LoadedModel>;
endpointer: StreamEndpointer;
/** Finalized segments awaiting decode, in order. */
segmentQueue: Float32Array[];
/** Latest in-progress segment audio awaiting a volatile partial decode (coalesced). */
pendingPartial: Float32Array | null;
/** Committed segment transcripts, joined for the final result. */
committed: string[];
segmentIndex: number;
pumping: boolean;
cancelled: boolean;
ended: boolean;
}
const sessions = new Map<string, StreamingSession>();
function startStreamingSession(
transport: SttTransport,
request: Extract<SttWorkerInbound, { type: "stream_start" }>,
): void {
const spec = getSttModelSpec(request.modelKey);
if (!spec) {
transport.send({ type: "error", id: request.id, error: `Unknown stt model: ${request.modelKey}` });
return;
}
sessions.set(request.id, {
id: request.id,
spec,
language: request.language,
model: loadModel(request.modelKey, transport, request.id),
endpointer: new StreamEndpointer(),
segmentQueue: [],
pendingPartial: null,
committed: [],
segmentIndex: 0,
pumping: false,
cancelled: false,
ended: false,
});
}
function ingestStreamEvents(session: StreamingSession, events: EndpointerEvent[]): void {
for (const event of events) {
if (event.kind === "segment") session.segmentQueue.push(event.audio);
else session.pendingPartial = event.audio;
}
}
/**
* Drain a session's pending work: finalized segments first (committed in order),
* then a single coalesced partial preview. Re-entrant-safe via `pumping`; new
* audio that arrives mid-decode is picked up when the current decode resolves.
*/
async function pumpSession(session: StreamingSession, transport: SttTransport): Promise<void> {
if (session.pumping) return;
session.pumping = true;
try {
const model = await session.model;
while (!session.cancelled) {
if (session.segmentQueue.length > 0) {
const audio = session.segmentQueue.shift()!;
// A fresh segment supersedes any queued preview for the prior one.
session.pendingPartial = null;
const text = await runOnModel(() => decodeSegment(model, session.spec, audio, session.language));
if (session.cancelled) return;
if (text.length > 0) {
session.committed.push(text);
transport.send({ type: "segment", id: session.id, index: session.segmentIndex++, text });
}
continue;
}
if (session.pendingPartial) {
const audio = session.pendingPartial;
session.pendingPartial = null;
const text = await runOnModel(() => decodeSegment(model, session.spec, audio, session.language));
if (session.cancelled) return;
// Skip a now-stale preview if a segment finalized mid-decode.
if (text.length > 0 && session.segmentQueue.length === 0) {
transport.send({ type: "partial", id: session.id, text });
}
continue;
}
break;
}
if (session.ended && !session.cancelled && session.segmentQueue.length === 0 && !session.pendingPartial) {
transport.send({ type: "stream_done", id: session.id, text: session.committed.join(" ") });
sessions.delete(session.id);
}
} catch (error) {
if (!session.cancelled) transport.send({ type: "error", id: session.id, error: errorText(error) });
sessions.delete(session.id);
} finally {
session.pumping = false;
}
}
function handleStreamMessage(
transport: SttTransport,
message: Extract<SttWorkerInbound, { type: "stream_start" | "stream_audio" | "stream_stop" | "stream_cancel" }>,
): void {
if (message.type === "stream_start") {
startStreamingSession(transport, message);
return;
}
const session = sessions.get(message.id);
if (!session || session.cancelled) return;
switch (message.type) {
case "stream_audio":
ingestStreamEvents(session, session.endpointer.push(message.audio));
void pumpSession(session, transport);
return;
case "stream_stop":
session.ended = true;
session.pendingPartial = null;
ingestStreamEvents(session, session.endpointer.flush());
void pumpSession(session, transport);
return;
case "stream_cancel":
session.cancelled = true;
sessions.delete(message.id);
return;
}
}
export function startSttWorker(transport: SttTransport): void {
transport.onMessage(message => {
switch (message.type) {
case "ping":
transport.send({ type: "pong", id: message.id });
return;
case "transcribe":
case "download":
void handleBatchRequest(transport, message);
return;
default:
handleStreamMessage(transport, message);
return;
}
});
}
+108 -48
View File
@@ -1,6 +1,10 @@
import { $which, logger } from "@oh-my-pi/pi-utils";
import { $ } from "bun";
import { resolvePython } from "./transcriber";
import * as fs from "node:fs/promises";
import * as path from "node:path";
import { getTinyModelsCacheDir } from "@oh-my-pi/pi-utils";
import { sttClient } from "./asr-client";
import type { SttProgressStatus } from "./asr-protocol";
import { resolveSttModelSpec } from "./models";
import { ensureRecorder } from "./recorder";
export interface DownloadProgress {
stage: string;
@@ -9,63 +13,119 @@ export interface DownloadProgress {
export interface EnsureOptions {
modelName?: string;
signal?: AbortSignal;
onProgress?: (progress: DownloadProgress) => void;
}
// ── Recording tool ─────────────────────────────────────────────────
// ── ONNX Whisper model ─────────────────────────────────────────────
async function ensureRecordingTool(options?: EnsureOptions): Promise<void> {
if ($which("sox")) return;
if ($which("ffmpeg")) return;
if (process.platform === "linux" && $which("arecord")) return;
// Windows: PowerShell mciSendString is always available as fallback
if (process.platform === "win32") {
// Try to get ffmpeg for better quality, but don't block on failure
options?.onProgress?.({ stage: "Trying to install FFmpeg via winget..." });
const result = await $`winget install --id Gyan.FFmpeg -e --accept-source-agreements --accept-package-agreements`
.quiet()
.nothrow();
if (result.exitCode === 0) {
logger.debug("FFmpeg installed via winget");
}
return;
}
throw new Error(
"No audio recording tool found. Install SoX: sudo apt install sox, or FFmpeg: sudo apt install ffmpeg",
);
/**
* Real-progress event for a speech-model download, surfaced to UI callers.
* `percent` is an integer 0–100 aggregated across all model files (encoder +
* decoder shards), so it advances monotonically toward completion.
*/
export interface SttDownloadProgress {
status: SttProgressStatus;
/** Integer 0–100 aggregated across files. */
percent: number;
/** Bytes downloaded so far across all files. */
loaded: number;
/** Total bytes across all files seen so far. */
total: number;
/** The file currently downloading, when known. */
file?: string;
repo: string;
label: string;
}
// ── Python whisper ─────────────────────────────────────────────────
async function ensurePythonWhisper(options?: EnsureOptions): Promise<void> {
const pythonCmd = resolvePython();
if (!pythonCmd) {
throw new Error("Python not found. Install Python 3.8+ from https://python.org");
/**
* Whether the selected model is already present in the local cache. For
* transformers.js Whisper tiers a complete download leaves `config.json` plus
* the `onnx/` weight files (a bare `config.json` from an interrupted fetch reads
* as not-cached); for sherpa-onnx tiers every model file (encoder/decoder/joiner
* + tokens) must be present (`.part` sidecars from an interrupted fetch are
* ignored).
*/
export async function isSttModelCached(key: string): Promise<boolean> {
const spec = resolveSttModelSpec(key);
const repoDir = path.join(getTinyModelsCacheDir(), spec.repo);
if (spec.engine === "sherpa") {
try {
const root = new Set(await fs.readdir(repoDir));
for (const role in spec.files) {
if (!root.has(spec.files[role as keyof typeof spec.files])) return false;
}
return true;
} catch {
return false;
}
}
try {
const root = await fs.readdir(repoDir);
if (!root.includes("config.json")) return false;
const onnxFiles = await fs.readdir(path.join(repoDir, "onnx")).catch(() => [] as string[]);
return onnxFiles.some(file => file.endsWith(".onnx"));
} catch {
return false;
}
}
// Check if whisper module is already importable
const check = Bun.spawnSync([pythonCmd, "-c", "import whisper"], {
stdout: "pipe",
stderr: "pipe",
/**
* Download (or warm from cache) the selected ONNX Whisper model via the speech
* worker, resolving once the model is fully present and loaded. Streams real
* Hub progress with an aggregated integer percent. Rejects if the worker cannot
* obtain the model. Safe to call non-interactively.
*/
export async function downloadSttModel(
key: string,
onProgress?: (progress: SttDownloadProgress) => void,
options?: { signal?: AbortSignal },
): Promise<void> {
const spec = resolveSttModelSpec(key);
const files = new Map<string, { loaded: number; total: number }>();
const ok = await sttClient.downloadModel(spec.key, {
signal: options?.signal,
onProgress: event => {
if ((event.status === "progress" || event.status === "progress_total") && event.file) {
if (typeof event.loaded === "number" && typeof event.total === "number" && event.total > 0) {
files.set(event.file, { loaded: event.loaded, total: event.total });
}
}
let loaded = 0;
let total = 0;
for (const file of files.values()) {
loaded += file.loaded;
total += file.total;
}
const settled = event.status === "ready" || event.status === "done";
const percent = total > 0 ? Math.min(100, Math.round((loaded / total) * 100)) : settled ? 100 : 0;
onProgress?.({
status: event.status,
percent,
loaded,
total,
file: event.file,
repo: spec.repo,
label: spec.label,
});
},
});
if (check.exitCode === 0) return;
options?.onProgress?.({ stage: "Installing openai-whisper (this may take a few minutes)..." });
logger.debug("Installing openai-whisper via pip");
const install = await $`${pythonCmd} -m pip install -q openai-whisper`.quiet().nothrow();
if (install.exitCode !== 0) {
const stderr = install.stderr.toString().trim();
throw new Error(`Failed to install openai-whisper: ${stderr.split("\n").pop()}`);
}
logger.debug("openai-whisper installed successfully");
if (!ok) throw new Error(`Failed to download speech model (${spec.repo}). Check your network connection.`);
}
// ── Public API ─────────────────────────────────────────────────────
export async function ensureSTTDependencies(options?: EnsureOptions): Promise<void> {
await ensureRecordingTool(options);
await ensurePythonWhisper(options);
await ensureRecorder(progress => options?.onProgress?.(progress), options?.signal);
await downloadSttModel(
resolveSttModelSpec(options?.modelName).key,
progress => {
const stage =
progress.status === "ready" || progress.status === "done"
? `Speech model ${progress.label} ready`
: `Downloading speech model ${progress.label}`;
options?.onProgress?.({ stage, percent: progress.percent });
},
{ signal: options?.signal },
);
}
+259
View File
@@ -0,0 +1,259 @@
/**
* Energy-based speech endpointer for live transcription.
*
* The on-device ASR models we ship are non-streaming: the sherpa-onnx Parakeet
* recognizer and the transformers.js Whisper pipelines both decode a complete
* waveform in one shot. To transcribe *while the user is still speaking*, this
* splits the continuous 16 kHz mono float stream into speech segments at natural
* pauses — each segment is decoded and committed as it finalizes, and the
* in-progress segment is re-decoded periodically for a volatile live preview.
*
* Segmentation is pure short-time-energy VAD with an adaptive noise floor, so it
* needs no extra model and is engine-agnostic (it runs the same way whether the
* downstream model is sherpa or transformers). It is deliberately simple and
* fully deterministic so it can be unit-tested with synthetic signals.
*/
/** Tunable thresholds for {@link StreamEndpointer}. All durations in ms. */
export interface EndpointerConfig {
/** Input sample rate (the recorder always delivers 16 kHz mono). */
sampleRate: number;
/** Short-time analysis frame size. */
frameMs: number;
/** Trailing silence inside a segment that finalizes (commits) it. */
endSilenceMs: number;
/** Shortest speech run that is committed; shorter runs are discarded as noise. */
minSpeechMs: number;
/** Hard cap on segment length so long pause-free speech still commits periodically. */
maxSegmentMs: number;
/** Audio retained before onset so the first phoneme of a segment is never clipped. */
preRollMs: number;
/** Cadence of volatile partial emissions for the in-progress segment. */
partialIntervalMs: number;
/** Speech threshold is `max(minThreshold, noiseFloor * energyRatio)`. */
energyRatio: number;
/** EMA weight tracking the ambient noise floor on non-speech frames. */
floorAttack: number;
/** Absolute RMS floor so a near-silent room never trips speech detection. */
minThreshold: number;
}
export const DEFAULT_ENDPOINTER_CONFIG: EndpointerConfig = {
sampleRate: 16_000,
frameMs: 30,
endSilenceMs: 600,
minSpeechMs: 200,
maxSegmentMs: 12_000,
preRollMs: 240,
partialIntervalMs: 450,
energyRatio: 2.5,
floorAttack: 0.05,
minThreshold: 0.008,
};
/**
* Emitted by {@link StreamEndpointer.push} / {@link StreamEndpointer.flush}.
* `partial` is the volatile in-progress segment (decode and show as preview,
* never commit); `segment` is a finalized run (decode and commit once).
*/
export type EndpointerEvent = { kind: "partial"; audio: Float32Array } | { kind: "segment"; audio: Float32Array };
/** Append-growable Float32 buffer (amortized O(1) push, no per-frame realloc). */
class FloatBuffer {
#data = new Float32Array(0);
#len = 0;
get length(): number {
return this.#len;
}
push(samples: Float32Array): void {
const needed = this.#len + samples.length;
if (needed > this.#data.length) {
const next = new Float32Array(Math.max(this.#data.length * 2, needed, 1 << 14));
next.set(this.#data.subarray(0, this.#len));
this.#data = next;
}
this.#data.set(samples, this.#len);
this.#len += samples.length;
}
/** Copy `[0, end)` into a fresh array the caller can retain. */
take(end = this.#len): Float32Array {
return this.#data.slice(0, Math.max(0, Math.min(end, this.#len)));
}
reset(): void {
this.#len = 0;
}
}
function rms(frame: Float32Array): number {
let sum = 0;
for (let i = 0; i < frame.length; i += 1) sum += frame[i]! * frame[i]!;
return Math.sqrt(sum / Math.max(1, frame.length));
}
export class StreamEndpointer {
readonly #cfg: EndpointerConfig;
readonly #frameSamples: number;
readonly #preRollSamples: number;
#leftover = new Float32Array(0);
#inSpeech = false;
#noiseFloor: number;
#silenceMs = 0;
#segmentMs = 0;
#msSincePartial = 0;
#partialDirty = false;
readonly #segment = new FloatBuffer();
/** Ring of the most recent pre-onset frames, used as segment pre-roll. */
readonly #preRoll = new FloatBuffer();
constructor(config: Partial<EndpointerConfig> = {}) {
this.#cfg = { ...DEFAULT_ENDPOINTER_CONFIG, ...config };
this.#frameSamples = Math.max(1, Math.round((this.#cfg.sampleRate * this.#cfg.frameMs) / 1000));
this.#preRollSamples = Math.max(0, Math.round((this.#cfg.sampleRate * this.#cfg.preRollMs) / 1000));
this.#noiseFloor = this.#cfg.minThreshold;
}
/** Feed newly-captured samples; returns ordered partial/segment events. */
push(samples: Float32Array): EndpointerEvent[] {
const events: EndpointerEvent[] = [];
// Prepend the carried-over tail, then consume whole frames.
let buf: Float32Array;
if (this.#leftover.length === 0) {
buf = samples;
} else {
buf = new Float32Array(this.#leftover.length + samples.length);
buf.set(this.#leftover, 0);
buf.set(samples, this.#leftover.length);
}
let offset = 0;
for (; offset + this.#frameSamples <= buf.length; offset += this.#frameSamples) {
this.#processFrame(buf.subarray(offset, offset + this.#frameSamples), events);
}
this.#leftover = buf.slice(offset);
return events;
}
/** End the stream; returns a trailing committed segment if one is pending. */
flush(): EndpointerEvent[] {
const events: EndpointerEvent[] = [];
if (this.#inSpeech && this.#leftover.length > 0) {
this.#segment.push(this.#leftover);
this.#segmentMs += (this.#leftover.length / this.#cfg.sampleRate) * 1000;
}
this.#leftover = new Float32Array(0);
if (this.#inSpeech) {
const speechMs = this.#segmentMs - this.#silenceMs;
if (speechMs >= this.#cfg.minSpeechMs) {
events.push({ kind: "segment", audio: this.#segment.take(this.#endpointKeep()) });
}
}
this.#reset();
return events;
}
#processFrame(frame: Float32Array, events: EndpointerEvent[]): void {
const energy = rms(frame);
const threshold = Math.max(this.#cfg.minThreshold, this.#noiseFloor * this.#cfg.energyRatio);
const voiced = energy > threshold;
// Track ambient noise on non-speech frames only, so loud speech never
// inflates the floor (which would make the tail of an utterance read as
// silence and clip the segment short).
if (!voiced) {
this.#noiseFloor = this.#noiseFloor * (1 - this.#cfg.floorAttack) + energy * this.#cfg.floorAttack;
}
if (!this.#inSpeech) {
this.#preRoll.push(frame);
// Keep only the most recent pre-roll window.
if (this.#preRoll.length > this.#preRollSamples) {
const tail = this.#preRoll.take().slice(this.#preRoll.length - this.#preRollSamples);
this.#preRoll.reset();
this.#preRoll.push(tail);
}
if (voiced) this.#beginSegment(frame);
return;
}
this.#segment.push(frame);
this.#segmentMs += this.#cfg.frameMs;
this.#msSincePartial += this.#cfg.frameMs;
if (voiced) {
this.#silenceMs = 0;
this.#partialDirty = true;
} else {
this.#silenceMs += this.#cfg.frameMs;
}
if (this.#silenceMs >= this.#cfg.endSilenceMs) {
this.#finalizeSegment(events);
return;
}
if (this.#segmentMs >= this.#cfg.maxSegmentMs) {
// Pause-free long speech: commit what we have and continue a fresh
// segment so output keeps flowing.
events.push({ kind: "segment", audio: this.#segment.take() });
this.#segment.reset();
this.#segmentMs = 0;
this.#silenceMs = 0;
this.#msSincePartial = 0;
this.#partialDirty = false;
return;
}
if (this.#partialDirty && this.#msSincePartial >= this.#cfg.partialIntervalMs) {
events.push({ kind: "partial", audio: this.#segment.take() });
this.#msSincePartial = 0;
this.#partialDirty = false;
}
}
#beginSegment(onsetFrame: Float32Array): void {
this.#inSpeech = true;
this.#segment.reset();
const preRoll = this.#preRoll.take();
if (preRoll.length > 0) this.#segment.push(preRoll);
this.#segment.push(onsetFrame);
this.#preRoll.reset();
this.#silenceMs = 0;
this.#segmentMs = (this.#segment.length / this.#cfg.sampleRate) * 1000;
this.#msSincePartial = 0;
this.#partialDirty = true;
}
#finalizeSegment(events: EndpointerEvent[]): void {
const speechMs = this.#segmentMs - this.#silenceMs;
if (speechMs >= this.#cfg.minSpeechMs) {
events.push({ kind: "segment", audio: this.#segment.take(this.#endpointKeep()) });
}
this.#inSpeech = false;
this.#segment.reset();
this.#silenceMs = 0;
this.#segmentMs = 0;
this.#msSincePartial = 0;
this.#partialDirty = false;
}
/** Samples to keep when committing on silence: drop most of the trailing
* silence but leave a short tail so the final word is not cut. */
#endpointKeep(): number {
const tailMs = Math.min(this.#silenceMs, 120);
const dropMs = Math.max(0, this.#silenceMs - tailMs);
const drop = Math.round((this.#cfg.sampleRate * dropMs) / 1000);
return Math.max(0, this.#segment.length - drop);
}
#reset(): void {
this.#inSpeech = false;
this.#segment.reset();
this.#preRoll.reset();
this.#silenceMs = 0;
this.#segmentMs = 0;
this.#msSincePartial = 0;
this.#partialDirty = false;
this.#noiseFloor = this.#cfg.minThreshold;
}
}
+5 -1
View File
@@ -1,3 +1,7 @@
export * from "./asr-client";
export * from "./asr-protocol";
export * from "./downloader";
export * from "./setup";
export * from "./models";
export * from "./stt-controller";
export * from "./transcriber";
export * from "./wav";
+150
View File
@@ -0,0 +1,150 @@
import type { TinyModelDtype } from "../tiny/dtype";
/**
* On-device speech-to-text model registry. Each tier maps a stable settings key
* onto a locally-runnable ASR model and the engine that loads it:
*
* - `transformers` — a transformers.js / ONNX Whisper repo, loaded by the
* `@huggingface/transformers` `automatic-speech-recognition` pipeline.
* - `sherpa` — a sherpa-onnx (Next-gen Kaldi) offline model, loaded by the
* native `sherpa-onnx-node` addon. Used for NVIDIA Parakeet, the Open ASR
* Leaderboard accuracy/speed leader.
*
* The worker resolves the spec by key and loads the model lazily (kept warm
* afterwards). Both engines run inside the hard-killed subprocess worker.
*/
/** ASR runtime that loads a given tier's model. */
export type SttEngine = "transformers" | "sherpa";
interface SttModelBase {
/** Stable key persisted in `stt.modelName` and sent over the worker protocol. */
key: string;
engine: SttEngine;
/** Hugging Face repo id (transformers.js ONNX repo, or sherpa-onnx model repo). */
repo: string;
/** English-only checkpoint: rejects a configured source `language`. */
englishOnly: boolean;
label: string;
description: string;
/** Approximate on-disk download size for the shipped weights (UI hint). */
sizeHint: string;
}
/** A Whisper-family tier loaded via the transformers.js ASR pipeline. */
export interface TransformersSttModelSpec extends SttModelBase {
engine: "transformers";
/** ONNX precision used unless overridden by `PI_TINY_DTYPE` / `providers.tinyModelDtype`. */
dtype: TinyModelDtype;
}
/** A sherpa-onnx offline tier (e.g. NeMo Parakeet transducer) loaded natively. */
export interface SherpaSttModelSpec extends SttModelBase {
engine: "sherpa";
/** sherpa-onnx offline model family (e.g. `nemo_transducer`). */
modelType: string;
/** Model files (relative to the repo root) fetched into the local cache. */
files: { encoder: string; decoder: string; joiner: string; tokens: string };
}
export type SttModelSpec = TransformersSttModelSpec | SherpaSttModelSpec;
/**
* Speech model tiers, ordered light → SoTA. Defaults to {@link DEFAULT_STT_MODEL_KEY}.
* `fast`/`balanced`/`turbo` are multilingual Whisper checkpoints on transformers.js;
* `parakeet` is NVIDIA Parakeet TDT 0.6B v3 on sherpa-onnx — the Open ASR
* Leaderboard leader (lower WER and far higher throughput than Whisper).
*/
export const STT_MODELS = [
{
key: "fast",
engine: "transformers",
repo: "onnx-community/whisper-base",
dtype: "q8",
englishOnly: false,
label: "Fast (Whisper base)",
description: "Whisper base, multilingual. Smallest + fastest; lowest accuracy. Best for low-resource machines.",
sizeHint: "~60 MB",
},
{
key: "balanced",
engine: "transformers",
repo: "onnx-community/whisper-small",
dtype: "q8",
englishOnly: false,
label: "Balanced (Whisper small)",
description: "Whisper small, multilingual. More accurate than Fast, still light on CPU/RAM.",
sizeHint: "~190 MB",
},
{
key: "turbo",
engine: "transformers",
repo: "onnx-community/whisper-large-v3-turbo",
dtype: "q4",
englishOnly: false,
label: "Turbo (Whisper large-v3)",
description: "Whisper large-v3-turbo, 99 languages. Widest language coverage; large download, slower.",
sizeHint: "~600 MB",
},
{
key: "parakeet",
engine: "sherpa",
repo: "csukuangfj/sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8",
modelType: "nemo_transducer",
files: {
encoder: "encoder.int8.onnx",
decoder: "decoder.int8.onnx",
joiner: "joiner.int8.onnx",
tokens: "tokens.txt",
},
englishOnly: false,
label: "Parakeet TDT v3 (SoTA)",
description:
"NVIDIA Parakeet TDT 0.6B v3, 25 languages. Open ASR Leaderboard leader — best accuracy and far fastest decoding. Default.",
sizeHint: "~680 MB",
},
] as const satisfies readonly SttModelSpec[];
/**
* SoTA default — NVIDIA Parakeet TDT 0.6B v3 (sherpa-onnx). Tops the Open ASR
* Leaderboard on accuracy while decoding ~20× faster than Whisper large-v3.
*/
export const DEFAULT_STT_MODEL_KEY = "parakeet";
export type SttModelKey = (typeof STT_MODELS)[number]["key"];
/** A concrete entry from {@link STT_MODELS}; `key` is the literal tier union. */
export type SttModel = (typeof STT_MODELS)[number];
export const STT_MODEL_VALUES = ["fast", "balanced", "turbo", "parakeet"] as const satisfies readonly SttModelKey[];
type MissingSttModelValue = Exclude<SttModelKey, (typeof STT_MODEL_VALUES)[number]>;
type ExtraSttModelValue = Exclude<(typeof STT_MODEL_VALUES)[number], SttModelKey>;
const STT_MODEL_VALUES_MATCH_REGISTRY: MissingSttModelValue extends never
? ExtraSttModelValue extends never
? true
: never
: never = true;
void STT_MODEL_VALUES_MATCH_REGISTRY;
export const STT_MODEL_OPTIONS = STT_MODELS.map(({ key, label, description }) => ({
value: key,
label,
description,
})) satisfies ReadonlyArray<{ value: SttModelKey; label: string; description: string }>;
export function isSttModelKey(value: string): value is SttModelKey {
return STT_MODELS.some(model => model.key === value);
}
export function getSttModelSpec(key: string): SttModel | undefined {
return STT_MODELS.find(model => model.key === key);
}
/**
* Resolve a (possibly stale or legacy) `stt.modelName` value onto a concrete
* spec, falling back to the SoTA default when the key is unknown.
*/
export function resolveSttModelSpec(key: string | undefined): SttModel {
return (key !== undefined ? getSttModelSpec(key) : undefined) ?? getSttModelSpec(DEFAULT_STT_MODEL_KEY)!;
}
+247 -60
View File
@@ -2,7 +2,9 @@ import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { $which, logger, Snowflake } from "@oh-my-pi/pi-utils";
import { $ } from "bun";
import { $, type Subprocess } from "bun";
import { ensureTool, getToolPath } from "../utils/tools-manager";
import { decodePcmS16LE } from "./wav";
export interface RecordingHandle {
stop(): Promise<void>;
@@ -14,18 +16,13 @@ const isWindows = process.platform === "win32";
* Returns available recording tools in priority order.
*/
export function detectRecordingTools(): string[] {
const tools: string[] = [];
if ($which("sox")) tools.push("sox");
if ($which("ffmpeg")) tools.push("ffmpeg");
if (!isWindows && $which("arecord")) tools.push("arecord");
if (isWindows) tools.push("powershell");
return tools;
return [...new Set(detectRecorders().map(recorder => recorder.tool))];
}
// ── ffmpeg dshow device detection ──────────────────────────────────
async function detectWindowsAudioDevice(): Promise<string> {
const result = await $`ffmpeg -f dshow -list_devices true -i dummy`.quiet().nothrow();
async function detectWindowsAudioDevice(bin: string): Promise<string> {
const result = await $`${bin} -f dshow -list_devices true -i dummy`.quiet().nothrow();
const output = result.stderr.toString();
const audioDevices: string[] = [];
const re = /"([^"]+)"\s*\(audio\)/gi;
@@ -41,11 +38,11 @@ async function detectWindowsAudioDevice(): Promise<string> {
// ── Recording implementations ──────────────────────────────────────
async function startSoxRecording(outputPath: string): Promise<RecordingHandle> {
async function startSoxRecording(bin: string, outputPath: string): Promise<RecordingHandle> {
// On Windows, "-d" (default device) often fails. Use "-t waveaudio 0" for the first input.
const inputArgs = isWindows ? ["-t", "waveaudio", "0"] : ["-d"];
const proc = Bun.spawn(["sox", ...inputArgs, "-r", "16000", "-c", "1", "-b", "16", "-t", "wav", outputPath], {
const proc = Bun.spawn([bin, ...inputArgs, "-r", "16000", "-c", "1", "-b", "16", "-t", "wav", outputPath], {
stdout: "pipe",
stderr: "ignore",
});
@@ -58,12 +55,12 @@ async function startSoxRecording(outputPath: string): Promise<RecordingHandle> {
};
}
async function startFFmpegRecording(outputPath: string): Promise<RecordingHandle> {
async function startFFmpegRecording(bin: string, outputPath: string): Promise<RecordingHandle> {
let args: string[];
if (isWindows) {
const device = await detectWindowsAudioDevice();
const device = await detectWindowsAudioDevice(bin);
args = [
"ffmpeg",
bin,
"-f",
"dshow",
"-i",
@@ -79,7 +76,7 @@ async function startFFmpegRecording(outputPath: string): Promise<RecordingHandle
];
} else if (process.platform === "darwin") {
args = [
"ffmpeg",
bin,
"-f",
"avfoundation",
"-i",
@@ -94,21 +91,7 @@ async function startFFmpegRecording(outputPath: string): Promise<RecordingHandle
outputPath,
];
} else {
args = [
"ffmpeg",
"-f",
"pulse",
"-i",
"default",
"-ar",
"16000",
"-ac",
"1",
"-sample_fmt",
"s16",
"-y",
outputPath,
];
args = [bin, "-f", "pulse", "-i", "default", "-ar", "16000", "-ac", "1", "-sample_fmt", "s16", "-y", outputPath];
}
const proc = Bun.spawn(args, {
@@ -133,8 +116,8 @@ async function startFFmpegRecording(outputPath: string): Promise<RecordingHandle
};
}
async function startArecordRecording(outputPath: string): Promise<RecordingHandle> {
const proc = Bun.spawn(["arecord", "-f", "S16_LE", "-r", "16000", "-c", "1", outputPath], {
async function startArecordRecording(bin: string, outputPath: string): Promise<RecordingHandle> {
const proc = Bun.spawn([bin, "-f", "S16_LE", "-r", "16000", "-c", "1", outputPath], {
stdout: "pipe",
stderr: "ignore",
});
@@ -277,7 +260,9 @@ async function startPowerShellRecording(outputPath: string): Promise<RecordingHa
// ── Health check ───────────────────────────────────────────────────
async function verifyProcessAlive(proc: ReturnType<typeof Bun.spawn>, tool: string): Promise<void> {
type RecorderProcess = Subprocess<"ignore" | "pipe", "pipe", "ignore">;
async function verifyProcessAlive(proc: RecorderProcess, tool: string): Promise<void> {
await Bun.sleep(300);
const exited = await Promise.race([proc.exited.then(code => code), Bun.sleep(0).then(() => "running" as const)]);
@@ -293,38 +278,101 @@ async function verifyProcessAlive(proc: ReturnType<typeof Bun.spawn>, tool: stri
// ── Public API ─────────────────────────────────────────────────────
export async function startRecording(outputPath: string): Promise<RecordingHandle> {
const tools = detectRecordingTools();
if (tools.length === 0) {
throw new Error(
isWindows
? "No audio recording tool found. Install FFmpeg or SoX and add to PATH."
: "No audio recording tool found. Install SoX: sudo apt install sox, or FFmpeg: sudo apt install ffmpeg",
);
export interface ResolvedRecorder {
tool: "sox" | "ffmpeg" | "arecord" | "powershell";
bin: string;
}
/**
* Resolve a usable recorder without triggering any download. Priority:
* sox (PATH) → ffmpeg (PATH or previously-downloaded static binary) →
* arecord (PATH, non-Windows) → PowerShell mci fallback (Windows) → none.
*/
function detectRecorders(): ResolvedRecorder[] {
const recorders: ResolvedRecorder[] = [];
const sox = $which("sox");
if (sox) recorders.push({ tool: "sox", bin: sox });
const pathFfmpeg = $which("ffmpeg");
if (pathFfmpeg) recorders.push({ tool: "ffmpeg", bin: pathFfmpeg });
const bundledFfmpeg = getToolPath("ffmpeg");
if (bundledFfmpeg && bundledFfmpeg !== pathFfmpeg) recorders.push({ tool: "ffmpeg", bin: bundledFfmpeg });
if (!isWindows) {
const arecord = $which("arecord");
if (arecord) recorders.push({ tool: "arecord", bin: arecord });
}
const errors: string[] = [];
for (const tool of tools) {
logger.debug("Trying audio recording", { tool, outputPath });
if (isWindows) recorders.push({ tool: "powershell", bin: "powershell" });
return recorders;
}
export function detectRecorder(): ResolvedRecorder | null {
return detectRecorders()[0] ?? null;
}
/**
* Ensure a recorder is available, downloading the static ffmpeg binary when
* nothing is already present. Returns the resolved recorder.
*/
export async function ensureRecorder(
onProgress?: (p: { stage: string; percent?: number }) => void,
signal?: AbortSignal,
): Promise<ResolvedRecorder> {
const existing = detectRecorder();
if (existing) return existing;
const bin = await ensureTool("ffmpeg", { signal, notify: m => onProgress?.({ stage: m }) });
if (bin) return { tool: "ffmpeg", bin };
if (isWindows) return { tool: "powershell", bin: "powershell" };
throw new Error(
"No audio recorder available and automatic ffmpeg download failed. " +
"Install SoX or FFmpeg manually and add it to PATH.",
);
}
function recorderFailure(recorder: ResolvedRecorder, error: unknown): string {
const message = error instanceof Error ? error.message : String(error);
return `${recorder.tool} (${recorder.bin}): ${message}`;
}
async function startRecordingWithRecorder(recorder: ResolvedRecorder, outputPath: string): Promise<RecordingHandle> {
logger.debug("Starting audio recording", { tool: recorder.tool, bin: recorder.bin, outputPath });
switch (recorder.tool) {
case "sox":
return startSoxRecording(recorder.bin, outputPath);
case "ffmpeg":
return startFFmpegRecording(recorder.bin, outputPath);
case "arecord":
return startArecordRecording(recorder.bin, outputPath);
case "powershell":
return startPowerShellRecording(outputPath);
}
}
export async function startRecording(outputPath: string): Promise<RecordingHandle> {
const recorders = detectRecorders();
if (recorders.length === 0) {
throw new Error("No audio recorder available — run `omp setup speech`");
}
const failures: string[] = [];
for (const recorder of recorders) {
try {
switch (tool) {
case "sox":
return await startSoxRecording(outputPath);
case "ffmpeg":
return await startFFmpegRecording(outputPath);
case "arecord":
return await startArecordRecording(outputPath);
case "powershell":
return await startPowerShellRecording(outputPath);
}
} catch (err) {
const msg = err instanceof Error ? err.message : String(err);
logger.debug(`Recording tool ${tool} failed, trying next`, { error: msg });
errors.push(`${tool}: ${msg}`);
return await startRecordingWithRecorder(recorder, outputPath);
} catch (error) {
const failure = recorderFailure(recorder, error);
failures.push(failure);
logger.warn("STT recorder failed to start; trying fallback", {
recorder: recorder.tool,
bin: recorder.bin,
error: failure,
});
}
}
throw new Error(`All recording tools failed:\n${errors.join("\n")}`);
throw new Error(`No audio recorder could start — run \`omp setup speech\`.\n${failures.join("\n")}`);
}
/**
@@ -349,3 +397,142 @@ export async function verifyRecordingFile(filePath: string): Promise<number> {
);
}
}
// ── Streaming (live) capture ───────────────────────────────────────
export interface StreamingRecordingHandle {
stop(): Promise<void>;
}
/** Build the argv for a recorder that emits raw 16 kHz mono s16le PCM to stdout. */
async function streamingRecorderArgs(recorder: ResolvedRecorder): Promise<string[]> {
const { tool, bin } = recorder;
switch (tool) {
case "sox": {
const input = isWindows ? ["-t", "waveaudio", "0"] : ["-d"];
return [bin, ...input, "-r", "16000", "-c", "1", "-b", "16", "-e", "signed-integer", "-t", "raw", "-"];
}
case "arecord":
return [bin, "-f", "S16_LE", "-r", "16000", "-c", "1", "-t", "raw", "-"];
case "ffmpeg": {
const input = isWindows
? ["-f", "dshow", "-i", `audio=${await detectWindowsAudioDevice(bin)}`]
: process.platform === "darwin"
? ["-f", "avfoundation", "-i", ":0"]
: ["-f", "pulse", "-i", "default"];
return [bin, ...input, "-ar", "16000", "-ac", "1", "-f", "s16le", "pipe:1"];
}
case "powershell":
throw new Error("PowerShell recorder cannot stream PCM to a pipe");
}
}
/**
* Start a recorder that streams raw 16 kHz mono s16le PCM to stdout, decoding it
* to float frames delivered through `onAudio` as they arrive. Returns `null`
* when the only available recorder (Windows PowerShell mci) records to a file
* and cannot pipe — the caller then falls back to file-based batch capture.
*/
async function startStreamingRecordingWithRecorder(
recorder: ResolvedRecorder,
onAudio: (samples: Float32Array) => void,
): Promise<StreamingRecordingHandle> {
const args = await streamingRecorderArgs(recorder);
logger.debug("Starting streaming audio recording", { tool: recorder.tool, bin: recorder.bin });
const proc = Bun.spawn(args, { stdin: "pipe", stdout: "pipe", stderr: "ignore" });
// Read s16le bytes off stdout, carrying any trailing odd byte across chunk
// boundaries so a sample is never split. Runs until the process closes stdout.
const reader = (proc.stdout as ReadableStream<Uint8Array>).getReader();
let leftover: Uint8Array | null = null;
const pump = async (): Promise<void> => {
try {
for (;;) {
const { done, value } = await reader.read();
if (done) break;
if (!value || value.length === 0) continue;
let bytes = value;
if (leftover) {
const merged = new Uint8Array(leftover.length + value.length);
merged.set(leftover, 0);
merged.set(value, leftover.length);
bytes = merged;
leftover = null;
}
const usable = bytes.length - (bytes.length % 2);
if (usable < bytes.length) leftover = bytes.slice(usable);
if (usable > 0) onAudio(decodePcmS16LE(bytes.subarray(0, usable)));
}
} catch (error) {
logger.debug("stt: streaming recorder read ended", {
error: error instanceof Error ? error.message : String(error),
});
}
};
void pump();
try {
await verifyProcessAlive(proc, recorder.tool);
} catch (error) {
try {
proc.kill("SIGKILL");
} catch {
// Already gone.
}
throw error;
}
let stopped = false;
return {
async stop() {
if (stopped) return;
stopped = true;
if (recorder.tool === "ffmpeg") {
try {
proc.stdin.write("q");
proc.stdin.end();
} catch {
// stdin may already be closed.
}
const killTimer = setTimeout(() => proc.kill(), 3000);
await proc.exited;
clearTimeout(killTimer);
} else {
proc.kill("SIGTERM");
await proc.exited;
}
try {
await reader.cancel();
} catch {
// Reader already released when stdout closed.
}
},
};
}
export async function startStreamingRecording(
onAudio: (samples: Float32Array) => void,
): Promise<StreamingRecordingHandle | null> {
const recorders = detectRecorders();
if (recorders.length === 0) {
throw new Error("No audio recorder available — run `omp setup speech`");
}
const streamingRecorders = recorders.filter(recorder => recorder.tool !== "powershell");
if (streamingRecorders.length === 0) return null;
const failures: string[] = [];
for (const recorder of streamingRecorders) {
try {
return await startStreamingRecordingWithRecorder(recorder, onAudio);
} catch (error) {
const failure = recorderFailure(recorder, error);
failures.push(failure);
logger.warn("STT streaming recorder failed to start; trying fallback", {
recorder: recorder.tool,
bin: recorder.bin,
error: failure,
});
}
}
throw new Error(`No streaming audio recorder could start.\n${failures.join("\n")}`);
}
-52
View File
@@ -1,52 +0,0 @@
import { detectRecordingTools } from "./recorder";
import { resolvePython } from "./transcriber";
const isWindows = process.platform === "win32";
export interface STTDependencyStatus {
recorder: { available: boolean; tool: string | null; installHint: string };
python: { available: boolean; path: string | null; installHint: string };
whisper: { available: boolean; installHint: string };
}
export async function checkDependencies(): Promise<STTDependencyStatus> {
const recorderTools = detectRecordingTools();
const recorderHint = isWindows
? "PowerShell fallback available. For better quality: install SoX or FFmpeg."
: "Install SoX: sudo apt install sox, or FFmpeg: sudo apt install ffmpeg";
const pythonCmd = resolvePython();
const pythonHint = "Install Python 3.8+ from https://python.org";
let whisperAvailable = false;
if (pythonCmd) {
const check = Bun.spawnSync([pythonCmd, "-c", "import whisper"], {
stdout: "pipe",
stderr: "pipe",
});
whisperAvailable = check.exitCode === 0;
}
const whisperHint = "Run 'omp setup stt' to auto-install, or: pip install openai-whisper";
return {
recorder: { available: recorderTools.length > 0, tool: recorderTools[0] ?? null, installHint: recorderHint },
python: { available: pythonCmd !== null, path: pythonCmd, installHint: pythonHint },
whisper: { available: whisperAvailable, installHint: whisperHint },
};
}
export function formatDependencyStatus(status: STTDependencyStatus): string {
const lines: string[] = ["STT Dependencies:"];
const check = (ok: boolean) => (ok ? "[ok]" : "[missing]");
lines.push(` Recorder: ${check(status.recorder.available)} ${status.recorder.tool ?? "none"}`);
if (!status.recorder.available) lines.push(` -> ${status.recorder.installHint}`);
lines.push(` Python: ${check(status.python.available)} ${status.python.path ?? "none"}`);
if (!status.python.available) lines.push(` -> ${status.python.installHint}`);
lines.push(` Whisper: ${check(status.whisper.available)}`);
if (!status.whisper.available) lines.push(` -> ${status.whisper.installHint}`);
return lines.join("\n");
}
+201 -22
View File
@@ -3,8 +3,17 @@ import * as os from "node:os";
import * as path from "node:path";
import { logger, Snowflake } from "@oh-my-pi/pi-utils";
import { settings } from "../config/settings";
import { type SttStreamHandle, sttClient } from "./asr-client";
import { ensureSTTDependencies } from "./downloader";
import { type RecordingHandle, startRecording, verifyRecordingFile } from "./recorder";
import { resolveSttModelSpec } from "./models";
import {
detectRecorder,
type RecordingHandle,
type StreamingRecordingHandle,
startRecording,
startStreamingRecording,
verifyRecordingFile,
} from "./recorder";
import { transcribe } from "./transcriber";
export type SttState = "idle" | "recording" | "transcribing";
@@ -13,21 +22,37 @@ interface ToggleOptions {
showWarning(msg: string): void;
showStatus(msg: string): void;
onStateChange(state: SttState): void;
/** Force a redraw after async edits to the composer (live segment/preview inserts). */
requestRender?(): void;
}
/** The slice of the composer editor the controller drives. */
interface Editor {
insertText(text: string): void;
setVolatileText(text: string): void;
clearVolatileText(): void;
commitVolatileText(text: string): void;
}
export class STTController {
#state: SttState = "idle";
#recordingHandle: RecordingHandle | null = null;
#tempFile: string | null = null;
#depsResolved = false;
#toggling = false;
#stopAfterStart = false;
#disposed = false;
// Batch (single-shot) capture.
#recordingHandle: RecordingHandle | null = null;
#tempFile: string | null = null;
#transcriptionAbort: AbortController | null = null;
// Live streaming capture.
#stream: SttStreamHandle | null = null;
#streamRecorder: StreamingRecordingHandle | null = null;
#streamEditor: Editor | null = null;
#streamCommitted = false;
#streamAbort: AbortController | null = null;
get state(): SttState {
return this.#state;
}
@@ -38,45 +63,192 @@ export class STTController {
}
async toggle(editor: Editor, options: ToggleOptions): Promise<void> {
if (this.#toggling) return;
if (this.#toggling) {
if (this.#state === "idle" || this.#state === "recording") this.#stopAfterStart = true;
return;
}
this.#toggling = true;
try {
switch (this.#state) {
case "idle":
await this.#startRecording(options);
await this.#start(editor, options);
break;
case "recording":
await this.#stopAndTranscribe(editor, options);
await this.#stop(editor, options);
break;
case "transcribing":
options.showStatus("Transcription in progress...");
break;
}
if (this.#stopAfterStart && this.#state === "recording") {
this.#stopAfterStart = false;
await this.#stop(editor, options);
} else if (this.#state !== "recording") {
this.#stopAfterStart = false;
}
} finally {
this.#toggling = false;
}
}
async #startRecording(options: ToggleOptions): Promise<void> {
if (!this.#depsResolved) {
try {
options.showStatus("Checking STT dependencies...");
await ensureSTTDependencies({
modelName: settings.get("stt.modelName") as string | undefined,
onProgress: p => options.showStatus(p.stage + (p.percent != null ? ` (${p.percent}%)` : "")),
});
options.showStatus("");
this.#depsResolved = true;
} catch (err) {
const msg = err instanceof Error ? err.message : "Failed to setup STT dependencies";
async #ensureDeps(options: ToggleOptions): Promise<boolean> {
if (this.#depsResolved) return true;
try {
options.showStatus("Checking STT dependencies...");
await ensureSTTDependencies({
modelName: settings.get("stt.modelName") as string | undefined,
onProgress: p => options.showStatus(p.stage + (p.percent != null ? ` (${p.percent}%)` : "")),
});
options.showStatus("");
this.#depsResolved = true;
return true;
} catch (err) {
const msg = err instanceof Error ? err.message : "Failed to setup STT dependencies";
options.showWarning(msg);
logger.error("STT dependency setup failed", { error: msg });
return false;
}
}
async #start(editor: Editor, options: ToggleOptions): Promise<void> {
if (!(await this.#ensureDeps(options))) return;
// Live transcription needs a recorder that can pipe PCM; the Windows
// PowerShell mci fallback records to a file, so it stays single-shot.
if (this.#recorderCanStream()) {
await this.#startStreaming(editor, options);
return;
}
await this.#startBatchRecording(options);
}
async #stop(editor: Editor, options: ToggleOptions): Promise<void> {
if (this.#stream) {
await this.#stopStreaming(options);
return;
}
await this.#stopBatch(editor, options);
}
// ── Live streaming ──────────────────────────────────────────────
#recorderCanStream(): boolean {
const recorder = detectRecorder();
return recorder !== null && recorder.tool !== "powershell";
}
/** Segment text gets a leading space once a prior segment is committed, so
* phrases join naturally; the first phrase is inserted at the cursor as-is. */
#prefixed(text: string): string {
const normalized = text.replace(/\s+/g, " ").trim();
if (!normalized) return "";
return this.#streamCommitted ? ` ${normalized}` : normalized;
}
async #startStreaming(editor: Editor, options: ToggleOptions): Promise<void> {
const modelKey = resolveSttModelSpec(settings.get("stt.modelName") as string | undefined).key;
const language = settings.get("stt.language") as string | undefined;
this.#streamEditor = editor;
this.#streamCommitted = false;
this.#streamAbort = new AbortController();
const stream = sttClient.startStream(modelKey, {
language: language || undefined,
signal: this.#streamAbort.signal,
onPartial: text => {
if (this.#disposed || this.#state !== "recording") return;
this.#streamEditor?.setVolatileText(this.#prefixed(text));
options.requestRender?.();
},
onSegment: text => {
if (this.#disposed) return;
const prefixed = this.#prefixed(text);
if (prefixed) {
this.#streamEditor?.commitVolatileText(prefixed);
this.#streamCommitted = true;
} else {
this.#streamEditor?.clearVolatileText();
}
options.requestRender?.();
},
});
this.#stream = stream;
let recorder: StreamingRecordingHandle | null = null;
try {
recorder = await startStreamingRecording(samples => stream.pushAudio(samples));
} catch (err) {
logger.warn("STT streaming recorder failed to start; falling back to batch recording", {
error: err instanceof Error ? err.message : String(err),
});
}
if (!recorder) {
stream.cancel();
this.#cleanupStream();
await this.#startBatchRecording(options);
return;
}
this.#streamRecorder = recorder;
this.#setState("recording", options);
logger.debug("STT live recording started", { modelKey });
}
async #stopStreaming(options: ToggleOptions): Promise<void> {
const stream = this.#stream;
const recorder = this.#streamRecorder;
if (!stream) {
this.#setState("idle", options);
return;
}
this.#setState("transcribing", options);
// Stop the mic first so no further audio is fed, then flush the worker.
try {
await recorder?.stop();
} catch (err) {
logger.debug("stt: streaming recorder stop failed", {
error: err instanceof Error ? err.message : String(err),
});
}
this.#streamRecorder = null;
let failed = false;
let finalText = "";
try {
finalText = (await stream.stop()).trim();
} catch (err) {
failed = true;
if (!this.#disposed) {
const msg = err instanceof Error ? err.message : "Transcription failed";
options.showWarning(msg);
logger.error("STT dependency setup failed", { error: msg });
return;
logger.error("STT live transcription failed", { error: msg });
}
}
if (this.#disposed) {
this.#cleanupStream();
return;
}
if (!this.#streamCommitted && finalText) {
this.#streamEditor?.commitVolatileText(this.#prefixed(finalText));
this.#streamCommitted = true;
} else {
this.#streamEditor?.clearVolatileText();
}
options.requestRender?.();
if (!failed) options.showStatus(this.#streamCommitted ? "" : "No speech detected.");
this.#cleanupStream();
this.#setState("idle", options);
}
#cleanupStream(): void {
this.#stream = null;
this.#streamRecorder = null;
this.#streamEditor = null;
this.#streamCommitted = false;
this.#streamAbort = null;
}
// ── Batch (single-shot) ─────────────────────────────────────────
async #startBatchRecording(options: ToggleOptions): Promise<void> {
const id = Snowflake.next();
this.#tempFile = path.join(os.tmpdir(), `omp-stt-${id}.wav`);
try {
this.#recordingHandle = await startRecording(this.#tempFile);
this.#setState("recording", options);
@@ -89,7 +261,7 @@ export class STTController {
}
}
async #stopAndTranscribe(editor: Editor, options: ToggleOptions): Promise<void> {
async #stopBatch(editor: Editor, options: ToggleOptions): Promise<void> {
const handle = this.#recordingHandle;
const tempFile = this.#tempFile;
this.#recordingHandle = null;
@@ -146,6 +318,13 @@ export class STTController {
this.#transcriptionAbort.abort();
this.#transcriptionAbort = null;
}
if (this.#streamAbort) {
this.#streamAbort.abort();
this.#streamAbort = null;
}
this.#stream?.cancel();
this.#streamRecorder?.stop().catch(() => {});
this.#cleanupStream();
if (this.#recordingHandle) {
this.#recordingHandle.stop().catch(() => {});
this.#recordingHandle = null;
+488
View File
@@ -0,0 +1,488 @@
import { describe, expect, it } from "bun:test";
import { SttClient } from "./asr-client";
import type { SttTransport, SttWorkerInbound, SttWorkerOutbound } from "./asr-protocol";
import { startSttWorker } from "./asr-worker";
import { StreamEndpointer } from "./endpointer";
import { DEFAULT_STT_MODEL_KEY, getSttModelSpec, isSttModelKey, resolveSttModelSpec } from "./models";
import { decodePcmS16LE, decodeWavToMono16k, resampleLinear, TARGET_SAMPLE_RATE } from "./wav";
// ── WAV building helpers ────────────────────────────────────────────
function writeAscii(view: DataView, offset: number, text: string): void {
for (let i = 0; i < text.length; i += 1) view.setUint8(offset + i, text.charCodeAt(i));
}
function buildWav(opts: {
sampleRate: number;
channels: number;
bitsPerSample: number;
format?: number;
body: Uint8Array;
}): ArrayBuffer {
const { sampleRate, channels, bitsPerSample, format = 1, body } = opts;
const blockAlign = channels * (bitsPerSample / 8);
const buffer = new ArrayBuffer(44 + body.length);
const view = new DataView(buffer);
writeAscii(view, 0, "RIFF");
view.setUint32(4, 36 + body.length, true);
writeAscii(view, 8, "WAVE");
writeAscii(view, 12, "fmt ");
view.setUint32(16, 16, true);
view.setUint16(20, format, true);
view.setUint16(22, channels, true);
view.setUint32(24, sampleRate, true);
view.setUint32(28, sampleRate * blockAlign, true);
view.setUint16(32, blockAlign, true);
view.setUint16(34, bitsPerSample, true);
writeAscii(view, 36, "data");
view.setUint32(40, body.length, true);
new Uint8Array(buffer).set(body, 44);
return buffer;
}
function bytesFromInt16(values: number[]): Uint8Array {
const out = new Uint8Array(values.length * 2);
const view = new DataView(out.buffer);
values.forEach((value, i) => {
view.setInt16(i * 2, value, true);
});
return out;
}
function bytesFromInt32(values: number[]): Uint8Array {
const out = new Uint8Array(values.length * 4);
const view = new DataView(out.buffer);
values.forEach((value, i) => {
view.setInt32(i * 4, value, true);
});
return out;
}
// ── WAV decoding ────────────────────────────────────────────────────
describe("decodeWavToMono16k", () => {
it("decodes 16-bit PCM to normalized [-1, 1] floats", () => {
const wav = buildWav({
sampleRate: TARGET_SAMPLE_RATE,
channels: 1,
bitsPerSample: 16,
body: bytesFromInt16([16_384, -16_384, 0]),
});
const audio = decodeWavToMono16k(wav);
expect(audio.length).toBe(3);
expect(audio[0]).toBeCloseTo(0.5, 4);
expect(audio[1]).toBeCloseTo(-0.5, 4);
expect(audio[2]).toBeCloseTo(0, 6);
});
it("decodes 8-bit unsigned PCM centered at 128", () => {
const wav = buildWav({
sampleRate: TARGET_SAMPLE_RATE,
channels: 1,
bitsPerSample: 8,
body: new Uint8Array([192, 64, 128]),
});
const audio = decodeWavToMono16k(wav);
expect(audio[0]).toBeCloseTo(0.5, 4);
expect(audio[1]).toBeCloseTo(-0.5, 4);
expect(audio[2]).toBeCloseTo(0, 6);
});
it("decodes 32-bit PCM", () => {
const wav = buildWav({
sampleRate: TARGET_SAMPLE_RATE,
channels: 1,
bitsPerSample: 32,
body: bytesFromInt32([1_073_741_824, -1_073_741_824]),
});
const audio = decodeWavToMono16k(wav);
expect(audio[0]).toBeCloseTo(0.5, 4);
expect(audio[1]).toBeCloseTo(-0.5, 4);
});
it("mixes stereo down to mono by averaging channels", () => {
// Two stereo frames: (L=1.0, R=0.0) -> 0.5, (L=-1.0, R=0.0) -> -0.5
const wav = buildWav({
sampleRate: TARGET_SAMPLE_RATE,
channels: 2,
bitsPerSample: 16,
body: bytesFromInt16([32_767, 0, -32_768, 0]),
});
const audio = decodeWavToMono16k(wav);
expect(audio.length).toBe(2);
expect(audio[0]).toBeCloseTo(0.5, 3);
expect(audio[1]).toBeCloseTo(-0.5, 3);
});
it("resamples to 16 kHz, preserving endpoints and bounds", () => {
const frames = 8;
const values: number[] = [];
for (let i = 0; i < frames; i += 1) values.push(Math.round(Math.sin((i / frames) * Math.PI) * 20_000));
const wav = buildWav({ sampleRate: 8_000, channels: 1, bitsPerSample: 16, body: bytesFromInt16(values) });
const audio = decodeWavToMono16k(wav);
// 8 kHz -> 16 kHz doubles the sample count.
expect(audio.length).toBe(frames * 2);
expect(audio[0]).toBeCloseTo(values[0]! / 32_768, 4);
expect(audio[audio.length - 1]).toBeCloseTo(values[frames - 1]! / 32_768, 4);
for (const sample of audio) expect(Math.abs(sample)).toBeLessThanOrEqual(1);
});
it("rejects non-RIFF input", () => {
expect(() => decodeWavToMono16k(new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]).buffer)).toThrow();
});
});
describe("resampleLinear", () => {
it("interpolates between samples", () => {
const out = resampleLinear(new Float32Array([0, 1]), 1, 2);
expect(out.length).toBe(4);
expect(out[0]).toBeCloseTo(0, 6);
expect(out[1]).toBeCloseTo(1 / 3, 5);
expect(out[2]).toBeCloseTo(2 / 3, 5);
expect(out[3]).toBeCloseTo(1, 6);
});
it("returns the input unchanged when rates match", () => {
const input = new Float32Array([0.1, 0.2]);
expect(resampleLinear(input, 16_000, 16_000)).toBe(input);
});
});
// ── Model key → repo mapping ────────────────────────────────────────
describe("stt model registry", () => {
it("maps tier keys onto their engine + repo", () => {
expect(getSttModelSpec("fast")?.repo).toBe("onnx-community/whisper-base");
expect(getSttModelSpec("fast")?.engine).toBe("transformers");
expect(getSttModelSpec("balanced")?.repo).toBe("onnx-community/whisper-small");
expect(getSttModelSpec("turbo")?.repo).toBe("onnx-community/whisper-large-v3-turbo");
expect(getSttModelSpec("turbo")?.engine).toBe("transformers");
expect(getSttModelSpec("parakeet")?.repo).toBe("csukuangfj/sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8");
expect(getSttModelSpec("parakeet")?.engine).toBe("sherpa");
expect(getSttModelSpec("nonexistent")).toBeUndefined();
});
it("defaults to the SoTA Parakeet tier", () => {
expect(DEFAULT_STT_MODEL_KEY).toBe("parakeet");
expect(getSttModelSpec(DEFAULT_STT_MODEL_KEY)?.engine).toBe("sherpa");
});
it("describes the sherpa tier's model files for the downloader/worker", () => {
const spec = getSttModelSpec("parakeet");
expect(spec?.engine).toBe("sherpa");
if (spec?.engine !== "sherpa") throw new Error("expected sherpa spec");
expect(spec.modelType).toBe("nemo_transducer");
expect(spec.files).toEqual({
encoder: "encoder.int8.onnx",
decoder: "decoder.int8.onnx",
joiner: "joiner.int8.onnx",
tokens: "tokens.txt",
});
});
it("flags multilingual checkpoints (no English-only tiers in the ladder)", () => {
expect(getSttModelSpec("fast")?.englishOnly).toBe(false);
expect(getSttModelSpec("turbo")?.englishOnly).toBe(false);
expect(getSttModelSpec("parakeet")?.englishOnly).toBe(false);
});
it("resolves legacy/unknown keys to the SoTA default", () => {
expect(resolveSttModelSpec("base.en").key).toBe("parakeet"); // legacy whisper size
expect(resolveSttModelSpec("accurate").key).toBe("parakeet"); // retired tier key
expect(resolveSttModelSpec("fast-en").key).toBe("parakeet"); // retired tier key
expect(resolveSttModelSpec(undefined).key).toBe("parakeet");
expect(resolveSttModelSpec("balanced").key).toBe("balanced");
});
it("recognizes valid tier keys", () => {
expect(isSttModelKey("balanced")).toBe(true);
expect(isSttModelKey("parakeet")).toBe(true);
expect(isSttModelKey("accurate")).toBe(false);
});
});
// ── Protocol round-trips ────────────────────────────────────────────
interface FakeWorker {
send(message: SttWorkerInbound): void;
onMessage(handler: (message: SttWorkerOutbound) => void): () => void;
onError(handler: (error: Error) => void): () => void;
terminate(): Promise<void>;
}
function makeFakeWorker(
onSend: (message: SttWorkerInbound, emit: (out: SttWorkerOutbound) => void) => void,
): FakeWorker {
let emit: (out: SttWorkerOutbound) => void = () => {};
return {
send(message) {
onSend(message, out => emit(out));
},
onMessage(handler) {
emit = handler;
return () => {
emit = () => {};
};
},
onError() {
return () => {};
},
async terminate() {},
};
}
describe("worker protocol (in-memory transport)", () => {
it("answers ping with pong via startSttWorker", () => {
const sent: SttWorkerOutbound[] = [];
let handler: ((message: SttWorkerInbound) => void) | undefined;
const transport: SttTransport = {
send: message => sent.push(message),
onMessage: h => {
handler = h;
return () => {};
},
};
startSttWorker(transport);
handler?.({ type: "ping", id: "p1" });
expect(sent).toEqual([{ type: "pong", id: "p1" }]);
});
});
describe("SttClient request/response correlation", () => {
it("resolves a transcribe request with the matching transcription", async () => {
const client = new SttClient(() =>
makeFakeWorker((message, emit) => {
if (message.type === "transcribe") emit({ type: "transcription", id: message.id, text: " hi there " });
}),
);
// Client returns the raw worker text (transcriber.ts owns trimming).
expect(await client.transcribe("fast", new Float32Array([0.1]), {})).toBe(" hi there ");
await client.terminate();
});
it("rejects with the worker error message", async () => {
const client = new SttClient(() =>
makeFakeWorker((message, emit) => {
if (message.type === "transcribe") emit({ type: "error", id: message.id, error: "model exploded" });
}),
);
await expect(client.transcribe("fast", new Float32Array([0]), {})).rejects.toThrow("model exploded");
await client.terminate();
});
it("throws synchronously on an already-aborted signal", async () => {
const controller = new AbortController();
controller.abort();
const client = new SttClient(() => makeFakeWorker(() => {}));
await expect(client.transcribe("fast", new Float32Array([0]), { signal: controller.signal })).rejects.toThrow();
await client.terminate();
});
it("correlates concurrent requests by id regardless of reply order", async () => {
const inbound: Array<Extract<SttWorkerInbound, { type: "transcribe" }>> = [];
let emit: ((out: SttWorkerOutbound) => void) | undefined;
const client = new SttClient(() =>
makeFakeWorker((message, e) => {
emit = e;
if (message.type === "transcribe") inbound.push(message);
}),
);
const first = client.transcribe("fast", new Float32Array([0]), { language: "a" });
const second = client.transcribe("fast", new Float32Array([0]), { language: "b" });
expect(inbound.length).toBe(2);
expect(inbound[0]!.id).not.toBe(inbound[1]!.id);
// Reply out of order.
emit?.({ type: "transcription", id: inbound[1]!.id, text: "second" });
emit?.({ type: "transcription", id: inbound[0]!.id, text: "first" });
expect(await first).toBe("first");
expect(await second).toBe("second");
await client.terminate();
});
it("downloadModel resolves true on a downloaded ack", async () => {
const client = new SttClient(() =>
makeFakeWorker((message, emit) => {
if (message.type === "download") emit({ type: "downloaded", id: message.id });
}),
);
expect(await client.downloadModel("fast")).toBe(true);
await client.terminate();
});
});
// ── Raw PCM decoding ────────────────────────────────────────────────
describe("decodePcmS16LE", () => {
it("decodes little-endian s16 frames to normalized [-1, 1] floats", () => {
const audio = decodePcmS16LE(bytesFromInt16([16_384, -16_384, 0, 32_767]));
expect(audio.length).toBe(4);
expect(audio[0]).toBeCloseTo(0.5, 4);
expect(audio[1]).toBeCloseTo(-0.5, 4);
expect(audio[2]).toBeCloseTo(0, 6);
expect(audio[3]).toBeCloseTo(32_767 / 32_768, 4);
});
it("ignores a trailing odd byte (callers carry it across chunks)", () => {
const bytes = new Uint8Array([0x00, 0x40, 0x7f]); // one whole sample + one stray byte
const audio = decodePcmS16LE(bytes.subarray(0, bytes.length - (bytes.length % 2)));
expect(audio.length).toBe(1);
expect(audio[0]).toBeCloseTo(0.5, 4);
});
});
// ── Energy endpointer (live segmentation) ───────────────────────────
function silenceFrames(ms: number, sampleRate = 16_000): Float32Array {
return new Float32Array(Math.round((sampleRate * ms) / 1000));
}
function toneFrames(ms: number, amplitude = 0.3, sampleRate = 16_000): Float32Array {
const n = Math.round((sampleRate * ms) / 1000);
const out = new Float32Array(n);
for (let i = 0; i < n; i += 1) out[i] = amplitude * Math.sin((2 * Math.PI * 220 * i) / sampleRate);
return out;
}
describe("StreamEndpointer", () => {
const base = {
frameMs: 10,
endSilenceMs: 50,
minSpeechMs: 20,
maxSegmentMs: 5_000,
preRollMs: 0,
partialIntervalMs: 10_000,
energyRatio: 2,
floorAttack: 0.1,
minThreshold: 0.01,
} as const;
it("commits one segment for a speech burst bounded by silence", () => {
const ep = new StreamEndpointer(base);
const events = [...ep.push(silenceFrames(80)), ...ep.push(toneFrames(120)), ...ep.push(silenceFrames(80))];
const segments = events.filter(e => e.kind === "segment");
expect(segments.length).toBe(1);
expect(segments[0]!.audio.length).toBeGreaterThan(0);
});
it("splits two phrases separated by a pause into two segments", () => {
const ep = new StreamEndpointer(base);
const events = [
...ep.push(silenceFrames(80)),
...ep.push(toneFrames(120)),
...ep.push(silenceFrames(80)),
...ep.push(toneFrames(120)),
...ep.push(silenceFrames(80)),
];
expect(events.filter(e => e.kind === "segment").length).toBe(2);
});
it("discards a sub-minimum blip as noise", () => {
const ep = new StreamEndpointer({ ...base, minSpeechMs: 80 });
const events = [...ep.push(silenceFrames(80)), ...ep.push(toneFrames(30)), ...ep.push(silenceFrames(80))];
expect(events.filter(e => e.kind === "segment").length).toBe(0);
});
it("flush commits in-progress speech with no trailing silence", () => {
const ep = new StreamEndpointer(base);
const live = [...ep.push(silenceFrames(40)), ...ep.push(toneFrames(120))];
expect(live.filter(e => e.kind === "segment").length).toBe(0);
expect(ep.flush().filter(e => e.kind === "segment").length).toBe(1);
});
it("emits volatile partials while a long phrase is still in progress", () => {
const ep = new StreamEndpointer({ ...base, partialIntervalMs: 30, endSilenceMs: 400 });
ep.push(silenceFrames(40));
const events = ep.push(toneFrames(120));
expect(events.filter(e => e.kind === "partial").length).toBeGreaterThanOrEqual(2);
});
});
// ── Streaming session correlation ───────────────────────────────────
describe("SttClient.startStream", () => {
it("routes partials/segments and resolves stop() with the final transcript", async () => {
let emit: ((out: SttWorkerOutbound) => void) | undefined;
let streamId: string | undefined;
const sent: SttWorkerInbound[] = [];
const client = new SttClient(() =>
makeFakeWorker((message, e) => {
emit = e;
sent.push(message);
if (message.type === "stream_start") streamId = message.id;
}),
);
const partials: string[] = [];
const segments: Array<{ text: string; index: number }> = [];
const handle = client.startStream("parakeet", {
onPartial: text => partials.push(text),
onSegment: (text, index) => segments.push({ text, index }),
});
expect(streamId).toBeDefined();
handle.pushAudio(new Float32Array([0.1, 0.2]));
emit?.({ type: "partial", id: streamId!, text: "hel" });
emit?.({ type: "segment", id: streamId!, index: 0, text: "hello" });
emit?.({ type: "segment", id: streamId!, index: 1, text: "world" });
const final = handle.stop();
emit?.({ type: "stream_done", id: streamId!, text: "hello world" });
expect(await final).toBe("hello world");
expect(partials).toEqual(["hel"]);
expect(segments).toEqual([
{ text: "hello", index: 0 },
{ text: "world", index: 1 },
]);
expect(sent.some(m => m.type === "stream_audio")).toBe(true);
expect(sent.some(m => m.type === "stream_stop")).toBe(true);
await client.terminate();
});
it("rejects stop() when the worker reports a stream error", async () => {
let emit: ((out: SttWorkerOutbound) => void) | undefined;
let streamId: string | undefined;
const client = new SttClient(() =>
makeFakeWorker((message, e) => {
emit = e;
if (message.type === "stream_start") streamId = message.id;
}),
);
const result = client.startStream("parakeet", {}).stop();
emit?.({ type: "error", id: streamId!, error: "decode blew up" });
await expect(result).rejects.toThrow("decode blew up");
await client.terminate();
});
it("cancels on an aborted signal, resolving stop() empty and emitting no segments", async () => {
const controller = new AbortController();
const sent: SttWorkerInbound[] = [];
const client = new SttClient(() =>
makeFakeWorker(message => {
sent.push(message);
}),
);
const segments: string[] = [];
const handle = client.startStream("parakeet", { signal: controller.signal, onSegment: t => segments.push(t) });
controller.abort();
expect(await handle.stop()).toBe("");
expect(segments).toEqual([]);
expect(sent.some(m => m.type === "stream_cancel")).toBe(true);
await client.terminate();
});
});
describe("worker streaming dispatch", () => {
it("rejects a streaming start for an unknown model", () => {
const sent: SttWorkerOutbound[] = [];
let handler: ((message: SttWorkerInbound) => void) | undefined;
const transport: SttTransport = {
send: message => sent.push(message),
onMessage: h => {
handler = h;
return () => {};
},
};
startSttWorker(transport);
handler?.({ type: "stream_start", id: "s1", modelKey: "bogus" } as unknown as SttWorkerInbound);
expect(sent.length).toBe(1);
expect(sent[0]).toMatchObject({ type: "error", id: "s1" });
expect((sent[0] as Extract<SttWorkerOutbound, { type: "error" }>).error).toContain("Unknown stt model");
});
});
@@ -1,70 +0,0 @@
"""Transcribe a WAV file using openai-whisper.
Reads WAV directly via Python's wave module (no ffmpeg needed).
Resamples to 16kHz mono float32 and passes to whisper as a numpy array.
Usage: python transcribe.py <audio.wav> <model_name> <language>
Prints transcribed text to stdout.
"""
import sys
import wave
import re
import numpy as np
import whisper
def load_wav(path: str) -> np.ndarray:
with wave.open(path, "rb") as wf:
rate = wf.getframerate()
channels = wf.getnchannels()
width = wf.getsampwidth()
n_frames = wf.getnframes()
raw = wf.readframes(n_frames)
if width == 2:
audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
elif width == 1:
audio = (np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128.0) / 128.0
elif width == 4:
audio = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2147483648.0
else:
raise ValueError(f"Unsupported sample width: {width}")
# Mix to mono
if channels > 1:
audio = audio.reshape(-1, channels).mean(axis=1)
# Resample to 16 kHz
if rate != 16000:
target_len = int(len(audio) * 16000 / rate)
audio = np.interp(
np.linspace(0, len(audio) - 1, target_len),
np.arange(len(audio)),
audio,
).astype(np.float32)
return audio
def main() -> None:
if len(sys.argv) < 2:
print("Usage: python transcribe.py <audio.wav> <model_name> <language>", file=sys.stderr)
sys.exit(1)
audio_path = sys.argv[1]
model_name = sys.argv[2] if len(sys.argv) > 2 else "base.en"
language = sys.argv[3] if len(sys.argv) > 3 else "en"
if not re.fullmatch(r"[A-Za-z]{2,3}(-[A-Za-z]{2})?", language):
print(f"Invalid language code: {language}", file=sys.stderr)
sys.exit(1)
audio = load_wav(audio_path)
model = whisper.load_model(model_name)
result = model.transcribe(audio, language=language)
print(result["text"].strip())
if __name__ == "__main__":
main()
+36 -67
View File
@@ -1,5 +1,7 @@
import { $which, logger } from "@oh-my-pi/pi-utils";
import transcribeScript from "./transcribe.py" with { type: "text" };
import { logger } from "@oh-my-pi/pi-utils";
import { sttClient } from "./asr-client";
import { resolveSttModelSpec } from "./models";
import { decodeWavToMono16k } from "./wav";
export interface TranscribeOptions {
modelName?: string;
@@ -10,82 +12,49 @@ export interface TranscribeOptions {
const TRANSCRIBE_TIMEOUT_MS = 120_000;
/**
* Find a usable Python command.
*/
export function resolvePython(): string | null {
for (const cmd of ["python", "py", "python3"]) {
if ($which(cmd)) return cmd;
}
return null;
}
/**
* Transcribe a WAV file using Python openai-whisper.
* Transcribe a WAV file using the local ONNX Whisper worker.
*
* Reads the WAV via Python's built-in `wave` module (no ffmpeg needed),
* resamples to 16 kHz mono, and passes the numpy array directly to whisper.
* Decodes the WAV to a 16 kHz mono Float32Array in-process (no Python, no
* ffmpeg) and routes it to the warm speech worker, which keeps the model loaded
* across calls. Honors `options.signal` (abort) and applies an internal timeout
* with the same semantics as the previous Python path.
*/
export async function transcribe(audioPath: string, options?: TranscribeOptions): Promise<string> {
const audioFile = Bun.file(audioPath);
if (audioFile.size < 100) {
throw new Error(`Audio file is empty or too small (${audioFile.size} bytes). Check microphone.`);
}
const pythonCmd = resolvePython();
if (!pythonCmd) {
throw new Error("Python not found. Install Python 3.8+ from https://python.org");
}
const modelName = options?.modelName ?? "base.en";
const language = options?.language ?? "en";
logger.debug("Transcribing with Python whisper", { pythonCmd, audioPath, modelName, language });
const proc = Bun.spawn([pythonCmd, "-c", transcribeScript, audioPath, modelName, language], {
stdout: "pipe",
stderr: "pipe",
});
if (options?.signal?.aborted) {
proc.kill();
options.signal.throwIfAborted();
}
const onAbort = () => proc.kill();
options?.signal?.addEventListener("abort", onAbort, { once: true });
let timedOut = false;
const killTimer = setTimeout(() => {
timedOut = true;
logger.error("Python whisper transcription timed out, killing process", { timeoutMs: TRANSCRIBE_TIMEOUT_MS });
proc.kill();
}, TRANSCRIBE_TIMEOUT_MS);
const exitCode = await proc.exited;
clearTimeout(killTimer);
options?.signal?.removeEventListener("abort", onAbort);
options?.signal?.throwIfAborted();
const stdout = await new Response(proc.stdout).text();
const stderr = await new Response(proc.stderr).text();
const spec = resolveSttModelSpec(options?.modelName);
const language = options?.language || undefined;
const audio = decodeWavToMono16k(await audioFile.arrayBuffer());
if (audio.length === 0) return "";
if (timedOut) {
throw new Error(`Transcription timed out after ${Math.round(TRANSCRIBE_TIMEOUT_MS / 1000)}s`);
}
logger.debug("Transcribing with local ONNX whisper", {
audioPath,
modelKey: spec.key,
repo: spec.repo,
language,
samples: audio.length,
});
if (exitCode !== 0) {
logger.error("Python whisper transcription failed", { exitCode, stderr: stderr.trim() });
if (stderr.includes("No module named 'whisper'")) {
throw new Error("openai-whisper not installed. Run: pip install openai-whisper");
// Bound runaway inference. Abort the request on timeout; the warm worker
// keeps the model loaded (the request promise just rejects).
const timeout = new AbortController();
const timer = setTimeout(() => timeout.abort(), TRANSCRIBE_TIMEOUT_MS);
const signal = options?.signal ? AbortSignal.any([options.signal, timeout.signal]) : timeout.signal;
try {
const text = (await sttClient.transcribe(spec.key, audio, { language, signal })).trim();
logger.debug("Transcription complete", { length: text.length });
return text;
} catch (error) {
if (timeout.signal.aborted && !options?.signal?.aborted) {
logger.error("Local whisper transcription timed out", { timeoutMs: TRANSCRIBE_TIMEOUT_MS });
throw new Error(`Transcription timed out after ${Math.round(TRANSCRIBE_TIMEOUT_MS / 1000)}s`);
}
// Show last line of stderr (the actual error, not the full traceback)
const lastLine = stderr.trim().split("\n").pop() ?? "";
throw new Error(`Transcription failed: ${lastLine}`);
throw error;
} finally {
clearTimeout(timer);
}
const text = stdout.trim();
logger.debug("Transcription complete", { length: text.length });
return text;
}
+173
View File
@@ -0,0 +1,173 @@
/**
* Minimal WAV (RIFF/PCM) decoder producing the Float32Array @ 16 kHz mono that
* transformers.js `automatic-speech-recognition` expects. Ports the decode/
* mono-mix/resample logic from the retired Python `transcribe.py` (which read
* via the stdlib `wave` module) so STT no longer shells out to Python.
*
* Supported sample formats: PCM uint8 (8-bit), int16 (16-bit), int32 (32-bit),
* and IEEE float32 (format tag 3). Any number of channels is mixed down to mono.
*/
/** transformers.js Whisper feature extractor operates at 16 kHz. */
export const TARGET_SAMPLE_RATE = 16_000;
const WAV_FORMAT_PCM = 1;
const WAV_FORMAT_IEEE_FLOAT = 3;
const WAV_FORMAT_EXTENSIBLE = 0xfffe;
interface WavData {
format: number;
channels: number;
sampleRate: number;
bitsPerSample: number;
/** Raw PCM/float bytes from the `data` chunk. */
samples: DataView;
}
function readFourCc(view: DataView, offset: number): string {
return String.fromCharCode(
view.getUint8(offset),
view.getUint8(offset + 1),
view.getUint8(offset + 2),
view.getUint8(offset + 3),
);
}
/** Parse the RIFF container, returning the `fmt ` parameters and `data` bytes. */
function parseWav(buffer: ArrayBuffer): WavData {
const view = new DataView(buffer);
if (buffer.byteLength < 12 || readFourCc(view, 0) !== "RIFF" || readFourCc(view, 8) !== "WAVE") {
throw new Error("Not a RIFF/WAVE file");
}
let format: number | undefined;
let channels = 0;
let sampleRate = 0;
let bitsPerSample = 0;
let samples: DataView | undefined;
// Chunks begin after the 12-byte RIFF/WAVE header; each is an 8-byte header
// (4-char id + uint32 LE size) followed by `size` bytes padded to even.
let offset = 12;
while (offset + 8 <= buffer.byteLength) {
const id = readFourCc(view, offset);
const size = view.getUint32(offset + 4, true);
const body = offset + 8;
if (id === "fmt ") {
format = view.getUint16(body, true);
channels = view.getUint16(body + 2, true);
sampleRate = view.getUint32(body + 4, true);
bitsPerSample = view.getUint16(body + 14, true);
// WAVE_FORMAT_EXTENSIBLE (ffmpeg & friends): the real codec is the
// first 2 bytes of the SubFormat GUID in the fmt extension.
if (format === WAV_FORMAT_EXTENSIBLE && size >= 40) format = view.getUint16(body + 24, true);
} else if (id === "data") {
const length = Math.min(size, buffer.byteLength - body);
samples = new DataView(buffer, body, length);
}
offset = body + size + (size % 2);
}
if (format === undefined || samples === undefined || channels < 1 || sampleRate < 1) {
throw new Error("WAV file missing fmt/data chunks");
}
return { format, channels, sampleRate, bitsPerSample, samples };
}
/** Decode raw PCM/float bytes into interleaved normalized [-1, 1] float samples. */
function decodeSamples(wav: WavData): Float32Array {
const { format, bitsPerSample, samples } = wav;
const view = samples;
if (format === WAV_FORMAT_IEEE_FLOAT && bitsPerSample === 32) {
const count = Math.floor(view.byteLength / 4);
const out = new Float32Array(count);
for (let i = 0; i < count; i += 1) out[i] = view.getFloat32(i * 4, true);
return out;
}
if (format !== WAV_FORMAT_PCM) {
throw new Error(`Unsupported WAV format tag: ${format}`);
}
if (bitsPerSample === 16) {
const count = Math.floor(view.byteLength / 2);
const out = new Float32Array(count);
for (let i = 0; i < count; i += 1) out[i] = view.getInt16(i * 2, true) / 32_768;
return out;
}
if (bitsPerSample === 8) {
// 8-bit PCM is unsigned, centered at 128.
const count = view.byteLength;
const out = new Float32Array(count);
for (let i = 0; i < count; i += 1) out[i] = (view.getUint8(i) - 128) / 128;
return out;
}
if (bitsPerSample === 32) {
const count = Math.floor(view.byteLength / 4);
const out = new Float32Array(count);
for (let i = 0; i < count; i += 1) out[i] = view.getInt32(i * 4, true) / 2_147_483_648;
return out;
}
throw new Error(`Unsupported PCM sample width: ${bitsPerSample} bits`);
}
/** Average interleaved channels down to a single mono track. */
function mixToMono(interleaved: Float32Array, channels: number): Float32Array {
if (channels <= 1) return interleaved;
const frames = Math.floor(interleaved.length / channels);
const out = new Float32Array(frames);
for (let frame = 0; frame < frames; frame += 1) {
let sum = 0;
for (let channel = 0; channel < channels; channel += 1) sum += interleaved[frame * channels + channel]!;
out[frame] = sum / channels;
}
return out;
}
/**
* Resample via linear interpolation, mirroring the Python `np.interp` over
* `linspace(0, n-1, targetLen)` against `arange(n)`.
*/
export function resampleLinear(input: Float32Array, fromRate: number, toRate: number): Float32Array {
if (fromRate === toRate || input.length === 0) return input;
const n = input.length;
const targetLen = Math.max(1, Math.floor((n * toRate) / fromRate));
const out = new Float32Array(targetLen);
if (targetLen === 1) {
out[0] = input[0]!;
return out;
}
const step = (n - 1) / (targetLen - 1);
for (let i = 0; i < targetLen; i += 1) {
const pos = i * step;
const lo = Math.floor(pos);
const hi = Math.min(lo + 1, n - 1);
const frac = pos - lo;
out[i] = input[lo]! * (1 - frac) + input[hi]! * frac;
}
return out;
}
/**
* Decode a WAV byte buffer into a 16 kHz mono Float32Array suitable for the
* transformers.js Whisper pipeline.
*/
export function decodeWavToMono16k(buffer: ArrayBuffer): Float32Array {
const wav = parseWav(buffer);
const interleaved = decodeSamples(wav);
const mono = mixToMono(interleaved, wav.channels);
return resampleLinear(mono, wav.sampleRate, TARGET_SAMPLE_RATE);
}
/**
* Decode interleaved little-endian signed 16-bit PCM bytes into normalized
* [-1, 1] mono float samples. The live recorder streams raw s16le frames from
* sox/ffmpeg/arecord stdout (no RIFF container), so this is the hot-path
* counterpart to {@link decodeWavToMono16k}. `bytes` MUST be 2-byte aligned;
* callers buffer any trailing odd byte across chunk boundaries.
*/
export function decodePcmS16LE(bytes: Uint8Array): Float32Array {
const count = bytes.length >>> 1;
const view = new DataView(bytes.buffer, bytes.byteOffset, count * 2);
const out = new Float32Array(count);
for (let i = 0; i < count; i += 1) out[i] = view.getInt16(i * 2, true) / 32_768;
return out;
}
+8
View File
@@ -23,6 +23,7 @@ import type { RenderResultOptions } from "../extensibility/custom-tools/types";
import type { ExtensionUISelectItem } from "../extensibility/extensions";
import { getMarkdownTheme, type Theme, theme } from "../modes/theme/theme";
import askDescription from "../prompts/tools/ask.md" with { type: "text" };
import { vocalizer } from "../tts/vocalizer";
import { framedBlock, renderStatusLine } from "../tui";
import type { ToolSession } from ".";
import { formatErrorMessage, formatMeta, formatTitle } from "./render-utils";
@@ -487,6 +488,13 @@ export class AskTool implements AgentTool<typeof askSchema, AskToolDetails> {
};
}
// Speak the question(s) aloud before surfacing them. Ask vocalizes in every
// mode — it's the assistant addressing the user — gated only by speech.enabled
// (the vocalizer re-checks the setting and no-ops when disabled).
if (this.session.settings.get("speech.enabled")) {
vocalizer.speak(params.questions.map(q => q.question).join("\n"));
}
const askQuestion = async (
q: AskParams["questions"][number],
options?: { previous?: QuestionResult; navigation?: NavigationControls },
+202 -91
View File
@@ -1,10 +1,17 @@
// Ported from NousResearch/hermes-agent (MIT) — tools/tts_tool.py L167-171, L896-959.
// The xAI Grok Voice path below is preserved intact; a local on-device neural TTS
// backend (Kokoro-82M via kokoro-js on the shared ONNX worker) is layered on behind
// the `providers.tts` switch.
import type { AgentToolResult } from "@oh-my-pi/pi-agent-core";
import { type ApiKey, ProviderHttpError, withAuth } from "@oh-my-pi/pi-ai";
import { z } from "zod/v4";
import { settings } from "../config/settings";
import type { CustomTool, CustomToolContext } from "../extensibility/custom-tools/types";
import { ohMyPiXAIUserAgent, resolveXAIHttpCredentials } from "../lib/xai-http";
import { DEFAULT_TTS_LOCAL_MODEL_KEY, DEFAULT_TTS_VOICE, isTtsLocalModelKey, KOKORO_VOICES } from "../tts/models";
import { ttsClient } from "../tts/tts-client";
import { encodeWav } from "../tts/wav";
import { formatPathRelativeToCwd, resolveToCwd } from "./path-utils";
// Hermes tts_tool.py L167-171
@@ -22,6 +29,7 @@ const formatVoiceList = (): string =>
XAI_BUILTIN_VOICES.map(v => (v === DEFAULT_XAI_VOICE_ID ? `${v} (default)` : v)).join(", ");
type TtsCodec = "mp3" | "wav";
type TtsBackend = "local" | "xai";
const ttsSchema = z.object({
text: z.string().min(1).max(XAI_MAX_TEXT_LENGTH),
@@ -36,16 +44,200 @@ interface TtsToolDetails {
bytes: number;
voiceId: string;
codec: TtsCodec;
backend: TtsBackend;
}
/**
* Pick the synthesis backend. Pure for testability.
*
* - `xai` / `local` are honored verbatim (the xAI path still surfaces its own
* "no credentials" error when creds are missing).
* - `auto` prefers the local on-device backend, except when the caller asked for
* an `.mp3` and xAI credentials exist — only the cloud path can emit MP3, so we
* route there to satisfy the requested container rather than substituting WAV.
*/
export function resolveTtsBackend(opts: { preference: string; wantsMp3: boolean; hasXaiCreds: boolean }): TtsBackend {
if (opts.preference === "xai") return "xai";
if (opts.preference === "local") return "local";
if (opts.wantsMp3 && opts.hasXaiCreds) return "xai";
return "local";
}
/**
* Resolve the on-disk path for local synthesis. Local output is always WAV (no
* MP3 encoder is bundled), so an `.mp3` (or any non-`.wav`) request is rewritten
* to a sibling `.wav` and flagged so the tool result can note the substitution.
*/
export function resolveLocalWavPath(outputPath: string): { wavPath: string; substituted: boolean } {
const lower = outputPath.toLowerCase();
if (lower.endsWith(".wav")) return { wavPath: outputPath, substituted: false };
const slash = Math.max(outputPath.lastIndexOf("/"), outputPath.lastIndexOf("\\"));
const dot = outputPath.lastIndexOf(".");
const base = dot > slash ? outputPath.slice(0, dot) : outputPath;
return { wavPath: `${base}.wav`, substituted: true };
}
function readStringSetting(key: "providers.tts" | "tts.localModel" | "tts.localVoice"): string | undefined {
try {
const value = settings.get(key);
return typeof value === "string" ? value : undefined;
} catch {
return undefined;
}
}
async function synthesizeXai(
params: z.infer<typeof ttsSchema>,
ctx: CustomToolContext,
outputPath: string,
displayPath: string,
codec: TtsCodec,
signal: AbortSignal | undefined,
): Promise<AgentToolResult<TtsToolDetails, typeof ttsSchema>> {
const creds = await resolveXAIHttpCredentials(ctx.modelRegistry);
if (!creds) {
return {
isError: true,
content: [
{
type: "text",
text: "No xAI credentials. Run /login → xAI Grok OAuth (SuperGrok Subscription) or set XAI_API_KEY.",
},
],
};
}
const voiceId = params.voice_id;
const language = params.language;
const sampleRate = params.sample_rate ?? DEFAULT_XAI_SAMPLE_RATE;
const bitRate = params.bit_rate ?? DEFAULT_XAI_BIT_RATE;
const payload: Record<string, unknown> = {
text: params.text,
voice_id: voiceId,
language,
};
// Hermes tts_tool.py L926-940 — only send output_format when caller overrides a default.
const codecOverridden = codec !== "mp3";
const sampleRateOverridden = sampleRate !== DEFAULT_XAI_SAMPLE_RATE;
const bitRateOverridden = codec === "mp3" && bitRate !== DEFAULT_XAI_BIT_RATE;
if (codecOverridden || sampleRateOverridden || bitRateOverridden) {
const fmt: Record<string, unknown> = { codec };
if (sampleRate) fmt.sample_rate = sampleRate;
if (codec === "mp3" && bitRate) fmt.bit_rate = bitRate;
payload.output_format = fmt;
}
// Compose the caller signal with a 60 s timeout fence.
const timeoutSignal = AbortSignal.timeout(60_000);
const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
const sessionId = ctx.sessionManager.getSessionId();
const apiKey: ApiKey = ctx.modelRegistry.resolver(creds.provider, {
sessionId,
baseUrl: creds.baseURL,
});
let response: Response;
try {
response = await withAuth(
apiKey,
async key => {
const resp = await fetch(`${creds.baseURL}/tts`, {
method: "POST",
headers: {
Authorization: `Bearer ${key}`,
"Content-Type": "application/json",
"User-Agent": ohMyPiXAIUserAgent(),
},
body: JSON.stringify(payload),
signal: combinedSignal,
});
if (!resp.ok) {
const detail = await resp.text();
throw new ProviderHttpError(`xAI TTS failed (${resp.status}): ${detail.slice(0, 300)}`, resp.status, {
headers: resp.headers,
});
}
return resp;
},
{ signal: combinedSignal },
);
} catch (error) {
const status = (error as { status?: unknown }).status;
if (error instanceof Error && typeof status === "number") {
return {
isError: true,
content: [{ type: "text", text: error.message }],
};
}
throw error;
}
const bytes = new Uint8Array(await response.arrayBuffer());
await Bun.write(outputPath, bytes);
return {
content: [
{
type: "text",
text: `Saved ${bytes.length} bytes to ${displayPath} (voice=${voiceId}, codec=${codec}, backend=xai).`,
},
],
details: { bytes: bytes.length, voiceId, codec, backend: "xai" },
};
}
async function synthesizeLocal(
params: z.infer<typeof ttsSchema>,
cwd: string,
outputPath: string,
signal: AbortSignal | undefined,
): Promise<AgentToolResult<TtsToolDetails, typeof ttsSchema>> {
const modelSetting = readStringSetting("tts.localModel");
const modelKey = modelSetting && isTtsLocalModelKey(modelSetting) ? modelSetting : DEFAULT_TTS_LOCAL_MODEL_KEY;
const voice = readStringSetting("tts.localVoice") || DEFAULT_TTS_VOICE;
const audio = await ttsClient.synthesize(modelKey, params.text, { voice, signal });
if (!audio) {
return {
isError: true,
content: [
{
type: "text",
text: `Local TTS synthesis failed (model=${modelKey}). The on-device worker may be unavailable or the model download was interrupted.`,
},
],
};
}
const { wavPath, substituted } = resolveLocalWavPath(outputPath);
const wav = encodeWav(audio.pcm, audio.sampleRate);
await Bun.write(wavPath, wav);
const displayPath = formatPathRelativeToCwd(wavPath, cwd);
const note = substituted
? ` No local MP3 encoder is bundled, so WAV (PCM16) was written instead of the requested container.`
: "";
return {
content: [
{
type: "text",
text: `Saved ${wav.length} bytes to ${displayPath} (voice=${modelKey}/${voice}, codec=wav, backend=local, ${audio.sampleRate} Hz).${note}`,
},
],
details: { bytes: wav.length, voiceId: `${modelKey}/${voice}`, codec: "wav", backend: "local" },
};
}
export const ttsTool: CustomTool<typeof ttsSchema, TtsToolDetails> = {
name: "tts",
label: "TextToSpeech",
label: "Speech Generation",
strict: false,
approval: "write",
description:
`Synthesize speech from text using xAI Grok Voice. Built-in voices: ${formatVoiceList()}. ` +
"Custom voice IDs also accepted. Output codec inferred from output_path suffix (.wav → wav, else mp3). " +
"Generate a speech audio file from text and write it to output_path. Two backends, selected by the providers.tts setting (auto|local|xai): " +
`local = on-device neural TTS (Kokoro-82M via the bundled ONNX runtime, no network, output is always WAV/PCM16; voice set by the tts.localVoice setting — ${KOKORO_VOICES.map(v => (v.id === DEFAULT_TTS_VOICE ? `${v.id} (default)` : v.id)).join(", ")}); ` +
`xai = xAI Grok Voice cloud (built-in voices: ${formatVoiceList()}; custom voice IDs accepted; MP3 or WAV). ` +
"auto prefers local, but routes an .mp3 request to xAI when credentials exist (only the cloud path emits MP3); " +
"otherwise an .mp3 path is written as a sibling .wav. xAI codec is inferred from the output_path suffix. " +
`Max ${XAI_MAX_TEXT_LENGTH.toLocaleString("en-US")} characters.`,
parameters: ttsSchema,
async execute(
@@ -55,99 +247,18 @@ export const ttsTool: CustomTool<typeof ttsSchema, TtsToolDetails> = {
ctx: CustomToolContext,
signal?: AbortSignal,
): Promise<AgentToolResult<TtsToolDetails, typeof ttsSchema>> {
const creds = await resolveXAIHttpCredentials(ctx.modelRegistry);
if (!creds) {
return {
isError: true,
content: [
{
type: "text",
text: "No xAI credentials. Run /login → xAI Grok OAuth (SuperGrok Subscription) or set XAI_API_KEY.",
},
],
};
}
const cwd = ctx.sessionManager.getCwd();
const outputPath = resolveToCwd(params.output_path, cwd);
const displayPath = formatPathRelativeToCwd(outputPath, cwd);
const codec: TtsCodec = outputPath.toLowerCase().endsWith(".wav") ? "wav" : "mp3";
const voiceId = params.voice_id;
const language = params.language;
const sampleRate = params.sample_rate ?? DEFAULT_XAI_SAMPLE_RATE;
const bitRate = params.bit_rate ?? DEFAULT_XAI_BIT_RATE;
const payload: Record<string, unknown> = {
text: params.text,
voice_id: voiceId,
language,
};
// Hermes tts_tool.py L926-940 — only send output_format when caller overrides a default.
const codecOverridden = codec !== "mp3";
const sampleRateOverridden = sampleRate !== DEFAULT_XAI_SAMPLE_RATE;
const bitRateOverridden = codec === "mp3" && bitRate !== DEFAULT_XAI_BIT_RATE;
if (codecOverridden || sampleRateOverridden || bitRateOverridden) {
const fmt: Record<string, unknown> = { codec };
if (sampleRate) fmt.sample_rate = sampleRate;
if (codec === "mp3" && bitRate) fmt.bit_rate = bitRate;
payload.output_format = fmt;
}
const preference = readStringSetting("providers.tts") ?? "auto";
// Only resolve xAI creds when they can affect routing (skip for an explicit local preference).
const hasXaiCreds =
preference === "local" ? false : (await resolveXAIHttpCredentials(ctx.modelRegistry)) !== null;
const backend = resolveTtsBackend({ preference, wantsMp3: codec === "mp3", hasXaiCreds });
// Compose the caller signal with a 60 s timeout fence.
const timeoutSignal = AbortSignal.timeout(60_000);
const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
const sessionId = ctx.sessionManager.getSessionId();
const apiKey: ApiKey = ctx.modelRegistry.resolver(creds.provider, {
sessionId,
baseUrl: creds.baseURL,
});
let response: Response;
try {
response = await withAuth(
apiKey,
async key => {
const resp = await fetch(`${creds.baseURL}/tts`, {
method: "POST",
headers: {
Authorization: `Bearer ${key}`,
"Content-Type": "application/json",
"User-Agent": ohMyPiXAIUserAgent(),
},
body: JSON.stringify(payload),
signal: combinedSignal,
});
if (!resp.ok) {
const detail = await resp.text();
throw new ProviderHttpError(`xAI TTS failed (${resp.status}): ${detail.slice(0, 300)}`, resp.status, {
headers: resp.headers,
});
}
return resp;
},
{ signal: combinedSignal },
);
} catch (error) {
const status = (error as { status?: unknown }).status;
if (error instanceof Error && typeof status === "number") {
return {
isError: true,
content: [{ type: "text", text: error.message }],
};
}
throw error;
}
const bytes = new Uint8Array(await response.arrayBuffer());
await Bun.write(outputPath, bytes);
return {
content: [
{
type: "text",
text: `Saved ${bytes.length} bytes to ${displayPath} (voice=${voiceId}, codec=${codec}).`,
},
],
details: { bytes: bytes.length, voiceId, codec },
};
if (backend === "local") return synthesizeLocal(params, cwd, outputPath, signal);
return synthesizeXai(params, ctx, outputPath, displayPath, codec, signal);
},
};
@@ -0,0 +1,64 @@
import * as fs from "node:fs/promises";
import * as path from "node:path";
import { getTinyModelsCacheDir } from "@oh-my-pi/pi-utils";
import { getTtsLocalModelSpec } from "./models";
import { isTtsRuntimeCached } from "./runtime";
import { ttsClient } from "./tts-client";
export interface TtsDownloadProgress {
stage: string;
/** Integer 0–100 download percent when known. */
percent?: number;
}
/**
* Whether the selected local TTS model and the side Kokoro runtime are already
* present. transformers.js stores `main`-revision files at
* `<cacheDir>/<repo>/...`, so any `.onnx` weight under the repo dir means the
* model weights can load without a network fetch; the Kokoro package runtime is
* version-keyed separately and must also exist before setup can report ready.
*/
export async function isTtsModelCached(modelKey: string): Promise<boolean> {
const spec = getTtsLocalModelSpec(modelKey);
if (!spec) return false;
const repoDir = path.join(getTinyModelsCacheDir(), ...spec.repo.split("/"));
try {
const entries = await fs.readdir(repoDir, { recursive: true });
const hasWeights = entries.some(entry => typeof entry === "string" && entry.endsWith(".onnx"));
return hasWeights && (await isTtsRuntimeCached());
} catch {
return false;
}
}
/**
* Ensure the selected local TTS model is downloaded into the transformers.js
* cache (and warm in the worker), streaming integer-percent Hub progress. The
* worker resolves the request once every model file is cached. Returns `false`
* if the worker is unavailable or the download failed.
*/
export async function downloadTtsModel(
modelKey: string,
onProgress?: (progress: TtsDownloadProgress) => void,
signal?: AbortSignal,
): Promise<boolean> {
const spec = getTtsLocalModelSpec(modelKey);
if (!spec) return false;
onProgress?.({ stage: `Preparing ${spec.label}...` });
return ttsClient.downloadModel(spec.key, {
signal,
onProgress: event => {
if (event.status === "ready" || event.status === "done") {
onProgress?.({ stage: `${spec.label} ready`, percent: 100 });
return;
}
const percent =
typeof event.total === "number" && event.total > 0 && typeof event.loaded === "number"
? Math.round((event.loaded / event.total) * 100)
: typeof event.progress === "number"
? Math.round(event.progress)
: undefined;
onProgress?.({ stage: `Downloading ${spec.label}`, percent });
},
});
}
+8
View File
@@ -0,0 +1,8 @@
export * from "./downloader";
export * from "./models";
export * from "./runtime";
export * from "./tts-client";
export * from "./tts-protocol";
export * from "./tts-worker";
export * from "./vocalizer";
export * from "./wav";
+137
View File
@@ -0,0 +1,137 @@
import type { TinyModelDtype } from "../tiny/dtype";
/**
* Voice exposed by a local TTS model. Kokoro ships a fixed catalog of named
* voices; a voice is just a stable id (e.g. `af_heart`) plus a display label.
* Selection is purely on-device — generating with a different voice needs no
* extra network fetch once the model weights are cached.
*/
export interface TtsLocalVoiceSpec {
id: string;
label: string;
}
/**
* A local (on-device, ONNX) text-to-speech model the worker can load. `repo` is
* the Hugging Face model id loaded through `kokoro-js`
* (`KokoroTTS.from_pretrained`), which runs on the same `@huggingface/transformers`
* + `onnxruntime` runtime as the rest of the tiny-model stack and bundles the
* misaki/espeak phonemizer Kokoro needs. `dtype` is the default ONNX precision
* (overridable via `providers.tinyModelDtype`/`PI_TINY_DTYPE`).
*/
export interface TtsLocalModelSpec {
key: string;
repo: string;
dtype: TinyModelDtype;
/** PCM sample rate the model emits; fallback only — the worker uses the value RawAudio reports. */
sampleRate: number;
label: string;
description: string;
/** First entry is the model's default voice. */
voices: readonly TtsLocalVoiceSpec[];
}
/**
* Curated Kokoro-82M voice catalog. Kokoro ships ~28 voices; we surface the
* higher-graded ones across American/British × female/male so the picker stays
* useful without listing every D/F-grade sample. `af_heart` (grade A) leads and
* is the default voice. Grades are Kokoro's own `overallGrade` ratings.
*/
export const KOKORO_VOICES: readonly TtsLocalVoiceSpec[] = [
{ id: "af_heart", label: "Heart (American female)" },
{ id: "af_bella", label: "Bella (American female)" },
{ id: "af_nicole", label: "Nicole (American female)" },
{ id: "af_aoede", label: "Aoede (American female)" },
{ id: "af_kore", label: "Kore (American female)" },
{ id: "af_sarah", label: "Sarah (American female)" },
{ id: "am_michael", label: "Michael (American male)" },
{ id: "am_fenrir", label: "Fenrir (American male)" },
{ id: "am_puck", label: "Puck (American male)" },
{ id: "bf_emma", label: "Emma (British female)" },
{ id: "bm_george", label: "George (British male)" },
{ id: "bm_fable", label: "Fable (British male)" },
] as const;
/** Default voice within the default model — Kokoro's flagship grade-A voice. */
export const DEFAULT_TTS_VOICE = "af_heart";
/** Default local TTS model used when `tts.localModel` is unset. */
export const DEFAULT_TTS_LOCAL_MODEL_KEY = "kokoro";
/**
* Local TTS model registry. Kokoro-82M is the on-device SoTA tiny TTS (tops the
* TTS Arena leaderboard); the `onnx-community` ONNX export runs through
* `kokoro-js` on the shared transformers.js/onnxruntime worker. q8 keeps the
* weights ~100 MB and CPU inference fast while preserving quality. One model
* spans every voice/accent — language selection is a voice choice, not a
* separate download.
*/
export const TTS_LOCAL_MODELS = [
{
key: "kokoro",
repo: "onnx-community/Kokoro-82M-v1.0-ONNX",
dtype: "q8",
sampleRate: 24_000,
label: "Kokoro-82M",
description: "Kokoro-82M neural TTS — SoTA on-device quality, multi-voice, fully local",
voices: KOKORO_VOICES,
},
] as const satisfies readonly TtsLocalModelSpec[];
export type TtsLocalModelKey = (typeof TTS_LOCAL_MODELS)[number]["key"];
export const TTS_LOCAL_MODEL_VALUES = ["kokoro"] as const;
type MissingTtsModelValue = Exclude<TtsLocalModelKey, (typeof TTS_LOCAL_MODEL_VALUES)[number]>;
type ExtraTtsModelValue = Exclude<(typeof TTS_LOCAL_MODEL_VALUES)[number], TtsLocalModelKey>;
const TTS_LOCAL_MODEL_VALUES_MATCH_REGISTRY: MissingTtsModelValue extends never
? ExtraTtsModelValue extends never
? true
: never
: never = true;
void TTS_LOCAL_MODEL_VALUES_MATCH_REGISTRY;
export const TTS_LOCAL_MODEL_OPTIONS = [
{
value: "kokoro",
label: "Kokoro-82M",
description: "Kokoro-82M neural TTS — SoTA on-device quality, multi-voice, fully local",
},
] as const satisfies ReadonlyArray<{ value: TtsLocalModelKey; label: string; description: string }>;
/** Voice options for the `tts.localVoice` setting picker (default model's catalog). */
export const TTS_LOCAL_VOICE_OPTIONS = KOKORO_VOICES.map(voice => ({
value: voice.id,
label: voice.label,
})) as ReadonlyArray<{ value: string; label: string }>;
/** Accepted `tts.localVoice` values (default model's catalog) for schema validation. */
export const TTS_LOCAL_VOICE_VALUES = KOKORO_VOICES.map(voice => voice.id) as readonly string[];
export function getTtsLocalModelSpec(key: string): TtsLocalModelSpec | undefined {
return TTS_LOCAL_MODELS.find(model => model.key === key);
}
export function isTtsLocalModelKey(value: string): value is TtsLocalModelKey {
return getTtsLocalModelSpec(value) !== undefined;
}
/** Resolve a model key (or the default) to its Hugging Face repo id. */
export function resolveTtsRepo(modelKey: string | undefined): string {
const spec = (modelKey && getTtsLocalModelSpec(modelKey)) || getTtsLocalModelSpec(DEFAULT_TTS_LOCAL_MODEL_KEY);
if (!spec) throw new Error(`No local TTS model registered for key: ${modelKey ?? DEFAULT_TTS_LOCAL_MODEL_KEY}`);
return spec.repo;
}
/**
* Resolve a requested voice id to a concrete voice the model supports, falling
* back to the model's default voice (first entry) when the id is unknown or the
* legacy `"default"` sentinel. The returned id is always a valid Kokoro voice.
*/
export function resolveTtsVoice(modelKey: string | undefined, voice: string | undefined): string {
const spec = (modelKey && getTtsLocalModelSpec(modelKey)) || getTtsLocalModelSpec(DEFAULT_TTS_LOCAL_MODEL_KEY);
const fallback = spec?.voices[0]?.id ?? DEFAULT_TTS_VOICE;
if (!spec || !voice) return fallback;
const match = spec.voices.find(v => v.id === voice);
return match ? match.id : fallback;
}
@@ -0,0 +1,60 @@
import { describe, expect, it } from "bun:test";
import { playerCommandsFor } from "./player";
describe("playerCommandsFor", () => {
const noTools = { which: () => null, ffmpeg: () => null };
it("uses the built-in afplay on darwin", () => {
expect(playerCommandsFor("darwin", "/tmp/a.wav", noTools)).toEqual([{ cmd: "afplay", args: ["/tmp/a.wav"] }]);
});
it("uses PowerShell SoundPlayer on win32", () => {
const cmds = playerCommandsFor("win32", "C:\\tmp\\a.wav", noTools);
expect(cmds).toEqual([
{
cmd: "powershell",
args: ["-NoProfile", "-Command", "(New-Object Media.SoundPlayer 'C:\\tmp\\a.wav').PlaySync()"],
},
]);
});
it("prefers paplay then aplay on linux", () => {
const cmds = playerCommandsFor("linux", "/tmp/a.wav", {
which: bin => (bin === "paplay" ? "/usr/bin/paplay" : bin === "aplay" ? "/usr/bin/aplay" : null),
ffmpeg: () => null,
});
expect(cmds).toEqual([
{ cmd: "/usr/bin/paplay", args: ["/tmp/a.wav"] },
{ cmd: "/usr/bin/aplay", args: ["/tmp/a.wav"] },
]);
});
it("skips missing players and keeps the ones present", () => {
const cmds = playerCommandsFor("linux", "/tmp/a.wav", {
which: bin => (bin === "aplay" ? "/usr/bin/aplay" : null),
ffmpeg: () => null,
});
expect(cmds).toEqual([{ cmd: "/usr/bin/aplay", args: ["/tmp/a.wav"] }]);
});
it("falls back to the bundled ffmpeg (pulse then alsa) on linux", () => {
const cmds = playerCommandsFor("linux", "/tmp/a.wav", {
which: () => null,
ffmpeg: () => "/tools/ffmpeg",
});
expect(cmds).toEqual([
{
cmd: "/tools/ffmpeg",
args: ["-loglevel", "error", "-nostdin", "-i", "/tmp/a.wav", "-f", "pulse", "default"],
},
{
cmd: "/tools/ffmpeg",
args: ["-loglevel", "error", "-nostdin", "-i", "/tmp/a.wav", "-f", "alsa", "default"],
},
]);
});
it("returns no commands when no linux player is available", () => {
expect(playerCommandsFor("linux", "/tmp/a.wav", noTools)).toEqual([]);
});
});
+137
View File
@@ -0,0 +1,137 @@
/**
* Cross-platform audio-file playback via the system's built-in players.
*
* The selection logic is split into a pure, injectable builder
* ({@link playerCommandsFor}) so it can be unit-tested without spawning a
* process or touching PATH, and a thin runtime wrapper ({@link playAudioFile})
* that walks the resulting fallback chain.
*/
import * as fs from "node:fs/promises";
import { $which } from "@oh-my-pi/pi-utils";
import { getToolPath } from "../utils/tools-manager";
export interface PlayerCommand {
cmd: string;
args: string[];
}
/** Injection seam for {@link playerCommandsFor} — defaults to real PATH/tools lookups. */
export interface PlayerLookup {
which?: (bin: string) => string | null;
ffmpeg?: () => string | null;
}
/**
* Build the ordered list of playback commands to try for `filePath` on the
* given platform. Pure + injectable so the selection logic is testable without
* spawning anything.
*
* - darwin: `afplay` (always present on macOS).
* - win32: PowerShell `Media.SoundPlayer.PlaySync()` (no extra deps).
* - linux/other POSIX: `paplay` (PulseAudio) → `aplay` (ALSA) → the bundled
* static `ffmpeg` (`-f pulse` then `-f alsa`). Empty result means nothing is
* available and the caller should surface an install hint.
*/
export function playerCommandsFor(
platform: NodeJS.Platform,
filePath: string,
lookup: PlayerLookup = {},
): PlayerCommand[] {
const which = lookup.which ?? $which;
const ffmpeg = lookup.ffmpeg ?? ((): string | null => getToolPath("ffmpeg"));
if (platform === "darwin") {
return [{ cmd: "afplay", args: [filePath] }];
}
if (platform === "win32") {
return [
{
cmd: "powershell",
args: ["-NoProfile", "-Command", `(New-Object Media.SoundPlayer '${filePath}').PlaySync()`],
},
];
}
// Linux and other POSIX desktops share the PulseAudio/ALSA fallback chain.
const commands: PlayerCommand[] = [];
const paplay = which("paplay");
if (paplay) commands.push({ cmd: paplay, args: [filePath] });
const aplay = which("aplay");
if (aplay) commands.push({ cmd: aplay, args: [filePath] });
const ffmpegBin = ffmpeg();
if (ffmpegBin) {
commands.push({
cmd: ffmpegBin,
args: ["-loglevel", "error", "-nostdin", "-i", filePath, "-f", "pulse", "default"],
});
commands.push({
cmd: ffmpegBin,
args: ["-loglevel", "error", "-nostdin", "-i", filePath, "-f", "alsa", "default"],
});
}
return commands;
}
export interface PlayAudioOptions {
signal?: AbortSignal;
}
function playbackAbortError(signal: AbortSignal): Error {
const reason = signal.reason;
return reason instanceof Error ? reason : new DOMException("Audio playback aborted", "AbortError");
}
/**
* Play `filePath` through the speakers, trying each candidate command in order
* and returning on the first clean exit. Throws an actionable Error if no
* player exists or every candidate fails (with the collected stderr).
*/
export async function playAudioFile(filePath: string, options: PlayAudioOptions = {}): Promise<void> {
const { signal } = options;
if (signal?.aborted) throw playbackAbortError(signal);
const commands = playerCommandsFor(process.platform, filePath);
if (commands.length === 0) {
throw new Error(
"No audio player available. Install PulseAudio (paplay) or ALSA (aplay), " +
"or run `omp setup speech` to download a bundled ffmpeg.",
);
}
const failures: string[] = [];
for (const command of commands) {
if (signal?.aborted) throw playbackAbortError(signal);
try {
const proc = Bun.spawn([command.cmd, ...command.args], { stdout: "ignore", stderr: "pipe" });
let killTimer: NodeJS.Timeout | undefined;
const abort = (): void => {
proc.kill("SIGTERM");
killTimer = setTimeout(() => proc.kill("SIGKILL"), 500);
killTimer.unref?.();
};
signal?.addEventListener("abort", abort, { once: true });
try {
const code = await proc.exited;
if (signal?.aborted) throw playbackAbortError(signal);
if (code === 0) return;
let stderr = "";
if (proc.stderr && typeof proc.stderr !== "number") {
stderr = await new Response(proc.stderr as ReadableStream).text();
}
failures.push(`${command.cmd} exited ${code}${stderr.trim() ? `: ${stderr.trim()}` : ""}`);
} finally {
signal?.removeEventListener("abort", abort);
if (killTimer) clearTimeout(killTimer);
}
} catch (err) {
if (signal?.aborted) throw playbackAbortError(signal);
failures.push(`${command.cmd}: ${err instanceof Error ? err.message : String(err)}`);
}
}
throw new Error(`Audio playback failed:\n${failures.join("\n")}`);
}
/** Best-effort temp-file cleanup used by callers after playback. */
export async function removeTempFile(filePath: string): Promise<void> {
await fs.unlink(filePath).catch(() => {});
}
+21
View File
@@ -0,0 +1,21 @@
import * as path from "node:path";
import { getTinyModelsCacheDir } from "@oh-my-pi/pi-utils";
export const KOKORO_PACKAGE = "kokoro-js";
export const KOKORO_VERSION = "1.2.1";
export const ONNXRUNTIME_NODE_PACKAGE = "onnxruntime-node";
export const ONNXRUNTIME_NODE_VERSION = "1.26.0";
export function getTtsRuntimeDir(): string {
const runtimeKey = KOKORO_VERSION.replace(/[^A-Za-z0-9._-]/g, "_");
return path.join(path.dirname(getTinyModelsCacheDir()), "tts-runtime", `kokoro-${runtimeKey}`);
}
export async function isTtsRuntimeCached(): Promise<boolean> {
try {
const pkg = await Bun.file(path.join(getTtsRuntimeDir(), "node_modules", KOKORO_PACKAGE, "package.json")).json();
return typeof pkg === "object" && pkg !== null && "version" in pkg && pkg.version === KOKORO_VERSION;
} catch {
return false;
}
}
@@ -0,0 +1,55 @@
import { describe, expect, it } from "bun:test";
import { streamingPlayerCommandsFor } from "./streaming-player";
// The exact argv matters: a wrong flag means ffmpeg/paplay/aplay produce no
// audio. These pin the raw-PCM stdin contract (32-bit float, mono, the chunk's
// rate) and the per-platform backend selection. macOS/Windows have no streaming
// backend (file-only players), so they fall back to per-file playback.
describe("streamingPlayerCommandsFor", () => {
const ffmpegOnly = { which: () => null, ffmpeg: () => "/opt/ffmpeg" };
it("has no streaming backend on macOS (afplay is file-only)", () => {
expect(streamingPlayerCommandsFor("darwin", 24_000, ffmpegOnly)).toEqual([]);
});
it("has no streaming backend on Windows", () => {
expect(streamingPlayerCommandsFor("win32", 24_000, ffmpegOnly)).toEqual([]);
});
it("streams raw f32le mono PCM via ffmpeg (pulse then alsa) on Linux", () => {
const cmds = streamingPlayerCommandsFor("linux", 24_000, { which: () => null, ffmpeg: () => "/usr/bin/ffmpeg" });
expect(cmds).toHaveLength(2);
expect(cmds[0]?.cmd).toBe("/usr/bin/ffmpeg");
expect(cmds[0]?.args).toEqual(
expect.arrayContaining(["-f", "f32le", "-ar", "24000", "-ac", "1", "-i", "pipe:0", "-f", "pulse"]),
);
expect(cmds[1]?.args).toContain("alsa");
});
it("prefers ffmpeg, then paplay/aplay raw fallbacks on Linux", () => {
const cmds = streamingPlayerCommandsFor("linux", 24_000, {
which: bin => (bin === "paplay" ? "/usr/bin/paplay" : bin === "aplay" ? "/usr/bin/aplay" : null),
ffmpeg: () => "/usr/bin/ffmpeg",
});
const labels = cmds.map(c =>
c.cmd === "/usr/bin/ffmpeg" ? (c.args.includes("pulse") ? "ffmpeg:pulse" : "ffmpeg:alsa") : c.cmd,
);
expect(labels).toEqual(["ffmpeg:pulse", "ffmpeg:alsa", "/usr/bin/paplay", "/usr/bin/aplay"]);
});
it("falls back to aplay raw when ffmpeg is absent on Linux", () => {
const cmds = streamingPlayerCommandsFor("linux", 48_000, {
which: bin => (bin === "aplay" ? "/usr/bin/aplay" : null),
ffmpeg: () => null,
});
expect(cmds).toHaveLength(1);
expect(cmds[0]?.cmd).toBe("/usr/bin/aplay");
expect(cmds[0]?.args).toEqual(["-q", "-f", "FLOAT_LE", "-r", "48000", "-c", "1", "-"]);
});
it("encodes the chunk's sample rate into the command", () => {
const cmds = streamingPlayerCommandsFor("linux", 16_000, { which: () => null, ffmpeg: () => "/usr/bin/ffmpeg" });
expect(cmds[0]?.args).toContain("16000");
expect(cmds[0]?.args).not.toContain("24000");
});
});
@@ -0,0 +1,266 @@
/**
* Gapless streaming audio output for assistant speech.
*
* Replaces the spawn-`afplay`-per-sentence approach (a fresh process per chunk
* meant audible gaps, per-spawn latency, and no way to interrupt a clip mid-play)
* with a single persistent player process fed raw 32-bit-float mono PCM over
* stdin. Chunks are queued and drained by one writer so sentences play back to
* back; writes are paced to stay only {@link LEAD_SECONDS} ahead of realtime so
* ducking and stop take effect promptly instead of after seconds of buffered
* audio. {@link StreamingAudioPlayer.stop} kills the process for instant silence.
*
* Where no streaming backend exists (Windows, or macOS without the bundled
* ffmpeg), it degrades to the per-file {@link playAudioFile} path so speech still
* works — just without gapless playback or mid-clip interruption.
*/
import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { $which, logger, Snowflake } from "@oh-my-pi/pi-utils";
import type { FileSink, Subprocess } from "bun";
import { getToolPath } from "../utils/tools-manager";
import { type PlayerCommand, playAudioFile } from "./player";
import { encodeWav } from "./wav";
/** Kokoro emits 24 kHz mono; used when a chunk does not declare a rate. */
const DEFAULT_SAMPLE_RATE = 24_000;
/** Cap how far ahead of realtime we buffer into the player so duck/stop are responsive. */
const LEAD_SECONDS = 0.6;
/** Output gain applied while ducked (the user is speaking over the assistant). */
export const DUCK_GAIN = 0.25;
/** Injection seam for {@link streamingPlayerCommandsFor} — defaults to real PATH/tools lookups. */
export interface StreamingPlayerLookup {
which?: (bin: string) => string | null;
ffmpeg?: () => string | null;
}
/**
* Ordered candidate commands for a persistent raw-PCM player on `platform`: each
* reads 32-bit-float little-endian mono PCM at `sampleRate` from stdin (`pipe:0`)
* and plays it to the default output device. An empty list means no streaming
* backend is available and the caller should fall back to per-file playback.
*
* - darwin: none; `afplay` is file-only, so macOS uses the interruptible
* per-file fallback.
* - linux/other POSIX: `ffmpeg` (`-f pulse` then `-f alsa`) → `paplay`/`aplay`
* raw fallbacks.
* - win32: none (PowerShell `SoundPlayer` is file-only).
*/
export function streamingPlayerCommandsFor(
platform: NodeJS.Platform,
sampleRate: number,
lookup: StreamingPlayerLookup = {},
): PlayerCommand[] {
const which = lookup.which ?? $which;
const ffmpeg = lookup.ffmpeg ?? ((): string | null => getToolPath("ffmpeg"));
const rate = String(sampleRate > 0 ? sampleRate : DEFAULT_SAMPLE_RATE);
const input = ["-loglevel", "error", "-nostdin", "-f", "f32le", "-ar", rate, "-ac", "1", "-i", "pipe:0"];
if (platform === "darwin") return [];
if (platform === "win32") {
return [];
}
const commands: PlayerCommand[] = [];
const ffmpegBin = ffmpeg();
if (ffmpegBin) {
commands.push({ cmd: ffmpegBin, args: [...input, "-f", "pulse", "default"] });
commands.push({ cmd: ffmpegBin, args: [...input, "-f", "alsa", "default"] });
}
const paplay = which("paplay");
if (paplay) commands.push({ cmd: paplay, args: ["--raw", `--rate=${rate}`, "--format=float32le", "--channels=1"] });
const aplay = which("aplay");
if (aplay) commands.push({ cmd: aplay, args: ["-q", "-f", "FLOAT_LE", "-r", rate, "-c", "1", "-"] });
return commands;
}
/**
* Single-session gapless player. Lifecycle: {@link start} once, {@link write}
* chunks in order, then {@link end} to drain or {@link stop} to abort. Not
* reusable after stop/end — create a new instance per utterance.
*/
export class StreamingAudioPlayer {
#queue: Float32Array[] = [];
#sampleRate = DEFAULT_SAMPLE_RATE;
#gain = 1;
#mode: "stream" | "file" = "file";
#proc: Subprocess<"pipe", "ignore", "ignore"> | null = null;
#sink: FileSink | null = null;
#writtenSec = 0;
#startedAt = 0;
#started = false;
#inputClosed = false;
#stopped = false;
#abortController = new AbortController();
#wake: (() => void) | null = null;
#drain: Promise<void> = Promise.resolve();
/** Pick a backend and begin draining. Idempotent; the first call's rate wins. */
start(sampleRate: number): void {
if (this.#started || this.#stopped) return;
this.#started = true;
this.#sampleRate = sampleRate > 0 ? sampleRate : DEFAULT_SAMPLE_RATE;
this.#mode = this.#spawnStream() ? "stream" : "file";
this.#startedAt = performance.now();
this.#drain = this.#drainLoop();
}
/** Queue a mono float32 PCM chunk for playback in arrival order. */
write(pcm: Float32Array): void {
if (this.#stopped) return;
this.#queue.push(pcm);
this.#signal();
}
/** Scale subsequent output (1 = normal, <1 = ducked). Applies within {@link LEAD_SECONDS}. */
setGain(gain: number): void {
this.#gain = gain < 0 ? 0 : gain;
}
/** Close the input; resolves once all queued audio has finished playing. */
async end(): Promise<void> {
this.#inputClosed = true;
this.#signal();
await this.#drain;
}
/** Stop immediately: kill the player, drop everything still queued. */
stop(): void {
if (this.#stopped) return;
this.#stopped = true;
this.#queue.length = 0;
this.#abortController.abort();
this.#signal();
try {
this.#sink?.end();
} catch {}
try {
this.#proc?.kill("SIGKILL");
} catch {}
}
#spawnStream(): boolean {
for (const command of streamingPlayerCommandsFor(process.platform, this.#sampleRate)) {
try {
const proc = Bun.spawn([command.cmd, ...command.args], {
stdin: "pipe",
stdout: "ignore",
stderr: "ignore",
});
this.#proc = proc;
this.#sink = proc.stdin;
return true;
} catch (error) {
logger.debug("tts: streaming player spawn failed", {
cmd: command.cmd,
error: error instanceof Error ? error.message : String(error),
});
}
}
return false;
}
#signal(): void {
const wake = this.#wake;
this.#wake = null;
wake?.();
}
async #drainLoop(): Promise<void> {
try {
while (!this.#stopped) {
const chunk = this.#queue.shift();
if (!chunk) {
if (this.#inputClosed) break;
await this.#waitForWork();
continue;
}
if (this.#mode === "stream") {
// Pace writes so the player buffers ~LEAD_SECONDS, no more, keeping
// ducking and stop responsive instead of locked behind buffered audio.
const ahead = this.#writtenSec - (performance.now() - this.#startedAt) / 1000;
if (ahead > LEAD_SECONDS) {
await Bun.sleep((ahead - LEAD_SECONDS) * 1000);
if (this.#stopped) return;
}
this.#writeStream(chunk);
this.#writtenSec += chunk.length / this.#sampleRate;
} else {
await this.#playFile(chunk);
}
}
if (!this.#stopped && this.#mode === "stream") {
try {
await this.#sink?.end();
} catch {}
if (this.#proc) {
try {
await this.#proc.exited;
} catch {}
}
}
} catch (error) {
logger.debug("tts: streaming player drain failed", {
error: error instanceof Error ? error.message : String(error),
});
}
}
/** Block until a chunk is queued, the input closes, or stop is called. */
#waitForWork(): Promise<void> {
const { promise, resolve } = Promise.withResolvers<void>();
this.#wake = resolve;
// Re-check after arming to close the gap between the empty shift and here.
if (this.#queue.length > 0 || this.#inputClosed || this.#stopped) {
this.#wake = null;
resolve();
}
return promise;
}
#writeStream(pcm: Float32Array): void {
const sink = this.#sink;
if (!sink) return;
try {
sink.write(this.#bytes(pcm));
sink.flush();
} catch (error) {
logger.debug("tts: streaming write failed", {
error: error instanceof Error ? error.message : String(error),
});
}
}
async #playFile(pcm: Float32Array): Promise<void> {
const wavPath = path.join(os.tmpdir(), `omp-speech-${Snowflake.next()}.wav`);
try {
await fs.writeFile(wavPath, encodeWav(this.#scaled(pcm), this.#sampleRate));
if (!this.#stopped) await playAudioFile(wavPath, { signal: this.#abortController.signal });
} catch (error) {
logger.debug("tts: file playback failed", {
error: error instanceof Error ? error.message : String(error),
});
} finally {
await fs.unlink(wavPath).catch(() => {});
}
}
/** Raw f32le bytes for the stream sink, applying gain only when ducked (avoids a copy at unity). */
#bytes(pcm: Float32Array): Uint8Array {
if (this.#gain === 1) return new Uint8Array(pcm.buffer, pcm.byteOffset, pcm.byteLength);
return new Uint8Array(this.#scaled(pcm).buffer);
}
#scaled(pcm: Float32Array): Float32Array {
if (this.#gain === 1) return pcm;
const out = new Float32Array(pcm.length);
for (let i = 0; i < pcm.length; i++) out[i] = (pcm[i] ?? 0) * this.#gain;
return out;
}
}
/** Factory the vocalizer calls; a function so tests can stub it without spawning a player. */
export function createStreamingPlayer(): StreamingAudioPlayer {
return new StreamingAudioPlayer();
}
+647
View File
@@ -0,0 +1,647 @@
import * as path from "node:path";
import { $env, isBunTestRuntime, isCompiledBinary, logger, workerHostEntry } from "@oh-my-pi/pi-utils";
import type { Subprocess } from "bun";
import { settings } from "../config/settings";
import { tinyWorkerEnvOverlay } from "../tiny/title-client";
import { isTtsLocalModelKey, type TtsLocalModelKey } from "./models";
import type { TtsProgressEvent, TtsWorkerInbound, TtsWorkerOutbound } from "./tts-protocol";
/** Decoded PCM returned by a local synthesis request. */
export interface TtsAudio {
pcm: Float32Array;
sampleRate: number;
}
/**
* Abstraction over the TTS subprocess. The runtime implementation is a Bun child
* process so `onnxruntime-node`'s NAPI finalizer never runs inside the main agent
* address space — that destructor segfaults Bun during shutdown (issue #1606).
*/
interface WorkerHandle {
send(message: TtsWorkerInbound): void;
onMessage(handler: (message: TtsWorkerOutbound) => void): () => void;
onError(handler: (error: Error) => void): () => void;
/** Re-reference the subprocess so a pending request keeps the parent event loop alive. */
ref(): void;
/** Drop the reference once the worker is idle so it never blocks process exit. */
unref(): void;
terminate(): Promise<void>;
}
type PendingRequest =
| { kind: "synthesize"; modelKey: TtsLocalModelKey; resolve: (audio: TtsAudio | null) => void }
| { kind: "download"; modelKey: TtsLocalModelKey; resolve: (ok: boolean) => void }
| { kind: "stream"; modelKey: TtsLocalModelKey; channel: AudioChunkChannel };
export interface TtsSynthesizeOptions {
voice?: string;
signal?: AbortSignal;
}
export interface TtsDownloadOptions {
signal?: AbortSignal;
onProgress?: (event: TtsProgressEvent) => void;
}
export interface TtsStreamOptions {
voice?: string;
signal?: AbortSignal;
}
/** One synthesized sentence of a streaming session, in emission order. */
export interface TtsAudioChunk {
index: number;
text: string;
pcm: Float32Array;
sampleRate: number;
}
/**
* A live streaming-synthesis session. Feed text incrementally with {@link push}
* and close the input with {@link end}; `chunks` yields each synthesized
* sentence's audio as soon as it is ready, then completes once the worker
* finishes draining the closed input.
*/
export interface TtsStreamHandle {
push(text: string): void;
end(): void;
chunks: AsyncIterableIterator<TtsAudioChunk>;
}
/**
* Single-producer/single-consumer async queue bridging the worker's IPC
* `audio-chunk` messages to an async iterator. Chunks pushed while no consumer
* is awaiting are buffered in order; {@link close} ends the iterator and
* {@link fail} surfaces an error to the awaiting (or next) consumer.
*/
class AudioChunkChannel {
#queue: TtsAudioChunk[] = [];
#waiters: Array<{
resolve: (result: IteratorResult<TtsAudioChunk>) => void;
reject: (error: Error) => void;
}> = [];
#error: Error | null = null;
#settled = false;
#onSettle: (() => void) | undefined;
constructor(onSettle?: () => void) {
this.#onSettle = onSettle;
}
push(chunk: TtsAudioChunk): void {
if (this.#settled) return;
const waiter = this.#waiters.shift();
if (waiter) waiter.resolve({ value: chunk, done: false });
else this.#queue.push(chunk);
}
close(): void {
this.#settle(null);
}
fail(error: Error): void {
this.#settle(error);
}
#settle(error: Error | null): void {
if (this.#settled) return;
this.#settled = true;
this.#error = error;
for (const waiter of this.#waiters) {
if (error) waiter.reject(error);
else waiter.resolve({ value: undefined, done: true });
}
this.#waiters = [];
this.#onSettle?.();
}
async *iterator(): AsyncIterableIterator<TtsAudioChunk> {
while (true) {
const buffered = this.#queue.shift();
if (buffered) {
yield buffered;
continue;
}
if (this.#error) throw this.#error;
if (this.#settled) return;
const { promise, resolve, reject } = Promise.withResolvers<IteratorResult<TtsAudioChunk>>();
this.#waiters.push({ resolve, reject });
const result = await promise;
if (result.done) return;
yield result.value;
}
}
}
// Cold-starting the worker from a compiled binary (decompress + module graph load)
// is slow on contended CI runners; the probe only proves the worker spawns and
// ponges, so a generous bound removes flakes without weakening the check.
const SMOKE_TEST_TIMEOUT_MS = 30_000;
/**
* Hidden subcommand on the main CLI that boots the TTS worker in the spawned
* subprocess. Kept in sync with the dispatch in `cli.ts` (Main-owned).
*/
export const TTS_WORKER_ARG = "__omp_tts_worker";
function readTinyModelSetting(path: "providers.tinyModelDevice" | "providers.tinyModelDtype"): string | undefined {
try {
const value = settings.get(path);
return typeof value === "string" ? value : undefined;
} catch {
// Settings may be uninitialized (e.g. `omp --smoke-test`); fall back to env/default.
return undefined;
}
}
/**
* Env handed to the TTS subprocess. The `PI_TINY_DEVICE` / `PI_TINY_DTYPE` env
* vars win; otherwise the persisted `providers.tinyModelDevice` /
* `providers.tinyModelDtype` settings are mapped onto those vars so the
* subprocess's env-based resolution governs speech the same way it governs the
* tiny LLM worker.
*/
function ttsWorkerEnv(): Record<string, string> {
const overlay = tinyWorkerEnvOverlay(
$env,
readTinyModelSetting("providers.tinyModelDevice"),
readTinyModelSetting("providers.tinyModelDtype"),
);
const base = $env as Record<string, string | undefined>;
const merged: Record<string, string> = {};
for (const key in base) {
const value = base[key];
if (typeof value === "string") merged[key] = value;
}
for (const key in overlay) merged[key] = overlay[key];
return merged;
}
interface TtsWorkerSpawnCommand {
cmd: string[];
cwd?: string;
}
/**
* Resolve the command used to relaunch the agent CLI into TTS-worker mode. In a
* compiled binary the entry point is the binary itself; otherwise re-enter the
* declared worker-host entry (cwd-relative for reliable Bun IPC), falling back
* to this package's own `src/cli.ts` when no host entry is declared (bun test).
*/
function ttsWorkerSpawnCmd(): TtsWorkerSpawnCommand {
if (isCompiledBinary()) return { cmd: [process.execPath, TTS_WORKER_ARG] };
const hostEntry = workerHostEntry();
if (hostEntry) {
return { cmd: [process.execPath, path.basename(hostEntry), TTS_WORKER_ARG], cwd: path.dirname(hostEntry) };
}
const packageRoot = path.resolve(import.meta.dir, "..", "..");
return { cmd: [process.execPath, "src/cli.ts", TTS_WORKER_ARG], cwd: packageRoot };
}
interface SpawnedSubprocess {
proc: Subprocess<"ignore", "ignore", "ignore">;
inbound: Set<(message: TtsWorkerOutbound) => void>;
errors: Set<(error: Error) => void>;
/** Flipped to `true` right before the deliberate SIGKILL so `onExit` can tell it apart from a crash. */
intentionalExit: { value: boolean };
}
/**
* Spawn the TTS worker as a subprocess. Exported for tests and the smoke probe;
* production callers go through {@link spawnTtsWorker}.
*/
export function createTtsSubprocess(): SpawnedSubprocess {
const inbound = new Set<(message: TtsWorkerOutbound) => void>();
const errors = new Set<(error: Error) => void>();
const intentionalExit = { value: false };
const spawnCommand = ttsWorkerSpawnCmd();
const proc = Bun.spawn({
cmd: spawnCommand.cmd,
cwd: spawnCommand.cwd,
env: ttsWorkerEnv(),
stdin: "ignore",
stdout: "ignore",
stderr: "ignore",
serialization: "advanced",
windowsHide: true,
ipc(message) {
for (const handler of inbound) handler(message as TtsWorkerOutbound);
},
onExit(_proc, exitCode, signalCode) {
if (exitCode === 0) return;
if (exitCode === null && intentionalExit.value) return;
const reason = exitCode !== null ? `code ${exitCode}` : `signal ${signalCode ?? "unknown"}`;
const err = new Error(`tts subprocess exited with ${reason}`);
for (const handler of errors) handler(err);
},
});
// Don't keep the parent event loop alive on an idle worker; the dispose path
// calls `terminate()` explicitly. Bun's test runner starves IPC for unref'd
// subprocesses, so keep it referenced only under tests.
if (!isBunTestRuntime()) proc.unref();
return { proc, inbound, errors, intentionalExit };
}
function wrapSubprocess({ proc, inbound, errors, intentionalExit }: SpawnedSubprocess): WorkerHandle {
return {
send(message) {
try {
proc.send(message);
} catch (error) {
logger.debug("tts: send to subprocess failed", {
error: error instanceof Error ? error.message : String(error),
});
}
},
onMessage(handler) {
inbound.add(handler);
return () => inbound.delete(handler);
},
onError(handler) {
errors.add(handler);
return () => errors.delete(handler);
},
ref() {
try {
proc.ref();
} catch {
// Already gone.
}
},
unref() {
try {
proc.unref();
} catch {
// Already gone.
}
},
async terminate() {
// SIGKILL: the point of subprocess isolation is that the parent never
// runs `onnxruntime-node`'s NAPI finalizer (it crashes Bun on Windows).
// Hard-kill instead; the OS reclaims the model memory.
intentionalExit.value = true;
try {
proc.kill("SIGKILL");
} catch {
// Already gone.
}
},
};
}
function spawnInlineUnavailableWorker(error: unknown): WorkerHandle {
const listeners = new Set<(message: TtsWorkerOutbound) => void>();
const errorMessage = error instanceof Error ? error.message : String(error);
const emit = (message: TtsWorkerOutbound): void => {
for (const listener of listeners) listener(message);
};
return {
send(message) {
queueMicrotask(() => {
if (message.type === "ping") {
emit({ type: "pong", id: message.id });
return;
}
emit({ type: "error", id: message.id, error: errorMessage });
});
},
onMessage(handler) {
listeners.add(handler);
return () => listeners.delete(handler);
},
onError() {
return () => {};
},
ref() {},
unref() {},
async terminate() {
listeners.clear();
},
};
}
function spawnTtsWorker(): WorkerHandle {
try {
return wrapSubprocess(createTtsSubprocess());
} catch (error) {
logger.warn("TTS worker spawn failed; local TTS disabled", {
error: error instanceof Error ? error.message : String(error),
});
return spawnInlineUnavailableWorker(error);
}
}
function logWorkerMessage(message: Extract<TtsWorkerOutbound, { type: "log" }>): void {
if (message.level === "debug") logger.debug(message.msg, message.meta);
else if (message.level === "warn") logger.warn(message.msg, message.meta);
else logger.error(message.msg, message.meta);
}
export class TtsClient {
#worker: WorkerHandle | null = null;
#unsubscribeMessage: (() => void) | null = null;
#unsubscribeError: (() => void) | null = null;
#pending = new Map<string, PendingRequest>();
#progressListeners = new Set<(event: TtsProgressEvent) => void>();
#nextRequestId = 0;
#refed = false;
#spawnWorker: () => WorkerHandle;
constructor(spawnWorker: () => WorkerHandle = spawnTtsWorker) {
this.#spawnWorker = spawnWorker;
}
onProgress(listener: (event: TtsProgressEvent) => void): () => void {
this.#progressListeners.add(listener);
return () => this.#progressListeners.delete(listener);
}
async synthesize(modelKey: string, text: string, options: TtsSynthesizeOptions = {}): Promise<TtsAudio | null> {
if (!isTtsLocalModelKey(modelKey)) return null;
if (options.signal?.aborted) return null;
try {
const worker = this.#ensureWorker();
const id = String(++this.#nextRequestId);
const { promise, resolve } = Promise.withResolvers<TtsAudio | null>();
this.#addPending(id, { kind: "synthesize", modelKey, resolve });
const abort = (): void => {
const pending = this.#pending.get(id);
if (pending?.kind !== "synthesize") return;
this.#deletePending(id);
pending.resolve(null);
};
options.signal?.addEventListener("abort", abort, { once: true });
try {
const request: TtsWorkerInbound = options.voice
? { type: "synthesize", id, modelKey, text, voice: options.voice }
: { type: "synthesize", id, modelKey, text };
worker.send(request);
return await promise;
} finally {
options.signal?.removeEventListener("abort", abort);
this.#deletePending(id);
}
} catch (error) {
logger.debug("tts: local synthesis failed", {
modelKey,
error: error instanceof Error ? error.message : String(error),
});
return null;
}
}
/**
* Open a streaming-synthesis session. Text is fed incrementally through the
* returned handle's `push`/`end`; audio is emitted one synthesized sentence at
* a time via `chunks`, so playback can begin before the full text is known.
* Returns an inert handle (immediately-ended `chunks`) for unknown models or
* an already-aborted signal, and fails the iterator if the worker cannot spawn.
*/
synthesizeStream(modelKey: string, options: TtsStreamOptions = {}): TtsStreamHandle {
if (!isTtsLocalModelKey(modelKey) || options.signal?.aborted) {
const channel = new AudioChunkChannel();
channel.close();
return { push: () => {}, end: () => {}, chunks: channel.iterator() };
}
let worker: WorkerHandle;
try {
worker = this.#ensureWorker();
} catch (error) {
logger.debug("tts: stream synthesis failed to start", {
modelKey,
error: error instanceof Error ? error.message : String(error),
});
const channel = new AudioChunkChannel();
channel.fail(error instanceof Error ? error : new Error(String(error)));
return { push: () => {}, end: () => {}, chunks: channel.iterator() };
}
const id = String(++this.#nextRequestId);
const signal = options.signal;
let closed = false;
let ended = false;
const abort = (): void => {
if (closed) return;
closed = true;
ended = true;
if (!this.#pending.has(id)) return;
this.#deletePending(id);
worker.send({ type: "stream-cancel", id });
channel.close();
};
const channel = new AudioChunkChannel(() => signal?.removeEventListener("abort", abort));
this.#addPending(id, { kind: "stream", modelKey, channel });
signal?.addEventListener("abort", abort, { once: true });
const start: TtsWorkerInbound = options.voice
? { type: "stream-start", id, modelKey, voice: options.voice }
: { type: "stream-start", id, modelKey };
worker.send(start);
return {
push: (text: string) => {
if (!closed && !ended) worker.send({ type: "stream-push", id, text });
},
end: () => {
if (closed || ended) return;
ended = true;
worker.send({ type: "stream-end", id });
},
chunks: channel.iterator(),
};
}
async downloadModel(modelKey: string, options: TtsDownloadOptions = {}): Promise<boolean> {
if (!isTtsLocalModelKey(modelKey)) return false;
if (options.signal?.aborted) return false;
const unsubscribe = options.onProgress ? this.onProgress(options.onProgress) : undefined;
try {
const worker = this.#ensureWorker();
const id = String(++this.#nextRequestId);
const { promise, resolve } = Promise.withResolvers<boolean>();
this.#addPending(id, { kind: "download", modelKey, resolve });
const abort = (): void => {
const pending = this.#pending.get(id);
if (pending?.kind !== "download") return;
this.#deletePending(id);
pending.resolve(false);
};
options.signal?.addEventListener("abort", abort, { once: true });
try {
worker.send({ type: "download", id, modelKey });
return await promise;
} finally {
options.signal?.removeEventListener("abort", abort);
this.#deletePending(id);
}
} catch (error) {
logger.debug("tts: local model download failed", {
modelKey,
error: error instanceof Error ? error.message : String(error),
});
return false;
} finally {
unsubscribe?.();
}
}
async terminate(): Promise<void> {
const worker = this.#worker;
this.#worker = null;
this.#unsubscribeMessage?.();
this.#unsubscribeMessage = null;
this.#unsubscribeError?.();
this.#unsubscribeError = null;
for (const pending of this.#pending.values()) {
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "synthesize") pending.resolve(null);
else if (pending.kind === "download") pending.resolve(false);
else pending.channel.close();
}
this.#pending.clear();
this.#refed = false;
try {
await worker?.terminate();
} catch {
// Already gone.
}
}
#ensureWorker(): WorkerHandle {
if (this.#worker) return this.#worker;
const worker = this.#spawnWorker();
this.#worker = worker;
this.#unsubscribeMessage = worker.onMessage(message => this.#handleMessage(message));
this.#unsubscribeError = worker.onError(error => this.#handleWorkerError(error));
return worker;
}
/** Register a pending request and keep the worker referenced while work is in flight. */
#addPending(id: string, request: PendingRequest): void {
this.#pending.set(id, request);
this.#syncWorkerRef();
}
/** Drop a pending request and unref the worker once nothing is in flight. */
#deletePending(id: string): void {
if (this.#pending.delete(id)) this.#syncWorkerRef();
}
/**
* The TTS subprocess is spawned `unref`'d so an idle worker never blocks
* process exit. A short-lived CLI command (`omp say`) awaiting a request would
* otherwise let the event loop drain and exit before the audio arrives, so we
* `ref` the worker exactly while at least one request is pending.
*/
#syncWorkerRef(): void {
const worker = this.#worker;
if (!worker) return;
const shouldRef = this.#pending.size > 0;
if (shouldRef === this.#refed) return;
this.#refed = shouldRef;
if (shouldRef) worker.ref();
else worker.unref();
}
#handleMessage(message: TtsWorkerOutbound): void {
if (message.type === "log") {
logWorkerMessage(message);
return;
}
if (message.type === "progress") {
this.#emitProgress(message.event);
return;
}
if (message.type === "pong") return;
const pending = this.#pending.get(message.id);
if (!pending) return;
// Streaming chunks are non-terminal: keep the session registered until
// `stream-done` (or an error) so later chunks still route to its channel.
if (message.type === "audio-chunk") {
if (pending.kind === "stream") {
pending.channel.push({
index: message.index,
text: message.text,
pcm: message.pcm,
sampleRate: message.sampleRate,
});
}
return;
}
this.#deletePending(message.id);
if (message.type === "stream-done") {
if (pending.kind === "stream") pending.channel.close();
return;
}
if (message.type === "audio") {
if (pending.kind === "synthesize") pending.resolve({ pcm: message.pcm, sampleRate: message.sampleRate });
return;
}
if (message.type === "downloaded") {
if (pending.kind === "download") pending.resolve(true);
return;
}
logger.debug("tts: worker returned error", { error: message.error });
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "synthesize") pending.resolve(null);
else if (pending.kind === "download") pending.resolve(false);
else pending.channel.fail(new Error(message.error));
void this.terminate();
}
#emitProgress(event: TtsProgressEvent): void {
for (const listener of this.#progressListeners) listener(event);
}
#handleWorkerError(error: Error): void {
logger.warn("tts: worker error", { error: error.message });
for (const pending of this.#pending.values()) {
this.#emitProgress({ modelKey: pending.modelKey, status: "error" });
if (pending.kind === "synthesize") pending.resolve(null);
else if (pending.kind === "download") pending.resolve(false);
else pending.channel.fail(error);
}
this.#pending.clear();
void this.terminate();
}
}
export const ttsClient = new TtsClient();
export async function shutdownTtsClient(): Promise<void> {
await ttsClient.terminate();
}
export async function smokeTestTtsWorker({
timeoutMs = SMOKE_TEST_TIMEOUT_MS,
}: {
timeoutMs?: number;
} = {}): Promise<void> {
const handle = wrapSubprocess(createTtsSubprocess());
const { promise, resolve, reject } = Promise.withResolvers<void>();
const timer = setTimeout(() => reject(new Error(`tts worker did not pong within ${timeoutMs}ms`)), timeoutMs);
const unsubscribeMessage = handle.onMessage(message => {
if (message.type === "pong") {
resolve();
return;
}
if (message.type === "log") return;
reject(new Error(`tts worker: expected pong, got ${JSON.stringify(message)}`));
});
const unsubscribeError = handle.onError(reject);
try {
handle.send({ type: "ping", id: "smoke" } satisfies TtsWorkerInbound);
await promise;
} finally {
clearTimeout(timer);
unsubscribeMessage();
unsubscribeError();
await handle.terminate();
}
}
@@ -0,0 +1,60 @@
import type { TtsLocalModelKey } from "./models";
export type TtsProgressStatus = "initiate" | "download" | "progress" | "progress_total" | "done" | "ready" | "error";
export interface TtsProgressFileState {
loaded: number;
total: number;
}
export interface TtsProgressEvent {
modelKey: TtsLocalModelKey;
status: TtsProgressStatus;
name?: string;
file?: string;
progress?: number;
loaded?: number;
total?: number;
files?: Record<string, TtsProgressFileState>;
task?: string;
model?: string;
}
export type TtsWorkerInbound =
| { type: "ping"; id: string }
| { type: "synthesize"; id: string; modelKey: TtsLocalModelKey; text: string; voice?: string }
| { type: "download"; id: string; modelKey: TtsLocalModelKey }
// Streaming synthesis: a session is opened with `stream-start`, fed incrementally
// with `stream-push`, and closed with `stream-end`. `stream-cancel` interrupts
// without a final drain. The worker emits an `audio-chunk` per synthesized
// sentence and a final `stream-done` only for non-cancelled sessions.
| { type: "stream-start"; id: string; modelKey: TtsLocalModelKey; voice?: string }
| { type: "stream-push"; id: string; text: string }
| { type: "stream-end"; id: string }
| { type: "stream-cancel"; id: string };
export type TtsWorkerOutbound =
| { type: "pong"; id: string }
| { type: "audio"; id: string; pcm: Float32Array; sampleRate: number }
| { type: "downloaded"; id: string }
| { type: "error"; id: string; error: string }
| { type: "progress"; id: string; event: TtsProgressEvent }
| { type: "log"; level: "debug" | "warn" | "error"; msg: string; meta?: Record<string, unknown> }
// One synthesized sentence of a streaming session, in emission order, followed
// by a single `stream-done` once the input stream is closed and drained.
| { type: "audio-chunk"; id: string; index: number; text: string; pcm: Float32Array; sampleRate: number }
| { type: "stream-done"; id: string };
/**
* Wire transport between the parent (`TtsClient`) and the local TTS subprocess.
* The parent owns the subprocess lifecycle (graceful work, hard SIGKILL on
* shutdown); the protocol carries no explicit close handshake — once the parent
* decides to terminate, it signals the OS to reap the child so
* `onnxruntime-node`'s NAPI finalizer never runs in the main agent address
* space (it segfaults Bun on shutdown — issue #1606). See `tts-client.ts` for
* the spawn/kill glue.
*/
export interface TtsTransport {
send(message: TtsWorkerOutbound): void;
onMessage(handler: (message: TtsWorkerInbound) => void): () => void;
}
+497
View File
@@ -0,0 +1,497 @@
import { createRequire } from "node:module";
import * as path from "node:path";
import type { ProgressInfo, RawAudio } from "@huggingface/transformers";
import {
ensureRuntimeInstalled,
getTinyModelsCacheDir,
installRuntimeModuleResolver,
resolveRuntimeModule,
} from "@oh-my-pi/pi-utils";
import { resolveTinyModelDevicePreference, type TinyModelDevice, tinyModelDeviceLoadOrder } from "../tiny/device";
import { resolveTinyModelDtypeOverride, type TinyModelDtype } from "../tiny/dtype";
import { getTtsLocalModelSpec, resolveTtsVoice, type TtsLocalModelKey, type TtsLocalModelSpec } from "./models";
import {
getTtsRuntimeDir,
KOKORO_PACKAGE,
KOKORO_VERSION,
ONNXRUNTIME_NODE_PACKAGE,
ONNXRUNTIME_NODE_VERSION,
} from "./runtime";
import type { TtsProgressEvent, TtsTransport, TtsWorkerInbound } from "./tts-protocol";
const TTS_TASK = "text-to-speech";
const TRANSFORMERS_PACKAGE = "@huggingface/transformers";
// kokoro-js is NEVER a dependency of the main tree: its transformers@3.8.1 +
// onnxruntime-node@1.21 graph must not pollute it (1.21 segfaults Bun on session
// creation). It is lazily `bun install`ed into a side runtime dir on first use,
// with onnxruntime-node force-pinned to the Bun-safe version the rest of the
// stack runs. Bump KOKORO_VERSION to roll the cached runtime + model wrapper.
const ttsDevicePreference = resolveTinyModelDevicePreference();
const ttsDtypeOverride = resolveTinyModelDtypeOverride();
/** Device values `kokoro-js` accepts; the tiny device order is mapped onto these. */
type KokoroDevice = "cpu" | "wasm" | "webgpu";
/** A loaded Kokoro voice synthesizer (subset of `kokoro-js`'s `KokoroTTS`). */
interface KokoroTtsInstance {
generate(text: string, options: { voice: string }): Promise<RawAudio>;
stream(
text: string | TextSplitterStreamInstance,
options: { voice: string },
): AsyncGenerator<{ text: string; phonemes: string; audio: RawAudio }, void, void>;
}
/**
* Incremental text source for {@link KokoroTtsInstance.stream} (subset of
* `kokoro-js`'s `TextSplitterStream`). Text pushed at any time is split into
* complete sentences; `close` flushes the trailing buffer and ends the stream.
*/
interface TextSplitterStreamInstance {
push(...texts: string[]): void;
close(): void;
}
/** `KokoroTTS` static surface used to load a model from the Hugging Face Hub. */
interface KokoroRuntime {
KokoroTTS: {
from_pretrained(
repo: string,
options: {
dtype: TinyModelDtype;
device: KokoroDevice;
progress_callback: (info: ProgressInfo) => void;
},
): Promise<KokoroTtsInstance>;
};
TextSplitterStream: new () => TextSplitterStreamInstance;
}
/**
* The `@huggingface/transformers` instance `kokoro-js` runs on. We only touch its
* `env` (cache dir + log level) and `LogLevel`; inference goes through Kokoro.
*/
interface TransformersEnv {
env: {
cacheDir?: string;
allowLocalModels?: boolean;
logLevel?: unknown;
};
LogLevel: {
ERROR: unknown;
};
}
const models = new Map<TtsLocalModelKey, Promise<KokoroTtsInstance>>();
let synthesizeQueue = Promise.resolve();
let kokoroRuntime: Promise<KokoroRuntime> | null = null;
/**
* In-flight streaming sessions keyed by request id. A session is created on
* `stream-start` and torn down when its generator finishes. Text pushed before
* the model finishes loading is held in `buffered` and flushed into the splitter
* once it exists; pushes after that go straight to the live splitter.
*/
interface StreamSession {
modelKey: TtsLocalModelKey;
voice: string | undefined;
buffered: string[];
splitter: TextSplitterStreamInstance | null;
ended: boolean;
cancelled: boolean;
}
const streamSessions = new Map<string, StreamSession>();
function errorText(error: unknown): string {
return error instanceof Error ? (error.stack ?? error.message) : String(error);
}
function errorMessage(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function sendLog(
transport: TtsTransport,
level: "debug" | "warn" | "error",
msg: string,
meta?: Record<string, unknown>,
): void {
transport.send({ type: "log", level, msg, meta });
}
function sendRuntimeInstallProgress(
transport: TtsTransport,
requestId: string,
modelKey: TtsLocalModelKey,
status: "initiate" | "download" | "done",
): void {
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status, name: `${KOKORO_PACKAGE}@${KOKORO_VERSION}` },
});
}
/**
* Map a tiny-model device onto the narrow set `kokoro-js` accepts. The worker
* always runs `kokoro-js` on Node, where `cpu` (onnxruntime-node) is the only
* safe option; `webgpu`/`wasm` are honored if explicitly requested.
*/
function toKokoroDevice(device: TinyModelDevice): KokoroDevice {
if (device === "wasm") return "wasm";
if (device === "webgpu" || device === "gpu") return "webgpu";
return "cpu";
}
function configureTransformers(transformers: TransformersEnv): void {
transformers.env.cacheDir = getTinyModelsCacheDir();
transformers.env.allowLocalModels = false;
transformers.env.logLevel = transformers.LogLevel.ERROR;
}
/**
* Lazily `bun install` `kokoro-js` into a side runtime dir (idempotent, version-
* keyed) and return its module, with the `@huggingface/transformers` instance it
* loads configured (cache dir + quiet logging). `kokoro-js` is NEVER a dependency
* of the main tree: its transformers@3.8.1 graph pulls onnxruntime-node@1.21,
* which segfaults Bun on session creation, so the runtime manifest force-pins
* onnxruntime-node to the Bun-safe version via `overrides`. `sharp` is stubbed —
* the TTS pipeline is audio-only, so the native image codec transformers eagerly
* requires is dead weight. Memoized so the runtime loads once per process.
*/
async function loadKokoroRuntime(
transport: TtsTransport,
requestId: string,
modelKey: TtsLocalModelKey,
): Promise<KokoroRuntime> {
if (kokoroRuntime) return kokoroRuntime;
kokoroRuntime = (async () => {
const runtimeDir = await ensureRuntimeInstalled({
runtimeDir: getTtsRuntimeDir(),
install: {
dependencies: { [KOKORO_PACKAGE]: KOKORO_VERSION },
overrides: { [ONNXRUNTIME_NODE_PACKAGE]: ONNXRUNTIME_NODE_VERSION },
trustedDependencies: [ONNXRUNTIME_NODE_PACKAGE],
},
probePackage: KOKORO_PACKAGE,
onPhase: phase => sendRuntimeInstallProgress(transport, requestId, modelKey, phase),
});
const nodeModules = path.join(runtimeDir, "node_modules");
const sharpStub = path.join(runtimeDir, "omp-sharp-stub.cjs");
await Bun.write(sharpStub, "module.exports = {};\n");
installRuntimeModuleResolver({ runtimeNodeModules: nodeModules, stubs: { sharp: sharpStub } });
const kokoroEntry = resolveRuntimeModule(nodeModules, KOKORO_PACKAGE);
if (!kokoroEntry) throw new Error(`Unable to resolve ${KOKORO_PACKAGE} in runtime at ${nodeModules}`);
const entryRequire = createRequire(kokoroEntry);
configureTransformers(entryRequire(TRANSFORMERS_PACKAGE) as TransformersEnv);
return entryRequire(kokoroEntry) as KokoroRuntime;
})().catch(error => {
kokoroRuntime = null;
throw error;
});
return kokoroRuntime;
}
function toProgressEvent(modelKey: TtsLocalModelKey, info: ProgressInfo): TtsProgressEvent {
if (info.status === "ready") {
return { modelKey, status: info.status, task: info.task, model: info.model };
}
if (info.status === "progress_total") {
return {
modelKey,
status: info.status,
name: info.name,
progress: info.progress,
loaded: info.loaded,
total: info.total,
files: info.files,
};
}
if (info.status === "progress") {
return {
modelKey,
status: info.status,
name: info.name,
file: info.file,
progress: info.progress,
loaded: info.loaded,
total: info.total,
};
}
return { modelKey, status: info.status, name: info.name, file: info.file };
}
function sendProgress(transport: TtsTransport, id: string, modelKey: TtsLocalModelKey, info: ProgressInfo): void {
transport.send({ type: "progress", id, event: toProgressEvent(modelKey, info) });
}
async function loadModelOnDevice(
runtime: KokoroRuntime,
spec: TtsLocalModelSpec,
modelKey: TtsLocalModelKey,
transport: TtsTransport,
requestId: string,
device: KokoroDevice,
): Promise<KokoroTtsInstance> {
return runtime.KokoroTTS.from_pretrained(spec.repo, {
device,
dtype: ttsDtypeOverride ?? spec.dtype,
progress_callback: info => sendProgress(transport, requestId, modelKey, info),
});
}
async function loadModelWithDeviceFallback(
runtime: KokoroRuntime,
spec: TtsLocalModelSpec,
modelKey: TtsLocalModelKey,
transport: TtsTransport,
requestId: string,
): Promise<{ model: KokoroTtsInstance; device: KokoroDevice }> {
const order = tinyModelDeviceLoadOrder(ttsDevicePreference);
if (order[0] !== ttsDevicePreference.device) {
sendLog(transport, "warn", "tts: requested device is unsafe in the worker; using CPU", {
modelKey,
repo: spec.repo,
requestedDevice: ttsDevicePreference.device,
device: order[0],
});
}
const devices: KokoroDevice[] = [];
for (const device of order) {
const mapped = toKokoroDevice(device);
if (!devices.includes(mapped)) devices.push(mapped);
}
for (let i = 0; i < devices.length; i += 1) {
const device = devices[i]!;
try {
return { model: await loadModelOnDevice(runtime, spec, modelKey, transport, requestId, device), device };
} catch (error) {
if (i === devices.length - 1) throw error;
const fallbackDevice = devices[i + 1]!;
sendLog(transport, "warn", "tts: accelerated device failed; falling back", {
modelKey,
repo: spec.repo,
device,
fallbackDevice,
error: errorMessage(error),
});
}
}
throw new Error("No TTS devices configured");
}
async function loadModel(
modelKey: TtsLocalModelKey,
transport: TtsTransport,
requestId: string,
): Promise<KokoroTtsInstance> {
const spec = getTtsLocalModelSpec(modelKey);
if (!spec) throw new Error(`Unknown local TTS model: ${modelKey}`);
const cached = models.get(modelKey);
if (cached) {
void cached
.then(() => {
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status: "ready", task: TTS_TASK, model: spec.repo },
});
})
.catch(() => undefined);
return cached;
}
const runtime = await loadKokoroRuntime(transport, requestId, modelKey);
const startedAt = performance.now();
const loaded = loadModelWithDeviceFallback(runtime, spec, modelKey, transport, requestId).then(
({ model, device }) => {
sendLog(transport, "debug", "tts: local model loaded", {
modelKey,
repo: spec.repo,
device,
requestedDevice: ttsDevicePreference.device,
dtype: ttsDtypeOverride ?? spec.dtype,
elapsedMs: Math.round(performance.now() - startedAt),
});
transport.send({
type: "progress",
id: requestId,
event: { modelKey, status: "ready", task: TTS_TASK, model: spec.repo },
});
return model;
},
error => {
models.delete(modelKey);
throw error;
},
);
models.set(modelKey, loaded);
return loaded;
}
async function synthesize(
transport: TtsTransport,
requestId: string,
modelKey: TtsLocalModelKey,
text: string,
voice: string | undefined,
): Promise<{ pcm: Float32Array; sampleRate: number }> {
const synthesizer = await loadModel(modelKey, transport, requestId);
const output = await synthesizer.generate(text, { voice: resolveTtsVoice(modelKey, voice) });
const spec = getTtsLocalModelSpec(modelKey);
const audio = Array.isArray(output.audio) ? output.audio[0] : output.audio;
if (!audio) throw new Error("Kokoro synthesis returned no audio samples");
return { pcm: audio, sampleRate: output.sampling_rate || spec?.sampleRate || 24_000 };
}
function enqueueRequest(
transport: TtsTransport,
request: Extract<TtsWorkerInbound, { type: "synthesize" | "download" }>,
): void {
synthesizeQueue = synthesizeQueue.then(
async () => {
await handleQueuedRequest(transport, request);
},
async () => {
await handleQueuedRequest(transport, request);
},
);
}
async function handleQueuedRequest(
transport: TtsTransport,
request: Extract<TtsWorkerInbound, { type: "synthesize" | "download" }>,
): Promise<void> {
try {
if (request.type === "download") {
await loadModel(request.modelKey, transport, request.id);
transport.send({ type: "downloaded", id: request.id });
return;
}
const { pcm, sampleRate } = await synthesize(
transport,
request.id,
request.modelKey,
request.text,
request.voice,
);
transport.send({ type: "audio", id: request.id, pcm, sampleRate });
} catch (error) {
transport.send({ type: "error", id: request.id, error: errorText(error) });
}
}
/**
* Drive one streaming session to completion: load the model, create the
* splitter, flush any text pushed before the model was ready, then emit one
* `audio-chunk` per synthesized sentence followed by a single `stream-done`.
* Serialized through {@link synthesizeQueue} so it never interleaves model
* access with a batch synthesize/download.
*/
async function runStreamSession(transport: TtsTransport, id: string, session: StreamSession): Promise<void> {
try {
if (session.cancelled) return;
const runtime = await loadKokoroRuntime(transport, id, session.modelKey);
if (session.cancelled) return;
const synthesizer = await loadModel(session.modelKey, transport, id);
if (session.cancelled) return;
const spec = getTtsLocalModelSpec(session.modelKey);
const splitter = new runtime.TextSplitterStream();
// Flush buffered text before exposing the splitter so a push racing this
// block can't slip ahead of the already-queued fragments.
for (const text of session.buffered) {
if (session.cancelled) return;
splitter.push(text);
}
session.buffered = [];
session.splitter = splitter;
if (session.ended || session.cancelled) splitter.close();
const voice = resolveTtsVoice(session.modelKey, session.voice);
let index = 0;
for await (const chunk of synthesizer.stream(splitter, { voice })) {
if (session.cancelled) break;
const audio = Array.isArray(chunk.audio.audio) ? chunk.audio.audio[0] : chunk.audio.audio;
if (!audio) continue;
transport.send({
type: "audio-chunk",
id,
index: index++,
text: chunk.text,
pcm: audio,
sampleRate: chunk.audio.sampling_rate || spec?.sampleRate || 24_000,
});
}
if (!session.cancelled) transport.send({ type: "stream-done", id });
} catch (error) {
if (!session.cancelled) transport.send({ type: "error", id, error: errorText(error) });
} finally {
streamSessions.delete(id);
}
}
function startStreamSession(
transport: TtsTransport,
message: Extract<TtsWorkerInbound, { type: "stream-start" }>,
): void {
const session: StreamSession = {
modelKey: message.modelKey,
voice: message.voice,
buffered: [],
splitter: null,
ended: false,
cancelled: false,
};
streamSessions.set(message.id, session);
synthesizeQueue = synthesizeQueue.then(
() => runStreamSession(transport, message.id, session),
() => runStreamSession(transport, message.id, session),
);
}
function pushToStreamSession(id: string, text: string): void {
const session = streamSessions.get(id);
if (!session || session.cancelled) return;
if (session.splitter) session.splitter.push(text);
else session.buffered.push(text);
}
function endStreamSession(id: string): void {
const session = streamSessions.get(id);
if (!session || session.cancelled) return;
session.ended = true;
session.splitter?.close();
}
function cancelStreamSession(id: string): void {
const session = streamSessions.get(id);
if (!session) return;
session.cancelled = true;
session.buffered = [];
session.splitter?.close();
streamSessions.delete(id);
}
export function startTtsWorker(transport: TtsTransport): void {
transport.onMessage(message => {
switch (message.type) {
case "ping":
transport.send({ type: "pong", id: message.id });
return;
case "stream-start":
startStreamSession(transport, message);
return;
case "stream-push":
pushToStreamSession(message.id, message.text);
return;
case "stream-end":
endStreamSession(message.id);
return;
case "stream-cancel":
cancelStreamSession(message.id);
return;
default:
enqueueRequest(transport, message);
return;
}
});
}
+360
View File
@@ -0,0 +1,360 @@
import { describe, expect, it } from "bun:test";
import { resolveLocalWavPath, resolveTtsBackend } from "../tools/tts";
import {
DEFAULT_TTS_LOCAL_MODEL_KEY,
DEFAULT_TTS_VOICE,
getTtsLocalModelSpec,
isTtsLocalModelKey,
KOKORO_VOICES,
resolveTtsRepo,
resolveTtsVoice,
TTS_LOCAL_VOICE_OPTIONS,
} from "./models";
import { type TtsAudioChunk, TtsClient } from "./tts-client";
import type { TtsTransport, TtsWorkerInbound, TtsWorkerOutbound } from "./tts-protocol";
import { startTtsWorker } from "./tts-worker";
import { encodeWav } from "./wav";
// ── Backend resolution (auto/local/xai × creds × mp3) ───────────────
describe("resolveTtsBackend", () => {
it("honors an explicit xai preference regardless of creds/codec", () => {
expect(resolveTtsBackend({ preference: "xai", wantsMp3: false, hasXaiCreds: false })).toBe("xai");
expect(resolveTtsBackend({ preference: "xai", wantsMp3: true, hasXaiCreds: true })).toBe("xai");
});
it("honors an explicit local preference regardless of creds/codec", () => {
expect(resolveTtsBackend({ preference: "local", wantsMp3: true, hasXaiCreds: true })).toBe("local");
expect(resolveTtsBackend({ preference: "local", wantsMp3: false, hasXaiCreds: false })).toBe("local");
});
it("auto prefers local for WAV output", () => {
expect(resolveTtsBackend({ preference: "auto", wantsMp3: false, hasXaiCreds: true })).toBe("local");
expect(resolveTtsBackend({ preference: "auto", wantsMp3: false, hasXaiCreds: false })).toBe("local");
});
it("auto routes mp3 to xai only when credentials exist", () => {
expect(resolveTtsBackend({ preference: "auto", wantsMp3: true, hasXaiCreds: true })).toBe("xai");
expect(resolveTtsBackend({ preference: "auto", wantsMp3: true, hasXaiCreds: false })).toBe("local");
});
it("treats an unknown preference like auto", () => {
expect(resolveTtsBackend({ preference: "garbage", wantsMp3: false, hasXaiCreds: true })).toBe("local");
expect(resolveTtsBackend({ preference: "garbage", wantsMp3: true, hasXaiCreds: true })).toBe("xai");
});
});
// ── Local output path / codec substitution ──────────────────────────
describe("resolveLocalWavPath", () => {
it("keeps a .wav path unchanged", () => {
expect(resolveLocalWavPath("voice.wav")).toEqual({ wavPath: "voice.wav", substituted: false });
expect(resolveLocalWavPath("dir/voice.WAV")).toEqual({ wavPath: "dir/voice.WAV", substituted: false });
});
it("rewrites .mp3 to a sibling .wav and flags substitution", () => {
expect(resolveLocalWavPath("voice.mp3")).toEqual({ wavPath: "voice.wav", substituted: true });
expect(resolveLocalWavPath("a/b/clip.mp3")).toEqual({ wavPath: "a/b/clip.wav", substituted: true });
});
it("appends .wav when there is no extension", () => {
expect(resolveLocalWavPath("voice")).toEqual({ wavPath: "voice.wav", substituted: true });
expect(resolveLocalWavPath("a.b/clip")).toEqual({ wavPath: "a.b/clip.wav", substituted: true });
});
});
// ── WAV container assembly ──────────────────────────────────────────
function readAscii(view: DataView, offset: number, length: number): string {
let out = "";
for (let i = 0; i < length; i += 1) out += String.fromCharCode(view.getUint8(offset + i));
return out;
}
describe("encodeWav", () => {
it("writes a canonical 16-bit mono PCM RIFF/WAVE header", () => {
const samples = new Float32Array([0, 0.5, -0.5, 1]);
const wav = encodeWav(samples, 16_000);
const view = new DataView(wav.buffer, wav.byteOffset, wav.byteLength);
// 44-byte header + 2 bytes per sample
expect(wav.length).toBe(44 + samples.length * 2);
expect(readAscii(view, 0, 4)).toBe("RIFF");
expect(view.getUint32(4, true)).toBe(36 + samples.length * 2);
expect(readAscii(view, 8, 4)).toBe("WAVE");
expect(readAscii(view, 12, 4)).toBe("fmt ");
expect(view.getUint32(16, true)).toBe(16); // PCM fmt chunk size
expect(view.getUint16(20, true)).toBe(1); // PCM format
expect(view.getUint16(22, true)).toBe(1); // mono
expect(view.getUint32(24, true)).toBe(16_000); // sample rate
expect(view.getUint32(28, true)).toBe(16_000 * 2); // byte rate (mono, 16-bit)
expect(view.getUint16(32, true)).toBe(2); // block align
expect(view.getUint16(34, true)).toBe(16); // bits per sample
expect(readAscii(view, 36, 4)).toBe("data");
expect(view.getUint32(40, true)).toBe(samples.length * 2);
});
it("quantizes and clamps float samples to signed 16-bit", () => {
const samples = new Float32Array([0, 1, -1, 0.5, 2, -2]);
const wav = encodeWav(samples, 24_000);
const view = new DataView(wav.buffer, wav.byteOffset, wav.byteLength);
const at = (i: number): number => view.getInt16(44 + i * 2, true);
expect(at(0)).toBe(0);
expect(at(1)).toBe(32_767); // +1 → max
expect(at(2)).toBe(-32_768); // -1 → min
expect(at(3)).toBe(Math.round(0.5 * 32_767)); // 16384
expect(at(4)).toBe(32_767); // +2 clamped
expect(at(5)).toBe(-32_768); // -2 clamped
expect(view.getUint32(24, true)).toBe(24_000);
});
});
// ── Model-key → repo / voice mapping ────────────────────────────────
describe("tts model registry", () => {
const KOKORO_REPO = "onnx-community/Kokoro-82M-v1.0-ONNX";
it("maps the default model key to the Kokoro ONNX repo at 24 kHz", () => {
expect(DEFAULT_TTS_LOCAL_MODEL_KEY).toBe("kokoro");
expect(resolveTtsRepo(DEFAULT_TTS_LOCAL_MODEL_KEY)).toBe(KOKORO_REPO);
expect(getTtsLocalModelSpec("kokoro")?.repo).toBe(KOKORO_REPO);
expect(getTtsLocalModelSpec("kokoro")?.sampleRate).toBe(24_000);
});
it("falls back to the default repo for unknown keys", () => {
expect(resolveTtsRepo("does-not-exist")).toBe(KOKORO_REPO);
expect(resolveTtsRepo(undefined)).toBe(KOKORO_REPO);
expect(getTtsLocalModelSpec("does-not-exist")).toBeUndefined();
});
it("narrows known keys with the type guard", () => {
expect(isTtsLocalModelKey("kokoro")).toBe(true);
expect(isTtsLocalModelKey("mms-en")).toBe(false);
});
it("defaults to the flagship voice and exposes it as the first catalog entry", () => {
expect(DEFAULT_TTS_VOICE).toBe("af_heart");
expect(KOKORO_VOICES[0]?.id).toBe("af_heart");
expect(TTS_LOCAL_VOICE_OPTIONS.map(o => o.value)).toEqual(KOKORO_VOICES.map(v => v.id));
});
it("resolves known voices verbatim and falls back to the default voice otherwise", () => {
expect(resolveTtsVoice("kokoro", "af_bella")).toBe("af_bella");
expect(resolveTtsVoice("kokoro", "bf_emma")).toBe("bf_emma");
expect(resolveTtsVoice("kokoro", "default")).toBe("af_heart");
expect(resolveTtsVoice("kokoro", undefined)).toBe("af_heart");
expect(resolveTtsVoice("kokoro", "not-a-voice")).toBe("af_heart");
expect(resolveTtsVoice(undefined, "am_michael")).toBe("am_michael");
});
});
// ── Protocol round-trip: in-memory transport + injected worker ──────
function transportPair(): {
workerTransport: TtsTransport;
sendToWorker: (message: TtsWorkerInbound) => void;
onParent: (handler: (message: TtsWorkerOutbound) => void) => void;
} {
const workerListeners = new Set<(message: TtsWorkerInbound) => void>();
const parentListeners = new Set<(message: TtsWorkerOutbound) => void>();
return {
workerTransport: {
send: message => {
for (const listener of parentListeners) listener(message);
},
onMessage: handler => {
workerListeners.add(handler);
return () => workerListeners.delete(handler);
},
},
sendToWorker: message => {
for (const listener of workerListeners) listener(message);
},
onParent: handler => {
parentListeners.add(handler);
},
};
}
interface FakeWorkerHandle {
send(message: TtsWorkerInbound): void;
onMessage(handler: (message: TtsWorkerOutbound) => void): () => void;
onError(handler: (error: Error) => void): () => void;
ref(): void;
unref(): void;
terminate(): Promise<void>;
}
function fakeWorker(
respond: (message: TtsWorkerInbound, emit: (out: TtsWorkerOutbound) => void) => void,
): FakeWorkerHandle {
const listeners = new Set<(message: TtsWorkerOutbound) => void>();
return {
send(message) {
queueMicrotask(() =>
respond(message, out => {
for (const listener of listeners) listener(out);
}),
);
},
onMessage(handler) {
listeners.add(handler);
return () => listeners.delete(handler);
},
onError() {
return () => {};
},
ref() {},
unref() {},
async terminate() {
listeners.clear();
},
};
}
describe("tts protocol round-trip", () => {
it("worker pongs to a ping over the transport (no model load)", () => {
const { workerTransport, sendToWorker, onParent } = transportPair();
startTtsWorker(workerTransport);
const received: TtsWorkerOutbound[] = [];
onParent(message => received.push(message));
sendToWorker({ type: "ping", id: "p1" });
expect(received).toEqual([{ type: "pong", id: "p1" }]);
});
it("client resolves a synthesize request to the worker's audio response", async () => {
const pcm = new Float32Array([0, 0.25, -0.25]);
const client = new TtsClient(() =>
fakeWorker((message, emit) => {
if (message.type === "synthesize") emit({ type: "audio", id: message.id, pcm, sampleRate: 24_000 });
}),
);
const audio = await client.synthesize("kokoro", "hello world", { voice: "af_heart" });
expect(audio).not.toBeNull();
expect(audio?.sampleRate).toBe(24_000);
expect(Array.from(audio?.pcm ?? [])).toEqual([0, 0.25, -0.25]);
await client.terminate();
});
it("client rejects unknown model keys without spawning a worker", async () => {
let spawned = false;
const client = new TtsClient(() => {
spawned = true;
return fakeWorker(() => {});
});
expect(await client.synthesize("not-a-model", "hi")).toBeNull();
expect(spawned).toBe(false);
await client.terminate();
});
it("client maps a worker error to null", async () => {
const client = new TtsClient(() =>
fakeWorker((message, emit) => {
if (message.type === "synthesize") emit({ type: "error", id: message.id, error: "boom" });
}),
);
expect(await client.synthesize("kokoro", "hi")).toBeNull();
await client.terminate();
});
it("client resolves a download request to true", async () => {
const client = new TtsClient(() =>
fakeWorker((message, emit) => {
if (message.type === "download") emit({ type: "downloaded", id: message.id });
}),
);
expect(await client.downloadModel("kokoro")).toBe(true);
await client.terminate();
});
it("client streams audio chunks in emission order and ends on stream-done", async () => {
const pushed: string[] = [];
const client = new TtsClient(() =>
fakeWorker((message, emit) => {
if (message.type === "stream-push") {
pushed.push(message.text);
} else if (message.type === "stream-end") {
// Simulate the worker synthesizing one sentence per pushed fragment.
for (const [index, text] of pushed.entries()) {
emit({
type: "audio-chunk",
id: message.id,
index,
text,
pcm: new Float32Array([index]),
sampleRate: 24_000,
});
}
emit({ type: "stream-done", id: message.id });
}
}),
);
const handle = client.synthesizeStream("kokoro", { voice: "af_heart" });
handle.push("one. ");
handle.push("two. ");
handle.end();
const chunks: TtsAudioChunk[] = [];
for await (const chunk of handle.chunks) chunks.push(chunk);
expect(chunks.map(c => c.text)).toEqual(["one. ", "two. "]);
expect(chunks.map(c => c.index)).toEqual([0, 1]);
expect(Array.from(chunks[0]?.pcm ?? [])).toEqual([0]);
await client.terminate();
});
it("sends stream-cancel and closes the iterator when aborted", async () => {
const controller = new AbortController();
const sent: TtsWorkerInbound[] = [];
const client = new TtsClient(() =>
fakeWorker(message => {
sent.push(message);
}),
);
const handle = client.synthesizeStream("kokoro", { signal: controller.signal });
handle.push("one. ");
controller.abort();
handle.push("two. ");
handle.end();
const chunks: TtsAudioChunk[] = [];
for await (const chunk of handle.chunks) chunks.push(chunk);
expect(chunks).toEqual([]);
expect(sent.some(message => message.type === "stream-cancel")).toBe(true);
expect(sent.some(message => message.type === "stream-end")).toBe(false);
await client.terminate();
});
it("client surfaces a worker error on the stream iterator", async () => {
const client = new TtsClient(() =>
fakeWorker((message, emit) => {
if (message.type === "stream-end") emit({ type: "error", id: message.id, error: "boom" });
}),
);
const handle = client.synthesizeStream("kokoro");
handle.end();
let caught: unknown;
try {
for await (const _chunk of handle.chunks) {
// drain — the iterator must reject before yielding anything
}
} catch (error) {
caught = error;
}
expect(caught).toBeInstanceOf(Error);
expect((caught as Error).message).toBe("boom");
await client.terminate();
});
it("client returns an inert, empty stream for unknown models without spawning", async () => {
let spawned = false;
const client = new TtsClient(() => {
spawned = true;
return fakeWorker(() => {});
});
const handle = client.synthesizeStream("not-a-model");
handle.push("hi");
handle.end();
const chunks: TtsAudioChunk[] = [];
for await (const chunk of handle.chunks) chunks.push(chunk);
expect(chunks).toEqual([]);
expect(spawned).toBe(false);
await client.terminate();
});
});
@@ -0,0 +1,169 @@
import { afterEach, beforeEach, describe, expect, it, type Mock, spyOn } from "bun:test";
import { resetSettingsForTest, Settings, settings } from "../config/settings";
import { DUCK_GAIN } from "./streaming-player";
import { type TtsAudioChunk, type TtsStreamOptions, ttsClient } from "./tts-client";
import { Vocalizer, type VocalizerPlayer } from "./vocalizer";
// Exercises the streaming contract: the vocalizer feeds assistant deltas
// straight into the engine's incremental text input (`synthesizeStream`) instead
// of pre-chunking in JS, so sentence segmentation lives in the worker. We spy
// `ttsClient.synthesizeStream` to observe the text pushed in, the input close,
// and the model/voice the session opens with, and drive the returned chunk
// iterator to assert ordered playback — no real worker, no audio device.
class FakePlayer implements VocalizerPlayer {
starts: number[] = [];
writes: Float32Array[] = [];
gains: number[] = [];
ended = false;
stopped = false;
start(sampleRate: number): void {
this.starts.push(sampleRate);
}
write(pcm: Float32Array): void {
this.writes.push(pcm);
}
setGain(gain: number): void {
this.gains.push(gain);
}
async end(): Promise<void> {
this.ended = true;
}
stop(): void {
this.stopped = true;
}
}
describe("vocalizer streaming", () => {
let stream: Mock<typeof ttsClient.synthesizeStream>;
let audio: TtsAudioChunk[];
let vocalizer: Vocalizer;
const probe: { pushed: string[]; ended: boolean; modelKey?: string; options?: TtsStreamOptions } = {
pushed: [],
ended: false,
};
const players: FakePlayer[] = [];
beforeEach(async () => {
resetSettingsForTest();
await Settings.init({ inMemory: true });
settings.override("speech.enabled", true);
settings.override("speech.voice", "af_heart");
settings.override("tts.localModel", "kokoro");
audio = [];
players.length = 0;
probe.pushed = [];
probe.ended = false;
probe.modelKey = undefined;
probe.options = undefined;
vocalizer = new Vocalizer(() => {
const player = new FakePlayer();
players.push(player);
return player;
});
stream = spyOn(ttsClient, "synthesizeStream").mockImplementation((modelKey, options) => {
probe.modelKey = modelKey;
probe.options = options;
async function* chunks(): AsyncIterableIterator<TtsAudioChunk> {
for (const chunk of audio) yield chunk;
}
return {
push: (text: string) => {
probe.pushed.push(text);
},
end: () => {
probe.ended = true;
},
chunks: chunks(),
};
});
vocalizer.clear();
});
afterEach(() => {
stream.mockRestore();
vocalizer.clear();
resetSettingsForTest();
});
it("streams deltas straight into the engine and closes the input on flush", async () => {
vocalizer.pushDelta("Hello wor");
vocalizer.pushDelta("ld. How are you?");
vocalizer.flush();
await vocalizer.idle();
expect(stream).toHaveBeenCalledTimes(1);
expect(probe.pushed).toEqual(["Hello wor", "ld. How are you?"]);
expect(probe.ended).toBe(true);
});
it("opens the stream with the configured local model and speech voice", async () => {
settings.override("speech.voice", "am_michael");
vocalizer.speak("Ready now.");
await vocalizer.idle();
expect(probe.modelKey).toBe("kokoro");
expect(probe.options?.voice).toBe("am_michael");
expect(probe.pushed).toEqual(["Ready now."]);
expect(probe.ended).toBe(true);
});
it("plays each synthesized sentence chunk in arrival order", async () => {
audio = [
{ index: 0, text: "one. ", pcm: new Float32Array([0]), sampleRate: 24_000 },
{ index: 1, text: "two.", pcm: new Float32Array([1]), sampleRate: 24_000 },
];
vocalizer.speak("one. two.");
await vocalizer.idle();
expect(players).toHaveLength(1);
expect(players[0]?.starts).toEqual([24_000]);
expect(players[0]?.writes.map(chunk => Array.from(chunk))).toEqual([[0], [1]]);
expect(players[0]?.ended).toBe(true);
});
it("stops playback when cleared", async () => {
audio = [{ index: 0, text: "one.", pcm: new Float32Array([0]), sampleRate: 24_000 }];
vocalizer.speak("one.");
await vocalizer.idle();
vocalizer.clear();
expect(players[0]?.stopped).toBe(true);
});
it("ducks the active player while the user speaks and restores afterward", async () => {
vocalizer.pushDelta("Talking over me.");
expect(players[0]?.gains).toEqual([1]);
vocalizer.duck();
vocalizer.unduck();
expect(players[0]?.gains).toEqual([1, DUCK_GAIN, 1]);
vocalizer.flush();
await vocalizer.idle();
});
it("opens a new session ducked while the user is still speaking", async () => {
vocalizer.duck();
vocalizer.pushDelta("Should start quiet.");
expect(players[0]?.gains).toEqual([DUCK_GAIN]);
vocalizer.flush();
await vocalizer.idle();
});
it("does not open a session when speech is disabled", async () => {
settings.override("speech.enabled", false);
vocalizer.pushDelta("This is a full sentence. ");
vocalizer.flush();
vocalizer.speak("Another sentence.");
await vocalizer.idle();
expect(stream).not.toHaveBeenCalled();
});
it("aborts the in-flight session on clear", async () => {
vocalizer.pushDelta("Interrupt me mid-turn.");
const signal = probe.options?.signal;
expect(signal?.aborted).toBe(false);
vocalizer.clear();
expect(signal?.aborted).toBe(true);
await vocalizer.idle();
});
});
+162
View File
@@ -0,0 +1,162 @@
/**
* Streaming assistant speech-vocalization.
*
* The vocalizer turns the assistant's STREAMING output into spoken audio as a
* side effect of the normal turn. Text deltas are streamed *straight into the
* TTS engine* ({@link Vocalizer.pushDelta} → the worker's incremental text
* input): the engine splits the running text at sentence boundaries and emits
* one audio chunk per sentence, which a single {@link StreamingAudioPlayer}
* plays back gaplessly. So the assistant starts speaking sentence 1 while later
* sentences are still being generated — low latency, never overlapping.
*
* Overspeech control:
* - {@link clear} stops playback instantly (kills the player) and aborts
* in-flight synthesis — wired to a new turn, an Esc/Ctrl+C interrupt, and a
* sent message.
* - {@link duck}/{@link unduck} lower/restore the volume while the user is
* speaking (push-to-talk), so the assistant doesn't talk over them.
* - Sessions are chained, so sequential utterances queue and drain in order
* rather than overlapping.
*
* Errors are swallowed (debug-logged) so a synthesis or playback failure never
* throws into the turn. A process-level singleton ({@link vocalizer}) is shared
* by the event controller (streaming deltas) and the ask tool (spoken questions).
*/
import { logger } from "@oh-my-pi/pi-utils";
import { settings } from "../config/settings";
import { DEFAULT_TTS_VOICE } from "./models";
import { createStreamingPlayer, DUCK_GAIN } from "./streaming-player";
import { type TtsStreamHandle, ttsClient } from "./tts-client";
export interface VocalizerPlayer {
start(sampleRate: number): void;
write(pcm: Float32Array): void;
setGain(gain: number): void;
end(): Promise<void>;
stop(): void;
}
export class Vocalizer {
/** Open stream session for the current utterance; null when none is active. */
#handle: TtsStreamHandle | null = null;
/** Aborts the in-flight session on {@link clear}; replaced per session. */
#abort: AbortController | null = null;
/** The current session's player; stopped on {@link clear}, gain-tracked for ducking. */
#player: VocalizerPlayer | null = null;
/** Serialized playback chain across sessions; awaited by {@link idle}. */
#chain: Promise<void> = Promise.resolve();
/** Whether the user is currently speaking; new sessions open ducked. */
#ducked = false;
#createPlayer: () => VocalizerPlayer;
constructor(createPlayer: () => VocalizerPlayer = createStreamingPlayer) {
this.#createPlayer = createPlayer;
}
/**
* Stream a delta of assistant text into the engine. No-op when vocalization
* is disabled. The engine buffers the running text and emits audio for each
* complete sentence; the trailing partial is flushed by {@link flush}.
*/
pushDelta(text: string): void {
if (!settings.get("speech.enabled")) return;
if (!text) return;
this.#ensureSession().push(text);
}
/**
* Close the current input stream (call at message/turn end). The engine
* flushes its trailing partial as a final chunk; the player keeps draining
* queued audio until it completes.
*/
flush(): void {
this.#handle?.end();
this.#handle = null;
}
/**
* Speak a complete piece of text in one shot (ask questions, yield-mode final
* message): stream it in and immediately close the input. No-op when disabled.
*/
speak(text: string): void {
if (!settings.get("speech.enabled")) return;
if (!text) return;
this.#ensureSession().push(text);
this.flush();
}
/**
* Interrupt and drop the current session, killing in-flight playback and
* synthesis (new turn / user message / Esc interrupt). Audio stops at once.
*/
clear(): void {
this.#handle = null;
this.#abort?.abort();
this.#abort = null;
this.#player?.stop();
this.#player = null;
}
/** Lower the volume while the user is speaking (push-to-talk), so speech doesn't drown them out. */
duck(): void {
this.#ducked = true;
this.#player?.setGain(DUCK_GAIN);
}
/** Restore full volume once the user stops speaking. */
unduck(): void {
this.#ducked = false;
this.#player?.setGain(1);
}
/** Resolve once the playback chain has drained (tests / shutdown). */
idle(): Promise<void> {
return this.#chain;
}
/**
* Open a streaming-synthesis session lazily on the first delta and chain its
* playback after any prior session's, so sequential utterances never overlap.
*/
#ensureSession(): TtsStreamHandle {
if (this.#handle) return this.#handle;
const modelKey = settings.get("tts.localModel");
const voice = settings.get("speech.voice") || DEFAULT_TTS_VOICE;
const abort = new AbortController();
this.#abort = abort;
const handle = ttsClient.synthesizeStream(modelKey, { voice, signal: abort.signal });
this.#handle = handle;
const player = this.#createPlayer();
player.setGain(this.#ducked ? DUCK_GAIN : 1);
this.#player = player;
this.#chain = this.#chain.then(() => this.#play(handle, player, abort.signal));
return handle;
}
/** Feed each synthesized sentence into the player in arrival order; abort stops it. */
async #play(handle: TtsStreamHandle, player: VocalizerPlayer, signal: AbortSignal): Promise<void> {
let started = false;
try {
for await (const chunk of handle.chunks) {
if (signal.aborted) break;
if (!started) {
player.start(chunk.sampleRate);
started = true;
}
player.write(chunk.pcm);
}
if (started && !signal.aborted) {
await player.end();
return;
}
} catch (error) {
logger.debug("vocalizer: stream failed", {
error: error instanceof Error ? error.message : String(error),
});
}
player.stop();
}
}
/** Process-level vocalizer shared by the event controller and the ask tool. */
export const vocalizer = new Vocalizer();
+58
View File
@@ -0,0 +1,58 @@
const WAV_HEADER_BYTES = 44;
const PCM16_FORMAT = 1;
const BITS_PER_SAMPLE = 16;
const INT16_MAX = 32_767;
const INT16_MIN = -32_768;
/**
* Assemble a mono PCM16 WAV byte buffer from Float32 PCM samples (the shape
* transformers.js `RawAudio` emits: normalized [-1, 1] amplitudes plus a sample
* rate). No external encoder is involved — we write a canonical 44-byte RIFF/
* WAVE header followed by little-endian signed 16-bit samples. Samples are
* clamped before quantization so out-of-range float values do not wrap.
*/
export function encodeWav(samples: Float32Array, sampleRate: number): Uint8Array {
const channels = 1;
const byteRate = sampleRate * channels * (BITS_PER_SAMPLE / 8);
const blockAlign = channels * (BITS_PER_SAMPLE / 8);
const dataBytes = samples.length * (BITS_PER_SAMPLE / 8);
const buffer = new ArrayBuffer(WAV_HEADER_BYTES + dataBytes);
const view = new DataView(buffer);
// RIFF chunk descriptor
writeAscii(view, 0, "RIFF");
view.setUint32(4, WAV_HEADER_BYTES - 8 + dataBytes, true); // file size minus the first 8 bytes
writeAscii(view, 8, "WAVE");
// fmt sub-chunk
writeAscii(view, 12, "fmt ");
view.setUint32(16, 16, true); // PCM fmt chunk size
view.setUint16(20, PCM16_FORMAT, true);
view.setUint16(22, channels, true);
view.setUint32(24, sampleRate, true);
view.setUint32(28, byteRate, true);
view.setUint16(32, blockAlign, true);
view.setUint16(34, BITS_PER_SAMPLE, true);
// data sub-chunk
writeAscii(view, 36, "data");
view.setUint32(40, dataBytes, true);
let offset = WAV_HEADER_BYTES;
for (let i = 0; i < samples.length; i += 1) {
const sample = samples[i]!;
const clamped = sample > 1 ? 1 : sample < -1 ? -1 : sample;
const quantized =
clamped < 0
? Math.max(INT16_MIN, Math.round(clamped * -INT16_MIN))
: Math.min(INT16_MAX, Math.round(clamped * INT16_MAX));
view.setInt16(offset, quantized, true);
offset += 2;
}
return new Uint8Array(buffer);
}
function writeAscii(view: DataView, offset: number, text: string): void {
for (let i = 0; i < text.length; i += 1) view.setUint8(offset + i, text.charCodeAt(i));
}
@@ -0,0 +1,25 @@
import { describe, expect, it } from "bun:test";
import { ffmpegAssetName } from "./tools-manager";
describe("ffmpegAssetName", () => {
it("maps supported platform/arch pairs to direct-binary asset names", () => {
expect(ffmpegAssetName("b6.1.1", "darwin", "arm64")).toBe("ffmpeg-darwin-arm64");
expect(ffmpegAssetName("b6.1.1", "darwin", "x64")).toBe("ffmpeg-darwin-x64");
expect(ffmpegAssetName("b6.1.1", "linux", "arm64")).toBe("ffmpeg-linux-arm64");
expect(ffmpegAssetName("b6.1.1", "linux", "x64")).toBe("ffmpeg-linux-x64");
expect(ffmpegAssetName("b6.1.1", "win32", "x64")).toBe("ffmpeg-win32-x64");
});
it("returns null for win32 on arm64 (no static asset published)", () => {
expect(ffmpegAssetName("b6.1.1", "win32", "arm64")).toBeNull();
});
it("returns null for unsupported arch", () => {
expect(ffmpegAssetName("b6.1.1", "darwin", "ia32")).toBeNull();
expect(ffmpegAssetName("b6.1.1", "linux", "ppc64")).toBeNull();
});
it("returns null for unsupported platform", () => {
expect(ffmpegAssetName("b6.1.1", "freebsd", "x64")).toBeNull();
});
});
@@ -16,6 +16,16 @@ interface ToolConfig {
getAssetName: (version: string, plat: string, architecture: string) => string | null;
}
// ffmpeg static-binary asset names (eugeneware/ffmpeg-static direct binaries).
// Maps node arch (arm64|x64) only; everything else is unsupported.
export function ffmpegAssetName(_version: string, plat: string, architecture: string): string | null {
if (architecture !== "arm64" && architecture !== "x64") return null;
if (plat === "darwin") return `ffmpeg-darwin-${architecture}`;
if (plat === "linux") return `ffmpeg-linux-${architecture}`;
if (plat === "win32") return architecture === "x64" ? "ffmpeg-win32-x64" : null;
return null;
}
const TOOLS: Record<string, ToolConfig> = {
sd: {
name: "sd",
@@ -72,6 +82,14 @@ const TOOLS: Record<string, ToolConfig> = {
return null;
},
},
ffmpeg: {
name: "ffmpeg",
repo: "eugeneware/ffmpeg-static",
binaryName: "ffmpeg",
tagPrefix: "",
isDirectBinary: true,
getAssetName: ffmpegAssetName,
},
};
// CLI packages installed via uv/pip
@@ -89,7 +107,7 @@ const PYTHON_TOOLS: Record<string, PythonPackageToolConfig> = {
},
};
export type ToolName = "sd" | "sg" | "yt-dlp" | "trafilatura";
export type ToolName = "sd" | "sg" | "yt-dlp" | "trafilatura" | "ffmpeg";
// Get the path to a tool (system-wide or in our tools dir)
export function getToolPath(tool: ToolName): string | null {