Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
21 kB · 600 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601import { spawnSync } from "node:child_process";import { createHash } from "node:crypto";import { chmodSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join, resolve } from "node:path";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { DEFAULT_RLM_EXTRA_IMPORT_NAMES, DEFAULT_RLM_EXTRA_UV_ARGS, ensureKernelPython, type KernelPythonSkill, kernelVenvPython, resolveRuntimeIdentity,} from "../src/core/kernel/bootstrap.js";
let tempDir = "";let originalEnv: NodeJS.ProcessEnv;let runtimeIdentity = "";
function pyprojectHash(pyprojectPath: string): string { return `sha256:${createHash("sha256").update(readFileSync(pyprojectPath)).digest("hex")}`;}
function writeExecutable(filePath: string, content: string): void { writeFileSync(filePath, content); chmodSync(filePath, 0o755);}
function writeBootstrapVersion(venv: string, pythonSkills: readonly KernelPythonSkill[] = []): void { writeFileSync( join(venv, ".bootstrap-version"), `${JSON.stringify({ schema: 9, runtime: runtimeIdentity, snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: pythonSkills.map((skill) => ({ importName: skill.importName, packagePath: skill.packagePath, pyprojectPath: skill.pyprojectPath, pyprojectHash: pyprojectHash(skill.pyprojectPath), })), })}\n`, );}
function createPythonSkill(name = "web-search"): KernelPythonSkill { const packagePath = join(tempDir, "skills", name); const importName = name.replaceAll("-", "_"); const pyprojectPath = join(packagePath, "pyproject.toml"); mkdirSync(join(packagePath, "src", importName), { recursive: true }); writeFileSync( pyprojectPath, `[project]name = "${name}"version = "0.1.0"`, ); writeFileSync(join(packagePath, "src", importName, "__init__.py"), "async def run():\n return 'ok'\n"); return { name, importName, packagePath, pyprojectPath, };}
function createPythonSkillWithDependency(name: string, dependencyName: string): KernelPythonSkill { const skill = createPythonSkill(name); writeFileSync( skill.pyprojectPath, `[project]name = "${name}"version = "0.1.0"dependencies = ["${dependencyName}"]`, ); return skill;}
function writeFakePython(filePath: string, importableModules: readonly string[]): void { const cases = importableModules.map((moduleName) => ` "import ${moduleName}") exit 0 ;;`).join("\n"); const runtimeCase = importableModules.includes("rlm") ? ' *"_harness_methods"*) exit 0 ;;' : ""; writeExecutable( filePath, [ "#!/bin/sh", 'if [ "$1" = "-c" ]; then', ' case "$2" in', cases, runtimeCase, " *) 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"); const extraImportCases = DEFAULT_RLM_EXTRA_IMPORT_NAMES.map((moduleName) => ` "import ${moduleName}") exit 0 ;;`); process.env.UV_LOG = logPath; process.env.PATH = `${binDir}${process.env.PATH ? `:${process.env.PATH}` : ""}`; writeExecutable( join(binDir, "uv"), [ "#!/bin/sh", "set -e", // Record every invocation plus the live venv's state as uv saw it, so a test // can assert what a concurrently booting kernel would have imported. 'printf "%s live=%s\\n" "$*" "$([ -e "$UV_LIVE/bin/python" ] && { [ -e "$UV_LIVE/.bootstrap-version" ] && echo complete || echo partial; } || echo absent)" >> "$UV_LOG"', 'if [ "$UV_GATE" != "" ]; then while [ ! -e "$UV_GATE" ]; do sleep 0.05; done; fi', '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 rlm") exit 0 ;;', ...extraImportCases, ' *"_harness_methods"*) [ "$UV_FAIL_HARNESS" = "1" ] && exit 1 || exit 0 ;;', " *) exit 1 ;;", " esac", "fi", "exit 0", "PY", ' chmod +x "$venv/bin/python"', " exit 0", "fi", 'if [ "$1" = "pip" ]; then', ' for arg in "$@"; do', ' if [ "$UV_FAIL_ARG" != "" ] && [ "$arg" = "$UV_FAIL_ARG" ]; then', " exit 1", " fi", " done", " exit 0", "fi", "exit 2", "", ].join("\n"), ); return logPath;}
describe("kernel bootstrap", () => { beforeEach(async () => { runtimeIdentity = await resolveRuntimeIdentity(); 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("bootstraps a missing venv with uv, the runtime, extras and editable Python skills", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const pythonSkill = createPythonSkill(); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(join(venv, "bin", "python"));
const log = readFileSync(logPath, "utf8"); expect(log).toContain("python install 3.11"); expect(log).toMatch(/venv .*kernel-venv\.staging-[^ ]+ --python 3\.11 --seed --relocatable/); expect(log).toContain("pip install --python"); expect(log).not.toContain("ipykernel"); expect(log).toContain("prime-agent-runtime"); expect(log).toContain("dill"); expect(log).toContain(`--editable ${pythonSkill.packagePath}`); for (const uvArg of DEFAULT_RLM_EXTRA_UV_ARGS) { expect(log).toContain(uvArg); }
// The on-disk manifest is what every later warm-start decision reads. const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version).toEqual({ schema: 9, runtime: runtimeIdentity, snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [ { importName: pythonSkill.importName, packagePath: pythonSkill.packagePath, pyprojectPath: pythonSkill.pyprojectPath, pyprojectHash: pyprojectHash(pythonSkill.pyprojectPath), }, ], }); expect(version.runtime).toMatch(/^sha256:/); });
it.each([ { name: "a sibling skill directory", packageName: undefined, dependency: "attach-image" }, { name: "a sibling whose package and directory names differ", packageName: "prime-agent-skill-attach-image", dependency: "prime-agent-skill-attach-image", }, { name: "a dependency declared with extras and a version bound", packageName: undefined, dependency: "attach-image[httpx]>4.0.0", }, ])("installs $name as an editable package", async ({ packageName, dependency }) => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const dependencySkill = createPythonSkill("attach-image"); if (packageName) { writeFileSync(dependencySkill.pyprojectPath, `[project]\nname = "${packageName}"\nversion = "0.1.0"\n`); } const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", dependency); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython({ pythonSkills: [dependentSkill] })).resolves.toBe(join(venv, "bin", "python"));
const log = readFileSync(logPath, "utf8"); expect(log).toContain(`--editable ${dependencySkill.packagePath}`); expect(log).toContain(`--editable ${dependentSkill.packagePath}`); });
it("syncs a warm venv when a Python skill pyproject changes", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); const pythonSkill = createPythonSkill(); mkdirSync(join(venv, "bin"), { recursive: true }); writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeBootstrapVersion(venv, [pythonSkill]); writeFileSync( pythonSkill.pyprojectPath, `[project]name = "${pythonSkill.name}"version = "0.1.0"dependencies = ["httpx"]`, ); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(python);
const log = readFileSync(logPath, "utf8"); expect(log).not.toContain("--python 3.11 --seed"); expect(log).toContain(`--editable ${pythonSkill.packagePath}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills[0].pyprojectHash).toBe(pyprojectHash(pythonSkill.pyprojectPath)); });
it("continues when a Python skill editable install fails and retries it next startup", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const goodSkill = createPythonSkill("good-skill"); const brokenSkill = createPythonSkill("broken-skill"); process.env.PRIME_AGENT_KERNEL_VENV = venv; process.env.UV_FAIL_ARG = brokenSkill.packagePath;
await expect(ensureKernelPython({ pythonSkills: [goodSkill, brokenSkill] })).resolves.toBe( join(venv, "bin", "python"), );
const log = readFileSync(logPath, "utf8"); expect(log).toContain(`--editable ${goodSkill.packagePath}`); expect(log).toContain(`--editable ${brokenSkill.packagePath}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills).toEqual([ { importName: goodSkill.importName, packagePath: goodSkill.packagePath, pyprojectPath: goodSkill.pyprojectPath, pyprojectHash: pyprojectHash(goodSkill.pyprojectPath), }, ]);
await expect(ensureKernelPython({ pythonSkills: [goodSkill, brokenSkill] })).resolves.toBe( join(venv, "bin", "python"), );
const retryLog = readFileSync(logPath, "utf8"); expect(retryLog.split("\n").filter((line) => line.startsWith("venv "))).toHaveLength(1); expect( retryLog.split("\n").filter((line) => line.includes(`--editable ${brokenSkill.packagePath}`)), ).toHaveLength(2); });
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 "))).toHaveLength(1); });
it("stages a build and swaps it in without a partial tree at the live path", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); const stale = JSON.stringify({ schema: 9, runtime: "sha256:stale", snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [], }); process.env.PRIME_AGENT_KERNEL_VENV = venv; process.env.UV_LIVE = venv; process.env.UV_GATE = join(tempDir, "build-gate");
const cold = ensureKernelPython(); // uv is gated mid-build: the path a booting kernel imports holds nothing yet, // rather than a half-installed tree. expect(existsSync(venv)).toBe(false); writeFileSync(process.env.UV_GATE, ""); await expect(cold).resolves.toBe(python); const coldLog = readFileSync(logPath, "utf8"); expect(coldLog).not.toContain("live=partial"); expect(coldLog).toContain(`venv ${venv}.staging-`); expect(existsSync(join(venv, ".bootstrap-version"))).toBe(true);
// Stale runtime identity: the next call rebuilds over the live venv while a // gated build runs, so a kernel booting then still imports the old complete tree. writeFileSync(join(venv, ".bootstrap-version"), `${stale}\n`); const warm = ensureKernelPython(); expect(spawnSync(python, ["-c", "import rlm"]).status).toBe(0); writeFileSync(process.env.UV_GATE, ""); await expect(warm).resolves.toBe(python); const warmLog = readFileSync(logPath, "utf8").slice(coldLog.length); expect(warmLog).not.toContain("live=absent"); expect(warmLog).not.toContain("live=partial"); expect(warmLog).toContain("live=complete"); });
it("discards a staged venv and preserves the live venv if the strict runtime check fails", async () => { installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); mkdirSync(join(venv, "bin"), { recursive: true }); const stale = JSON.stringify({ schema: 9, runtime: "sha256:stale", snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [], }); process.env.PRIME_AGENT_KERNEL_VENV = venv; writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeFileSync(join(venv, ".bootstrap-version"), `${stale}\n`); process.env.UV_FAIL_HARNESS = "1"; await expect(ensureKernelPython()).rejects.toThrow("does not satisfy the prime-agent-runtime check"); expect(existsSync(venv) && readFileSync(join(venv, ".bootstrap-version"), "utf8")).toBe(`${stale}\n`); });
it("waits for another process's build lock instead of importing a partial venv", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const lock = `${venv}.bootstrap.lock`; writeFileSync(lock, `${process.pid}\n`); process.env.PRIME_AGENT_KERNEL_VENV = venv; const progress: string[] = [];
const python = await ensureKernelPython({ onProgress: (message) => { progress.push(message); // Standing in for the other process: it finishes while we wait on its lock. if (message.includes("waiting for it to finish")) rmSync(lock, { force: true }); }, });
expect(python).toBe(join(venv, "bin", "python")); expect(progress.filter((message) => message.includes("waiting for it to finish"))).toHaveLength(1); expect(readFileSync(logPath, "utf8")).toContain("--python 3.11 --seed --relocatable"); });
it.each([ { name: "reuses a current warm venv without invoking uv", rebuilds: false, prepare: (venv: string, python: string) => { writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeBootstrapVersion(venv); }, }, { name: "rebuilds when the recorded runtime hash no longer matches local source", rebuilds: true, prepare: (venv: string, python: string) => { writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeFileSync( join(venv, ".bootstrap-version"), `${JSON.stringify({ schema: 9, runtime: "sha256:stale", snapshot: "dill", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [], })}\n`, ); }, }, { name: "rebuilds a legacy unhashed Python skill manifest", rebuilds: true, prepare: (venv: string, python: string) => { const pythonSkill = createPythonSkill(); writeFakePython(python, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeFileSync( join(venv, ".bootstrap-version"), `${JSON.stringify({ schema: 4, runtime: "prime-agent-runtime", extraUvArgs: DEFAULT_RLM_EXTRA_UV_ARGS, pythonSkills: [ { importName: pythonSkill.importName, packagePath: pythonSkill.packagePath, pyprojectPath: pythonSkill.pyprojectPath, }, ], })}\n`, ); }, }, { name: "rebuilds a venv whose rlm runtime is stale", rebuilds: true, prepare: (venv: string, python: string) => { // Imports rlm but fails the runtime-capability probe. writeExecutable( python, [ "#!/bin/sh", 'if [ "$1" = "-c" ]; then', ' case "$2" in', ' "import rlm") exit 0 ;;', " *) exit 1 ;;", " esac", "fi", "exit 0", "", ].join("\n"), ); writeBootstrapVersion(venv); }, }, { name: "rebuilds a venv whose interpreter is missing entirely", rebuilds: true, prepare: (venv: string) => { writeBootstrapVersion(venv); }, }, ])("$name", async ({ rebuilds, prepare }) => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); mkdirSync(join(venv, "bin"), { recursive: true }); prepare(venv, python); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython()).resolves.toBe(python);
// A warm venv that needs no work never runs uv, so the log may not exist at all. const log = existsSync(logPath) ? readFileSync(logPath, "utf8") : ""; expect(log.includes("--python 3.11 --seed")).toBe(rebuilds); if (rebuilds) { expect(JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")).runtime).toBe(runtimeIdentity); } });
it.each([ { name: "accepts an interpreter carrying the runtime and every default extra", write: (path: string) => writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]), withSkill: false, error: undefined, }, { name: "accepts an interpreter that cannot import the Python skills", write: (path: string) => writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]), withSkill: true, error: undefined, }, { name: "rejects an interpreter missing a default extra package", write: (path: string) => writeFakePython(path, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES.filter((module) => module !== "yaml")]), withSkill: false, error: /default Python packages \(yaml \(PyYAML\)\)/, }, { name: "rejects an interpreter with a stale rlm runtime", write: (path: string) => writeFakePython(path, ["dill"]), withSkill: false, error: /current prime-agent-runtime with callable rlm\.spawn/, }, { name: "rejects an interpreter exposing only the legacy harness API", write: (path: string) => writeExecutable( path, [ "#!/bin/sh", 'if [ "$1" = "-c" ]; then', ' case "$2" in', ' "import rlm") exit 0 ;;', ' *"_harness_methods"*) exit 1 ;;', " *\"assert not hasattr(rlm.rlm, 'background')\"*) exit 0 ;;", " *) exit 1 ;;", " esac", "fi", "exit 0", "", ].join("\n"), ), withSkill: false, error: /current prime-agent-runtime with callable rlm\.spawn/, }, { name: "rejects an interpreter with a pre-progress-note rlm runtime", write: (path: string) => writeExecutable( path, [ "#!/bin/sh", 'if [ "$1" = "-c" ]; then', ' case "$2" in', ' "import rlm") exit 0 ;;', ' *"progress_note"*) exit 1 ;;', ' *"_harness_methods"*) exit 0 ;;', " *) exit 1 ;;", " esac", "fi", "exit 0", "", ].join("\n"), ), withSkill: false, error: /current prime-agent-runtime with callable rlm\.spawn, rlm\.create_session, rlm\.host_request, rlm\.progress_note/, }, { name: "rejects an interpreter missing the runtime, without bootstrapping a venv", write: (path: string) => writeFakePython(path, []), withSkill: false, error: /PRIME_AGENT_KERNEL_PYTHON points to a Python missing/, }, ])("PRIME_AGENT_KERNEL_PYTHON $name", async ({ write, withSkill, error }) => { const overridePython = join(tempDir, "override-python"); write(overridePython); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython; const options = withSkill ? { pythonSkills: [createPythonSkill()] } : {};
if (error) { await expect(ensureKernelPython(options)).rejects.toThrow(error); } else { await expect(ensureKernelPython(options)).resolves.toBe(overridePython); } });
it("resolves the venv python under Scripts\\python.exe on win32 (uv layout)", () => { const venv = join(tempDir, "kernel-venv"); expect(kernelVenvPython(venv, "win32")).toBe(join(venv, "Scripts", "python.exe")); expect(kernelVenvPython(venv, "linux")).toBe(join(venv, "bin", "python")); });
it("postinstall bootstraps the kernel venv with bundled python skills", async () => { installFakeUv(); const venv = join(tempDir, "kernel-venv"); process.env.PRIME_AGENT_KERNEL_VENV = venv; process.env.PRIME_AGENT_BOOTSTRAP_KERNEL_ON_INSTALL = "1"; delete process.env.PRIME_AGENT_BOOTSTRAP_TOOLS_ON_INSTALL; const postPath = resolve(import.meta.dirname, "..", "src", "postinstall.ts"); const tsxCli = resolve(import.meta.dirname, "..", "..", "..", "node_modules", "tsx", "dist", "cli.mjs"); const res = spawnSync(process.execPath, [tsxCli, postPath], { env: process.env, encoding: "utf8" }); expect(res.status, res.stderr).toBe(0); const ver = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(ver.pythonSkills.some((s: { importName: string }) => s.importName === "git")).toBe(true); });});