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";

/**
 * What the UI may request. "auto" (the default) lets the worker probe
 * `navigator.gpu` and pick WebGPU with automatic WASM fallback (D-004).
 */
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;
}

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;
}

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 WebGPU detection / fallback, D-004).
 */
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 };

/** Incremental transcript fragment — append to the text streamed so far. */
export interface PartialMessage {
  type: "partial";
  text: 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[];
}

export interface ErrorMessage {
  type: "error";
  stage: "load" | "transcribe";
  message: string;
}

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

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

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

const DEVICES: readonly string[] = ["webgpu", "wasm"];
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 &&
        (value.language === undefined || typeof value.language === "string") &&
        (value.task === undefined ||
          value.task === "transcribe" ||
          value.task === "translate")
      );
    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);
        default:
          return false;
      }
    case "partial":
      return typeof value.text === "string";
    case "result":
      return typeof value.text === "string" && Array.isArray(value.chunks);
    case "error":
      return (
        typeof value.message === "string" &&
        (value.stage === "load" || value.stage === "transcribe")
      );
    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 } };
}