src/spreadsheetbench_env/executor.py

"""Bounded subprocess execution for one SpreadsheetBench session."""

from __future__ import annotations

from dataclasses import dataclass, replace
import os
from pathlib import Path
import re
import shutil
import signal
import subprocess
import sys
import textwrap
import threading


_PATH_ASSIGNMENT = re.compile(
    r"^\s*(INPUT_PATH|OUTPUT_PATH)\s*=\s*.+$",
    re.MULTILINE,
)
_SAFE_CASE_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
_MAX_CAPTURE_CHARS = 50_000
_UID_LEASE_LOCK = threading.Lock()
_LEASED_UIDS: set[int] = set()


def _lease_uid(base: int) -> int:
    with _UID_LEASE_LOCK:
        candidate = base
        while candidate in _LEASED_UIDS:
            candidate += 1
        _LEASED_UIDS.add(candidate)
        return candidate


def _release_uid(uid: int) -> None:
    with _UID_LEASE_LOCK:
        _LEASED_UIDS.discard(uid)


@dataclass(frozen=True, slots=True)
class ExecutorLimits:
    """Wall-clock and POSIX resource limits for generated code."""

    timeout_s: float = 120.0
    cpu_seconds: int = 120
    memory_bytes: int = 2 * 1024 * 1024 * 1024
    file_bytes: int = 256 * 1024 * 1024
    process_count: int = 64
    open_files: int = 128
    run_uid: int | None = None
    run_uid_base: int | None = None
    run_gid: int | None = None


@dataclass(frozen=True, slots=True)
class ExecutionResult:
    """One generated-script execution result."""

    ok: bool
    output_path: Path | None
    stdout: str
    stderr: str
    error: str
    timed_out: bool
    returncode: int | None


def _bounded_text(value: str) -> str:
    if len(value) <= _MAX_CAPTURE_CHARS:
        return value
    half = _MAX_CAPTURE_CHARS // 2
    return value[:half] + "\n...[subprocess output truncated]...\n" + value[-half:]


def _runner_source(code: str, input_path: Path, output_path: Path) -> str:
    cleaned = _PATH_ASSIGNMENT.sub("", code)
    indented = textwrap.indent(cleaned, "    ") or "    pass"
    return (
        "import traceback\n"
        f"INPUT_PATH = {str(input_path)!r}\n"
        f"OUTPUT_PATH = {str(output_path)!r}\n"
        "try:\n"
        f"{indented}\n"
        "except Exception:\n"
        "    traceback.print_exc()\n"
        "    raise SystemExit(2)\n"
    )


def _safe_environment(case_dir: Path, scratch_dir: Path) -> dict[str, str]:
    environment = {
        "PATH": os.environ.get("PATH") or os.defpath,
        "HOME": str(case_dir),
        "TMPDIR": str(scratch_dir),
        "PYTHONUNBUFFERED": "1",
        "OPENBLAS_NUM_THREADS": "1",
        "OMP_NUM_THREADS": "1",
        "MKL_NUM_THREADS": "1",
        "NUMEXPR_NUM_THREADS": "1",
    }
    for key in ("LANG", "LC_ALL", "LC_CTYPE", "PYTHONIOENCODING", "PYTHONUTF8"):
        value = os.environ.get(key)
        if value:
            environment[key] = value
    return environment


def _launcher_source(limits: ExecutorLimits, runner_path: Path) -> str:
    """Apply limits after exec, avoiding unsafe preexec_fn work in the server."""
    return (
        "import os\n"
        "import resource\n"
        "import sys\n"
        "resource.setrlimit(resource.RLIMIT_CORE, (0, 0))\n"
        f"resource.setrlimit(resource.RLIMIT_CPU, ({limits.cpu_seconds}, {limits.cpu_seconds}))\n"
        f"memory_bytes = {limits.memory_bytes}\n"
        "if sys.platform != 'darwin':\n"
        "    resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))\n"
        f"resource.setrlimit(resource.RLIMIT_FSIZE, ({limits.file_bytes}, {limits.file_bytes}))\n"
        f"resource.setrlimit(resource.RLIMIT_NOFILE, ({limits.open_files}, {limits.open_files}))\n"
        f"process_count = {limits.process_count}\n"
        "if sys.platform != 'darwin' and hasattr(resource, 'RLIMIT_NPROC'):\n"
        "    resource.setrlimit(resource.RLIMIT_NPROC, (process_count, process_count))\n"
        f"run_gid = {limits.run_gid!r}\n"
        "if run_gid is not None:\n"
        "    os.setgroups([])\n"
        "    os.setgid(run_gid)\n"
        f"run_uid = {limits.run_uid!r}\n"
        "if run_uid is not None:\n"
        "    os.setuid(run_uid)\n"
        f"runner_path = {str(runner_path)!r}\n"
        "os.execv(sys.executable, [sys.executable, runner_path])\n"
    )


class ScriptExecutor:
    """Own the private working tree for one connection-owned episode."""

    def __init__(self, root: str | Path, *, limits: ExecutorLimits | None = None) -> None:
        self._root = Path(root).resolve()
        configured = limits or ExecutorLimits()
        if configured.run_uid is not None and configured.run_uid_base is not None:
            raise ValueError("configure either run_uid or run_uid_base, not both")
        self._leased_uid = (
            None
            if configured.run_uid_base is None
            else _lease_uid(configured.run_uid_base)
        )
        self._limits = replace(configured, run_uid=self._leased_uid or configured.run_uid)
        self._closed = False
        try:
            self._root.mkdir(parents=True, exist_ok=False, mode=0o700)
        except Exception:
            if self._leased_uid is not None:
                _release_uid(self._leased_uid)
            raise

    @property
    def root(self) -> Path:
        return self._root

    @property
    def leased_uid(self) -> int | None:
        return self._leased_uid or self._limits.run_uid

    def run(
        self,
        code: str,
        input_path: str | Path,
        *,
        case_name: str,
    ) -> ExecutionResult:
        if self._closed:
            raise RuntimeError("script executor is closed")
        if (
            not _SAFE_CASE_NAME.fullmatch(case_name)
            or case_name in {".", ".."}
            or ".." in case_name
        ):
            raise ValueError(f"unsafe case name: {case_name!r}")
        source = Path(input_path).resolve(strict=True)
        case_dir = self._root / case_name
        case_dir.mkdir(mode=0o700)
        scratch_dir = case_dir / "tmp"
        scratch_dir.mkdir(mode=0o700)
        private_input = case_dir / "input.xlsx"
        output_path = case_dir / "output.xlsx"
        runner_path = case_dir / "runner.py"
        launcher_path = case_dir / "launcher.py"
        stdout_path = case_dir / "stdout.log"
        stderr_path = case_dir / "stderr.log"
        shutil.copy2(source, private_input)
        runner_path.write_text(
            _runner_source(code, private_input, output_path),
            encoding="utf-8",
        )
        launcher_path.write_text(
            _launcher_source(self._limits, runner_path),
            encoding="utf-8",
        )

        if self._limits.run_uid is not None or self._limits.run_gid is not None:
            uid = -1 if self._limits.run_uid is None else self._limits.run_uid
            gid = -1 if self._limits.run_gid is None else self._limits.run_gid
            os.chmod(self._root, 0o711)
            for path in (case_dir, scratch_dir, private_input, runner_path, launcher_path):
                os.chown(path, uid, gid)

        with stdout_path.open("w", encoding="utf-8") as stdout_file, stderr_path.open(
            "w", encoding="utf-8"
        ) as stderr_file:
            process = subprocess.Popen(
                [sys.executable, str(launcher_path)],
                cwd=case_dir,
                env=_safe_environment(case_dir, scratch_dir),
                stdout=stdout_file,
                stderr=stderr_file,
                text=True,
                start_new_session=True,
            )
            timed_out = False
            try:
                process.wait(timeout=self._limits.timeout_s)
            except subprocess.TimeoutExpired:
                timed_out = True
                try:
                    os.killpg(process.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                process.wait()

        stdout = stdout_path.read_text(encoding="utf-8", errors="replace")
        stderr = stderr_path.read_text(encoding="utf-8", errors="replace")

        stdout = _bounded_text(stdout)
        stderr = _bounded_text(stderr)
        if timed_out:
            return ExecutionResult(
                ok=False,
                output_path=None,
                stdout=stdout,
                stderr=stderr,
                error=f"timeout after {self._limits.timeout_s:g}s",
                timed_out=True,
                returncode=process.returncode,
            )
        if process.returncode != 0:
            error = (stdout + "\n" + stderr).strip() or (
                f"generated script exited with status {process.returncode}"
            )
            return ExecutionResult(
                ok=False,
                output_path=None,
                stdout=stdout,
                stderr=stderr,
                error=error,
                timed_out=False,
                returncode=process.returncode,
            )
        if not output_path.is_file():
            return ExecutionResult(
                ok=False,
                output_path=None,
                stdout=stdout,
                stderr=stderr,
                error="output file was not created",
                timed_out=False,
                returncode=process.returncode,
            )
        return ExecutionResult(
            ok=True,
            output_path=output_path,
            stdout=stdout,
            stderr=stderr,
            error="",
            timed_out=False,
            returncode=process.returncode,
        )

    def remove_case(self, case_name: str) -> None:
        if _SAFE_CASE_NAME.fullmatch(case_name) and ".." not in case_name:
            shutil.rmtree(self._root / case_name, ignore_errors=True)

    def close(self) -> None:
        if self._closed:
            return
        shutil.rmtree(self._root, ignore_errors=True)
        if self._leased_uid is not None:
            _release_uid(self._leased_uid)
        self._closed = True