whisper_web/tests/engine.test.ts

import { describe, expect, test } from "bun:test";
import {
  BrowserEngine,
  ENGINE_CHOICE_KEY,
  ServerEngine,
  fetchServerEngineStatus,
  isTranscribeStreamLine,
  loadEngineChoice,
  saveEngineChoice,
  type StorageLike,
  type WorkerLike,
} from "../src/lib/engine";
import type { FetchLike } from "../src/lib/authClient";
import type { TranscriptChunk } from "../src/worker/protocol";

// D-010: both engines behind one interface, testable without a DOM/worker/server.

// ---------------------------------------------------------------------------
// BrowserEngine — requestId correlation against a fake worker
// ---------------------------------------------------------------------------

class FakeWorker implements WorkerLike {
  posted: { message: Record<string, unknown>; transfer?: Transferable[] }[] = [];
  private listeners = { message: [] as ((e: { data: unknown }) => void)[], error: [] as ((e: { message?: string }) => void)[] };

  postMessage(message: unknown, transfer?: Transferable[]): void {
    this.posted.push({ message: message as Record<string, unknown>, transfer });
  }

  addEventListener(
    type: "message" | "error",
    listener: (event: { data?: unknown; message?: string }) => void,
  ): void {
    this.listeners[type].push(listener);
  }

  emit(data: unknown): void {
    for (const listener of this.listeners.message) listener({ data });
  }

  emitError(message: string): void {
    for (const listener of this.listeners.error) listener({ message });
  }

  lastRequestId(): string {
    return this.posted.at(-1)!.message.requestId as string;
  }
}

const CHUNKS: TranscriptChunk[] = [{ text: " hi", timestamp: [0, 1] }];

describe("BrowserEngine", () => {
  test("posts a transcribe message with a requestId and transfers the buffer", () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    const audio = new Float32Array([0.1, 0.2]);
    void engine.transcribe(audio, { language: "da", task: "translate" });

    expect(worker.posted).toHaveLength(1);
    const { message, transfer } = worker.posted[0]!;
    expect(message.type).toBe("transcribe");
    expect(message.language).toBe("da");
    expect(message.task).toBe("translate");
    expect(typeof message.requestId).toBe("string");
    expect(transfer).toEqual([audio.buffer as ArrayBuffer]);
  });

  test('omits language "auto" and the modelId (worker holds the model)', () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    void engine.transcribe(new Float32Array(4), {
      language: "auto",
      modelId: "onnx-community/whisper-tiny.en",
    });
    const { message } = worker.posted[0]!;
    expect("language" in message).toBe(false);
    expect("modelId" in message).toBe(false);
  });

  test("forwards the vocabulary hint as `prompt`, and omits an empty one (D-WS2a)", () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    void engine.transcribe(new Float32Array(4), { prompt: "Maxwell\nVoid" });
    void engine.transcribe(new Float32Array(4), { prompt: "" });
    void engine.transcribe(new Float32Array(4));
    expect(worker.posted[0]!.message.prompt).toBe("Maxwell\nVoid");
    expect("prompt" in worker.posted[1]!.message).toBe(false);
    expect("prompt" in worker.posted[2]!.message).toBe(false);
  });

  test("forwards the long-form request only when set (D-WS3a)", () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    void engine.transcribe(new Float32Array(4), { longForm: true });
    void engine.transcribe(new Float32Array(4));
    expect(worker.posted[0]!.message.longForm).toBe(true);
    expect("longForm" in worker.posted[1]!.message).toBe(false);
  });

  test("asks the worker to stream only when a consumer listens", () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    void engine.transcribe(new Float32Array(4));
    void engine.transcribe(new Float32Array(4), { onPartial: () => {} });
    void engine.transcribe(new Float32Array(4), { onProgress: () => {} });
    expect("stream" in worker.posted[0]!.message).toBe(false);
    expect(worker.posted[1]!.message.stream).toBe(true);
    expect(worker.posted[2]!.message.stream).toBe(true);
  });

  test("routes transcribe-progress to onProgress, ignoring other ids", async () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    const updates: { processedSeconds: number; totalSeconds: number }[] = [];
    const promise = engine.transcribe(new Float32Array(4), {
      onProgress: (progress) => updates.push(progress),
    });
    const requestId = worker.lastRequestId();

    worker.emit({ type: "transcribe-progress", processedSeconds: 5, totalSeconds: 60 });
    worker.emit({
      type: "transcribe-progress",
      processedSeconds: 7,
      totalSeconds: 60,
      requestId: "other",
    });
    worker.emit({ type: "transcribe-progress", processedSeconds: 0, totalSeconds: 90, requestId });
    worker.emit({ type: "transcribe-progress", processedSeconds: 30, totalSeconds: 90, requestId });
    worker.emit({ type: "result", text: "done", chunks: [], requestId });

    expect((await promise).text).toBe("done");
    expect(updates).toEqual([
      { processedSeconds: 0, totalSeconds: 90 },
      { processedSeconds: 30, totalSeconds: 90 },
    ]);
  });

  test("resolves on the matching result and routes partials to onPartial", async () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    const fragments: string[] = [];
    const promise = engine.transcribe(new Float32Array(4), {
      onPartial: (fragment) => fragments.push(fragment),
    });
    const requestId = worker.lastRequestId();

    // Noise the engine must ignore: other ids, id-less file-flow messages,
    // load progress.
    worker.emit({ type: "partial", text: "IGNORED", requestId: "other" });
    worker.emit({ type: "partial", text: "IGNORED" });
    worker.emit({ type: "load-progress", status: "ready", device: "wasm" });

    worker.emit({ type: "partial", text: "Hel", requestId });
    worker.emit({ type: "partial", text: "lo", requestId });
    worker.emit({ type: "result", text: "Hello", chunks: CHUNKS, requestId });
    worker.emit({ type: "result", text: "STALE", chunks: [], requestId: "other" });

    expect(await promise).toEqual({ text: "Hello", chunks: CHUNKS });
    expect(fragments).toEqual(["Hel", "lo"]);
  });

  test("rejects on a matching error message", async () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    const promise = engine.transcribe(new Float32Array(4));
    worker.emit({
      type: "error",
      stage: "transcribe",
      message: "out of memory",
      requestId: worker.lastRequestId(),
    });
    expect(promise).rejects.toThrow("out of memory");
  });

  test("a worker crash rejects every in-flight request", async () => {
    const worker = new FakeWorker();
    const engine = new BrowserEngine(worker);
    const first = engine.transcribe(new Float32Array(4));
    const second = engine.transcribe(new Float32Array(4));
    worker.emitError("worker exploded");
    expect(first).rejects.toThrow("worker exploded");
    expect(second).rejects.toThrow("worker exploded");
  });

  test("two engines can share one worker without stealing responses", async () => {
    const worker = new FakeWorker();
    const engineA = new BrowserEngine(worker);
    const engineB = new BrowserEngine(worker);
    const promiseA = engineA.transcribe(new Float32Array(4));
    const promiseB = engineB.transcribe(new Float32Array(4));
    const [idA, idB] = worker.posted.map(
      (post) => post.message.requestId as string,
    );
    expect(idA).not.toBe(idB);
    worker.emit({ type: "result", text: "for A", chunks: [], requestId: idA });
    worker.emit({ type: "result", text: "for B", chunks: [], requestId: idB });
    expect((await promiseA).text).toBe("for A");
    expect((await promiseB).text).toBe("for B");
  });
});

// ---------------------------------------------------------------------------
// ServerEngine — ndjson stream over an injected fetch
// ---------------------------------------------------------------------------

function ndjsonResponse(lines: unknown[], status = 200): Response {
  const body = lines.map((line) => `${JSON.stringify(line)}\n`).join("");
  return new Response(body, {
    status,
    headers: { "Content-Type": "application/x-ndjson" },
  });
}

describe("ServerEngine", () => {
  test("streams segments to onPartial and resolves with the result line", async () => {
    const calls: { url: string; init?: RequestInit }[] = [];
    const fetchFn: FetchLike = async (url, init) => {
      calls.push({ url, init });
      return ndjsonResponse([
        { type: "segment", text: " Hello", start: 0, end: 1.5 },
        { type: "segment", text: " world.", start: 1.5, end: 3 },
        { type: "result", text: "Hello world.", chunks: CHUNKS },
      ]);
    };
    const engine = new ServerEngine(fetchFn);
    const fragments: string[] = [];
    const audio = new Float32Array([0.5, -0.5, 0.25]);
    const result = await engine.transcribe(audio, {
      modelId: "onnx-community/whisper-tiny.en",
      language: "en",
      task: "transcribe",
      onPartial: (fragment) => fragments.push(fragment),
    });

    expect(result).toEqual({ text: "Hello world.", chunks: CHUNKS });
    expect(fragments).toEqual([" Hello", " world."]);

    const call = calls[0]!;
    const url = new URL(call.url, "http://localhost");
    expect(url.pathname).toBe("/api/transcribe");
    expect(url.searchParams.get("model")).toBe("onnx-community/whisper-tiny.en");
    expect(url.searchParams.get("language")).toBe("en");
    expect(url.searchParams.get("task")).toBe("transcribe");
    expect(call.init?.method).toBe("POST");
    expect(call.init?.credentials).toBe("same-origin");
    // The body must be the raw PCM bytes of the Float32Array.
    const sent = new Float32Array((call.init?.body as Uint8Array).buffer.slice(0));
    expect([...sent]).toEqual([...audio]);
  });

  test('drops language "auto" from the query', async () => {
    let url = "";
    const engine = new ServerEngine(async (input) => {
      url = input;
      return ndjsonResponse([{ type: "result", text: "", chunks: [] }]);
    });
    await engine.transcribe(new Float32Array(4), {
      modelId: "onnx-community/whisper-tiny",
      language: "auto",
    });
    expect(url).not.toContain("language");
  });

  test("sends the vocabulary hint as the `prompt` query param, only when non-empty", async () => {
    const urls: string[] = [];
    const engine = new ServerEngine(async (input) => {
      urls.push(input);
      return ndjsonResponse([{ type: "result", text: "", chunks: [] }]);
    });
    await engine.transcribe(new Float32Array(4), {
      modelId: "onnx-community/whisper-tiny",
      prompt: "Århus, Voices of the Void",
    });
    await engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny", prompt: "" });
    expect(new URL(urls[0]!, "http://localhost").searchParams.get("prompt")).toBe(
      "Århus, Voices of the Void",
    );
    expect(urls[1]).not.toContain("prompt");
  });

  test("requires a modelId", () => {
    const engine = new ServerEngine(async () => ndjsonResponse([]));
    expect(engine.transcribe(new Float32Array(4))).rejects.toThrow(/model id/);
  });

  test("rejects with the in-stream error line's message", () => {
    const engine = new ServerEngine(async () =>
      ndjsonResponse([
        { type: "segment", text: " partial", start: 0, end: 1 },
        { type: "error", message: "sidecar exploded" },
      ]),
    );
    expect(
      engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny" }),
    ).rejects.toThrow("sidecar exploded");
  });

  test("maps HTTP statuses to UI-safe messages", async () => {
    const cases: [number, RegExp][] = [
      [401, /sign in/i],
      [404, /not available/i],
      [413, /too large|limit/i],
      [503, /unavailable/i],
    ];
    for (const [status, pattern] of cases) {
      const engine = new ServerEngine(async () => new Response("{}", { status }));
      expect(
        engine.transcribe(new Float32Array(4), {
          modelId: "onnx-community/whisper-tiny",
        }),
      ).rejects.toThrow(pattern);
    }
  });

  test("surfaces the server's {error} detail on 503", () => {
    const engine = new ServerEngine(async () =>
      new Response(JSON.stringify({ error: "faster-whisper is not installed" }), {
        status: 503,
      }),
    );
    expect(
      engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny" }),
    ).rejects.toThrow(/faster-whisper is not installed/);
  });

  test("a stream that ends without a result line is a lost connection", () => {
    const engine = new ServerEngine(async () =>
      ndjsonResponse([{ type: "segment", text: " cut", start: 0, end: 1 }]),
    );
    expect(
      engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny" }),
    ).rejects.toThrow(/connection was lost/);
  });

  test("network failure gets a friendly message", () => {
    const engine = new ServerEngine(async () => {
      throw new TypeError("fetch failed");
    });
    expect(
      engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny" }),
    ).rejects.toThrow(/could not reach the server/i);
  });

  test("passes the abort signal to fetch and rethrows AbortError as-is", async () => {
    // Cancel (D-018 amendment): the UI aborts the controller; the rejection
    // must stay an AbortError, not be rewrapped as a network failure.
    let seenSignal: AbortSignal | undefined;
    const engine = new ServerEngine(async (_url, init) => {
      seenSignal = init?.signal ?? undefined;
      throw new DOMException("The user aborted a request.", "AbortError");
    });
    const controller = new AbortController();
    const promise = engine.transcribe(new Float32Array(4), {
      modelId: "onnx-community/whisper-tiny",
      signal: controller.signal,
    });
    expect(promise).rejects.toMatchObject({ name: "AbortError" });
    await promise.catch(() => {});
    expect(seenSignal).toBe(controller.signal);
  });

  test("an abort mid-stream surfaces as the stream's rejection", () => {
    // Once the 200 arrived, an abort tears down the body stream; whatever
    // error the reader throws must propagate out of transcribe().
    const body = new ReadableStream<Uint8Array>({
      start(streamController) {
        streamController.enqueue(
          new TextEncoder().encode(
            `${JSON.stringify({ type: "segment", text: " a", start: 0, end: 1 })}\n`,
          ),
        );
        streamController.error(
          new DOMException("The user aborted a request.", "AbortError"),
        );
      },
    });
    const engine = new ServerEngine(async () => new Response(body, { status: 200 }));
    expect(
      engine.transcribe(new Float32Array(4), { modelId: "onnx-community/whisper-tiny" }),
    ).rejects.toMatchObject({ name: "AbortError" });
  });
});

// ---------------------------------------------------------------------------
// Wire-format guard
// ---------------------------------------------------------------------------

describe("isTranscribeStreamLine", () => {
  test("accepts the three line shapes", () => {
    expect(
      isTranscribeStreamLine({ type: "segment", text: " a", start: 0, end: 1 }),
    ).toBe(true);
    expect(isTranscribeStreamLine({ type: "result", text: "a", chunks: [] })).toBe(true);
    expect(isTranscribeStreamLine({ type: "error", message: "x" })).toBe(true);
  });

  test("rejects malformed lines", () => {
    expect(isTranscribeStreamLine(null)).toBe(false);
    expect(isTranscribeStreamLine({ type: "segment", text: " a" })).toBe(false);
    expect(isTranscribeStreamLine({ type: "result", text: "a" })).toBe(false);
    expect(isTranscribeStreamLine({ type: "unknown" })).toBe(false);
  });
});

// ---------------------------------------------------------------------------
// Availability probe
// ---------------------------------------------------------------------------

describe("fetchServerEngineStatus", () => {
  test("passes through an available status", async () => {
    const status = await fetchServerEngineStatus(async () =>
      new Response(JSON.stringify({ available: true }), { status: 200 }),
    );
    expect(status).toEqual({ available: true });
  });

  test("passes through unavailable + reason", async () => {
    const status = await fetchServerEngineStatus(async () =>
      new Response(
        JSON.stringify({ available: false, reason: "sidecar missing" }),
        { status: 200 },
      ),
    );
    expect(status).toEqual({ available: false, reason: "sidecar missing" });
  });

  test("404 (deployment without the feature) resolves unavailable", async () => {
    const status = await fetchServerEngineStatus(async () =>
      new Response("Not Found", { status: 404 }),
    );
    expect(status.available).toBe(false);
  });

  test("never rejects on network failure", async () => {
    const status = await fetchServerEngineStatus(async () => {
      throw new TypeError("offline");
    });
    expect(status.available).toBe(false);
  });
});

// ---------------------------------------------------------------------------
// Engine choice persistence
// ---------------------------------------------------------------------------

function fakeStorage(initial: Record<string, string> = {}): StorageLike & {
  data: Record<string, string>;
} {
  const data = { ...initial };
  return {
    data,
    getItem: (key) => (key in data ? data[key]! : null),
    setItem: (key, value) => {
      data[key] = value;
    },
  };
}

describe("engine choice persistence", () => {
  test("round-trips through storage", () => {
    const storage = fakeStorage();
    saveEngineChoice("server", storage);
    expect(storage.data[ENGINE_CHOICE_KEY]).toBe("server");
    expect(loadEngineChoice(storage)).toBe("server");
  });

  test('defaults to "browser" for missing or garbage values', () => {
    expect(loadEngineChoice(fakeStorage())).toBe("browser");
    expect(
      loadEngineChoice(fakeStorage({ [ENGINE_CHOICE_KEY]: "cloud" })),
    ).toBe("browser");
    expect(loadEngineChoice(null)).toBe("browser");
  });

  test("throwing storage is survived", () => {
    const storage: StorageLike = {
      getItem: () => {
        throw new Error("denied");
      },
      setItem: () => {
        throw new Error("denied");
      },
    };
    expect(loadEngineChoice(storage)).toBe("browser");
    expect(() => saveEngineChoice("server", storage)).not.toThrow();
  });
});