mnote/scripts/lib/pi-lab-warmup.js

"use strict";

const fs = require("node:fs");
const path = require("node:path");

function isEnabledEnv(value) {
  const normalized = String(value || "").trim().toLowerCase();
  return normalized === "1" || normalized === "true" || normalized === "yes" || normalized === "on";
}

function actorSegment(actorId) {
  const value = String(actorId || "").trim().replace(/[/:\\]/g, "_");
  return value || "mnote-e2e";
}

function rootUriFromPath(rootPath) {
  return `file://${path.resolve(rootPath)}`;
}

function resolvePiLabWarmupPlan(env = process.env, rootDir = path.resolve(__dirname, "../..")) {
  const actorId = String(env.MNOTE_PAGE_AI_PI_WARMUP_ACTOR_ID || env.MNOTE_E2E_ACTOR_ID || "mnote-e2e").trim();
  const defaultRootPath = String(
    env.MNOTE_PAGE_AI_PI_WARMUP_ROOT_PATH ||
      env.MNOTE_E2E_ROOT_PATH ||
      "/mnt/Data1T/Mnote_data/users/mnote-e2e/workspaces/my-space",
  ).trim();
  const rootUri = String(
    env.MNOTE_PAGE_AI_PI_WARMUP_ROOT_URI ||
      env.MNOTE_E2E_ROOT_URI ||
      rootUriFromPath(fs.existsSync(defaultRootPath) ? defaultRootPath : rootDir),
  ).trim();
  const baseUrl = String(
    env.MNOTE_PAGE_AI_PI_WARMUP_BASE_URL ||
      env.MNOTE_UI_BASE_URL ||
      env.MNOTE_WEB_PUBLIC_URL ||
      "http://127.0.0.1:3000",
  ).replace(/\/+$/, "");

  return {
    enabled: isEnabledEnv(env.MNOTE_PAGE_AI_PI_WARMUP),
    sendEnabled: isEnabledEnv(env.MNOTE_PAGE_AI_PI_WARMUP_SEND),
    baseUrl,
    actorId,
    auth: String(env.MNOTE_PAGE_AI_PI_WARMUP_AUTH || "Bearer pi-lab-smoke").trim(),
    sessionId: String(
      env.MNOTE_PAGE_AI_PI_WARMUP_SESSION_ID || `pi_lab_dev_warm_${actorSegment(actorId)}`,
    ).trim(),
    rootUri,
    workspaceId: String(
      env.MNOTE_PAGE_AI_PI_WARMUP_WORKSPACE_ID ||
        env.MNOTE_E2E_WORKSPACE_ID ||
        `local-ws:${actorId}:my-space`,
    ).trim(),
    pagePath: String(env.MNOTE_PAGE_AI_PI_WARMUP_PAGE_PATH || "pi-lab-warmup.md").trim(),
    pageTitle: String(env.MNOTE_PAGE_AI_PI_WARMUP_PAGE_TITLE || "Pi Lab dev warmup").trim(),
    modelProvider: String(
      env.MNOTE_PAGE_AI_PI_WARMUP_MODEL_PROVIDER || env.MNOTE_PAGE_AI_PI_DEFAULT_MODEL_PROVIDER || "",
    ).trim(),
    modelId: String(
      env.MNOTE_PAGE_AI_PI_WARMUP_MODEL_ID ||
        env.MNOTE_PAGE_AI_PI_DEFAULT_MODEL ||
        env.MNOTE_PAGE_AI_PI_DEFAULT_MODEL_ID ||
        "",
    ).trim(),
    thinkingLevel: String(
      env.MNOTE_PAGE_AI_PI_WARMUP_THINKING ||
        env.MNOTE_PAGE_AI_PI_THINKING ||
        "medium",
    ).trim(),
    permissionMode: String(env.MNOTE_PAGE_AI_PI_WARMUP_PERMISSION_MODE || "plan").trim(),
    prompt: String(env.MNOTE_PAGE_AI_PI_WARMUP_PROMPT || "只回复 OK，不要解释。").trim(),
    readyTimeoutMs: Number.parseInt(env.MNOTE_PAGE_AI_PI_WARMUP_READY_TIMEOUT_MS || "180000", 10),
    requestTimeoutMs: Number.parseInt(env.MNOTE_PAGE_AI_PI_WARMUP_REQUEST_TIMEOUT_MS || "30000", 10),
    sendTimeoutMs: Number.parseInt(env.MNOTE_PAGE_AI_PI_WARMUP_SEND_TIMEOUT_MS || "120000", 10),
  };
}

function warmupHeaders(plan) {
  return {
    "Content-Type": "application/json",
    Accept: "application/json",
    Authorization: plan.auth,
    "x-mnote-actor-id": plan.actorId,
  };
}

async function sleep(ms) {
  await new Promise((resolve) => setTimeout(resolve, ms));
}

async function fetchJson(url, options, timeoutMs) {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);
  try {
    const response = await fetch(url, { ...options, signal: controller.signal });
    const text = await response.text();
    let body = {};
    try {
      body = text ? JSON.parse(text) : {};
    } catch {
      body = { raw: text };
    }
    return { response, body };
  } finally {
    clearTimeout(timer);
  }
}

async function waitForStatus(plan, log) {
  const startedAt = Date.now();
  const statusUrl = `${plan.baseUrl}/api/page-ai/pi/status`;
  let lastError = "";
  while (Date.now() - startedAt <= plan.readyTimeoutMs) {
    try {
      const result = await fetchJson(statusUrl, {
        method: "GET",
        headers: warmupHeaders(plan),
      }, Math.min(plan.requestTimeoutMs, 5000));
      if (result.response.ok) {
        return result.body;
      }
      lastError = `HTTP ${result.response.status}`;
    } catch (error) {
      lastError = error && error.message ? error.message : String(error);
    }
    await sleep(1000);
  }
  throw new Error(`等待 Pi Lab status 超时：${lastError || statusUrl}`);
}

async function waitForWarmupReply(plan, log) {
  const startedAt = Date.now();
  const eventsUrl = `${plan.baseUrl}/api/page-ai/pi/sessions/${encodeURIComponent(plan.sessionId)}/events?limit=200`;
  while (Date.now() - startedAt <= plan.sendTimeoutMs) {
    const result = await fetchJson(eventsUrl, {
      method: "GET",
      headers: warmupHeaders(plan),
    }, Math.min(plan.requestTimeoutMs, 5000));
    if (result.response.ok) {
      const events = Array.isArray(result.body.events) ? result.body.events : [];
      const terminal = events.some((event) => {
        if (!event || event.eventType !== "pi_rpc_event") return false;
        const payload = event.payload || {};
        const type = payload.type || (payload.assistantMessageEvent && payload.assistantMessageEvent.type);
        return type === "agent_end" || type === "message_end" || type === "done" || type === "error";
      });
      const firstText = events.some((event) => {
        const payload = event && event.payload ? event.payload : {};
        const msg = payload.assistantMessageEvent || payload;
        return msg.type === "text_delta" || msg.type === "text_start" || payload.type === "message";
      });
      if (terminal || firstText) return { terminal, firstText };
    }
    await sleep(1000);
  }
  log("warmup prompt 已发送，但等待首个模型事件超时；运行时仍保持可复用。");
  return { terminal: false, firstText: false, timeout: true };
}

async function runPiLabWarmup(env = process.env, options = {}) {
  const rootDir = options.rootDir || path.resolve(__dirname, "../..");
  const log = options.log || (() => undefined);
  const plan = resolvePiLabWarmupPlan(env, rootDir);
  if (!plan.enabled) return { ok: true, skipped: true, reason: "disabled" };

  const startedAt = Date.now();
  log(`等待 Pi Lab status：${plan.baseUrl}`);
  const status = await waitForStatus(plan, log);
  if (status.enabled === false) {
    log("Pi Lab 未启用，跳过 warmup。");
    return { ok: true, skipped: true, reason: "disabled_by_backend" };
  }

  const startBody = {
    sessionId: plan.sessionId,
    rootUri: plan.rootUri,
    workspaceId: plan.workspaceId,
    pagePath: plan.pagePath,
    pageTitle: plan.pageTitle,
    thinkingLevel: plan.thinkingLevel,
    permissionMode: plan.permissionMode,
  };
  if (plan.modelProvider) startBody.modelProvider = plan.modelProvider;
  if (plan.modelId) startBody.modelId = plan.modelId;

  log(`启动 Pi Lab warm session：${plan.sessionId}`);
  const start = await fetchJson(`${plan.baseUrl}/api/page-ai/pi/start`, {
    method: "POST",
    headers: warmupHeaders(plan),
    body: JSON.stringify(startBody),
  }, plan.requestTimeoutMs);
  if (!start.response.ok || start.body.ok !== true) {
    throw new Error(`Pi Lab warmup start 失败：HTTP ${start.response.status} ${JSON.stringify(start.body).slice(0, 500)}`);
  }

  if (plan.sendEnabled) {
    log("发送 Pi Lab warmup prompt；这会产生一次真实模型请求。");
    const send = await fetchJson(`${plan.baseUrl}/api/page-ai/pi/send`, {
      method: "POST",
      headers: warmupHeaders(plan),
      body: JSON.stringify({
        sessionId: plan.sessionId,
        message: plan.prompt,
      }),
    }, plan.requestTimeoutMs);
    if (!send.response.ok || send.body.accepted !== true) {
      throw new Error(`Pi Lab warmup send 失败：HTTP ${send.response.status} ${JSON.stringify(send.body).slice(0, 500)}`);
    }
    await waitForWarmupReply(plan, log);
  } else {
    log("已预启动 Pi runtime / MCP cache；未发送模型请求。");
  }

  return {
    ok: true,
    skipped: false,
    sessionId: plan.sessionId,
    sendEnabled: plan.sendEnabled,
    elapsedMs: Date.now() - startedAt,
  };
}

if (require.main === module) {
  runPiLabWarmup(process.env, {
    rootDir: path.resolve(__dirname, "../.."),
    log: (message) => console.log(`[pi-warmup] ${message}`),
  }).then((result) => {
    if (!result.skipped) {
      console.log(`[pi-warmup] 完成：session=${result.sessionId} elapsedMs=${result.elapsedMs}`);
    }
  }).catch((error) => {
    console.error(`[pi-warmup] 失败：${error.message}`);
    process.exit(1);
  });
}

module.exports = {
  isEnabledEnv,
  resolvePiLabWarmupPlan,
  runPiLabWarmup,
};