pi-ext

import { createTestSession } from "@marcfargas/pi-test-harness";
import rtk from "../index.ts";

const ROUNDS = Number(process.env.RTK_BENCH_ROUNDS ?? 10);
const ITERATIONS = Number(process.env.RTK_BENCH_ITERATIONS ?? 1);
let sink = 0;

function event(command: string, id: string) {
	return {
		type: "tool_call" as const,
		toolName: "bash" as const,
		toolCallId: id,
		input: { command },
	};
}

async function main(): Promise<void> {
	if (!Number.isInteger(ROUNDS) || ROUNDS < 1)
		throw new Error("RTK_BENCH_ROUNDS must be a positive integer");
	if (!Number.isInteger(ITERATIONS) || ITERATIONS < 1)
		throw new Error("RTK_BENCH_ITERATIONS must be a positive integer");

	const testSession = await createTestSession({ extensionFactories: [rtk] });
	const runner = testSession.session.extensionRunner;
	const rewrite = async (command: string, id: string) => {
		const call = event(command, id);
		await runner.emitToolCall(call);
		if (!call.input.command.startsWith("rtk "))
			throw new Error(`RTK did not rewrite: ${command}`);
		sink ^= call.input.command.length;
	};

	try {
		await rewrite("git status", "warmup");
		console.log(`goos: ${process.platform}`);
		console.log(`goarch: ${process.arch}`);
		console.log("pkg: pi-ext/rtk");
		for (const name of ["RepeatedRewrite", "UniqueRewrite"] as const) {
			for (let round = 0; round < ROUNDS; round++) {
				const started = process.hrtime.bigint();
				for (let index = 0; index < ITERATIONS; index++) {
					const command =
						name === "RepeatedRewrite"
							? "git status"
							: `git status # rtk-bench-${round}-${index}`;
					await rewrite(command, `${name}-${round}-${index}`);
				}
				const elapsed = Number(process.hrtime.bigint() - started);
				console.log(
					`Benchmark${name}\t${ITERATIONS}\t${(elapsed / ITERATIONS).toFixed(0)} ns/op`,
				);
			}
		}
	} finally {
		await runner.emit({ type: "session_shutdown", reason: "quit" });
		testSession.dispose();
	}
}

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