pi-ext

import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { hostname, tmpdir } from "node:os";
import { join, resolve } from "node:path";
import type {
	ExtensionAPI,
	ExtensionContext,
	ResourcesDiscoverEvent,
	ResourcesDiscoverResult,
} from "@earendil-works/pi-coding-agent";
import { createTestSession } from "@marcfargas/pi-test-harness";
import skillz from "../index.ts";

const COUNT = Number(process.env.SKILLZ_BENCH_SKILLS ?? 256);
const ROUNDS = Number(process.env.SKILLZ_BENCH_ROUNDS ?? 12);
const ITERATIONS = Number(process.env.SKILLZ_BENCH_ITERATIONS ?? 5);
const KIND = process.env.SKILLZ_BENCH_KIND ?? "both";
if (!["both", "discover", "reload"].includes(KIND))
	throw new Error("Unknown SKILLZ_BENCH_KIND");
if (
	![COUNT, ROUNDS, ITERATIONS].every((n) => Number.isSafeInteger(n) && n > 0) ||
	COUNT % 4 !== 0
)
	throw new Error(
		"SKILLZ_BENCH_SKILLS must be divisible by four; all parameters must be positive integers",
	);

const root = mkdtempSync(join(tmpdir(), "skillz-bench-"));
const agentDir = join(root, "agent");
const body = "x".repeat(4096);
const extension = resolve(import.meta.dirname, "../index.ts");
const previousAgentDir = process.env.PI_CODING_AGENT_DIR;
let session: Awaited<ReturnType<typeof createTestSession>> | undefined;

try {
	writeFileSync(join(root, "marker"), "");
	for (let index = 0; index < COUNT; index++) {
		const dir = join(
			agentDir,
			"skillz",
			`skill-${String(index).padStart(4, "0")}`,
		);
		mkdirSync(dir, { recursive: true });
		const condition = [
			`host: ${hostname()}`,
			"path: marker",
			"host: definitely-not-this-host",
			"os: [lnx, mac, win32]",
		][index % 4];
		writeFileSync(
			join(dir, "SKILL.md"),
			`---\nname: skill-${String(index).padStart(4, "0")}\ndescription: Benchmark skill ${index}\n${condition}\n---\n\n${body}\n`,
		);
	}
	process.env.PI_CODING_AGENT_DIR = agentDir;
	session = await createTestSession({ cwd: root, extensions: [extension] });
	const current = session.session;
	let handler:
		| ((
				event: ResourcesDiscoverEvent,
				ctx: ExtensionContext,
		  ) => ResourcesDiscoverResult | Promise<ResourcesDiscoverResult>)
		| undefined;
	skillz({
		on: (_event: string, callback: typeof handler) => {
			handler = callback;
		},
		getCommands: () => [],
	} as unknown as ExtensionAPI);
	if (!handler) throw new Error("No discovery handler registered");
	const discoverHandler = handler;
	const context = {
		mode: "json",
		ui: { notify: () => {} },
	} as unknown as ExtensionContext;
	const discover = async () => {
		const result = await discoverHandler(
			{ type: "resources_discover", cwd: root, reason: "reload" },
			context,
		);
		if (result.skillPaths?.length !== (COUNT * 3) / 4)
			throw new Error(
				`Unexpected eligible skill count: ${result.skillPaths?.length}`,
			);
	};
	const reload = async () => {
		await current.reload();
		const count = current.resourceLoader
			.getSkills()
			.skills.filter((skill) =>
				skill.filePath.startsWith(join(agentDir, "skillz")),
			).length;
		if (count !== (COUNT * 3) / 4)
			throw new Error(`Unexpected matching skill count: ${count}`);
	};
	for (let index = 0; index < 20; index++) {
		if (KIND !== "reload") await discover();
		if (KIND !== "discover") await reload();
	}
	console.log(`goos: ${process.platform}`);
	console.log(`goarch: ${process.arch}`);
	console.log("pkg: pi-ext/skillz");
	for (const [name, operation] of [
		["Discover", discover],
		["Reload", reload],
	] as const) {
		if (KIND !== "both" && KIND.toLowerCase() !== name.toLowerCase()) continue;
		for (let round = 0; round < ROUNDS; round++) {
			const start = process.hrtime.bigint();
			for (let index = 0; index < ITERATIONS; index++) await operation();
			const elapsed = Number(process.hrtime.bigint() - start) / ITERATIONS;
			console.log(
				`Benchmark${name}${COUNT}\t${ITERATIONS}\t${elapsed.toFixed(0)} ns/op`,
			);
		}
	}
} finally {
	session?.dispose();
	if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR;
	else process.env.PI_CODING_AGENT_DIR = previousAgentDir;
	rmSync(root, { recursive: true, force: true });
}