pi-ext

import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { visibleWidth } from "@earendil-works/pi-tui";
import { afterEach, describe, expect, it } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import askExtension, { __test } from "../index.ts";

let t: TestSession | undefined;

const KEYS: Record<string, string> = {
	"tui.select.confirm": "\r",
	"tui.select.cancel": "\x1b",
	"tui.select.up": "\x1b[A",
	"tui.select.down": "\x1b[B",
	"tui.input.tab": "\t",
};

function tuiSession() {
	let component:
		| {
				handleInput: (data: string) => void;
				handleMouse?: (event: {
					type: "press" | "click";
					button: "left";
					x: number;
					y: number;
					screenX: number;
					screenY: number;
					width: number;
					height: number;
					shift: boolean;
					alt: boolean;
					ctrl: boolean;
				}) => { handled?: boolean } | undefined;
				render: (width: number) => string[];
		  }
		| undefined;
	const custom = (
		factory: (
			tui: unknown,
			theme: unknown,
			keybindings: unknown,
			done: (value: never) => void,
		) => unknown,
	) =>
		new Promise((resolve) => {
			component = factory(
				{ requestRender() {}, terminal: { rows: 40 } },
				{
					fg: (_style: string, text: string) => text,
					bold: (text: string) => text,
				},
				{
					matches: (data: string, action: string) => KEYS[action] === data,
				},
				resolve,
			) as typeof component;
		});
	const ctx = {
		hasUI: true,
		mode: "tui",
		ui: { custom },
	} as unknown as ExtensionContext;
	return { ctx, component: () => component };
}

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

function fakeCtx(
	ui: {
		select?: (
			title: string,
			entries: string[],
			options?: unknown,
		) => Promise<string | undefined>;
		input?: (
			title: string,
			placeholder?: string,
			options?: unknown,
		) => Promise<string | undefined>;
		confirm?: (
			title: string,
			message: string,
			options?: unknown,
		) => Promise<boolean>;
	},
	overrides: Record<string, unknown> = {},
): ExtensionContext {
	return {
		hasUI: true,
		mode: "rpc",
		ui: {
			select: ui.select ?? (async () => undefined),
			input: ui.input ?? (async () => undefined),
			confirm: ui.confirm ?? (async () => false),
		},
		...overrides,
	} as ExtensionContext;
}

function mouseEvent(type: "press" | "click", y: number, width = 80) {
	return {
		type,
		button: "left" as const,
		x: 1,
		y,
		screenX: 1,
		screenY: y,
		width,
		height: 40,
		shift: false,
		alt: false,
		ctrl: false,
	};
}

function singleQuestion() {
	return __test.normalizeQuestions([
		{
			id: "db",
			header: "Database",
			question: "Which database?",
			options: [
				{ label: "Postgres (Recommended)", description: "Relational." },
				{ label: "SQLite", description: "Embedded." },
			],
		},
	]);
}

describe("ask extension", () => {
	it("loads in real Pi extension runtime and registers ask_user_question", async () => {
		t = await createTestSession({ extensionFactories: [askExtension] });
		const tool =
			t.session.extensionRunner.extensions[0].tools.get(
				"ask_user_question",
			).definition;
		expect(tool).toBeDefined();
		expect(tool.description).toContain(
			"Non-interactive sessions return unavailable",
		);
		expect(
			tool.promptGuidelines.every((guideline: string) =>
				guideline.includes("ask_user_question"),
			),
		).toBe(true);
	});

	it("terminates on cancel and reports missing UI at tool level", async () => {
		t = await createTestSession({ extensionFactories: [askExtension] });
		const tool =
			t.session.extensionRunner.extensions[0].tools.get("ask_user_question");
		const definition = tool?.definition;
		expect(definition).toBeDefined();
		const params = {
			questions: [
				{
					question: "Q?",
					options: [
						{ label: "A", description: "a" },
						{ label: "B", description: "b" },
					],
				},
			],
		};
		const cancelled = await definition.execute(
			"id",
			params,
			undefined,
			undefined,
			fakeCtx({}),
		);
		expect(cancelled.terminate).toBe(true);
		expect(JSON.parse(cancelled.content[0].text).status).toBe("cancelled");

		const unavailable = await definition.execute(
			"id",
			params,
			undefined,
			undefined,
			fakeCtx({}, { hasUI: false }),
		);
		expect(JSON.parse(unavailable.content[0].text).status).toBe("unavailable");
	});
});

describe("normalizeQuestions", () => {
	it("derives key and header and strips model-supplied Other options", () => {
		const [question] = __test.normalizeQuestions([
			{
				question: "Which database?",
				options: [
					{ label: "Postgres (Recommended)", description: "Relational." },
					{ label: "Other", description: "ignored" },
				],
			},
		]);
		expect(question.key).toBe("Which database?");
		expect(question.header).toBe("Which databa");
		expect(question.options).toHaveLength(1);
		expect(question.options[0].label).toBe("Postgres (Recommended)");
	});

	it("prefers id for key and header, dedupes keys, defaults multiSelect", () => {
		const [first, second] = __test.normalizeQuestions([
			{
				id: "deploy",
				header: "Target",
				question: "Deploy where?",
				options: [
					{ label: "A", description: "a" },
					{ label: "B", description: "b" },
				],
			},
			{
				id: "deploy",
				question: "Again?",
				options: [
					{ label: "C", description: "c" },
					{ label: "D", description: "d" },
				],
			},
		]);
		expect(first.key).toBe("deploy");
		expect(first.header).toBe("Target");
		expect(first.multiSelect).toBe(false);
		expect(second.key).toBe("deploy-2");
	});

	it("rejects empty questions and questions without usable options", () => {
		expect(() => __test.normalizeQuestions([])).toThrow(/1-4 questions/);
		expect(() =>
			__test.normalizeQuestions([
				{
					question: "Q?",
					options: [{ label: "Other", description: "x" }],
				},
			]),
		).toThrow(/no usable options/);
	});
});

describe("runAskFlow", () => {
	it("answers single-select with the chosen label and offers Other", async () => {
		const calls: Array<[string, string[]]> = [];
		const result = await __test.runAskFlow(
			singleQuestion(),
			fakeCtx({
				select: async (title, entries) => {
					calls.push([title, entries]);
					return entries[0];
				},
			}),
		);
		expect(result).toEqual({
			status: "answered",
			answers: { db: { answers: ["Postgres (Recommended)"] } },
		});
		expect(calls[0][0]).toBe("Question 1/1 · Database");
		expect(calls[0][1]).toContain(__test.OTHER_LABEL);
	});

	it("routes Other through free-text input and trims it", async () => {
		const result = await __test.runAskFlow(
			singleQuestion(),
			fakeCtx({
				select: async (_title, entries) => entries.at(-1) ?? __test.OTHER_LABEL,
				input: async () => "  Whatever fits  ",
			}),
		);
		expect(result).toEqual({
			status: "answered",
			answers: { db: { answers: ["Whatever fits"] } },
		});
	});

	it("cancels the whole call when any dialog is dismissed", async () => {
		const result = await __test.runAskFlow(
			singleQuestion(),
			fakeCtx({
				select: async () => undefined,
				input: async () => "x",
			}),
		);
		expect(result).toEqual({ status: "cancelled" });
	});

	it("collects multi-select confirms, then Done", async () => {
		const picks = [true, false];
		const questions = __test.normalizeQuestions([
			{
				id: "feat",
				question: "Which features?",
				multiSelect: true,
				options: [
					{ label: "Auth", description: "a" },
					{ label: "Audit", description: "b" },
				],
			},
		]);
		const result = await __test.runAskFlow(
			questions,
			fakeCtx({
				select: async (_title, entries) =>
					entries[0] === "Done" ? "Done" : undefined,
				confirm: async () => picks.shift() ?? false,
			}),
		);
		expect(result).toEqual({
			status: "answered",
			answers: { feat: { answers: ["Auth"] } },
		});
	});
});

describe("runAskFlow in TUI mode", () => {
	const keys = Object.values(KEYS);

	it("selects an option with enter", async () => {
		const session = tuiSession();
		const flow = __test.runAskFlow(singleQuestion(), session.ctx);
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(KEYS["tui.select.confirm"]);
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: { db: { answers: ["SQLite"] } },
		});
	});

	it("routes Other through the free-text input", async () => {
		const session = tuiSession();
		const flow = __test.runAskFlow(singleQuestion(), session.ctx);
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(KEYS["tui.select.confirm"]);
		for (const key of ["C", "u", "s", "t", "o", "m", "\r"]) {
			session.component()?.handleInput(key);
		}
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: { db: { answers: ["Custom"] } },
		});
	});

	it("attaches notes through tab and submits the highlighted option", async () => {
		const session = tuiSession();
		const flow = __test.runAskFlow(singleQuestion(), session.ctx);
		session.component()?.handleInput(KEYS["tui.input.tab"]);
		for (const key of ["l", "a", "t", "e", "r", "\r"]) {
			session.component()?.handleInput(key);
		}
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: { db: { answers: ["Postgres (Recommended)"], notes: "later" } },
		});
	});

	it("wraps questions, options, and free-text input to the available width", async () => {
		const questions = __test.normalizeQuestions([
			{
				id: "wrap",
				question:
					"Which database should handle this unusually demanding workload?",
				options: [
					{
						label: "Postgres with replication (Recommended)",
						description: "Supports relational queries and resilient failover.",
					},
					{ label: "SQLite", description: "Embedded and deliberately simple." },
				],
			},
		]);
		const session = tuiSession();
		const flow = __test.runAskFlow(questions, session.ctx);
		const initial = session.component()?.render(24) ?? [];
		expect(initial.every((line) => visibleWidth(line) <= 24)).toBe(true);
		expect(initial.join("\n")).toContain("demanding workload?");
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(KEYS["tui.select.confirm"]);
		for (const key of "A free-form answer long enough to wrap") {
			session.component()?.handleInput(key);
		}
		const withInput = session.component()?.render(24) ?? [];
		expect(withInput.every((line) => visibleWidth(line) <= 24)).toBe(true);
		expect(
			withInput.filter(
				(line) => line.includes("free-form") || line.includes("enough to wrap"),
			),
		).toHaveLength(2);
		session.component()?.handleInput("\r");
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: {
				wrap: { answers: ["A free-form answer long enough to wrap"] },
			},
		});
	});

	it("selects an option with the mouse in fullscreen TUI mode", async () => {
		const session = tuiSession();
		const flow = __test.runAskFlow(singleQuestion(), session.ctx);
		const lines = session.component()?.render(80) ?? [];
		const sqliteRow = lines.findIndex((line) => line.includes("SQLite"));
		expect(sqliteRow).toBeGreaterThanOrEqual(0);
		session.component()?.handleMouse?.(mouseEvent("press", sqliteRow));
		session.component()?.handleMouse?.(mouseEvent("click", sqliteRow));
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: { db: { answers: ["SQLite"] } },
		});
	});

	it("cancels from the option list with escape", async () => {
		const session = tuiSession();
		const flow = __test.runAskFlow(singleQuestion(), session.ctx);
		session.component()?.handleInput(KEYS["tui.select.cancel"]);
		await expect(flow).resolves.toEqual({ status: "cancelled" });
	});

	it("toggles multi-select options with space", async () => {
		const questions = __test.normalizeQuestions([
			{
				id: "feat",
				question: "Which features?",
				multiSelect: true,
				options: [
					{ label: "Auth", description: "a" },
					{ label: "Audit", description: "b" },
				],
			},
		]);
		const session = tuiSession();
		const flow = __test.runAskFlow(questions, session.ctx);
		session.component()?.handleInput(" ");
		session.component()?.handleInput(KEYS["tui.select.down"]);
		session.component()?.handleInput(" ");
		const rendered = session.component()?.render(80)?.join("\n") ?? "";
		session.component()?.handleInput(KEYS["tui.select.confirm"]);
		await expect(flow).resolves.toEqual({
			status: "answered",
			answers: { feat: { answers: ["Auth", "Audit"] } },
		});
		expect(rendered).toContain("Which features?");
		expect(rendered).toContain("Other");
		expect(keys).toContain(KEYS["tui.select.confirm"]);
	});
});