Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/klaus/src/setup-token.ts

Raw
import { spawn } from "node:child_process";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { getCapabilities, Input } from "@earendil-works/pi-tui";
import { dbg, diagnosticKind } from "./debug.js";
import { openUrl } from "./pi-ext-browser-open.js";
import {
	type ClipboardState,
	renderScene,
	SCENE_INTRO_TICKS,
	SCENE_TICK_MS,
	type ScenePhase,
	type SceneState,
	sceneOutcomeTicks,
} from "./setup-scene.js";

export { openerCandidates } from "./pi-ext-browser-open.js";

// Same public client as Claude Code and Pi's own Anthropic OAuth flow.
const CLIENT_ID = atob("OWQxYzI1MGEtZTYxYi00NGQ5LTg4ZWQtNTk0NGQxOTYyZjVl");
export const AUTHORIZE_URL = "https://claude.com/cai/oauth/authorize";
export const TOKEN_URL = "https://platform.claude.com/v1/oauth/token";
export const REDIRECT_URI = "https://platform.claude.com/oauth/code/callback";
const SCOPE = "user:inference";
export const EXPIRES_IN_SECONDS = 365 * 24 * 60 * 60;

export type Pkce = { verifier: string; challenge: string; state: string };

function base64url(bytes: Uint8Array): string {
	return Buffer.from(bytes).toString("base64url");
}

export async function generatePkce(): Promise<Pkce> {
	const verifier = base64url(crypto.getRandomValues(new Uint8Array(32)));
	const digest = await crypto.subtle.digest(
		"SHA-256",
		new TextEncoder().encode(verifier),
	);
	return {
		verifier,
		challenge: base64url(new Uint8Array(digest)),
		state: base64url(crypto.getRandomValues(new Uint8Array(32))),
	};
}

export function buildAuthorizeUrl(pkce: Pkce): string {
	const params = new URLSearchParams({
		code: "true",
		client_id: CLIENT_ID,
		response_type: "code",
		redirect_uri: REDIRECT_URI,
		scope: SCOPE,
		code_challenge: pkce.challenge,
		code_challenge_method: "S256",
		state: pkce.state,
	});
	return `${AUTHORIZE_URL}?${params}`;
}

export type ParsedAuthorization = { code: string; state?: string };

/** Accepts the callback page's `code#state`, a redirect URL, a query string, or a bare code. */
export function parseAuthorizationInput(
	input: string,
): ParsedAuthorization | undefined {
	const value = input.trim();
	if (!value) return undefined;
	try {
		const url = new URL(value);
		const code = url.searchParams.get("code");
		return code
			? { code, state: url.searchParams.get("state") ?? undefined }
			: undefined;
	} catch {}
	if (value.includes("#")) {
		const [code, state] = value.split("#", 2);
		return code ? { code, state: state || undefined } : undefined;
	}
	if (value.includes("code=")) {
		const params = new URLSearchParams(value);
		const code = params.get("code");
		return code ? { code, state: params.get("state") ?? undefined } : undefined;
	}
	return { code: value };
}

export type FetchLike = (url: string, init: RequestInit) => Promise<Response>;

export type MintedToken = { token: string; expiresInSeconds?: number };

export async function exchangeCode(
	parsed: ParsedAuthorization,
	pkce: Pkce,
	fetchFn: FetchLike = fetch,
): Promise<MintedToken> {
	if (parsed.state && parsed.state !== pkce.state)
		throw new Error("OAuth state mismatch; start over.");
	dbg?.("setupToken.exchange.start");
	const response = await fetchFn(TOKEN_URL, {
		method: "POST",
		headers: { "Content-Type": "application/json", Accept: "application/json" },
		body: JSON.stringify({
			grant_type: "authorization_code",
			code: parsed.code,
			redirect_uri: REDIRECT_URI,
			client_id: CLIENT_ID,
			code_verifier: pkce.verifier,
			state: pkce.state,
			expires_in: EXPIRES_IN_SECONDS,
		}),
		signal: AbortSignal.timeout(30_000),
	});
	const body = await response.text();
	dbg?.("setupToken.exchange.response");
	if (!response.ok) {
		throw new Error(
			response.status === 401
				? "Anthropic rejected the authorization code."
				: `Token exchange failed (${response.status}): ${body.slice(0, 200)}`,
		);
	}
	let data: { access_token?: unknown; expires_in?: unknown };
	try {
		data = JSON.parse(body) as typeof data;
	} catch {
		throw new Error("Token exchange returned invalid JSON.");
	}
	if (typeof data.access_token !== "string" || !data.access_token)
		throw new Error("Token exchange returned no access token.");
	return {
		token: data.access_token,
		expiresInSeconds:
			typeof data.expires_in === "number" ? data.expires_in : undefined,
	};
}

export function clipboardCandidates(
	platform: NodeJS.Platform = process.platform,
): string[][] {
	switch (platform) {
		case "darwin":
			return [["pbcopy"]];
		case "win32":
			return [["clip"]];
		case "android":
			return [["termux-clipboard-set"]];
		default:
			return [
				["wl-copy"],
				["xclip", "-selection", "clipboard"],
				["xsel", "--clipboard", "--input"],
			];
	}
}

export type ClipboardRunner = (
	command: string[],
	text: string,
) => Promise<boolean>;

export async function copyToClipboard(
	text: string,
	run: ClipboardRunner,
	platform: NodeJS.Platform = process.platform,
): Promise<boolean> {
	for (const command of clipboardCandidates(platform)) {
		if (await run(command, text)) return true;
	}
	return false;
}

export function defaultRunner(
	command: string[],
	text: string,
): Promise<boolean> {
	return new Promise((resolve) => {
		const child = spawn(command[0] ?? "", command.slice(1), {
			stdio: ["pipe", "ignore", "ignore"],
		});
		child.on("error", () => resolve(false));
		child.on("close", (code) => resolve(code === 0));
		child.stdin?.on("error", () => resolve(false));
		child.stdin?.end(text);
	});
}

async function openBrowser(url: string): Promise<boolean> {
	if (process.env.NODE_ENV === "test" || process.env.VITEST) return false;
	return (await openUrl(url, {})) !== undefined;
}

function fail(message: string): string {
	return `Klaus could not collect a token: ${message}`;
}

export type SetupTokenDeps = {
	pkce?: () => Promise<Pkce>;
	fetchFn?: FetchLike;
	copy?: typeof copyToClipboard;
	run?: ClipboardRunner;
	openUrl?: (url: string) => Promise<boolean>;
};

let inFlight = false;

export async function runSetupTokenCommand(
	ctx: ExtensionContext,
	deps: SetupTokenDeps = {},
): Promise<void> {
	if (!ctx.hasUI) return;
	if (inFlight) {
		ctx.ui.notify("Klaus is already minting a token.", "warning");
		return;
	}
	inFlight = true;
	dbg?.("setupToken.command.start");
	try {
		if (ctx.mode === "tui") {
			await runCounterScene(ctx, deps);
			return;
		}
		await runDialogFlow(ctx, deps);
	} finally {
		inFlight = false;
	}
}

/** RPC and other UI hosts without a TUI: plain dialogs, no scene. */
async function runDialogFlow(
	ctx: ExtensionContext,
	deps: SetupTokenDeps,
): Promise<void> {
	const pkce = await (deps.pkce ?? generatePkce)();
	const url = buildAuthorizeUrl(pkce);
	const answer = await ctx.ui.input(
		`Klaus: sign in at ${url} and paste the code`,
		"code#state",
	);
	const parsed =
		answer === undefined ? undefined : parseAuthorizationInput(answer);
	if (!parsed) {
		dbg?.("setupToken.command.cancelled");
		ctx.ui.notify("Klaus: application withdrawn.", "info");
		return;
	}
	ctx.ui.setStatus("klaus", "minting token…");
	let minted: MintedToken;
	try {
		minted = await exchangeCode(parsed, pkce, deps.fetchFn);
	} catch (error) {
		ctx.ui.notify(
			fail(error instanceof Error ? error.message : String(error)),
			"error",
		);
		return;
	} finally {
		ctx.ui.setStatus("klaus", undefined);
	}
	const copied = await (deps.copy ?? copyToClipboard)(
		minted.token,
		deps.run ?? defaultRunner,
	);
	dbg?.("setupToken.command.minted", { copied });
	ctx.ui.notify(
		copied ? "Token in clipboard." : `No clipboard. Token: ${minted.token}`,
		copied ? "info" : "warning",
	);
}

/**
 * TUI: one focused component owns the whole visit, from filing the application
 * through the sign-in code to the stamp, so the animation frames the flow
 * instead of interrupting it.
 */
async function runCounterScene(
	ctx: ExtensionContext,
	deps: SetupTokenDeps,
): Promise<void> {
	let pkce = await (deps.pkce ?? generatePkce)();
	let url = buildAuthorizeUrl(pkce);
	void (deps.openUrl ?? openBrowser)(url);

	const outcome = await ctx.ui.custom<"cancelled" | "closed">(
		(tui, theme, keybindings, done) => {
			const field = new Input({
				prompt: "Code: ",
				placeholder: "code#state",
				placeholderStyle: (text) => theme.fg("dim", text),
			});
			field.focused = true;
			let tick = 0;
			let phaseTick = 0;
			let phase: ScenePhase = "signin";
			let clipboard: ClipboardState = "idle";
			let token: string | undefined;
			let tokenCopied = false;
			let expiresInSeconds: number | undefined;
			let failure: string | undefined;
			let settled = false;
			let timer: NodeJS.Timeout | undefined;
			let interval = 0;

			const outcomeTicks = () => sceneOutcomeTicks(sceneState([], ""));
			const outcomeSettled = () =>
				(phase === "approved" || phase === "denied") &&
				phaseTick >= outcomeTicks();

			function sceneState(input: string[], hints: string): SceneState {
				return {
					tick,
					phaseTick,
					phase,
					input,
					hyperlinks: getCapabilities().hyperlinks,
					url,
					clipboard,
					hints,
					token,
					tokenCopied,
					expiresInSeconds,
					error: failure,
				};
			}

			function hints(): string {
				if (phase === "signin")
					return [
						"enter submit",
						field.getValue() === "" ? "c copy link" : "",
						"esc cancel",
					]
						.filter(Boolean)
						.join(" · ");
				if (phase === "exchanging") return "esc cancel";
				if (!outcomeSettled()) return "any key to skip";
				return phase === "denied" ? "r retry · esc close" : "any key to close";
			}

			/** Fast during animation, slow while the human is signing in, off when static. */
			function schedule(): void {
				const desired =
					tick < SCENE_INTRO_TICKS
						? SCENE_TICK_MS
						: phase === "exchanging"
							? 90
							: phase === "signin"
								? 400
								: outcomeSettled()
									? 0
									: SCENE_TICK_MS;
				if (desired === interval) return;
				interval = desired;
				if (timer) clearInterval(timer);
				timer = undefined;
				if (desired === 0) return;
				timer = setInterval(() => {
					tick++;
					phaseTick++;
					schedule();
					tui.requestRender();
				}, desired);
				timer.unref?.();
			}

			function stop(): void {
				if (timer) clearInterval(timer);
				timer = undefined;
				interval = 0;
			}

			function finish(result: "cancelled" | "closed"): void {
				if (settled) return;
				settled = true;
				stop();
				done(result);
			}

			function enterPhase(next: ScenePhase): void {
				phase = next;
				phaseTick = 0;
				schedule();
				tui.requestRender();
			}

			async function submit(value: string): Promise<void> {
				const parsed = parseAuthorizationInput(value);
				if (!parsed) return;
				enterPhase("exchanging");
				try {
					const minted = await exchangeCode(parsed, pkce, deps.fetchFn);
					token = minted.token;
					expiresInSeconds = minted.expiresInSeconds;
					tokenCopied = await (deps.copy ?? copyToClipboard)(
						minted.token,
						deps.run ?? defaultRunner,
					);
					dbg?.("setupToken.command.minted", { copied: tokenCopied });
					enterPhase("approved");
				} catch (error) {
					failure = error instanceof Error ? error.message : String(error);
					dbg?.("setupToken.command.denied", { kind: diagnosticKind(error) });
					enterPhase("denied");
				}
			}

			async function retry(): Promise<void> {
				pkce = await (deps.pkce ?? generatePkce)();
				url = buildAuthorizeUrl(pkce);
				failure = undefined;
				clipboard = "idle";
				field.setValue("");
				void (deps.openUrl ?? openBrowser)(url);
				enterPhase("signin");
			}

			function copyLink(): void {
				void (deps.copy ?? copyToClipboard)(
					url,
					deps.run ?? defaultRunner,
				).then((ok) => {
					clipboard = ok ? "copied" : "failed";
					tui.requestRender();
				});
			}

			field.onSubmit = (value) => {
				if (value.trim()) void submit(value);
			};
			field.onEscape = () => finish("cancelled");
			schedule();

			return {
				render(width: number) {
					// Leave room for the centering indent the scene adds around the form block.
					const input =
						phase === "signin"
							? field.render(Math.max(Math.min(width, 52), 12))
							: [];
					return renderScene(sceneState(input, hints()), width, theme);
				},
				invalidate() {
					field.invalidate();
				},
				handleInput(data: string) {
					const cancel =
						keybindings.matches(data, "app.interrupt") ||
						keybindings.matches(data, "tui.select.cancel");
					if (phase === "signin") {
						if (cancel) {
							finish("cancelled");
							return;
						}
						if (data === "c" && field.getValue() === "") {
							copyLink();
							return;
						}
						field.handleInput(data);
						schedule();
						tui.requestRender();
						return;
					}
					if (phase === "exchanging") {
						if (cancel) finish("cancelled");
						return;
					}
					if (!outcomeSettled()) {
						phaseTick = outcomeTicks();
						schedule();
						tui.requestRender();
						return;
					}
					if (phase === "denied" && data === "r") {
						void retry();
						return;
					}
					finish("closed");
				},
				dispose() {
					stop();
				},
			};
		},
	);

	if (outcome === "cancelled") {
		dbg?.("setupToken.command.cancelled");
		ctx.ui.notify("Klaus: application withdrawn.", "info");
	}
}