client/modules/speechToText.test.ts

import {
  DictationError,
  type DictationSession,
  getDictationEngine,
  resampleTo16k,
  setWorkerFactory,
  startDictation,
} from "@/modules/speechToText";
import type { WorkerResponse } from "@/modules/speechToTextWorkerProtocol";

class FakeWorker {
  static instances: FakeWorker[] = [];
  onmessage: ((event: MessageEvent) => void) | null = null;
  onerror: ((event: unknown) => void) | null = null;
  terminated = false;
  messages: unknown[] = [];

  constructor() {
    FakeWorker.instances.push(this);
  }

  postMessage(message: unknown, _transfer?: Transferable[]) {
    this.messages.push(message);
  }

  terminate() {
    this.terminated = true;
  }

  respond(response: WorkerResponse) {
    this.onmessage?.({ data: response } as MessageEvent);
  }
}

class FakeMediaStreamTrack {
  stopped = false;
  stop() {
    this.stopped = true;
  }
}

function makeMediaStream() {
  const tracks = [new FakeMediaStreamTrack(), new FakeMediaStreamTrack()];
  return {
    tracks,
    stream: {
      getTracks: () => tracks,
    } as unknown as MediaStream,
  };
}

function stubWasmSupport(getUserMedia: ReturnType<typeof vi.fn>) {
  vi.stubGlobal("Worker", FakeWorker);
  vi.stubGlobal("WebAssembly", { instantiate: vi.fn() });
  vi.stubGlobal(
    "AudioContext",
    class {
      sampleRate = 48000;
      createMediaStreamSource() {
        return { connect: vi.fn(), disconnect: vi.fn() };
      }
      createScriptProcessor() {
        return {
          onaudioprocess: null,
          connect: vi.fn(),
          disconnect: vi.fn(),
        };
      }
      close() {
        return Promise.resolve();
      }
    },
  );
  Object.defineProperty(navigator, "mediaDevices", {
    configurable: true,
    value: { getUserMedia },
  });
}

const originalMediaDevices = navigator.mediaDevices;

afterEach(() => {
  vi.unstubAllGlobals();
  Object.defineProperty(navigator, "mediaDevices", {
    configurable: true,
    value: originalMediaDevices,
  });
  delete (window as unknown as Record<string, unknown>).webkitSpeechRecognition;
  setWorkerFactory(() => {
    throw new Error("no worker factory configured");
  });
  FakeWorker.instances = [];
});

describe("getDictationEngine", () => {
  function installRecognitionFake() {
    class FakeRecognition {
      start = vi.fn();
    }
    (window as unknown as Record<string, unknown>).webkitSpeechRecognition =
      FakeRecognition;
  }

  it("reports the wasm engine when it is preferred and every piece is present", () => {
    stubWasmSupport(vi.fn());
    expect(getDictationEngine(true)).toBe("wasm");
  });

  it("falls back to web-speech when the wasm pieces are missing", () => {
    installRecognitionFake();
    expect(getDictationEngine(true)).toBe("web-speech");
  });

  it("prefers web-speech over a capable wasm path when the local model is off", () => {
    stubWasmSupport(vi.fn());
    installRecognitionFake();
    expect(getDictationEngine(false)).toBe("web-speech");
  });

  it("keeps the wasm engine when the local model is off but no recognizer exists", () => {
    // The preference reorders engines, it never empties the list: a browser
    // without `SpeechRecognition` (Firefox, for one) must still get a
    // working Dictate button with the setting off.
    stubWasmSupport(vi.fn());
    expect(getDictationEngine(false)).toBe("wasm");
  });

  it("reports no engine when neither path is available", () => {
    expect(getDictationEngine(true)).toBeNull();
    expect(getDictationEngine(false)).toBeNull();
  });
});

describe("resampleTo16k", () => {
  it("passes 16 kHz audio through untouched", () => {
    const input = new Float32Array([0.5, -0.5]);
    expect(resampleTo16k(input, 16000)).toBe(input);
  });

  it("resamples 48 kHz audio to a third of the length with interpolated values", () => {
    const input = new Float32Array([0, 1, 2, 3, 4, 5]);
    const output = resampleTo16k(input, 48000);
    expect(output.length).toBe(2);
    expect(output[0]).toBeCloseTo(0);
    expect(output[1]).toBeCloseTo(3);
  });
});

describe("startDictation with the wasm engine", () => {
  it("loads the model, streams the transcript and releases everything on stop", async () => {
    const { tracks, stream } = makeMediaStream();
    const getUserMedia = vi.fn().mockResolvedValue(stream);
    stubWasmSupport(getUserMedia);
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);

    const transcripts: string[] = [];
    const progress: [number, number | undefined][] = [];

    const sessionPromise = startDictation(
      {
        onTranscript: (text) => transcripts.push(text),
        onProgress: (loaded, total) => progress.push([loaded, total]),
      },
      true,
    );

    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    expect(worker.messages[0]).toMatchObject({ type: "load" });

    worker.respond({ type: "progress", loaded: 1000, total: 2000 });
    worker.respond({ type: "loaded" });

    const session: DictationSession = await sessionPromise;

    worker.respond({ type: "transcript", text: "hello" });
    worker.respond({ type: "transcript", text: "hello world" });

    expect(transcripts).toEqual(["hello", "hello world"]);
    expect(progress).toEqual([[1000, 2000]]);

    await session.stop();

    expect(tracks.every((track) => track.stopped)).toBe(true);
    expect(worker.terminated).toBe(true);
    expect(worker.messages).toContainEqual({ type: "stop" });
  });

  it("reports a permission error when the microphone is denied", async () => {
    const getUserMedia = vi
      .fn()
      .mockRejectedValue(
        new DOMException("Permission denied", "NotAllowedError"),
      );
    stubWasmSupport(getUserMedia);
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    FakeWorker.instances[0].respond({ type: "loaded" });

    await expect(sessionPromise).rejects.toMatchObject({ kind: "permission" });
    expect(getUserMedia).toHaveBeenCalledWith({ audio: true });
  });

  /** Minimal `SpeechRecognition` stand-in; reports whether it was started. */
  function installLocalRecognitionFake() {
    let started = false;
    class FakeRecognition {
      continuous = false;
      interimResults = false;
      lang = "";
      onresult: unknown = null;
      onerror: unknown = null;
      onend: unknown = null;
      start() {
        started = true;
      }
      stop() {}
    }
    (window as unknown as Record<string, unknown>).webkitSpeechRecognition =
      FakeRecognition;
    return () => started;
  }

  it("falls back to the browser recognizer when the local engine fails", async () => {
    const { stream } = makeMediaStream();
    stubWasmSupport(vi.fn().mockResolvedValue(stream));
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const started = installLocalRecognitionFake();

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    // A 502 from `/dictation-models/` is "cannot run here", not "no engine".
    FakeWorker.instances[0].respond({
      type: "error",
      message: "model download failed",
    });

    await sessionPromise;
    expect(started()).toBe(true);
  });

  it("reports a denied microphone instead of falling back", async () => {
    stubWasmSupport(
      vi
        .fn()
        .mockRejectedValue(
          new DOMException("Permission denied", "NotAllowedError"),
        ),
    );
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const started = installLocalRecognitionFake();

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    FakeWorker.instances[0].respond({ type: "loaded" });

    await expect(sessionPromise).rejects.toMatchObject({ kind: "permission" });
    // Falling back here would ask again and, in Chrome, send the audio away.
    expect(started()).toBe(false);
  });

  it("fails the load when the engine dies while permission is pending", async () => {
    const { tracks, stream } = makeMediaStream();
    let grantMicrophone: ((stream: MediaStream) => void) | undefined;
    stubWasmSupport(
      vi.fn().mockReturnValue(
        new Promise((resolve) => {
          grantMicrophone = resolve as (stream: MediaStream) => void;
        }),
      ),
    );
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const started = installLocalRecognitionFake();

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "loaded" });

    // The prompt can sit open for minutes, and the transcriber is already
    // running by now. Rejecting the settled load promise would be a no-op.
    worker.respond({ type: "error", message: "the transcriber died" });
    grantMicrophone?.(stream);

    await sessionPromise;
    expect(tracks.every((track) => track.stopped)).toBe(true);
    expect(worker.terminated).toBe(true);
    // Treated as any other engine failure, so the browser recognizer takes it.
    expect(started()).toBe(true);
  });

  it("fails the load when the worker dies uncaught while permission is pending", async () => {
    const { tracks, stream } = makeMediaStream();
    let grantMicrophone: ((stream: MediaStream) => void) | undefined;
    stubWasmSupport(
      vi.fn().mockReturnValue(
        new Promise((resolve) => {
          grantMicrophone = resolve as (stream: MediaStream) => void;
        }),
      ),
    );
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const started = installLocalRecognitionFake();

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "loaded" });

    // An uncaught error rather than a posted one: the worker's own try/catch
    // converts most failures, so this is the channel it cannot cover.
    worker.onerror?.(new ErrorEvent("error"));
    grantMicrophone?.(stream);

    await sessionPromise;
    expect(tracks.every((track) => track.stopped)).toBe(true);
    expect(worker.terminated).toBe(true);
    expect(started()).toBe(true);
  });

  it("ignores a transcript that arrives after stop", async () => {
    const { stream } = makeMediaStream();
    stubWasmSupport(vi.fn().mockResolvedValue(stream));
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const onTranscript = vi.fn();

    const sessionPromise = startDictation({ onTranscript }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "loaded" });
    const session = await sessionPromise;

    worker.respond({ type: "transcript", text: "hello world" });
    expect(onTranscript).toHaveBeenCalledWith("hello world");

    const stopping = session.stop();
    // Both engines emit one last result after being told to stop. The caller
    // has already forgotten what it appended, so this would be appended twice.
    worker.respond({ type: "transcript", text: "hello world" });
    worker.respond({ type: "stopped" });
    await stopping;

    expect(onTranscript).toHaveBeenCalledTimes(1);
  });

  it("waits for the worker to drain before terminating it", async () => {
    const { stream } = makeMediaStream();
    stubWasmSupport(vi.fn().mockResolvedValue(stream));
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "loaded" });
    const session = await sessionPromise;

    const stopping = session.stop();
    await Promise.resolve();
    // The transcriber runs synchronously inside the worker's `onmessage`, so a
    // backlog can leave it seconds behind; terminating now drops the tail.
    expect(worker.terminated).toBe(false);

    worker.respond({ type: "stopped" });
    await stopping;
    expect(worker.terminated).toBe(true);
  });

  it("reports a worker failure that happens after the engine loaded", async () => {
    const { stream } = makeMediaStream();
    stubWasmSupport(vi.fn().mockResolvedValue(stream));
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    const onError = vi.fn();

    const sessionPromise = startDictation(
      {
        onTranscript: vi.fn(),
        onError,
      },
      true,
    );
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "loaded" });
    await sessionPromise;

    // The load promise has settled, so rejecting it again is a no-op: without
    // a live channel the UI listens forever with the microphone open.
    worker.respond({ type: "error", message: "the transcriber died" });

    expect(onError).toHaveBeenCalledTimes(1);
    expect(onError.mock.calls[0][0]).toMatchObject({ kind: "engine" });
  });

  it("never opens the microphone when the model fails to load", async () => {
    const { tracks, stream } = makeMediaStream();
    const getUserMedia = vi.fn().mockResolvedValue(stream);
    stubWasmSupport(getUserMedia);
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);

    const sessionPromise = startDictation({ onTranscript: vi.fn() }, true);
    await vi.waitFor(() => expect(FakeWorker.instances.length).toBe(1));
    const worker = FakeWorker.instances[0];
    worker.respond({ type: "error", message: "model download failed" });

    await expect(sessionPromise).rejects.toBeInstanceOf(DictationError);
    await expect(
      sessionPromise.catch((error: DictationError) => error.kind),
    ).resolves.toBe("engine");
    // The load runs first, so a failed download never lights the recording
    // indicator at all.
    expect(getUserMedia).not.toHaveBeenCalled();
    expect(tracks.every((track) => track.stopped)).toBe(false);
    expect(worker.terminated).toBe(true);
  });
});

describe("startDictation with the local model preference off", () => {
  let recognitionStarted = false;

  function installRecognitionFake(
    startImpl: () => void = () => {
      recognitionStarted = true;
    },
  ) {
    class FakeRecognition {
      continuous = false;
      interimResults = false;
      lang = "";
      onresult: unknown = null;
      onerror: unknown = null;
      onend: unknown = null;
      start = startImpl;
      stop = () => {};
    }
    (window as unknown as Record<string, unknown>).webkitSpeechRecognition =
      FakeRecognition;
  }

  beforeEach(() => {
    recognitionStarted = false;
  });

  it("never constructs the model worker on a fully wasm-capable browser", async () => {
    // The whole point of the setting: with it off, the ~51 MB model must
    // not be downloaded even where it could run.
    const { stream } = makeMediaStream();
    stubWasmSupport(vi.fn().mockResolvedValue(stream));
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    installRecognitionFake();

    const session = await startDictation({ onTranscript: vi.fn() }, false);

    expect(FakeWorker.instances.length).toBe(0);
    expect(recognitionStarted).toBe(true);
    await session.stop();
  });

  it("does not fall back to the wasm engine when the recognizer fails", async () => {
    // A reverse fallback would start the ~51 MB download right after the
    // user turned that model off; the failure must surface instead.
    stubWasmSupport(vi.fn());
    setWorkerFactory(() => new FakeWorker() as unknown as Worker);
    installRecognitionFake(() => {
      throw new Error("recognizer refused to start");
    });

    await expect(
      startDictation({ onTranscript: vi.fn() }, false),
    ).rejects.toMatchObject({ kind: "unavailable" });
    expect(FakeWorker.instances.length).toBe(0);
  });
});

describe("startDictation with the web speech fallback", () => {
  interface RecognitionFake {
    continuous: boolean;
    interimResults: boolean;
    lang: string;
    onresult:
      | ((event: {
          resultIndex: number;
          results: ArrayLike<{ isFinal: boolean; 0: { transcript: string } }>;
        }) => void)
      | null;
    onerror: ((event: { error: string }) => void) | null;
    onend: (() => void) | null;
    start: () => void;
    stop: () => void;
    stopped: boolean;
  }

  let instance: RecognitionFake | undefined;

  function installRecognitionFake(startImpl: () => void = () => {}) {
    class FakeRecognition {
      continuous = false;
      interimResults = false;
      lang = "";
      onresult: RecognitionFake["onresult"] = null;
      onerror: RecognitionFake["onerror"] = null;
      onend: (() => void) | null = null;
      stopped = false;
      start: () => void;
      stop = () => {
        this.stopped = true;
      };
      constructor() {
        this.start = startImpl;
        instance = this as unknown as RecognitionFake;
      }
    }
    (window as unknown as Record<string, unknown>).webkitSpeechRecognition =
      FakeRecognition;
  }

  it("transcribes interim results and stops on request", async () => {
    installRecognitionFake();

    const transcripts: string[] = [];
    const session = await startDictation(
      {
        onTranscript: (text) => transcripts.push(text),
      },
      true,
    );

    expect(instance?.continuous).toBe(true);
    expect(instance?.interimResults).toBe(true);

    instance?.onresult?.({
      resultIndex: 0,
      results: [{ isFinal: false, 0: { transcript: "search for" } }],
    });
    instance?.onresult?.({
      resultIndex: 0,
      results: [
        { isFinal: true, 0: { transcript: "search for" } },
        { isFinal: false, 0: { transcript: " moonshine" } },
      ],
    });

    expect(transcripts).toEqual(["search for", "search for moonshine"]);

    await session.stop();
    expect(instance?.stopped).toBe(true);
  });

  it("reports a permission error raised before the session resolves", async () => {
    installRecognitionFake(() => {
      instance?.onerror?.({ error: "not-allowed" });
    });

    await expect(
      startDictation({ onTranscript: vi.fn() }, true),
    ).rejects.toMatchObject({ kind: "permission" });
  });

  it("reports a denial that arrives after the session resolved", async () => {
    // What a real browser does: `start()` returns, and the denial arrives on a
    // later task. Rejecting the already-resolved promise would be a no-op, so
    // without the `onError` route the user is told nothing at all.
    installRecognitionFake();
    const onError = vi.fn();

    const session = await startDictation(
      {
        onTranscript: vi.fn(),
        onError,
      },
      true,
    );

    instance?.onerror?.({ error: "not-allowed" });

    expect(onError).toHaveBeenCalledTimes(1);
    expect(onError.mock.calls[0][0]).toMatchObject({ kind: "permission" });
    await session.stop();
  });

  it("ends the session quietly when the recognizer stops on its own", async () => {
    installRecognitionFake();
    const onError = vi.fn();
    const onEnd = vi.fn();

    const session = await startDictation(
      {
        onTranscript: vi.fn(),
        onEnd,
        onError,
      },
      true,
    );

    // Chrome ends after silence even with `continuous`. The session is over,
    // but there is nothing the user needs told about.
    instance?.onend?.();

    expect(onEnd).toHaveBeenCalledTimes(1);
    expect(onError).not.toHaveBeenCalled();
    await session.stop();
  });

  it("says nothing when the recognizer ends because it was stopped", async () => {
    installRecognitionFake();
    const onError = vi.fn();
    const onEnd = vi.fn();

    const session = await startDictation(
      {
        onTranscript: vi.fn(),
        onEnd,
        onError,
      },
      true,
    );

    await session.stop();
    instance?.onend?.();

    expect(onError).not.toHaveBeenCalled();
    expect(onEnd).not.toHaveBeenCalled();
  });
});

describe("startDictation without any engine", () => {
  it("rejects with an unavailable error", async () => {
    await expect(
      startDictation({ onTranscript: vi.fn() }, true),
    ).rejects.toMatchObject({ kind: "unavailable" });
  });
});