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 { AUTO_ID, BACKEND_INFO, BACKENDS } from "../src/catalog.ts";

/**
 * Real `grts/auto` against real free endpoints. This is the cascade the pinned
 * per-backend suite does not exercise: it proves whichever backend routing
 * lands on parses its own native stream, so the answer is not tied to one API.
 */

const CAPACITY =
	/\b(404|429|503)\b|not found|high demand|rate[_ -]?limit|overloaded|unavailable|aborted|no free backend/i;

// Every key present in the ambient env: the widest cascade this machine can run.
const ENV = Object.fromEntries(
	BACKENDS.flatMap((backend) => {
		const info = BACKEND_INFO[backend];
		return [info.gratisEnv, info.nativeEnv].flatMap((name) => {
			const value = process.env[name] ?? "";
			return value === "" ? [] : [[name, value]];
		});
	}),
) as Record<string, string>;

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

describe("gratis live auto", () => {
	it("grts/auto answers, and the serving backend parses its own stream", async () => {
		const agentDir = mkdtempSync(path.join(tmpdir(), "gratis-auto-live-"));
		const session = await createTestSession({
			extensionFactories: [(await import("../index.ts")).default],
			env: { PI_CODING_AGENT_DIR: agentDir, ...ENV },
		});
		try {
			const registry = session.session.extensionRunner.createContext()
				.modelRegistry as Registry;
			await registry.refresh({ allowNetwork: true, providers: ["grts"] });
			const auto = registry
				.getAll()
				.find((entry) => entry.provider === "grts" && entry.id === AUTO_ID);
			expect(auto, "grts/auto not registered").toBeDefined();
			if (!auto) return;

			const message = await registry
				.streamSimple(
					auto,
					{
						messages: [
							{
								role: "user",
								content: "Reply with the word ok.",
								timestamp: Date.now(),
							},
						],
					},
					{ maxTokens: 1_024, signal: AbortSignal.timeout(120_000) },
				)
				.result();

			if (
				(message.stopReason === "error" || message.stopReason === "aborted") &&
				CAPACITY.test(message.errorMessage ?? message.stopReason)
			)
				return; // Every free backend at capacity says nothing about routing.

			expect(
				["stop", "length"],
				`auto → ${message.responseModel ?? "?"}: ${message.errorMessage ?? message.stopReason}`,
			).toContain(message.stopReason);
			expect(message.provider).toBe("grts");
			// responseModel records the real backend/model that answered.
			expect(
				message.responseModel,
				"auto did not record a serving backend",
			).toBeTruthy();
		} finally {
			session.dispose();
			rmSync(agentDir, { recursive: true, force: true });
		}
	});
});