Jugar en el navegador

// Messages exchanged with `src/workers/decoder.worker.ts`. Both sides import this module, so a
// change to a payload breaks the build on the side that did not follow.
//
// The flow is always the same: `init` (download + session) answers with `progress`* and then
// `ready` or `error`; `logits` answers with `logits` or `error`; `dispose` answers with
// `disposed`. Every request carries an `id` the answer echoes, so the client can correlate
// replies even though the worker already serialises the work.

/** ONNX Runtime Web version; `scripts/copy-assets.mjs` copies that release into `public/ort/`. */
export const ORT_VERSION = '1.30.0';

/** Where the self-hosted ORT runtime (`.wasm` + loader) is served from. */
export const ORT_BASE = `/ort/${ORT_VERSION}/`;

/**
 * Cache API bucket for downloaded models. The suffix is bumped whenever what is stored (or how
 * it is keyed) changes, so an old browser cache is never read with new expectations.
 */
export const MODEL_CACHE = 'rukh-models-v1';

export type Backend = 'webgpu' | 'wasm';

export interface InitRequest {
  type: 'init';
  id: number;
  /** Registry stage id, only used for diagnostics and the cache log. */
  stage: string;
  /** Absolute or same-origin URL of the `.onnx` file. */
  url: string;
  /** Expected size in bytes; used for the progress bar when there is no `content-length`. */
  sizeBytes: number;
  /** Context window the stage was trained with (`DecoderConfig.block`), from the registry. */
  block: number;
  /** Vocabulary size the caller's tokenizer has; the model's output width must match it. */
  vocab: number;
}

export interface LogitsRequest {
  type: 'logits';
  id: number;
  /** Token ids of the prompt, already cropped to the context by the caller. */
  ids: number[];
}

export interface DisposeRequest {
  type: 'dispose';
  id: number;
}

export type WorkerRequest = InitRequest | LogitsRequest | DisposeRequest;

export interface ProgressMessage {
  type: 'progress';
  id: number;
  loaded: number;
  total: number;
}

export interface ReadyMessage {
  type: 'ready';
  id: number;
  backend: Backend;
  /** Why WebGPU was not used, when `backend` is `wasm` and a reason is known. */
  fallbackReason?: string;
  loadMs: number;
  /** The contract the session was checked against; `buildPrompt` crops to `block`. */
  block: number;
  vocab: number;
}

export interface LogitsMessage {
  type: 'logits';
  id: number;
  data: Float32Array;
  inferMs: number;
}

export interface DisposedMessage {
  type: 'disposed';
  id: number;
}

export interface ErrorMessage {
  type: 'error';
  id: number;
  message: string;
}

export type WorkerResponse =
  ProgressMessage | ReadyMessage | LogitsMessage | DisposedMessage | ErrorMessage;

const RESPONSE_TYPES = ['progress', 'ready', 'logits', 'disposed', 'error'] as const;

/**
 * Narrows whatever arrived through `postMessage` to a known response, or `null`. The worker is
 * ours, but the message channel is not typed at runtime and a stray message must not crash the
 * page.
 */
export function asResponse(data: unknown): WorkerResponse | null {
  if (typeof data !== 'object' || data === null) return null;
  const message = data as { type?: unknown; id?: unknown };
  if (typeof message.id !== 'number') return null;
  const type = message.type;
  if (typeof type !== 'string') return null;
  if (!(RESPONSE_TYPES as readonly string[]).includes(type)) return null;
  return data as WorkerResponse;
}

/** Percentage (0-100) of a progress message; `total` of 0 reports 0 rather than `NaN`. */
export function progressPercent(message: Pick<ProgressMessage, 'loaded' | 'total'>): number {
  if (message.total <= 0) return 0;
  return Math.min(100, Math.round((message.loaded / message.total) * 100));
}

/** Bytes as MB with one decimal, the unit the consent step and the progress bar speak. */
export function megabytes(bytes: number): number {
  return Math.round((bytes / 1_000_000) * 10) / 10;
}

// What the demo demands of a model file before it is allowed to play, and the checks that prove
// it. A downloaded `.onnx` is just bytes: nothing stops a stale cache entry, a half-renamed Hub
// file or a model trained with another vocabulary from loading cleanly and then answering with a
// logits vector that means something else. Every id the sampler reads out of that vector would be
// a different move, so the game would look plausible and be nonsense.
//
// Two things have to hold and both are checked:
//
//   * the **width** of the output equals the tokenizer's vocabulary (2030 UCI tokens). ORT Web
//     does not expose `metadata_props`, so the `rukh_vocab_size` the exporter writes into the file
//     is unreachable from the browser; what is reachable is the output shape ORT reports
//     (`outputMetadata`) and, failing that, the length of the first real answer;
//   * the **context** (`block`) the prompt is cropped to. Nothing in the session says it, so it
//     comes from the registry entry and travels back to the main thread in `ready`, which is what
//     keeps `buildPrompt` from hardcoding 200.
//
// The messages are in Spanish on purpose: they end up in the model panel, in front of a player.

/** The contract a loaded session is playing under. */
export interface ModelContract {
  /** Context window in tokens (`DecoderConfig.block`); the prompt is cropped to it. */
  block: number;
  /** Width of the logits vector, which must equal the tokenizer's vocabulary size. */
  vocab: number;
}

/** The part of ORT's `InferenceSession.ValueMetadata` this check needs. */
export interface OutputMetadata {
  name: string;
  isTensor?: boolean;
  /** Dimensions; a string entry is a symbolic axis (`batch`, `sequence`). */
  shape?: readonly (number | string)[];
}

/** The part of `InferenceSession` this check needs, so the tests can hand it a plain object. */
export interface SessionLike {
  readonly outputNames: readonly string[];
  readonly outputMetadata?: readonly OutputMetadata[];
}

/** Smallest context that can still hold the three header tokens plus one move. */
export const MIN_BLOCK = 4;

export const DECLARED_SOURCE = 'la forma declarada en el fichero';
export const RUN_SOURCE = 'la primera respuesta del modelo';

/**
 * The fixed last dimension of the model's only output, or `null` when ORT reports it as symbolic
 * (or reports nothing at all, which is what the WASM backend does for some graphs).
 */
export function declaredVocab(session: SessionLike): number | null {
  const name = session.outputNames[0];
  const meta = session.outputMetadata?.find((entry) => entry.name === name);
  const shape = meta?.shape;
  if (!shape || shape.length === 0) return null;
  const last = shape[shape.length - 1];
  return typeof last === 'number' && Number.isInteger(last) && last > 0 ? last : null;
}

/** Throws with a message a player can read when the logits are not the tokenizer's width. */
export function assertVocab(width: number, expected: number, source: string): void {
  if (width === expected) return;
  throw new Error(
    `El modelo no encaja con el tokenizador: ${source} da ${width} valores y el vocabulario UCI ` +
      `tiene ${expected}. No se juega con este fichero; borra los modelos descargados y vuelve a ` +
      `intentarlo.`,
  );
}

/**
 * The contract of a freshly created session: `block` from the registry (nothing in the session
 * knows it) and `vocab` from the declared output shape when ORT gives a fixed one, checked against
 * the tokenizer right away. When the shape is symbolic the width is still unknown here and
 * `assertVocab` runs on the first `run` instead — see `decoder.worker.ts`.
 */
export function readContract(
  session: SessionLike,
  expectedVocab: number,
  block: number,
): ModelContract {
  if (!Number.isInteger(block) || block < MIN_BLOCK) {
    throw new Error(`El contexto declarado para esta etapa (${block}) no es utilizable.`);
  }
  const declared = declaredVocab(session);
  if (declared !== null) assertVocab(declared, expectedVocab, DECLARED_SOURCE);
  return { block, vocab: declared ?? expectedVocab };
}

// Dedicated worker that owns the ONNX Runtime session: it downloads the model (streaming the
// progress back), keeps it in the Cache API and answers `logits` requests with the last step's
// logits. The main thread never touches ORT, so a 28 MB WASM runtime and a 40 MB model never
// block the board.
//
// Two serialisation rules come from ORT itself and are not negotiable: sessions are created one
// at a time and `run` calls never overlap (the JSEP/asyncify build cannot re-enter an async
// call). Both go through `serial`, a single promise chain — the download included, so two `init`
// messages in flight cannot end up fetching twice and racing to install their session.
//
// The file is never trusted: `readContract` checks the declared output width against the
// tokenizer's vocabulary as soon as the session exists, and `assertVocab` checks the real width
// of every answer. See `src/lib/contract.ts` for why.

import * as ort from 'onnxruntime-web/webgpu';
import { RUN_SOURCE, assertVocab, readContract, type ModelContract } from '../lib/contract';
import { downloadModel } from '../lib/download';
import {
  MODEL_CACHE,
  ORT_BASE,
  type Backend,
  type WorkerRequest,
  type WorkerResponse,
} from '../lib/worker-protocol';

/** The worker global, typed with just what this file uses (avoids pulling in the webworker lib). */
interface WorkerScope {
  postMessage(message: WorkerResponse, transfer?: Transferable[]): void;
  addEventListener(type: 'message', listener: (event: MessageEvent) => void): void;
}

const ctx = self as unknown as WorkerScope;

// Self-hosted runtime: `public/ort/<version>/` is filled by `scripts/copy-assets.mjs` so the
// page never reaches a CDN (the CSP would not allow it either).
ort.env.wasm.wasmPaths = ORT_BASE;
// One thread, always. WASM is only the fallback here (WebGPU is the fast path and this decoder
// is 40 MB), and asking for more threads under cross-origin isolation makes ORT reach for the
// threaded/proxy artefacts that `scripts/copy-assets.mjs` deliberately does not publish. Keeping
// it at one means the runtime needs exactly the two files we do publish, isolated or not.
ort.env.wasm.numThreads = 1;
// No proxy worker either: it would be a third artefact (`ort-wasm-proxy-worker`) and this code
// already runs off the main thread, which is the only thing the proxy buys.
ort.env.wasm.proxy = false;
ort.env.logLevel = 'error';

let session: ort.InferenceSession | null = null;
/** Block and vocabulary the live session was accepted under; null while there is no session. */
let contract: ModelContract | null = null;

/** The single chain every ORT call is queued on: no two sessions or runs are ever in flight. */
let chain: Promise<unknown> = Promise.resolve();

function serial<T>(work: () => Promise<T>): Promise<T> {
  const next = chain.then(work, work);
  chain = next.then(
    () => undefined,
    () => undefined,
  );
  return next;
}

function reply(message: WorkerResponse, transfer?: Transferable[]): void {
  ctx.postMessage(message, transfer);
}

/** `downloadModel` wired to this worker's globals, reporting progress back to the main thread. */
function fetchModel(id: number, url: string, sizeBytes: number): Promise<Uint8Array> {
  return downloadModel(url, sizeBytes, {
    fetch: globalThis.fetch.bind(globalThis),
    caches: typeof caches === 'undefined' ? undefined : caches,
    cacheName: MODEL_CACHE,
    onProgress: (loaded, total) => reply({ type: 'progress', id, loaded, total }),
  });
}

/** The WebGPU adapter, or the reason there is none. Never throws. */
async function webgpuAdapter(): Promise<{ adapter: unknown } | { reason: string }> {
  const gpu = (globalThis as { navigator?: { gpu?: { requestAdapter(): Promise<unknown> } } })
    .navigator?.gpu;
  if (!gpu) return { reason: 'este navegador no expone WebGPU' };
  try {
    const adapter = await gpu.requestAdapter();
    if (!adapter) return { reason: 'WebGPU no ha ofrecido ningún adaptador' };
    return { adapter };
  } catch (cause) {
    return { reason: `WebGPU ha fallado al pedir el adaptador: ${describe(cause)}` };
  }
}

function describe(cause: unknown): string {
  return cause instanceof Error ? cause.message : String(cause);
}

/**
 * WebGPU whenever `requestAdapter()` hands out an adapter; WASM only when it does not, or when
 * creating the session on WebGPU throws. Both fallbacks carry the reason, which the panel shows
 * under the badge so a slow game on WASM is never a mystery.
 */
async function createSession(
  bytes: Uint8Array,
): Promise<{ session: ort.InferenceSession; backend: Backend; reason: string | null }> {
  const options: ort.InferenceSession.SessionOptions = {
    graphOptimizationLevel: 'all',
    executionMode: 'sequential',
  };
  const probe = await webgpuAdapter();
  let reason = 'reason' in probe ? probe.reason : null;
  if (!reason) {
    try {
      const created = await ort.InferenceSession.create(bytes, {
        ...options,
        executionProviders: ['webgpu'],
      });
      return { session: created, backend: 'webgpu', reason: null };
    } catch (cause) {
      // An adapter that cannot compile the graph is a normal outcome, not a bug: say so and
      // carry on with WASM.
      reason = `WebGPU no ha podido crear la sesión: ${describe(cause)}`;
    }
  }
  const created = await ort.InferenceSession.create(bytes, {
    ...options,
    executionProviders: ['wasm'],
  });
  return { session: created, backend: 'wasm', reason };
}

/**
 * Downloads the weights and hands them to ORT. The buffer is a local of this frame on purpose:
 * when it returns, the only copy of the model still alive is the one inside the session, instead
 * of 80 MB sitting next to it for as long as the worker lives.
 */
async function loadSession(request: Extract<WorkerRequest, { type: 'init' }>) {
  const bytes = await fetchModel(request.id, request.url, request.sizeBytes);
  return createSession(bytes);
}

async function init(request: Extract<WorkerRequest, { type: 'init' }>): Promise<void> {
  const started = performance.now();
  // The download runs inside the chain too: outside it, a second `init` would start its own fetch
  // while the first was still tearing the old session down, and both would race to install one.
  const ready = await serial(async () => {
    await session?.release();
    session = null;
    contract = null;
    const created = await loadSession(request);
    const checked = readContract(created.session, request.vocab, request.block);
    session = created.session;
    contract = checked;
    return { backend: created.backend, reason: created.reason, contract: checked };
  });
  reply({
    type: 'ready',
    id: request.id,
    backend: ready.backend,
    fallbackReason: ready.reason ?? undefined,
    loadMs: Math.round(performance.now() - started),
    block: ready.contract.block,
    vocab: ready.contract.vocab,
  });
}

async function logits(request: Extract<WorkerRequest, { type: 'logits' }>): Promise<void> {
  const data = await serial(async () => {
    const current = session;
    const live = contract;
    if (!current || !live) throw new Error('el modelo todavía no está cargado');
    const ids = BigInt64Array.from(request.ids, (id) => BigInt(id));
    const feeds: Record<string, ort.Tensor> = {
      [current.inputNames[0]]: new ort.Tensor('int64', ids, [1, request.ids.length]),
    };
    const started = performance.now();
    const output = await current.run(feeds);
    const inferMs = performance.now() - started;
    const tensor = output[current.outputNames[0]];
    const raw = tensor.data;
    if (!(raw instanceof Float32Array)) {
      throw new Error(`la salida del modelo no es float32 (${tensor.type})`);
    }
    // The declared shape is often symbolic, so the real width is only knowable here. A model
    // whose logits are not the tokenizer's vocabulary would map every id to a different move.
    assertVocab(raw.length, live.vocab, RUN_SOURCE);
    // Copy out of the ORT arena: the tensor's buffer may be reused by the next run.
    return { values: raw.slice(), inferMs };
  });
  reply({ type: 'logits', id: request.id, data: data.values, inferMs: data.inferMs }, [
    data.values.buffer,
  ]);
}

async function dispose(request: Extract<WorkerRequest, { type: 'dispose' }>): Promise<void> {
  await serial(async () => {
    await session?.release();
    session = null;
    contract = null;
  });
  reply({ type: 'disposed', id: request.id });
}

ctx.addEventListener('message', (event: MessageEvent) => {
  const request = event.data as WorkerRequest;
  const run = async () => {
    switch (request.type) {
      case 'init':
        return init(request);
      case 'logits':
        return logits(request);
      case 'dispose':
        return dispose(request);
      default:
        return undefined;
    }
  };
  void run().catch((error: unknown) => {
    reply({
      type: 'error',
      id: request.id,
      message: error instanceof Error ? error.message : String(error),
    });
  });
});

// Main-thread side of the decoder worker. The `new Worker(new URL(...), { type: 'module' })`
// literal lives here, at the call site, because that is the only shape Vite recognises to emit
// the worker as its own chunk (`tests/worker-chunk.test.ts` guards it).
import {
  asResponse,
  type Backend,
  type InitRequest,
  type ProgressMessage,
  type ReadyMessage,
  type WorkerRequest,
} from '../lib/worker-protocol';

export interface LoadReport {
  backend: Backend;
  /** Why WebGPU was not used, when it was not. */
  fallbackReason?: string;
  loadMs: number;
  /** Context window the session was accepted under; `buildPrompt` crops the prompt to it. */
  block: number;
  /** Width of the logits the model answers with, already checked against the tokenizer. */
  vocab: number;
}

/** What `init` needs to know about a stage; the worker's `id` is added by `send`. */
export type LoadRequest = Omit<InitRequest, 'type' | 'id'>;

export interface LogitsReport {
  data: Float32Array;
  inferMs: number;
}

export interface Decoder {
  /** Downloads (or reads from the cache) the model, creates the session and checks its contract. */
  init(request: LoadRequest, onProgress?: (progress: ProgressMessage) => void): Promise<LoadReport>;
  /** Logits of the last step for a prompt already cropped to the context window. */
  logits(ids: number[]): Promise<LogitsReport>;
  /** Releases the session and terminates the worker; the handle is unusable afterwards. */
  dispose(): Promise<void>;
}

/** `Omit` over a union collapses it, so the distribution is written out explicitly. */
type Unidentified<T> = T extends WorkerRequest ? Omit<T, 'id'> : never;

interface Pending {
  resolve: (value: never) => void;
  reject: (error: Error) => void;
  onProgress?: (progress: ProgressMessage) => void;
}

export function createDecoder(): Decoder {
  const worker = new Worker(new URL('./decoder.worker.ts', import.meta.url), { type: 'module' });
  const pending = new Map<number, Pending>();
  let nextId = 1;
  let terminated = false;

  const fail = (error: Error) => {
    for (const entry of pending.values()) entry.reject(error);
    pending.clear();
  };

  worker.addEventListener('message', (event: MessageEvent) => {
    const message = asResponse(event.data);
    if (!message) return;
    const entry = pending.get(message.id);
    if (!entry) return;
    if (message.type === 'progress') {
      entry.onProgress?.(message);
      return;
    }
    pending.delete(message.id);
    if (message.type === 'error') {
      entry.reject(new Error(message.message));
      return;
    }
    (entry.resolve as (value: unknown) => void)(message);
  });

  worker.addEventListener('error', (event: ErrorEvent) => {
    fail(new Error(event.message || 'el worker del modelo ha fallado'));
  });

  function send<T>(
    request: Unidentified<WorkerRequest>,
    onProgress?: (progress: ProgressMessage) => void,
  ): Promise<T> {
    if (terminated) return Promise.reject(new Error('el worker del modelo ya se ha cerrado'));
    const id = nextId++;
    return new Promise<T>((resolve, reject) => {
      pending.set(id, { resolve: resolve as (value: never) => void, reject, onProgress });
      worker.postMessage({ ...request, id } as WorkerRequest);
    });
  }

  return {
    async init(request, onProgress) {
      const ready = await send<ReadyMessage>({ type: 'init', ...request }, onProgress);
      return {
        backend: ready.backend,
        fallbackReason: ready.fallbackReason,
        loadMs: ready.loadMs,
        block: ready.block,
        vocab: ready.vocab,
      };
    },
    async logits(ids) {
      const answer = await send<{ data: Float32Array; inferMs: number }>({ type: 'logits', ids });
      return { data: answer.data, inferMs: answer.inferMs };
    },
    async dispose() {
      if (terminated) return;
      try {
        await send<unknown>({ type: 'dispose' });
      } catch {
        /* the worker may already be gone; terminating below is enough */
      }
      terminated = true;
      fail(new Error('el worker del modelo se ha cerrado'));
      worker.terminate();
    },
  };
}

# Included in every location: nginx does not inherit add_header into a block that declares
# its own add_header, so each location pulls the full set from here.
#
# The page-level CSP (script/style hashes) lives in the <meta> tag generated by Astro;
# frame-ancestors cannot be set from <meta>, so it is added as a header.
add_header Content-Security-Policy "frame-ancestors 'none'" always;
add_header Strict-Transport-Security "max-age=31536000; includeSubDomains" always;
add_header X-Content-Type-Options "nosniff" always;
add_header Referrer-Policy "strict-origin-when-cross-origin" always;
add_header Permissions-Policy "camera=(), microphone=(), geolocation=(), payment=(), usb=()" always;
# Cross-origin isolation: required for multi-threaded WASM (onnxruntime-web, P2).
add_header Cross-Origin-Opener-Policy "same-origin" always;
add_header Cross-Origin-Embedder-Policy "require-corp" always;