pi-ext

import { createServer } from "node:http";
import type { AddressInfo } from "node:net";
import type {
	Api,
	AssistantMessageEventStream,
	Model,
	Provider,
	TranscriptContext,
} from "@earendil-works/pi-ai";
import { openAICompletionsApi } from "@earendil-works/pi-ai/compat";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { createGratis } from "../index.ts";
import { parseLiveCatalog } from "../src/catalog.ts";

const ROUNDS = 12;
const MINIMUM_ROUND_NS = 200_000_000;
const STREAM_CHUNKS = 512;
let sink = 0;

/** Catalog shaped like live Kilo/OpenRouter: ~400 entries, a few dozen free. */
function livePayload(total: number, free: number) {
	return {
		data: Array.from({ length: total }, (_, index) => ({
			id:
				index === 0
					? "kilo-auto/free"
					: index === 1
						? "openrouter/free"
						: index < free
							? `vendor/model-${index}:free`
							: `vendor/model-${index}`,
			name: `Model ${index}`,
			context_length: 262_144,
			pricing:
				index < free
					? { prompt: "0", completion: "0" }
					: { prompt: "0.000001", completion: "0.000002" },
			architecture: {
				input_modalities: ["text", "image"],
				output_modalities: ["text"],
			},
			top_provider: { context_length: 262_144, max_completion_tokens: 32_768 },
			supported_parameters: ["tools", "reasoning", "max_tokens"],
			description: "x".repeat(400),
		})),
	};
}

const kiloPayload = livePayload(400, 16);
const openrouterPayload = livePayload(460, 17);

/**
 * Discrete SSE frames per backend API, written one `data:` line at a time like
 * the test suite's fake upstream. Each parser must see its own native format:
 * OpenAI ends with a `[DONE]` sentinel; the Google adapter (`@google/genai`)
 * has no `[DONE]` guard and JSON.parses every data line, so feeding it OpenAI
 * frames made `grts/auto`'s first hop (google) throw and abort the cascade.
 */
function openaiFrames(chunks: number): string[] {
	const frame = (delta: unknown, finish: string | null) =>
		JSON.stringify({
			id: "c",
			object: "chat.completion.chunk",
			created: 0,
			model: "m",
			choices: [{ index: 0, delta, finish_reason: finish }],
			...(finish
				? {
						usage: {
							prompt_tokens: 3,
							completion_tokens: chunks,
							total_tokens: 3 + chunks,
						},
					}
				: {}),
		});
	const frames = [frame({ role: "assistant", content: "" }, null)];
	for (let index = 0; index < chunks; index++)
		frames.push(frame({ content: "token " }, null));
	frames.push(frame({}, "stop"), "[DONE]");
	return frames;
}

function googleFrames(chunks: number): string[] {
	const frame = (text: string, finish: boolean) =>
		JSON.stringify({
			candidates: [
				{
					content: { parts: [{ text }], role: "model" },
					index: 0,
					...(finish ? { finishReason: "STOP" } : {}),
				},
			],
			...(finish
				? {
						usageMetadata: {
							promptTokenCount: 3,
							candidatesTokenCount: chunks,
							totalTokenCount: 3 + chunks,
						},
					}
				: {}),
		});
	const frames: string[] = [];
	for (let index = 0; index < chunks; index++)
		frames.push(frame("token ", false));
	frames.push(frame("", true));
	return frames;
}

async function startUpstream() {
	const openai = openaiFrames(STREAM_CHUNKS);
	const google = googleFrames(STREAM_CHUNKS);
	const server = createServer((request, response) => {
		request.resume();
		request.on("end", () => {
			// The Google adapter targets `:streamGenerateContent`; everyone else is OpenAI-shaped.
			const frames = (request.url ?? "").includes("streamGenerateContent")
				? google
				: openai;
			response.writeHead(200, { "content-type": "text/event-stream" });
			for (const data of frames) response.write(`data: ${data}\n\n`);
			response.end();
		});
	});
	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
	const base = `http://127.0.0.1:${(server.address() as AddressInfo).port}`;
	return { base, close: () => server.close() };
}

/** Load gratis against a minimal Pi API and seed its live catalogs. */
async function loadProvider(resolveUrl: (url: string) => string) {
	let provider: Provider | undefined;
	const pi = {
		registerProvider(registered: Provider) {
			provider = registered;
		},
		registerCommand() {},
		on() {},
	} as unknown as ExtensionAPI;
	createGratis({ resolveUrl })(pi);
	if (!provider?.refreshModels) throw new Error("gratis did not register");
	const stored = [
		...parseLiveCatalog("kilo", kiloPayload),
		...parseLiveCatalog("openrouter", openrouterPayload),
	];
	await provider.refreshModels({
		stored: { models: stored },
		allowNetwork: false,
		signal: new AbortController().signal,
		publish: async (publication) => {
			publication.update?.();
			return true;
		},
	});
	return provider;
}

function benchmark(name: string, operation: () => void): void {
	for (let index = 0; index < 3; index++) operation();
	let iterations = 1;
	while (true) {
		const started = process.hrtime.bigint();
		for (let index = 0; index < iterations; index++) operation();
		if (Number(process.hrtime.bigint() - started) >= MINIMUM_ROUND_NS) break;
		iterations *= 2;
	}
	for (let round = 0; round < ROUNDS; round++) {
		const started = process.hrtime.bigint();
		for (let index = 0; index < iterations; index++) operation();
		const elapsed = Number(process.hrtime.bigint() - started);
		report(name, iterations, elapsed);
	}
}

async function benchmarkAsync(
	name: string,
	operation: () => Promise<void>,
): Promise<void> {
	for (let index = 0; index < 3; index++) await operation();
	let iterations = 1;
	while (true) {
		const started = process.hrtime.bigint();
		for (let index = 0; index < iterations; index++) await operation();
		if (Number(process.hrtime.bigint() - started) >= MINIMUM_ROUND_NS) break;
		iterations *= 2;
	}
	for (let round = 0; round < ROUNDS; round++) {
		const started = process.hrtime.bigint();
		for (let index = 0; index < iterations; index++) await operation();
		const elapsed = Number(process.hrtime.bigint() - started);
		report(name, iterations, elapsed);
	}
}

function report(name: string, iterations: number, elapsed: number): void {
	console.log(
		`Benchmark${name}\t${iterations}\t${(elapsed / iterations).toFixed(0)} ns/op`,
	);
}

async function drain(stream: AssistantMessageEventStream): Promise<void> {
	const message = await stream.result();
	if (message.stopReason !== "stop")
		throw new Error(`stream failed: ${message.errorMessage}`);
	sink ^= message.content.length;
}

// Every backend keyed: the largest model set gratis can publish.
for (const name of [
	"GEMINI_API_KEY",
	"OPENROUTER_API_KEY",
	"GROQ_API_KEY",
	"MISTRAL_API_KEY",
])
	process.env[name] = "bench";

const upstream = await startUpstream();
const resolveUrl = (url: string) =>
	url.replace(/^https:\/\/[^/]+/, upstream.base);
const provider = await loadProvider(resolveUrl);
const models = provider.getModels() as Model<Api>[];
const find = (id: string) => {
	const model = models.find((entry) => entry.id === id);
	if (!model) throw new Error(`missing ${id}`);
	return model;
};
const pinnedModel = find("kilo/kilo-auto/free");
const autoModel = find("auto");
const context = {
	messages: [{ role: "user", content: "hi", timestamp: 0 }],
} as unknown as TranscriptContext;
const direct = openAICompletionsApi();

console.log(`goos: ${process.platform}`);
console.log(`goarch: ${process.arch}`);
console.log("pkg: pi-ext/gratis");
benchmark("GetModels", () => {
	sink ^= provider.getModels().length;
});
benchmark("ParseLiveCatalog", () => {
	sink ^= parseLiveCatalog("kilo", kiloPayload).length;
});
await benchmarkAsync("StreamDirect", () =>
	drain(
		direct.streamSimple(
			{ ...pinnedModel, id: "kilo-auto/free", baseUrl: upstream.base },
			context,
			{ apiKey: "bench" },
		),
	),
);
await benchmarkAsync("StreamPinned", () =>
	drain(provider.streamSimple(pinnedModel, context, {})),
);
await benchmarkAsync("StreamAuto", () =>
	drain(provider.streamSimple(autoModel, context, {})),
);
upstream.close();
if (sink === 42) console.error(sink);