fix(stt): never hold a chunk on the word aligner's download

Address the review of #954:
- the aligner downloads in the background; chunks before it lands keep
  whisper's times, a cached copy is verified in place; a 30 s stall, Cancel
  and quit abort it, and a failed one is retried on the next transcription
- accept an emissions region ending within two frames of a stretch that runs
  to the end of the upload
- the harness skips (evaluate) or names (real-check) unreadable emissions
- test-whisper-stt says which aligner file it looked for; the doc lists the
  aligner checks it runs

Author: EtienneLescot

electron/stt/ctcAlign.test.ts

@@ -174,6 +174,17 @@ describe("alignWordsOnEmissions", () => {
 		expect(out[1].endSec).toBe(out[0].endSec);
 	});
 
+	it("aligns a stretch that runs to the end of the upload, past the last frame", () => {
+		// 1 s of audio gives 49 frames (0.98 s): the stretch ends after them.
+		const tokens = [...gap(30), "a", "b", ...gap(17)];
+		const out = alignWordsOnEmissions(
+			[word("ab", 0.65, 1, 0.7)],
+			[{ startSec: 0.5, endSec: 1 }],
+			emissions(tokens),
+		);
+		expect(r(out[0].startSec)).toBe(r(30 * STRIDE + PAUSE_START_OFFSET_SEC));
+	});
+
 	it("keeps the helper's times where it cannot align", () => {
 		const words = [word("ab", 0.3, 0.5, 0.35), word("c", 1.5, 1.7, 1.6)];
 		// The second word is outside every stretch; the first one's letters do not fit.

electron/stt/ctcAlign.ts

@@ -198,8 +198,11 @@ export function alignWordsOnEmissions<W extends AlignableWord>(
 		const until = Math.min(offset + TAIL_SEC, speech[i + 1]?.startSec ?? Number.POSITIVE_INFINITY);
 		const owned: number[] = [];
 		for (; k < out.length && words[k].anchorSec < until; k++) owned.push(k);
+		// A region's frames stop up to one receptive field (25 ms) short of the audio
+		// it was cut from, so a stretch that runs to the end of the upload ends
+		// past them; two frames of slack keep it. `hi` below stays on the frames.
 		const region = emissions.regions.find(
-			(r) => r.startSec <= onset && r.startSec + r.frames * strideSec >= offset,
+			(r) => r.startSec <= onset && r.startSec + (r.frames + 2) * strideSec >= offset,
 		);
 		if (!region || owned.length === 0) continue;
 		// The frames of this stretch: its margin, stopping at the neighbours' speech.

electron/stt/index.test.ts

@@ -35,6 +35,8 @@ vi.mock("./whisperServer", () => {
 vi.mock("./modelManager", () => ({
 	ensureModels: vi.fn(async () => undefined),
 	ensureAligner: vi.fn(async ({ language }: { language: string }) => `/fake/${language}.gguf`),
+	cachedAligner: vi.fn(async () => null),
+	alignerPath: vi.fn((base: string, language: string): string | null => `${base}/${language}.gguf`),
 	modelPaths: (base: string) => ({
 		whisper: `${base}/whisper-ggml/ggml-small-q8_0.bin`,
 	}),
@@ -471,39 +473,108 @@ describe("SttManager", () => {
 		expect(fakeWhisperServer.start).toHaveBeenCalledOnce();
 	});
 
-	it("hands the chunks an aligner, tried once per run when its download fails", async () => {
-		const { ensureAligner } = await import("./modelManager");
-		const mocked = vi.mocked(ensureAligner);
-		mocked.mockClear();
-		mocked.mockRejectedValueOnce(new Error("offline"));
-		// Each chunk asks for the aligner of the language it detected, as the helper does.
-		const asked: Array<string | null> = [];
-		fakeWhisperServer.transcribe.mockImplementation(
-			async (opts: { alignerFor?: (language: string) => Promise<string | null> }) => {
-				asked.push((await opts.alignerFor?.("fr")) ?? null);
-				return {
-					segments: [],
-					wordSegments: [],
-					detectedLanguage: "fr",
-					backend: "whispercpp-cpu" as const,
-				};
-			},
-		);
-		const mgr = new SttManager();
-		await mgr.init({ modelsBaseDir: "/tmp/fake-stt-models" });
+	describe("word aligner", () => {
+		/** Chunks report what `alignerFor("fr")` gave them, as the helper would use it. */
+		function recordAligner() {
+			const asked: Array<string | null> = [];
+			fakeWhisperServer.transcribe.mockImplementation(
+				async (opts: { alignerFor?: (language: string) => Promise<string | null> }) => {
+					asked.push((await opts.alignerFor?.("fr")) ?? null);
+					return {
+						segments: [],
+						wordSegments: [],
+						detectedLanguage: "fr",
+						backend: "whispercpp-cpu" as const,
+					};
+				},
+			);
+			return asked;
+		}
 		const long = new Float32Array(16_000 * 400);
-		await mgr.transcribe({ samples: long, language: "fr" });
-		// Failed once: the rest of the run keeps whisper's times without retrying.
-		expect(asked.length).toBeGreaterThan(1);
-		expect(asked.every((x) => x === null)).toBe(true);
-		expect(mocked).toHaveBeenCalledOnce();
-
-		// The next run tries again, and a success is kept for the session.
-		asked.length = 0;
-		await mgr.transcribe({ samples: long, language: "fr" });
-		await mgr.transcribe({ samples: new Float32Array(16_000), language: "fr" });
-		expect(asked.every((x) => x === "/fake/fr.gguf")).toBe(true);
-		expect(mocked).toHaveBeenCalledTimes(2);
+
+		async function mocks() {
+			const mm = await import("./modelManager");
+			const ensure = vi.mocked(mm.ensureAligner);
+			const cached = vi.mocked(mm.cachedAligner);
+			ensure.mockReset();
+			cached.mockReset();
+			cached.mockResolvedValue(null);
+			return { ensure, cached };
+		}
+
+		it("never holds a chunk on the download, and uses the aligner once it lands", async () => {
+			const { ensure } = await mocks();
+			let land: (file: string) => void = () => undefined;
+			ensure.mockImplementation(() => new Promise((resolve) => (land = resolve)));
+			const asked = recordAligner();
+			const mgr = new SttManager();
+			await mgr.init({ modelsBaseDir: "/tmp/fake-stt-models" });
+
+			// The download never finishes during this run: every chunk goes on without it.
+			await mgr.transcribe({ samples: long, language: "fr" });
+			expect(asked.length).toBeGreaterThan(1);
+			expect(asked.every((x) => x === null)).toBe(true);
+			expect(ensure).toHaveBeenCalledOnce();
+
+			land("/fake/fr.gguf");
+			await new Promise((resolve) => setTimeout(resolve, 0));
+			asked.length = 0;
+			await mgr.transcribe({ samples: new Float32Array(16_000), language: "fr" });
+			expect(asked).toEqual(["/fake/fr.gguf"]);
+			expect(ensure).toHaveBeenCalledOnce();
+		});
+
+		it("uses a verified copy already on disk without downloading", async () => {
+			const { ensure, cached } = await mocks();
+			cached.mockResolvedValue("/fake/fr.gguf");
+			const asked = recordAligner();
+			const mgr = new SttManager();
+			await mgr.init({ modelsBaseDir: "/tmp/fake-stt-models" });
+			await mgr.transcribe({ samples: long, language: "fr" });
+			expect(asked.every((x) => x === "/fake/fr.gguf")).toBe(true);
+			expect(cached).toHaveBeenCalledOnce();
+			expect(ensure).not.toHaveBeenCalled();
+		});
+
+		it("aborts the download on cancel and on shutdown, then retries on a later run", async () => {
+			const { ensure } = await mocks();
+			const signals: AbortSignal[] = [];
+			ensure.mockImplementation(
+				({ signal }: { signal?: AbortSignal }) =>
+					new Promise((_, reject) => {
+						if (signal) signals.push(signal);
+						signal?.addEventListener("abort", () => reject(signal.reason));
+					}),
+			);
+			recordAligner();
+			const mgr = new SttManager();
+			await mgr.init({ modelsBaseDir: "/tmp/fake-stt-models" });
+			await mgr.transcribe({ samples: long, language: "fr" });
+			expect(signals).toHaveLength(1);
+
+			mgr.cancel();
+			expect(signals[0].aborted).toBe(true);
+			await new Promise((resolve) => setTimeout(resolve, 0));
+
+			// Not cached as failed for good: the next transcription starts it again.
+			await mgr.transcribe({ samples: long, language: "fr" });
+			expect(signals).toHaveLength(2);
+			expect(signals[1].aborted).toBe(false);
+			await mgr.shutdown();
+			expect(signals[1].aborted).toBe(true);
+		});
+
+		it("gives a language without an aligner nothing, and downloads nothing", async () => {
+			const { ensure } = await mocks();
+			const mm = await import("./modelManager");
+			vi.mocked(mm.alignerPath).mockReturnValueOnce(null);
+			const asked = recordAligner();
+			const mgr = new SttManager();
+			await mgr.init({ modelsBaseDir: "/tmp/fake-stt-models" });
+			await mgr.transcribe({ samples: new Float32Array(16_000), language: "fr" });
+			expect(asked).toEqual([null]);
+			expect(ensure).not.toHaveBeenCalled();
+		});
 	});
 
 	it("fans status out to every sink, and detaching one leaves the others", async () => {

electron/stt/index.ts

@@ -2,7 +2,13 @@ import path from "node:path";
 import { app, type IpcMain } from "electron";
 import { planChunks } from "./chunking";
 import { extractMono16kPcm } from "./extractAudio";
-import { ensureAligner, ensureModels, modelPaths } from "./modelManager";
+import {
+	alignerPath,
+	cachedAligner,
+	ensureAligner,
+	ensureModels,
+	modelPaths,
+} from "./modelManager";
 import type {
 	SttPhraseSegment,
 	SttStatusEvent,
@@ -106,36 +112,58 @@ export class SttManager {
 	private cancelEpoch = 0;
 
 	/**
-	 * The word aligner per language (ctcAlign.ts), resolved once per session: its
-	 * file is hashed on the way in, and that is not free at 350 MB. A language
-	 * whose download failed is left out for the rest of the run rather than
-	 * retried per chunk, and tried again on the next one.
+	 * The word aligner per language (ctcAlign.ts). A chunk never waits for a
+	 * download: it is an improvement, and a 350 MB fetch inside a chunk would hold
+	 * the helper's queue for as long as the network takes, past any request
+	 * ceiling and out of Cancel's reach. So:
+	 * - a copy already on disk is verified in place the first time a language is
+	 *   asked for in a session (a local read, no network), then reused;
+	 * - otherwise the download starts in the background and the chunks keep
+	 *   whisper's times until it lands; later chunks and transcriptions use it;
+	 * - Cancel and quit abort it, and so does a stall (`ensureAligner`). A failed
+	 *   one is not retried for the rest of the run, only on the next transcription.
 	 */
-	private readonly aligners = new Map<string, Promise<string | null>>();
+	private readonly alignersReady = new Map<string, string>();
+	private readonly alignersChecked = new Set<string>();
+	private readonly alignerDownloads = new Map<string, AbortController>();
 	private readonly alignersFailed = new Set<string>();
 
-	private readonly alignerFor = (language: string): Promise<string | null> => {
-		if (this.alignersFailed.has(language)) return Promise.resolve(null);
-		let resolved = this.aligners.get(language);
-		if (!resolved) {
-			resolved = ensureAligner({
-				baseDir: this.getModelsDir(),
-				language,
-				onProgress: (downloadedBytes, totalBytes) =>
-					this.emit({ phase: "model", model: "ctc-aligner", downloadedBytes, totalBytes }),
-			}).catch((error: unknown) => {
+	private readonly alignerFor = async (language: string): Promise<string | null> => {
+		const ready = this.alignersReady.get(language);
+		if (ready) return ready;
+		if (!alignerPath(this.getModelsDir(), language)) return null;
+		if (!this.alignersChecked.has(language)) {
+			this.alignersChecked.add(language);
+			const cached = await cachedAligner(this.getModelsDir(), language);
+			if (cached) {
+				this.alignersReady.set(language, cached);
+				return cached;
+			}
+		}
+		if (!this.alignersFailed.has(language) && !this.alignerDownloads.has(language)) {
+			this.downloadAligner(language);
+		}
+		return null;
+	};
+
+	private downloadAligner(language: string): void {
+		const controller = new AbortController();
+		this.alignerDownloads.set(language, controller);
+		console.info(`[stt] downloading the "${language}" word aligner in the background`);
+		ensureAligner({ baseDir: this.getModelsDir(), language, signal: controller.signal })
+			.then((file) => {
+				if (file) this.alignersReady.set(language, file);
+				console.info(`[stt] "${language}" word aligner ready`);
+			})
+			.catch((error: unknown) => {
+				this.alignersFailed.add(language);
 				console.warn(
 					`[stt] no word aligner for "${language}", keeping whisper's word times: ` +
 						`${error instanceof Error ? error.message : String(error)}`,
 				);
-				this.alignersFailed.add(language);
-				this.aligners.delete(language);
-				return null;
-			});
-			this.aligners.set(language, resolved);
-		}
-		return resolved;
-	};
+			})
+			.finally(() => this.alignerDownloads.delete(language));
+	}
 
 	/**
 	 * The extraction in flight, if any. `cancelEpoch` alone stops the CHUNK loop, which is
@@ -176,6 +204,7 @@ export class SttManager {
 	cancel(): void {
 		this.cancelEpoch++;
 		this.extraction?.abort();
+		for (const download of this.alignerDownloads.values()) download.abort(cancelledError());
 	}
 
 	/**
@@ -470,7 +499,7 @@ export class SttManager {
 	async shutdown(): Promise<void> {
 		if (this.shuttingDown) return;
 		this.shuttingDown = true;
-		this.cancelEpoch++;
+		this.cancel();
 		await this.server.shutdown();
 	}
 }

electron/stt/modelManager.test.ts

@@ -8,6 +8,7 @@ import {
 	alignerPath,
 	areModelsPresent,
 	CTC_ALIGNERS,
+	cachedAligner,
 	ensureAligner,
 	ensureModels,
 	modelPaths,
@@ -283,17 +284,14 @@ describe("modelManager", () => {
 			const bytes = Buffer.from("gguf weights");
 			const original = CTC_ALIGNERS.fr.expectedSha256;
 			CTC_ALIGNERS.fr.expectedSha256 = createHash("sha256").update(bytes).digest("hex");
-			const progress: number[] = [];
 			try {
 				const file = await ensureAligner({
 					baseDir: dir,
 					language: "fr",
 					fetcher: async () => new Response(bytes, { status: 200 }),
-					onProgress: (done) => progress.push(done),
 				});
 				expect(file).toBe(alignerPath(dir, "fr"));
 				expect(await readFile(file as string)).toEqual(bytes);
-				expect(progress[progress.length - 1]).toBe(bytes.length);
 			} finally {
 				CTC_ALIGNERS.fr.expectedSha256 = original;
 			}
@@ -309,5 +307,72 @@ describe("modelManager", () => {
 			).rejects.toThrow(/SHA-256 mismatch/);
 			expect(existsSync(alignerPath(dir, "en") as string)).toBe(false);
 		});
+
+		/** A body that sends a few bytes, then nothing, ever. */
+		const stalledBody = () =>
+			new Response(
+				new ReadableStream<Uint8Array>({
+					start(controller) {
+						controller.enqueue(new Uint8Array(8));
+					},
+				}),
+				{ status: 200 },
+			);
+
+		it("aborts a download that stops sending, and leaves no partial file", async () => {
+			const file = alignerPath(dir, "fr") as string;
+			await expect(
+				ensureAligner({
+					baseDir: dir,
+					language: "fr",
+					stallMs: 50,
+					fetcher: async () => stalledBody(),
+				}),
+			).rejects.toThrow(/stalled/);
+			expect(existsSync(file)).toBe(false);
+			expect(existsSync(`${file}.partial`)).toBe(false);
+		});
+
+		it("stops when its signal aborts, mid-body or between attempts", async () => {
+			const controller = new AbortController();
+			const download = ensureAligner({
+				baseDir: dir,
+				language: "fr",
+				signal: controller.signal,
+				fetcher: async () => stalledBody(),
+			});
+			setTimeout(() => controller.abort(new Error("cancelled")), 20);
+			await expect(download).rejects.toThrow("cancelled");
+
+			// A 503 puts it in a backoff of seconds: the abort must not wait for it.
+			const again = new AbortController();
+			const started = Date.now();
+			const retrying = ensureAligner({
+				baseDir: dir,
+				language: "fr",
+				signal: again.signal,
+				fetcher: async () => new Response("busy", { status: 503 }),
+			});
+			setTimeout(() => again.abort(new Error("quit")), 20);
+			await expect(retrying).rejects.toThrow("quit");
+			expect(Date.now() - started).toBeLessThan(1000);
+		});
+
+		it("verifies a cached copy without the network", async () => {
+			expect(await cachedAligner(dir, "fr")).toBeNull();
+			const file = alignerPath(dir, "fr") as string;
+			const bytes = Buffer.from("gguf weights");
+			await mkdir(path.dirname(file), { recursive: true });
+			await writeFile(file, bytes);
+			// Present but not the pinned bytes: not usable.
+			expect(await cachedAligner(dir, "fr")).toBeNull();
+			const original = CTC_ALIGNERS.fr.expectedSha256;
+			CTC_ALIGNERS.fr.expectedSha256 = createHash("sha256").update(bytes).digest("hex");
+			try {
+				expect(await cachedAligner(dir, "fr")).toBe(file);
+			} finally {
+				CTC_ALIGNERS.fr.expectedSha256 = original;
+			}
+		});
 	});
 });