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.
5.4 kB · 141 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142import { 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(); expect(code).toContain('async def find_models(self, query="", limit=8)'); expect(code).toContain("async def list_subagents(self)"); expect(code).toContain("async def delete_subagent(self, target)"); expect(code).toContain("self._raise_missing()"); });
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 missing-runtime stub fallback", () => { const code = buildRlmBootstrapCode(); expect(code).toContain("bash = _prime_agent_rlm_module.bash"); expect(code).toContain("def bash(command):"); expect(code).toContain("rlm._raise_missing()"); });
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"); });});
/** 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);});