pi-ext

import { spawn, spawnSync } from "node:child_process";
import { readdir } from "node:fs/promises";
import { dirname, join, resolve } from "node:path";
import { fileURLToPath } from "node:url";

// End-to-end metric: process spawn until get_state confirms RPC readiness after
// extension factories and startup handlers. Go benchmark output supports benchstat.
const root = resolve(dirname(fileURLToPath(import.meta.url)), "..");
const rounds = Number(process.env.EXTENSION_STARTUP_BENCH_ROUNDS ?? 20);
const warmups = Number(process.env.EXTENSION_STARTUP_BENCH_WARMUPS ?? 3);
const timeoutMs = Number(
	process.env.EXTENSION_STARTUP_BENCH_TIMEOUT_MS ?? 15_000,
);

for (const [name, value] of [
	["EXTENSION_STARTUP_BENCH_ROUNDS", rounds],
	["EXTENSION_STARTUP_BENCH_WARMUPS", warmups],
	["EXTENSION_STARTUP_BENCH_TIMEOUT_MS", timeoutMs],
]) {
	if (
		!Number.isSafeInteger(value) ||
		value < (name.endsWith("WARMUPS") ? 0 : 1)
	) {
		throw new Error(`${name} must be a valid benchmark count`);
	}
}

// Comma-separated extension directory names to leave out, e.g. to isolate one extension's cost.
const excluded = new Set(
	(process.env.EXTENSION_STARTUP_BENCH_EXCLUDE ?? "")
		.split(",")
		.map((name) => name.trim())
		.filter(Boolean),
);

const extensions = (
	await readdir(join(root, "extensions"), {
		withFileTypes: true,
	})
)
	.filter((entry) => entry.isDirectory() && !excluded.has(entry.name))
	.map((entry) => join(root, "extensions", entry.name, "index.ts"))
	.sort();

const args = [
	"--mode",
	"rpc",
	"--no-session",
	"--approve",
	"--no-extensions",
	"--no-context-files",
	"--no-skills",
	"--no-prompt-templates",
	"--no-themes",
	...extensions.flatMap((extension) => ["--extension", extension]),
];

async function measureStartup() {
	const started = process.hrtime.bigint();
	const child = spawn("pi", args, {
		cwd: root,
		env: { ...process.env, PI_OFFLINE: "1" },
		stdio: ["pipe", "pipe", "pipe"],
	});
	let stdout = "";
	let stderr = "";
	let elapsed;
	let settled = false;

	return new Promise((resolveMeasurement, rejectMeasurement) => {
		const timeout = setTimeout(() => {
			child.kill("SIGKILL");
			rejectMeasurement(
				new Error(`Pi did not become RPC-ready within ${timeoutMs} ms`),
			);
		}, timeoutMs);

		const finish = (error) => {
			if (settled) return;
			settled = true;
			clearTimeout(timeout);
			if (error) {
				child.kill("SIGKILL");
				rejectMeasurement(error);
			} else {
				resolveMeasurement(elapsed);
			}
		};

		child.on("error", (error) => finish(error));
		child.stderr.setEncoding("utf8");
		child.stderr.on("data", (chunk) => {
			stderr += chunk;
		});
		child.stdout.setEncoding("utf8");
		child.stdout.on("data", (chunk) => {
			stdout += chunk;
			for (;;) {
				const newline = stdout.indexOf("\n");
				if (newline < 0) break;
				const line = stdout.slice(0, newline);
				stdout = stdout.slice(newline + 1);
				let event;
				try {
					event = JSON.parse(line.endsWith("\r") ? line.slice(0, -1) : line);
				} catch (error) {
					finish(new Error(`Invalid Pi RPC output: ${error.message}`));
					return;
				}
				if (event.type === "response" && event.command === "get_state") {
					if (!event.success) {
						finish(new Error(`Pi get_state failed: ${event.error}`));
						return;
					}
					elapsed = process.hrtime.bigint() - started;
					return;
				}
			}
		});
		child.on("exit", (code, signal) => {
			if (settled) return;
			if (elapsed !== undefined) {
				finish();
				return;
			}
			finish(
				new Error(
					`Pi exited before RPC readiness (code=${code}, signal=${signal})${stderr ? `: ${stderr.trim()}` : ""}`,
				),
			);
		});
		child.stdin.end('{"type":"get_state"}\n');
	});
}

for (let index = 0; index < warmups; index++) await measureStartup();

const piVersion = spawnSync("pi", ["--version"], { encoding: "utf8" });
if (piVersion.status !== 0) {
	throw new Error(`Unable to read Pi version: ${piVersion.stderr.trim()}`);
}

console.log(`goos: ${process.platform}`);
console.log(`goarch: ${process.arch}`);
console.log("pkg: pi-ext/startup");
console.log(`node: ${process.version}`);
console.log(`pi: ${piVersion.stdout.trim()}`);
console.log(`extensions: ${extensions.length}`);
for (let round = 0; round < rounds; round++) {
	const elapsed = await measureStartup();
	console.log(`BenchmarkExtensionStartup\t1\t${elapsed} ns/op`);
}