385 lines
15 KiB
TypeScript
385 lines
15 KiB
TypeScript
import { describe, test, expect } from "bun:test";
|
||
import {
|
||
applyLoadProgress,
|
||
deviceAttempts,
|
||
dtypeRequiresF16,
|
||
initialLoadProgress,
|
||
isUiToWorkerMessage,
|
||
isWorkerToUiMessage,
|
||
TranscribeProgressTracker,
|
||
type LoadProgressMessage,
|
||
type UiToWorkerMessage,
|
||
type WorkerToUiMessage,
|
||
} from "../src/worker/protocol";
|
||
|
||
// Protocol handling only — no worker, no network, no model weights.
|
||
|
||
describe("isUiToWorkerMessage", () => {
|
||
test("accepts valid load and transcribe messages", () => {
|
||
const loads: UiToWorkerMessage[] = [
|
||
{ type: "load", modelId: "onnx-community/whisper-tiny.en", device: "wasm" },
|
||
{ type: "load", modelId: "onnx-community/whisper-tiny", device: "webgpu" },
|
||
{ type: "load", modelId: "onnx-community/whisper-base", device: "auto" },
|
||
{ type: "load", modelId: "onnx-community/whisper-base", device: "auto", dtype: "q4" },
|
||
{
|
||
type: "load",
|
||
modelId: "onnx-community/whisper-large-v3-turbo",
|
||
device: "webgpu",
|
||
dtype: { encoder_model: "fp16", decoder_model_merged: "q4" },
|
||
},
|
||
];
|
||
const transcribes: UiToWorkerMessage[] = [
|
||
{ type: "transcribe", audio: new Float32Array(8) },
|
||
{ type: "transcribe", audio: new Float32Array(8), language: "da", task: "transcribe" },
|
||
{ type: "transcribe", audio: new Float32Array(8), language: "auto", task: "translate" },
|
||
{ type: "transcribe", audio: new Float32Array(8), requestId: "live-3" },
|
||
{ type: "transcribe", audio: new Float32Array(8), requestId: "engine-1-1", stream: true },
|
||
{ type: "transcribe", audio: new Float32Array(8), stream: false },
|
||
];
|
||
for (const message of [...loads, ...transcribes]) {
|
||
expect(isUiToWorkerMessage(message)).toBe(true);
|
||
}
|
||
});
|
||
|
||
test("rejects malformed messages", () => {
|
||
expect(isUiToWorkerMessage(null)).toBe(false);
|
||
expect(isUiToWorkerMessage("load")).toBe(false);
|
||
expect(isUiToWorkerMessage({ type: "load" })).toBe(false);
|
||
expect(isUiToWorkerMessage({ type: "load", modelId: "x", device: "cuda" })).toBe(false);
|
||
expect(isUiToWorkerMessage({ type: "load", modelId: "x", device: "auto", dtype: "int3" })).toBe(false);
|
||
expect(
|
||
isUiToWorkerMessage({
|
||
type: "load",
|
||
modelId: "x",
|
||
device: "auto",
|
||
dtype: { encoder_model: "fp16" },
|
||
}),
|
||
).toBe(false);
|
||
expect(isUiToWorkerMessage({ type: "transcribe", audio: [1, 2, 3] })).toBe(false);
|
||
expect(
|
||
isUiToWorkerMessage({ type: "transcribe", audio: new Float32Array(4), task: "detect" }),
|
||
).toBe(false);
|
||
expect(
|
||
isUiToWorkerMessage({ type: "transcribe", audio: new Float32Array(4), requestId: 7 }),
|
||
).toBe(false);
|
||
expect(
|
||
isUiToWorkerMessage({ type: "transcribe", audio: new Float32Array(4), stream: "yes" }),
|
||
).toBe(false);
|
||
expect(isUiToWorkerMessage({ type: "result", text: "", chunks: [] })).toBe(false);
|
||
});
|
||
});
|
||
|
||
describe("isWorkerToUiMessage", () => {
|
||
test("accepts every worker → UI message shape", () => {
|
||
const messages: WorkerToUiMessage[] = [
|
||
{ type: "load-progress", status: "initiate", file: "model.onnx" },
|
||
{ type: "load-progress", status: "download", file: "model.onnx" },
|
||
{
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: "model.onnx",
|
||
loaded: 10,
|
||
total: 100,
|
||
progress: 10,
|
||
},
|
||
{ type: "load-progress", status: "done", file: "model.onnx" },
|
||
{ type: "load-progress", status: "ready", device: "wasm" },
|
||
{ type: "load-progress", status: "ready", device: "wasm", threads: 11 },
|
||
{ type: "load-progress", status: "ready", device: "webgpu" },
|
||
{ type: "load-progress", status: "ready", device: "webgpu", threads: 1 },
|
||
{ type: "partial", text: " hello" },
|
||
{ type: "partial", text: " hello", requestId: "live-1" },
|
||
{ type: "transcribe-progress", processedSeconds: 0, totalSeconds: 90 },
|
||
{ type: "transcribe-progress", processedSeconds: 42.5, totalSeconds: 90, requestId: "live-1" },
|
||
{ type: "result", text: "hello world", chunks: [{ text: "hello world", timestamp: [0, 1.2] }] },
|
||
{ type: "result", text: "hello world", chunks: [], requestId: "live-1" },
|
||
{ type: "error", stage: "load", message: "boom" },
|
||
{ type: "error", stage: "transcribe", message: "boom", requestId: "live-2" },
|
||
{
|
||
type: "load-restart",
|
||
modelId: "onnx-community/whisper-medium-ONNX",
|
||
preference: "auto",
|
||
device: "wasm",
|
||
},
|
||
{
|
||
type: "load-restart",
|
||
modelId: "onnx-community/whisper-tiny.en",
|
||
preference: "webgpu",
|
||
device: "wasm",
|
||
},
|
||
];
|
||
for (const message of messages) {
|
||
expect(isWorkerToUiMessage(message)).toBe(true);
|
||
}
|
||
});
|
||
|
||
test("rejects malformed messages", () => {
|
||
expect(isWorkerToUiMessage(undefined)).toBe(false);
|
||
expect(isWorkerToUiMessage({})).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "load-progress", status: "progress", file: "f" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "load-progress", status: "ready" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "load-progress", status: "ready", device: "cuda" })).toBe(false);
|
||
expect(
|
||
isWorkerToUiMessage({ type: "load-progress", status: "ready", device: "wasm", threads: "12" }),
|
||
).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "partial" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "partial", text: "x", requestId: 1 })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "transcribe-progress", processedSeconds: 1 })).toBe(false);
|
||
expect(
|
||
isWorkerToUiMessage({ type: "transcribe-progress", processedSeconds: "1", totalSeconds: 2 }),
|
||
).toBe(false);
|
||
expect(
|
||
isWorkerToUiMessage({
|
||
type: "transcribe-progress",
|
||
processedSeconds: 1,
|
||
totalSeconds: 2,
|
||
requestId: 9,
|
||
}),
|
||
).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "result", text: "x", chunks: "nope" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "error", stage: "decode", message: "m" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "load", modelId: "x", device: "wasm" })).toBe(false);
|
||
expect(isWorkerToUiMessage({ type: "load-restart", modelId: "x" })).toBe(false);
|
||
expect(
|
||
isWorkerToUiMessage({ type: "load-restart", modelId: "x", preference: "auto", device: "auto" }),
|
||
).toBe(false);
|
||
expect(
|
||
isWorkerToUiMessage({ type: "load-restart", preference: "auto", device: "wasm" }),
|
||
).toBe(false);
|
||
});
|
||
});
|
||
|
||
describe("dtypeRequiresF16 (D-025)", () => {
|
||
test("plain dtypes: only fp16 needs the shader-f16 feature", () => {
|
||
expect(dtypeRequiresF16("fp16")).toBe(true);
|
||
expect(dtypeRequiresF16("fp32")).toBe(false);
|
||
expect(dtypeRequiresF16("q8")).toBe(false);
|
||
expect(dtypeRequiresF16("q4")).toBe(false);
|
||
});
|
||
|
||
test("per-part dtypes: fp16 in either part needs it", () => {
|
||
expect(
|
||
dtypeRequiresF16({ encoder_model: "fp16", decoder_model_merged: "q4" }),
|
||
).toBe(true);
|
||
expect(
|
||
dtypeRequiresF16({ encoder_model: "fp32", decoder_model_merged: "fp16" }),
|
||
).toBe(true);
|
||
expect(
|
||
dtypeRequiresF16({ encoder_model: "fp32", decoder_model_merged: "fp32" }),
|
||
).toBe(false);
|
||
expect(
|
||
dtypeRequiresF16({ encoder_model: "fp32", decoder_model_merged: "q4" }),
|
||
).toBe(false);
|
||
});
|
||
});
|
||
|
||
describe("deviceAttempts (D-004, D-017)", () => {
|
||
test("auto resolves WebNN → WebGPU → WASM when everything is available", () => {
|
||
expect(deviceAttempts("auto", { webnn: true, webgpu: true })).toEqual([
|
||
"webnn",
|
||
"webgpu",
|
||
"wasm",
|
||
]);
|
||
});
|
||
|
||
test("auto skips backends the probe ruled out, in order", () => {
|
||
expect(deviceAttempts("auto", { webnn: false, webgpu: true })).toEqual([
|
||
"webgpu",
|
||
"wasm",
|
||
]);
|
||
expect(deviceAttempts("auto", { webnn: true, webgpu: false })).toEqual([
|
||
"webnn",
|
||
"wasm",
|
||
]);
|
||
expect(deviceAttempts("auto", { webnn: false, webgpu: false })).toEqual(["wasm"]);
|
||
});
|
||
|
||
test("explicit webgpu falls back to WASM only", () => {
|
||
expect(deviceAttempts("webgpu", { webnn: true, webgpu: true })).toEqual([
|
||
"webgpu",
|
||
"wasm",
|
||
]);
|
||
expect(deviceAttempts("webgpu", { webnn: true, webgpu: false })).toEqual(["wasm"]);
|
||
});
|
||
|
||
test("explicit webnn falls back to WASM only (no WebGPU detour)", () => {
|
||
expect(deviceAttempts("webnn", { webnn: true, webgpu: true })).toEqual([
|
||
"webnn",
|
||
"wasm",
|
||
]);
|
||
expect(deviceAttempts("webnn", { webnn: false, webgpu: true })).toEqual(["wasm"]);
|
||
});
|
||
|
||
test("wasm is honored as-is and every chain ends in WASM", () => {
|
||
expect(deviceAttempts("wasm", { webnn: true, webgpu: true })).toEqual(["wasm"]);
|
||
for (const preference of ["auto", "webnn", "webgpu", "wasm"] as const) {
|
||
for (const webnn of [true, false]) {
|
||
for (const webgpu of [true, false]) {
|
||
const attempts = deviceAttempts(preference, { webnn, webgpu });
|
||
expect(attempts.length).toBeGreaterThan(0);
|
||
expect(attempts.at(-1)).toBe("wasm");
|
||
}
|
||
}
|
||
}
|
||
});
|
||
});
|
||
|
||
describe("TranscribeProgressTracker", () => {
|
||
// Production shape: 30 s windows advancing by 20 s (30 − 2·5, the worker's
|
||
// chunk/stride settings), like the chunked ASR pipeline.
|
||
const WINDOW = 30;
|
||
const JUMP = 20;
|
||
|
||
test("maps window-relative timestamps to absolute seconds per window", () => {
|
||
const tracker = new TranscribeProgressTracker(90, WINDOW, JUMP);
|
||
expect(tracker.timestamp(0)).toBeNull(); // 0 does not advance past 0
|
||
expect(tracker.timestamp(4.2)).toBe(4.2);
|
||
expect(tracker.timestamp(29.5)).toBe(29.5);
|
||
expect(tracker.finalizeWindow()).toBeNull(); // 1·20 < 29.5 already seen
|
||
// Second window: timestamps are relative again, offset by one jump.
|
||
expect(tracker.timestamp(3)).toBeNull(); // 20 + 3 < 29.5
|
||
expect(tracker.timestamp(12)).toBe(32);
|
||
});
|
||
|
||
test("is monotone across overlapping windows", () => {
|
||
const tracker = new TranscribeProgressTracker(90, WINDOW, JUMP);
|
||
expect(tracker.timestamp(28)).toBe(28);
|
||
tracker.finalizeWindow();
|
||
// The next window revisits overlapped audio — earlier absolute times
|
||
// than what was already reported are swallowed, never regressions.
|
||
expect(tracker.timestamp(5)).toBeNull(); // 20 + 5 = 25 < 28
|
||
expect(tracker.timestamp(10)).toBe(30);
|
||
});
|
||
|
||
test("clamps within the window and to the total duration", () => {
|
||
const tracker = new TranscribeProgressTracker(35, WINDOW, JUMP);
|
||
// A stray timestamp beyond the window length clamps to the window end.
|
||
expect(tracker.timestamp(45)).toBe(30);
|
||
expect(tracker.timestamp(-1)).toBeNull(); // negative clamps to window start
|
||
tracker.finalizeWindow();
|
||
// Second (final) window: 20 + 16 = 36 clamps to the 35 s total.
|
||
expect(tracker.timestamp(16)).toBe(35);
|
||
expect(tracker.timestamp(17)).toBeNull(); // already at the end
|
||
});
|
||
|
||
test("finalizeWindow advances to the next window start", () => {
|
||
const tracker = new TranscribeProgressTracker(90, WINDOW, JUMP);
|
||
expect(tracker.finalizeWindow()).toBe(20);
|
||
expect(tracker.finalizeWindow()).toBe(40);
|
||
});
|
||
|
||
test("single-window audio never reports past its duration", () => {
|
||
const tracker = new TranscribeProgressTracker(8, WINDOW, JUMP);
|
||
expect(tracker.timestamp(6)).toBe(6);
|
||
expect(tracker.timestamp(8)).toBe(8);
|
||
expect(tracker.finalizeWindow()).toBeNull(); // 20 clamps to 8, not an advance
|
||
});
|
||
|
||
describe("forWindows (variable-length VAD chunks, D-WS3a)", () => {
|
||
// Two speech chunks with a skipped 15 s silence between them: [0–20 s]
|
||
// and [35–60 s] of a 70 s file.
|
||
const windows = [
|
||
{ toAbsolute: (t: number) => t, endSeconds: 20, nextStartSeconds: 35 },
|
||
{ toAbsolute: (t: number) => 35 + t, endSeconds: 60, nextStartSeconds: 70 },
|
||
];
|
||
|
||
test("maps chunk-relative times through each window", () => {
|
||
const tracker = TranscribeProgressTracker.forWindows(70, windows);
|
||
expect(tracker.timestamp(5)).toBe(5);
|
||
// Finishing a chunk counts the skipped silence up to the next one.
|
||
expect(tracker.finalizeWindow()).toBe(35);
|
||
expect(tracker.timestamp(10)).toBe(45);
|
||
expect(tracker.timestamp(2)).toBeNull(); // 37 < 45: no regression
|
||
expect(tracker.finalizeWindow()).toBe(70);
|
||
expect(tracker.finalizeWindow()).toBeNull(); // past the last window
|
||
});
|
||
|
||
test("clamps a stray timestamp to the window's end", () => {
|
||
const tracker = TranscribeProgressTracker.forWindows(70, windows);
|
||
expect(tracker.timestamp(400)).toBe(20);
|
||
});
|
||
|
||
test("no windows (no speech) → nothing to report but the end", () => {
|
||
const tracker = TranscribeProgressTracker.forWindows(70, []);
|
||
expect(tracker.timestamp(3)).toBeNull();
|
||
expect(tracker.finalizeWindow()).toBe(70);
|
||
});
|
||
});
|
||
});
|
||
|
||
describe("applyLoadProgress", () => {
|
||
test("tracks per-file byte progress without mutating previous state", () => {
|
||
const p1: LoadProgressMessage = { type: "load-progress", status: "initiate", file: "a.onnx" };
|
||
const p2: LoadProgressMessage = {
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: "a.onnx",
|
||
loaded: 50,
|
||
total: 200,
|
||
progress: 25,
|
||
};
|
||
|
||
const s1 = applyLoadProgress(initialLoadProgress, p1);
|
||
const s2 = applyLoadProgress(s1, p2);
|
||
|
||
expect(initialLoadProgress.files).toEqual({});
|
||
expect(s1.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 0, total: 0, done: false });
|
||
expect(s1.files["a.onnx"]!.done).toBe(false); // s1 untouched by s2
|
||
expect(s2.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 50, total: 200, done: false });
|
||
expect(s2.ready).toBe(false);
|
||
});
|
||
|
||
test("tracks multiple files independently", () => {
|
||
let state = initialLoadProgress;
|
||
state = applyLoadProgress(state, {
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: "a.onnx",
|
||
loaded: 10,
|
||
total: 100,
|
||
progress: 10,
|
||
});
|
||
state = applyLoadProgress(state, {
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: "b.json",
|
||
loaded: 5,
|
||
total: 5,
|
||
progress: 100,
|
||
});
|
||
expect(Object.keys(state.files)).toEqual(["a.onnx", "b.json"]);
|
||
expect(state.files["a.onnx"]!.loaded).toBe(10);
|
||
expect(state.files["b.json"]!.total).toBe(5);
|
||
});
|
||
|
||
test("done snaps a file's loaded bytes to its total", () => {
|
||
let state = applyLoadProgress(initialLoadProgress, {
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: "a.onnx",
|
||
loaded: 199,
|
||
total: 200,
|
||
progress: 99.5,
|
||
});
|
||
state = applyLoadProgress(state, { type: "load-progress", status: "done", file: "a.onnx" });
|
||
expect(state.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 200, total: 200, done: true });
|
||
});
|
||
|
||
test("ready flips the ready flag, records the device, keeps file state", () => {
|
||
let state = applyLoadProgress(initialLoadProgress, {
|
||
type: "load-progress",
|
||
status: "done",
|
||
file: "a.onnx",
|
||
});
|
||
expect(state.device).toBeNull();
|
||
state = applyLoadProgress(state, {
|
||
type: "load-progress",
|
||
status: "ready",
|
||
device: "webgpu",
|
||
});
|
||
expect(state.ready).toBe(true);
|
||
expect(state.device).toBe("webgpu");
|
||
expect(state.files["a.onnx"]!.done).toBe(true);
|
||
});
|
||
});
|