Model/tools/code_interpreter.py

"""
Frox AI — Code Interpreter Tool

Runs Python in a separate subprocess with a wall-clock timeout and
(on POSIX) memory/CPU resource limits via `resource`. This is a
reasonable baseline for a local dev repo; the production backend
architecture document specs full container isolation (gVisor runtime,
no network, tmpfs workspace — Section 4.5) for multi-tenant serving.
Don't point this at untrusted multi-tenant traffic without that layer.
"""
from __future__ import annotations

import subprocess
import sys
import tempfile
import textwrap
from pathlib import Path
from typing import Optional

from tools.registry import tool, ToolContext


DEFAULT_TIMEOUT_S = 10.0
MAX_OUTPUT_CHARS = 8000
MEMORY_LIMIT_MB = 512


def _build_preexec_fn():
    """POSIX-only resource limiting, applied inside the child process before exec."""
    if sys.platform == "win32":
        return None

    def _limit():
        import resource
        mem_bytes = MEMORY_LIMIT_MB * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes))
        resource.setrlimit(resource.RLIMIT_CPU, (int(DEFAULT_TIMEOUT_S) + 2, int(DEFAULT_TIMEOUT_S) + 2))

    return _limit


def run_python(code: str, timeout_s: float = DEFAULT_TIMEOUT_S) -> dict:
    """
    Execute Python code in an isolated subprocess.

    Isolation properties:
      - Separate process (a crash or infinite loop can't take down the caller)
      - Wall-clock timeout via subprocess.run(timeout=...)
      - POSIX: memory (RLIMIT_AS) and CPU-time (RLIMIT_CPU) limits
      - No shared filesystem state beyond a throwaway temp file for the script itself

    NOT provided here (see module docstring): network isolation,
    filesystem sandboxing beyond the OS's normal permissions, or
    protection against a determined adversary — that needs container-
    level isolation for untrusted multi-tenant use.
    """
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(code)
        script_path = f.name

    try:
        result = subprocess.run(
            [sys.executable, script_path],
            capture_output=True, text=True, timeout=timeout_s,
            preexec_fn=_build_preexec_fn(),
        )
        stdout = result.stdout[:MAX_OUTPUT_CHARS]
        stderr = result.stderr[:MAX_OUTPUT_CHARS]
        return {
            "success": result.returncode == 0,
            "stdout": stdout,
            "stderr": stderr,
            "return_code": result.returncode,
            "stdout_truncated": len(result.stdout) > MAX_OUTPUT_CHARS,
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False, "stdout": "", "stderr": f"Execution timed out after {timeout_s}s",
            "return_code": None, "stdout_truncated": False,
        }
    finally:
        Path(script_path).unlink(missing_ok=True)


@tool(
    name="code_interpreter",
    description="Execute Python code and return its output",
    permissions=["pro", "team", "api"],   # gated — not on the free tier by default
    timeout=DEFAULT_TIMEOUT_S + 5.0,       # tool-level timeout > subprocess timeout, for cleanup headroom
)
def code_interpreter(ctx: ToolContext, code: str, timeout_s: float = DEFAULT_TIMEOUT_S) -> dict:
    """
    Args:
        code: Python source to execute.
        timeout_s: Wall-clock limit in seconds (max 30).

    Plain `def`, not `async def`: subprocess.run() blocks regardless,
    so this is thread-offloaded by the registry rather than pretending
    to be non-blocking with `async def` and freezing the event loop
    for the subprocess's whole runtime.
    """
    timeout_s = max(1.0, min(timeout_s, 30.0))
    return run_python(code, timeout_s=timeout_s)