import { delimiter, join } from "node:path"; import { type BashSpawnContext, createBashTool, type ExtensionAPI, } from "@earendil-works/pi-coding-agent"; import { closeDebug, dbg, span } from "./src/debug.ts"; import { findExecutable } from "./src/pi-ext-executable.ts"; const INTERCEPTED_COMMANDS_PATH = join( import.meta.dirname, "intercepted-commands", ); 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 = process.env, platform = process.platform, ): Promise { 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 prependInterceptedCommands( ctx: BashSpawnContext, platform = process.platform, ): BashSpawnContext { if (platform === "win32") return ctx; const key = pathEnvKey(ctx.env); const path = [INTERCEPTED_COMMANDS_PATH, ctx.env[key]] .filter(Boolean) .join(delimiter); return { ...ctx, env: { ...ctx.env, [key]: 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, prependInterceptedCommands, uvInstalled, }; export default function uvExtension(pi: ExtensionAPI) { 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++; }); const bash = createBashTool(process.cwd(), { spawnHook: (ctx) => { const message = getBlockedCommandMessage(ctx.command); if (message) throw new Error(message); return prependInterceptedCommands(ctx); }, }); const promptMetadata = bash as typeof bash & { promptGuidelines?: readonly string[]; }; Object.assign(promptMetadata, { description: `${bash.description} This replacement routes PATH-based Python through uv and rejects pip, Poetry, python -m pip, python -m venv, and python -m py_compile.`, promptGuidelines: [ ...(promptMetadata.promptGuidelines ?? []).map((guideline) => guideline === "You can inspect PI_* environment variables for current model and session details." ? "Use bash to inspect PI_* environment variables for current model and session details." : guideline, ), "Use bash with uv instead of pip, pip3, Poetry, python -m pip, or python -m venv for Python package and environment operations; use uv run python -m ast PATH >/dev/null instead of python -m py_compile for syntax checks.", ], }); pi.registerTool(bash); }