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.
11 kB · 243 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244import { spawn, spawnSync } from "node:child_process";import { existsSync, mkdirSync, mkdtempSync, realpathSync, rmSync, writeFileSync } from "node:fs";import { homedir, tmpdir } from "node:os";import { join, resolve } from "node:path";import { afterAll, describe, expect, it } from "vitest";import { ReplKernelManager } from "../src/core/kernel/index.js";import { buildRlmBootstrapCode } from "../src/core/tools/ipython.js";
describe("RLM bootstrap", () => { it("pre-imports asyncio so the prompt's subagent patterns work without a manual import", () => { expect(buildRlmBootstrapCode()).toMatch(/^import asyncio$/m); });
it("gives subagent registry operations the actionable missing-runtime fallback", () => { const code = buildRlmBootstrapCode(); // Role routing was replaced by `rlm.models()` discovery, so the stub set is the spawn // and registry surface the prompt actually calls. expect(code).toContain("async def spawn(self, prompt, **kwargs)"); expect(code).toContain("async def models(self)"); expect(code).toContain("async def list_subagents(self)"); expect(code).toContain("async def delete_subagent(self, target)"); expect(code).toContain("self._rebind()"); });
it("disables colored output for subprocesses launched by the kernel", () => { expect(buildRlmBootstrapCode()).toContain('_prime_agent_os.environ["NO_COLOR"] = "1"'); });
it("binds bash from the runtime with a retrying missing-runtime stub fallback", () => { const code = buildRlmBootstrapCode(); expect(code).toContain("bash=_prime_agent_rlm_module.bash"); // The stub accepts the real call shape, so a cell that passes a timeout gets the // runtime's own signature check rather than the stub's arity. expect(code).toContain("def bash(command, *args, **kwargs):"); expect(code).toContain('_prime_agent_require_runtime()["bash"](command, *args, **kwargs)'); });
it("does not install a kernel shutdown hook (the runtime closes MCP on shutdown itself)", () => { expect(buildRlmBootstrapCode()).not.toContain("install_shutdown_hook"); });
it("guards Python skill imports so a broken skill does not abort bootstrap", () => { const code = buildRlmBootstrapCode([ { name: "broken-skill", importName: "broken_skill", packagePath: "/tmp/broken-skill", pyprojectPath: "/tmp/broken-skill/pyproject.toml", }, ]);
expect(code).toContain("except Exception as _prime_agent_skill_error"); expect(code).toContain("_PrimeAgentUnavailableSkill"); expect(code).toContain("_PRIME_AGENT_SKILL_IMPORT_ERRORS"); expect(code).toContain("globals()[_prime_agent_skill_name] = _PrimeAgentUnavailableSkill"); });});
const FAKE_RUNTIME_SOURCE = `rlm = object()
def bash(command, *args, **kwargs): return f"real:{command}:{sorted(kwargs)}"`;
describe("RLM bootstrap recovery", () => { const python = process.env.PRIME_AGENT_KERNEL_PYTHON ?? "python3"; const dir = mkdtempSync(join(tmpdir(), "prime-agent-rlm-recovery-"));
afterAll(() => rmSync(dir, { recursive: true, force: true }));
it("waits out a build lock, then rebinds the runtime on the next call", async () => { const site = join(dir, "site"); const venv = join(dir, "kernel-venv"); const lock = `${venv}.bootstrap.lock`; const missing = join(site, "rlm", "mcp.py"); mkdirSync(join(site, "rlm"), { recursive: true }); writeFileSync(join(site, "rlm", "__init__.py"), FAKE_RUNTIME_SOURCE); writeFileSync(join(site, "remote.py"), "MARK = 'remote'\n"); writeFileSync(lock, `${process.pid}\n`); const script = [ buildRlmBootstrapCode([{ name: "flaky", importName: "flaky", packagePath: site, pyprojectPath: lock }]), 'print("lock_held_after_wait:", _prime_agent_build_in_progress(), flush=True)', 'print("ready:", _PRIME_AGENT_RUNTIME_READY, flush=True)', "try:", ' bash("too-early")', "except RuntimeError as error:", ' print("terminal:", "prime-agent-runtime is not installed" in str(error), flush=True)', 'print("awaiting-go", flush=True)', "input()", 'print("bash:", bash("echo", timeout=5), flush=True)', 'print("mcp:", mcp.MARK, remote.MARK, flush=True)', 'print("skill:", asyncio.run(flaky.run()), flush=True)', ].join("\n"); const child = spawn(python, ["-c", script], { env: { ...process.env, PYTHONPATH: site, PRIME_AGENT_KERNEL_VENV: venv, PRIME_AGENT_KERNEL_BOOTSTRAP_WAIT_SECONDS: "10", }, }); let stdout = ""; child.stdout.setEncoding("utf8"); child.stdout.on("data", (chunk: string) => { stdout += chunk; // Act only on the kernel's own progress: release the lock it is waiting on, // then publish the module whose absence failed its first import. if (stdout.includes("waiting up to")) rmSync(lock, { force: true }); if (stdout.includes("awaiting-go")) { writeFileSync(missing, "MARK = 'mcp'\n"); // The skill's editable install landed with the same rebuild. writeFileSync(join(site, "flaky.py"), "async def run():\n return 'skill-ok'\n"); child.stdin.write("go\n"); } }); const exitCode = await new Promise<number | null>((settle) => child.on("close", settle));
expect(exitCode).toBe(0); expect(stdout).toContain("lock_held_after_wait: False"); expect(stdout).toContain("waiting up to 10s for the kernel venv build"); expect(stdout).toContain("ready: False"); expect(stdout).toContain("terminal: True"); expect(stdout).toContain("bash: real:echo:['timeout']"); expect(stdout).toContain("mcp: mcp remote"); expect(stdout).toContain("skill: skill-ok"); });
it("treats a lock with an alive builder pid as in-progress beyond stale age and clears dead locks", () => { const venv = join(dir, "pid-lock-venv"); const lock = `${venv}.bootstrap.lock`; const script = `${buildRlmBootstrapCode()}\nimport os, time\nopen('${lock}', 'w').write(f'{os.getpid()}:tok\\n')\nos.utime('${lock}', (time.time() - 60, time.time() - 60))\nassert _prime_agent_build_in_progress() is True\nopen('${lock}', 'w').write('999999999:dead\\n')\nassert _prime_agent_build_in_progress() is False and not os.path.exists('${lock}')`; expect(spawnSync(python, ["-c", script], { env: { ...process.env, PRIME_AGENT_KERNEL_VENV: venv } }).status).toBe( 0, ); });
it("isolates optional remote import failure so rlm, bash, and mcp remain available", () => { const site = join(dir, "isolated-site"); mkdirSync(join(site, "rlm"), { recursive: true }); writeFileSync(join(site, "rlm", "__init__.py"), FAKE_RUNTIME_SOURCE); writeFileSync(join(site, "rlm", "mcp.py"), "MARK = 'mcp'\n"); const script = `${buildRlmBootstrapCode()}\nassert _PRIME_AGENT_RUNTIME_READY and callable(bash) and hasattr(mcp, 'MARK') and isinstance(globals().get('remote'), _PrimeAgentUnavailableSkill)`; expect(spawnSync(python, ["-c", script], { env: { ...process.env, PYTHONPATH: site } }).status).toBe(0); });
it("placeholder objects self-heal on attribute access and report teaching error on failure", () => { const venv = join(dir, "placeholder-venv"); const script = `${buildRlmBootstrapCode([{ name: "git", importName: "git", packagePath: "/tmp/git", pyprojectPath: "/tmp/git/pyproject.toml" }])}\ntry: rlm.progress_note('t')\nexcept RuntimeError as e: assert 'kernel-venv' in str(e) or '${venv}' in str(e)\ntry: git.status()\nexcept RuntimeError as e: assert 'is unavailable in this kernel' in str(e) and '${venv}' in str(e)`; expect(spawnSync(python, ["-c", script], { env: { ...process.env, PRIME_AGENT_KERNEL_VENV: venv } }).status).toBe( 0, ); });});
/** Find a python with a current rlm runtime, or null to skip. */function resolveKernelPython(): string | null { const candidates = [ process.env.PRIME_AGENT_KERNEL_PYTHON, resolve(__dirname, "..", "..", "..", "prime-agent-runtime", ".venv", "bin", "python"), join(homedir(), ".prime", "agent", "kernel-venv", "bin", "python"), ].filter((p): p is string => Boolean(p)); for (const python of candidates) { if (!existsSync(python)) continue; const check = spawnSync(python, ["-c", "import rlm.repl, rlm; assert callable(rlm.emit)"], { encoding: "utf8", }); if (check.status === 0) return python; } return null;}
const python = resolveKernelPython();const describeIfKernel = python ? describe : describe.skip;
describeIfKernel("RLM bootstrap (real kernel)", () => { const dir = mkdtempSync(join(tmpdir(), "prime-agent-bootstrap-"));
afterAll(() => { rmSync(dir, { recursive: true, force: true }); });
it("binds asyncio in the user namespace", async () => { const manager = new ReplKernelManager({ python: python as string, cwd: dir }); try { await manager.start(); const bootstrap = await manager.execute(buildRlmBootstrapCode()); expect(bootstrap.status).toBe("ok");
const result = await manager.execute("_t = asyncio.create_task(asyncio.sleep(0))\nprint(type(_t).__name__)"); expect(result.status).toBe("ok"); expect(result.stdout).toContain("Task");
const envResult = await manager.execute('import os\nprint(os.environ["NO_COLOR"])'); expect(envResult.status).toBe("ok"); expect(envResult.stdout.trim()).toBe("1"); } finally { await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }, 60_000);
it("emits canonical paths for edits after the kernel changes directories", async () => { const firstDir = join(dir, "first"); const secondDir = join(dir, "second"); mkdirSync(firstDir, { recursive: true }); mkdirSync(secondDir, { recursive: true }); writeFileSync(join(firstDir, "same.txt"), "old"); writeFileSync(join(secondDir, "same.txt"), "old"); const editSkillRoot = join(process.cwd(), "skills", "edit"); const manager = new ReplKernelManager({ python: python as string, cwd: dir, env: { PYTHONPATH: join(editSkillRoot, "src") }, }); try { await manager.start(); const bootstrap = await manager.execute( buildRlmBootstrapCode([ { name: "edit", importName: "edit", packagePath: editSkillRoot, pyprojectPath: join(editSkillRoot, "pyproject.toml"), }, ]), ); expect(bootstrap.status).toBe("ok");
const first = await manager.execute( 'import os\nos.chdir("first")\nawait edit(path="same.txt", old_str="old", new_str="new")', ); const second = await manager.execute( 'os.chdir("../second")\nawait edit(path="same.txt", old_str="old", new_str="new")', );
expect(first.diffs?.[0]?.path).toBe(realpathSync(join(firstDir, "same.txt"))); expect(second.diffs?.[0]?.path).toBe(realpathSync(join(secondDir, "same.txt"))); expect(first.diffs?.[0]?.path).not.toBe(second.diffs?.[0]?.path); } finally { await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }, 60_000);});