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.
12 kB · 397 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398import { spawn } from "node:child_process";import { DEFAULT_RLM_EXTRA_UV_ARGS, type KernelPythonSkill, resolveRuntimeIdentity, resolveRuntimeSourceDir,} from "./bootstrap.js";import { errorMessage } from "./shared.js";
const DEFAULT_REMOTE_CACHE_BASE = "$HOME/.cache/prime-agent/runtimes";const STATE_SNAPSHOT_REQUIREMENT = "dill";
/** Quote one literal POSIX shell argument, optionally preserving the default $HOME path expression. */export function quotePosixShellArg(value: string, options: { expandHome?: boolean } = {}): string { if (value.includes("\0")) { throw new Error("POSIX shell arguments cannot contain NUL bytes"); } if (options.expandHome && value.startsWith("$HOME")) { return `$HOME${quotePosixShellArg(value.slice("$HOME".length))}`; } return `'${value.replaceAll("'", "'\\''")}'`;}
export interface RemoteHostConfig { /** Remote SSH destination (e.g. "chernobog" or "mayer@chernobog") */ host: string; /** Path to an existing OpenSSH ControlMaster socket if any */ controlPath?: string; /** Optional custom SSH port */ port?: number; /** Optional identity file */ identityFile?: string; /** Extra options to pass to SSH */ sshOptions?: string[]; /** Base directory for runtimes on the remote host. Defaults to "$HOME/.cache/prime-agent/runtimes" */ cacheDir?: string; /** Optional explicit remote Python binary override */ pythonCmd?: string;}
export interface RemoteKernelRuntimeInfo { /** The content hash / identity of prime-agent-runtime */ runtimeHash: string; /** Remote directory containing the cached runtime ($HOME/.cache/prime-agent/runtimes/<hash>) */ remoteDir: string; /** Path to the python interpreter on the remote host ($HOME/.cache/prime-agent/runtimes/<hash>/.venv/bin/python) */ remotePython: string; /** Whether this execution hit the cache (.ready was already present on the remote host) */ cacheHit: boolean;}
export interface EnsureRemoteKernelRuntimeOptions { hostConfig: RemoteHostConfig; pythonSkills?: readonly KernelPythonSkill[]; onProgress?: (message: string) => void; signal?: AbortSignal;}
export interface RemoteExecResult { exitCode: number; stdout: string; stderr: string;}
/** * Builds the arguments array for the `ssh` command. */export function buildSshArgs(hostConfig: RemoteHostConfig, remoteCommand?: string): string[] { const args: string[] = []; if (hostConfig.controlPath) { args.push("-S", hostConfig.controlPath); } else if (!hostConfig.sshOptions?.some((opt) => opt.includes("ControlMaster"))) { args.push("-o", "ControlMaster=auto", "-o", "ControlPersist=30m", "-o", "ControlPath=~/.ssh/cm-%C"); } if (hostConfig.port !== undefined) { args.push("-p", String(hostConfig.port)); } if (hostConfig.identityFile) { args.push("-i", hostConfig.identityFile); } if (hostConfig.sshOptions) { args.push(...hostConfig.sshOptions); } args.push(hostConfig.host); if (remoteCommand !== undefined) { args.push(remoteCommand); } return args;}
/** * Computes the remote directory path keyed by content hash. */export function getRemoteRuntimeDir(cacheBase: string | undefined, runtimeHash: string): string { const base = cacheBase ?? DEFAULT_REMOTE_CACHE_BASE; return `${base}/${runtimeHash}`;}
/** * Executes a remote command via SSH and returns its exit code and output. */export async function execRemote( hostConfig: RemoteHostConfig, remoteCommand: string, options: { signal?: AbortSignal } = {},): Promise<RemoteExecResult> { const args = buildSshArgs(hostConfig, remoteCommand); return new Promise<RemoteExecResult>((resolve, reject) => { const child = spawn("ssh", args, { stdio: ["ignore", "pipe", "pipe"], signal: options.signal, });
let stdout = ""; let stderr = "";
child.stdout?.on("data", (chunk: Buffer) => { stdout += chunk.toString("utf8"); }); child.stderr?.on("data", (chunk: Buffer) => { stderr += chunk.toString("utf8"); });
child.on("error", (err) => { reject(new Error(`Failed to execute SSH command on ${hostConfig.host}: ${errorMessage(err)}`)); });
child.on("close", (code) => { resolve({ exitCode: code ?? 1, stdout, stderr, }); }); });}
/** * Fast path: probes whether the content-addressed runtime already exists and is ready on the remote host. */export async function probeRemoteKernelRuntime( hostConfig: RemoteHostConfig, remoteDir: string, signal?: AbortSignal, expandHome = false,): Promise<boolean> { const readyFile = `${remoteDir}/.ready`; const cmd = `test -f ${quotePosixShellArg(readyFile, { expandHome })}`; const result = await execRemote(hostConfig, cmd, { signal }); return result.exitCode === 0;}
/** * Streams a gzipped tarball of the local runtime source directory directly over SSH to the remote host. */export async function transferRuntimeSource( hostConfig: RemoteHostConfig, sourceDir: string, remoteSrcDir: string, signal?: AbortSignal, expandHome = false,): Promise<void> { const quotedRemoteSrcDir = quotePosixShellArg(remoteSrcDir, { expandHome }); const sshExtractCmd = `mkdir -p ${quotedRemoteSrcDir} && tar -xzf - -C ${quotedRemoteSrcDir}`; const sshArgs = buildSshArgs(hostConfig, sshExtractCmd);
const tarArgs = [ "-czf", "-", "--exclude=.venv", "--exclude=__pycache__", "--exclude=.git", "--exclude=*.pyc", "--exclude=test", "-C", sourceDir, ".", ];
return new Promise<void>((resolve, reject) => { const tarChild = spawn("tar", tarArgs, { stdio: ["ignore", "pipe", "pipe"], signal, });
const sshChild = spawn("ssh", sshArgs, { stdio: ["pipe", "pipe", "pipe"], signal, });
let sshStderr = ""; let tarStderr = "";
sshChild.stderr?.on("data", (chunk: Buffer) => { sshStderr += chunk.toString("utf8"); }); tarChild.stderr?.on("data", (chunk: Buffer) => { tarStderr += chunk.toString("utf8"); });
tarChild.on("error", (err) => { sshChild.kill(); reject(new Error(`tar packaging failed: ${errorMessage(err)}`)); });
sshChild.on("error", (err) => { tarChild.kill(); reject(new Error(`SSH transfer failed to ${hostConfig.host}: ${errorMessage(err)}`)); });
// Pipe tar stdout to ssh stdin tarChild.stdout?.pipe(sshChild.stdin);
let tarClosed = false; let sshClosed = false; let tarExitCode: number | null = null; let sshExitCode: number | null = null;
const checkDone = () => { if (!tarClosed || !sshClosed) return; if (tarExitCode !== 0) { reject(new Error(`tar command exited with code ${tarExitCode}: ${tarStderr.trim()}`)); } else if (sshExitCode !== 0) { reject(new Error(`Remote extraction exited with code ${sshExitCode}: ${sshStderr.trim()}`)); } else { resolve(); } };
tarChild.on("close", (code) => { tarClosed = true; tarExitCode = code; checkDone(); });
sshChild.on("close", (code) => { sshClosed = true; sshExitCode = code; checkDone(); }); });}
/** * Builds the shell script to bootstrap a virtualenv and install dependencies on the remote host. */export function buildRemoteBootstrapScript( remoteDir: string, extraUvArgs: readonly string[] = DEFAULT_RLM_EXTRA_UV_ARGS, expandHome = false,): string { const packages = ["-e", "src", STATE_SNAPSHOT_REQUIREMENT, ...extraUvArgs] .map((arg) => quotePosixShellArg(arg)) .join(" "); const quotedRemoteDir = quotePosixShellArg(remoteDir, { expandHome }); return `set -euTARGET=${quotedRemoteDir}cd "$TARGET"if command -v uv >/dev/null 2>&1; then uv venv .venv uv pip install ${packages}elif command -v nix >/dev/null 2>&1; then nix shell nixpkgs#uv -c uv venv .venv nix shell nixpkgs#uv -c uv pip install ${packages}else python3 -m venv .venv .venv/bin/pip install ${packages}fitouch "$TARGET/.ready"`;}
/** * Ensures that a matching content-addressed prime-agent-runtime exists on the remote host. * * Fast path: probes for `$HOME/.cache/prime-agent/runtimes/<hash>/.ready`. If found, returns immediately. * Cold path: streams the runtime source tarball, creates a virtualenv, installs runtime dependencies, * and marks it ready. */export async function ensureRemoteKernelRuntime( options: EnsureRemoteKernelRuntimeOptions,): Promise<RemoteKernelRuntimeInfo> { const { hostConfig, onProgress, signal } = options;
const remotePythonOverride = hostConfig.pythonCmd ?? process.env.PRIME_AGENT_REMOTE_PYTHON; if (remotePythonOverride) { return { runtimeHash: "custom", remoteDir: "", remotePython: remotePythonOverride, cacheHit: true, }; }
const runtimeHash = await resolveRuntimeIdentity(); const remoteDir = getRemoteRuntimeDir(hostConfig.cacheDir, runtimeHash); const remotePython = `${remoteDir}/.venv/bin/python`;
// Fast path: probe for .ready file const expandHome = hostConfig.cacheDir === undefined; const isReady = await probeRemoteKernelRuntime(hostConfig, remoteDir, signal, expandHome); if (isReady) { if (options.pythonSkills && options.pythonSkills.length > 0) { await syncRemotePythonSkills(hostConfig, remoteDir, options.pythonSkills, signal, onProgress, expandHome); } return { runtimeHash, remoteDir, remotePython, cacheHit: true, }; }
// Cold path: resolve local source directory const sourceDir = await resolveRuntimeSourceDir(); if (!sourceDir) { throw new Error("Cannot find local prime-agent-runtime source directory to sync to remote host"); }
onProgress?.(`Transferring prime-agent-runtime (${runtimeHash.slice(0, 12)}) to ${hostConfig.host}...`); const remoteSrcDir = `${remoteDir}/src`; await transferRuntimeSource(hostConfig, sourceDir, remoteSrcDir, signal, expandHome);
onProgress?.(`Bootstrapping remote kernel environment on ${hostConfig.host}...`); const bootstrapScript = buildRemoteBootstrapScript(remoteDir, DEFAULT_RLM_EXTRA_UV_ARGS, expandHome); const bootstrapResult = await execRemote(hostConfig, bootstrapScript, { signal }); if (bootstrapResult.exitCode !== 0) { throw new Error( `Failed to bootstrap remote kernel on ${hostConfig.host} (exit ${bootstrapResult.exitCode}): ` + `${bootstrapResult.stderr || bootstrapResult.stdout}`, ); }
if (options.pythonSkills && options.pythonSkills.length > 0) { await syncRemotePythonSkills(hostConfig, remoteDir, options.pythonSkills, signal, onProgress, expandHome); }
onProgress?.(`Remote kernel runtime ready on ${hostConfig.host}`); return { runtimeHash, remoteDir, remotePython, cacheHit: false, };}
/** * Syncs user/workspace Python skills to the remote host and installs them in the remote kernel environment. */export async function syncRemotePythonSkills( hostConfig: RemoteHostConfig, remoteDir: string, pythonSkills: readonly KernelPythonSkill[], signal?: AbortSignal, onProgress?: (message: string) => void, expandHome = false,): Promise<void> { if (!pythonSkills || pythonSkills.length === 0) return;
for (const skill of pythonSkills) { const remoteSkillDir = `${remoteDir}/skills/${skill.importName}`; const quotedRemoteSkillDir = quotePosixShellArg(remoteSkillDir, { expandHome }); const readyFile = quotePosixShellArg(`${remoteSkillDir}/.prime-agent-installed`, { expandHome }); const probeCmd = `test -f ${readyFile} && echo 1 || echo 0`; const probe = await execRemote(hostConfig, probeCmd, { signal }); if (probe.exitCode === 0 && probe.stdout.trim() === "1") { continue; }
onProgress?.(`Syncing Python skill ${skill.name || skill.importName} to ${hostConfig.host}...`); try { await transferRuntimeSource(hostConfig, skill.packagePath, remoteSkillDir, signal, expandHome); const installCmd = `set -euTARGET=${quotePosixShellArg(remoteDir, { expandHome })}SKILL=${quotedRemoteSkillDir}if command -v uv >/dev/null 2>&1; then uv pip install --python "$TARGET/.venv" -e "$SKILL"elif command -v nix >/dev/null 2>&1; then nix shell nixpkgs#uv -c uv pip install --python "$TARGET/.venv" -e "$SKILL"else "$TARGET/.venv/bin/pip" install -e "$SKILL"fitouch "$SKILL/.prime-agent-installed"`; const installResult = await execRemote(hostConfig, installCmd, { signal }); if (installResult.exitCode !== 0) { throw new Error( `remote install exited with code ${installResult.exitCode}: ${installResult.stderr || installResult.stdout}`, ); } } catch (error) { throw new Error( `Failed to install Python skill ${skill.name || skill.importName} on ${hostConfig.host}: ${errorMessage(error)}`, ); } }}