CoolFace
Apppublic

apodex/frontier-agent-demo

sourceHugging Faceapache-2.0updated 1mo agoView on Hugging Face
14likes
run_python_code.py471 linesDownload Raw Back to tools
1"""Tool: run_python_code — execute Python code in an isolated sandbox."""2from __future__ import annotations3 4import ast5import asyncio6import logging7import re8import uuid9from pathlib import Path10 11from frontier_agent.core.tool import tool12from frontier_agent.infra.usage_meter import record_api_request13from plugins.tools._code_sanitize import sanitize_code14from plugins.tools._net_guard import ensure_guard_file, guard_env_prefix15from plugins.tools._overflow import maybe_overflow16from plugins.tools._sandbox import (17    aget_sandbox,18    arun_sandbox_cmd,19    is_e2b_sandbox,20    remote_exec_prefix,21)22from plugins.tools.bash import BASH_STDERR_SEPARATOR23 24logger = logging.getLogger(__name__)25 26_MAX_OUTPUT = 10_00027 28# Tail of captured stdout surfaced when a run is killed at the timeout, so a29# batched crawl that printed progress leaves salvageable output instead of a30# total loss (observed: 2×600s OpenAlex pagination crawls, zero output kept).31_PARTIAL_TAIL = 3_00032 33# Detect-only (2026-06-04): log when agent code does raw HTTP so we can size34# how often crawls bypass the governed web tools (proxy cache, retries,35# metering) before deciding whether to clamp their timeout. See the network36# discipline section of the sub-agent research prompt.37_NET_LIB_RE = re.compile(38    r"^\s*(?:import|from)\s+(requests|httpx|aiohttp|urllib3|urllib|socket)\b",39    re.MULTILINE,40)41 42_ML_MODULES = frozenset({43    "torch",44    "torchvision",45    "torchaudio",46    "transformers",47    "datasets",48    "huggingface_hub",49    "sentence_transformers",50    "diffusers",51    "accelerate",52    "tensorflow",53    "keras",54})55_ML_INSTALL_PACKAGES = _ML_MODULES | {56    "huggingface-hub",57    "sentence-transformers",58}59_ML_DOWNLOAD_ATTRS = frozenset({60    "from_pretrained",61    "snapshot_download",62    "hf_hub_download",63    "load_state_dict_from_url",64})65_SUBPROCESS_CALLS = frozenset({66    "subprocess.run",67    "subprocess.call",68    "subprocess.check_call",69    "subprocess.check_output",70    "subprocess.Popen",71})72 73# Hard-deny heavyweight ML/model-download paths in the generic research74# sandbox. These libraries are not part of run_python_code's advertised75# package set, and their convenience APIs can pull multi-GB weights into RAM or76# disk cache before the Python process sees a clean MemoryError.77_ML_IMPORT_RE = re.compile(78    r"^\s*(?:import|from)\s+("79    r"torch|torchvision|torchaudio|transformers|datasets|huggingface_hub|"80    r"sentence_transformers|diffusers|accelerate|tensorflow|keras"81    r")(?:\b|\.)",82    re.MULTILINE,83)84_ML_DOWNLOAD_API_RE = re.compile(85    r"\b("86    r"torch\.hub\.(?:load|download_url_to_file)|"87    r"load_state_dict_from_url|"88    r"(?:from_pretrained|snapshot_download|hf_hub_download)\s*\("89    r")"90)91_ML_INSTALL_RE = re.compile(92    r"\b(?:pip|python\s+-m\s+pip|uv\s+pip)\s+install\b[^\n;]*\b("93    r"torch|torchvision|torchaudio|transformers|datasets|huggingface-hub|"94    r"sentence-transformers|diffusers|accelerate|tensorflow|keras"95    r")\b",96    re.IGNORECASE,97)98 99_ML_BLOCK_MESSAGE = (100    "Error: run_python_code blocks PyTorch/HuggingFace/transformers-style "101    "model loading and downloads in the generic research sandbox. These "102    "paths can fetch multi-GB weights or metadata caches and OOM the worker. "103    "Use lightweight structured APIs, aggregate/count endpoints, or the "104    "governed web_search/web_fetch tools instead."105)106 107_OFFLINE_DOWNLOAD_ENV = (108    "HF_HUB_OFFLINE=1 "109    "TRANSFORMERS_OFFLINE=1 "110    "HF_DATASETS_OFFLINE=1 "111    "HF_HUB_DISABLE_TELEMETRY=1 "112    "TORCH_HOME=/tmp/frontier_agent_no_torch_cache "113    "HF_HOME=/tmp/frontier_agent_no_hf_cache "114    "TRANSFORMERS_CACHE=/tmp/frontier_agent_no_hf_cache "115)116 117# Streaming recipe shared by the OOM (-1) and MemoryError messages. Precise on118# purpose: a vague "use chunked processing" hint led models to chunk only the119# pandas parse while still buffering the whole download via120# ``requests.get(url).content`` — the retry then OOM'd identically.121_MEM_RECIPE = (122    "- pd.read_csv(url, chunksize=10000) streams the download AND the parse; "123    "filter each chunk, keep only needed columns/rows.\n"124    "- Or requests.get(url, stream=True) + iterate lines; never touch "125    "response.content / .text on a large body (it buffers everything, and "126    ".decode() doubles it).\n"127    "- Write the filtered subset to a file in the current working directory "128    "first, then analyze that small file in a second run."129)130 131 132def _default_timeout() -> int:133    """Per-exec wall-clock default from config (run_python_timeout_s)."""134    try:135        from frontier_agent.infra.config import get_config136        return int(get_config().run_python_timeout_s)137    except Exception:138        return 300139 140 141def _max_timeout() -> int:142    """Hard ceiling on an agent-supplied timeout (run_python_max_timeout_s)."""143    try:144        from frontier_agent.infra.config import get_config145        return int(get_config().run_python_max_timeout_s)146    except Exception:147        return 300148 149 150def _root_module(name: str) -> str:151    return name.split(".", 1)[0]152 153 154def _literal_str(node: ast.AST) -> str | None:155    if isinstance(node, ast.Constant) and isinstance(node.value, str):156        return node.value157    return None158 159 160def _call_name(node: ast.AST) -> str:161    parts: list[str] = []162    cur = node163    while isinstance(cur, ast.Attribute):164        parts.append(cur.attr)165        cur = cur.value166    if isinstance(cur, ast.Name):167        parts.append(cur.id)168    return ".".join(reversed(parts))169 170 171def _iter_literal_strings(node: ast.AST) -> list[str]:172    if not isinstance(node, (ast.List, ast.Tuple)):173        return []174    out: list[str] = []175    for elt in node.elts:176        text = _literal_str(elt)177        if text is not None:178            out.append(text)179    return out180 181 182def _is_blocked_pip_install_args(args: list[str]) -> bool:183    lowered = [arg.lower() for arg in args]184    if "install" not in lowered:185        return False186    installer = lowered[:3]187    if not (188        any(arg.endswith("pip") or arg in {"pip", "pip3"} for arg in installer)189        or (len(installer) >= 3 and installer[1:3] == ["-m", "pip"])190        or installer[:2] == ["uv", "pip"]191    ):192        return False193    install_at = lowered.index("install")194    packages = {195        arg.split("==", 1)[0].split(">=", 1)[0].split("<=", 1)[0]196        for arg in lowered[install_at + 1 :]197        if arg and not arg.startswith("-")198    }199    return any(pkg in _ML_INSTALL_PACKAGES for pkg in packages)200 201 202class _MLDownloadBlockVisitor(ast.NodeVisitor):203    def __init__(self) -> None:204        self.blocked = False205 206    def visit_Import(self, node: ast.Import) -> None:207        if any(_root_module(alias.name) in _ML_MODULES for alias in node.names):208            self.blocked = True209            return210        self.generic_visit(node)211 212    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:213        if node.module and _root_module(node.module) in _ML_MODULES:214            self.blocked = True215            return216        self.generic_visit(node)217 218    def visit_Call(self, node: ast.Call) -> None:219        name = _call_name(node.func)220        if name == "__import__" and node.args:221            mod = _literal_str(node.args[0])222            if mod and _root_module(mod) in _ML_MODULES:223                self.blocked = True224                return225        if name == "importlib.import_module" and node.args:226            mod = _literal_str(node.args[0])227            if mod and _root_module(mod) in _ML_MODULES:228                self.blocked = True229                return230        if (231            name in {"torch.hub.load", "torch.hub.download_url_to_file"}232            or name.rsplit(".", 1)[-1] in _ML_DOWNLOAD_ATTRS233        ):234            self.blocked = True235            return236        if name == "getattr" and len(node.args) >= 2:237            attr = _literal_str(node.args[1])238            if attr in _ML_DOWNLOAD_ATTRS:239                self.blocked = True240                return241        if (242            name in _SUBPROCESS_CALLS243            and node.args244            and _is_blocked_pip_install_args(_iter_literal_strings(node.args[0]))245        ):246            self.blocked = True247            return248        self.generic_visit(node)249 250 251def _blocked_ml_download_ast(code: str) -> bool:252    try:253        tree = ast.parse(code)254    except SyntaxError:255        return False256    visitor = _MLDownloadBlockVisitor()257    visitor.visit(tree)258    return visitor.blocked259 260 261def _blocked_ml_download_reason(code: str) -> str | None:262    """Return a user-facing block reason for heavyweight ML download paths."""263    if _ML_IMPORT_RE.search(code):264        return _ML_BLOCK_MESSAGE265    if _ML_DOWNLOAD_API_RE.search(code):266        return _ML_BLOCK_MESSAGE267    if _ML_INSTALL_RE.search(code):268        return _ML_BLOCK_MESSAGE269    if _blocked_ml_download_ast(code):270        return _ML_BLOCK_MESSAGE271    return None272 273 274@tool275async def run_python_code(code: str, timeout: int = 0) -> str:276    """Execute Python code in an isolated sandbox environment.277 278    Pre-installed packages: numpy, pandas, scipy, sympy, mpmath, networkx,279    plotly, Pillow, openpyxl, tabulate.280 281    The code MUST terminate on its own within seconds. NEVER submit code that282    runs forever: no unbounded ``while True:`` refresh loops, no live curses /283    TUI dashboards, no servers, daemons, or ``input()`` waits. If you are284    testing a long-running program (e.g. a monitoring dashboard that samples285    repeatedly), drive it with a BOUNDED test harness instead — call its286    collect/render function 2-3 times with a short (<=1s) sleep, or run it with287    an explicit ``--iterations N`` / ``--once`` flag — so it exits quickly. A288    non-terminating program is killed only at the timeout, wasting the whole289    budget; repeated kills can stall the entire task.290 291    Network rules: to read the CONTENT of a web page or PDF use ``web_fetch``,292    not requests/httpx. Calling a structured API (OpenAlex, Crossref, EDGAR,293    ...) from code is allowed ONLY for aggregate/count queries (e.g. OpenAlex294    ``group_by``, ``meta.count``) — a handful of requests that each return a295    small statistical answer. Bulk-paginating a full result set (downloading296    every record's metadata page by page) is FORBIDDEN: a corpus-scale crawl297    cannot be verified item-by-item anyway, so sample instead — get the total298    via an aggregate endpoint, pull ≤2 pages as a representative sample, and299    reason from count + sample. Set a short per-request timeout and write any300    intermediate data to files in the current working directory, not stdout.301    If the plan seems to need more than ~10 requests, narrow the query302    server-side.303 304    Scratch files: use RELATIVE paths. The working directory is private to305    this agent and persists across calls, so a follow-up run sees what this306    one wrote. Absolute ``/tmp`` and ``/workspace`` paths are shared with307    every other agent on this task — a fixed absolute name can be overwritten308    by a concurrent agent, and you would read back their data as if it were309    yours.310 311    Output rules: NEVER print a full dataset or raw API responses to stdout —312    print counts, aggregates, and at most ~20 sample rows. Anything larger313    belongs in a file in the working directory (analyze it in a follow-up314    run). Stdout is capped; a full-corpus dump gets truncated AND bloats every315    downstream consumer of this conversation.316 317    Memory rules: the sandbox has LIMITED RAM (~512MB) and an out-of-memory318    kill loses the whole run. Any file/dataset over ~20MB MUST be streamed,319    never buffered: use ``pd.read_csv(url, chunksize=10000)`` (streams both320    download and parse; filter each chunk) or ``requests.get(url,321    stream=True)`` + line iteration. NEVER call ``response.content`` /322    ``response.text`` on a large body — it buffers the full payload and323    ``.decode()`` doubles it; loading the result into a DataFrame multiplies324    it again. Filter early, keep only needed columns, write the reduced325    subset to a file in the working directory and analyze that instead.326 327    Args:328        code: Python source code to execute. Must self-terminate.329        timeout: Maximum execution time in seconds. 0 (default) uses the330            server-configured default (run_python_timeout_s). Hard-capped at331            run_python_max_timeout_s (300s) — a larger request is clamped, so332            a single exec can't pin a scarce sandbox slot for many minutes.333 334    Returns:335        stdout + stderr from the execution, or an error message.336    """337    if not code or not code.strip():338        return "Error: empty code provided."339 340    if timeout <= 0:341        timeout = _default_timeout()342    # Clamp an agent-supplied timeout to the hard ceiling regardless of what343    # the model asked for (it can ask for less, never more). Stops a runaway344    # data-collection exec from burning many minutes on one sandbox slot.345    max_timeout = _max_timeout()346    if timeout > max_timeout:347        timeout = max_timeout348 349    # Normalise Unicode math symbols copied from problem statements350    # (``∫ Σ π ≤`` → ``integral sum pi <=``) before executing — LLMs351    # frequently echo these and Python rejects them with SyntaxError.352    code = sanitize_code(code)353 354    blocked = _blocked_ml_download_reason(code)355    if blocked:356        return blocked357 358    net_libs = sorted(set(_NET_LIB_RE.findall(code)))359    if net_libs:360        logger.info(361            "run_python_code: raw HTTP libs in agent code: %s "362            "(timeout=%ss, code_len=%d)",363            ",".join(net_libs), timeout, len(code),364        )365 366    try:367        sandbox = await aget_sandbox()368    except RuntimeError as exc:369        return f"Error: sandbox unavailable — {exc}"370 371    filename = f"/tmp/exec_{uuid.uuid4().hex[:8]}.py"372    # E2B / Docker ``commands.run(timeout=...)`` only times out the SDK client373    # wait — it does NOT kill the in-container process, so a brute-force374    # enumeration keeps burning a scarce sandbox slot well past ``timeout``375    # (issue #221: observed up to 600s on a 120s budget). Wrap remote execs in376    # a coreutils ``timeout -s KILL`` for an OS-level hard kill that even a377    # numpy/C loop can't ignore.378    # Socket-level download cap (sitecustomize injection): bounds how many379    # bytes any python process in this exec tree — including pip children —380    # can receive per connection, whether buffered or streamed to disk. See381    # ``plugins/tools/_net_guard.py`` (smoke-memwt-003: GB-scale dataset382    # downloads inside the sandbox).383    await ensure_guard_file(sandbox)384    net_guard_env = guard_env_prefix()385    base_exec_cmd = f"python3 {filename}"386    # Add the per-exec memory cap + single-thread math-lib env so a buffered387    # parse fails inside the sandbox instead of OOM-killing its environment.388    exec_cmd = (389        f"{remote_exec_prefix()}{net_guard_env}{_OFFLINE_DOWNLOAD_ENV}"390        f"timeout -s KILL {timeout}s {base_exec_cmd}"391    )392    cmd_timeout = timeout + 30  # let the inner OS timeout fire first393    try:394        if hasattr(sandbox, "files"):395            await asyncio.to_thread(sandbox.files.write, filename, code)396        else:397            await asyncio.to_thread(398                Path(filename).write_text, code, encoding="utf-8",399            )400 401        # Count one E2B execution (lifetime is metered at the402        # sandbox create/kill sites in ``_sandbox.py``). Bwrap/Current/Docker403        # facades bill nobody — they all carry a ``sandbox_id``404        # too, so discriminate by implementing module instead.405        if is_e2b_sandbox(sandbox):406            record_api_request("e2b")407        result = await arun_sandbox_cmd(408            sandbox, exec_cmd, timeout=cmd_timeout,409            # Match ``bash``: both run model-authored code, so denying the410            # network here while allowing it there only means the same snippet411            # succeeds under ``bash -c 'python3 …'`` and fails via this tool —412            # an asymmetry with no security value and a hard-to-place error.413            # The bound is the socket cap injected above, not the namespace;414            # on E2B/CurrentSandbox this path already had network anyway, so415            # only the bwrap backend changes.416            allow_net=True,417        )418    except TimeoutError:419        return f"Error: execution timed out after {timeout}s."420    except Exception as exc:421        return f"Error: {type(exc).__name__}: {exc}"422 423    # coreutils ``timeout`` exit codes: 124 = TERM expired, 137 = 128+SIGKILL.424    # Surface a consistent timeout message so the LLM425    # gets a consistent signal instead of a bare "[exit code 137]". Append the426    # tail of whatever stdout the process streamed before the kill — a batched427    # crawl that printed progress / checkpointed partial results can be resumed428    # or narrowed instead of being a total loss.429    if result.exit_code in (124, 137):430        msg = f"Error: execution timed out after {timeout}s."431        partial = (result.stdout or "").strip()432        if partial:433            msg += (434                "\nPartial stdout before the kill (salvage it: resume from the "435                "last checkpoint, or narrow the query / use an aggregate "436                f"endpoint instead of re-running as-is):\n{partial[-_PARTIAL_TAIL:]}"437            )438        return msg439 440    # E2B reports a process that died without a normal exit status (OOM-killed,441    # envd-side failure) as exit code -1, usually with an empty stderr442    # (observed: swarm_gv 2026-06-03, repeated -1 bursts on one sandbox). Give443    # the model an actionable signal instead of a bare "[exit code -1]".444    if result.exit_code == -1:445        detail = (result.stderr or result.stdout or "").strip()446        suffix = f"\n{detail}" if detail else ""447        return (448            "Error: sandbox process died unexpectedly (exit code -1, likely "449            "out-of-memory or a sandbox-side failure). The sandbox has very "450            "limited RAM — do NOT retry the same code. To process a large "451            "file/dataset, stream it end-to-end instead of buffering it:\n"452            f"{_MEM_RECIPE}{suffix}"453        )454 455    output = result.stdout or ""456    if result.stderr:457        output += f"{BASH_STDERR_SEPARATOR}{result.stderr}"458    if result.exit_code != 0:459        output = f"[exit code {result.exit_code}]\n{output}"460        # The per-exec ``ulimit -v`` cap converts a would-be VM OOM kill into461        # a clean MemoryError traceback — steer the retry toward streaming462        # instead of letting the model shrink the workload and re-buffer.463        if "MemoryError" in (result.stderr or ""):464            output += (465                "\n[hint] The process hit the per-exec memory cap. Do NOT "466                "retry the same approach with a smaller slice — stream "467                f"end-to-end instead of buffering:\n{_MEM_RECIPE}"468            )469 470    return maybe_overflow("run_python_code", output)471