From a48108bd1d99d403f190bfb9749bbba2a3ca6d09 Mon Sep 17 00:00:00 2001 From: Kevin Thomas Date: Fri, 8 May 2026 18:53:51 -0700 Subject: [PATCH] auto bootstrap the kernel runtime (#19) * add auto bootstrap for kernel runtime * guard kernel bootstrap concurrency --- package-lock.json | 1 + packages/coding-agent/CHANGELOG.md | 4 + packages/coding-agent/README.md | 3 + .../docs/kernel-and-rlm-recursion.md | 23 +- packages/coding-agent/package.json | 4 +- packages/coding-agent/postinstall.cjs | 14 + .../src/core/kernel/bootstrap-cli.ts | 13 + .../coding-agent/src/core/kernel/bootstrap.ts | 352 ++++++++++++++++++ .../coding-agent/src/core/kernel/index.ts | 51 +-- .../coding-agent/src/core/tools/ipython.ts | 14 +- packages/coding-agent/src/postinstall.ts | 15 + .../test/kernel-bootstrap.test.ts | 173 +++++++++ scripts/setup-kernel-venv.sh | 26 +- 13 files changed, 614 insertions(+), 79 deletions(-) create mode 100644 packages/coding-agent/postinstall.cjs create mode 100644 packages/coding-agent/src/core/kernel/bootstrap-cli.ts create mode 100644 packages/coding-agent/src/core/kernel/bootstrap.ts create mode 100644 packages/coding-agent/src/postinstall.ts create mode 100644 packages/coding-agent/test/kernel-bootstrap.test.ts diff --git a/package-lock.json b/package-lock.json index 77b8d2cef..8c7a4b7ed 100644 --- a/package-lock.json +++ b/package-lock.json @@ -7399,6 +7399,7 @@ "packages/coding-agent": { "name": "@earendil-works/pi-coding-agent", "version": "0.74.0", + "hasInstallScript": true, "license": "MIT", "dependencies": { "@earendil-works/pi-agent-core": "^0.74.0", diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 20d4b0ced..c9e32387c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added automatic IPython kernel runtime bootstrap with uv-managed Python, `ipykernel`, and `prime-agent-runtime`. + ### Changed - Changed the default active built-in tool set to `ipython`. diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 11127a4d8..2cdac6133 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -93,6 +93,8 @@ pi Then just talk to pi. By default, pi gives the model one tool: `ipython`. The model uses the persistent kernel to read files, run commands, edit code, and inspect data. Add capabilities via [skills](#skills), [prompt templates](#prompt-templates), [extensions](#extensions), or [pi packages](#pi-packages). +The Python kernel runtime is set up automatically on first invocation. Set `PRIME_AGENT_KERNEL_PYTHON` to use an existing Python environment with `ipykernel`. + **Platform notes:** [Windows](docs/windows.md) | [Termux (Android)](docs/termux.md) | [tmux](docs/tmux.md) | [Terminal setup](docs/terminal-setup.md) | [Shell aliases](docs/shell-aliases.md) --- @@ -631,6 +633,7 @@ pi --thinking high "Solve this complex problem" | `PI_SKIP_VERSION_CHECK` | Skip the Pi version update check at startup. This prevents the `pi.dev` latest-version request | | `PI_TELEMETRY` | Override install/update telemetry. Use `1`/`true`/`yes` to enable or `0`/`false`/`no` to disable. This does not disable update checks | | `PI_CACHE_RETENTION` | Set to `long` for extended prompt cache (Anthropic: 1h, OpenAI: 24h) | +| `PRIME_AGENT_KERNEL_PYTHON` | Use an existing Python environment with `ipykernel` instead of auto-bootstrapping `~/.prime/agent/kernel-venv` | | `VISUAL`, `EDITOR` | External editor for Ctrl+G | --- diff --git a/packages/coding-agent/docs/kernel-and-rlm-recursion.md b/packages/coding-agent/docs/kernel-and-rlm-recursion.md index c68a5985b..611b0f3ec 100644 --- a/packages/coding-agent/docs/kernel-and-rlm-recursion.md +++ b/packages/coding-agent/docs/kernel-and-rlm-recursion.md @@ -69,11 +69,11 @@ packages/coding-agent/src/core/rlm-runtime.ts TypeScript request/result types for the rlm.run bridge. prime-agent-runtime/src/rlm/__init__.py - Python shim installed into .kernel-venv. Exposes rlm, rlm.run(), + Python shim installed into ~/.prime/agent/kernel-venv. Exposes rlm, rlm.run(), RLMResult, and TokenUsage. scripts/setup-kernel-venv.sh - Creates .kernel-venv and installs ipykernel plus prime-agent-runtime. + Thin wrapper around the automatic kernel bootstrap. ``` ## ZeroMQ Jupyter Kernel Setup @@ -173,6 +173,11 @@ Startup sequence: ```text KernelManager.start() + | + | ensureKernelPython() + | bootstraps ~/.prime/agent/kernel-venv with uv if needed + v +resolved python with ipykernel and prime-agent-runtime | v makeConnection() @@ -255,14 +260,16 @@ Process-wide cleanup is registered through `registerSessionResourceCleanup()`. K ## IPython Tool Bootstrapping -The ipython tool constructs a `KernelManager` lazily on first use. It resolves Python in this order: +The ipython tool constructs a `KernelManager` lazily on first use. `KernelManager.start()` resolves Python in this order: ```text -PRIME_AGENT_KERNEL_PYTHON -~/.prime-agent/kernel-venv/bin/python -/.kernel-venv/bin/python +PRIME_AGENT_KERNEL_PYTHON, if set and able to import ipykernel +~/.prime/agent/kernel-venv/bin/python, bootstrapped by uv if needed +$XDG_DATA_HOME/prime/agent/kernel-venv/bin/python, only if ~/.prime is not writable ``` +The automatic bootstrap installs Python 3.11, `ipykernel`, and `prime-agent-runtime`, then writes `.bootstrap-version` in the venv. Missing or stale markers cause a rebuild. + After the kernel starts, the tool bootstraps `rlm` into the user namespace: ```python @@ -595,9 +602,9 @@ Those stay in TypeScript so children share the parent harness implementation: Common failures and where they surface: ```text -prime-agent-runtime missing from .kernel-venv +prime-agent-runtime missing from the kernel environment -> kernel startup succeeds - -> rlm.run raises RuntimeError with setup instructions when called + -> rlm.run raises RuntimeError with rebuild or PRIME_AGENT_KERNEL_PYTHON instructions when called RLM_DEPTH >= RLM_MAX_DEPTH -> Python RuntimeError before comm_open diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index 28e8a7c1f..f5e6b0adb 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -25,6 +25,7 @@ "dist", "docs", "examples", + "postinstall.cjs", "CHANGELOG.md" ], "scripts": { @@ -32,9 +33,10 @@ "dev": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput", "build": "tsgo -p tsconfig.build.json && shx chmod +x dist/cli.js && npm run copy-assets", "build:binary": "npm --prefix ../tui run build && npm --prefix ../ai run build && npm --prefix ../agent run build && npm run build && bun build --compile ./dist/bun/cli.js --outfile dist/pi && npm run copy-binary-assets", - "copy-assets": "shx mkdir -p dist/modes/interactive/theme && shx cp src/modes/interactive/theme/*.json dist/modes/interactive/theme/ && shx mkdir -p dist/modes/interactive/assets && shx cp src/modes/interactive/assets/*.png dist/modes/interactive/assets/ && shx mkdir -p dist/core/export-html/vendor && shx cp src/core/export-html/template.html src/core/export-html/template.css src/core/export-html/template.js dist/core/export-html/ && shx cp src/core/export-html/vendor/*.js dist/core/export-html/vendor/", + "copy-assets": "shx mkdir -p dist/modes/interactive/theme && shx cp src/modes/interactive/theme/*.json dist/modes/interactive/theme/ && shx mkdir -p dist/modes/interactive/assets && shx cp src/modes/interactive/assets/*.png dist/modes/interactive/assets/ && shx mkdir -p dist/core/export-html/vendor && shx cp src/core/export-html/template.html src/core/export-html/template.css src/core/export-html/template.js dist/core/export-html/ && shx cp src/core/export-html/vendor/*.js dist/core/export-html/vendor/ && shx rm -rf dist/prime-agent-runtime && shx cp -r ../../prime-agent-runtime dist/prime-agent-runtime", "copy-binary-assets": "shx cp package.json dist/ && shx cp README.md dist/ && shx cp CHANGELOG.md dist/ && shx mkdir -p dist/theme && shx cp src/modes/interactive/theme/*.json dist/theme/ && shx mkdir -p dist/assets && shx cp src/modes/interactive/assets/*.png dist/assets/ && shx mkdir -p dist/export-html/vendor && shx cp src/core/export-html/template.html dist/export-html/ && shx cp src/core/export-html/vendor/*.js dist/export-html/vendor/ && shx cp -r docs dist/ && shx cp -r examples dist/ && shx cp ../../node_modules/@silvia-odwyer/photon-node/photon_rs_bg.wasm dist/", "test": "vitest --run", + "postinstall": "node postinstall.cjs", "prepublishOnly": "npm run clean && npm run build" }, "dependencies": { diff --git a/packages/coding-agent/postinstall.cjs b/packages/coding-agent/postinstall.cjs new file mode 100644 index 000000000..cbe986b37 --- /dev/null +++ b/packages/coding-agent/postinstall.cjs @@ -0,0 +1,14 @@ +const { spawnSync } = require("node:child_process"); +const { existsSync } = require("node:fs"); +const { join } = require("node:path"); + +const script = join(__dirname, "dist", "postinstall.js"); +if (!existsSync(script)) { + process.exit(0); +} + +const result = spawnSync(process.execPath, [script], { stdio: "inherit" }); +if (result.error) { + console.error(`prime-agent: python kernel setup skipped: ${result.error.message}`); +} +process.exit(0); diff --git a/packages/coding-agent/src/core/kernel/bootstrap-cli.ts b/packages/coding-agent/src/core/kernel/bootstrap-cli.ts new file mode 100644 index 000000000..923aee69f --- /dev/null +++ b/packages/coding-agent/src/core/kernel/bootstrap-cli.ts @@ -0,0 +1,13 @@ +import { ensureKernelPython } from "./bootstrap.js"; + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +try { + const python = await ensureKernelPython(); + console.log(`kernel python: ${python}`); +} catch (error) { + console.error(errorMessage(error)); + process.exit(1); +} diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts new file mode 100644 index 000000000..0a7c245b2 --- /dev/null +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -0,0 +1,352 @@ +import { spawn } from "node:child_process"; +import { constants, existsSync } from "node:fs"; +import { access, mkdir, readFile, rm, stat, writeFile } from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { setTimeout as sleep } from "node:timers/promises"; +import { fileURLToPath } from "node:url"; + +const BOOTSTRAP_SCHEMA = 1; +const PYTHON_VERSION = "3.11"; +const IPYKERNEL_REQUIREMENT = "ipykernel"; +const RUNTIME_REQUIREMENT = "prime-agent-runtime"; +const BOOTSTRAP_VERSION_FILE = ".bootstrap-version"; +const BOOTSTRAP_LOCK_NAME = ".bootstrap.lock"; +const BOOTSTRAP_LOCK_RETRY_MS = 100; +const BOOTSTRAP_LOCK_STALE_WITHOUT_PID_MS = 30_000; + +let inFlightEnsureKernelPython: { key: string; promise: Promise } | null = null; + +interface BootstrapVersion { + schema: number; + ipykernel?: string; + runtime?: string; +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isNodeError(error: unknown, code: string): boolean { + return error instanceof Error && "code" in error && error.code === code; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +async function exists(filePath: string): Promise { + try { + await access(filePath); + return true; + } catch { + return false; + } +} + +async function isExecutable(filePath: string): Promise { + try { + await access(filePath, constants.X_OK); + return true; + } catch { + return false; + } +} + +function expandHome(filePath: string): string { + if (filePath === "~") return os.homedir(); + if (filePath.startsWith("~/")) return path.join(os.homedir(), filePath.slice(2)); + return filePath; +} + +function ensureKernelPythonKey(): string { + return [ + process.env.PRIME_AGENT_KERNEL_PYTHON ?? "", + process.env.PRIME_AGENT_KERNEL_VENV ?? "", + process.env.HOME ?? "", + process.env.XDG_DATA_HOME ?? "", + ].join("\0"); +} + +export function getKernelVenvDir(): string { + const override = process.env.PRIME_AGENT_KERNEL_VENV; + if (override) return path.resolve(expandHome(override)); + return path.join(os.homedir(), ".prime", "agent", "kernel-venv"); +} + +function getXdgKernelVenvDir(): string { + const dataHome = process.env.XDG_DATA_HOME + ? path.resolve(expandHome(process.env.XDG_DATA_HOME)) + : path.join(os.homedir(), ".local", "share"); + return path.join(dataHome, "prime", "agent", "kernel-venv"); +} + +async function resolveWritableKernelVenvDir(): Promise { + const primary = getKernelVenvDir(); + try { + await mkdir(path.dirname(primary), { recursive: true }); + return primary; + } catch (primaryError) { + if (process.env.PRIME_AGENT_KERNEL_VENV) { + throw new Error(`couldn't create kernel venv parent directory for ${primary}: ${errorMessage(primaryError)}`); + } + + const fallback = getXdgKernelVenvDir(); + try { + await mkdir(path.dirname(fallback), { recursive: true }); + return fallback; + } catch (fallbackError) { + throw new Error( + `couldn't create kernel venv directory at ${primary} or ${fallback}; set PRIME_AGENT_KERNEL_PYTHON to a python with ipykernel installed. ${errorMessage(fallbackError)}`, + ); + } + } +} + +function run(command: string, args: string[], options: { stdio?: "ignore" | "inherit" } = {}): Promise { + return new Promise((resolve, reject) => { + const child = spawn(command, args, { + env: process.env, + stdio: options.stdio ?? "ignore", + }); + child.on("error", reject); + child.on("exit", (code, signal) => { + if (code === 0) { + resolve(); + return; + } + const reason = signal ? `signal ${signal}` : `exit code ${code}`; + reject(new Error(`${command} ${args.join(" ")} failed with ${reason}`)); + }); + }); +} + +async function pythonImports(python: string, moduleName: string): Promise { + try { + await run(python, ["-c", `import ${moduleName}`], { stdio: "ignore" }); + return true; + } catch { + return false; + } +} + +async function hasIpykernel(python: string): Promise { + return pythonImports(python, "ipykernel"); +} + +async function hasPrimeAgentRuntime(python: string): Promise { + return pythonImports(python, "rlm"); +} + +function bootstrapLockDir(venv: string): string { + return path.join(path.dirname(venv), `${path.basename(venv)}${BOOTSTRAP_LOCK_NAME}`); +} + +function processIsRunning(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + return isNodeError(error, "EPERM"); + } +} + +async function readLockPid(lockDir: string): Promise { + try { + const raw = await readFile(path.join(lockDir, "pid"), "utf8"); + const pid = Number.parseInt(raw.trim(), 10); + return Number.isInteger(pid) && pid > 0 ? pid : null; + } catch { + return null; + } +} + +async function lockMissingPidIsStale(lockDir: string): Promise { + try { + const lockStat = await stat(lockDir); + return Date.now() - lockStat.mtimeMs > BOOTSTRAP_LOCK_STALE_WITHOUT_PID_MS; + } catch { + return false; + } +} + +async function acquireBootstrapLock(venv: string): Promise<() => Promise> { + const lockDir = bootstrapLockDir(venv); + await mkdir(path.dirname(lockDir), { recursive: true }); + + for (;;) { + try { + await mkdir(lockDir); + await writeFile(path.join(lockDir, "pid"), `${process.pid}\n`, "utf8"); + return () => rm(lockDir, { recursive: true, force: true }); + } catch (error) { + if (!isNodeError(error, "EEXIST")) throw error; + + const pid = await readLockPid(lockDir); + if (pid === null ? await lockMissingPidIsStale(lockDir) : !processIsRunning(pid)) { + await rm(lockDir, { recursive: true, force: true }); + continue; + } + + await sleep(BOOTSTRAP_LOCK_RETRY_MS); + } + } +} + +async function findExecutable(name: string): Promise { + const pathValue = process.env.PATH; + if (!pathValue) return null; + const candidates = process.platform === "win32" ? [name, `${name}.exe`] : [name]; + for (const dir of pathValue.split(path.delimiter)) { + if (!dir) continue; + for (const candidate of candidates) { + const fullPath = path.join(dir, candidate); + if (await isExecutable(fullPath)) return fullPath; + } + } + return null; +} + +async function ensureUv(): Promise { + const fromPath = await findExecutable("uv"); + if (fromPath) return fromPath; + + const localUv = path.join(os.homedir(), ".local", "bin", process.platform === "win32" ? "uv.exe" : "uv"); + if (await isExecutable(localUv)) return localUv; + + process.stderr.write("› installing uv (one-time)…\n"); + try { + await run("sh", ["-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"], { stdio: "inherit" }); + } catch (error) { + throw new Error( + `couldn't install uv from astral.sh; install it yourself: curl -LsSf https://astral.sh/uv/install.sh | sh, then re-run prime-agent. ${errorMessage(error)}`, + ); + } + + if (await isExecutable(localUv)) return localUv; + const installedFromPath = await findExecutable("uv"); + if (installedFromPath) return installedFromPath; + throw new Error("uv install completed but binary not found at ~/.local/bin/uv"); +} + +async function readBootstrapVersion(venv: string): Promise { + try { + const raw = await readFile(path.join(venv, BOOTSTRAP_VERSION_FILE), "utf8"); + const parsed: unknown = JSON.parse(raw); + if (!isRecord(parsed) || typeof parsed.schema !== "number") return null; + return { + schema: parsed.schema, + ipykernel: typeof parsed.ipykernel === "string" ? parsed.ipykernel : undefined, + runtime: typeof parsed.runtime === "string" ? parsed.runtime : undefined, + }; + } catch { + return null; + } +} + +function bootstrapVersionCurrent(version: BootstrapVersion | null): boolean { + return ( + version?.schema === BOOTSTRAP_SCHEMA && + version.ipykernel === IPYKERNEL_REQUIREMENT && + version.runtime === RUNTIME_REQUIREMENT + ); +} + +async function writeBootstrapVersion(venv: string): Promise { + const version: BootstrapVersion = { + schema: BOOTSTRAP_SCHEMA, + ipykernel: IPYKERNEL_REQUIREMENT, + runtime: RUNTIME_REQUIREMENT, + }; + await writeFile(path.join(venv, BOOTSTRAP_VERSION_FILE), `${JSON.stringify(version)}\n`, "utf8"); +} + +function runtimeCandidateDirs(): string[] { + const moduleDir = path.dirname(fileURLToPath(import.meta.url)); + return [ + path.resolve(moduleDir, "..", "..", "prime-agent-runtime"), + path.resolve(moduleDir, "..", "..", "..", "..", "..", "prime-agent-runtime"), + ]; +} + +async function resolveRuntimeRequirement(): Promise { + for (const candidate of runtimeCandidateDirs()) { + if (await exists(path.join(candidate, "pyproject.toml"))) { + return candidate; + } + } + return RUNTIME_REQUIREMENT; +} + +async function bootstrapVenv(venv: string): Promise { + await mkdir(path.dirname(venv), { recursive: true }); + const uv = await ensureUv(); + const python = path.join(venv, "bin", "python"); + const runtimeRequirement = await resolveRuntimeRequirement(); + + await run(uv, ["python", "install", PYTHON_VERSION]); + await run(uv, ["venv", venv, "--python", PYTHON_VERSION, "--seed"]); + await run(uv, ["pip", "install", "--python", python, IPYKERNEL_REQUIREMENT, runtimeRequirement]); + await writeBootstrapVersion(venv); +} + +async function kernelReady(python: string, venv: string): Promise { + return ( + (await hasIpykernel(python)) && + (await hasPrimeAgentRuntime(python)) && + bootstrapVersionCurrent(await readBootstrapVersion(venv)) + ); +} + +function formatBootstrapFailure(error: unknown): Error { + return new Error( + `Failed to set up the Python kernel runtime. ${errorMessage(error)}\n` + + "First-time setup needs internet to install uv, Python, ipykernel, and prime-agent-runtime; once set up, prime-agent runs offline. " + + "Set PRIME_AGENT_KERNEL_PYTHON to a Python with ipykernel installed to skip auto-bootstrap.", + ); +} + +async function ensureKernelPythonUncached(): Promise { + const override = process.env.PRIME_AGENT_KERNEL_PYTHON; + if (override) { + const python = path.resolve(expandHome(override)); + if (await hasIpykernel(python)) return python; + throw new Error(`PRIME_AGENT_KERNEL_PYTHON points to a Python that cannot import ipykernel: ${python}`); + } + + const venv = await resolveWritableKernelVenvDir(); + const python = path.join(venv, "bin", "python"); + if (await kernelReady(python, venv)) return python; + + const releaseLock = await acquireBootstrapLock(venv); + try { + if (await kernelReady(python, venv)) return python; + + const hadVenv = existsSync(venv); + process.stderr.write("› setting up python kernel (one-time, ~30s)…\n"); + if (hadVenv) { + process.stderr.write("rebuilding kernel venv\n"); + await rm(venv, { recursive: true, force: true }); + } + + await bootstrapVenv(venv); + } catch (error) { + throw formatBootstrapFailure(error); + } finally { + await releaseLock().catch(() => undefined); + } + + process.stderr.write("✓ ready\n"); + return python; +} + +export function ensureKernelPython(): Promise { + const key = ensureKernelPythonKey(); + if (inFlightEnsureKernelPython?.key === key) return inFlightEnsureKernelPython.promise; + + const promise = ensureKernelPythonUncached().finally(() => { + if (inFlightEnsureKernelPython?.promise === promise) inFlightEnsureKernelPython = null; + }); + inFlightEnsureKernelPython = { key, promise }; + return promise; +} diff --git a/packages/coding-agent/src/core/kernel/index.ts b/packages/coding-agent/src/core/kernel/index.ts index 321afca90..94951d562 100644 --- a/packages/coding-agent/src/core/kernel/index.ts +++ b/packages/coding-agent/src/core/kernel/index.ts @@ -1,7 +1,7 @@ // TODO: reconsider persistent kernel vs stateless `python -c` once RLM-1 weights land. import { type ChildProcess, spawn } from "node:child_process"; import { createHmac, randomBytes } from "node:crypto"; -import { existsSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { setTimeout as sleep } from "node:timers/promises"; @@ -9,6 +9,7 @@ import { registerSessionResourceCleanup } from "@earendil-works/pi-ai"; import { v4 as uuid } from "uuid"; import { Dealer, Subscriber } from "zeromq"; import type { RlmRunHandler, RlmRunResult } from "../rlm-runtime.js"; +import { ensureKernelPython } from "./bootstrap.js"; const DELIM = Buffer.from(""); const PROTOCOL_VERSION = "5.3"; @@ -18,8 +19,8 @@ const DEFAULT_MAX_OUTPUT_CHARS = 65536; const RLM_DISPOSE_TIMEOUT_MS = 5000; export interface KernelManagerOptions { - /** Python interpreter that has `ipykernel` available. */ - python: string; + /** Python interpreter that has `ipykernel` available. Defaults to the auto-bootstrapped kernel. */ + python?: string; cwd?: string; env?: Record; sessionId?: string; @@ -207,8 +208,8 @@ function installSignalHandlersOnce(): void { // ---- kernel manager ------------------------------------------------------ export class KernelManager { - private readonly options: Required> & - Pick; + private readonly options: Pick & + Required>; private readonly session = uuid(); private readonly commTargets = new Map(); private readonly handledRlmCommIds = new Set(); @@ -227,9 +228,6 @@ export class KernelManager { private startPromise?: Promise; constructor(options: KernelManagerOptions) { - if (!existsSync(options.python)) { - throw new Error(`Python interpreter not found: ${options.python}`); - } this.options = { python: options.python, cwd: options.cwd, @@ -256,11 +254,20 @@ export class KernelManager { this.state = "starting"; installSignalHandlersOnce(); + let python: string; + try { + python = this.options.python ?? (await ensureKernelPython()); + this.options.python = python; + } catch (error) { + this.state = "idle"; + throw error; + } + const { info, path: connectionPath, tempDir } = makeConnection(); this.connection = info; this.tempDir = tempDir; - const kernel = spawn(this.options.python, ["-m", "ipykernel_launcher", "-f", connectionPath], { + const kernel = spawn(python, ["-m", "ipykernel_launcher", "-f", connectionPath], { cwd: this.options.cwd, env: this.options.env ? { ...process.env, ...this.options.env } : process.env, stdio: ["ignore", "pipe", "pipe"], @@ -682,29 +689,3 @@ export class KernelManager { return this.state === "running"; } } - -// ---- Python interpreter resolution --------------------------------------- - -/** - * Resolve the Python interpreter to use for the kernel. Searched in order: - * 1. PRIME_AGENT_KERNEL_PYTHON env var - * 2. ~/.prime-agent/kernel-venv/bin/python (canonical user-install location) - * 3. /.kernel-venv/bin/python (development; created by scripts/setup-kernel-venv.sh) - */ -export function resolveKernelPython(): string | null { - const envOverride = process.env.PRIME_AGENT_KERNEL_PYTHON; - if (envOverride && existsSync(envOverride)) return envOverride; - - const home = process.env.HOME; - if (home) { - const canonical = join(home, ".prime-agent", "kernel-venv", "bin", "python"); - if (existsSync(canonical)) return canonical; - } - - for (let dir = process.cwd(); dir !== "/"; dir = join(dir, "..")) { - const candidate = join(dir, ".kernel-venv", "bin", "python"); - if (existsSync(candidate)) return candidate; - } - - return null; -} diff --git a/packages/coding-agent/src/core/tools/ipython.ts b/packages/coding-agent/src/core/tools/ipython.ts index a22c6c1b9..298be271a 100644 --- a/packages/coding-agent/src/core/tools/ipython.ts +++ b/packages/coding-agent/src/core/tools/ipython.ts @@ -2,7 +2,7 @@ import type { AgentTool } from "@earendil-works/pi-agent-core"; import { type Static, Type } from "typebox"; import type { ToolDefinition } from "../extensions/types.js"; -import { KernelManager, resolveKernelPython } from "../kernel/index.js"; +import { KernelManager } from "../kernel/index.js"; import type { RlmRunHandler } from "../rlm-runtime.js"; import { wrapToolDefinition } from "./tool-definition-wrapper.js"; @@ -17,7 +17,7 @@ except Exception as _prime_agent_rlm_error: async def run(self, prompt, **kwargs): raise RuntimeError( "prime-agent-runtime is not installed in this IPython kernel. " - "Run ./scripts/setup-kernel-venv.sh from the repo root, or set " + "Remove ~/.prime/agent/kernel-venv so prime-agent can rebuild it, or set " "PRIME_AGENT_KERNEL_PYTHON to a kernel environment with prime-agent-runtime installed. " f"Import error: {_PRIME_AGENT_RLM_IMPORT_ERROR}" ) @@ -45,7 +45,7 @@ export interface IpythonToolDetails { } export interface IpythonToolOptions { - /** Defaults to {@link resolveKernelPython}. Must have `ipykernel` installed. */ + /** Python override. Must have `ipykernel` installed. */ python?: string; env?: Record; sessionId?: string; @@ -69,14 +69,8 @@ export function createIpythonToolDefinition( function getManager(): Promise { if (!managerPromise) { managerPromise = (async () => { - const python = options?.python ?? resolveKernelPython(); - if (!python) { - throw new Error( - "No Python interpreter with `ipykernel` was found. Run `./scripts/setup-kernel-venv.sh` from the repo root, or set PRIME_AGENT_KERNEL_PYTHON to point at a python that has ipykernel installed.", - ); - } const m = new KernelManager({ - python, + python: options?.python, cwd, env: options?.env, sessionId: options?.sessionId, diff --git a/packages/coding-agent/src/postinstall.ts b/packages/coding-agent/src/postinstall.ts new file mode 100644 index 000000000..ad85e6bc0 --- /dev/null +++ b/packages/coding-agent/src/postinstall.ts @@ -0,0 +1,15 @@ +import { ensureKernelPython } from "./core/kernel/bootstrap.js"; + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function oneLine(message: string): string { + return message.replace(/\s+/g, " ").trim(); +} + +try { + await ensureKernelPython(); +} catch (error) { + console.error(`prime-agent: python kernel setup skipped: ${oneLine(errorMessage(error))}`); +} diff --git a/packages/coding-agent/test/kernel-bootstrap.test.ts b/packages/coding-agent/test/kernel-bootstrap.test.ts new file mode 100644 index 000000000..a78c453f5 --- /dev/null +++ b/packages/coding-agent/test/kernel-bootstrap.test.ts @@ -0,0 +1,173 @@ +import { chmodSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { ensureKernelPython, getKernelVenvDir } from "../src/core/kernel/bootstrap.js"; + +let tempDir = ""; +let originalEnv: NodeJS.ProcessEnv; + +function writeExecutable(filePath: string, content: string): void { + writeFileSync(filePath, content); + chmodSync(filePath, 0o755); +} + +function writeBootstrapVersion(venv: string): void { + writeFileSync( + join(venv, ".bootstrap-version"), + `${JSON.stringify({ schema: 1, ipykernel: "ipykernel", runtime: "prime-agent-runtime" })}\n`, + ); +} + +function writeFakePython(filePath: string, importableModules: readonly string[]): void { + const cases = importableModules.map((moduleName) => ` "import ${moduleName}") exit 0 ;;`).join("\n"); + writeExecutable( + filePath, + [ + "#!/bin/sh", + 'if [ "$1" = "-c" ]; then', + ' case "$2" in', + cases, + " *) exit 1 ;;", + " esac", + "fi", + "exit 0", + "", + ].join("\n"), + ); +} + +function installFakeUv(): string { + const binDir = join(tempDir, "bin"); + mkdirSync(binDir, { recursive: true }); + const logPath = join(tempDir, "uv.log"); + process.env.UV_LOG = logPath; + process.env.PATH = `${binDir}${process.env.PATH ? `:${process.env.PATH}` : ""}`; + writeExecutable( + join(binDir, "uv"), + [ + "#!/bin/sh", + "set -eu", + 'printf "%s\\n" "$*" >> "$UV_LOG"', + 'if [ "$1" = "python" ]; then', + " exit 0", + "fi", + 'if [ "$1" = "venv" ]; then', + ' venv="$2"', + ' mkdir -p "$venv/bin"', + " cat > \"$venv/bin/python\" <<'PY'", + "#!/bin/sh", + 'if [ "$1" = "-c" ]; then', + ' case "$2" in', + ' "import ipykernel"|"import rlm") exit 0 ;;', + " *) exit 1 ;;", + " esac", + "fi", + "exit 0", + "PY", + ' chmod +x "$venv/bin/python"', + " exit 0", + "fi", + 'if [ "$1" = "pip" ]; then', + " exit 0", + "fi", + "exit 2", + "", + ].join("\n"), + ); + return logPath; +} + +describe("kernel bootstrap", () => { + beforeEach(() => { + originalEnv = { ...process.env }; + tempDir = mkdtempSync(join(tmpdir(), "prime-agent-kernel-bootstrap-")); + process.env.HOME = tempDir; + process.env.PATH = originalEnv.PATH ?? ""; + delete process.env.PRIME_AGENT_KERNEL_PYTHON; + delete process.env.PRIME_AGENT_KERNEL_VENV; + delete process.env.XDG_DATA_HOME; + }); + + afterEach(() => { + process.env = originalEnv; + if (tempDir) { + rmSync(tempDir, { recursive: true, force: true }); + tempDir = ""; + } + }); + + it("returns the configured kernel venv directory", () => { + const venv = join(tempDir, "custom-venv"); + process.env.PRIME_AGENT_KERNEL_VENV = venv; + + expect(getKernelVenvDir()).toBe(venv); + }); + + it("bootstraps a missing venv with uv, ipykernel, and prime-agent-runtime", async () => { + const logPath = installFakeUv(); + const venv = join(tempDir, "kernel-venv"); + process.env.PRIME_AGENT_KERNEL_VENV = venv; + + await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); + + const log = readFileSync(logPath, "utf8"); + expect(log).toContain("python install 3.11"); + expect(log).toContain(`venv ${venv} --python 3.11 --seed`); + expect(log).toContain("pip install --python"); + expect(log).toContain("ipykernel"); + expect(log).toContain("prime-agent-runtime"); + expect(readFileSync(join(venv, ".bootstrap-version"), "utf8")).toContain('"schema":1'); + }); + + it("shares concurrent bootstrap work in one process", async () => { + const logPath = installFakeUv(); + const venv = join(tempDir, "kernel-venv"); + const python = join(venv, "bin", "python"); + process.env.PRIME_AGENT_KERNEL_VENV = venv; + + await expect(Promise.all([ensureKernelPython(), ensureKernelPython()])).resolves.toEqual([python, python]); + + const log = readFileSync(logPath, "utf8"); + expect(log.split("\n").filter((line) => line.startsWith(`venv ${venv} `))).toHaveLength(1); + }); + + it("reuses a current warm venv without invoking uv", async () => { + const venv = join(tempDir, "kernel-venv"); + const python = join(venv, "bin", "python"); + mkdirSync(join(venv, "bin"), { recursive: true }); + writeFakePython(python, ["ipykernel", "rlm"]); + writeBootstrapVersion(venv); + process.env.PRIME_AGENT_KERNEL_VENV = venv; + + await expect(ensureKernelPython()).resolves.toBe(python); + }); + + it("rebuilds a broken venv", async () => { + const logPath = installFakeUv(); + const venv = join(tempDir, "kernel-venv"); + mkdirSync(join(venv, "bin"), { recursive: true }); + writeBootstrapVersion(venv); + process.env.PRIME_AGENT_KERNEL_VENV = venv; + + await expect(ensureKernelPython()).resolves.toBe(join(venv, "bin", "python")); + + expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); + }); + + it("uses PRIME_AGENT_KERNEL_PYTHON as an override contract", async () => { + const overridePython = join(tempDir, "override-python"); + writeFakePython(overridePython, ["ipykernel"]); + process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; + + await expect(ensureKernelPython()).resolves.toBe(overridePython); + }); + + it("fails an invalid PRIME_AGENT_KERNEL_PYTHON without bootstrapping", async () => { + const overridePython = join(tempDir, "override-python"); + writeFakePython(overridePython, []); + process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; + + await expect(ensureKernelPython()).rejects.toThrow(/cannot import ipykernel/); + }); +}); diff --git a/scripts/setup-kernel-venv.sh b/scripts/setup-kernel-venv.sh index 523e8c5d6..4a2df7314 100755 --- a/scripts/setup-kernel-venv.sh +++ b/scripts/setup-kernel-venv.sh @@ -1,28 +1,4 @@ #!/usr/bin/env bash -# Set up the Python venv used by the IPython kernel. -# -# Creates `.kernel-venv` at the repo root with `ipykernel` installed. -# The `ipython` tool's `resolveKernelPython()` finds this venv automatically -# when developing inside the repo; in user installs, the canonical location -# is ~/.prime-agent/kernel-venv (handled in a future PR). set -euo pipefail - cd "$(dirname "$0")/.." - -if ! command -v uv >/dev/null 2>&1; then - echo "uv is required. Install: https://docs.astral.sh/uv/" >&2 - exit 1 -fi - -echo ">>> creating .kernel-venv with Python 3.12" -uv venv --python 3.12 .kernel-venv - -echo ">>> installing ipykernel" -uv pip install --python .kernel-venv/bin/python ipykernel - -echo ">>> installing prime-agent-runtime" -uv pip install --python .kernel-venv/bin/python -e ./prime-agent-runtime - -echo -echo "Done. The ipython tool will use this venv automatically." -echo "Override with PRIME_AGENT_KERNEL_PYTHON=/path/to/python." +npx tsx packages/coding-agent/src/core/kernel/bootstrap-cli.ts -- 2.54.0