pi-ext

import { describe, expect, it } from "vitest";
import { createPerf, MIN_DEADLINE_MS } from "../src/perf.ts";

const MAX = 10_000;
const ok = (firstTokenMs: number, tokensPerSecond?: number) =>
	({ kind: "ok", firstTokenMs, tokensPerSecond }) as const;
const fail = { kind: "fail" } as const;

function clock() {
	let time = 0;
	return {
		now: () => time,
		advance: (ms: number) => {
			time += ms;
		},
	};
}

describe("gratis perf", () => {
	it("keeps the default deadline until a model has enough timed samples", () => {
		const perf = createPerf(clock().now);
		perf.record("groq", "groq/a", ok(500));
		perf.record("groq", "groq/a", ok(500));

		expect(perf.deadline("groq/a", MAX)).toBe(MAX);
		expect(perf.deadline("groq/unknown", MAX)).toBe(MAX);
	});

	it("scales the deadline to three times the typical first token, within bounds", () => {
		const perf = createPerf(clock().now);
		for (const ms of [1_500, 1_500, 1_500])
			perf.record("groq", "groq/a", ok(ms));
		for (const ms of [200, 200, 200]) perf.record("groq", "groq/fast", ok(ms));
		for (const ms of [8_000, 8_000, 8_000])
			perf.record("groq", "groq/slow", ok(ms));

		expect(perf.deadline("groq/a", MAX)).toBe(4_500);
		expect(perf.deadline("groq/fast", MAX)).toBe(MIN_DEADLINE_MS);
		expect(perf.deadline("groq/slow", MAX)).toBe(MAX);
		expect(perf.deadline("groq/a", 1_000)).toBe(1_000);
	});

	it("moves a repeatedly failing model behind its healthy siblings, stably", () => {
		const perf = createPerf(clock().now);
		const candidates = [{ id: "groq/a" }, { id: "groq/b" }, { id: "groq/c" }];
		for (let i = 0; i < 3; i++) perf.record("groq", "groq/a", fail);
		perf.record("groq", "groq/b", fail);

		expect(perf.orderModels(candidates).map((model) => model.id)).toEqual([
			"groq/b",
			"groq/c",
			"groq/a",
		]);
	});

	it("restores a demoted model as its failures age or it succeeds", () => {
		const time = clock();
		const perf = createPerf(time.now);
		const candidates = [{ id: "groq/a" }, { id: "groq/b" }];
		for (let i = 0; i < 3; i++) perf.record("groq", "groq/a", fail);
		expect(perf.orderModels(candidates)[0]?.id).toBe("groq/b");

		time.advance(30 * 60_000);
		expect(perf.orderModels(candidates)[0]?.id).toBe("groq/a");

		const fresh = createPerf(clock().now);
		for (let i = 0; i < 3; i++) fresh.record("groq", "groq/a", fail);
		fresh.record("groq", "groq/a", ok(500));
		expect(fresh.orderModels(candidates)[0]?.id).toBe("groq/a");
	});

	it("demotes a chronically failing backend behind the others", () => {
		const perf = createPerf(clock().now);
		for (let i = 0; i < 3; i++) perf.record("nvidia", "nvidia/x", fail);

		expect(perf.orderBackends(["nvidia", "google", "kilo"])).toEqual([
			"google",
			"kilo",
			"nvidia",
		]);
	});

	it("reports measured models, most used first", () => {
		const perf = createPerf(clock().now);
		perf.record("kilo", "kilo/a", ok(1_000, 40));
		perf.record("kilo", "kilo/a", ok(2_000, 60));
		perf.record("groq", "groq/b", fail);

		expect(perf.snapshot()).toEqual([
			{
				id: "kilo/a",
				samples: 2,
				failureRate: 0,
				firstTokenMs: 1_300,
				tokensPerSecond: 46,
				demoted: false,
			},
			{
				id: "groq/b",
				samples: 1,
				failureRate: expect.closeTo(0.3, 5),
				demoted: false,
			},
		]);
	});
});

describe("gratis perf votes, promotion, persistence", () => {
	it("demotes a model after two votes, past its cooldown, until votes fade", () => {
		const time = clock();
		const perf = createPerf(time.now);
		const candidates = [{ id: "kilo/a" }, { id: "kilo/b" }];
		perf.vote("kilo/a");
		expect(perf.orderModels(candidates)[0]?.id).toBe("kilo/a");

		perf.vote("kilo/a");
		time.advance(61 * 60_000);
		expect(perf.orderModels(candidates)[0]?.id).toBe("kilo/b");

		time.advance(7 * 24 * 60 * 60_000);
		expect(perf.orderModels(candidates)[0]?.id).toBe("kilo/a");
	});

	it("moves a proven backend one place past an unreliable neighbor, no further", () => {
		const perf = createPerf(clock().now);
		for (const sample of [fail, fail, ok(500)])
			perf.record("shaky", "shaky/x", sample);
		for (let i = 0; i < 20; i++) perf.record("solid", "solid/x", ok(500));
		for (let i = 0; i < 20; i++) perf.record("steady", "steady/x", ok(500));

		expect(perf.orderBackends(["top", "shaky", "solid", "steady"])).toEqual([
			"top",
			"solid",
			"shaky",
			"steady",
		]);
		expect(perf.orderBackends(["solid", "top"])).toEqual(["solid", "top"]);
	});

	it("never promotes on speed alone", () => {
		const perf = createPerf(clock().now);
		for (const sample of [fail, fail, ok(500)])
			perf.record("shaky", "shaky/x", sample);
		for (let i = 0; i < 5; i++) perf.record("fast", "fast/x", ok(50));

		expect(perf.orderBackends(["shaky", "fast"])).toEqual(["shaky", "fast"]);
	});

	it("round-trips its state and merges without overriding fresh measurements", () => {
		const time = clock();
		const first = createPerf(time.now);
		for (let i = 0; i < 3; i++) first.record("groq", "groq/a", fail);
		first.vote("groq/b");
		const saved = JSON.parse(JSON.stringify(first.export()));

		const second = createPerf(time.now);
		second.record("groq", "groq/b", ok(900));
		second.import(saved);

		expect(
			second.orderModels([{ id: "groq/a" }, { id: "groq/c" }])[0]?.id,
		).toBe("groq/c");
		expect(second.snapshot().find((m) => m.id === "groq/b")?.samples).toBe(1);
	});

	it("ignores corrupt or foreign state", () => {
		const perf = createPerf(clock().now);
		for (const state of [
			undefined,
			"nope",
			{ version: 2, models: {} },
			{ version: 1, models: { "groq/a": { failure: "x" } } },
		])
			perf.import(state);

		expect(perf.snapshot()).toEqual([]);
	});
});