pi-ext

import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai";
import { describe, expect, it } from "vitest";
import { createTestSession } from "../../../test/harness";
import gratis from "../index.ts";
import { BACKEND_INFO, BACKENDS, type Backend } from "../src/catalog.ts";

/**
 * Real endpoints, spending real (free) quota: tiny pinned requests per backend.
 * A keyed backend runs only when one of its env keys is set; keyless Kilo
 * always runs. Keys come from the process env only and are never printed.
 */

/**
 * Capacity errors and listed-but-unusable models (404) say nothing about
 * gratis; like grts/auto, try the next model, and skip if none answers.
 */
const CAPACITY =
	/\b(404|429|503)\b|not found|high demand|rate[_ -]?limit|overloaded|unavailable|aborted/i;
/** Free queues can hold a model for minutes; give each model this long. */
const PER_MODEL_MS = 30_000;
/** Try this many of a backend's models before calling it at capacity. */
const MODELS_PER_BACKEND = 3;

// Captured before any session: the harness swaps process.env per session, and
// a stuck test must not hide keys from the tests after it.
const KEYS = Object.fromEntries(
	BACKENDS.map((backend) => {
		const info = BACKEND_INFO[backend];
		const env = Object.fromEntries(
			[info.gratisEnv, info.nativeEnv].flatMap((name) => {
				const value = process.env[name] ?? "";
				return value === "" ? [] : [[name, value]];
			}),
		);
		return [backend, env];
	}),
) as Record<Backend, Record<string, string>>;

type Registry = {
	refresh(options: object): Promise<unknown>;
	getAll(): Model<Api>[];
	streamSimple(
		model: Model<Api>,
		context: unknown,
		options?: object,
	): { result(): Promise<AssistantMessage> };
};

async function askBackend(backend: Backend) {
	const agentDir = mkdtempSync(path.join(tmpdir(), "gratis-live-"));
	const session = await createTestSession({
		extensionFactories: [gratis],
		env: { PI_CODING_AGENT_DIR: agentDir, ...KEYS[backend] },
	});
	try {
		const registry = session.session.extensionRunner.createContext()
			.modelRegistry as Registry;
		await registry.refresh({ allowNetwork: true, providers: ["grts"] });
		const models = registry
			.getAll()
			.filter(
				(entry) =>
					entry.provider === "grts" && entry.id.startsWith(`${backend}/`),
			);
		expect(models.length, `no ${backend} models registered`).toBeGreaterThan(0);
		const tried: { model: Model<Api>; message: AssistantMessage }[] = [];
		for (const model of models.slice(0, MODELS_PER_BACKEND)) {
			const message = await registry
				.streamSimple(
					model,
					{
						messages: [
							{
								role: "user",
								content: "Reply with the word ok.",
								timestamp: Date.now(),
							},
						],
					},
					{ maxTokens: 1_024, signal: AbortSignal.timeout(PER_MODEL_MS) },
				)
				.result();
			tried.push({ model, message });
			if (message.stopReason !== "error" && message.stopReason !== "aborted")
				break;
			if (!CAPACITY.test(message.errorMessage ?? message.stopReason)) break;
		}
		return tried;
	} finally {
		session.dispose();
		rmSync(agentDir, { recursive: true, force: true });
	}
}

describe("gratis live backends", () => {
	for (const backend of BACKENDS) {
		const keyed = Object.keys(KEYS[backend]).length > 0;
		it.runIf(BACKEND_INFO[backend].keyless || keyed)(
			`${backend} lists models and answers a pinned request`,
			async (context) => {
				const tried = await askBackend(backend);
				const last = tried.at(-1);
				if (!last) throw new Error("no model tried");
				const { model, message } = last;
				if (
					(message.stopReason === "error" ||
						message.stopReason === "aborted") &&
					CAPACITY.test(message.errorMessage ?? message.stopReason)
				)
					context.skip(
						`${tried.map((entry) => entry.model.id).join(", ")} at upstream capacity`,
					);

				expect(model.cost).toEqual({
					input: 0,
					output: 0,
					cacheRead: 0,
					cacheWrite: 0,
				});
				expect(
					["stop", "length"],
					`${model.id}: ${message.errorMessage ?? message.stopReason}`,
				).toContain(message.stopReason);
				expect(message.provider).toBe("grts");
			},
		);
	}
});