ive harnessed the harness
Something went wrong. Try again.
34 kB · 697 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698"""Target-side dispatcher for opaque auxiliary-namespace operations.
The host transports ``operation`` and JSON data without interpreting either.This module is the one vocabulary owner on the target. It executes onlyexplicit operations, returns JSON-compatible bounded values, and never usessource text as a fallback for trusted closure shipment."""from __future__ import annotations
import astimport asyncioimport base64import contextlibimport importlibimport importlib.utilimport inspectimport mathimport jsonimport osimport pathlibimport shutilimport sysimport timeimport tracebackfrom dataclasses import dataclassfrom typing import Any, Mapping, Sequence
MAX_FRAME_BYTES = 1_048_576ENVELOPE_HEADROOM = 8_576MAX_VALUE_BYTES = MAX_FRAME_BYTES - ENVELOPE_HEADROOMMAX_ERROR_BYTES = 16_384MAX_ARTIFACT_CHANNEL_BYTES = 16 * 1024 * 1024
class RemoteDispatchError(RuntimeError): """A typed target-side operation failure for the relay runner."""
def __init__(self, code: str, message: str): encoded = str(message).encode("utf-8") if len(encoded) > MAX_ERROR_BYTES: self.code = "error_too_large" message = f"remote error exceeds {MAX_ERROR_BYTES}-byte limit" else: self.code = code self.message = message super().__init__(message)
@dataclass(frozen=True, slots=True)class _CapturedStream: data: bytes exceeded: bool
@dataclass(frozen=True, slots=True)class _ProcessOutcome: exit_code: int | None stdout: _CapturedStream stderr: _CapturedStream duration: float timed_out: bool
async def _write_stdin(process: asyncio.subprocess.Process, data: bytes | None) -> None: stream = process.stdin if stream is None: return try: if data: stream.write(data) await stream.drain() finally: with contextlib.suppress(OSError): stream.close()
async def _spawn( *, cwd: str, argv: Sequence[str] | None = None, shell: str | None = None, stdin: bytes | None = None, timeout: float | None = None,) -> _ProcessOutcome: """Run one process and drain it completely. Callers decide what the outcome means.
Bounded pipes, one deadline and one kill path live here, so a leaked process or an undrained pipe has a single cause instead of one per caller. Stdin is fed in the background: writing it in the foreground can deadlock against a full stdout pipe. """ assert (argv is None) != (shell is None), "spawn takes exactly one of argv or shell" started = time.monotonic() streams = { "stdin": asyncio.subprocess.PIPE if stdin is not None else asyncio.subprocess.DEVNULL, "stdout": asyncio.subprocess.PIPE, "stderr": asyncio.subprocess.PIPE, "start_new_session": True, } process = ( await asyncio.create_subprocess_shell(shell, cwd=cwd, **streams) if shell is not None else await asyncio.create_subprocess_exec(*argv, cwd=cwd, **streams) ) tasks = ( asyncio.create_task(_read_capped(process.stdout, process)), asyncio.create_task(_read_capped(process.stderr, process)), asyncio.create_task(_write_stdin(process, stdin)), ) timed_out = False try: try: await asyncio.wait_for(process.wait(), timeout=timeout) except asyncio.TimeoutError: timed_out = True with contextlib.suppress(ProcessLookupError): process.kill() await process.wait() stdout, stderr, _ = await asyncio.gather(*tasks) except asyncio.CancelledError: with contextlib.suppress(ProcessLookupError): process.kill() with contextlib.suppress(Exception): await process.wait() for task in tasks: task.cancel() raise return _ProcessOutcome(process.returncode, stdout, stderr, time.monotonic() - started, timed_out)
def _wire_encode(value: Any) -> Any: if isinstance(value, (bytes, bytearray)): return {"__klbr_bytes__": base64.b64encode(bytes(value)).decode("ascii")} if isinstance(value, pathlib.Path): return str(value) if isinstance(value, Mapping): return {str(key): _wire_encode(item) for key, item in value.items()} if isinstance(value, (list, tuple)): return [_wire_encode(item) for item in value] if value is None or isinstance(value, (str, int, float, bool)): return value raise RemoteDispatchError("unsupported_value", f"result type is not JSON-compatible: {type(value).__name__}")
def _wire_decode(value: Any) -> Any: if isinstance(value, Mapping): if set(value) == {"__klbr_bytes__"} and isinstance(value["__klbr_bytes__"], str): try: return base64.b64decode(value["__klbr_bytes__"], validate=True) except (ValueError, TypeError) as error: raise RemoteDispatchError("invalid_payload", "invalid byte encoding") from error return {str(key): _wire_decode(item) for key, item in value.items()} if isinstance(value, list): return [_wire_decode(item) for item in value] return value
def _bounded(value: Any) -> Any: try: encoded = json.dumps(_wire_encode(value), ensure_ascii=False, allow_nan=False).encode("utf-8") except (TypeError, ValueError) as error: raise RemoteDispatchError("unsupported_value", f"result is not JSON-compatible: {error}") from error if len(encoded) > MAX_VALUE_BYTES: raise RemoteDispatchError("result_too_large", f"encoded result exceeds {MAX_VALUE_BYTES} bytes") return value
def _require_mapping(payload: Any) -> Mapping[str, Any]: if not isinstance(payload, Mapping): raise RemoteDispatchError("invalid_payload", "operation payload must be an object") return payload
def _text(payload: Mapping[str, Any], key: str, *, required: bool = True, maximum: int = MAX_VALUE_BYTES) -> str | None: value = payload.get(key) if value is None and not required: return None if not isinstance(value, str): raise RemoteDispatchError("invalid_payload", f"{key} must be text") if len(value.encode("utf-8")) > maximum: raise RemoteDispatchError("payload_too_large", f"{key} exceeds {maximum} bytes") return value
class RemoteDispatcher: """Persistent target namespace for one auxiliary connection."""
def __init__(self, workspace: str | os.PathLike[str] | None = None): self.workspace = pathlib.Path(workspace).resolve() if workspace else None self._cwd = self.workspace or pathlib.Path.cwd() self.globals: dict[str, Any] = { "__name__": "__klbr_remote__", "__builtins__": __builtins__, } self._imports: set[str] = set() self._active: dict[str, asyncio.Task[Any]] = {} self._dispatch_handlers = { "clear": lambda _: self.clear(), "capabilities": lambda _: {"capabilities": self.capabilities}, "imports": lambda p: RemoteDispatcher._imports(self, p), "execute": self._execute, "eval": self._eval, "reflect": self._reflect, "skill.call": self._skill_call, "ship": self._ship, "shell": self._shell, } self._namespace_handlers = ( ("path.", self._path), ("os.", self._os), ("sys.path.", self._sys_path), ("shutil.", self._shutil), ("workspace.", self._workspace), )
@property def capabilities(self) -> list[dict[str, Any]]: dill_available = importlib.util.find_spec("dill") is not None return [ {"capability": "shell", "shell": os.environ.get("SHELL")}, {"capability": "file_system"}, {"capability": "imports"}, {"capability": "closures", "dill_available": dill_available}, {"capability": "skill_manifest", "revision": "remote-dispatch-v1"}, ]
async def dispatch(self, operation: str, payload: Mapping[str, Any], operation_id: str | None = None) -> Any: if not isinstance(operation, str) or not operation or any(part in {"", ".", ".."} for part in operation.split(".")): raise RemoteDispatchError("invalid_operation", "operation must be a dotted name") payload = _require_mapping(_wire_decode(payload)) task = asyncio.current_task() if operation_id and task is not None: self._active[operation_id] = task try: result = await self._dispatch(operation, payload) _bounded(result) return _wire_encode(result) finally: if operation_id: self._active.pop(operation_id, None)
async def cancel(self, operation_id: str) -> dict[str, str]: task = self._active.get(operation_id) if task is None: return {"state": "unknown", "reason": "operation is not active"} task.cancel() try: await asyncio.wait_for(asyncio.shield(task), timeout=1.0) except (asyncio.CancelledError, Exception): pass if task.done(): return {"state": "confirmed"} return {"state": "requested", "outstanding": "target task cancellation is pending"}
async def clear(self) -> dict[str, Any]: """Clear target namespace globals and imports.
Namespace clear cannot and does not cancel background/active tasks. Live tasks retain their references and continue running. """ self.globals = { "__name__": "__klbr_remote__", "__builtins__": __builtins__, } self._imports.clear() return { "cleared": True, "tasks_cancelled": 0, "active_tasks": len(self._active), }
async def _dispatch(self, operation: str, payload: Mapping[str, Any]) -> Any: handler = self._dispatch_handlers.get(operation) if handler is not None: result = handler(payload) return await result if inspect.isawaitable(result) else result for prefix, handler in self._namespace_handlers: if operation.startswith(prefix): return await handler(operation[len(prefix):], payload) raise RemoteDispatchError("unsupported", f"unknown remote operation: {operation}")
async def _imports(self, payload: Mapping[str, Any]) -> None: modules = payload.get("modules") if not isinstance(modules, list) or not all(isinstance(module, str) for module in modules): raise RemoteDispatchError("invalid_payload", "modules must be a list of names") reload = payload.get("reload", False) if type(reload) is not bool: raise RemoteDispatchError("invalid_payload", "reload must be boolean") for name in modules: if not name or not all(part.isidentifier() for part in name.split(".")): raise RemoteDispatchError("invalid_payload", f"invalid import name: {name!r}") try: module = importlib.import_module(name) if reload: module = importlib.reload(module) except Exception as error: raise RemoteDispatchError("import_failed", f"{name}: {error}") from error root = importlib.import_module(name.split(".", 1)[0]) self.globals[name.split(".", 1)[0]] = root self.globals[name.rsplit(".", 1)[-1]] = module self._imports.add(name) return None
async def _execute(self, payload: Mapping[str, Any]) -> None: source = _text(payload, "code", maximum=262_144) assert source is not None filename = "<klbr-remote>" try: tree = compile(source, filename, "exec", flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT | ast.PyCF_ONLY_AST, dont_inherit=True) code = compile(tree, filename, "exec", flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT, dont_inherit=True) result = eval(code, self.globals) if inspect.isawaitable(result): await result except asyncio.CancelledError: raise except BaseException as error: raise RemoteDispatchError("execution", "".join(traceback.format_exception(error))) from error return None
async def _eval(self, payload: Mapping[str, Any]) -> Any: expression = _text(payload, "expression", maximum=262_144) assert expression is not None try: result = eval(expression, self.globals) if inspect.isawaitable(result): result = await result return result except BaseException as error: raise RemoteDispatchError("evaluation", "".join(traceback.format_exception(error))) from error
def _resolve(self, path: Sequence[str]) -> Any: if not path or not all(isinstance(part, str) and part.isidentifier() for part in path): raise RemoteDispatchError("invalid_payload", "callable path is invalid") try: value = self.globals[path[0]] for part in path[1:]: value = getattr(value, part) return value except (KeyError, AttributeError) as error: raise RemoteDispatchError("unsupported", f"callable is not installed: {'.'.join(path)}") from error
async def _reflect(self, payload: Mapping[str, Any]) -> dict[str, Any]: path = payload.get("path") if not isinstance(path, list): raise RemoteDispatchError("invalid_payload", "reflection path must be a list") value = self._resolve(path) qualified = ".".join(path) try: signature = str(inspect.signature(value)) except (TypeError, ValueError): signature = None try: source = inspect.getsource(value) except (OSError, TypeError): source = None return { "qualified_name": qualified, "signature": signature, "source": source, "documentation": inspect.getdoc(value), }
async def _skill_call(self, payload: Mapping[str, Any]) -> Any: path = payload.get("path") args = payload.get("args", []) kwargs = payload.get("kwargs", {}) if not isinstance(path, list) or not isinstance(args, list) or not isinstance(kwargs, Mapping): raise RemoteDispatchError("invalid_payload", "skill.call fields are malformed") value = self._resolve(path) if not callable(value): raise RemoteDispatchError("unsupported", "skill target is not callable") try: result = value(*args, **dict(kwargs)) if inspect.isawaitable(result): result = await result return result except RemoteDispatchError: raise except BaseException as error: raise RemoteDispatchError("skill_call", str(error)) from error
async def _ship(self, payload: Mapping[str, Any]) -> Any: if payload.get("trusted") is not True or payload.get("serialization") != "dill" or payload.get("python_only") is not True: raise RemoteDispatchError("unsupported", "trusted closure shipment requires explicit dill mode") target_version = payload.get("python_version") local_version = f"{sys.version_info.major}.{sys.version_info.minor}" if target_version is not None and str(target_version) != local_version: raise RemoteDispatchError("unsupported", f"interpreter incompatible: target requires {target_version}, worker is {local_version}") dependencies = payload.get("dependencies", []) if not isinstance(dependencies, list) or not all(isinstance(dep, str) and dep for dep in dependencies): raise RemoteDispatchError("invalid_payload", "dependencies must be a list of names") missing = [] for dep in dependencies: try: available = importlib.util.find_spec(dep) is not None except (ImportError, ModuleNotFoundError, ValueError): available = False if not available: missing.append(dep) if missing: raise RemoteDispatchError("unsupported", f"dependencies unavailable in worker namespace: {missing!r}") try: import dill except ImportError as error: raise RemoteDispatchError("unsupported", "dill is not importable in the worker namespace") from error try: function = dill.loads(_decode_b64(payload, "function")) args = dill.loads(_decode_b64(payload, "args")) kwargs = dill.loads(_decode_b64(payload, "kwargs")) except Exception as error: raise RemoteDispatchError("invalid_payload", f"dill shipment could not be loaded: {error}") from error if not callable(function) or not isinstance(args, tuple) or not isinstance(kwargs, dict): raise RemoteDispatchError("invalid_payload", "dill shipment has invalid callable arguments") try: result = function(*args, **kwargs) if inspect.isawaitable(result): result = await result return result except BaseException as error: raise RemoteDispatchError("closure_call", str(error)) from error
def _path_value(self, payload: Mapping[str, Any]) -> pathlib.Path: value = _text(payload, "path", maximum=32_768) assert value is not None candidate = pathlib.Path(value) if self.workspace is None: raise RemoteDispatchError("workspace_unbound", "filesystem operations require a checked workspace") resolved = (candidate if candidate.is_absolute() else self._cwd / candidate).resolve() if resolved != self.workspace and self.workspace not in resolved.parents: raise RemoteDispatchError("path_escape", "path is outside the checked remote workspace") return resolved
async def _path(self, operation: str, payload: Mapping[str, Any]) -> Any: path = self._path_value(payload) try: if operation == "exists": return path.exists() if operation == "is_file": return path.is_file() if operation == "is_dir": return path.is_dir() if operation == "read_text": if path.stat().st_size > MAX_VALUE_BYTES: raise RemoteDispatchError("result_too_large", f"encoded result exceeds {MAX_VALUE_BYTES} bytes") raw = path.read_bytes() return raw.decode(encoding=_text(payload, "encoding", required=False) or "utf-8", errors=_text(payload, "errors", required=False) or "strict") if operation == "read_bytes": # Base64 markers add 4/3 overhead before the frame check. if path.stat().st_size > (MAX_VALUE_BYTES * 3) // 4: raise RemoteDispatchError("result_too_large", f"encoded byte result exceeds {MAX_VALUE_BYTES} bytes") return path.read_bytes() if operation == "write_text": data = _text(payload, "data") assert data is not None return path.write_text(data, encoding=_text(payload, "encoding", required=False) or "utf-8", errors=_text(payload, "errors", required=False) or "strict") if operation == "write_bytes": data = payload.get("data") if not isinstance(data, (bytes, bytearray)): raise RemoteDispatchError("invalid_payload", "write_bytes data must be bytes") return path.write_bytes(bytes(data)) if operation == "stat": value = path.stat() return {"mode": value.st_mode, "size": value.st_size, "mtime_ns": value.st_mtime_ns, "atime_ns": value.st_atime_ns, "ctime_ns": value.st_ctime_ns} if operation == "mkdir": path.mkdir(mode=int(payload.get("mode", 0o777)), parents=bool(payload.get("parents", False)), exist_ok=bool(payload.get("exist_ok", False))); return None if operation == "unlink": path.unlink(missing_ok=bool(payload.get("missing_ok", False))); return None if operation == "rmdir": path.rmdir(); return None if operation == "resolve": return str(path.resolve()) if operation == "iterdir": return [item.name for item in path.iterdir()] if operation == "glob": return [str(self._path_value({"path": str(item)})) for item in path.glob(_text(payload, "pattern") or "")] except RemoteDispatchError: raise except (OSError, ValueError) as error: raise RemoteDispatchError("filesystem", str(error)) from error raise RemoteDispatchError("unsupported", f"unknown path operation: {operation}")
async def _os(self, operation: str, payload: Mapping[str, Any]) -> Any: if operation == "environ.get": return os.environ.get(_text(payload, "key", maximum=256) or "", payload.get("default")) if operation == "environ.set": os.environ[_text(payload, "key", maximum=256) or ""] = _text(payload, "value") or ""; return None if operation == "environ.delete": os.environ.pop(_text(payload, "key", maximum=256) or "", None); return None if operation == "environ.update": values = payload.get("values") if not isinstance(values, Mapping): raise RemoteDispatchError("invalid_payload", "environment values must be an object") os.environ.update({str(k): str(v) for k, v in values.items()}); return None if operation == "environ.snapshot": return dict(os.environ) if operation == "chdir": candidate = self._path_value(payload) if not candidate.is_dir(): raise RemoteDispatchError("filesystem", "cwd is not an existing directory") self._cwd = candidate return None if operation == "getcwd": return str(self._cwd) if operation == "listdir": return os.listdir(self._path_value(payload)); if operation == "mkdir": os.mkdir(self._path_value(payload), int(payload.get("mode", 0o777))); return None if operation == "makedirs": os.makedirs(self._path_value(payload), int(payload.get("mode", 0o777)), exist_ok=bool(payload.get("exist_ok", False))); return None if operation == "remove": os.remove(self._path_value(payload)); return None raise RemoteDispatchError("unsupported", f"unknown os operation: {operation}")
async def _sys_path(self, operation: str, payload: Mapping[str, Any]) -> Any: if operation == "get": return list(sys.path) path = _text(payload, "path", maximum=32_768) assert path is not None if operation == "append": sys.path.append(path); return None if operation == "prepend": sys.path.insert(0, path); return None raise RemoteDispatchError("unsupported", f"unknown sys.path operation: {operation}")
async def _shutil(self, operation: str, payload: Mapping[str, Any]) -> Any: if self.workspace is None: raise RemoteDispatchError("workspace_unbound", "filesystem operations require a checked workspace") if operation == "copy": return shutil.copy(self._path_value({"path": payload.get("src")}), self._path_value({"path": payload.get("dst")})) if operation == "copytree": return shutil.copytree(self._path_value({"path": payload.get("src")}), self._path_value({"path": payload.get("dst")}), **dict(payload.get("kwargs", {}))) if operation == "rmtree": shutil.rmtree(self._path_value(payload), ignore_errors=bool(payload.get("ignore_errors", False))); return None if operation == "move": return shutil.move(self._path_value({"path": payload.get("src")}), self._path_value({"path": payload.get("dst")})) if operation == "which": return shutil.which(_text(payload, "command", maximum=256) or "") raise RemoteDispatchError("unsupported", f"unknown shutil operation: {operation}")
async def _workspace(self, operation: str, payload: Mapping[str, Any]) -> Any: """The same five verbs a local cell has, resolved against this target's tree.
The semantics live in `klbr.workspace`, which is importable beside this module, so a remote cell and a local cell cannot drift into two ideas of what `delete` refuses or what `ls` prunes. Nothing here restates those rules. """ if self.workspace is None: raise RemoteDispatchError("workspace_unbound", "tree operations require a checked workspace") from klbr.workspace import Workspace
tree = Workspace(self.workspace) relative = payload.get("path", ".") if not isinstance(relative, str): raise RemoteDispatchError("invalid_payload", "path must be a string") if operation in {"exists", "read", "write", "replace", "hashlines", "edit"}: handle = await tree.open_file(relative) if operation == "exists": return await handle.exists() if operation == "read": if (await handle.read()).__len__() > (MAX_VALUE_BYTES * 3) // 4: raise RemoteDispatchError("result_too_large", f"encoded byte result exceeds {MAX_VALUE_BYTES} bytes") return await handle.read() if operation == "hashlines": return await handle.hashlines() if operation == "write": data = payload.get("data") if not isinstance(data, (str, bytes, bytearray)): raise RemoteDispatchError("invalid_payload", "write data must be text or bytes") encoding = payload.get("encoding", "utf-8") await handle.write(bytes(data) if isinstance(data, (bytes, bytearray)) else data, encoding=str(encoding)) return None if operation == "replace": old, new = payload.get("old"), payload.get("new") if not isinstance(old, str) or not isinstance(new, str): raise RemoteDispatchError("invalid_payload", "replace needs string old and new") await handle.replace(old, new, encoding=str(payload.get("encoding", "utf-8"))) return None if operation == "edit": anchor = payload.get("anchor") if not isinstance(anchor, str): raise RemoteDispatchError("invalid_payload", "edit needs a string anchor") new = payload.get("new") if new is not None and not isinstance(new, str): raise RemoteDispatchError("invalid_payload", "edit new must be a string or absent") await handle.edit(anchor, new, after=bool(payload.get("after", False)), encoding=str(payload.get("encoding", "utf-8"))) return None if operation == "ls": listing = await tree.ls(relative, recursive=bool(payload.get("recursive", False)), match=payload.get("match") or None) return { "root": str(listing.root), "path": str(listing.path), "entries": [ {"path": str(entry.path), "relative": entry.relative, "is_dir": entry.is_dir, "size": entry.size} for entry in listing.entries ], "pruned": list(listing.pruned), "truncated": listing.truncated, } if operation == "search": needle = payload.get("needle") if not isinstance(needle, str) or not needle: raise RemoteDispatchError("invalid_payload", "search needs a needle") found = await tree.search( needle, relative, encoding=str(payload.get("encoding", "utf-8")), max_results=int(payload.get("max_results", 100)), match=payload.get("match") or None, ) return { "matches": [ {"path": str(match.path), "relative": match.relative, "line": match.line, "column": match.column, "text": match.text} for match in found.matches ], "truncated": found.truncated, "files_searched": found.files_searched, "files_unreadable": found.files_unreadable, "pruned": list(found.pruned), } if operation == "delete": await tree.delete(relative, force=bool(payload.get("force", False)), recursive=bool(payload.get("recursive", False))) return None if operation == "run": # A local `run` asks the *host* to own the process group. Here the target's own # dispatcher is that owner, so this spawns directly instead of round-tripping to # the caller's host. return await self._workspace_run(payload) raise RemoteDispatchError("unsupported", f"unknown workspace operation: {operation}")
def _tree_cwd(self, value: Any) -> pathlib.Path: if not isinstance(value, str): raise RemoteDispatchError("invalid_payload", "cwd must be a string") candidate = pathlib.Path(value) if candidate.is_absolute(): raise RemoteDispatchError("invalid_payload", "cwd must be relative to the target tree") assert self.workspace is not None resolved = (self.workspace / candidate).resolve(strict=False) try: resolved.relative_to(self.workspace) except ValueError as error: raise RemoteDispatchError("path_escapes_workspace", "cwd escapes the target tree") from error if not resolved.is_dir(): raise RemoteDispatchError("no_such_directory", f"cwd is not a directory: {value}") return resolved
async def _workspace_run(self, payload: Mapping[str, Any]) -> dict[str, Any]: argv = payload.get("argv") if not isinstance(argv, list) or not argv or not all(isinstance(item, str) and item for item in argv): raise RemoteDispatchError("invalid_payload", "argv must be a non-empty list of non-empty strings") timeout = payload.get("timeout") if timeout is not None and (type(timeout) not in {int, float} or not math.isfinite(timeout) or timeout <= 0): raise RemoteDispatchError("invalid_payload", "timeout must be positive and finite") cwd = self._tree_cwd(payload.get("cwd", ".")) outcome = await _spawn(argv=argv, cwd=str(cwd), timeout=timeout) return { "argv": list(argv), "cwd": str(cwd), "exit_code": outcome.exit_code, "stdout": outcome.stdout.data, "stderr": outcome.stderr.data, "duration": outcome.duration, "completed": not outcome.timed_out, "timed_out": outcome.timed_out, "cancelled": False, "stdout_truncated": outcome.stdout.exceeded, "stderr_truncated": outcome.stderr.exceeded, }
async def _shell(self, payload: Mapping[str, Any]) -> dict[str, Any]: if self.workspace is None: raise RemoteDispatchError("workspace_unbound", "shell operations require a checked workspace") command = _text(payload, "command", maximum=262_144) assert command is not None stdin = payload.get("stdin") if stdin is not None and not isinstance(stdin, (str, bytes, bytearray)): raise RemoteDispatchError("invalid_payload", "stdin must be text or bytes") raw_stdin = stdin.encode() if isinstance(stdin, str) else (bytes(stdin) if stdin is not None else None) outcome = await _spawn(shell=command, cwd=str(self._cwd), stdin=raw_stdin) if outcome.stdout.exceeded or outcome.stderr.exceeded: raise RemoteDispatchError("output_too_large", f"shell output exceeded {MAX_ARTIFACT_CHANNEL_BYTES} bytes per channel") return { "exit_code": outcome.exit_code, "stdout": outcome.stdout.data.decode("utf-8", "replace"), "stderr": outcome.stderr.data.decode("utf-8", "replace"), "duration": outcome.duration, }
def _decode_b64(payload: Mapping[str, Any], key: str) -> bytes: value = payload.get(key) if not isinstance(value, str): raise RemoteDispatchError("invalid_payload", f"{key} must be base64 text") try: return base64.b64decode(value, validate=True) except ValueError as error: raise RemoteDispatchError("invalid_payload", f"{key} is not valid base64") from error
async def _read_capped(stream: asyncio.StreamReader | None, process: asyncio.subprocess.Process) -> _CapturedStream: if stream is None: return _CapturedStream(b"", False) data = bytearray() exceeded = False while True: chunk = await stream.read(64 * 1024) if not chunk: break room = MAX_ARTIFACT_CHANNEL_BYTES - len(data) if room <= 0: if not exceeded: exceeded = True with contextlib.suppress(ProcessLookupError): process.kill() else: data.extend(chunk[:room]) if len(chunk) > room: exceeded = True with contextlib.suppress(ProcessLookupError): process.kill() return _CapturedStream(bytes(data), exceeded)
__all__ = ["ENVELOPE_HEADROOM", "MAX_ERROR_BYTES", "MAX_FRAME_BYTES", "MAX_VALUE_BYTES", "RemoteDispatchError", "RemoteDispatcher"]