@oh-my-pi/pi-coding-agent 16.1.13 → 16.1.15

package/src/stt/asr-worker.ts

@@ -10,14 +10,23 @@ import type {
 import {
 	ensureRuntimeInstalled,
 	getTinyModelsCacheDir,
-	installRuntimeModuleResolver,
 	isCompiledBinary,
 	resolveRuntimeModule,
 } from "@oh-my-pi/pi-utils";
 import packageJson from "../../package.json" with { type: "json" };
+import {
+	errorMessage,
+	errorText,
+	getTransformersVersionSpec,
+	loadTransformersRuntime,
+	MemoizedRuntime,
+	replayCachedReady,
+	sendLog,
+	sendProgress,
+} from "../subprocess/worker-runtime";
 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 { SttTransport, SttWorkerInbound } from "./asr-protocol";
 import { type EndpointerEvent, StreamEndpointer } from "./endpointer";
 import {
 	getSttModelSpec,
@@ -28,9 +37,7 @@ import {
 } 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;
@@ -136,34 +143,8 @@ function runOnModel<T>(work: () => Promise<T>): Promise<T> {
 	);
 	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;
-}
+const transformersRuntime = new MemoizedRuntime<TransformersRuntime>();
+const sherpaRuntime = new MemoizedRuntime<SherpaRuntime>();
 
 let cachedSherpaVersionSpec: string | undefined;
 function resolveSherpaVersionSpec(): string {
@@ -181,23 +162,6 @@ function getSherpaVersionSpec(): string {
 	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}`);
@@ -208,73 +172,6 @@ function getSherpaRuntimeDir(): string {
 	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
@@ -283,69 +180,25 @@ async function loadTransformers(
  * 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 () => {
+function loadSherpaRuntime(transport: SttTransport, requestId: string, modelKey: SttModelKey): Promise<SherpaRuntime> {
+	return sherpaRuntime.load(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()}`,
-				),
+				transport.send({
+					type: "progress",
+					id: requestId,
+					event: { modelKey, status: phase, name: `${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(
@@ -407,7 +260,13 @@ async function loadTransformersModel(
 	transport: SttTransport,
 	requestId: string,
 ): Promise<LoadedModel> {
-	const transformers = await loadTransformers(transport, requestId, modelKey);
+	const transformers = await loadTransformersRuntime(
+		transformersRuntime,
+		transport,
+		requestId,
+		modelKey,
+		getSttRuntimeDir,
+	);
 	const startedAt = performance.now();
 	const { pipeline, device } = await loadPipelineWithDeviceFallback(
 		transformers,
@@ -548,19 +407,8 @@ async function loadSherpaModel(
 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 cached = replayCachedReady(models, modelKey, transport, requestId, ASR_TASK, spec.repo);
+	if (cached) return cached;
 
 	const loading =
 		spec.engine === "sherpa"

package/src/subprocess/worker-client.ts

@@ -0,0 +1,297 @@
+import * as path from "node:path";
+import { $env, isBunTestRuntime, isCompiledBinary, logger, workerHostEntry } from "@oh-my-pi/pi-utils";
+import type { Subprocess } from "bun";
+
+/**
+ * Shared lifecycle scaffolding for the ONNX inference subprocess clients
+ * (mnemopi embeddings, speech-to-text, tiny-model titles/completions, TTS).
+ * Each runs `onnxruntime-node` inside a dedicated Bun child process so the NAPI
+ * constructor/finalizer never executes in the main agent address space — those
+ * destructors segfault Bun on shutdown (issues #1606 / #1607 / #3031).
+ *
+ * Only the genuinely identical pieces live here: the worker-handle shape, the
+ * spawn-command resolution, the parent-env snapshot, the `Bun.spawn` wiring,
+ * the inline "worker unavailable" stub, and the ping/pong smoke probe. Each
+ * client keeps its own divergent request/response correlation, streaming, and
+ * teardown semantics.
+ */
+
+/** Minimal inbound contract shared by every worker: a correlated `ping`. */
+export type WorkerInboundBase = { type: "ping"; id: string };
+
+/** Structured log line forwarded from a worker to the parent logger. */
+export type WorkerLogMessage = {
+	type: "log";
+	level: "debug" | "warn" | "error";
+	msg: string;
+	meta?: Record<string, unknown>;
+};
+
+/** Minimal outbound contract shared by every worker: `pong`, `error`, `log`. */
+export type WorkerOutboundBase =
+	| { type: "pong"; id: string }
+	| { type: "error"; id: string; error: string }
+	| WorkerLogMessage;
+
+/**
+ * Parent-side view of a worker subprocess: send typed inbound messages,
+ * subscribe to outbound messages and worker errors, and hard-terminate.
+ */
+export interface WorkerHandle<Inbound, Outbound> {
+	send(message: Inbound): void;
+	onMessage(handler: (message: Outbound) => void): () => void;
+	onError(handler: (error: Error) => void): () => void;
+	terminate(): Promise<void>;
+}
+
+/**
+ * A {@link WorkerHandle} that can also be (un)referenced so a pending request
+ * keeps the parent event loop alive while an idle worker never blocks exit.
+ */
+export interface RefCountedWorkerHandle<Inbound, Outbound> extends WorkerHandle<Inbound, Outbound> {
+	/** 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;
+}
+
+/** The raw spawned subprocess plus the parent-side fan-out sets. */
+export interface SpawnedSubprocess<Outbound> {
+	proc: Subprocess<"ignore", "ignore", "ignore">;
+	inbound: Set<(message: Outbound) => void>;
+	errors: Set<(error: Error) => void>;
+	/**
+	 * Flipped to `true` right before the deliberate SIGKILL so `onExit` can
+	 * distinguish the expected hard-kill from a crash (SIGSEGV from a native
+	 * fault, OOM SIGKILL, operator `kill -9`). Only the latter surfaces as a
+	 * worker error so callers don't await forever.
+	 */
+	intentionalExit: { value: boolean };
+}
+
+export interface WorkerSpawnCommand {
+	cmd: string[];
+	cwd?: string;
+}
+
+/**
+ * Cold-starting a 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 it.
+ */
+export const SMOKE_TEST_TIMEOUT_MS = 30_000;
+
+/**
+ * Resolve the command used to relaunch the agent CLI into 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 (bun test, SDK
+ * embedding).
+ */
+export function resolveWorkerSpawnCmd(workerArg: string): WorkerSpawnCommand {
+	if (isCompiledBinary()) return { cmd: [process.execPath, workerArg] };
+	const hostEntry = workerHostEntry();
+	if (hostEntry) {
+		return { cmd: [process.execPath, path.basename(hostEntry), workerArg], cwd: path.dirname(hostEntry) };
+	}
+	const packageRoot = path.resolve(import.meta.dir, "..", "..");
+	return { cmd: [process.execPath, "src/cli.ts", workerArg], cwd: packageRoot };
+}
+
+/**
+ * Snapshot the parent environment for the child. `process.env` carries
+ * `undefined` slots that `Bun.spawn` rejects, so filter them out; an optional
+ * `overlay` (e.g. the tiny-model device/dtype vars) wins over inherited keys.
+ */
+export function workerEnvFromParent(overlay?: Record<string, string>): Record<string, string> {
+	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;
+	}
+	if (overlay) {
+		for (const key in overlay) merged[key] = overlay[key];
+	}
+	return merged;
+}
+
+/**
+ * Spawn an inference worker subprocess and wire its IPC fan-out. The child
+ * inherits no stdio (native model runtimes may otherwise print progress or
+ * decoded text and corrupt the chat scrollback) and is `unref`'d outside `bun
+ * test` so an idle worker never blocks process exit. `exitLabel` prefixes the
+ * worker-error message surfaced for an unexpected (non-intentional) exit.
+ */
+export function createWorkerSubprocess<Outbound>(options: {
+	spawnCommand: WorkerSpawnCommand;
+	env: Record<string, string>;
+	exitLabel: string;
+}): SpawnedSubprocess<Outbound> {
+	const inbound = new Set<(message: Outbound) => void>();
+	const errors = new Set<(error: Error) => void>();
+	const intentionalExit = { value: false };
+	const proc = Bun.spawn({
+		cmd: options.spawnCommand.cmd,
+		cwd: options.spawnCommand.cwd,
+		env: options.env,
+		stdin: "ignore",
+		stdout: "ignore",
+		stderr: "ignore",
+		serialization: "advanced",
+		windowsHide: true,
+		ipc(message) {
+			for (const handler of inbound) handler(message as Outbound);
+		},
+		onExit(_proc, exitCode, signalCode) {
+			if (exitCode === 0) return;
+			// Swallow only the expected SIGKILL from `terminate()`; every other
+			// signal exit (SIGSEGV from a native fault, OOM SIGKILL, operator
+			// `kill -9`) 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(`${options.exitLabel} 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 };
+}
+
+/**
+ * Wrap a {@link SpawnedSubprocess} as a {@link WorkerHandle}. The `send`
+ * strategy is injected so each client keeps its exact IPC-send behaviour (e.g.
+ * `safeSend` vs an inline guarded `proc.send`). `terminate()` SIGKILLs: the
+ * point of subprocess isolation is that the parent never runs
+ * `onnxruntime-node`'s NAPI finalizer (it crashes Bun on Windows), so the OS
+ * reclaims the model memory instead. The intentional-exit flag is flipped
+ * *before* the kill so `onExit` can tell it apart from a native crash.
+ */
+export function createWorkerHandle<Inbound, Outbound>(
+	spawned: SpawnedSubprocess<Outbound>,
+	send: (message: Inbound) => void,
+): WorkerHandle<Inbound, Outbound> {
+	const { proc, inbound, errors, intentionalExit } = spawned;
+	return {
+		send,
+		onMessage(handler) {
+			inbound.add(handler);
+			return () => inbound.delete(handler);
+		},
+		onError(handler) {
+			errors.add(handler);
+			return () => errors.delete(handler);
+		},
+		async terminate() {
+			intentionalExit.value = true;
+			try {
+				proc.kill("SIGKILL");
+			} catch {
+				// Already gone.
+			}
+		},
+	};
+}
+
+/**
+ * A stand-in handle used when the worker subprocess cannot be spawned. It
+ * ponges `ping` (so the smoke probe and readiness checks still resolve) and
+ * answers every other request with the spawn error so callers fail fast
+ * instead of awaiting forever.
+ */
+export function createUnavailableWorker<
+	Inbound extends { type: string; id: string },
+	Outbound extends { type: string },
+>(error: unknown): WorkerHandle<Inbound, Outbound> {
+	const listeners = new Set<(message: Outbound) => void>();
+	const errorMessage = error instanceof Error ? error.message : String(error);
+	const emit = (message: WorkerOutboundBase): void => {
+		// The stub only ever emits pong/error — members of every concrete worker
+		// Outbound union — but the generic cannot prove it, hence the assertion.
+		for (const listener of listeners) listener(message as unknown as Outbound);
+	};
+	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();
+		},
+	};
+}
+
+/**
+ * Spawn a worker handle, falling back to {@link createUnavailableWorker} (after
+ * a warning) when the subprocess cannot be created so the feature degrades
+ * gracefully instead of throwing into callers.
+ */
+export function spawnWorkerOrUnavailable<Handle>(
+	spawn: () => Handle,
+	unavailable: (error: unknown) => Handle,
+	warnMessage: string,
+): Handle {
+	try {
+		return spawn();
+	} catch (error) {
+		logger.warn(warnMessage, { error: error instanceof Error ? error.message : String(error) });
+		return unavailable(error);
+	}
+}
+
+/** Forward a worker's structured `log` message to the matching logger level. */
+export function logWorkerMessage(message: WorkerLogMessage): 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);
+}
+
+/**
+ * Drive the ping/pong readiness probe wired into `omp --smoke-test`: send one
+ * `ping`, resolve on the first `pong` (ignoring `log` chatter), and reject on
+ * any other message, a worker error, or the timeout. Always tears the handle
+ * down on the way out. `label` prefixes the failure messages.
+ */
+export async function smokeTestWorker<Inbound extends { type: string; id: string }, Outbound extends { type: string }>(
+	handle: WorkerHandle<Inbound, Outbound>,
+	label: string,
+	timeoutMs: number,
+): Promise<void> {
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	const timer = setTimeout(() => reject(new Error(`${label} 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(`${label}: expected pong, got ${JSON.stringify(message)}`));
+	});
+	const unsubscribeError = handle.onError(reject);
+	try {
+		handle.send({ type: "ping", id: "smoke" } as Inbound);
+		await promise;
+	} finally {
+		clearTimeout(timer);
+		unsubscribeMessage();
+		unsubscribeError();
+		await handle.terminate();
+	}
+}