pi-ext

import { mkdtempSync, readdirSync, readFileSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import modelInfoExtension, { __test } from "../index.ts";
import { shortModelLabel } from "../model-label.ts";

describe("model-info pi-test-harness", () => {
	let t: TestSession | undefined;

	afterEach(() => {
		t?.dispose();
		t = undefined;
	});

	it("loads in real Pi extension runtime", async () => {
		t = await createTestSession({ extensionFactories: [modelInfoExtension] });

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

	it("writes enabled lifecycle debug records in an isolated sandbox", async () => {
		const root = mkdtempSync(join(tmpdir(), "pi-ext-model-info-debug-"));
		const original = {
			debug: process.env.PI_MODEL_INFO_DEBUG,
			state: process.env.XDG_STATE_HOME,
		};
		try {
			process.env.PI_MODEL_INFO_DEBUG = "1";
			process.env.XDG_STATE_HOME = join(root, "state");
			vi.resetModules();
			const { default: extension } = await import("../index.ts");
			t = await createTestSession({
				cwd: root,
				env: {
					PI_MODEL_INFO_DEBUG: "1",
					XDG_STATE_HOME: join(root, "state"),
				},
				extensionFactories: [extension],
			});
			t.session.extensionRunner.setUIContext(
				t.session.extensionRunner.getUIContext(),
				"tui",
			);
			await t.session.extensionRunner.emit({
				type: "session_start",
				reason: "startup",
			});
			await t.session.extensionRunner.emit({
				type: "session_shutdown",
				reason: "quit",
			});
			const events = readdirSync(
				join(root, "state", "pi-ext", "debug", "model-info"),
			).flatMap((file) =>
				readFileSync(
					join(root, "state", "pi-ext", "debug", "model-info", 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 {
			t?.dispose();
			t = undefined;
			if (original.debug === undefined) delete process.env.PI_MODEL_INFO_DEBUG;
			else process.env.PI_MODEL_INFO_DEBUG = original.debug;
			if (original.state === undefined) delete process.env.XDG_STATE_HOME;
			else process.env.XDG_STATE_HOME = original.state;
			rmSync(root, { recursive: true, force: true });
		}
	});

	it("uses acknowledged footer segments without transient fallback status", async () => {
		let accepted = 0;
		const footer = (pi: ExtensionAPI) => {
			pi.events.on("footer:segment", (data: unknown) => {
				const accept = (data as { accept?: unknown })?.accept;
				if (typeof accept !== "function") return;
				accept();
				accepted++;
			});
		};
		t = await createTestSession({
			extensionFactories: [footer, modelInfoExtension],
		});
		const runner = t.session.extensionRunner;
		runner.setUIContext(runner.getUIContext(), "tui");
		await runner.emit({ type: "session_start", reason: "startup" });

		const statuses = t.events
			.uiCallsFor("setStatus")
			.filter((call) => call.args[0] === "model-info");
		expect(accepted).toBeGreaterThan(0);
		expect(statuses.at(-1)?.args[1]).toBeUndefined();
	});

	it("keeps fallback status when only removals are acknowledged", async () => {
		let removals = 0;
		const partialFooter = (pi: ExtensionAPI) => {
			pi.events.on("footer:segment", (data: unknown) => {
				const segment = data as { text?: unknown; accept?: unknown };
				if (segment.text !== undefined || typeof segment.accept !== "function")
					return;
				segment.accept();
				removals++;
			});
		};
		t = await createTestSession({
			extensionFactories: [partialFooter, modelInfoExtension],
		});
		const runner = t.session.extensionRunner;
		runner.setUIContext(runner.getUIContext(), "tui");
		await runner.emit({ type: "session_start", reason: "startup" });

		const statuses = t.events
			.uiCallsFor("setStatus")
			.filter((call) => call.args[0] === "model-info");
		expect(removals).toBeGreaterThan(0);
		expect(statuses.at(-1)?.args[1]).toMatch(/:/);
	});

	it("replaces producer fallbacks with the combined route", async () => {
		const producer = (pi: ExtensionAPI) => {
			pi.on("session_start", (_event, ctx) => {
				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" },
				});
			});
		};
		t = await createTestSession({
			extensionFactories: [producer, modelInfoExtension],
		});

		const statuses = t.events.uiCallsFor("setStatus");
		const latest = (id: string) =>
			statuses.filter((call) => call.args[0] === id).at(-1)?.args[1];
		expect(latest("fast")).toBeUndefined();
		expect(latest("angel")).toBeUndefined();
		expect(latest("model-info")).toMatch(/→ 󰧑 astra:max$/u);
	});

	it("shows the model a router resolved to after a reply, until the model changes", async () => {
		t = await createTestSession({ extensionFactories: [modelInfoExtension] });
		const runner = t.session.extensionRunner;
		const model = runner.createContext().model;
		const latest = () =>
			t?.events
				.uiCallsFor("setStatus")
				.filter((call) => call.args[0] === "model-info")
				.at(-1)?.args[1];

		await runner.emit({
			type: "message_end",
			message: {
				role: "assistant",
				provider: model.provider,
				model: model.id,
				responseModel: "kilo/dots-studio/dots-3-note-preview:free",
			},
		});
		expect(latest()).toMatch(/ → dots-3-note-preview:free$/);

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

	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" },
			sessionManager: {
				getBranch: () => [
					{ type: "thinking_level_change", thinkingLevel: "high" },
				],
			},
		} as any;

		expect(__test.modelInfoLabel(ctx)).toBe("sol: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:medium", true, {
				modelId: "gpt-6-astra",
				thinkingLevel: "max",
			}),
		).toBe("sol:medium 󱐋 → 󰧑 astra:max");
	});

	it.each([
		["gpt-6-astra", "astra"],
		["gpt-5.5", "gpt-5.5"],
		["claude-fable-5-1", "fable"],
		["claude-opus-5", "opus"],
		["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);
	});
});