whisper_web/src/worker/protocol.ts

/**
 * Typed postMessage protocol between the UI and the inference Web Worker
 * (Decisions.md D-003). Shared by `src/worker/inference.worker.ts` and the UI.
 *
 * Everything in this file must stay free of browser/worker globals so it can
 * be unit-tested with `bun test`.
 */

/** Devices the worker can actually run on (D-004). */
export type InferenceDevice = "webgpu" | "wasm" | "webnn";

/**
 * What the UI may request. "auto" (the default) lets the worker probe the
 * available backends and pick in priority order WebNN → WebGPU → WASM
 * (D-004/D-017), with automatic fallback down the same chain if pipeline
 * creation fails on the picked device.
 */
export type DevicePreference = InferenceDevice | "auto";

/** Quantization levels used by the curated catalog (D-006). */
export type WhisperDtype = "fp32" | "fp16" | "q8" | "q4";

/**
 * Either one dtype for the whole model, or per-part dtypes — the shape
 * Transformers.js accepts. Whisper on WebGPU commonly needs a mixed setup
 * (e.g. fp16 encoder + q4 merged decoder), so per-part must be expressible.
 */
export type DtypeSpec =
  | WhisperDtype
  | { encoder_model: WhisperDtype; decoder_model_merged: WhisperDtype };

/** Whisper generation task (multilingual models only). */
export type WhisperTask = "transcribe" | "translate";

// ---------------------------------------------------------------------------
// UI → worker
// ---------------------------------------------------------------------------

export interface LoadMessage {
  type: "load";
  /** Hugging Face Hub model id, e.g. "onnx-community/whisper-tiny.en". */
  modelId: string;
  /** Device preference; the worker resolves "auto" and reports the outcome. */
  device: DevicePreference;
  /**
   * Optional dtype override. When omitted the worker resolves the catalog
   * default for the device it actually selected (src/lib/models.ts).
   */
  dtype?: DtypeSpec;
  /**
   * Load the encoder/decoder ONNX sessions one at a time instead of
   * concurrently (D-027/D-029). Sequential halves the peak JS-heap/WASM
   * memory a load needs (avoiding mobile tab crashes on larger models) at
   * the cost of losing the CDN-concurrency benefit a fast connection would
   * otherwise get from two simultaneous downloads — worth paying only where
   * the memory-crash risk is real. The UI computes this (`isMobileLoadRisk`
   * in App.tsx); omitted/false loads in parallel.
   */
  sequentialSessionLoad?: boolean;
}

export interface TranscribeMessage {
  type: "transcribe";
  /**
   * Mono PCM at 16 kHz (see src/lib/audio.ts). The underlying ArrayBuffer
   * should be passed as a transferable for zero-copy handoff.
   */
  audio: Float32Array;
  /**
   * ISO-ish Whisper language code (e.g. "en", "da"). Omit (or "auto") to let
   * Whisper detect the language. Only meaningful for multilingual models.
   */
  language?: string;
  /** Only meaningful for multilingual models. Defaults to "transcribe". */
  task?: WhisperTask;
  /**
   * Optional correlation id, echoed on the partial/result/error messages this
   * request produces (D-016). The live-transcription loop uses it to match
   * responses to segment passes; the file flow omits it and is unchanged.
   */
  requestId?: string;
  /**
   * Stream partial tokens and transcribe-progress for this request. Defaults
   * to `requestId === undefined` (the historical rule: id-less file-flow
   * passes stream, live segment passes resolve wholesale via "result").
   * `BrowserEngine` sets it explicitly so a file pass routed through the
   * engine (which correlates via requestId, D-010) still streams.
   */
  stream?: boolean;
}

export type UiToWorkerMessage = LoadMessage | TranscribeMessage;

// ---------------------------------------------------------------------------
// Worker → UI
// ---------------------------------------------------------------------------

/**
 * Model download / initialization progress. Mirrors the Transformers.js
 * `progress_callback` events we care about; `status: "ready"` marks the end
 * of a successful load and reports the backend that actually got selected
 * (after device detection / fallback, D-004, D-017).
 */
export type LoadProgressMessage =
  | { type: "load-progress"; status: "initiate"; file: string }
  | { type: "load-progress"; status: "download"; file: string }
  | {
      type: "load-progress";
      status: "progress";
      file: string;
      /** Bytes downloaded so far. */
      loaded: number;
      /** Total bytes for this file. */
      total: number;
      /** Percentage in [0, 100]. */
      progress: number;
    }
  | { type: "load-progress"; status: "done"; file: string }
  | {
      type: "load-progress";
      status: "ready";
      device: InferenceDevice;
      /**
       * WASM thread count onnxruntime actually ended up with (D-025) — ort
       * clamps to 1 when cross-origin isolation is missing, so this is field
       * observability for "is multithreading really on?". Accelerator devices
       * report 1 (their WASM side must be single-threaded, D-023).
       */
      threads?: number;
    };

/** Incremental transcript fragment — append to the text streamed so far. */
export interface PartialMessage {
  type: "partial";
  text: string;
  /** Echo of the originating transcribe request's id, if it carried one. */
  requestId?: string;
}

/**
 * How far a file-flow pass has progressed through the audio, in seconds.
 * Posted once at pass start (0 / total) and then on every Whisper segment
 * boundary, so slow models/backends show movement long before the first
 * partial token. Values are monotone and clamped to [0, totalSeconds]
 * (see TranscribeProgressTracker). Live passes (requestId set) don't report
 * progress — the live UI has its own committed/tentative feedback (D-016).
 */
export interface TranscribeProgressMessage {
  type: "transcribe-progress";
  processedSeconds: number;
  totalSeconds: number;
  /** Echo of the originating transcribe request's id, if it carried one. */
  requestId?: string;
}

/** One timestamped segment of the final transcript. */
export interface TranscriptChunk {
  text: string;
  /** [start, end] in seconds; end may be null if Whisper never closed the segment. */
  timestamp: [number, number | null];
}

export interface ResultMessage {
  type: "result";
  text: string;
  chunks: TranscriptChunk[];
  /** Echo of the originating transcribe request's id, if it carried one. */
  requestId?: string;
}

export interface ErrorMessage {
  type: "error";
  stage: "load" | "transcribe";
  message: string;
  /** Echo of the failing transcribe request's id, if it carried one. */
  requestId?: string;
}

/**
 * The worker cannot finish this load usefully in-place and asks the UI to
 * respawn it (D-025). Posted when an accelerator attempt already initialized
 * onnxruntime's WASM runtime — whose thread pool is fixed at first init —
 * single-threaded, so the WASM fallback in *this* worker would run on one
 * thread for the worker's whole life. The UI terminates the worker, spawns a
 * fresh one, and re-issues the load with the explicit `device`, remembering
 * `modelId`+`preference` so later loads skip the doomed attempt entirely.
 */
export interface LoadRestartMessage {
  type: "load-restart";
  /** Echo of the aborted load's model id. */
  modelId: string;
  /** Echo of the aborted load's device preference. */
  preference: DevicePreference;
  /** Explicit device the fresh worker should load with. */
  device: InferenceDevice;
}

export type WorkerToUiMessage =
  | LoadProgressMessage
  | PartialMessage
  | TranscribeProgressMessage
  | ResultMessage
  | ErrorMessage
  | LoadRestartMessage;

// ---------------------------------------------------------------------------
// Device selection policy (D-004, D-017) — pure so it is unit-testable; the
// worker probes the browser capabilities and this decides the attempt order.
// ---------------------------------------------------------------------------

/** What the worker's capability probe found (see inference.worker.ts). */
export interface DeviceCapabilities {
  /** `navigator.ml` exists (WebNN). */
  webnn: boolean;
  /** `navigator.gpu.requestAdapter()` returned an adapter (WebGPU). */
  webgpu: boolean;
  /**
   * The WebGPU adapter supports the "shader-f16" feature. Firefox ships
   * WebGPU without it, so fp16 weights cannot run there (D-025); the worker
   * skips the WebGPU attempt when the resolved dtype needs fp16 and this is
   * false. Optional so `deviceAttempts` callers that don't probe it (tests)
   * stay valid.
   */
  webgpuF16?: boolean;
}

/**
 * Whether a dtype spec involves fp16 weights anywhere — such a model cannot
 * run on a WebGPU adapter without the "shader-f16" feature (D-025).
 */
export function dtypeRequiresF16(dtype: DtypeSpec): boolean {
  if (typeof dtype === "string") return dtype === "fp16";
  return dtype.encoder_model === "fp16" || dtype.decoder_model_merged === "fp16";
}

/**
 * Resolve a device preference against probed capabilities into the ordered
 * list of devices to attempt pipeline creation on. The first entry is the
 * selected device; later entries are automatic fallbacks for when creation
 * throws (driver quirks, unsupported ops, out-of-memory). Never empty — WASM
 * is always the last resort.
 *
 * - "auto": WebNN → WebGPU → WASM (D-017), skipping backends the probe ruled
 *   out, so an auto-detected WebNN that fails still tries WebGPU before WASM.
 * - Explicit "webnn"/"webgpu": that device (if available) → WASM. An explicit
 *   choice deliberately does not detour through the other accelerator.
 * - "wasm": WASM only.
 */
export function deviceAttempts(
  preference: DevicePreference,
  capabilities: DeviceCapabilities,
): InferenceDevice[] {
  switch (preference) {
    case "wasm":
      return ["wasm"];
    case "webgpu":
      return capabilities.webgpu ? ["webgpu", "wasm"] : ["wasm"];
    case "webnn":
      return capabilities.webnn ? ["webnn", "wasm"] : ["wasm"];
    case "auto": {
      const attempts: InferenceDevice[] = [];
      if (capabilities.webnn) attempts.push("webnn");
      if (capabilities.webgpu) attempts.push("webgpu");
      attempts.push("wasm");
      return attempts;
    }
  }
}

// ---------------------------------------------------------------------------
// Runtime guards (postMessage data arrives untyped)
// ---------------------------------------------------------------------------

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null;
}

function isOptionalString(value: unknown): value is string | undefined {
  return value === undefined || typeof value === "string";
}

const DEVICES: readonly string[] = ["webgpu", "wasm", "webnn"];
const DTYPES: readonly string[] = ["fp32", "fp16", "q8", "q4"];

export function isInferenceDevice(value: unknown): value is InferenceDevice {
  return typeof value === "string" && DEVICES.includes(value);
}

export function isDevicePreference(value: unknown): value is DevicePreference {
  return value === "auto" || isInferenceDevice(value);
}

export function isDtypeSpec(value: unknown): value is DtypeSpec {
  if (typeof value === "string") return DTYPES.includes(value);
  return (
    isRecord(value) &&
    typeof value.encoder_model === "string" &&
    DTYPES.includes(value.encoder_model) &&
    typeof value.decoder_model_merged === "string" &&
    DTYPES.includes(value.decoder_model_merged)
  );
}

export function isUiToWorkerMessage(value: unknown): value is UiToWorkerMessage {
  if (!isRecord(value)) return false;
  switch (value.type) {
    case "load":
      return (
        typeof value.modelId === "string" &&
        isDevicePreference(value.device) &&
        (value.dtype === undefined || isDtypeSpec(value.dtype))
      );
    case "transcribe":
      return (
        value.audio instanceof Float32Array &&
        isOptionalString(value.language) &&
        (value.task === undefined ||
          value.task === "transcribe" ||
          value.task === "translate") &&
        isOptionalString(value.requestId) &&
        (value.stream === undefined || typeof value.stream === "boolean")
      );
    default:
      return false;
  }
}

export function isWorkerToUiMessage(value: unknown): value is WorkerToUiMessage {
  if (!isRecord(value)) return false;
  switch (value.type) {
    case "load-progress":
      switch (value.status) {
        case "initiate":
        case "download":
        case "done":
          return typeof value.file === "string";
        case "progress":
          return (
            typeof value.file === "string" &&
            typeof value.loaded === "number" &&
            typeof value.total === "number" &&
            typeof value.progress === "number"
          );
        case "ready":
          return (
            isInferenceDevice(value.device) &&
            (value.threads === undefined || typeof value.threads === "number")
          );
        default:
          return false;
      }
    case "partial":
      return typeof value.text === "string" && isOptionalString(value.requestId);
    case "transcribe-progress":
      return (
        typeof value.processedSeconds === "number" &&
        typeof value.totalSeconds === "number" &&
        isOptionalString(value.requestId)
      );
    case "result":
      return (
        typeof value.text === "string" &&
        Array.isArray(value.chunks) &&
        isOptionalString(value.requestId)
      );
    case "error":
      return (
        typeof value.message === "string" &&
        (value.stage === "load" || value.stage === "transcribe") &&
        isOptionalString(value.requestId)
      );
    case "load-restart":
      return (
        typeof value.modelId === "string" &&
        isDevicePreference(value.preference) &&
        isInferenceDevice(value.device)
      );
    default:
      return false;
  }
}

// ---------------------------------------------------------------------------
// Load-progress accumulation (pure; drives the download progress UI)
// ---------------------------------------------------------------------------

export interface FileProgress {
  file: string;
  loaded: number;
  total: number;
  done: boolean;
}

export interface LoadProgressState {
  /** Per-file progress, keyed by file name, in first-seen order. */
  files: Record<string, FileProgress>;
  /** True once the pipeline reported ready. */
  ready: boolean;
  /** Backend the worker actually selected; null until ready (D-004). */
  device: InferenceDevice | null;
}

export const initialLoadProgress: LoadProgressState = {
  files: {},
  ready: false,
  device: null,
};

/** Fold one load-progress message into the accumulated state. Pure. */
export function applyLoadProgress(
  state: LoadProgressState,
  message: LoadProgressMessage,
): LoadProgressState {
  if (message.status === "ready") {
    return { ...state, ready: true, device: message.device };
  }
  const previous = state.files[message.file];
  const next: FileProgress = {
    file: message.file,
    loaded: previous?.loaded ?? 0,
    total: previous?.total ?? 0,
    done: previous?.done ?? false,
  };
  if (message.status === "progress") {
    next.loaded = message.loaded;
    next.total = message.total;
  } else if (message.status === "done") {
    next.done = true;
    if (next.total > 0) next.loaded = next.total;
  }
  return { ...state, files: { ...state.files, [message.file]: next } };
}

// ---------------------------------------------------------------------------
// Transcription progress accumulation (pure; drives the in-progress UI)
// ---------------------------------------------------------------------------

/**
 * Maps the streamer's window-relative timestamps to absolute audio progress.
 *
 * The chunked ASR pipeline slices audio into windows of `windowSeconds`,
 * advancing by `jumpSeconds` (window − 2·stride) per window, and runs one
 * generate pass per window. Whisper's timestamp tokens — surfaced via
 * `WhisperTextStreamer`'s `on_chunk_start`/`on_chunk_end` — are relative to
 * the current window's start, and `on_finalize` fires when a window's pass
 * ends. This tracker turns those into absolute seconds that are monotone
 * (windows overlap, so a later window's early timestamps can precede an
 * earlier window's late ones — regressions are swallowed) and clamped to
 * [0, totalSeconds].
 */
export class TranscribeProgressTracker {
  private windowIndex = 0;
  private reported = 0;

  constructor(
    private readonly totalSeconds: number,
    private readonly windowSeconds: number,
    private readonly jumpSeconds: number,
  ) {}

  /**
   * Absolute progress for a window-relative timestamp, or null when it does
   * not advance past what was already reported.
   */
  timestamp(withinWindowSeconds: number): number | null {
    const clamped = Math.min(Math.max(withinWindowSeconds, 0), this.windowSeconds);
    return this.advance(this.windowIndex * this.jumpSeconds + clamped);
  }

  /**
   * The current window's generate pass finished; progress moves to the start
   * of the next window (the amount of audio no later window revisits).
   */
  finalizeWindow(): number | null {
    this.windowIndex += 1;
    return this.advance(this.windowIndex * this.jumpSeconds);
  }

  private advance(processedSeconds: number): number | null {
    const clamped = Math.min(processedSeconds, this.totalSeconds);
    if (clamped <= this.reported) return null;
    this.reported = clamped;
    return clamped;
  }
}