pi-ext

import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
	getCapabilities,
	resetCapabilitiesCache,
	setCapabilities,
} from "@earendil-works/pi-tui";
import { afterEach, describe, expect, it } from "vitest";
import {
	AUTHORIZE_URL,
	buildAuthorizeUrl,
	type ClipboardRunner,
	clipboardCandidates,
	copyToClipboard,
	EXPIRES_IN_SECONDS,
	exchangeCode,
	generatePkce,
	openerCandidates,
	type Pkce,
	parseAuthorizationInput,
	REDIRECT_URI,
	runSetupTokenCommand,
	TOKEN_URL,
} from "../src/setup-token";

const LONG_LIVED = "sk-ant-oat01-abcdefghijklmnopqrstuvwxyz0123456789";
const PKCE: Pkce = {
	verifier: "verifier",
	challenge: "challenge",
	state: "st",
};

const KEYS: Record<string, string> = {
	"tui.select.cancel": "\x1b",
	"app.interrupt": "\x03",
};

function jsonResponse(body: unknown, status = 200): Response {
	return new Response(typeof body === "string" ? body : JSON.stringify(body), {
		status,
	});
}

type Component = {
	render: (width: number) => string[];
	handleInput?: (data: string) => void;
};

function tuiCtx() {
	const components: Component[] = [];
	const notifications: Array<{ message: string; level: string }> = [];
	const statuses: Array<string | undefined> = [];
	const ctx = {
		hasUI: true,
		mode: "tui",
		ui: {
			custom: (
				factory: (
					tui: unknown,
					theme: unknown,
					keybindings: unknown,
					done: (value: unknown) => void,
				) => Component,
			) =>
				new Promise((resolve) => {
					components.push(
						factory(
							{ requestRender() {} },
							{
								fg: (_style: string, text: string) => text,
								bold: (text: string) => text,
								getColorMode: () => "256color",
							},
							{
								matches: (data: string, action: string) =>
									KEYS[action] === data,
							},
							resolve,
						),
					);
				}),
			notify: (message: string, level: string) => {
				notifications.push({ message, level });
			},
			setStatus: (_key: string, value: string | undefined) => {
				statuses.push(value);
			},
		},
	} as unknown as ExtensionContext;
	return { ctx, components, notifications, statuses };
}

async function waitFor<T>(
	probe: () => T | undefined,
	timeoutMs = 6000,
): Promise<T> {
	const deadline = Date.now() + timeoutMs;
	for (;;) {
		const value = probe();
		if (value !== undefined) return value;
		if (Date.now() > deadline) throw new Error("timed out");
		await new Promise((resolve) => setTimeout(resolve, 10));
	}
}

describe("Klaus setup-token authorization", () => {
	it("builds a PKCE authorize URL with the setup-token scope", async () => {
		const pkce = await generatePkce();
		expect(pkce.verifier).not.toBe(pkce.state);
		const url = new URL(buildAuthorizeUrl(pkce));
		expect(`${url.origin}${url.pathname}`).toBe(AUTHORIZE_URL);
		expect(url.searchParams.get("scope")).toBe("user:inference");
		expect(url.searchParams.get("redirect_uri")).toBe(REDIRECT_URI);
		expect(url.searchParams.get("code_challenge")).toBe(pkce.challenge);
		expect(url.searchParams.get("code_challenge_method")).toBe("S256");
		expect(url.searchParams.get("state")).toBe(pkce.state);
		expect(url.searchParams.get("client_id")).toMatch(/^[0-9a-f-]{36}$/);
	});

	it("parses code#state, redirect URLs, query strings, and bare codes", () => {
		expect(parseAuthorizationInput(" abc#st ")).toEqual({
			code: "abc",
			state: "st",
		});
		expect(
			parseAuthorizationInput(`${REDIRECT_URI}?code=abc&state=st`),
		).toEqual({ code: "abc", state: "st" });
		expect(parseAuthorizationInput("code=abc&state=st")).toEqual({
			code: "abc",
			state: "st",
		});
		expect(parseAuthorizationInput("abc")).toEqual({ code: "abc" });
		expect(parseAuthorizationInput("   ")).toBeUndefined();
		expect(parseAuthorizationInput("#st")).toBeUndefined();
	});
});

describe("Klaus setup-token exchange", () => {
	it("posts the long-lived exchange request and returns the token", async () => {
		const requests: Array<{ url: string; body: unknown }> = [];
		const minted = await exchangeCode(
			{ code: "abc", state: "st" },
			PKCE,
			async (url, init) => {
				requests.push({ url, body: JSON.parse(String(init.body)) });
				return jsonResponse({
					access_token: LONG_LIVED,
					expires_in: EXPIRES_IN_SECONDS,
				});
			},
		);
		expect(minted).toEqual({
			token: LONG_LIVED,
			expiresInSeconds: EXPIRES_IN_SECONDS,
		});
		expect(requests[0]?.url).toBe(TOKEN_URL);
		expect(requests[0]?.body).toMatchObject({
			grant_type: "authorization_code",
			code: "abc",
			redirect_uri: REDIRECT_URI,
			code_verifier: "verifier",
			state: "st",
			expires_in: EXPIRES_IN_SECONDS,
		});
	});

	it("rejects state mismatches before any request", async () => {
		let called = false;
		await expect(
			exchangeCode({ code: "abc", state: "other" }, PKCE, async () => {
				called = true;
				return jsonResponse({});
			}),
		).rejects.toThrow(/state mismatch/);
		expect(called).toBe(false);
	});

	it("reports HTTP failures, invalid JSON, and missing tokens", async () => {
		await expect(
			exchangeCode({ code: "abc" }, PKCE, async () =>
				jsonResponse("nope", 401),
			),
		).rejects.toThrow(/rejected the authorization code/);
		await expect(
			exchangeCode({ code: "abc" }, PKCE, async () =>
				jsonResponse("<html>", 500),
			),
		).rejects.toThrow(/500/);
		await expect(
			exchangeCode({ code: "abc" }, PKCE, async () => jsonResponse("<html>")),
		).rejects.toThrow(/invalid JSON/);
		await expect(
			exchangeCode({ code: "abc" }, PKCE, async () => jsonResponse({})),
		).rejects.toThrow(/no access token/);
	});
});

describe("Klaus setup-token counter scene", () => {
	afterEach(() => {
		resetCapabilitiesCache();
	});

	function withHyperlinks(hyperlinks: boolean): void {
		setCapabilities({ ...getCapabilities(), hyperlinks });
	}

	async function openScene(
		deps: Parameters<typeof runSetupTokenCommand>[1] = {},
		hyperlinks = true,
	) {
		withHyperlinks(hyperlinks);
		const session = tuiCtx();
		const run = runSetupTokenCommand(session.ctx, {
			pkce: async () => PKCE,
			openUrl: async () => true,
			copy: async () => true,
			...deps,
		});
		const scene = await waitFor(() => session.components[0]);
		return { ...session, scene, run };
	}

	function type(
		scene: { handleInput?: (data: string) => void },
		value: string,
	): void {
		for (const key of [...value]) scene.handleInput?.(key);
	}

	async function waitForText(
		scene: { render: (width: number) => string[] },
		needle: string,
	): Promise<string> {
		return waitFor(() => {
			const rendered = scene.render(100).join("\n");
			return rendered.includes(needle) ? rendered : undefined;
		});
	}

	it("runs filing, sign-in, exchange, and stamp in one component", async () => {
		const copied: string[] = [];
		const opened: string[] = [];
		const { scene, run, notifications, components } = await openScene({
			openUrl: async (url) => {
				opened.push(url);
				return true;
			},
			fetchFn: async () =>
				jsonResponse({ access_token: LONG_LIVED, expires_in: 3600 }),
			copy: async (text) => {
				copied.push(text);
				return true;
			},
		});

		expect(opened).toEqual([buildAuthorizeUrl(PKCE)]);
		const signin = await waitForText(scene, "waiting for your code");
		expect(signin).toContain("BUREAU OF TOKENS");
		expect(signin).toContain("Open the Claude sign-in page");
		expect(signin).toContain("Code: ");

		type(scene, "abc#st");
		scene.handleInput?.("\r");
		const stamped = await waitForText(scene, "A P P R O V E D");
		expect(copied).toEqual([LONG_LIVED]);
		expect(stamped).toContain("VALID UNTIL");
		expect(stamped).not.toContain("Code: ");
		expect(stamped).toContain("any key to skip");

		scene.handleInput?.("x");
		const settled = scene.render(100).join("\n");
		expect(settled).toContain("✔ Token in clipboard.");
		expect(settled).toContain("any key to close");

		scene.handleInput?.("x");
		await run;
		expect(components).toHaveLength(1);
		expect(notifications).toEqual([]);
	});

	it("cancels on escape and on interrupt without any exchange", async () => {
		for (const key of ["\x1b", "\x03"]) {
			let fetched = false;
			const { scene, run, notifications } = await openScene({
				fetchFn: async () => {
					fetched = true;
					return jsonResponse({});
				},
			});
			scene.handleInput?.(key);
			await run;
			expect(fetched).toBe(false);
			expect(notifications.at(-1)?.message).toContain("application withdrawn");
		}
	});

	it("stamps DENIED in place and retries on r", async () => {
		const opened: string[] = [];
		let attempt = 0;
		const { scene, run, notifications } = await openScene({
			openUrl: async (url) => {
				opened.push(url);
				return true;
			},
			fetchFn: async () => {
				attempt++;
				return attempt === 1
					? jsonResponse("bad", 401)
					: jsonResponse({ access_token: LONG_LIVED });
			},
		});

		await waitForText(scene, "Code: ");
		type(scene, "abc");
		scene.handleInput?.("\r");
		const denied = await waitForText(scene, "rejected the authorization code");
		expect(denied).toContain("D E N I E D");
		expect(notifications).toEqual([]);

		scene.handleInput?.("x");
		expect(scene.render(100).join("\n")).toContain("r retry · esc close");
		scene.handleInput?.("r");
		const retried = await waitForText(scene, "Code: ");
		expect(retried).toContain("Open the Claude sign-in page");
		expect(retried).toContain("enter submit");
		expect(retried).not.toContain("D E N I E D");
		expect(opened).toHaveLength(2);

		type(scene, "abc");
		scene.handleInput?.("\r");
		await waitForText(scene, "A P P R O V E D");
		scene.handleInput?.("x");
		scene.handleInput?.("x");
		await run;
	});

	it("links the sign-in page with OSC 8 and falls back to the raw URL", async () => {
		const linked = await openScene({}, true);
		await waitForText(linked.scene, "Open the Claude sign-in page");
		const link =
			linked.scene.render(80).find((line) => line.includes("\x1b]8;;")) ?? "";
		expect(link).toContain(`\x1b]8;;${buildAuthorizeUrl(PKCE)}\x1b\\`);
		expect(link).toContain("Open the Claude sign-in page");
		linked.scene.handleInput?.("\x1b");
		await linked.run;

		const plain = await openScene({}, false);
		await waitForText(plain.scene, "code_challenge");
		const rendered = plain.scene.render(80).join("");
		expect(rendered).not.toContain("\x1b]8;;");
		expect(rendered).toContain("code_challenge");
		plain.scene.handleInput?.("\x1b");
		await plain.run;
	});

	it("copies the link on c while the code field is empty", async () => {
		const copies: string[] = [];
		const { scene, run } = await openScene({
			copy: async (text) => {
				copies.push(text);
				return true;
			},
		});
		await waitForText(scene, "Open the Claude sign-in page");
		scene.handleInput?.("c");
		await waitForText(scene, "Link copied to clipboard.");
		expect(copies).toEqual([buildAuthorizeUrl(PKCE)]);

		type(scene, "abc");
		scene.handleInput?.("c");
		expect(copies).toHaveLength(1);
		expect(scene.render(100).join("\n")).toContain("abcc");
		scene.handleInput?.("\x1b");
		await run;
	});

	it("uses the plain input dialog outside the TUI and never opens a browser", async () => {
		const notifications: Array<{ message: string; level: string }> = [];
		const titles: string[] = [];
		let opened = false;
		const ctx = {
			hasUI: true,
			mode: "rpc",
			ui: {
				input: async (title: string) => {
					titles.push(title);
					return "abc#st";
				},
				notify: (message: string, level: string) => {
					notifications.push({ message, level });
				},
				setStatus() {},
			},
		} as unknown as ExtensionContext;
		await runSetupTokenCommand(ctx, {
			pkce: async () => PKCE,
			openUrl: async () => {
				opened = true;
				return true;
			},
			fetchFn: async () => jsonResponse({ access_token: LONG_LIVED }),
			copy: async () => false,
		});
		expect(opened).toBe(false);
		expect(titles[0]).toContain(AUTHORIZE_URL);
		expect(notifications.at(-1)).toEqual({
			message: `No clipboard. Token: ${LONG_LIVED}`,
			level: "warning",
		});
	});
});

describe("Klaus setup-token clipboard", () => {
	it("prefers wayland, then xclip, then xsel on linux", () => {
		expect(clipboardCandidates("linux")).toEqual([
			["wl-copy"],
			["xclip", "-selection", "clipboard"],
			["xsel", "--clipboard", "--input"],
		]);
	});

	it("uses the platform tool on darwin, win32, and android", () => {
		expect(clipboardCandidates("darwin")).toEqual([["pbcopy"]]);
		expect(clipboardCandidates("win32")).toEqual([["clip"]]);
		expect(clipboardCandidates("android")).toEqual([["termux-clipboard-set"]]);
	});

	it("stops at the first working tool", async () => {
		const attempted: string[][] = [];
		const runner: ClipboardRunner = async (command, text) => {
			attempted.push(command);
			return command[0] === "xclip" && text === LONG_LIVED;
		};
		await expect(copyToClipboard(LONG_LIVED, runner, "linux")).resolves.toBe(
			true,
		);
		expect(attempted).toEqual([
			["wl-copy"],
			["xclip", "-selection", "clipboard"],
		]);
	});

	it("reports failure when every tool fails", async () => {
		const runner: ClipboardRunner = async () => false;
		await expect(copyToClipboard(LONG_LIVED, runner, "linux")).resolves.toBe(
			false,
		);
	});
});

describe("Klaus browser opener", () => {
	it("prefers termux, macOS, and Windows openers by platform", () => {
		expect(
			openerCandidates("https://x", {
				platform: "linux",
				env: { TERMUX_VERSION: "1" },
			}),
		).toEqual([{ command: "termux-open-url", args: ["https://x"] }]);
		expect(
			openerCandidates("https://x", { platform: "darwin", env: {} }),
		).toEqual([{ command: "open", args: ["https://x"] }]);
		expect(
			openerCandidates("https://x", { platform: "linux", env: {} }).map(
				(candidate) => candidate.command,
			),
		).toEqual(["xdg-open", "wslview", "gio"]);
	});
});