Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
_execution.py386 linesDownload Raw Back to middleware
1"""Execution policies for the persistent shell middleware."""2 3from __future__ import annotations4 5import abc6import json7import os8import shutil9import subprocess10import sys11import typing12from collections.abc import Mapping, Sequence13from dataclasses import dataclass, field14from pathlib import Path15 16try:  # pragma: no cover - optional dependency on POSIX platforms17    import resource18 19    _HAS_RESOURCE = True20except ImportError:  # pragma: no cover - non-POSIX systems21    _HAS_RESOURCE = False22 23 24SHELL_TEMP_PREFIX = "langchain-shell-"25 26 27def _launch_subprocess(28    command: Sequence[str],29    *,30    env: Mapping[str, str],31    cwd: Path,32    preexec_fn: typing.Callable[[], None] | None,33    start_new_session: bool,34) -> subprocess.Popen[str]:35    return subprocess.Popen(  # noqa: S60336        list(command),37        stdin=subprocess.PIPE,38        stdout=subprocess.PIPE,39        stderr=subprocess.PIPE,40        cwd=cwd,41        text=True,42        encoding="utf-8",43        errors="replace",44        bufsize=1,45        env=env,46        preexec_fn=preexec_fn,  # noqa: PLW150947        start_new_session=start_new_session,48    )49 50 51if typing.TYPE_CHECKING:52    from collections.abc import Mapping, Sequence53    from pathlib import Path54 55 56@dataclass57class BaseExecutionPolicy(abc.ABC):58    """Configuration contract for persistent shell sessions.59 60    Concrete subclasses encapsulate how a shell process is launched and constrained.61 62    Each policy documents its security guarantees and the operating environments in63    which it is appropriate. Use `HostExecutionPolicy` for trusted, same-host execution;64    `CodexSandboxExecutionPolicy` when the Codex CLI sandbox is available and you want65    additional syscall restrictions; and `DockerExecutionPolicy` for container-level66    isolation using Docker.67    """68 69    command_timeout: float = 30.070    startup_timeout: float = 30.071    termination_timeout: float = 10.072    max_output_lines: int = 10073    max_output_bytes: int | None = None74 75    def __post_init__(self) -> None:76        if self.max_output_lines <= 0:77            msg = "max_output_lines must be positive."78            raise ValueError(msg)79 80    @abc.abstractmethod81    def spawn(82        self,83        *,84        workspace: Path,85        env: Mapping[str, str],86        command: Sequence[str],87    ) -> subprocess.Popen[str]:88        """Launch the persistent shell process."""89 90 91@dataclass92class HostExecutionPolicy(BaseExecutionPolicy):93    """Run the shell directly on the host process.94 95    This policy is best suited for trusted or single-tenant environments (CI jobs,96    developer workstations, pre-sandboxed containers) where the agent must access the97    host filesystem and tooling without additional isolation. Enforces optional CPU and98    memory limits to prevent runaway commands but offers **no** filesystem or network99    sandboxing; commands can modify anything the process user can reach.100 101    On Linux platforms resource limits are applied with `resource.prlimit` after the102    shell starts. On macOS, where `prlimit` is unavailable, limits are set in a103    `preexec_fn` before `exec`. In both cases the shell runs in its own process group104    so timeouts can terminate the full subtree.105    """106 107    cpu_time_seconds: int | None = None108    memory_bytes: int | None = None109    create_process_group: bool = True110 111    _limits_requested: bool = field(init=False, repr=False, default=False)112 113    def __post_init__(self) -> None:114        super().__post_init__()115        if self.cpu_time_seconds is not None and self.cpu_time_seconds <= 0:116            msg = "cpu_time_seconds must be positive if provided."117            raise ValueError(msg)118        if self.memory_bytes is not None and self.memory_bytes <= 0:119            msg = "memory_bytes must be positive if provided."120            raise ValueError(msg)121        self._limits_requested = any(122            value is not None for value in (self.cpu_time_seconds, self.memory_bytes)123        )124        if self._limits_requested and not _HAS_RESOURCE:125            msg = (126                "HostExecutionPolicy cpu/memory limits require the Python 'resource' module. "127                "Either remove the limits or run on a POSIX platform."128            )129            raise RuntimeError(msg)130 131    def spawn(132        self,133        *,134        workspace: Path,135        env: Mapping[str, str],136        command: Sequence[str],137    ) -> subprocess.Popen[str]:138        process = _launch_subprocess(139            list(command),140            env=env,141            cwd=workspace,142            preexec_fn=self._create_preexec_fn(),143            start_new_session=self.create_process_group,144        )145        self._apply_post_spawn_limits(process)146        return process147 148    def _create_preexec_fn(self) -> typing.Callable[[], None] | None:149        if not self._limits_requested or self._can_use_prlimit():150            return None151 152        def _configure() -> None:  # pragma: no cover - depends on OS153            if self.cpu_time_seconds is not None:154                limit = (self.cpu_time_seconds, self.cpu_time_seconds)155                resource.setrlimit(resource.RLIMIT_CPU, limit)156            if self.memory_bytes is not None:157                limit = (self.memory_bytes, self.memory_bytes)158                if hasattr(resource, "RLIMIT_AS"):159                    resource.setrlimit(resource.RLIMIT_AS, limit)160                elif hasattr(resource, "RLIMIT_DATA"):161                    resource.setrlimit(resource.RLIMIT_DATA, limit)162 163        return _configure164 165    def _apply_post_spawn_limits(self, process: subprocess.Popen[str]) -> None:166        if not self._limits_requested or not self._can_use_prlimit():167            return168        if not _HAS_RESOURCE:  # pragma: no cover - defensive169            return170        pid = process.pid171        try:172            prlimit = typing.cast("typing.Any", resource).prlimit173            if self.cpu_time_seconds is not None:174                prlimit(pid, resource.RLIMIT_CPU, (self.cpu_time_seconds, self.cpu_time_seconds))175            if self.memory_bytes is not None:176                limit = (self.memory_bytes, self.memory_bytes)177                if hasattr(resource, "RLIMIT_AS"):178                    prlimit(pid, resource.RLIMIT_AS, limit)179                elif hasattr(resource, "RLIMIT_DATA"):180                    prlimit(pid, resource.RLIMIT_DATA, limit)181        except OSError as exc:  # pragma: no cover - depends on platform support182            msg = "Failed to apply resource limits via prlimit."183            raise RuntimeError(msg) from exc184 185    @staticmethod186    def _can_use_prlimit() -> bool:187        return _HAS_RESOURCE and hasattr(resource, "prlimit") and sys.platform.startswith("linux")188 189 190@dataclass191class CodexSandboxExecutionPolicy(BaseExecutionPolicy):192    """Launch the shell through the Codex CLI sandbox.193 194    Ideal when you have the Codex CLI installed and want the additional syscall and195    filesystem restrictions provided by Anthropic's Seatbelt (macOS) or Landlock/seccomp196    (Linux) profiles. Commands still run on the host, but within the sandbox requested by197    the CLI. If the Codex binary is unavailable or the runtime lacks the required198    kernel features (e.g., Landlock inside some containers), process startup fails with a199    `RuntimeError`.200 201    Configure sandbox behavior via `config_overrides` to align with your Codex CLI202    profile. This policy does not add its own resource limits; combine it with203    host-level guards (cgroups, container resource limits) as needed.204    """205 206    binary: str = "codex"207    platform: typing.Literal["auto", "macos", "linux"] = "auto"208    config_overrides: Mapping[str, typing.Any] = field(default_factory=dict)209 210    def spawn(211        self,212        *,213        workspace: Path,214        env: Mapping[str, str],215        command: Sequence[str],216    ) -> subprocess.Popen[str]:217        full_command = self._build_command(command)218        return _launch_subprocess(219            full_command,220            env=env,221            cwd=workspace,222            preexec_fn=None,223            start_new_session=False,224        )225 226    def _build_command(self, command: Sequence[str]) -> list[str]:227        binary = self._resolve_binary()228        platform_arg = self._determine_platform()229        full_command: list[str] = [binary, "sandbox", platform_arg]230        for key, value in sorted(dict(self.config_overrides).items()):231            full_command.extend(["-c", f"{key}={self._format_override(value)}"])232        full_command.append("--")233        full_command.extend(command)234        return full_command235 236    def _resolve_binary(self) -> str:237        path = shutil.which(self.binary)238        if path is None:239            msg = (240                "Codex sandbox policy requires the '%s' CLI to be installed and available on PATH."241            )242            raise RuntimeError(msg % self.binary)243        return path244 245    def _determine_platform(self) -> str:246        if self.platform != "auto":247            return self.platform248        if sys.platform.startswith("linux"):249            return "linux"250        if sys.platform == "darwin":  # type: ignore[unreachable, unused-ignore]251            return "macos"252        msg = (  # type: ignore[unreachable, unused-ignore]253            "Codex sandbox policy could not determine a supported platform; "254            "set 'platform' explicitly."255        )256        raise RuntimeError(msg)257 258    @staticmethod259    def _format_override(value: typing.Any) -> str:260        try:261            return json.dumps(value)262        except TypeError:263            return str(value)264 265 266@dataclass267class DockerExecutionPolicy(BaseExecutionPolicy):268    """Run the shell inside a dedicated Docker container.269 270    Choose this policy when commands originate from untrusted users or you require271    strong isolation between sessions. By default the workspace is bind-mounted only272    when it refers to an existing non-temporary directory; ephemeral sessions run273    without a mount to minimise host exposure. The container's network namespace is274    disabled by default (`--network none`) and you can enable further hardening via275    `read_only_rootfs` and `user`.276 277    The security guarantees depend on your Docker daemon configuration. Run the agent on278    a host where Docker is locked down (rootless mode, AppArmor/SELinux, etc.) and279    review any additional volumes or capabilities passed through ``extra_run_args``. The280    default image is `python:3.12-alpine3.19`; supply a custom image if you need281    preinstalled tooling.282    """283 284    binary: str = "docker"285    image: str = "python:3.12-alpine3.19"286    remove_container_on_exit: bool = True287    network_enabled: bool = False288    extra_run_args: Sequence[str] | None = None289    memory_bytes: int | None = None290    cpu_time_seconds: typing.Any | None = None291    cpus: str | None = None292    read_only_rootfs: bool = False293    user: str | None = None294 295    def __post_init__(self) -> None:296        super().__post_init__()297        if self.memory_bytes is not None and self.memory_bytes <= 0:298            msg = "memory_bytes must be positive if provided."299            raise ValueError(msg)300        if self.cpu_time_seconds is not None:301            msg = (302                "DockerExecutionPolicy does not support cpu_time_seconds; configure CPU limits "303                "using Docker run options such as '--cpus'."304            )305            raise RuntimeError(msg)306        if self.cpus is not None and not self.cpus.strip():307            msg = "cpus must be a non-empty string when provided."308            raise ValueError(msg)309        if self.user is not None and not self.user.strip():310            msg = "user must be a non-empty string when provided."311            raise ValueError(msg)312        self.extra_run_args = tuple(self.extra_run_args or ())313 314    def spawn(315        self,316        *,317        workspace: Path,318        env: Mapping[str, str],319        command: Sequence[str],320    ) -> subprocess.Popen[str]:321        full_command = self._build_command(workspace, env, command)322        host_env = os.environ.copy()323        return _launch_subprocess(324            full_command,325            env=host_env,326            cwd=workspace,327            preexec_fn=None,328            start_new_session=False,329        )330 331    def _build_command(332        self,333        workspace: Path,334        env: Mapping[str, str],335        command: Sequence[str],336    ) -> list[str]:337        binary = self._resolve_binary()338        full_command: list[str] = [binary, "run", "-i"]339        if self.remove_container_on_exit:340            full_command.append("--rm")341        if not self.network_enabled:342            full_command.extend(["--network", "none"])343        if self.memory_bytes is not None:344            full_command.extend(["--memory", str(self.memory_bytes)])345        if self._should_mount_workspace(workspace):346            host_path = str(workspace)347            full_command.extend(["-v", f"{host_path}:{host_path}"])348            full_command.extend(["-w", host_path])349        else:350            full_command.extend(["-w", "/"])351        if self.read_only_rootfs:352            full_command.append("--read-only")353        for key, value in env.items():354            full_command.extend(["-e", f"{key}={value}"])355        if self.cpus is not None:356            full_command.extend(["--cpus", self.cpus])357        if self.user is not None:358            full_command.extend(["--user", self.user])359        if self.extra_run_args:360            full_command.extend(self.extra_run_args)361        full_command.append(self.image)362        full_command.extend(command)363        return full_command364 365    @staticmethod366    def _should_mount_workspace(workspace: Path) -> bool:367        return not workspace.name.startswith(SHELL_TEMP_PREFIX)368 369    def _resolve_binary(self) -> str:370        path = shutil.which(self.binary)371        if path is None:372            msg = (373                "Docker execution policy requires the '%s' CLI to be installed"374                " and available on PATH."375            )376            raise RuntimeError(msg % self.binary)377        return path378 379 380__all__ = [381    "BaseExecutionPolicy",382    "CodexSandboxExecutionPolicy",383    "DockerExecutionPolicy",384    "HostExecutionPolicy",385]386 
codekingpro/portable-devtools · Team Ai