242 lines
8.5 KiB
JavaScript
242 lines
8.5 KiB
JavaScript
"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,
|
||
};
|