src/worker/engine.ts

// Engine suy luận MiniCPM5-2B (ONNX q4f16) chạy trên WebGPU bằng Transformers.js.
// Toàn bộ nằm trong worker: nạp model, dựng prompt từ chat template, lấy mẫu, parse tool call.

import {
  AutoModelForCausalLM,
  AutoTokenizer,
  InterruptableStoppingCriteria,
  TextStreamer,
  env,
} from '@huggingface/transformers';
import type { AgentMessage, AssistantMessage, DeviceInfo, ToolDef } from '../shared/types';
import type { WorkerEvent } from '../shared/protocol';
import { MODEL_HOST, MODEL_REPO, MODEL_REVISION } from './manifest';
import { openModelCache } from './modelCache';
import { RateMeter, createTopPProcessor } from './sampling';
import { parseAssistantOutput, splitThinking, toTemplateMessages, type TemplateMessage } from './toolParser';

/** Ngân sách prompt + output, và số token tối đa cho mỗi lần gọi model. */
export const CONTEXT_WINDOW = 8192;
export const MAX_NEW_TOKENS = 2048;
export const EOS_TOKEN_IDS = [1, 130073];
export const TOP_P = 0.95;

export const ORT_VERSION = '1.26.0-dev.20260416-b7804b056c';

export const MODEL_DISPLAY_NAME = 'MiniCPM5-2B · q4f16';

export interface GenerateRequest {
  systemPrompt: string;
  messages: AgentMessage[];
  tools: ToolDef[];
  signal?: AbortSignal;
  emit: (event: WorkerEvent) => void;
  /** Gọi mỗi khi có thêm chữ mới (đã gồm phần suy nghĩ và câu trả lời hiển thị được). */
  onPartial: (message: AssistantMessage) => void;
}

export interface Engine {
  readonly loaded: boolean;
  readonly device: DeviceInfo | undefined;
  load(options: { cachedOnly?: boolean }, signal: AbortSignal, emit: (e: WorkerEvent) => void): Promise<DeviceInfo>;
  generate(req: GenerateRequest): Promise<AssistantMessage>;
  stop(): void;
}

// ───────────────────────── Dựng prompt trong ngân sách context ─────────────────────────

export interface TokenizedInputs {
  input_ids: { dims: readonly number[] };
  [key: string]: unknown;
}

/**
 * Áp chat template; nếu prompt dài hơn `maxInput` token thì bỏ dần các lượt cũ nhất
 * (tới trước tin nhắn người dùng kế tiếp) cho tới khi vừa.
 */
export function buildInputs<T extends TokenizedInputs>(
  systemPrompt: string,
  messages: AgentMessage[],
  apply: (messages: TemplateMessage[]) => T,
  maxInput: number,
): { inputs: T; dropped: number } {
  const remaining = [...messages];
  let dropped = 0;
  for (;;) {
    const inputs = apply(toTemplateMessages(systemPrompt, remaining));
    if (inputs.input_ids.dims[1] <= maxInput) return { inputs, dropped };
    const next = remaining.findIndex((m, i) => i > 0 && m.role === 'user');
    if (next < 0) {
      throw new Error(
        'This task exceeds the browser context budget. Start a new chat or use smaller files and shorter tool output.',
      );
    }
    remaining.splice(0, next);
    dropped += next;
  }
}

// ───────────────────────── Đường dẫn runtime ONNX (wasm) ─────────────────────────

async function resolveRuntimeBase(): Promise<string> {
  const candidates: string[] = [];
  const override = import.meta.env?.VITE_ORT_WASM_BASE as string | undefined;
  if (override) candidates.push(override.endsWith('/') ? override : override + '/');
  // Nếu bạn tự đặt các file ort-wasm-simd-threaded.asyncify.{mjs,wasm} vào public/runtime/ thì ưu tiên dùng.
  candidates.push(new URL(`${import.meta.env?.BASE_URL ?? '/'}runtime/`, self.location.href).href);
  candidates.push(`https://cdn.jsdelivr.net/npm/onnxruntime-web@${ORT_VERSION}/dist/`);
  candidates.push(`https://unpkg.com/onnxruntime-web@${ORT_VERSION}/dist/`);

  for (const base of candidates) {
    try {
      const res = await fetch(base + 'ort-wasm-simd-threaded.asyncify.mjs', {
        method: 'HEAD',
        signal: AbortSignal.timeout(8000),
      });
      const type = res.headers.get('content-type') ?? '';
      if (res.ok && !type.includes('html')) return base;
    } catch {
      /* thử nguồn tiếp theo */
    }
  }
  // Cuối cùng: để Transformers.js tự chọn (jsDelivr).
  return candidates[candidates.length - 2];
}

// ───────────────────────── Engine ─────────────────────────

export class MiniCpmEngine implements Engine {
  private tokenizer: any;
  private model: any;
  private loading: Promise<DeviceInfo> | undefined;
  private deviceInfo: DeviceInfo | undefined;
  private readonly stopping = new InterruptableStoppingCriteria();

  get loaded() {
    return Boolean(this.model);
  }
  get device() {
    return this.deviceInfo;
  }

  stop() {
    this.stopping.interrupt();
  }

  load(
    { cachedOnly = false }: { cachedOnly?: boolean },
    signal: AbortSignal,
    emit: (e: WorkerEvent) => void,
  ): Promise<DeviceInfo> {
    if (this.model && this.deviceInfo) return Promise.resolve(this.deviceInfo);
    if (this.loading) return this.loading;

    this.loading = (async () => {
      const gpu = (navigator as Navigator & { gpu?: any }).gpu;
      const adapter = await gpu?.requestAdapter({ powerPreference: 'high-performance' });
      if (!adapter) {
        throw new Error('WebGPU is unavailable. Try an up-to-date browser with GPU acceleration enabled.');
      }
      if (!adapter.features.has('shader-f16')) {
        throw new Error('This model requires WebGPU shader-f16 support on your device.');
      }
      const device: DeviceInfo = {
        vendor: adapter.info?.vendor,
        architecture: adapter.info?.architecture,
        maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
        features: [...adapter.features],
      };
      this.deviceInfo = device;
      emit({ type: 'device', device });

      // Cấu hình Transformers.js: đọc model từ cache OPFS tuỳ biến (đã xác minh SHA-256).
      env.allowRemoteModels = true;
      env.allowLocalModels = false;
      env.remoteHost = MODEL_HOST;
      env.remotePathTemplate = `{model}/resolve/${MODEL_REVISION}/`;
      const baseUrl = env.remoteHost + env.remotePathTemplate.replace('{model}', MODEL_REPO);

      const cache = await openModelCache(baseUrl, {
        signal,
        cachedOnly,
        onProgress: (p) => emit({ type: 'load_progress', ...p }),
      });
      env.customCache = cache as any;
      env.useCustomCache = true;
      env.useBrowserCache = false;

      const onnx = (env.backends as any).onnx;
      onnx.wasm.wasmPaths = await resolveRuntimeBase();
      onnx.wasm.numThreads = 1;
      signal.throwIfAborted();

      emit({ type: 'load_progress', phase: 'compile' });
      [this.tokenizer, this.model] = await Promise.all([
        AutoTokenizer.from_pretrained(MODEL_REPO),
        AutoModelForCausalLM.from_pretrained(MODEL_REPO, { device: 'webgpu', dtype: 'q4f16' } as any),
      ]);
      signal.throwIfAborted();

      emit({ type: 'load_progress', phase: 'warmup' });
      const warm = this.tokenizer('Hello');
      await this.model.generate({ ...warm, max_new_tokens: 1, do_sample: false, top_k: 0 });
      signal.throwIfAborted();

      const gpuDevice = onnx.webgpu?.device as any;
      device.maxStorageBufferBindingSize = gpuDevice?.limits.maxStorageBufferBindingSize ?? device.maxStorageBufferBindingSize;
      gpuDevice?.lost.then((info: { reason: string }) => {
        if (info.reason !== 'destroyed') {
          emit({ type: 'fatal', error: 'GPU device was lost. Reload this page to reload the cached model.' });
        }
      });
      emit({ type: 'loaded', device, cachedBytes: cache.cachedBytes });
      return device;
    })()
      .catch(async (err) => {
        try {
          await this.model?.dispose?.();
        } catch {
          /* bỏ qua lỗi khi dọn dẹp */
        }
        this.model = undefined;
        this.tokenizer = undefined;
        if (signal.aborted) throw new DOMException('Stopped.', 'AbortError');
        throw err;
      })
      .finally(() => {
        this.loading = undefined;
      });
    return this.loading;
  }

  async generate(req: GenerateRequest): Promise<AssistantMessage> {
    const { systemPrompt, messages, tools, signal, emit, onPartial } = req;
    const message: AssistantMessage = {
      role: 'assistant',
      content: [],
      usage: { input: 0, output: 0, totalTokens: 0 },
      timestamp: Date.now(),
    };

    const meter = new RateMeter();
    let rateTimer: ReturnType<typeof setInterval> | undefined;
    let lastUsageAt = 0;
    const reportUsage = (force = false) => {
      const now = performance.now();
      if (!force && now - lastUsageAt < 80) return;
      lastUsageAt = now;
      emit({ type: 'context_usage', inputTokens: message.usage.input, outputTokens: message.usage.output });
    };
    const interrupt = () => this.stopping.interrupt();

    try {
      if (!this.model) throw new Error('Load the model first.');
      this.stopping.reset();
      signal?.throwIfAborted();
      signal?.addEventListener('abort', interrupt, { once: true });
      emit({ type: 'inference_activity', phase: 'prefill' });

      const maxNew = MAX_NEW_TOKENS;
      const templateTools = tools.map((t) => ({
        type: 'function',
        function: { name: t.name, description: t.description, parameters: t.parameters },
      }));
      const { inputs, dropped } = buildInputs(
        systemPrompt,
        messages,
        (msgs) =>
          this.tokenizer.apply_chat_template(msgs, {
            tools: templateTools,
            enable_thinking: true,
            add_generation_prompt: true,
            return_dict: true,
          }),
        CONTEXT_WINDOW - maxNew,
      );
      if (dropped) emit({ type: 'context_trim', dropped });

      const promptLength: number = inputs.input_ids.dims[1];
      message.usage.input = promptLength;
      message.usage.totalTokens = promptLength;
      reportUsage(true);

      let raw = '';
      const startedAt = performance.now();
      let firstTokenMs: number | undefined;

      const streamer = new TextStreamer(this.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: false,
        token_callback_function: (tokens: bigint[]) => {
          if (!tokens.length) return;
          const now = performance.now();
          firstTokenMs ??= now - startedAt;
          meter.add(tokens.length, now);
          if (rateTimer === undefined) {
            emit({ type: 'inference_activity', phase: 'decode', rate: null, outputTokens: meter.total });
            rateTimer = setInterval(
              () =>
                emit({
                  type: 'inference_activity',
                  phase: 'decode',
                  rate: meter.rate(performance.now()),
                  outputTokens: meter.total,
                }),
              250,
            );
          }
          message.usage.output += tokens.length;
          message.usage.totalTokens = message.usage.input + message.usage.output;
          reportUsage();
        },
        callback_function: (piece: string) => {
          raw += piece;
          const split = splitThinking(raw, { thinkingPrefilled: true });
          // Khi chưa xong phần suy nghĩ, giữ lại vài ký tự cuối phòng đó là đầu thẻ </think>.
          const thinking = split.complete ? split.thinking : split.thinking.slice(0, Math.max(0, split.thinking.length - 8));
          const content: AssistantMessage['content'] = [];
          if (thinking) content.push({ type: 'thinking', thinking });
          if (split.complete) {
            const visible = split.answer.split('<')[0]; // ẩn XML của lời gọi tool
            if (visible) content.push({ type: 'text', text: visible });
          }
          message.content = content;
          onPartial(message);
        },
      });

      const output = await this.model.generate({
        ...inputs,
        max_new_tokens: maxNew,
        do_sample: true,
        temperature: 1,
        top_p: TOP_P,
        top_k: 0,
        repetition_penalty: 1,
        logits_processor: [createTopPProcessor(TOP_P)],
        eos_token_id: EOS_TOKEN_IDS,
        streamer,
        stopping_criteria: this.stopping,
      });

      const generated: number[] = output.tolist()[0].slice(promptLength).map(Number);
      message.usage.output = generated.length;
      message.usage.totalTokens = promptLength + generated.length;
      reportUsage(true);
      signal?.throwIfAborted();

      const finalText: string = this.tokenizer.decode(generated, { skip_special_tokens: false });
      const content = parseAssistantOutput(finalText, tools, { thinkingPrefilled: true });
      const hasToolCall = content.some((c) => c.type === 'toolCall');
      const endedWithEos = EOS_TOKEN_IDS.includes(generated.at(-1) as number);
      if (hasToolCall && !endedWithEos) {
        throw new Error('The tool response exceeded the output limit. No tool was executed. Try a smaller edit.');
      }
      message.content = content;
      message.stopReason = hasToolCall ? 'toolUse' : endedWithEos ? 'stop' : 'length';

      emit({
        type: 'generation',
        inputTokens: promptLength,
        outputTokens: generated.length,
        elapsedMs: performance.now() - startedAt,
        firstTokenMs,
        stopReason: message.stopReason,
      });
      return message;
    } catch (err) {
      if (message.usage.input) reportUsage(true);
      message.content = message.content.filter((c) => c.type !== 'toolCall');
      const aborted = Boolean(signal?.aborted);
      message.stopReason = aborted ? 'aborted' : 'error';
      message.errorMessage = aborted ? 'Stopped.' : String((err as Error)?.message ?? err);
      return message;
    } finally {
      if (rateTimer !== undefined) clearInterval(rateTimer);
      emit({ type: 'inference_activity', phase: 'end' });
      signal?.removeEventListener('abort', interrupt);
    }
  }
}