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.
9.2 kB · 261 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262import { EventEmitter } from "node:events";import { describe, expect, it, vi } from "vitest";import { buildRemoteBootstrapScript, buildSshArgs, ensureRemoteKernelRuntime, getRemoteRuntimeDir, probeRemoteKernelRuntime, quotePosixShellArg, type RemoteHostConfig, syncRemotePythonSkills,} from "../src/core/kernel/remote-bootstrap.js";
vi.mock("node:child_process", () => ({ spawn: vi.fn(),}));
import { spawn } from "node:child_process";
const mockedSpawn = vi.mocked(spawn);
function createMockChildProcess(exitCode = 0, stdoutText = "", stderrText = "") { const child = new EventEmitter() as any; child.stdout = new EventEmitter(); child.stderr = new EventEmitter(); child.kill = vi.fn();
process.nextTick(() => { if (stdoutText) child.stdout.emit("data", Buffer.from(stdoutText)); if (stderrText) child.stderr.emit("data", Buffer.from(stderrText)); child.emit("close", exitCode); });
return child;}
describe("remote-bootstrap", () => { describe("buildSshArgs", () => { it("builds ssh command with default ControlMaster persistence", () => { const config: RemoteHostConfig = { host: "example-box" }; expect(buildSshArgs(config)).toEqual([ "-o", "ControlMaster=auto", "-o", "ControlPersist=30m", "-o", "ControlPath=~/.ssh/cm-%C", "example-box", ]); });
it("includes controlPath, port, identityFile, and sshOptions", () => { const config: RemoteHostConfig = { host: "example-box", controlPath: "/tmp/cm.sock", port: 2222, identityFile: "~/.ssh/id_ed25519", sshOptions: ["-o", "StrictHostKeyChecking=no"], }; expect(buildSshArgs(config, "uname -a")).toEqual([ "-S", "/tmp/cm.sock", "-p", "2222", "-i", "~/.ssh/id_ed25519", "-o", "StrictHostKeyChecking=no", "example-box", "uname -a", ]); }); });
describe("quotePosixShellArg", () => { it("quotes hostile POSIX path characters as one literal argument", () => { const path = "/remote dir/it's $cache\\runtime"; expect(quotePosixShellArg(path)).toBe("'/remote dir/it'\\''s $cache\\runtime'"); });
it("preserves the default HOME expression without exposing the path suffix", () => { expect(quotePosixShellArg("$HOME/cache dir", { expandHome: true })).toBe("$HOME'/cache dir'"); }); });
describe("getRemoteRuntimeDir", () => { it("keys path by content hash under default cache dir", () => { const dir = getRemoteRuntimeDir(undefined, "abc12345"); expect(dir).toBe("$HOME/.cache/prime-agent/runtimes/abc12345"); });
it("respects custom cache dir", () => { const dir = getRemoteRuntimeDir("/opt/cache", "abc12345"); expect(dir).toBe("/opt/cache/abc12345"); }); });
describe("buildRemoteBootstrapScript", () => { it("contains uv, nix fallback, python3 -m venv fallback, and touches .ready", () => { const script = buildRemoteBootstrapScript("/remote/path/hash123"); expect(script).toContain("TARGET='/remote/path/hash123'"); expect(script).toContain("uv venv .venv"); expect(script).toContain("nix shell nixpkgs#uv -c uv venv .venv"); expect(script).toContain("python3 -m venv .venv"); expect(script).toContain('touch "$TARGET/.ready"'); });
it("quotes every caller-provided package argument", () => { const script = buildRemoteBootstrapScript("/remote/path/hash123", ["safe-package", "package name;$(boom)"]); expect(script).toContain("'safe-package'"); expect(script).toContain("'package name;$(boom)'"); }); });
describe("probeRemoteKernelRuntime", () => { it("returns true when .ready probe exits 0", async () => { mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0)); const isReady = await probeRemoteKernelRuntime({ host: "remote-host" }, "/remote/dir"); expect(isReady).toBe(true); expect(mockedSpawn).toHaveBeenCalledWith( "ssh", [ "-o", "ControlMaster=auto", "-o", "ControlPersist=30m", "-o", "ControlPath=~/.ssh/cm-%C", "remote-host", "test -f '/remote/dir/.ready'", ], expect.anything(), ); });
it("quotes hostile runtime paths in the ready probe", async () => { mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0)); await expect( probeRemoteKernelRuntime({ host: "remote-host" }, "/remote dir/it's $cache\\runtime"), ).resolves.toBe(true); expect(mockedSpawn).toHaveBeenLastCalledWith( "ssh", expect.arrayContaining(["test -f '/remote dir/it'\\''s $cache\\runtime/.ready'"]), expect.anything(), ); });
it("returns false when .ready probe exits non-zero", async () => { mockedSpawn.mockImplementationOnce(() => createMockChildProcess(1)); const isReady = await probeRemoteKernelRuntime({ host: "remote-host" }, "/remote/dir"); expect(isReady).toBe(false); }); });
describe("ensureRemoteKernelRuntime", () => { it("returns immediately on fast-path cache hit without uploading or building", async () => { mockedSpawn.mockClear(); // Fast path probe succeeds (exit code 0) mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0));
const progress: string[] = []; const result = await ensureRemoteKernelRuntime({ hostConfig: { host: "remote-host" }, onProgress: (msg) => progress.push(msg), });
expect(result.cacheHit).toBe(true); expect(result.remotePython).toContain("/.venv/bin/python"); expect(result.remoteDir).toContain("$HOME/.cache/prime-agent/runtimes/"); expect(mockedSpawn).toHaveBeenCalledTimes(1); // Only the probe! expect(progress).toHaveLength(0); });
it("respects custom pythonCmd without probing remote host", async () => { mockedSpawn.mockClear(); const result = await ensureRemoteKernelRuntime({ hostConfig: { host: "remote-host", pythonCmd: "/custom/python" }, }); expect(result.cacheHit).toBe(true); expect(result.remotePython).toBe("/custom/python"); expect(mockedSpawn).toHaveBeenCalledTimes(0); });
it("respects PRIME_AGENT_REMOTE_PYTHON environment variable override", async () => { mockedSpawn.mockClear(); const prev = process.env.PRIME_AGENT_REMOTE_PYTHON; try { process.env.PRIME_AGENT_REMOTE_PYTHON = "/nix/profile/bin/python"; const result = await ensureRemoteKernelRuntime({ hostConfig: { host: "remote-host" }, }); expect(result.cacheHit).toBe(true); expect(result.remotePython).toBe("/nix/profile/bin/python"); expect(mockedSpawn).toHaveBeenCalledTimes(0); } finally { if (prev === undefined) delete process.env.PRIME_AGENT_REMOTE_PYTHON; else process.env.PRIME_AGENT_REMOTE_PYTHON = prev; } });
it("transfers source and runs bootstrap on cache miss", async () => { mockedSpawn.mockClear(); // 1. Probe fails (exit code 1) mockedSpawn.mockImplementationOnce(() => createMockChildProcess(1)); // 2. tar process (exit code 0) mockedSpawn.mockImplementationOnce(() => { const p = createMockChildProcess(0); p.stdout.pipe = vi.fn(); return p; }); // 3. ssh extract process (exit code 0) mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0)); // 4. ssh bootstrap process (exit code 0) mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0));
const progress: string[] = []; const result = await ensureRemoteKernelRuntime({ hostConfig: { host: "remote-host" }, onProgress: (msg) => progress.push(msg), });
expect(result.cacheHit).toBe(false); expect(result.remotePython).toContain("/.venv/bin/python"); expect(progress).toContainEqual(expect.stringContaining("Transferring prime-agent-runtime")); expect(progress).toContainEqual(expect.stringContaining("Bootstrapping remote kernel environment")); expect(progress).toContainEqual(expect.stringContaining("Remote kernel runtime ready")); }); });
describe("syncRemotePythonSkills", () => { it("surfaces a remote editable install failure", async () => { const skill = { name: "hostile skill", importName: "skill$module", packagePath: "/local skill/it's $source\\package", pyprojectPath: "/local skill/pyproject.toml", }; mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0, "0\n")); mockedSpawn.mockImplementationOnce(() => { const child = createMockChildProcess(0); child.stdout.pipe = vi.fn(); return child; }); mockedSpawn.mockImplementationOnce(() => createMockChildProcess(0)); mockedSpawn.mockImplementationOnce(() => createMockChildProcess(7, "", "install failed")); const firstCall = mockedSpawn.mock.calls.length;
await expect( syncRemotePythonSkills({ host: "remote-host" }, "/remote dir/it's $cache\\runtime", [skill]), ).rejects.toThrow("Failed to install Python skill hostile skill on remote-host"); const remoteCommands = mockedSpawn.mock.calls .slice(firstCall) .filter(([command]) => command === "ssh") .map(([, args]) => (args as string[]).at(-1)); expect(remoteCommands[0]).toContain( "test -f '/remote dir/it'\\''s $cache\\runtime/skills/skill$module/.prime-agent-installed'", ); expect(remoteCommands.at(-1)).toContain('touch "$SKILL/.prime-agent-installed"'); }); });});