trade_pool/executor.py

"""Safely turn agent-written Python into a callable strategy(view, pos) -> float.

In-process restricted exec (NOT a per-rollout sandbox — the hackathon docs warn that
per-rollout sandboxes burn credits fast). The strategy surface is deliberately tiny:
a pure function over a feature dict + current position. No imports, no IO, no builtins
beyond a safe math allowlist. A strategy that tries to escape returns an error string,
which the rubric maps to reward 0.
"""

from __future__ import annotations

import math
from typing import Callable

_SAFE_BUILTINS = {
    "abs": abs, "min": min, "max": max, "round": round, "float": float,
    "int": int, "bool": bool, "len": len, "sum": sum, "pow": pow,
    "True": True, "False": False, "None": None,
}
_SAFE_MATH = {k: getattr(math, k) for k in (
    "sqrt", "log", "log10", "exp", "fabs", "floor", "ceil", "tanh", "isnan", "isfinite", "pi", "e"
)}

_FORBIDDEN = ("import", "__", "open(", "eval(", "exec(", "globals(", "locals(",
              "compile(", "getattr(", "setattr(", "input(", "lambda os")


def compile_strategy(code: str) -> tuple[Callable | None, str | None]:
    """Compile agent code that must define `def strategy(features, position):`.

    Returns (callable, None) on success or (None, error_message) on failure.
    """
    if not code or "def strategy" not in code:
        return None, "code must define `def strategy(features, position):`"
    lowered = code.lower()
    for tok in _FORBIDDEN:
        if tok in lowered:
            return None, f"forbidden token in strategy: {tok!r}"

    sandbox_globals = {"__builtins__": _SAFE_BUILTINS, "math": _SAFE_MATH, **_SAFE_MATH}
    try:
        compiled = compile(code, "<strategy>", "exec")
        exec(compiled, sandbox_globals)  # noqa: S102 - restricted globals, no builtins
    except Exception as e:  # noqa: BLE001
        return None, f"compile/exec error: {type(e).__name__}: {e}"

    fn = sandbox_globals.get("strategy")
    if not callable(fn):
        return None, "`strategy` is not callable"

    def wrapped(view: dict, pos: float) -> float:
        return fn(view, pos)

    return wrapped, None


def extract_code(completion: str) -> str:
    """Pull a python code block out of a model completion (```python ... ``` or raw)."""
    if "```" in completion:
        parts = completion.split("```")
        for i in range(1, len(parts), 2):
            block = parts[i]
            if block.startswith("python"):
                block = block[len("python"):]
            if "def strategy" in block:
                return block.strip()
    return completion.strip()