harness-tools/extensions/update/updater.ts

import { createHash } from "node:crypto";
import {
	access,
	chmod,
	cp,
	mkdir,
	lstat,
	mkdtemp,
	readFile,
	readdir,
	readlink,
	realpath,
	rename,
	rm,
	writeFile,
} from "node:fs/promises";
import { existsSync } from "node:fs";
import { homedir, tmpdir } from "node:os";
import path from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const PI_PACKAGE = "@earendil-works/pi-coding-agent";
const BREW_FORMULA = "pi-coding-agent";
const RELEASE_API = "https://api.github.com/repos/earendil-works/pi/releases/latest";
const MAX_ATTEMPTS = 3;
const COMMAND_TIMEOUT_MS = 15 * 60_000;

export type InstallMethod = "vp" | "bun" | "npm" | "brew" | "native";

export interface UpdateReporter {
	info(message: string): void;
	retry(message: string): void;
}

export interface UpdateResult {
	method: InstallMethod;
	beforeVersion: string;
	afterVersion: string;
	extensionsChecked: boolean;
	requiresShutdown: boolean;
}

export class UpdateFailure extends Error {
	constructor(message: string, readonly requiresShutdown: boolean) {
		super(message);
		this.name = "UpdateFailure";
	}
}

class NativeReplacementFailure extends Error {
	readonly requiresShutdown = true;

	constructor(message: string, options?: ErrorOptions) {
		super(message, options);
		this.name = "NativeReplacementFailure";
	}
}

interface CommandSpec {
	command: string;
	args: string[];
	display: string;
}

interface NativeUpdateResult {
	version: string;
	requiresShutdown: boolean;
}

interface GitHubAsset {
	name: string;
	browser_download_url: string;
}

interface GitHubRelease {
	tag_name: string;
	assets: GitHubAsset[];
}

function normalizePath(value: string | undefined): string {
	return (value ?? "").replaceAll("\\", "/").toLowerCase();
}

function isGenericRuntime(executable: string): boolean {
	return /^(node|bun)(\.exe)?$/i.test(path.basename(executable));
}

async function pathExists(filePath: string): Promise<boolean> {
	try {
		await access(filePath);
		return true;
	} catch {
		return false;
	}
}

async function resolveIfPossible(filePath: string | undefined): Promise<string | undefined> {
	if (!filePath) return undefined;
	try {
		return await realpath(filePath);
	} catch {
		return filePath;
	}
}

function findLauncherOnPath(): string | undefined {
	const names = process.platform === "win32" ? ["pi.exe", "pi.cmd", "pi.bat", "pi"] : ["pi"];
	for (const directory of (process.env.PATH ?? "").split(path.delimiter)) {
		if (!directory) continue;
		for (const name of names) {
			const candidate = path.join(directory, name);
			if (existsSync(candidate)) return candidate;
		}
	}
	return undefined;
}

export function detectInstallMethodFromPaths(input: {
	launcher?: string;
	resolvedLauncher?: string;
	currentScript?: string;
	execPath?: string;
	vpHome?: string;
	isBunRuntime?: boolean;
}): InstallMethod {
	const paths = [input.currentScript, input.execPath, input.resolvedLauncher, input.launcher]
		.map(normalizePath)
		.filter(Boolean);
	const joined = paths.join("\n");
	const vpHome = normalizePath(input.vpHome);

	if (
		(vpHome && paths.some((candidate) => candidate.startsWith(`${vpHome}/`))) ||
		joined.includes("/.vite-plus/packages/") ||
		joined.includes("/.vite-plus/bin/")
	) {
		return "vp";
	}
	if (joined.includes("/cellar/pi-coding-agent/") || joined.includes("/homebrew/cellar/pi-coding-agent/")) {
		return "brew";
	}
	if (
		joined.includes("/.bun/install/global/") ||
		joined.includes("/bun/install/global/") ||
		(input.isBunRuntime === true && path.basename(input.execPath ?? "").toLowerCase() === "bun")
	) {
		return "bun";
	}
	if (joined.includes("/node_modules/")) {
		return "npm";
	}
	return "native";
}

export async function detectInstallMethod(): Promise<InstallMethod> {
	const launcher = findLauncherOnPath();
	const resolvedLauncher = await resolveIfPossible(launcher);
	const resolvedCurrentScript = await resolveIfPossible(process.argv[1]);
	const resolvedExecPath = await resolveIfPossible(process.execPath);
	const normalizedLauncher = normalizePath(resolvedLauncher);
	const launcherBelongsToCurrentProcess =
		normalizedLauncher === normalizePath(resolvedCurrentScript) || normalizedLauncher === normalizePath(resolvedExecPath);
	return detectInstallMethodFromPaths({
		launcher: launcherBelongsToCurrentProcess ? launcher : undefined,
		resolvedLauncher: launcherBelongsToCurrentProcess ? resolvedLauncher : undefined,
		currentScript: resolvedCurrentScript,
		execPath: resolvedExecPath,
		vpHome: process.env.VP_HOME ?? path.join(homedir(), ".vite-plus"),
		isBunRuntime: Boolean(process.versions.bun),
	});
}

export function isTransientRegistryError(output: string): boolean {
	return /(?:(?:HTTP|status(?: code)?)\s*(?:408|425|429|500|502|503|504)\b|EAI_AGAIN|ECONNRESET|ECONNREFUSED|ETIMEDOUT|ENETDOWN|ENETUNREACH|EHOSTUNREACH|ERR_SOCKET_TIMEOUT|UND_ERR_CONNECT_TIMEOUT|socket hang up|network timeout|timed out|aborted due to timeout|temporary failure in name resolution|temporarily unavailable)/i.test(
		output,
	);
}

function delay(milliseconds: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, milliseconds));
}

function quoteArgument(argument: string): string {
	return /\s/.test(argument) ? JSON.stringify(argument) : argument;
}

function commandSpec(command: string, args: string[]): CommandSpec {
	return { command, args, display: [command, ...args].map(quoteArgument).join(" ") };
}

export function getPiInvocation(args: string[]): CommandSpec {
	const currentScript = process.argv[1];
	const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
	if (currentScript && !isBunVirtualScript && existsSync(currentScript)) {
		return commandSpec(process.execPath, [currentScript, ...args]);
	}
	if (!isGenericRuntime(process.execPath)) {
		return commandSpec(process.execPath, args);
	}
	return commandSpec(process.platform === "win32" ? "pi.exe" : "pi", args);
}

async function runWithRetry(
	pi: ExtensionAPI,
	spec: CommandSpec,
	reporter: UpdateReporter,
): Promise<{ stdout: string; stderr: string }> {
	let lastOutput = "";
	for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt += 1) {
		const result = await pi.exec(spec.command, spec.args, { timeout: COMMAND_TIMEOUT_MS });
		if (result.code === 0 && !result.killed) return { stdout: result.stdout, stderr: result.stderr };

		lastOutput = [result.stdout, result.stderr, result.killed ? "Process timed out and was killed" : ""]
			.filter(Boolean)
			.join("\n")
			.trim();
		if (result.killed) {
			throw new Error(`${spec.display} timed out and was killed; refusing to replay a mutating update.\n${lastOutput}`.trim());
		}
		if (attempt === MAX_ATTEMPTS || !isTransientRegistryError(lastOutput)) {
			throw new Error(`${spec.display} failed with exit code ${result.code ?? "unknown"}\n${lastOutput}`.trim());
		}
		const waitSeconds = 2 ** (attempt - 1);
		reporter.retry(`Transient registry error; retrying ${spec.display} in ${waitSeconds}s (${attempt}/${MAX_ATTEMPTS})`);
		await delay(waitSeconds * 1_000);
	}
	throw new Error(`${spec.display} failed\n${lastOutput}`.trim());
}

async function readPiVersion(pi: ExtensionAPI): Promise<string> {
	const invocation = getPiInvocation(["--version"]);
	const result = await pi.exec(invocation.command, invocation.args, { timeout: 30_000 });
	if (result.code !== 0 || result.killed) {
		throw new Error(`Could not read Pi version: ${result.stderr || result.stdout || "process timed out"}`.trim());
	}
	const version = result.stdout.trim().replace(/^v/, "");
	if (!version) throw new Error("Could not read Pi version: empty output");
	return version;
}

function describeError(error: unknown): string {
	if (!(error instanceof Error)) return String(error);
	const cause = "cause" in error ? (error as Error & { cause?: unknown }).cause : undefined;
	return cause === undefined ? error.message : `${error.message}: ${describeError(cause)}`;
}

async function fetchWithRetry(url: string, reporter: UpdateReporter): Promise<Response> {
	let lastError = "";
	for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt += 1) {
		try {
			const response = await fetch(url, {
				headers: { Accept: "application/vnd.github+json", "User-Agent": "pi-update-extension" },
				signal: AbortSignal.timeout(30_000),
			});
			if (response.ok) return response;
			lastError = `HTTP ${response.status} ${response.statusText}`;
			if (!isTransientRegistryError(String(response.status))) throw new Error(lastError);
		} catch (error) {
			lastError = describeError(error);
			if (attempt === MAX_ATTEMPTS || !isTransientRegistryError(lastError)) throw error;
		}
		if (attempt < MAX_ATTEMPTS) {
			const waitSeconds = 2 ** (attempt - 1);
			reporter.retry(`Transient download error; retrying in ${waitSeconds}s (${attempt}/${MAX_ATTEMPTS})`);
			await delay(waitSeconds * 1_000);
		}
	}
	throw new Error(`Download failed: ${lastError}`);
}

async function downloadFile(url: string, destination: string, reporter: UpdateReporter): Promise<void> {
	const response = await fetchWithRetry(url, reporter);
	const data = new Uint8Array(await response.arrayBuffer());
	await writeFile(destination, data, { mode: 0o600 });
}

function releaseAssetName(): string {
	const platform = process.platform === "win32" ? "windows" : process.platform;
	const architecture = process.arch === "x64" || process.arch === "arm64" ? process.arch : undefined;
	if (!architecture || !["linux", "darwin", "windows"].includes(platform)) {
		throw new Error(`No native Pi release is available for ${process.platform}/${process.arch}`);
	}
	return `pi-${platform}-${architecture}.${process.platform === "win32" ? "zip" : "tar.gz"}`;
}

async function sha256(filePath: string): Promise<string> {
	const data = await readFile(filePath);
	return createHash("sha256").update(data).digest("hex");
}

async function verifyReleaseChecksum(archivePath: string, assetName: string, checksumsPath: string): Promise<void> {
	const checksums = await readFile(checksumsPath, "utf8");
	const expected = checksums
		.split(/\r?\n/)
		.map((line) => line.trim().split(/\s+/))
		.find((parts) => parts.at(-1)?.replace(/^\*/, "") === assetName)?.[0];
	if (!expected) throw new Error(`SHA256SUMS does not contain ${assetName}`);
	const actual = await sha256(archivePath);
	if (actual.toLowerCase() !== expected.toLowerCase()) {
		throw new Error(`Checksum mismatch for ${assetName}: expected ${expected}, received ${actual}`);
	}
}

const NATIVE_BUNDLE_ENTRIES = new Set([
	"pi",
	"pi.exe",
	"package.json",
	"README.md",
	"CHANGELOG.md",
	"photon_rs_bg.wasm",
	"theme",
	"assets",
	"export-html",
	"docs",
	"examples",
	"node_modules",
	"native",
]);

async function validateNativeBundle(bundleDir: string, executableName: string, expectedVersion?: string): Promise<void> {
	const manifestPath = path.join(bundleDir, "package.json");
	let manifest: { name?: string; version?: string };
	try {
		const manifestJson = await readFile(manifestPath, "utf8");
		manifest = JSON.parse(manifestJson) as { name?: string; version?: string };
	} catch (error) {
		throw new Error(`Could not validate native Pi manifest at ${manifestPath}`, { cause: error });
	}
	if (manifest.name !== PI_PACKAGE) throw new Error(`Unexpected native bundle package: ${manifest.name ?? "missing"}`);
	if (expectedVersion && manifest.version !== expectedVersion) {
		throw new Error(`Unexpected native bundle version: ${manifest.version ?? "missing"}`);
	}
	for (const requiredPath of [executableName, "theme", "assets", "export-html", "node_modules"]) {
		if (!(await pathExists(path.join(bundleDir, requiredPath)))) {
			throw new Error(`Native Pi bundle is missing ${requiredPath}`);
		}
	}
	const unexpectedEntries = (await readdir(bundleDir)).filter((entry) => !NATIVE_BUNDLE_ENTRIES.has(entry));
	if (unexpectedEntries.length > 0) {
		throw new Error(`Native Pi bundle contains unowned entries: ${unexpectedEntries.join(", ")}`);
	}
}

async function runOnce(pi: ExtensionAPI, spec: CommandSpec): Promise<void> {
	const result = await pi.exec(spec.command, spec.args, { timeout: COMMAND_TIMEOUT_MS });
	if (result.code !== 0 || result.killed) {
		const output = [result.stdout, result.stderr, result.killed ? "Process timed out and was killed" : ""]
			.filter(Boolean)
			.join("\n")
			.trim();
		throw new Error(`${spec.display} failed with exit code ${result.code ?? "unknown"}\n${output}`.trim());
	}
}

async function extractNativeArchive(pi: ExtensionAPI, archivePath: string, extractDir: string): Promise<string> {
	await mkdir(extractDir, { recursive: true });
	await runOnce(pi, commandSpec("tar", ["-xzf", archivePath, "-C", extractDir]));
	return path.join(extractDir, "pi");
}

async function validateNewNativeExecutable(pi: ExtensionAPI, bundleDir: string, version: string): Promise<void> {
	const executable = path.join(bundleDir, process.platform === "win32" ? "pi.exe" : "pi");
	if (process.platform !== "win32") await chmod(executable, 0o755);
	const result = await pi.exec(executable, ["--version"], { timeout: 30_000 });
	if (result.code !== 0 || result.killed || result.stdout.trim().replace(/^v/, "") !== version) {
		throw new Error(`Downloaded native Pi failed validation: ${result.stderr || result.stdout}`.trim());
	}
}

function safeVersion(version: string): string {
	return version.replace(/[^0-9A-Za-z._-]/g, "-");
}

function compareSemver(left: string, right: string): number {
	const parse = (value: string) => {
		const match = /^(\d+)\.(\d+)\.(\d+)(?:-([0-9A-Za-z.-]+))?/.exec(value);
		if (!match) throw new Error(`Invalid semantic version: ${value}`);
		return { numbers: match.slice(1, 4).map(Number), prerelease: match[4] };
	};
	const a = parse(left);
	const b = parse(right);
	for (let index = 0; index < 3; index += 1) {
		const difference = (a.numbers[index] ?? 0) - (b.numbers[index] ?? 0);
		if (difference !== 0) return Math.sign(difference);
	}
	if (a.prerelease === b.prerelease) return 0;
	if (a.prerelease === undefined) return 1;
	if (b.prerelease === undefined) return -1;
	return a.prerelease.localeCompare(b.prerelease, undefined, { numeric: true });
}

async function bundleInventory(root: string): Promise<Map<string, string>> {
	const inventory = new Map<string, string>();
	async function visit(relativeDir: string): Promise<void> {
		const absoluteDir = path.join(root, relativeDir);
		for (const entry of await readdir(absoluteDir)) {
			const relativePath = path.join(relativeDir, entry);
			const absolutePath = path.join(root, relativePath);
			const metadata = await lstat(absolutePath);
			if (metadata.isDirectory()) {
				inventory.set(relativePath, "directory");
				await visit(relativePath);
			} else if (metadata.isSymbolicLink()) {
				inventory.set(relativePath, `symlink:${await readlink(absolutePath)}`);
			} else if (metadata.isFile()) {
				inventory.set(relativePath, "file");
			} else {
				inventory.set(relativePath, "unsupported");
			}
		}
	}
	await visit("");
	return inventory;
}

async function assertReplacementPreservesOwnedPaths(currentBundle: string, newBundle: string): Promise<void> {
	const [currentInventory, newInventory] = await Promise.all([
		bundleInventory(currentBundle),
		bundleInventory(newBundle),
	]);
	const removedOrChanged = [...currentInventory].filter(
		([relativePath, kind]) => newInventory.get(relativePath) !== kind,
	);
	if (removedOrChanged.length > 0) {
		const preview = removedOrChanged.slice(0, 10).map(([relativePath]) => relativePath).join(", ");
		throw new Error(
			`Native update would remove or change unrecognized bundle paths: ${preview}` +
				(removedOrChanged.length > 10 ? ` (+${removedOrChanged.length - 10} more)` : ""),
		);
	}
}

async function replaceNativeBundleUnix(currentBundle: string, extractedBundle: string, oldVersion: string): Promise<void> {
	await assertReplacementPreservesOwnedPaths(currentBundle, extractedBundle);
	const parent = path.dirname(currentBundle);
	const stagingRoot = await mkdtemp(path.join(parent, ".pi-update-"));
	const stagedBundle = path.join(stagingRoot, "bundle");
	const backup = `${currentBundle}.previous-${safeVersion(oldVersion)}-${Date.now()}`;
	try {
		await cp(extractedBundle, stagedBundle, { recursive: true, force: false, errorOnExist: true });
		await rename(currentBundle, backup);
		try {
			await rename(stagedBundle, currentBundle);
		} catch (error) {
			try {
				await rename(backup, currentBundle);
			} catch (rollbackError) {
				throw new NativeReplacementFailure(
					`Native Pi replacement and rollback both failed. Backup remains at ${backup}: ${describeError(rollbackError)}`,
					{ cause: error },
				);
			}
			throw error;
		}
	} finally {
		await rm(stagingRoot, { recursive: true, force: true }).catch(() => undefined);
	}
}

async function updateNativePi(
	pi: ExtensionAPI,
	currentVersion: string,
	reporter: UpdateReporter,
): Promise<NativeUpdateResult> {
	if (isGenericRuntime(process.execPath)) {
		throw new Error("Native Pi detection selected a generic runtime; refusing an unsafe binary replacement.");
	}
	if (process.platform === "win32") {
		throw new Error(
			"Automatic updates for a running native Windows bundle are unsafe. Download and verify the latest Windows archive from https://github.com/earendil-works/pi/releases/latest.",
		);
	}
	const currentBundle = path.dirname(await realpath(process.execPath));
	try {
		await validateNativeBundle(currentBundle, "pi", currentVersion);
	} catch (error) {
		throw new Error(
			`The native Pi executable is not inside an exact dedicated release bundle (${currentBundle}). ` +
				`Automatic replacement is unsafe: ${error instanceof Error ? error.message : String(error)}`,
		);
	}

	const release = (await (await fetchWithRetry(RELEASE_API, reporter)).json()) as GitHubRelease;
	const version = release.tag_name.replace(/^v/, "");
	if (compareSemver(version, currentVersion) <= 0) return { version: currentVersion, requiresShutdown: false };

	const assetName = releaseAssetName();
	const archive = release.assets.find((asset) => asset.name === assetName);
	const checksums = release.assets.find((asset) => asset.name === "SHA256SUMS");
	if (!archive || !checksums) throw new Error(`Release ${release.tag_name} is missing ${assetName} or SHA256SUMS`);

	const workingDir = await mkdtemp(path.join(tmpdir(), "pi-native-update-"));
	try {
		const archivePath = path.join(workingDir, assetName);
		const checksumsPath = path.join(workingDir, "SHA256SUMS");
		reporter.info(`Downloading verified native Pi ${version}`);
		await downloadFile(archive.browser_download_url, archivePath, reporter);
		await downloadFile(checksums.browser_download_url, checksumsPath, reporter);
		await verifyReleaseChecksum(archivePath, assetName, checksumsPath);
		const extractedBundle = await extractNativeArchive(pi, archivePath, path.join(workingDir, "extract"));
		await validateNativeBundle(extractedBundle, "pi", version);
		await validateNewNativeExecutable(pi, extractedBundle, version);
		await replaceNativeBundleUnix(currentBundle, extractedBundle, currentVersion);
		return { version, requiresShutdown: true };
	} finally {
		await rm(workingDir, { recursive: true, force: true });
	}
}

async function updateExtensions(
	pi: ExtensionAPI,
	reporter: UpdateReporter,
	projectTrusted: boolean,
): Promise<void> {
	reporter.info("Updating configured Pi extension packages");
	const args = ["update", "--extensions", projectTrusted ? "--approve" : "--no-approve"];
	await runWithRetry(pi, getPiInvocation(args), reporter);
}

async function updateSelf(
	pi: ExtensionAPI,
	method: InstallMethod,
	beforeVersion: string,
	reporter: UpdateReporter,
): Promise<NativeUpdateResult> {
	reporter.info(`Updating Pi (${method} install)`);
	switch (method) {
		case "vp":
			await runWithRetry(pi, commandSpec("vp", ["update", "-g", `${PI_PACKAGE}@latest`]), reporter);
			return { version: await readPiVersion(pi), requiresShutdown: false };
		case "brew":
			await runWithRetry(pi, commandSpec("brew", ["upgrade", BREW_FORMULA]), reporter);
			return { version: await readPiVersion(pi), requiresShutdown: false };
		case "bun":
		case "npm":
			await runWithRetry(pi, getPiInvocation(["update", "--self"]), reporter);
			return { version: await readPiVersion(pi), requiresShutdown: false };
		case "native":
			return updateNativePi(pi, beforeVersion, reporter);
		default:
			throw new Error(`Unsupported Pi install method: ${String(method)}`);
	}
}

export async function updatePiAndExtensions(
	pi: ExtensionAPI,
	reporter: UpdateReporter,
	options: { projectTrusted?: boolean } = {},
): Promise<UpdateResult> {
	if (/^(?:1|true|yes)$/i.test(process.env.PI_OFFLINE ?? "")) {
		throw new Error("Pi updates are disabled because PI_OFFLINE is enabled.");
	}
	const method = await detectInstallMethod();
	const beforeVersion = await readPiVersion(pi);
	let extensionError: Error | undefined;
	let selfError: Error | undefined;
	let extensionsChecked = false;
	let selfResult: NativeUpdateResult = { version: beforeVersion, requiresShutdown: false };

	try {
		await updateExtensions(pi, reporter, options.projectTrusted === true);
		extensionsChecked = true;
	} catch (error) {
		extensionError = error instanceof Error ? error : new Error(String(error));
	}
	try {
		selfResult = await updateSelf(pi, method, beforeVersion, reporter);
	} catch (error) {
		selfError = error instanceof Error ? error : new Error(String(error));
		if (error instanceof NativeReplacementFailure) {
			selfResult = { version: beforeVersion, requiresShutdown: true };
		}
	}

	if (extensionError || selfError) {
		const extensionReport = extensionError
			? `Extension update failed: ${extensionError.message}`
			: "Configured extension packages were checked successfully.";
		let selfReport: string;
		if (selfError) {
			selfReport = `Pi update failed: ${selfError.message}`;
		} else if (beforeVersion === selfResult.version) {
			selfReport = `Pi remains at v${selfResult.version}.`;
		} else {
			selfReport = `Pi updated: v${beforeVersion} → v${selfResult.version}.`;
		}
		const report = [extensionReport, selfReport];
		throw new UpdateFailure(report.join("\n\n"), selfResult.requiresShutdown);
	}

	return {
		method,
		beforeVersion,
		afterVersion: selfResult.version,
		extensionsChecked,
		requiresShutdown: selfResult.requiresShutdown,
	};
}