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.
20 kB · 561 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562import { createHash } from "node:crypto";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, vi } from "vitest";import { DEFAULT_RLM_EXTRA_IMPORT_NAMES, DEFAULT_RLM_EXTRA_UV_ARGS, ensureKernelPython, getKernelVenvDir, 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", '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 rlm") exit 0 ;;', ...extraImportCases, ' *"_harness_methods"*) 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("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, prime-agent-runtime, and default extra packages", 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).not.toContain("ipykernel"); expect(log).toContain("prime-agent-runtime"); expect(log).toContain("dill"); for (const uvArg of DEFAULT_RLM_EXTRA_UV_ARGS) { expect(log).toContain(uvArg); } 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: [], }); expect(version.runtime).toMatch(/^sha256:/); });
it("routes bootstrap progress through the provided callback", async () => { installFakeUv(); const venv = join(tempDir, "kernel-venv"); const progress: string[] = []; process.env.PRIME_AGENT_KERNEL_VENV = venv; const stderrWrite = vi.spyOn(process.stderr, "write").mockImplementation(() => true);
try { await expect(ensureKernelPython({ onProgress: (message) => progress.push(message) })).resolves.toBe( join(venv, "bin", "python"), ); } finally { stderrWrite.mockRestore(); }
expect(progress).toEqual(expect.arrayContaining(["› setting up python kernel (one-time, ~30s)…", "✓ ready"])); expect(stderrWrite).not.toHaveBeenCalledWith(expect.stringContaining("setting up python kernel")); expect(stderrWrite).not.toHaveBeenCalledWith(expect.stringContaining("ready")); });
it("installs Python skills into the bootstrapped venv", 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(`--editable ${pythonSkill.packagePath}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills).toEqual([ { importName: pythonSkill.importName, packagePath: pythonSkill.packagePath, pyprojectPath: pythonSkill.pyprojectPath, pyprojectHash: pyprojectHash(pythonSkill.pyprojectPath), }, ]); });
it("installs sibling Python skill dependencies with dependent editable packages", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const dependencySkill = createPythonSkill("agent-observe"); const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", "agent-observe"); 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}`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.pythonSkills).toEqual([ { importName: dependencySkill.importName, packagePath: dependencySkill.packagePath, pyprojectPath: dependencySkill.pyprojectPath, pyprojectHash: pyprojectHash(dependencySkill.pyprojectPath), }, { importName: dependentSkill.importName, packagePath: dependentSkill.packagePath, pyprojectPath: dependentSkill.pyprojectPath, pyprojectHash: pyprojectHash(dependentSkill.pyprojectPath), }, ]); });
it("installs sibling Python skill dependencies when package and directory names differ", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const dependencySkill = createPythonSkill("attach-image"); writeFileSync( dependencySkill.pyprojectPath, `[project]name = "prime-agent-skill-attach-image"version = "0.1.0"`, ); const dependentSkill = createPythonSkillWithDependency( "orchestration-heartbeat", "prime-agent-skill-attach-image", ); 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("parses Python skill dependencies with extras", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const dependencySkill = createPythonSkill("gidgethub"); const dependentSkill = createPythonSkillWithDependency("orchestration-heartbeat", "gidgethub[httpx]>4.0.0"); 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(`venv ${venv} --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 ${venv} `))).toHaveLength(1); expect( retryLog.split("\n").filter((line) => line.includes(`--editable ${brokenSkill.packagePath}`)), ).toHaveLength(2); });
it("rebuilds a warm venv with legacy unhashed Python skill manifest entries", 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]); 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`, ); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython()).resolves.toBe(python);
expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); });
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, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); writeBootstrapVersion(venv); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython()).resolves.toBe(python); });
it("rebuilds a warm venv whose recorded runtime hash no longer matches local source", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); mkdirSync(join(venv, "bin"), { recursive: true }); 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`, ); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython()).resolves.toBe(python);
expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); const version = JSON.parse(readFileSync(join(venv, ".bootstrap-version"), "utf8")); expect(version.runtime).toBe(runtimeIdentity); });
it("rebuilds a warm venv with a stale rlm runtime", async () => { const logPath = installFakeUv(); const venv = join(tempDir, "kernel-venv"); const python = join(venv, "bin", "python"); mkdirSync(join(venv, "bin"), { recursive: true }); 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); process.env.PRIME_AGENT_KERNEL_VENV = venv;
await expect(ensureKernelPython()).resolves.toBe(python);
expect(readFileSync(logPath, "utf8")).toContain(`venv ${venv} --python 3.11 --seed`); });
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, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython;
await expect(ensureKernelPython()).resolves.toBe(overridePython); });
it("allows PRIME_AGENT_KERNEL_PYTHON missing Python skill imports", async () => { const overridePython = join(tempDir, "override-python"); const pythonSkill = createPythonSkill(); writeFakePython(overridePython, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES]); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython;
await expect(ensureKernelPython({ pythonSkills: [pythonSkill] })).resolves.toBe(overridePython); });
it("rejects PRIME_AGENT_KERNEL_PYTHON missing default extra packages", async () => { const overridePython = join(tempDir, "override-python"); writeFakePython(overridePython, ["rlm", ...DEFAULT_RLM_EXTRA_IMPORT_NAMES.filter((name) => name !== "yaml")]); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython;
await expect(ensureKernelPython()).rejects.toThrow(/default Python packages \(yaml \(PyYAML\)\)/); });
it("rejects PRIME_AGENT_KERNEL_PYTHON with a stale rlm runtime", async () => { const overridePython = join(tempDir, "override-python"); writeFakePython(overridePython, ["dill"]); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython;
await expect(ensureKernelPython()).rejects.toThrow(/current prime-agent-runtime with callable rlm\.run/); });
it("rejects PRIME_AGENT_KERNEL_PYTHON with a legacy harness API", async () => { const overridePython = join(tempDir, "override-python"); writeExecutable( overridePython, [ "#!/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"), ); process.env.PRIME_AGENT_KERNEL_PYTHON = overridePython;
await expect(ensureKernelPython()).rejects.toThrow(/current prime-agent-runtime with callable rlm\.run/); });
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(/PRIME_AGENT_KERNEL_PYTHON points to a Python missing/); });
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")); });});