pi-ext

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";
import { findExecutable } from "./src/pi-ext-executable.ts";
import { isShellToolCall } from "./src/pi-ext-shell-tool.ts";

const PIP_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?pip\s*(?:$|\s)/m;
const PIP3_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?pip3\s*(?:$|\s)/m;
const POETRY_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?poetry\s*(?:$|\s)/m;
const PYTHON_COMMAND = String.raw`(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?python(?:3(?:\.\d+)?)?\b[^\n;|&]*`;
const PYTHON_PIP_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*pip\\b|\\s-mpip\\b)`,
	"m",
);
const PYTHON_VENV_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*venv\\b|\\s-mvenv\\b)`,
	"m",
);
const PYTHON_COMPILE_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*py_compile\\b|\\s-mpy_compile\\b)`,
	"m",
);

async function uvInstalled(
	env: NodeJS.ProcessEnv = process.env,
	platform: NodeJS.Platform = process.platform,
): Promise<boolean> {
	return Boolean(
		findExecutable(
			platform === "win32" ? ["uv.exe", "uv.cmd", "uv.bat"] : ["uv"],
			{ path: env[pathEnvKey(env, platform)] ?? "", platform },
		),
	);
}

function pathEnvKey(
	env: NodeJS.ProcessEnv,
	platform = process.platform,
): string {
	if (platform !== "win32") return "PATH";
	return Object.keys(env).find((key) => key.toLowerCase() === "path") ?? "PATH";
}

function blocked(message: string, ...body: string[]): string {
	return [
		`Error: ${message}. Use uv instead:`,
		"",
		...body.map((line) => `  ${line}`),
		"",
	].join("\n");
}

const PIP_BLOCKED = blocked(
	"pip is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const PIP3_BLOCKED = blocked(
	"pip3 is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const POETRY_BLOCKED = blocked(
	"poetry is disabled",
	"To initialize a project: uv init",
	"To add a dependency: uv add PACKAGE",
	"To sync dependencies: uv sync",
	"To run commands: uv run COMMAND",
);
const PYTHON_PIP_BLOCKED = blocked(
	"'python -m pip' is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const PYTHON_VENV_BLOCKED = blocked(
	"'python -m venv' is disabled",
	"To create a virtual environment: uv venv",
);
const PYTHON_COMPILE_BLOCKED = [
	"Error: 'python -m py_compile' is disabled because it writes .pyc files to __pycache__.",
	"",
	"  To verify syntax without bytecode output: uv run python -m ast path/to/file.py >/dev/null",
	"",
].join("\n");

function getBlockedCommandMessage(command: string): string | null {
	if (PIP_COMMAND.test(command)) return PIP_BLOCKED;
	if (PIP3_COMMAND.test(command)) return PIP3_BLOCKED;
	if (POETRY_COMMAND.test(command)) return POETRY_BLOCKED;
	if (PYTHON_PIP_COMMAND.test(command)) return PYTHON_PIP_BLOCKED;
	if (PYTHON_VENV_COMMAND.test(command)) return PYTHON_VENV_BLOCKED;
	if (PYTHON_COMPILE_COMMAND.test(command)) return PYTHON_COMPILE_BLOCKED;
	return null;
}

export const __test: {
	getBlockedCommandMessage: typeof getBlockedCommandMessage;
	uvInstalled: typeof uvInstalled;
} = {
	getBlockedCommandMessage,
	uvInstalled,
};

export default function uvExtension(pi: ExtensionAPI): void {
	let generation = 0;

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start");
		const ticket = ++generation;
		const finish = span?.("availability.check");
		void uvInstalled().then(
			(installed) => {
				finish?.("finish", { installed });
				if (ticket === generation && !installed && ctx.hasUI)
					ctx.ui.notify("uv extension: uv not found on PATH", "warning");
			},
			(error) => {
				finish?.("error");
				throw error;
			},
		);
	});
	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
		generation++;
	});

	pi.on("tool_call", (event) => {
		if (!isShellToolCall(event, pi)) return;
		const message = getBlockedCommandMessage(event.input.command);
		if (!message) return;
		dbg?.("command.block");
		return { block: true, reason: message };
	});
}