SkateSpots/src/workers/translate.worker.ts

/**
 * On-device translation worker (all app languages).
 *
 * - One multilingual model (M2M100, MIT) covers every app language except
 *   Basque. Basque uses a small Opus-MT eu↔en pair and pivots through English.
 * - Model files come from the self-hosted relay (public ONNX bytes).
 * - Message text is received via postMessage and never put into a URL/body.
 * - The network guard blocks every request except the model host (GET/HEAD)
 *   and same-origin ONNX runtime assets.
 * - No logging of message text; errors are returned as codes only.
 */
import { env, pipeline } from '@xenova/transformers';
import { createGuardedFetch } from '@/lib/translationGuard';
import { clearStaleBasqueTokenizers } from '@/lib/translationModelCache';

const MODEL_ORIGIN = 'https://relay.skate-spots.org';
const MODEL_PREFIX = '/models/';
const APP_ORIGIN = self.location.origin;

const M2M100 = 'Xenova/m2m100_418M';
const EU_EN = 'Xenova/opus-mt-eu-en';
const EN_EU = 'Xenova/opus-mt-en-eu';

env.allowLocalModels = false;
env.allowRemoteModels = true;
env.remoteHost = `${MODEL_ORIGIN}${MODEL_PREFIX}`;
env.useBrowserCache = true;

const onnxWasm = (env.backends.onnx as unknown as {
  wasm?: { wasmPaths?: string; numThreads?: number };
} | undefined)?.wasm;
if (onnxWasm) {
  onnxWasm.wasmPaths = new URL('../ort/', self.location.href).href;
  onnxWasm.numThreads = 1;
}

globalThis.fetch = createGuardedFetch(
  globalThis.fetch.bind(globalThis),
  {
    modelOrigin: MODEL_ORIGIN,
    modelPathPrefix: MODEL_PREFIX,
    appOrigin: APP_ORIGIN,
    appAssetPathPrefix: '/ort/',
  },
);

type Incoming = {
  type: 'translate';
  id: number;
  text: string;
  source: string;
  target: string;
};

type Outgoing =
  | { type: 'progress'; id: number; status?: string; progress?: number; file?: string }
  | { type: 'result'; id: number; text: string }
  | { type: 'error'; id: number; code: 'FAILED' };

const ctx = self as unknown as {
  postMessage: (message: Outgoing) => void;
  onmessage: ((event: MessageEvent<Incoming>) => void) | null;
};

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const translators = new Map<string, Promise<any>>();
let basqueCacheChecked = false;

/** Evict the pre-fix Basque tokenizer from browser cache exactly once. */
async function ensureBasqueCache(): Promise<void> {
  if (basqueCacheChecked) return;
  basqueCacheChecked = true;
  await clearStaleBasqueTokenizers();
}

/** Collapse UI variants onto the model's language tokens. */
function baseLanguage(code: string): string {
  if (code === 'pt-PT') return 'pt';
  if (code === 'ko-KP') return 'ko';
  if (code === 'zh-CN' || code === 'zh-TW') return 'zh';
  return code;
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
function getTranslator(modelId: string, requestId: number): Promise<any> {
  const existing = translators.get(modelId);
  if (existing) return existing;

  const loading = pipeline('translation', modelId, {
    quantized: true,
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    progress_callback: (progress: any) => {
      ctx.postMessage({
        type: 'progress',
        id: requestId,
        status: typeof progress?.status === 'string' ? progress.status : undefined,
        progress: typeof progress?.progress === 'number' ? progress.progress : undefined,
        file: typeof progress?.file === 'string' ? progress.file : undefined,
      });
    },
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
  }).catch((error: any) => {
    translators.delete(modelId);
    throw error;
  });

  translators.set(modelId, loading);
  return loading;
}

async function run(
  modelId: string,
  text: string,
  requestId: number,
  options: Record<string, unknown> = {},
): Promise<string> {
  const translator = await getTranslator(modelId, requestId);
  const output = await translator(text, { max_new_tokens: 512, ...options });
  const translated = Array.isArray(output)
    ? output[0]?.translation_text
    : output?.translation_text;
  if (typeof translated !== 'string' || !translated) {
    throw new Error('empty translation');
  }
  return translated;
}

async function translate(
  text: string,
  source: string,
  target: string,
  requestId: number,
): Promise<string> {
  const s = baseLanguage(source);
  const t = baseLanguage(target);
  if (s === t) return text;

  // Basque is not in M2M100: use Opus eu↔en and pivot through English.
  if (s === 'eu' || t === 'eu') {
    await ensureBasqueCache();
    if (s === 'eu' && t === 'en') return run(EU_EN, text, requestId);
    if (s === 'en' && t === 'eu') return run(EN_EU, text, requestId);
    if (s === 'eu') {
      const english = await run(EU_EN, text, requestId);
      return run(M2M100, english, requestId, { src_lang: 'en', tgt_lang: t });
    }
    const english = await run(M2M100, text, requestId, { src_lang: s, tgt_lang: 'en' });
    return run(EN_EU, english, requestId);
  }

  return run(M2M100, text, requestId, { src_lang: s, tgt_lang: t });
}

ctx.onmessage = (event) => {
  const data = event.data;
  if (!data || data.type !== 'translate' || typeof data.id !== 'number') return;

  void (async () => {
    try {
      const translated = await translate(data.text, data.source, data.target, data.id);
      ctx.postMessage({ type: 'result', id: data.id, text: translated });
    } catch {
      // Never forward the underlying error: it may contain message text.
      ctx.postMessage({ type: 'error', id: data.id, code: 'FAILED' });
    }
  })();
};