pi-ext

import { spawn, spawnSync } from "node:child_process";
import { readdir, readFile } 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 piPackage = new URL(
	"../package.json",
	import.meta.resolve("@earendil-works/pi-coding-agent"),
);
const piMetadata = JSON.parse(await readFile(piPackage, "utf8")) as {
	bin: { pi: string };
};
const piCli = fileURLToPath(new URL(piMetadata.bin.pi, piPackage));
const variant = process.env["EXTENSION_STARTUP_BENCH_SOURCE"] ?? "source";
if (variant !== "source" && variant !== "dist")
	throw new Error("EXTENSION_STARTUP_BENCH_SOURCE must be source or dist");
const sdkHook = `import { registerHooks } from 'node:module';
registerHooks({ load(url, context, nextLoad) {
	const loaded = nextLoad(url, context);
	if (url.replaceAll('\\\\', '/').endsWith('/pi-coding-agent/dist/index.js')) process.stderr.write('STARTUP_SDK_LOAD=' + url + '\\n');
	return loaded;
}});`;
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,
);

const counts: [string, number][] = [
	["EXTENSION_STARTUP_BENCH_ROUNDS", rounds],
	["EXTENSION_STARTUP_BENCH_WARMUPS", warmups],
	["EXTENSION_STARTUP_BENCH_TIMEOUT_MS", timeoutMs],
];
for (const [name, value] of counts) {
	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) =>
		variant === "dist"
			? join(root, "dist", entry.name)
			: 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(): Promise<bigint> {
	const started = process.hrtime.bigint();
	const child = spawn(
		process.execPath,
		[
			"--import",
			`data:text/javascript,${encodeURIComponent(sdkHook)}`,
			piCli,
			...args,
		],
		{
			cwd: root,
			env: { ...process.env, PI_OFFLINE: "1" },
			stdio: ["pipe", "pipe", "pipe"],
		},
	);
	let stdout = "";
	let stderr = "";
	let elapsed: bigint | undefined;
	let settled = false;

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

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

		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: {
					type?: string;
					command?: string;
					success?: boolean;
					error?: string;
				};
				try {
					event = JSON.parse(line.endsWith("\r") ? line.slice(0, -1) : line);
				} catch (error) {
					finish(
						new Error(
							`Invalid Pi RPC output: ${error instanceof Error ? error.message : String(error)}`,
						),
					);
					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("close", (code, signal) => {
			if (settled) return;
			if (elapsed !== undefined) {
				if (stderr.includes("STARTUP_SDK_LOAD=")) {
					finish(
						new Error(
							`Extra native Pi SDK evaluated during ${variant} startup: ${stderr.trim()}`,
						),
					);
					return;
				}
				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(process.execPath, [piCli, "--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}`);
console.log(`source: ${variant}`);
console.log("extra-native-sdk-loads: 0 (checked per sample)");
for (let round = 0; round < rounds; round++) {
	const elapsed = await measureStartup();
	console.log(`BenchmarkExtensionStartup\t1\t${elapsed} ns/op`);
}