code/execution/python_runner.py

"""Sandboxed Python code execution.

Runs user Python code in a subprocess with resource limits,
captures stdout/stderr, and saves matplotlib figures.
"""

from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import textwrap
from dataclasses import dataclass
from pathlib import Path

from code.config.constants import (
    MAX_STDIO_CHARS,
    OUTPUT_PNG,
    PY_MEM_LIMIT_MB,
    PY_TIMEOUT_S,
)


@dataclass
class PythonExecutionResult:
    """Result of a sandboxed Python execution."""
    stdout: str
    stderr: str
    image_path: str | None
    returncode: int | None
    timed_out: bool = False


def _apply_subprocess_limits() -> None:
    """Set resource limits for the subprocess (Linux only)."""
    import resource
    mem_bytes = PY_MEM_LIMIT_MB * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes))
    resource.setrlimit(resource.RLIMIT_CPU, (PY_TIMEOUT_S, PY_TIMEOUT_S))


def _python_runner_source() -> str:
    """Return the source code of the runner script that wraps user code."""
    return textwrap.dedent(
        f"""
        import os
        import runpy
        import sys
        import traceback

        os.environ.setdefault("MPLBACKEND", "Agg")
        exit_code = 0
        try:
            runpy.run_path(os.path.join(os.getcwd(), "user_code.py"), run_name="__main__")
        except SystemExit as exc:
            code = exc.code
            exit_code = code if isinstance(code, int) else 1
        except Exception:
            traceback.print_exc()
            exit_code = 1
        finally:
            try:
                import matplotlib
                matplotlib.use("Agg", force=True)
                import matplotlib.pyplot as plt
                if plt.get_fignums():
                    plt.savefig(os.environ["OUTPUT_PNG"], bbox_inches="tight")
            except ModuleNotFoundError as exc:
                if exc.name != "matplotlib":
                    traceback.print_exc()
            except Exception:
                traceback.print_exc()

        raise SystemExit(exit_code)
        """
    ).strip()


def _truncate_output(text: str) -> str:
    """Truncate output to MAX_STDIO_CHARS with a note."""
    if len(text) <= MAX_STDIO_CHARS:
        return text
    remaining = len(text) - MAX_STDIO_CHARS
    return text[:MAX_STDIO_CHARS] + f"\n\n... truncated {remaining} characters ..."


def _decode_timeout_output(value: str | bytes | None) -> str:
    """Safely decode subprocess output from timeout exceptions."""
    if value is None:
        return ""
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    return value


def run_python(code: str) -> PythonExecutionResult:
    """Execute Python code in a sandboxed subprocess.

    Returns a PythonExecutionResult with stdout, stderr, image path, and status.
    """
    with tempfile.TemporaryDirectory(prefix="fullstack_run_") as tmp:
        workdir = Path(tmp)
        runner_path = workdir / "runner.py"
        user_path = workdir / "user_code.py"
        image_path = workdir / OUTPUT_PNG

        runner_path.write_text(_python_runner_source(), encoding="utf-8")
        user_path.write_text(code, encoding="utf-8")

        env = {
            "PATH": "/usr/bin:/bin",
            "HOME": str(workdir),
            "TMPDIR": str(workdir),
            "MPLBACKEND": "Agg",
            "MPLCONFIGDIR": str(workdir / ".matplotlib"),
            "OUTPUT_PNG": str(image_path),
            "PYTHONIOENCODING": "utf-8",
            "PYTHONNOUSERSITE": "1",
            "PYTHONUNBUFFERED": "1",
            "LANG": "C.UTF-8",
            "OPENBLAS_NUM_THREADS": "1",
            "OMP_NUM_THREADS": "1",
            "MKL_NUM_THREADS": "1",
            "NUMEXPR_NUM_THREADS": "1",
        }

        try:
            completed = subprocess.run(
                [sys.executable, "-I", str(runner_path)],
                cwd=workdir,
                env=env,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=PY_TIMEOUT_S,
                preexec_fn=_apply_subprocess_limits if sys.platform == "linux" else None,
                check=False,
            )
            stdout = _truncate_output(completed.stdout)
            stderr = _truncate_output(completed.stderr)

            if completed.returncode and not stderr:
                stderr = f"Process exited with status {completed.returncode}."

            saved_image: str | None = None
            if image_path.exists() and image_path.stat().st_size > 0:
                saved = tempfile.NamedTemporaryFile(
                    prefix="fullstack_plot_", suffix=".png", delete=False
                )
                saved.close()
                Path(saved.name).write_bytes(image_path.read_bytes())
                saved_image = saved.name

            return PythonExecutionResult(
                stdout=stdout,
                stderr=stderr,
                image_path=saved_image,
                returncode=completed.returncode,
            )
        except subprocess.TimeoutExpired as exc:
            stdout = _truncate_output(_decode_timeout_output(exc.stdout))
            stderr = _truncate_output(_decode_timeout_output(exc.stderr))
            timeout_note = f"Timed out after {PY_TIMEOUT_S} seconds; the process was killed."
            stderr = f"{stderr}\n{timeout_note}".strip()
            return PythonExecutionResult(
                stdout=stdout,
                stderr=stderr,
                image_path=None,
                returncode=None,
                timed_out=True,
            )