src/worker.ts

import {
  InterruptableStoppingCriteria,
  pipeline,
  TextStreamer,
} from "@huggingface/transformers";
import { MODEL_ID, type Precision, type GenerationSettings } from "./model";
import type { WorkerRequest, WorkerResponse } from "./messages";

type ChatMessage = {
  role: "system" | "user";
  content: string;
};

type ChatTokenizer = ConstructorParameters<typeof TextStreamer>[0] & {
  apply_chat_template: (
    conversation: ChatMessage[],
    options: { tokenize: false; add_generation_prompt: boolean },
  ) => string;
  chat_template?: unknown;
};

type TextGenerationPipeline = {
  tokenizer: ChatTokenizer;
  (input: string, options: Record<string, unknown>): Promise<unknown>;
};

type ErrorData = {
  name?: string;
  message: string;
  stack?: string;
  cause?: unknown;
  raw?: unknown;
};

const instances = new Map<Precision, Promise<TextGenerationPipeline>>();
let activeJobId = 0;
let stoppingCriteria: InstanceType<typeof InterruptableStoppingCriteria> | null = null;

function post(message: WorkerResponse) {
  self.postMessage(message);
}

function debug(jobId: number, message: string, data?: unknown) {
  post({ type: "debug", jobId, message, data });
}

function errorToData(error: unknown): ErrorData {
  if (error instanceof Error) {
    return {
      name: error.name,
      message: error.message,
      stack: error.stack,
      cause: error.cause instanceof Error ? errorToData(error.cause) : error.cause,
    };
  }
  return { message: String(error), raw: error };
}

function formatError(error: unknown) {
  const data = errorToData(error);
  if (typeof data.stack === "string") return data.stack;
  if (typeof data.message === "string") return data.message;
  return String(error);
}

function extractGeneratedText(result: unknown) {
  if (!Array.isArray(result) || !result[0] || typeof result[0] !== "object") return "";
  const generated = (result[0] as Record<string, unknown>).generated_text;
  return typeof generated === "string" ? generated : "";
}

function getProgressName(progress: unknown) {
  if (!progress || typeof progress !== "object") return "model files";
  const record = progress as Record<string, unknown>;
  return String(record.file ?? record.name ?? record.url ?? "model files");
}

function getProgressPercent(progress: unknown) {
  if (!progress || typeof progress !== "object") return 0;
  const record = progress as Record<string, unknown>;
  const raw = record.progress;
  return typeof raw === "number" && Number.isFinite(raw) ? Math.max(0, Math.min(100, raw)) : 0;
}

function getProgressStatus(progress: unknown) {
  if (!progress || typeof progress !== "object") return "loading";
  const record = progress as Record<string, unknown>;
  return String(record.status ?? "loading");
}

function workerEnvironment() {
  const nav = self.navigator;
  return {
    hasWebGpu: "gpu" in nav,
    userAgent: nav.userAgent,
    crossOriginIsolated: self.crossOriginIsolated,
    origin: self.location.origin,
    href: self.location.href,
  };
}

async function getGenerator(settings: GenerationSettings, jobId: number) {
  const existing = instances.get(settings.precision);
  if (existing) {
    debug(jobId, "Reusing cached pipeline", { precision: settings.precision });
    return existing;
  }

  post({
    type: "status",
    jobId,
    message: `Loading ${MODEL_ID} (${settings.precision.toUpperCase()})`,
  });
  debug(jobId, "Creating Transformers.js pipeline", {
    model: MODEL_ID,
    task: "text-generation",
    device: "webgpu",
    dtype: settings.precision,
    environment: workerEnvironment(),
  });

  const instancePromise = pipeline("text-generation", MODEL_ID, {
    device: "webgpu",
    dtype: settings.precision,
    progress_callback: (progress: unknown) => {
      if (jobId !== activeJobId) return;
      const data = {
        file: getProgressName(progress),
        progress: getProgressPercent(progress),
        status: getProgressStatus(progress),
      };
      post({ type: "progress", jobId, ...data });
      debug(jobId, "Model load progress", data);
    },
  }) as Promise<TextGenerationPipeline>;

  instances.set(settings.precision, instancePromise);

  try {
    const instance = await instancePromise;
    debug(jobId, "Pipeline loaded", {
      precision: settings.precision,
      hasTokenizer: Boolean(instance.tokenizer),
      hasChatTemplate: typeof instance.tokenizer.chat_template === "string",
    });
    return instance;
  } catch (error) {
    instances.delete(settings.precision);
    throw error;
  }
}

async function generate(request: Extract<WorkerRequest, { type: "generate" }>) {
  activeJobId = request.jobId;
  const started = performance.now();

  debug(request.jobId, "Generate request received", {
    settings: request.settings,
    systemPromptChars: request.systemPrompt.length,
    promptChars: request.prompt.length,
    environment: workerEnvironment(),
  });

  if (!("gpu" in self.navigator)) {
    post({
      type: "error",
      jobId: request.jobId,
      message: "WebGPU is not available in this browser. Use a recent Chrome or Edge build.",
      data: workerEnvironment(),
    });
    return;
  }

  const generator = await getGenerator(request.settings, request.jobId);
  if (request.jobId !== activeJobId) return;

  stoppingCriteria = new InterruptableStoppingCriteria();
  stoppingCriteria.reset();

  post({ type: "status", jobId: request.jobId, message: "Rendering prompt" });

  const messages: ChatMessage[] = [
    { role: "system", content: request.systemPrompt },
    { role: "user", content: request.prompt },
  ];

  const promptText = generator.tokenizer.apply_chat_template(messages, {
    tokenize: false,
    add_generation_prompt: true,
  }) as string;

  debug(request.jobId, "Prompt rendered", {
    promptChars: promptText.length,
    promptPreview: promptText.slice(0, 500),
    hasChatTemplate: typeof generator.tokenizer.chat_template === "string",
  });

  post({ type: "status", jobId: request.jobId, message: "Generating" });

  let streamedText = "";
  const streamer = new TextStreamer(generator.tokenizer, {
    skip_prompt: true,
    skip_special_tokens: true,
    callback_function: (text: string) => {
      if (request.jobId !== activeJobId) return;
      streamedText += text;
      post({ type: "token", jobId: request.jobId, text });
    },
  });

  const generationOptions = {
    add_special_tokens: false,
    max_new_tokens: request.settings.maxNewTokens,
    temperature: request.settings.temperature,
    top_k: request.settings.topK,
    repetition_penalty: request.settings.repetitionPenalty,
    do_sample: request.settings.temperature > 0,
    streamer,
    stopping_criteria: stoppingCriteria,
    return_full_text: false,
  };

  debug(request.jobId, "Calling generator", {
    ...generationOptions,
    streamer: "TextStreamer",
    stopping_criteria: "InterruptableStoppingCriteria",
  });

  const result = await generator(promptText, generationOptions);

  if (request.jobId !== activeJobId) return;

  const fallbackText = extractGeneratedText(result);
  debug(request.jobId, "Generation completed", {
    streamedChars: streamedText.length,
    fallbackChars: fallbackText.length,
    elapsedMs: performance.now() - started,
    resultShape: Array.isArray(result) ? result.map((item) => typeof item) : typeof result,
  });

  post({
    type: "complete",
    jobId: request.jobId,
    text: streamedText || fallbackText,
    elapsedMs: performance.now() - started,
  });
}

self.addEventListener("message", (event: MessageEvent<WorkerRequest>) => {
  const request = event.data;
  if (request.type === "reset") {
    activeJobId = request.jobId;
    stoppingCriteria?.interrupt();
    debug(request.jobId, "Generation interrupted by user");
    post({ type: "status", jobId: request.jobId, message: "Stopped" });
    return;
  }

  generate(request).catch((error: unknown) => {
    post({
      type: "error",
      jobId: request.jobId,
      message: formatError(error),
      data: {
        error: errorToData(error),
        environment: workerEnvironment(),
        settings: request.settings,
      },
    });
  });
});

self.addEventListener("error", (event) => {
  post({
    type: "error",
    jobId: activeJobId,
    message: event.error ? formatError(event.error) : event.message,
    data: {
      filename: event.filename,
      lineno: event.lineno,
      colno: event.colno,
      environment: workerEnvironment(),
    },
  });
});

self.addEventListener("unhandledrejection", (event) => {
  post({
    type: "error",
    jobId: activeJobId,
    message: formatError(event.reason),
    data: {
      error: errorToData(event.reason),
      environment: workerEnvironment(),
    },
  });
});