pi-ext

import type { Api, Model } from "@earendil-works/pi-ai";
import type {
	ExtensionAPI,
	ExtensionContext,
	ModelRoute,
	ModelRouteRequest,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";

const PROVIDER = "grts";
const ID = "auto";
const MAX_TRIED = 8;
const PRIORITY = ["openrouter", "google", "mistral", "nvidia"] as const;
const MISTRAL_FREE_IDS = new Set([
	"labs-devstral-small-2512",
	"labs-leanstral-1-5-1",
]);

type GratisState = { tried: string[] };
type Request = ModelRouteRequest<GratisState>;
type ChatModel = Model<Api>;

function zeroPrice(model: ChatModel): boolean {
	const cost = model.cost;
	return (
		cost?.input === 0 &&
		cost.output === 0 &&
		cost.cacheRead === 0 &&
		cost.cacheWrite === 0
	);
}

function eligible(model: ChatModel): boolean {
	if (!zeroPrice(model)) return false;
	if (model.provider === "openrouter")
		return model.id === "openrouter/free" || model.id.endsWith(":free");
	if (model.provider === "google") return model.id.startsWith("gemma-");
	if (model.provider === "mistral") return MISTRAL_FREE_IDS.has(model.id);
	return model.provider === "nvidia";
}

function candidates(ctx: ExtensionContext): ChatModel[] {
	const available = ctx.modelRegistry.getAvailable();
	const groups = PRIORITY.map((provider) => {
		const models = available.filter(
			(model) =>
				model.provider === provider &&
				model.api !== "pi-virtual" &&
				eligible(model),
		);
		const pool = models.findIndex((model) => model.id === "openrouter/free");
		if (pool > 0) models.unshift(...models.splice(pool, 1));
		return models;
	});
	const selected: ChatModel[] = [];
	// Alternate providers before trying their next model: quotas are often provider-wide.
	for (let index = 0; index < MAX_TRIED; index++) {
		for (const group of groups) {
			const model = group[index];
			if (model) selected.push(model);
			if (selected.length === MAX_TRIED) return selected;
		}
	}
	return selected;
}

function key(model: ChatModel): string {
	return `${model.provider}/${model.id}`;
}

function routeTo(
	request: Request,
	model: ChatModel,
	state: GratisState | undefined,
): ModelRoute<GratisState> {
	return {
		model,
		thinkingLevel: request.thinkingLevel,
		...(state === undefined ? {} : { state }),
	};
}

function select(
	request: Request,
	ctx: ExtensionContext,
): ModelRoute<GratisState> {
	const end = span?.("route", { reason: request.reason });
	const available = candidates(ctx);
	if (available.length === 0) {
		end?.("error", { outcome: "none" });
		throw new Error("Gratis: no eligible free native models are available.");
	}
	const previousKey = request.previous && key(request.previous.model);
	const previous = available.find((model) => key(model) === previousKey);
	if (request.reason === "continuation" && previous) {
		end?.("finish", { outcome: "sticky" });
		return routeTo(request, previous, request.state);
	}
	const tried =
		request.reason === "retry"
			? [
					...new Set([
						...(request.state?.tried ?? []),
						...(request.failed ? [key(request.failed.model)] : []),
					]),
				].slice(-MAX_TRIED)
			: [];
	const model = available.find((candidate) => !tried.includes(key(candidate)));
	if (!model) {
		end?.("error", { outcome: "exhausted" });
		throw new Error(
			"Gratis: all eligible free native models were exhausted by Pi retries.",
		);
	}
	const state =
		request.reason === "direct" ? undefined : { tried: [...tried, key(model)] };
	end?.("finish", { outcome: "selected", candidates: available.length });
	return routeTo(request, model, state);
}

export default function gratis(pi: ExtensionAPI): void {
	pi.registerVirtualModel<GratisState>({
		provider: PROVIDER,
		id: ID,
		name: "Gratis auto",
		thinkingLevels: ["off", "low", "medium", "high"],
		route: select,
	});
	pi.on("session_start", () => dbg?.("session.start"));
	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
	});
}