pi-ext

import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import type {
	BuildSystemPromptOptions,
	ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { createTestSession } from "@marcfargas/pi-test-harness";
import pfui from "../index.ts";

const ROUNDS = 12;
const WARMUP_ITERATIONS = 25_000;
const ITERATIONS = 25_000;
const cwd = mkdtempSync(path.join(tmpdir(), "pfui-bench-"));
mkdirSync(path.join(cwd, ".pi", "agent"), { recursive: true });

const skills = Array.from({ length: 40 }, (_, index) => ({
	name: `skill-${index}`,
	description: `Use for focused workflow ${index} when its specialized guidance applies.`,
	filePath: `/home/user/.pi/agent/skills/skill-${index}/SKILL.md`,
	baseDir: `/home/user/.pi/agent/skills/skill-${index}`,
	sourceInfo: {
		path: `/home/user/.pi/agent/skills/skill-${index}/SKILL.md`,
		source: "local",
		scope: "user" as const,
		origin: "top-level" as const,
	},
}));
const options: BuildSystemPromptOptions = {
	cwd: "/home/user/work/project",
	selectedTools: ["read", "grep", "ls", "bash", "edit", "write"],
	toolSnippets: {
		read: "Read text files and images",
		grep: "Search file contents",
		ls: "List directory contents",
		bash: "Execute a shell command",
		edit: "Apply exact text replacements",
		write: "Create or overwrite a file",
	},
	promptGuidelines: [
		"Use read before editing files.",
		"Prefer exact, minimal changes.",
		"Preserve unrelated dirty work.",
		"Run the narrowest relevant checks.",
		"Use read before editing files.",
	],
	appendSystemPrompt: "# Response style\n\nAnswer directly.\n".repeat(8),
	contextFiles: Array.from({ length: 4 }, (_, index) => ({
		path: `/home/user/work/project${index ? `/module-${index}` : ""}/AGENTS.md`,
		content: `# Project guidance ${index}\n\n${"Keep changes focused and verify observable behavior.\n".repeat(40)}`,
	})),
	skills,
};
let sink = 0;

function consume(prompt: string | undefined): void {
	if (!prompt) throw new Error("PFUI did not compose a fallback prompt");
	sink ^= prompt.length;
	sink ^= prompt.charCodeAt(prompt.length >> 1);
}

async function main(): Promise<void> {
	let renderedPrompt: string | undefined;
	const observePrompt = (pi: ExtensionAPI) => {
		pi.on("before_agent_start", (event) => {
			renderedPrompt = event.systemPrompt;
		});
	};
	const testSession = await createTestSession({
		cwd,
		extensionFactories: [pfui, observePrompt],
	});
	const runner = testSession.session.extensionRunner;
	const compose = async () => {
		renderedPrompt = undefined;
		const result = await runner.emitBeforeAgentStart(
			"benchmark prompt",
			undefined,
			options,
		);
		if (!result.systemPromptOptions.customPrompt)
			throw new Error("PFUI did not select the fallback prompt");
		consume(renderedPrompt);
	};

	try {
		for (let index = 0; index < WARMUP_ITERATIONS; index++) await compose();
		console.log(`goos: ${process.platform}`);
		console.log(`goarch: ${process.arch}`);
		console.log("pkg: pi-ext/pfui");
		for (let round = 0; round < ROUNDS; round++) {
			const started = process.hrtime.bigint();
			for (let index = 0; index < ITERATIONS; index++) await compose();
			const elapsed = Number(process.hrtime.bigint() - started);
			console.log(
				`BenchmarkFallbackComposition\t${ITERATIONS}\t${(elapsed / ITERATIONS).toFixed(0)} ns/op`,
			);
		}
	} finally {
		await runner.emit({ type: "session_shutdown", reason: "quit" });
		testSession.dispose();
		rmSync(cwd, { recursive: true, force: true });
	}
}

await main();
if (sink === 42) console.error(sink);