pi-ext

import {
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import {
	type AssistantMessage,
	createFauxCore,
	fauxAssistantMessage,
} from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

const extensionPath = fileURLToPath(new URL("../index.ts", import.meta.url));
const free = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
const paid = { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 };

function models(provider: string) {
	return [
		{
			id:
				provider === "openrouter"
					? "qwen/test:free"
					: provider === "google"
						? "gemma-3-27b-it"
						: "test-free",
			name: `${provider} free`,
			reasoning: false,
			input: ["text" as const],
			cost: free,
			contextWindow: 32_000,
			maxTokens: 1_024,
		},
	];
}

function nativeProvider(
	provider: string,
	cost = free,
	id?: string,
	additionalIds: string[] = [],
) {
	const core = createFauxCore({ api: "openai", provider });
	const model = {
		...models(provider)[0]!,
		id: id ?? models(provider)[0]!.id,
		cost,
	};
	return {
		core,
		factory: (pi: ExtensionAPI) =>
			pi.registerProvider(provider, {
				api: "openai-completions",
				apiKey: "test-key",
				baseUrl: "https://faux.invalid",
				models: [model, ...additionalIds.map((id) => ({ ...model, id }))],
				streamSimple: (requestModel, context, options) =>
					core.streamSimple(requestModel, context, options),
			}),
	};
}

let session: PiSession | undefined;
let debugRoot: string | undefined;

afterEach(async () => {
	await session?.dispose();
	session = undefined;
	if (debugRoot) rmSync(debugRoot, { recursive: true, force: true });
	debugRoot = undefined;
});

async function setup(providers: ReturnType<typeof nativeProvider>[]) {
	session = await createPiSession({
		env: {
			GEMINI_API_KEY: undefined,
			GOOGLE_API_KEY: undefined,
			MISTRAL_API_KEY: undefined,
			NVIDIA_API_KEY: undefined,
			OPENROUTER_API_KEY: undefined,
		},
		extensionPaths: [extensionPath],
		extensionFactories: providers.map((provider) => provider.factory),
	});
	expect(session.errors).toEqual([]);
	const virtual = session.session.extensionRunner
		.createContext()
		.modelRegistry.find("grts", "auto");
	if (!virtual) throw new Error("missing grts/auto");
	return virtual;
}

function physicalMessages() {
	return session!.session.agent.state.messages.filter(
		(message) => message.role === "assistant",
	) as AssistantMessage[];
}

describe("gratis extension contract", () => {
	it("loads path-loaded, registers grts/auto, and dispatches a native physical model", async () => {
		const openrouter = nativeProvider("openrouter");
		const virtual = await setup([openrouter]);
		openrouter.core.setResponses([fauxAssistantMessage("ok")]);

		const message = await session!.session.extensionRunner
			.createContext()
			.modelRegistry.streamSimple(virtual, { messages: [] })
			.result();

		expect(message).toMatchObject({
			provider: "openrouter",
			model: "qwen/test:free",
		});
		expect(openrouter.core.state.callCount).toBe(1);
	});

	it.each([
		[free, "auto"],
		[paid, "qwen/paid:free"],
		[{ ...free, cacheWrite: 1 }, "qwen/cache-priced:free"],
	])("excludes paid or non-free OpenRouter entry %j/%s", async (cost, id) => {
		const virtual = await setup([nativeProvider("openrouter", cost, id)]);

		await expect(
			session!.session.extensionRunner
				.createContext()
				.modelRegistry.streamSimple(virtual, { messages: [] })
				.result(),
		).resolves.toMatchObject({
			stopReason: "error",
			errorMessage: "Gratis: no eligible free native models are available.",
		});
	});

	it.each([
		["openrouter", "openrouter/free"],
		["google", "gemma-4-31b-it"],
		["mistral", "labs-devstral-small-2512"],
		["mistral", "labs-leanstral-1-5-1"],
		["nvidia", "openai/gpt-oss-20b"],
	])("dispatches eligible %s/%s through native auth", async (provider, id) => {
		const physical = nativeProvider(provider, free, id);
		const virtual = await setup([physical]);
		physical.core.setResponses([fauxAssistantMessage("ok")]);
		const message = await session!.session.extensionRunner
			.createContext()
			.modelRegistry.streamSimple(virtual, { messages: [] })
			.result();
		expect(message).toMatchObject({ provider, model: id, stopReason: "stop" });
	});

	it.each([
		["faux", "model"],
		["zai", "glm-5.3-highspeed"],
		["mistral", "unknown-zero-priced"],
		["openrouter", "unknown-zero-priced"],
		["google", "gemini-paid-tier"],
	])("rejects unsupported free inference for %s/%s", async (provider, id) => {
		const virtual = await setup([nativeProvider(provider, free, id)]);
		const message = await session!.session.extensionRunner
			.createContext()
			.modelRegistry.streamSimple(virtual, { messages: [] })
			.result();
		expect(message).toMatchObject({ stopReason: "error" });
	});

	it("keeps tool follow-ups on the previous eligible physical model", async () => {
		const openrouter = nativeProvider("openrouter");
		const google = nativeProvider("google");
		const virtual = await setup([openrouter, google]);
		const file = path.join(session!.cwd, "note.txt");
		writeFileSync(file, "test evidence");
		openrouter.core.setResponses([
			fauxAssistantMessage(
				[
					{
						type: "toolCall",
						id: "read-note",
						name: "read",
						arguments: { path: file },
					},
				],
				{ stopReason: "toolUse" },
			),
			fauxAssistantMessage("finished"),
		]);
		await session!.session.setModel(virtual);
		await session!.session.prompt("read the note");
		expect(openrouter.core.state.callCount).toBe(2);
		expect(google.core.state.callCount).toBe(0);
		expect(session!.session.model).toMatchObject({
			provider: "grts",
			id: "auto",
		});
		expect(physicalMessages().at(-1)).toMatchObject({ provider: "openrouter" });
	});

	it("does not fail over when Pi retries are disabled", async () => {
		const openrouter = nativeProvider("openrouter");
		const google = nativeProvider("google");
		const virtual = await setup([openrouter, google]);
		openrouter.core.setResponses([
			fauxAssistantMessage("down", {
				stopReason: "error",
				errorMessage: "429 rate limit",
			}),
		]);
		session!.session.setAutoRetryEnabled(false);
		await session!.session.setModel(virtual);
		await session!.session.prompt("one request");
		expect(openrouter.core.state.callCount).toBe(1);
		expect(google.core.state.callCount).toBe(0);
	});

	it("advances only when Pi retries", async () => {
		const openrouter = nativeProvider("openrouter");
		const google = nativeProvider("google");
		const virtual = await setup([openrouter, google]);
		openrouter.core.setResponses([
			fauxAssistantMessage("down", {
				stopReason: "error",
				errorMessage: "429 rate limit",
			}),
		]);
		google.core.setResponses([fauxAssistantMessage("recovered")]);

		session!.session.setAutoRetryEnabled(true);
		await session!.session.setModel(virtual);
		await session!.session.sendUserMessage("retry once");
		await session!.session.waitForIdle();

		expect(openrouter.core.state.callCount).toBe(1);
		expect(google.core.state.callCount).toBe(1);
		expect(physicalMessages().at(-1)).toMatchObject({
			provider: "google",
			model: "gemma-3-27b-it",
		});
		const before = session!.session.sessionManager.getBranch();
		openrouter.core.setResponses([fauxAssistantMessage("direct")]);
		await session!.session.extensionRunner
			.createContext()
			.modelRegistry.streamSimple(virtual, { messages: [] })
			.result();
		expect(session!.session.sessionManager.getBranch()).toEqual(before);
		openrouter.core.setResponses([fauxAssistantMessage("new turn")]);
		await session!.session.prompt("fresh user turn");
		expect(physicalMessages().at(-1)).toMatchObject({ provider: "openrouter" });
	});

	it("alternates providers and exhausts a pool of at most eight candidates", async () => {
		const openrouter = nativeProvider(
			"openrouter",
			free,
			"openrouter/free",
			Array.from({ length: 12 }, (_, index) => `test/${index}:free`),
		);
		const google = nativeProvider("google");
		const virtual = await setup([openrouter, google]);
		const failure = () =>
			fauxAssistantMessage("down", {
				stopReason: "error",
				errorMessage: "429 rate limit",
			});
		openrouter.core.setResponses(Array.from({ length: 12 }, failure));
		google.core.setResponses([failure()]);
		session!.session.settingsManager.applyOverrides({
			retry: {
				enabled: true,
				maxRetries: 10,
				baseDelayMs: 1,
				maxAgentDelayMs: 1,
			},
		});
		await session!.session.setModel(virtual);
		await session!.session.prompt("bounded failure");
		expect(openrouter.core.state.callCount).toBe(7);
		expect(google.core.state.callCount).toBe(1);
		expect(physicalMessages().at(-1)).toMatchObject({
			errorMessage:
				"Gratis: all eligible free native models were exhausted by Pi retries.",
		});
	});

	it("reports exhausted retries and does not leak tried state into another session", async () => {
		const openrouter = nativeProvider("openrouter");
		const virtual = await setup([openrouter]);
		openrouter.core.setResponses([
			fauxAssistantMessage("down", {
				stopReason: "error",
				errorMessage: "429 rate limit",
			}),
		]);
		session!.session.setAutoRetryEnabled(true);
		await session!.session.setModel(virtual);
		await session!.session.sendUserMessage("exhaust");
		await session!.session.waitForIdle();
		expect(physicalMessages().at(-1)).toMatchObject({
			errorMessage:
				"Gratis: all eligible free native models were exhausted by Pi retries.",
		});
		await session!.dispose();
		session = undefined;

		const next = nativeProvider("openrouter");
		await setup([next]);
		next.core.setResponses([fauxAssistantMessage("fresh")]);
		const nextVirtual = session!.session.extensionRunner
			.createContext()
			.modelRegistry.find("grts", "auto");
		if (!nextVirtual) throw new Error("missing grts/auto");
		await session!.session.setModel(nextVirtual);
		await session!.session.sendUserMessage("fresh branch");
		await session!.session.waitForIdle();
		expect(next.core.state.callCount).toBe(1);
	});

	it("writes lifecycle debug records under sandboxed LOCALAPPDATA", async () => {
		const openrouter = nativeProvider("openrouter");
		debugRoot = mkdtempSync(path.join(tmpdir(), "gratis-local-app-data-"));
		session = await createPiSession({
			env: { PI_GRATIS_DEBUG: "1", LOCALAPPDATA: debugRoot },
			extensionPaths: [extensionPath],
			extensionFactories: [openrouter.factory],
		});
		const localAppData = debugRoot;
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});
		await session.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		const directory =
			process.platform === "win32"
				? path.join(localAppData, "pi-ext", "debug", "gratis")
				: process.platform === "darwin"
					? path.join(
							session.cwd,
							".home",
							"Library",
							"Logs",
							"pi-ext",
							"gratis",
						)
					: path.join(
							session.env["XDG_STATE_HOME"]!,
							"pi-ext",
							"debug",
							"gratis",
						);
		const records = readFileSync(
			path.join(directory, readdirSync(directory)[0]!),
			"utf8",
		);
		expect(records).toContain('"event":"session.start"');
		expect(records).toContain('"event":"session.shutdown"');
	});
});