pi-ext

import { mkdtempSync, readdirSync, readFileSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import type { MessageEndEvent } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, type MockInstance, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { __test } from "../index.ts";
import { shortModelLabel } from "../model-label.ts";
import type { ProbeHooks } from "./probe.ts";

function extensionPath(): string {
	return fileURLToPath(new URL("../index.ts", import.meta.url));
}

function probePath(): string {
	return fileURLToPath(new URL("./probe.ts", import.meta.url));
}

function setProbe(hooks: ProbeHooks): void {
	(globalThis as { __modelInfoProbe?: ProbeHooks }).__modelInfoProbe = hooks;
}

function clearProbe(): void {
	delete (globalThis as { __modelInfoProbe?: ProbeHooks }).__modelInfoProbe;
}

function captureStatus(
	session: PiSession,
	mode: "tui" | "print",
): MockInstance {
	const runner = session.session.extensionRunner;
	const ui = { ...runner.getUIContext() };
	const spy = vi.spyOn(ui, "setStatus");
	runner.setUIContext(ui, mode);
	return spy;
}

function latestStatus(spy: MockInstance, id: string): unknown {
	const calls = spy.mock.calls as Array<[string, unknown?]>;
	return calls.filter((call) => call[0] === id).at(-1)?.[1];
}

describe("model-info contract", () => {
	let session: PiSession | undefined;

	afterEach(async () => {
		await session?.dispose();
		session = undefined;
		clearProbe();
		vi.restoreAllMocks();
	});

	it("loads in real Pi extension runtime", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });

		expect(session.extensions).toHaveLength(1);
	});

	it("writes enabled lifecycle debug records in an isolated sandbox", async () => {
		const root = mkdtempSync(join(tmpdir(), "pi-ext-model-info-debug-"));
		try {
			session = await createPiSession({
				cwd: root,
				extensionPaths: [extensionPath()],
				env: {
					PI_MODEL_INFO_DEBUG: "1",
					XDG_STATE_HOME: join(root, "state"),
					LOCALAPPDATA: join(root, "state"),
				},
			});
			const runner = session.session.extensionRunner;
			runner.setUIContext({ ...runner.getUIContext() }, "tui");
			await runner.emit({ type: "session_start", reason: "startup" });
			await runner.emit({ type: "session_shutdown", reason: "quit" });
			const debugDir =
				process.platform === "darwin"
					? join(root, ".home", "Library", "Logs", "pi-ext", "model-info")
					: join(root, "state", "pi-ext", "debug", "model-info");
			const events = readdirSync(debugDir).flatMap((file) =>
				readFileSync(join(debugDir, file), "utf8")
					.trim()
					.split("\n")
					.map((line) => JSON.parse(line)),
			);
			expect(events).toEqual(
				expect.arrayContaining([
					expect.objectContaining({
						extension: "model-info",
						event: "session.start",
					}),
					expect.objectContaining({
						extension: "model-info",
						event: "session.shutdown",
					}),
					expect.objectContaining({
						event: "route.refresh",
						outcome: "fallback",
					}),
				]),
			);
		} finally {
			await session?.dispose();
			session = undefined;
			rmSync(root, { recursive: true, force: true });
		}
	});

	it("uses acknowledged footer segments without transient fallback status", async () => {
		let accepted = 0;
		setProbe({
			onFooterSegment: (data: unknown) => {
				const accept = (data as { accept?: unknown })?.accept;
				if (typeof accept !== "function") return;
				(accept as () => void)();
				accepted++;
			},
		});
		session = await createPiSession({
			extensionPaths: [probePath(), extensionPath()],
		});
		const spy = captureStatus(session, "tui");
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});

		expect(accepted).toBeGreaterThan(0);
		expect(latestStatus(spy, "model-info")).toBeUndefined();
	});

	it("keeps fallback status when only removals are acknowledged", async () => {
		let removals = 0;
		setProbe({
			onFooterSegment: (data: unknown) => {
				const segment = data as { text?: unknown; accept?: unknown };
				if (segment.text !== undefined || typeof segment.accept !== "function")
					return;
				(segment.accept as () => void)();
				removals++;
			},
		});
		session = await createPiSession({
			extensionPaths: [probePath(), extensionPath()],
		});
		const spy = captureStatus(session, "tui");
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});

		expect(removals).toBeGreaterThan(0);
		expect(latestStatus(spy, "model-info")).toMatch(/:/);
	});

	it("replaces producer fallbacks with the combined route", async () => {
		setProbe({
			onSessionStart: (ctx, pi) => {
				ctx.ui.setStatus("fast", "fast fallback");
				pi.events.emit("model-route:fast", { active: true });
				ctx.ui.setStatus("angel", "angel fallback");
				pi.events.emit("model-route:angel", {
					advisor: { modelId: "gpt-6-astra", thinkingLevel: "max" },
				});
			},
		});
		session = await createPiSession({
			extensionPaths: [probePath(), extensionPath()],
		});
		const spy = captureStatus(session, "print");
		await session.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});

		expect(latestStatus(spy, "fast")).toBeUndefined();
		expect(latestStatus(spy, "angel")).toBeUndefined();
		expect(latestStatus(spy, "model-info")).toMatch(/→ 󰧑 astra-6:max$/u);
	});

	it("shows the model a router resolved to after a reply, until the model changes", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		const runner = session.session.extensionRunner;
		const spy = captureStatus(session, "print");
		const model = runner.createContext().model;
		if (!model) throw new Error("fixture model missing");

		await runner.emitMessageEnd({
			type: "message_end",
			message: {
				role: "assistant",
				provider: model.provider,
				model: model.id,
				responseModel: "kilo/dots-studio/dots-3-note-preview:free",
			} as unknown as MessageEndEvent["message"],
		});
		expect(latestStatus(spy, "model-info")).toMatch(
			/ → dots-3-note-preview:free$/,
		);

		await runner.emit({
			type: "model_select",
			model,
			previousModel: model,
			source: "set",
		});
		expect(latestStatus(spy, "model-info")).not.toMatch(/→/);
	});

	it("reports native virtual dispatch without synthetic responseModel metadata", () => {
		const ctx = {
			model: { api: "pi-virtual", provider: "grts", id: "auto" },
		} as never;
		const reply = {
			role: "assistant",
			provider: "google",
			model: "gemma-4-31b-it",
			stopReason: "stop",
		};
		expect(__test.servedModel(reply, ctx)).toBe("google/gemma-4-31b-it");
		expect(
			__test.servedModel({ ...reply, stopReason: "error" }, ctx),
		).toBeUndefined();
		expect(
			__test.servedModel({ ...reply, stopReason: "aborted" }, ctx),
		).toBeUndefined();
		expect(
			__test.servedModel({ ...reply, provider: "grts", model: "auto" }, ctx),
		).toBeUndefined();
		expect(
			__test.servedModel(
				{
					...reply,
					provider: "openrouter",
					model: "openrouter/free",
					responseModel: "qwen/test:free",
				},
				ctx,
			),
		).toBe("openrouter/qwen/test:free");
	});

	it("reports a served model only when it differs from the selected one", () => {
		const ctx = { model: { provider: "grts", id: "auto" } } as never;
		const reply = (extra: object) => ({
			role: "assistant",
			provider: "grts",
			model: "auto",
			...extra,
		});

		expect(
			__test.servedModel(reply({ responseModel: "kilo/x:free" }), ctx),
		).toBe("kilo/x:free");
		expect(__test.servedModel(reply({}), ctx)).toBeUndefined();
		expect(
			__test.servedModel(reply({ responseModel: "auto" }), ctx),
		).toBeUndefined();
		expect(
			__test.servedModel(
				reply({ provider: "openrouter", responseModel: "a/b" }),
				ctx,
			),
		).toBeUndefined();
		expect(
			__test.servedModel({ role: "user", responseModel: "a/b" }, ctx),
		).toBeUndefined();
	});

	it("restores the served model from the last assistant reply in the branch", () => {
		const ctx = (branch: object[]) =>
			({
				model: { provider: "grts", id: "auto" },
				sessionManager: { getBranch: () => branch },
			}) as never;
		const assistant = (responseModel?: string) => ({
			type: "message",
			message: {
				role: "assistant",
				provider: "grts",
				model: "auto",
				responseModel,
			},
		});

		expect(
			__test.servedFromBranch(
				ctx([
					assistant("kilo/old:free"),
					assistant("google/gemini-3.8-flash"),
					{ type: "message", message: { role: "user" } },
				]),
			),
		).toBe("google/gemini-3.8-flash");
		expect(__test.servedFromBranch(ctx([assistant()]))).toBeUndefined();
	});

	it("formats model and thinking", () => {
		const ctx = {
			cwd: "/tmp",
			model: { provider: "test", id: "openai-codex/gpt-5.6-sol" },
		} as any;

		expect(__test.modelInfoLabel(ctx, "high")).toBe("sol-5.6:high");
		expect(__test.providerLabel(ctx)).toBe("test");
		expect(__test.providerIcon(ctx)).toBeUndefined();
	});

	it.each([
		["openai-codex", ""],
		["anthropic", ""],
		["klaus", ""],
		["github-copilot", ""],
		["google-gemini-cli", ""],
		["hetzner", "󰰀"],
		["grts", "󰙺"],
	])("uses the verified brand icon for %s", (provider, icon) => {
		expect(__test.providerIcon({ model: { provider } } as any)).toBe(icon);
	});

	it("detects declared OAuth subscriptions and Kimi Coding", () => {
		const modelRegistry = {
			isUsingOAuth: () => true,
			getProvider: () => ({ auth: { oauth: { isSubscription: true } } }),
		};
		expect(
			__test.usesSubscription({
				model: { provider: "anthropic" },
				modelRegistry,
			} as any),
		).toBe(true);
		expect(
			__test.usesSubscription({
				model: { provider: "kimi-coding" },
				modelRegistry: {},
			} as any),
		).toBe(true);
	});

	it("combines executor, transport, and advisor into one route", () => {
		expect(
			__test.statusLabel("sol-5.6:medium", true, {
				modelId: "gpt-6-astra",
				thinkingLevel: "max",
			}),
		).toBe("sol-5.6:medium 󱐋 → 󰧑 astra-6:max");
	});

	it.each([
		["gpt-6-astra", "astra-6"],
		["gpt-5.6-sol", "sol-5.6"],
		["gpt-5.5", "gpt-5.5"],
		["claude-fable-5-1", "fable-5.1"],
		["claude-opus-5", "opus-5"],
		["claude-opus-5-5", "opus-5.5"],
		["claude-opus-4-5-20251101", "opus-4.5"],
		["gemini-3.1-pro-preview", "3.1-pro"],
		["grok-4.6", "grok-4.6"],
		["deepseek-v4-flash-vision-exp", "v4-flash"],
		["MiniMax-M3", "M3"],
		["kimi-k2.7-code-highspeed", "k2.7-code"],
		["glm-5.3-highspeed", "glm-5.3"],
		["mistral-medium-2604", "mistral-medium"],
		["new-family-7-alpha", "new-family-7-alpha"],
	])("shortens %s by model-family grammar", (modelId, expected) => {
		expect(shortModelLabel(modelId)).toBe(expected);
	});
});