pi-ext

import {
	type AssistantMessage,
	type AssistantMessageEventStream,
	createAssistantMessageEventStream,
	type TranscriptContext,
} from "@earendil-works/pi-ai";
import { describe, expect, it } from "vitest";
import { AUTO_ID, type GratisModel, gratisModel } from "../src/catalog.ts";
import { createRouter } from "../src/router.ts";

// No module mocks: tests share a worker (isolate: false), so the router takes
// fake transports through its `apis` seam instead.

const model = gratisModel("kilo", {
	id: "vendor/model:free",
	name: "Model",
	api: "openai-completions",
	baseUrl: "https://api.kilo.ai/api/gateway",
	contextWindow: 262_144,
	maxTokens: 32_768,
	input: ["text"],
	reasoning: false,
});
const auto: GratisModel = { ...model, id: AUTO_ID };
const context = { messages: [] } as unknown as TranscriptContext;

function router(streamSimple: () => AssistantMessageEventStream) {
	return createRouter({
		key: () => undefined,
		candidates: (backend) => (backend === "kilo" ? [model] : []),
		resolveUrl: (url) => url,
		apis: { "openai-completions": { stream: streamSimple, streamSimple } },
	});
}

function message(stopReason: AssistantMessage["stopReason"]): AssistantMessage {
	return {
		role: "assistant",
		content: [],
		api: "openai-completions",
		provider: "kilo",
		model: "vendor/model:free",
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason,
		timestamp: 0,
	};
}

/** Upstream that starts streaming text, then ends with the given terminal event. */
function upstreamEnding(ending: "done" | "error" | "aborted") {
	return () => {
		const stream = createAssistantMessageEventStream();
		const partial = message("pending" as AssistantMessage["stopReason"]);
		stream.push({ type: "start", partial });
		partial.content.push({ type: "text", text: "thinking" });
		stream.push({ type: "text_start", contentIndex: 0, partial });
		if (ending === "done")
			stream.push({ type: "done", reason: "stop", message: message("stop") });
		else {
			const failed = message(ending);
			failed.errorMessage = "Provider finish_reason: content_filter";
			stream.push({ type: "error", reason: ending, error: failed });
		}
		stream.end();
		return stream;
	};
}

describe("gratis router", () => {
	it("turns a synchronous transport throw into a grts error message", async () => {
		const result = await router(() => {
			throw new Error("transport setup failed");
		})
			.streamSimple(model, context)
			.result();

		expect(result).toMatchObject({
			stopReason: "error",
			errorMessage: "transport setup failed",
			provider: "grts",
			model: "kilo/vendor/model:free",
			api: "openai-completions",
		});
	});

	for (const [route, requested] of [
		["pinned", model],
		["auto", auto],
	] as const) {
		it.each([
			["error", "error"],
			["aborted", "aborted"],
			["done", "ok"],
		] as const)(
			`records a ${route} stream ending in %s as %s`,
			async (ending, outcome) => {
				const instance = router(upstreamEnding(ending));

				await instance.streamSimple(requested, context).result();

				await expect
					.poll(() => instance.status().last)
					.toMatchObject({ route, outcome });
			},
		);
	}
});