NexusOS/synapse/code_run.py

"""Run a short code snippet in an ephemeral working directory.

This is the *second* track of the code-preview feature and deliberately not the
first. `render_preview` (synapse/tools.py) executes nothing server-side: it
validates markup and the chat UI renders it inside an opaque-origin iframe.
That model fits HTML/SVG/JSX and cannot fit C, Rust or Erlang, which need a real
toolchain. So those go through here instead, and the result is shown as terminal
output rather than a rendered document.

WHAT THIS IS NOT
----------------
Not a security sandbox. Snippets run as the current user on the host. What this
module actually provides is *containment by limits*, layered:

  1. Consent      run_snippet is an ACTION tool, so it is withheld entirely
                  unless `action_tool_policy` is "ask"/"allow" — and on "ask"
                  every call waits for the user's Approve/Deny in chat.
  2. Screening    `critique` rejects the obvious-abuse shapes (sockets, process
                  spawning, absolute paths) before anything is written to disk.
  3. Isolation    cwd is a fresh temp dir that is deleted afterwards; home and
                  temp environment variables point at it; the environment is
                  scrubbed to a small allowlist.
  4. Limits       wall-clock timeout, RLIMIT_CPU/AS/FSIZE/NPROC on POSIX,
                  truncated output.
  5. Network      on Linux, `unshare -rn` when unprivileged user namespaces are
                  available (probed once, see `_net_isolation`). Nowhere else —
                  macOS and Windows get no network isolation at all.

Layer 2 is a tripwire, not a boundary: arbitrary Python can evade any regex
trivially. It exists to catch a model that reaches for `requests` by habit, not
an adversary. The load-bearing layers are 1 and 3–5. Anyone wanting a real
boundary should run this under a container or a VM; that is a deployment
decision this module does not make for them.
"""
from __future__ import annotations

import os
import re
import shutil
import subprocess
import sys
import tempfile
from functools import lru_cache
from pathlib import Path

# Wall clock. Compilation gets its own, larger budget: rustc on a cold cache
# routinely spends longer than any snippet is allowed to *run*, and killing a
# compile at the run timeout would look like the code hung when it never started.
RUN_TIMEOUT = 5.0
COMPILE_TIMEOUT = 25.0

MAX_SOURCE = 100_000      # chars of source accepted
MAX_OUTPUT = 20_000       # chars of stdout/stderr returned, per stream
MAX_STDIN = 10_000

# Applied to the run step only. A compiler legitimately needs more address space
# than a snippet does and forks a linker, so imposing these on the compile step
# breaks the toolchain rather than containing the snippet.
_MEM_BYTES = 512 * 1024 * 1024
_MAX_PROCS = 64
_MAX_FILE_BYTES = 8 * 1024 * 1024

# RLIMIT_NPROC is counted per real UID, not per run. An absolute "512" is
# therefore 512 minus whatever the user already has — enough on a quiet host,
# zero on a busy macOS desktop. Erlang needs relative headroom for BEAM's
# boot process tree; other languages keep the absolute _MAX_PROCS (blocking
# fork is the intent there).
_ERLANG_PROC_HEADROOM = 256


# ---------------------------------------------------------------------------
# Static screening
# ---------------------------------------------------------------------------

def _shared_issues(source: str) -> list[str]:
    stripped = source.strip()
    if len(stripped) < 8:
        return ["source is empty or too short to run."]
    issues: list[str] = []
    # Only flag URLs inside string literals. A citation in a comment
    # (`/* see https://… */`) is not a network reach and bouncing it costs a
    # useless retry round — this layer is a tripwire, not a parser.
    if re.search(r"""['"][^'"]*https?://[^'"]*['"]""", source):
        issues.append(
            "source contains an http(s) URL string — the runner has no network "
            "access. Inline the data you need."
        )
    if re.search(r"\b(TODO|FIXME|your code here|implement this)\b", source, re.I):
        issues.append("source still contains a placeholder — send the finished code.")
    # An absolute path is either reaching outside the ephemeral dir or is a
    # machine-specific guess that will not exist. Relative paths are fine: cwd
    # is the temp dir and goes away with it.
    # Windows: match `C:\…` and `C:\\…` (raw / escaped). Unix: common roots.
    if re.search(
        r"""['"](?:/(?:etc|home|root|usr|var|proc|sys)/|[A-Za-z]:[/\\])""",
        source,
    ):
        issues.append(
            "source references an absolute filesystem path. The snippet runs in a "
            "throwaway directory — use relative paths, or inline the data."
        )
    return issues


def _deny(source: str, rules: list[tuple[str, str]]) -> list[str]:
    return [msg for pattern, msg in rules if re.search(pattern, source)]


_NO_NET = "the runner has no network access"
_NO_SPAWN = "the runner does not allow spawning other processes"

_PY_RULES = [
    (r"\b(?:import|from)\s+(?:socket|ssl|ftplib|smtplib|telnetlib|urllib|http)\b",
     f"networking module imported — {_NO_NET}."),
    (r"\bimport\s+(?:requests|httpx|aiohttp|urllib3)\b",
     f"HTTP client imported — {_NO_NET}."),
    (r"\b(?:import|from)\s+(?:subprocess|multiprocessing)\b",
     f"subprocess/multiprocessing imported — {_NO_SPAWN}."),
    (r"\bos\.(?:system|popen|exec[lv]|fork|spawn|kill)\b",
     f"os process call — {_NO_SPAWN}."),
    (r"\b(?:import|from)\s+ctypes\b",
     "ctypes imported — the runner does not allow native calls."),
]

_C_RULES = [
    (r"#\s*include\s*<(?:sys/socket|netinet/|arpa/|netdb)",
     f"socket header included — {_NO_NET}."),
    (r"\b(?:system|popen|fork|execv?[lpe]*)\s*\(",
     f"process call — {_NO_SPAWN}."),
]

_RUST_RULES = [
    (r"\bstd::net\b", f"std::net used — {_NO_NET}."),
    (r"\bstd::process::(?:Command|abort)\b", f"std::process::Command used — {_NO_SPAWN}."),
]

_ERL_RULES = [
    (r"\b(?:gen_tcp|gen_udp|httpc|inets|ssl)\b", f"networking module used — {_NO_NET}."),
    (r"\bos:cmd\b", f"os:cmd/1 used — {_NO_SPAWN}."),
]


def _entry_point(pattern: str, message: str):
    """Most of these languages fail with a linker/loader error rather than a
    useful one when the entry point is missing, so name it up front."""
    def check(source: str) -> list[str]:
        return [] if re.search(pattern, source) else [message]
    return check


def _critique(rules: list[tuple[str, str]], entry=None):
    def check(source: str) -> list[str]:
        issues = _shared_issues(source)
        if issues and issues[0].startswith("source is empty"):
            return issues  # nothing else is worth saying about an empty body
        issues += _deny(source, rules)
        if entry:
            issues += entry(source)
        return issues
    return check


# ---------------------------------------------------------------------------
# Drivers
# ---------------------------------------------------------------------------

def _source_name(default: str):
    return lambda _source: default


def _erlang_source_name(source: str) -> str:
    """escript reads the same file two different ways, chosen by extension: a
    `.erl` file is compiled as a module (needs -module/-export), anything else
    is a plain script (needs only main/1). Models write both, so let the source
    pick its own filename instead of forcing one dialect."""
    return "main.erl" if re.search(r"^\s*-module\s*\(", source, re.M) else "main.escript"


def _erlang_write_source(source: str) -> str:
    """OTP 28+ escript rejects a shebang-less `.escript` with 'Premature end of
    file'. Module-form `.erl` files do not need one. Inject only when missing so
    a model that already wrote `#!/usr/bin/env escript` is left alone."""
    if _erlang_source_name(source).endswith(".escript"):
        stripped = source.lstrip()
        if not stripped.startswith("#!"):
            return "#!/usr/bin/env escript\n" + source
    return source


@lru_cache(maxsize=None)
def _compiled_tool(lang: str, candidates: tuple[str, ...]) -> str | None:
    """Resolve a compiler, verifying the complete Windows toolchain once.

    A compiler executable alone is not a usable toolchain on Windows: rustc's
    MSVC target also needs Microsoft's linker, and an MSYS2 driver can remain on
    PATH after one of its runtime DLLs has broken. Both cases otherwise make the
    capability monitor say "ready" and turn every snippet into a compile error.
    POSIX keeps the cheap historical which(1) check; the release hosts there
    install compiler packages atomically.
    """
    found = [tool for name in candidates if (tool := shutil.which(name))]
    if sys.platform != "win32":
        return found[0] if found else None
    for tool in found:
        if _probe_compiled_tool(lang, tool):
            return tool
    return None


def _probe_compiled_tool(lang: str, tool: str) -> bool:
    """Compile a minimal known-good program with the runner's real child env."""
    source = {
        "c": "int main(void){return 0;}",
        "cpp": "int main(){return 0;}",
        "rust": "fn main() {}",
    }[lang]
    try:
        with tempfile.TemporaryDirectory(prefix="nexus-toolchain-") as tmp:
            workdir = Path(tmp)
            entry = RUN_LANGS[lang]
            src = workdir / entry["source_name"](source)
            src.write_text(source, encoding="utf-8")
            exe = str(workdir / "probe.exe")
            built = _spawn(
                entry["compile"](tool, str(src), exe),
                workdir,
                _child_env(lang, workdir),
                COMPILE_TIMEOUT,
                constrain_memory=False,
            )
            return built.returncode == 0 and Path(exe).is_file()
    except (OSError, subprocess.SubprocessError):
        return False


# The one place that says which languages can be executed. Each entry owns that
# language's screening, toolchain probe and argv. The tool schema's `lang` enum,
# the capability line in the system prompt and the dispatch below are all derived
# from these keys rather than repeating them.
#
# The frontend keeps a matching registry (RUN_LANGS in
# interface/web/src/preview/run-langs.js) because the two sides need different
# things per language — this side executes, that side labels and displays — and
# neither should depend on the other at runtime. tests/test_tools.py asserts the
# key sets stay equal, so drift fails the check gate instead of silently showing
# a run result with the wrong language on it.
RUN_LANGS: dict[str, dict] = {
    "python": {
        "summary": "Python script (stdlib only)",
        "tool": lambda: sys.executable,
        "install": "Python is bundled with NexusOS; this should not happen.",
        "source_name": _source_name("main.py"),
        # -I is isolated mode: ignores PYTHON* env vars, the user site-packages
        # dir and the script's own directory on sys.path.
        "compile": None,
        "run": lambda tool, src, _exe: [tool, "-I", src],
        "critique": _critique(_PY_RULES),
    },
    "c": {
        "summary": "single-file C program (C11, libm linked)",
        "tool": lambda: _compiled_tool("c", ("cc", "gcc", "clang")),
        "install": "install a C compiler (clang or gcc)",
        "source_name": _source_name("main.c"),
        "compile": lambda cc, src, exe: [cc, "-std=c11", "-O0", "-Wall", "-o", exe, src, "-lm"],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_C_RULES, _entry_point(
            r"\bmain\s*\(", "no main() — a C program needs `int main(void)`.")),
    },
    "cpp": {
        "summary": "single-file C++ program (C++17)",
        "tool": lambda: _compiled_tool("cpp", ("c++", "g++", "clang++")),
        "install": "install a C++ compiler (clang++ or g++)",
        "source_name": _source_name("main.cpp"),
        "compile": lambda cc, src, exe: [cc, "-std=c++17", "-O0", "-Wall", "-o", exe, src],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_C_RULES, _entry_point(
            r"\bmain\s*\(", "no main() — a C++ program needs `int main()`.")),
    },
    "rust": {
        "summary": "single-file Rust program (2021 edition, std only)",
        "tool": lambda: _compiled_tool("rust", ("rustc",)),
        "install": "install Rust (https://rustup.rs)",
        "source_name": _source_name("main.rs"),
        # Debug build: -O roughly triples compile time for snippets that run for
        # milliseconds either way.
        "compile": lambda cc, src, exe: [cc, "--edition", "2021", "-o", exe, src],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_RUST_RULES, _entry_point(
            r"\bfn\s+main\s*\(", "no main() — a Rust program needs `fn main()`.")),
    },
    "erlang": {
        "summary": "escript program with a main/1 entry point",
        "tool": lambda: shutil.which("escript"),
        "install": "install Erlang/OTP (provides escript)",
        "source_name": _erlang_source_name,
        "compile": None,
        "run": lambda tool, src, _exe: [tool, src],
        "critique": _critique(_ERL_RULES, _entry_point(
            r"\bmain\s*\(", "no main/1 — escript calls `main(Args)`.")),
    },
}

# Spellings a model reaches for that aren't the registry key. Resolved before
# lookup so `c++` and `py` work without doubling the tool schema's enum.
ALIASES = {
    "c++": "cpp", "cc": "c", "py": "python", "python3": "python",
    "rs": "rust", "erl": "erlang", "escript": "erlang",
}


def resolve_lang(lang: str) -> str:
    key = (lang or "").strip().lower()
    return ALIASES.get(key, key)


def lang_prose() -> str:
    """'python, c or rust' — the runnable languages as a phrase for prompts."""
    names = list(RUN_LANGS)
    if len(names) < 2:
        return names[0] if names else ""
    return f"{', '.join(names[:-1])} or {names[-1]}"


# ---------------------------------------------------------------------------
# Execution
# ---------------------------------------------------------------------------

# Environment variables the child keeps. Everything else is dropped: the snippet
# has no business seeing API keys, proxy settings or the user's shell config, and
# PYTHON*/LD_* in particular would let ambient config change how it runs.
_ENV_KEEP = ("PATH", "LANG", "LC_ALL", "TERM", "SYSTEMROOT", "COMSPEC")

# rustc is usually a rustup shim, and a shim with HOME rewritten cannot find its
# own toolchain. Passing these through is what makes Rust work at all here; they
# point at read-only toolchain data, not at anything the snippet should write.
_ENV_KEEP_BY_LANG = {"rust": ("RUSTUP_HOME", "CARGO_HOME", "RUSTUP_TOOLCHAIN")}


def _child_env(lang: str, workdir: Path) -> dict:
    env = {k: os.environ[k] for k in _ENV_KEEP if k in os.environ}
    for k in _ENV_KEEP_BY_LANG.get(lang, ()):
        if k in os.environ:
            env[k] = os.environ[k]
    if lang == "rust" and "RUSTUP_HOME" not in env:
        # HOME is about to be rewritten, so resolve rustup's default location
        # against the real home while we still know it.
        default = Path.home() / ".rustup"
        if default.is_dir():
            env["RUSTUP_HOME"] = str(default)
            env.setdefault("CARGO_HOME", str(Path.home() / ".cargo"))
    scratch = str(workdir)
    env["HOME"] = scratch
    env["TMPDIR"] = scratch
    # Windows ignores HOME/TMPDIR in its standard path helpers. Without these,
    # expanduser() reaches the real profile and GetTempPath() falls back to the
    # Windows directory; GCC and rustc then either escape the scratch directory
    # or fail because a normal user cannot write there.
    env["TEMP"] = scratch
    env["TMP"] = scratch
    if os.name == "nt":
        env["USERPROFILE"] = scratch
    env.setdefault("LC_ALL", "C.UTF-8")
    return env


def _user_process_count() -> int | None:
    """How many processes this UID already owns. None when we cannot count
    (Windows, or psutil missing) — callers must not invent an absolute cap."""
    if os.name == "nt" or not hasattr(os, "getuid"):
        return None
    try:
        import psutil  # optional; process extra
    except ImportError:
        return None
    uid = os.getuid()
    n = 0
    for proc in psutil.process_iter(["uids"]):
        try:
            uids = proc.info.get("uids")
            if uids is not None and getattr(uids, "real", None) == uid:
                n += 1
        except (psutil.Error, TypeError, AttributeError):
            continue
    return n


def _max_procs_for(lang: str) -> int | None:
    """Soft RLIMIT_NPROC for this run, or None to leave the limit unset.

    Erlang: current per-UID count + headroom (RLIMIT_NPROC is UID-scoped).
    Everything else: the absolute _MAX_PROCS tripwire against fork bombs.
    """
    if lang == "erlang":
        current = _user_process_count()
        if current is None:
            return None
        return current + _ERLANG_PROC_HEADROOM
    return _MAX_PROCS


def _limits(constrain_memory: bool, max_procs: int | None = None):
    """preexec_fn applying POSIX rlimits, or None where they don't exist.

    RLIMIT_CPU is a backstop for the wall-clock timeout: a snippet that ignores
    SIGTERM still loses the CPU. The memory and process caps are skipped for
    compilation — see _MEM_BYTES.

    `max_procs=None` with constrain_memory means "do not set RLIMIT_NPROC"
    (used when we cannot compute a relative Erlang ceiling). Passing an int
    always sets it.
    """
    if sys.platform == "win32":
        return None
    try:
        import resource
    except ImportError:  # pragma: no cover - POSIX only
        return None

    cpu = int(COMPILE_TIMEOUT if not constrain_memory else RUN_TIMEOUT) + 1
    wanted = [("RLIMIT_CPU", cpu), ("RLIMIT_FSIZE", _MAX_FILE_BYTES), ("RLIMIT_CORE", 0)]
    if constrain_memory:
        wanted.append(("RLIMIT_AS", _MEM_BYTES))
        # Distinguish "caller omitted" (use default) from "explicitly skip"
        # by requiring the kw to be passed — see _spawn.
        if max_procs is not None:
            wanted.append(("RLIMIT_NPROC", max_procs))

    def apply():  # runs in the forked child, between fork and exec
        # Every limit is set independently and failure is swallowed. Which of
        # these exist, and which can be lowered, varies by platform (macOS has
        # no usable RLIMIT_AS, RLIMIT_NPROC is absent on some POSIX systems) —
        # and an exception raised here does not "skip a limit", it aborts the
        # spawn entirely. Partial limits are the right failure mode; no run at
        # all is not.
        for name, soft in wanted:
            which = getattr(resource, name, None)
            if which is None:
                continue
            try:
                _, hard = resource.getrlimit(which)
                if hard != resource.RLIM_INFINITY:
                    soft = min(soft, hard)
                resource.setrlimit(which, (soft, hard))
            except (ValueError, OSError):
                continue

    return apply


_net_isolation_cache: list | None = None


def _net_isolation() -> list[str]:
    """argv prefix that drops the child into an empty network namespace, or [].

    Linux only, and only where unprivileged user namespaces are enabled — which
    is a kernel/distro setting we can't change and shouldn't fail over. Probed
    once and cached; an empty list means the run simply has host networking, and
    callers must not treat this as a guarantee either way.
    """
    global _net_isolation_cache
    if _net_isolation_cache is not None:
        return _net_isolation_cache
    _net_isolation_cache = []
    if sys.platform.startswith("linux") and shutil.which("unshare"):
        try:
            probe = subprocess.run(
                ["unshare", "-rn", "true"],
                capture_output=True, timeout=5,
            )
            if probe.returncode == 0:
                _net_isolation_cache = ["unshare", "-rn"]
        except (OSError, subprocess.SubprocessError):
            pass
    return _net_isolation_cache


def _clip(raw: bytes) -> str:
    text = raw.decode("utf-8", errors="replace")
    if len(text) <= MAX_OUTPUT:
        return text
    return text[:MAX_OUTPUT] + f"\n... [truncated at {MAX_OUTPUT} characters]"


def _spawn(argv: list[str], workdir: Path, env: dict, timeout: float,
           stdin: str = "", constrain_memory: bool = True,
           max_procs: int | None = _MAX_PROCS):
    """Spawn a child. `max_procs` defaults to `_MAX_PROCS`; pass `None` to skip
    RLIMIT_NPROC entirely (Erlang, when a relative ceiling cannot be computed)."""
    return subprocess.run(
        argv,
        cwd=str(workdir),
        env=env,
        input=stdin.encode("utf-8"),
        capture_output=True,
        timeout=timeout,
        preexec_fn=_limits(constrain_memory, max_procs=max_procs),
    )


def run(lang: str, source: str, stdin: str = "") -> dict:
    """Compile (if needed) and run `source`. Blocking — call from a thread.

    Returns {ok, lang, stage, exit_code, stdout, stderr, error}. `ok` is False
    only when the snippet could not be run at all (missing toolchain, compile
    error, timeout); a program that runs and exits non-zero is a successful run
    with a non-zero exit_code, because its stderr is the answer the user wants.
    """
    key = resolve_lang(lang)
    entry = RUN_LANGS.get(key)
    if entry is None:
        return {"ok": False, "lang": lang, "stage": "lang",
                "error": f"cannot run {lang!r} — use {lang_prose()}."}

    tool = entry["tool"]()
    if not tool:
        return {"ok": False, "lang": key, "stage": "toolchain",
                "error": f"no toolchain for {key} on this machine — {entry['install']}. "
                         "Show the code instead of running it."}

    with tempfile.TemporaryDirectory(prefix="nexus-run-") as tmp:
        workdir = Path(tmp)
        body = _erlang_write_source(source) if key == "erlang" else source
        src = workdir / entry["source_name"](source)
        src.write_text(body, encoding="utf-8")
        exe = str(workdir / ("program.exe" if sys.platform == "win32" else "program"))
        env = _child_env(key, workdir)
        max_procs = _max_procs_for(key)

        if entry["compile"]:
            try:
                built = _spawn(entry["compile"](tool, str(src), exe), workdir, env,
                               COMPILE_TIMEOUT, constrain_memory=False)
            except subprocess.TimeoutExpired:
                return {"ok": False, "lang": key, "stage": "compile",
                        "error": f"compilation timed out after {COMPILE_TIMEOUT:g}s."}
            except OSError as e:
                return {"ok": False, "lang": key, "stage": "compile",
                        "error": f"could not start the compiler: {e}"}
            if built.returncode != 0:
                return {"ok": False, "lang": key, "stage": "compile",
                        "exit_code": built.returncode,
                        "stdout": _clip(built.stdout), "stderr": _clip(built.stderr),
                        "error": "compilation failed — read stderr, fix the source, "
                                 "and call run_snippet again."}

        argv = _net_isolation() + entry["run"](tool, str(src), exe)
        try:
            done = _spawn(argv, workdir, env, RUN_TIMEOUT, stdin=stdin[:MAX_STDIN],
                          max_procs=max_procs)
        except subprocess.TimeoutExpired as e:
            return {"ok": False, "lang": key, "stage": "run",
                    "stdout": _clip(e.stdout or b""), "stderr": _clip(e.stderr or b""),
                    "error": f"the program did not finish within {RUN_TIMEOUT:g}s — "
                             "it is probably looping. Bound the work and try again."}
        except OSError as e:
            return {"ok": False, "lang": key, "stage": "run",
                    "error": f"could not start the program: {e}"}

    stdout = _clip(done.stdout)
    stderr = _clip(done.stderr)
    # BEAM failing to fork under RLIMIT_NPROC looks like a snippet bug if we
    # report ok=True. Call it out as a runner limit so the model does not keep
    # rewriting a correct program.
    if (
        key == "erlang"
        and done.returncode != 0
        and "Resource temporarily unavailable" in stderr
    ):
        return {
            "ok": False,
            "lang": key,
            "stage": "run",
            "exit_code": done.returncode,
            "stdout": stdout,
            "stderr": stderr,
            "error": (
                "Erlang could not start under the process limit (host already has "
                "many processes for this user). Free some processes and retry, or "
                "show the code instead of running it."
            ),
        }

    return {
        "ok": True,
        "lang": key,
        "stage": "run",
        "exit_code": done.returncode,
        "stdout": stdout,
        "stderr": stderr,
    }


def critique(lang: str, source: str) -> list[str]:
    """Static screening for one snippet. See the module docstring on what this
    is worth: a tripwire against habitual network/process reaches, not a
    boundary."""
    key = resolve_lang(lang)
    entry = RUN_LANGS.get(key)
    if entry is None:
        return [f"cannot run {lang!r} — use {lang_prose()}."]
    if len(source) > MAX_SOURCE:
        return [f"source is over {MAX_SOURCE} characters — send a single focused snippet."]
    return entry["critique"](source)