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.
42 kB · 1221 lines
Python
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222"""Remote kernel management and execution for Prime Agent."""
from __future__ import annotations
import astimport asyncioimport base64import jsonimport osimport timeimport uuidfrom pathlib import PurePosixPathfrom typing import Any, AsyncIterator, Awaitable, Callable, Generator, Generic, Sequence, TypeVar
from rlm import host_requestfrom rlm.bash import BashResult
T = TypeVar("T")
KNOWN_PRIME_SKILLS: frozenset[str] = frozenset({ "agent_message", "agent_observe", "atproto", "attach_image", "browser", "compact", "dependency_recon", "edit", "git", "goal", "ideonomy_plain", "ideonomy_rich", "jj", "linear", "lucid_code", "lucid_review", "nix", "notion", "prime_agent_update", "prime_intellect", "refine", "remote", "repo", "review_fanout", "rlm_heartbeat", "search", "skill_creator", "ssh_control", "tangled", "tangled_lexicon_codegen", "tangled_linux_validation", "tangled_spindle_ops", "threads", "websearch", "yt",})
def _decode_response_text(raw: str) -> str: if raw.startswith(("'", '"')): try: val = ast.literal_eval(raw) if isinstance(val, str): return val except (ValueError, SyntaxError): pass return raw
def _coerce_path(path: Any, kernel: RemoteKernelHandle) -> str: if isinstance(path, RemotePath): if path._kernel is not kernel: raise ValueError("RemotePath belongs to a different remote connection") return str(path._path) if isinstance(path, os.PathLike): fspath = os.fspath(path) if isinstance(fspath, bytes): raise TypeError("Bytes paths are not supported") return str(fspath) if isinstance(path, bytes): raise TypeError("Bytes paths are not supported") if isinstance(path, str): return path raise TypeError(f"Expected str, os.PathLike, or RemotePath, got {type(path).__name__}")
class RemoteExecutionError(RuntimeError): """Raised when remote code execution fails."""
def __init__(self, message: str, ename: str = "", evalue: str = "", traceback: list[str] | None = None): super().__init__(message) self.ename = ename self.evalue = evalue self.traceback = traceback or []
class RemoteHandle(Generic[T]): """Live handle to remote execution; behaves like BashHandle."""
def __init__( self, req_id: str, task: asyncio.Task[T], interrupt_fn: Callable[[str], asyncio.Future[None] | None], stdout_chunks: list[str], ): self._req_id = req_id self._task = task self._interrupt_fn = interrupt_fn self._stdout_chunks = stdout_chunks self._started = time.monotonic()
@property def running(self) -> bool: return not self._task.done()
def output(self) -> str: return "".join(self._stdout_chunks)
def tail(self, n: int = 50) -> str: lines = self.output().splitlines(keepends=True) return "".join(lines[-n:]) if lines else ""
def poll(self) -> T | None: if self._task.done(): return self._task.result() return None
def kill(self) -> None: if not self._task.done(): self._interrupt_fn(self._req_id) self._task.cancel()
def __await__(self) -> Generator[Any, None, T]: return self._task.__await__()
def __repr__(self) -> str: state = "running" if self.running else "done" return f"<RemoteHandle {self._req_id[:8]} ({state})>"
class RemoteSkillProxy: """Proxy for remote skill execution."""
def __init__(self, kernel: RemoteKernelHandle, skill: str, parts: tuple[str, ...] = ()): self._kernel = kernel self._skill = skill self._parts = parts
def __getattr__(self, name: str) -> RemoteSkillProxy: if name.startswith("_"): raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}") return RemoteSkillProxy(self._kernel, self._skill, (*self._parts, name))
def __call__(self, *args: Any, **kwargs: Any) -> RemoteHandle[Any]: return self._kernel._call_skill(self._skill, self._parts, *args, **kwargs)
def __repr__(self) -> str: target = f"{self._skill}.{'.'.join(self._parts)}" if self._parts else self._skill return f"<RemoteSkillProxy {target}>"
class RemotePath: """Async remote path modeled after AnyIO Path."""
def __init__(self, path: PathArg, kernel: RemoteKernelHandle): self._kernel = kernel if isinstance(path, RemotePath): if path._kernel is not kernel: raise ValueError("RemotePath belongs to a different remote connection") self._path = path._path elif isinstance(path, PurePosixPath): self._path = path else: self._path = PurePosixPath(_coerce_path(path, kernel))
@property def name(self) -> str: return self._path.name
@property def stem(self) -> str: return self._path.stem
@property def suffix(self) -> str: return self._path.suffix
@property def suffixes(self) -> list[str]: return self._path.suffixes
@property def parent(self) -> RemotePath: return RemotePath(self._path.parent, self._kernel)
@property def parts(self) -> tuple[str, ...]: return self._path.parts
def with_name(self, name: str) -> RemotePath: return RemotePath(self._path.with_name(name), self._kernel)
def with_suffix(self, suffix: str) -> RemotePath: return RemotePath(self._path.with_suffix(suffix), self._kernel)
def with_stem(self, stem: str) -> RemotePath: return RemotePath(self._path.with_stem(stem), self._kernel)
def joinpath(self, *other: PathArg) -> RemotePath: subs = [_coerce_path(o, self._kernel) for o in other] return RemotePath(self._path.joinpath(*subs), self._kernel)
def __truediv__(self, other: PathArg) -> RemotePath: sub = _coerce_path(other, self._kernel) return RemotePath(self._path / sub, self._kernel)
def __str__(self) -> str: return str(self._path)
def __repr__(self) -> str: return f"RemotePath({str(self._path)!r}, host={self._kernel.host!r})"
def __eq__(self, other: Any) -> bool: if isinstance(other, RemotePath): return self._kernel is other._kernel and self._path == other._path return NotImplemented
def __hash__(self) -> int: return hash((self._path, id(self._kernel)))
async def exists(self) -> bool: def _remote(p: str) -> bool: import os return os.path.exists(p) return await self._kernel.call(_remote, str(self._path))
async def is_file(self) -> bool: def _remote(p: str) -> bool: import os return os.path.isfile(p) return await self._kernel.call(_remote, str(self._path))
async def is_dir(self) -> bool: def _remote(p: str) -> bool: import os return os.path.isdir(p) return await self._kernel.call(_remote, str(self._path))
async def read_text(self, encoding: str = "utf-8", errors: str = "strict") -> str: def _remote(p: str, enc: str, err: str) -> str: with open(p, "r", encoding=enc, errors=err) as f: return f.read() return await self._kernel.call(_remote, str(self._path), encoding, errors)
async def read_bytes(self) -> bytes: def _remote(p: str) -> bytes: with open(p, "rb") as f: return f.read() return await self._kernel.call(_remote, str(self._path))
async def write_text(self, data: str, encoding: str = "utf-8", errors: str = "strict") -> int: def _remote(p: str, content: str, enc: str, err: str) -> int: with open(p, "w", encoding=enc, errors=err) as f: return f.write(content) return await self._kernel.call(_remote, str(self._path), data, encoding, errors)
async def write_bytes(self, data: bytes) -> int: def _remote(p: str, content: bytes) -> int: with open(p, "wb") as f: return f.write(content) return await self._kernel.call(_remote, str(self._path), data)
async def stat(self) -> dict[str, Any]: def _remote(p: str) -> dict[str, Any]: import os st = os.stat(p) return { "st_size": st.st_size, "st_mode": st.st_mode, "st_mtime": st.st_mtime, "st_atime": st.st_atime, "st_ctime": st.st_ctime, "st_uid": st.st_uid, "st_gid": st.st_gid, } return await self._kernel.call(_remote, str(self._path))
async def mkdir(self, mode: int = 0o777, parents: bool = False, exist_ok: bool = False) -> None: def _remote(p: str, m: int, par: bool, ok: bool) -> None: import pathlib pathlib.Path(p).mkdir(mode=m, parents=par, exist_ok=ok) await self._kernel.call(_remote, str(self._path), mode, parents, exist_ok)
async def unlink(self, missing_ok: bool = False) -> None: def _remote(p: str, ok: bool) -> None: import pathlib pathlib.Path(p).unlink(missing_ok=ok) await self._kernel.call(_remote, str(self._path), missing_ok)
async def rmdir(self) -> None: def _remote(p: str) -> None: import os os.rmdir(p) await self._kernel.call(_remote, str(self._path))
async def resolve(self) -> RemotePath: def _remote(p: str) -> str: import pathlib return str(pathlib.Path(p).resolve()) resolved = await self._kernel.call(_remote, str(self._path)) return RemotePath(resolved, self._kernel)
async def iterdir(self) -> AsyncIterator[RemotePath]: def _remote(p: str) -> list[str]: import os return os.listdir(p) entries = await self._kernel.call(_remote, str(self._path)) for name in entries: yield RemotePath(self._path / name, self._kernel)
async def glob(self, pattern: str) -> AsyncIterator[RemotePath]: def _remote(p: str, pat: str) -> list[str]: import pathlib return [str(match) for match in pathlib.Path(p).glob(pat)] matches = await self._kernel.call(_remote, str(self._path), pattern) for path_str in matches: yield RemotePath(path_str, self._kernel)
PathArg = str | os.PathLike[str] | RemotePath
class RemotePathlibFacade: """Facade for remote pathlib access."""
def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel self.PurePosixPath = PurePosixPath
def Path(self, *parts: PathArg) -> RemotePath: if not parts: return RemotePath(".", self._kernel) str_parts = [_coerce_path(p, self._kernel) for p in parts] return RemotePath(PurePosixPath(*str_parts), self._kernel)
def __call__(self, *parts: PathArg) -> RemotePath: return self.Path(*parts)
class RemoteSysPathFacade: def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel
async def get(self) -> list[str]: def _remote() -> list[str]: import sys return list(sys.path) return await self._kernel.call(_remote)
async def append(self, dir_path: PathArg) -> None: p = _coerce_path(dir_path, self._kernel) def _remote(target: str) -> None: import sys if target not in sys.path: sys.path.append(target) await self._kernel.call(_remote, p)
async def prepend(self, dir_path: PathArg) -> None: p = _coerce_path(dir_path, self._kernel) def _remote(target: str) -> None: import sys if target in sys.path: sys.path.remove(target) sys.path.insert(0, target) await self._kernel.call(_remote, p)
class RemoteSysFacade: """Remote process metadata and state facade."""
def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel self.path = RemoteSysPathFacade(kernel)
@property def platform(self) -> str: return self._kernel._sys_platform
@property def executable(self) -> str: return self._kernel._sys_executable
@property def version(self) -> str: return self._kernel._sys_version
class RemoteEnvironFacade: def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel
async def get(self, key: str, default: str | None = None) -> str | None: def _remote(k: str, d: str | None) -> str | None: import os return os.environ.get(k, d) return await self._kernel.call(_remote, key, default)
async def set(self, key: str, val: str) -> None: def _remote(k: str, v: str) -> None: import os os.environ[k] = str(v) await self._kernel.call(_remote, key, val)
async def delete(self, key: str) -> None: def _remote(k: str) -> None: import os os.environ.pop(k, None) await self._kernel.call(_remote, key)
async def update(self, mapping: dict[str, str]) -> None: def _remote(m: dict[str, str]) -> None: import os os.environ.update({str(k): str(v) for k, v in m.items()}) await self._kernel.call(_remote, mapping)
async def snapshot(self) -> dict[str, str]: def _remote() -> dict[str, str]: import os return dict(os.environ) return await self._kernel.call(_remote)
class RemoteOsFacade: """Remote process and filesystem primitives facade."""
def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel self.environ = RemoteEnvironFacade(kernel)
async def chdir(self, path: PathArg) -> None: p = _coerce_path(path, self._kernel) def _remote(target: str) -> None: import os os.chdir(target) await self._kernel.call(_remote, p)
async def getcwd(self) -> str: def _remote() -> str: import os return os.getcwd() return await self._kernel.call(_remote)
async def listdir(self, path: PathArg = ".") -> list[str]: p = _coerce_path(path, self._kernel) def _remote(target: str) -> list[str]: import os return os.listdir(target) return await self._kernel.call(_remote, p)
async def mkdir(self, path: PathArg, mode: int = 0o777) -> None: p = _coerce_path(path, self._kernel) def _remote(target: str, m: int) -> None: import os os.mkdir(target, mode=m) await self._kernel.call(_remote, p, mode)
async def makedirs(self, path: PathArg, mode: int = 0o777, exist_ok: bool = False) -> None: p = _coerce_path(path, self._kernel) def _remote(target: str, m: int, ok: bool) -> None: import os os.makedirs(target, mode=m, exist_ok=ok) await self._kernel.call(_remote, p, mode, exist_ok)
async def remove(self, path: PathArg) -> None: p = _coerce_path(path, self._kernel) def _remote(target: str) -> None: import os os.remove(target) await self._kernel.call(_remote, p)
class RemoteShutilFacade: """Remote high-level filesystem primitives facade."""
def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel
async def copy(self, src: PathArg, dst: PathArg) -> str: s = _coerce_path(src, self._kernel) d = _coerce_path(dst, self._kernel) def _remote(src_p: str, dst_p: str) -> str: import shutil return shutil.copy(src_p, dst_p) return await self._kernel.call(_remote, s, d)
async def copytree(self, src: PathArg, dst: PathArg, **kwargs: Any) -> str: s = _coerce_path(src, self._kernel) d = _coerce_path(dst, self._kernel) def _remote(src_p: str, dst_p: str, kw: dict[str, Any]) -> str: import shutil return shutil.copytree(src_p, dst_p, **kw) return await self._kernel.call(_remote, s, d, kwargs)
async def rmtree(self, path: PathArg, ignore_errors: bool = False) -> None: p = _coerce_path(path, self._kernel) def _remote(target: str, ign: bool) -> None: import shutil shutil.rmtree(target, ignore_errors=ign) await self._kernel.call(_remote, p, ignore_errors)
async def move(self, src: PathArg, dst: PathArg) -> str: s = _coerce_path(src, self._kernel) d = _coerce_path(dst, self._kernel) def _remote(src_p: str, dst_p: str) -> str: import shutil return shutil.move(src_p, dst_p) return await self._kernel.call(_remote, s, d)
async def which(self, cmd: str) -> str | None: def _remote(c: str) -> str | None: import shutil return shutil.which(c) return await self._kernel.call(_remote, cmd)
class RemoteControlPlane: """Connection and execution plumbing for a remote kernel."""
def __init__(self, kernel: RemoteKernelHandle): self._kernel = kernel
@property def capabilities(self) -> frozenset[str]: return frozenset(self._kernel._capabilities)
def execute(self, code: str) -> RemoteHandle[Any]: return self._kernel._execute_code(code)
def eval(self, expr: str) -> RemoteHandle[Any]: return self._kernel._eval_expr(expr)
def ship(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]: return self._kernel._call_fn(fn, *args, **kwargs)
call = ship
async def aclose(self) -> None: await self._kernel.close()
close = aclose
async def refresh_capabilities(self) -> frozenset[str]: return await self._kernel._discover_capabilities()
@property def closed(self) -> bool: return self._kernel.closed
class RemoteKernelHandle: """Active connection to an IPython/RLM kernel running on a remote host."""
def __init__( self, host: str, proc: asyncio.subprocess.Process, remote_python: str, remote_dir: str, capabilities: Sequence[str] | set[str] | None = None, ): self.host = host self._proc = proc self.remote_python = remote_python self.remote_dir = remote_dir or "" self._closed = False self._lock = asyncio.Lock() self._imported_modules: set[str] = set() self._capabilities: set[str] = set(capabilities) if capabilities is not None else set() self._sys_platform: str = "" self._sys_executable: str = self.remote_python self._sys_version: str = "" self._close_task: asyncio.Task[None] | None = None self.control = RemoteControlPlane(self) self.pathlib = RemotePathlibFacade(self) self.sys = RemoteSysFacade(self) self.os = RemoteOsFacade(self) self.shutil = RemoteShutilFacade(self)
@property def closed(self) -> bool: return self._closed or self._proc.returncode is not None
async def import_(self, *modules: str, reload: bool = False) -> RemoteKernelHandle: lines: list[str] = [] new_modules: list[str] = [] for mod in modules: mod = mod.strip() if not mod: continue if not reload and mod in self._imported_modules: continue if reload: base_name = mod.split()[0] lines.append(f"import sys; sys.modules.pop({repr(base_name)}, None)") if mod.startswith("import ") or mod.startswith("from "): lines.append(mod) else: lines.append(f"import {mod}") new_modules.append(mod)
if lines: import_code = "\n".join(lines) await self._execute_raw(import_code) self._imported_modules.update(new_modules) return self
imports = import_
def call(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]: return self.control.ship(fn, *args, **kwargs)
ship = call
def bash(self, command: str, input: Any = None) -> RemoteHandle[BashResult]: if hasattr(input, "text") and callable(input.text): input_payload = input.text() elif hasattr(input, "output"): input_payload = input.output() if callable(input.output) else input.output else: input_payload = input code = f"""import jsonfrom rlm.bash import bash_h = bash({repr(command)}, input={repr(input_payload)})_res = await _hjson.dumps({{"exit_code": _res.exit_code, "output": _res.output, "duration": _res.duration, "stdout": _res.stdout, "stderr": _res.stderr}})""" req_id = uuid.uuid4().hex stdout_chunks: list[str] = []
async def _run() -> BashResult: raw = await self._execute_request(req_id, code, stdout_chunks) raw = _decode_response_text(raw) data = json.loads(raw) return BashResult( exit_code=data["exit_code"], output=data["output"], duration=data.get("duration", 0.0), stdout=data.get("stdout", data["output"]), stderr=data.get("stderr", ""), )
task = asyncio.create_task(_run()) return RemoteHandle( req_id=req_id, task=task, interrupt_fn=self._send_interrupt, stdout_chunks=stdout_chunks, )
def _execute_code(self, code: str) -> RemoteHandle[Any]: req_id = uuid.uuid4().hex stdout_chunks: list[str] = []
async def _run() -> Any: raw = await self._execute_request(req_id, code, stdout_chunks) return _decode_response_text(raw)
task = asyncio.create_task(_run()) return RemoteHandle( req_id=req_id, task=task, interrupt_fn=self._send_interrupt, stdout_chunks=stdout_chunks, )
def _eval_expr(self, expr: str) -> RemoteHandle[Any]: code = f"""import dill, base64, inspect, json_val = eval({repr(expr)}, globals())if inspect.isawaitable(_val): _val = await _valtry: _dumped = json.dumps({{"status": "ok", "result": _val}})except Exception: _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_val)).decode()}})_dumped""" req_id = uuid.uuid4().hex stdout_chunks: list[str] = []
async def _run() -> Any: raw = await self._execute_request(req_id, code, stdout_chunks) raw = _decode_response_text(raw) parsed = json.loads(raw) status = parsed.get("status") if status == "ok": return parsed.get("result") if status == "dill": import dill return dill.loads(base64.b64decode(parsed.get("result"))) return parsed
task = asyncio.create_task(_run()) return RemoteHandle( req_id=req_id, task=task, interrupt_fn=self._send_interrupt, stdout_chunks=stdout_chunks, )
def _call_fn(self, fn: Callable, *args: Any, **kwargs: Any) -> RemoteHandle[Any]: import dill import importlib import types
def _save_mod(pickler, obj): pickler.save_reduce(importlib.import_module, (obj.__name__,), obj=obj)
try: dill.register(types.ModuleType)(_save_mod) for sub in types.ModuleType.__subclasses__(): dill.register(sub)(_save_mod) except Exception: pass
is_builtin = isinstance(fn, (types.BuiltinFunctionType, types.BuiltinMethodType)) orig_module = getattr(fn, "__module__", None) can_patch_module = not is_builtin and orig_module not in (None, "builtins", "posix", "nt", "os") try: if can_patch_module: fn.__module__ = "__main__" fn_b64 = base64.b64encode(dill.dumps(fn, recurse=True)).decode() finally: if can_patch_module and orig_module is not None: try: fn.__module__ = orig_module except Exception: pass
args_b64 = base64.b64encode(dill.dumps(args, recurse=True)).decode() kwargs_b64 = base64.b64encode(dill.dumps(kwargs, recurse=True)).decode()
code = f"""import dill, base64, inspect, json_fn = dill.loads(base64.b64decode({repr(fn_b64)}))_args = dill.loads(base64.b64decode({repr(args_b64)}))_kwargs = dill.loads(base64.b64decode({repr(kwargs_b64)}))
if hasattr(_fn, "__globals__"): _fn.__globals__.update(globals()) if "bash" not in _fn.__globals__: try: from rlm.bash import bash as _prime_agent_bash _fn.__globals__["bash"] = _prime_agent_bash except ImportError: pass
_res = _fn(*_args, **_kwargs)if inspect.isawaitable(_res): _res = await _res
try: _dumped = json.dumps({{"status": "ok", "result": _res}})except Exception: _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_res)).decode()}})_dumped""" req_id = uuid.uuid4().hex stdout_chunks: list[str] = []
async def _run() -> Any: raw = await self._execute_request(req_id, code, stdout_chunks) raw = _decode_response_text(raw) parsed = json.loads(raw) if parsed.get("status") == "ok": return parsed.get("result") if parsed.get("status") == "dill": import dill return dill.loads(base64.b64decode(parsed.get("result"))) return parsed
task = asyncio.create_task(_run()) return RemoteHandle( req_id=req_id, task=task, interrupt_fn=self._send_interrupt, stdout_chunks=stdout_chunks, )
def _call_skill(self, skill: str, parts: tuple[str, ...], *args: Any, **kwargs: Any) -> RemoteHandle[Any]: import dill import importlib import types
def _save_mod(pickler, obj): pickler.save_reduce(importlib.import_module, (obj.__name__,), obj=obj)
try: dill.register(types.ModuleType)(_save_mod) for sub in types.ModuleType.__subclasses__(): dill.register(sub)(_save_mod) except Exception: pass
args_b64 = base64.b64encode(dill.dumps(args, recurse=True)).decode() kwargs_b64 = base64.b64encode(dill.dumps(kwargs, recurse=True)).decode()
code = f"""import dill, base64, inspect, json, importlib
_mod = importlib.import_module({repr(skill)})_target = _modfor _p in {repr(parts)}: _target = getattr(_target, _p)
if not callable(_target) and hasattr(_target, "run") and callable(_target.run): _fn = _target.runelse: _fn = _target
_args = dill.loads(base64.b64decode({repr(args_b64)}))_kwargs = dill.loads(base64.b64decode({repr(kwargs_b64)}))
_res = _fn(*_args, **_kwargs)if inspect.isawaitable(_res): _res = await _res
try: _dumped = json.dumps({{"status": "ok", "result": _res}})except Exception: _dumped = json.dumps({{"status": "dill", "result": base64.b64encode(dill.dumps(_res)).decode()}})_dumped""" req_id = uuid.uuid4().hex stdout_chunks: list[str] = []
async def _run() -> Any: raw = await self._execute_request(req_id, code, stdout_chunks) raw = _decode_response_text(raw) parsed = json.loads(raw) if parsed.get("status") == "ok": return parsed.get("result") if parsed.get("status") == "dill": import dill return dill.loads(base64.b64decode(parsed.get("result"))) return parsed
task = asyncio.create_task(_run()) return RemoteHandle( req_id=req_id, task=task, interrupt_fn=self._send_interrupt, stdout_chunks=stdout_chunks, )
async def _discover_capabilities(self) -> frozenset[str]: if self.closed: return frozenset(self._capabilities) code = f"""import importlib.util, json, os, sys
_known = {sorted(list(KNOWN_PRIME_SKILLS))!r}_skills = set()
for _name in _known: try: if _name in sys.modules or importlib.util.find_spec(_name) is not None: _skills.add(_name) except Exception: pass
_remote_dir = os.path.expandvars(os.path.expanduser({repr(self.remote_dir)}))if _remote_dir: _skills_path = os.path.join(_remote_dir, "skills") if os.path.isdir(_skills_path): for _entry in os.listdir(_skills_path): if not _entry.startswith((".", "_")) and os.path.isdir(os.path.join(_skills_path, _entry)): _norm = _entry.replace("-", "_") try: if importlib.util.find_spec(_norm) is not None: _skills.add(_norm) except Exception: pass
json.dumps({{ "skills": sorted(list(_skills)), "platform": sys.platform, "executable": sys.executable, "version": sys.version,}})""" try: raw = await self._execute_raw(code) raw = _decode_response_text(raw) data = json.loads(raw) if isinstance(data, dict): skills = data.get("skills", []) if isinstance(skills, list): self._capabilities = set(skills) self._sys_platform = data.get("platform", "") self._sys_executable = data.get("executable", self.remote_python) self._sys_version = data.get("version", "") elif isinstance(data, list): self._capabilities = set(data) except Exception: pass return frozenset(self._capabilities)
def __getattr__(self, name: str) -> RemoteSkillProxy: if name.startswith("_"): raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}") if name not in self._capabilities: raise AttributeError( f"'{name}' is not an available remote skill. Available: {sorted(self._capabilities)}" ) return RemoteSkillProxy(self, name)
def __getitem__(self, skill_name: str) -> RemoteSkillProxy: norm = skill_name.replace("-", "_") if norm in self._capabilities: return RemoteSkillProxy(self, norm) if skill_name in self._capabilities: return RemoteSkillProxy(self, skill_name) raise KeyError( f"'{skill_name}' is not an available remote skill. Available: {sorted(self._capabilities)}" )
def __contains__(self, key: str) -> bool: norm = key.replace("-", "_") return norm in self._capabilities or key in self._capabilities
def __dir__(self) -> list[str]: attrs = set(super().__dir__()) attrs.update(self._capabilities) return sorted(attrs)
def _send_interrupt(self, req_id: str) -> None: if self.closed or not self._proc.stdin: return msg = (json.dumps({"type": "interrupt", "id": req_id}) + "\n").encode("utf-8") try: self._proc.stdin.write(msg) except Exception: pass
async def _execute_raw(self, code: str) -> str: req_id = uuid.uuid4().hex stdout_chunks: list[str] = [] return await self._execute_request(req_id, code, stdout_chunks)
async def _dispatch_host_request(self, event: dict[str, Any]) -> None: h_id = event.get("id") h_data = event.get("data", {}) try: reply_res = await host_request(h_data.get("type", "unknown"), h_data) reply = {"type": "host_reply", "id": h_id, "data": {"status": "ok", "result": reply_res}} except Exception as exc: reply = {"type": "host_reply", "id": h_id, "data": {"status": "error", "error": str(exc)}} if self._proc.stdin and not self._proc.stdin.is_closing(): self._proc.stdin.write((json.dumps(reply) + "\n").encode("utf-8")) await self._proc.stdin.drain()
async def _drain_until_done(self, req_id: str) -> None: assert self._proc.stdout is not None while True: line = await self._proc.stdout.readline() if not line: raise RuntimeError(f"Remote kernel on {self.host} closed unexpectedly during drain") try: event = json.loads(line.decode("utf-8")) except Exception: continue
ev_type = event.get("event") if ev_type == "host_request": await self._dispatch_host_request(event) continue
if event.get("id") == req_id and ev_type == "done": return
async def _execute_request(self, req_id: str, code: str, stdout_chunks: list[str]) -> str: if self.closed: raise RuntimeError(f"Remote kernel connection to {self.host} is closed")
async with self._lock: req = {"type": "execute", "id": req_id, "code": code} data = (json.dumps(req) + "\n").encode("utf-8")
assert self._proc.stdin is not None assert self._proc.stdout is not None
result_text: str | None = None error_data: dict[str, Any] | None = None submitted = False
try: self._proc.stdin.write(data) submitted = True await self._proc.stdin.drain()
while True: line = await self._proc.stdout.readline() if not line: raise RuntimeError(f"Remote kernel on {self.host} closed unexpectedly during execution")
try: event = json.loads(line.decode("utf-8")) except Exception: continue
ev_type = event.get("event") ev_id = event.get("id")
if ev_type == "host_request": await self._dispatch_host_request(event) continue
if ev_id != req_id: continue
if ev_type == "stdout": text = event.get("text", "") stdout_chunks.append(text) elif ev_type == "result": result_text = event.get("text", "") elif ev_type == "error": error_data = event elif ev_type == "done": break except asyncio.CancelledError: if submitted: async def _cancel_cleanup() -> None: self._send_interrupt(req_id) try: await asyncio.wait_for(self._drain_until_done(req_id), timeout=3.0) except (asyncio.TimeoutError, Exception): await self._poison_and_close()
cleanup_task = asyncio.create_task(_cancel_cleanup()) cancelled = True while not cleanup_task.done(): try: await asyncio.shield(cleanup_task) except asyncio.CancelledError: cancelled = True if cancelled: raise asyncio.CancelledError() raise
if error_data is not None: ename = error_data.get("ename", "Error") evalue = error_data.get("evalue", "") tb = error_data.get("traceback", []) msg = f"{ename}: {evalue}\n" + "".join(tb) raise RemoteExecutionError(msg, ename=ename, evalue=evalue, traceback=tb)
if result_text is not None: return result_text return "".join(stdout_chunks)
async def _do_close(self) -> None: try: if self._proc.stdin: self._proc.stdin.close() await self._proc.stdin.wait_closed() except Exception: pass try: self._proc.terminate() await asyncio.wait_for(self._proc.wait(), timeout=2.0) except Exception: try: self._proc.kill() await self._proc.wait() except Exception: pass
async def _poison_and_close(self) -> None: self._closed = True if self._close_task is None: self._close_task = asyncio.create_task(self._do_close()) task = self._close_task cancelled = False while not task.done(): try: await asyncio.shield(task) except asyncio.CancelledError: cancelled = True if cancelled: raise asyncio.CancelledError()
async def close(self) -> None: await self._poison_and_close()
aclose = close
async def __aenter__(self) -> RemoteKernelHandle: return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: await self.close()
class _ConnectContext: def __init__(self, coro: Awaitable[RemoteKernelHandle]): self._coro = coro self._handle: RemoteKernelHandle | None = None
def __await__(self) -> Generator[Any, None, RemoteKernelHandle]: return self._coro.__await__()
async def __aenter__(self) -> RemoteKernelHandle: self._handle = await self._coro return await self._handle.__aenter__()
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: if self._handle is not None: await self._handle.__aexit__(exc_type, exc_val, exc_tb)
def connect( host: str, *, port: int | None = None, identity_file: str | None = None, ssh_options: Sequence[str] | None = None, cache_dir: str | None = None, python_cmd: str | None = None,) -> _ConnectContext: """Connect to or spawn an on-demand remote kernel on a remote host via SSH."""
async def _do_connect() -> RemoteKernelHandle: ensure_payload: dict[str, Any] = {"host": host} if port is not None: ensure_payload["port"] = port if identity_file is not None: ensure_payload["identityFile"] = identity_file if ssh_options: ensure_payload["sshOptions"] = list(ssh_options) if cache_dir is not None: ensure_payload["cacheDir"] = cache_dir if python_cmd is not None: ensure_payload["pythonCmd"] = python_cmd
info = await host_request("remote.ensure", ensure_payload) remote_python = info.get("remotePython", "python3") remote_dir = info.get("remoteDir", "")
ssh_args = [ "ssh", "-o", "ControlMaster=auto", "-o", "ControlPersist=30m", "-o", "ControlPath=~/.ssh/cm-%C", ] if port is not None: ssh_args.extend(["-p", str(port)]) if identity_file is not None: ssh_args.extend(["-i", identity_file]) if ssh_options: ssh_args.extend(ssh_options) ssh_args.append(host) ssh_args.append(f"{remote_python} -m rlm.repl")
proc = await asyncio.create_subprocess_exec( *ssh_args, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, )
try: assert proc.stdout is not None ready_line = await proc.stdout.readline() if not ready_line: stderr_out = await proc.stderr.read() if proc.stderr else b"" raise RuntimeError(f"Failed to start remote kernel on {host}: {stderr_out.decode('utf-8')}")
try: ready_data = json.loads(ready_line.decode("utf-8")) if ready_data.get("event") != "ready": raise RuntimeError(f"Unexpected initial message from remote kernel on {host}: {ready_line.decode('utf-8')}") except json.JSONDecodeError: raise RuntimeError(f"Corrupt startup output from remote kernel on {host}: {ready_line.decode('utf-8')}")
handle = RemoteKernelHandle(host=host, proc=proc, remote_python=remote_python, remote_dir=remote_dir) try: await handle.control.refresh_capabilities() except Exception: pass return handle except BaseException: try: proc.kill() await proc.wait() except BaseException: pass raise
return _ConnectContext(_do_connect())
__all__ = [ "KNOWN_PRIME_SKILLS", "PathArg", "RemoteControlPlane", "RemoteEnvironFacade", "RemoteExecutionError", "RemoteHandle", "RemoteKernelHandle", "RemoteOsFacade", "RemotePath", "RemotePathlibFacade", "RemoteShutilFacade", "RemoteSkillProxy", "RemoteSysFacade", "RemoteSysPathFacade", "connect",]