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 · 276 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277import { spawnSync } from "node:child_process";import { existsSync, mkdtempSync, readFileSync, rmSync } from "node:fs";import { homedir, tmpdir } from "node:os";import { join, resolve } from "node:path";import { afterEach, beforeEach, describe, expect, it } from "vitest";import { AGENT_MESSAGE_DISPLAY_MIME, ATTACHMENT_DISPLAY_MIME, CONTENT_DISPLAY_MIME, DIFF_DISPLAY_MIME, type HostRequestHandlers, ReplKernelManager,} from "../src/core/kernel/index.js";
function resolveReplPython(): 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, dill"], { encoding: "utf8" }); if (check.status === 0) return python; } return null;}
const python = resolveReplPython();const describeIf = python ? describe : describe.skip;
describeIf("ReplKernelManager execute (real runtime)", () => { let dir = ""; let manager: ReplKernelManager | undefined;
beforeEach(() => { dir = mkdtempSync(join(tmpdir(), "prime-agent-repl-execute-")); });
afterEach(async () => { await manager?.shutdown({ snapshot: true, drainHostRequests: true }); manager = undefined; if (dir) { rmSync(dir, { recursive: true, force: true }); dir = ""; } });
it("streams stdout/stderr, returns results, and persists state across cells", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const chunks: { name: string; text: string }[] = []; const first = await manager.execute("import sys\nx = 21\nprint('to-out')\nsys.stderr.write('to-err\\n')", { onStream: (chunk, name) => chunks.push({ name, text: chunk }), }); expect(first.status).toBe("ok"); expect(first.stdout).toContain("to-out"); expect(first.stderr).toContain("to-err"); expect(chunks.some((c) => c.name === "stdout" && c.text.includes("to-out"))).toBe(true);
const second = await manager.execute("x * 2"); expect(second.status).toBe("ok"); expect(second.result).toBe("42"); }, 30_000);
it("returns the generated skill-surface digest the prompt publishes", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); await manager.start();
const digest = await manager.skillDigest(["rlm"]);
// Generated from the live modules, so it names the runtime's own calls and the // settled/pending convention rather than any hand-written copy. expect(digest).toContain("# skill surface"); expect(digest).toContain("bash("); }, 30_000);
it("reports cell errors with a clean traceback", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const r = await manager.execute("def boom():\n raise ValueError('nope')\nboom()"); expect(r.status).toBe("error"); expect(r.error?.ename).toBe("ValueError"); expect(r.error?.evalue).toBe("nope"); expect(r.error?.traceback.join("")).toContain("raise ValueError('nope')"); }, 30_000);
it("parses emitted display payloads into diffs, attachments, and sent messages", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const code = [ "from rlm import emit", `emit({${JSON.stringify(DIFF_DISPLAY_MIME)}: {"path": "/tmp/f.py", "old_str": "a", "new_str": "b", "start_line": 3}})`, `emit({${JSON.stringify(CONTENT_DISPLAY_MIME)}: {"content": "diff --git a/f.py b/f.py", "title": "my reading diff", "format": "diff"}})`, `emit({${JSON.stringify(ATTACHMENT_DISPLAY_MIME)}: {"mime_type": "image/png", "data": "aGVsbG8=", "path": "/tmp/i.png"}})`, `emit({${JSON.stringify(AGENT_MESSAGE_DISPLAY_MIME)}: {"id": "m1", "message": "hi", "deliveryStatus": "delivered", "receiverRole": "parent", "target": {"activeSessionId": "a", "sessionId": "s"}}})`, ].join("\n"); const r = await manager.execute(code); expect(r.status).toBe("ok"); expect(r.diffs).toEqual([{ path: "/tmp/f.py", oldStr: "a", newStr: "b", startLine: 3 }]); expect(r.displays).toEqual([{ content: "diff --git a/f.py b/f.py", title: "my reading diff", format: "diff" }]); expect(r.attachments).toEqual([{ mimeType: "image/png", data: "aGVsbG8=", path: "/tmp/i.png" }]); expect(r.sentAgentMessages).toEqual([ { id: "m1", message: "hi", deliveryStatus: "delivered", receiverRole: "parent", target: { activeSessionId: "a", sessionId: "s" }, }, ]); }, 30_000);
it("round-trips host requests through hostHandlers, including error replies", async () => { const hostHandlers: HostRequestHandlers = { "test.echo": async (payload) => ({ echoed: payload.value, cell: payload.cellSourceCode }), "test.fail": async () => { throw new Error("handler exploded"); }, }; manager = new ReplKernelManager({ python: python as string, cwd: dir, hostHandlers });
const ok = await manager.execute( "import rlm\nreply = await rlm.host_request('test.echo', {'value': 7})\nreply['echoed']", ); expect(ok.status).toBe("ok"); expect(ok.result).toBe("7");
const cellSource = await manager.execute("reply['cell']"); expect(cellSource.status).toBe("ok"); expect(cellSource.result).toContain("test.echo");
const failed = await manager.execute("import rlm\nawait rlm.host_request('test.fail')"); expect(failed.status).toBe("error"); expect(failed.error?.ename).toBe("RuntimeError"); expect(failed.error?.evalue).toBe("handler exploded");
const unknown = await manager.execute("import rlm\nawait rlm.host_request('test.unknown')"); expect(unknown.status).toBe("error"); expect(unknown.error?.evalue).toContain('host request type "test.unknown" is not available'); }, 30_000);
it("spawns through rlm.spawn over the unchanged rlm.run wire type, requires a child name, and refuses a direct rlm call", async () => { const requestTypes: string[] = []; const hostHandlers: HostRequestHandlers = { "rlm.run": async (payload) => { requestTypes.push("rlm.run"); expect(payload.prompt).toMatch(/^You are child subagent 'worker' at depth \d+\.[\s\S]*\nTASK:\nchild work/); expect(payload.kwargs).toEqual({ name: "worker", model: "anthropic/claude-opus-4-7" }); return { rlm_child_id: "sub-a1b2c3d4", name: "worker", session_dir: dir, model: "anthropic/claude-opus-4-7", }; }, }; manager = new ReplKernelManager({ python: python as string, cwd: dir, hostHandlers });
const spawned = await manager.execute( "import rlm\nhandle = await rlm.spawn('child work', name='worker', model='anthropic/claude-opus-4-7')\nhandle.name", ); expect(spawned.status).toBe("ok"); expect(spawned.result).toBe("'worker'"); expect(requestTypes).toEqual(["rlm.run"]);
const nameless = await manager.execute("await rlm.spawn('child work')"); expect(nameless.status).toBe("error"); expect(nameless.error?.ename).toBe("TypeError"); expect(nameless.error?.evalue).toContain("required keyword-only arguments"); expect(nameless.error?.evalue).toContain("'name'"); expect(nameless.error?.evalue).toContain("'model'"); expect(requestTypes).toEqual(["rlm.run"]);
const called = await manager.execute("await rlm('child work')"); expect(called.status).toBe("error"); expect(called.error?.ename).toBe("TypeError"); expect(called.error?.evalue).toContain("await rlm.spawn("); }, 30_000);
it("dispose sends the protocol shutdown so live bash children die with the kernel", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const r = await manager.execute("from rlm import bash\nh = bash('sleep 600')\nh.pid"); expect(r.status).toBe("ok"); const pid = Number(r.result); expect(Number.isInteger(pid)).toBe(true); await manager.shutdown({ snapshot: true, drainHostRequests: true }); let alive = true; for (let i = 0; i < 100; i++) { try { process.kill(pid, 0); } catch { alive = false; break; } await new Promise((resolve) => setTimeout(resolve, 50)); } expect(alive).toBe(false); }, 30_000);
it("stops a cell at its deadline, keeps partial output, and keeps serving cells", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const r = await manager.execute("import time\nprint('before-the-deadline', flush=True)\ntime.sleep(5)", { timeoutMs: 1500, }); expect(r.status).toBe("aborted"); expect(r.timedOut).toBe(true); expect(r.stdout).toContain("before-the-deadline"); expect(r.durationMs).toBeGreaterThanOrEqual(1500);
// The kernel survived the interrupt, and so did its namespace. const again = await manager.execute("x = 41\nx + 1"); expect(again.status).toBe("ok"); expect(again.result).toBe("42"); }, 30_000);
it("ends a cell blocked on a one-shot await bash() at its deadline", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const pidFile = join(dir, "child.pid"); // The audited 61-minute blocking await, now bounded by the cell deadline. const r = await manager.execute( `from rlm import bash\nawait bash(${JSON.stringify(`echo $$ > ${pidFile}; sleep 600`)})\nprint('never reached', flush=True)`, { timeoutMs: 2000 }, ); expect(r.status).toBe("aborted"); expect(r.timedOut).toBe(true); expect(r.stdout).not.toContain("never reached");
// `await bash(cmd)` owns its process: the cancelled await kills the group. const pid = Number(readFileSync(pidFile, "utf8").trim()); expect(Number.isInteger(pid)).toBe(true); let alive = true; for (let i = 0; i < 60; i++) { try { process.kill(pid, 0); } catch { alive = false; break; } await new Promise((resolve) => setTimeout(resolve, 100)); } expect(alive).toBe(false); }, 30_000);
it("runs past any deadline when the caller arms no timer", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); // Deliberately longer than the 1.5s deadline above: without timeoutMs there is no timer. const slow = await manager.execute("import time\ntime.sleep(2.0)\n'finished'"); expect(slow.status).toBe("ok"); expect(slow.result).toBe("'finished'"); expect(slow.timedOut).toBeUndefined();
// State ops keep their own explicit bound instead of inheriting a cell deadline. const internal = await manager.execute("'state-op'", { internal: true }); expect(internal.status).toBe("ok"); expect(internal.result).toBe("'state-op'"); expect(internal.timedOut).toBeUndefined(); }, 30_000);
it("surfaces unattributed background output separately from cell stdout", async () => { manager = new ReplKernelManager({ python: python as string, cwd: dir }); const first = await manager.execute( [ "import threading, time", "def late():", " time.sleep(0.5)", " print('SECRET-thread', flush=True)", "threading.Thread(target=late, daemon=True).start()", ].join("\n"), ); expect(first.status).toBe("ok");
const second = await manager.execute("import time\ntime.sleep(1.0)\nprint('own-output')"); expect(second.status).toBe("ok"); expect(second.stdout).toContain("own-output"); expect(second.stdout).not.toContain("SECRET-thread"); expect(second.backgroundOutput ?? "").toContain("SECRET-thread"); }, 30_000);});