"""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 only explicit operations, returns JSON-compatible bounded values, and never uses source text as a fallback for trusted closure shipment. """ from __future__ import annotations import ast import asyncio import base64 import contextlib import importlib import importlib.util import inspect import math import json import os import pathlib import shutil import sys import time import traceback from dataclasses import dataclass from typing import Any, Mapping, Sequence MAX_FRAME_BYTES = 1_048_576 ENVELOPE_HEADROOM = 8_576 MAX_VALUE_BYTES = MAX_FRAME_BYTES - ENVELOPE_HEADROOM MAX_ERROR_BYTES = 16_384 MAX_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 = "" 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"]