trading-agent/tools/sandbox.py

"""Analyst script sandbox — run_analysis.

Lets an agent write and run a real Python analysis script (pandas/numpy/yfinance/matplotlib)
like a research analyst: fetch data, compute indicators/stats/backtests, save charts. Runs in
a throwaway temp workdir with a **scrubbed environment** (no ALPACA_/FINNHUB_/OMLX_ secrets —
model-generated code must never see keys), a wall-clock timeout, and network access (for data).

Security posture (honest): running model-written networked code is a real surface. Mitigations:
scrubbed env (no secrets), isolated temp workdir (no repo writes), wall-clock timeout, truncated
output. Acceptable on a personal machine + local model; not a hardened multi-tenant sandbox.
"""

from __future__ import annotations

import os
import subprocess
import sys
import tempfile
from pathlib import Path

from langchain_core.tools import tool

from config.settings import get_settings

_STDOUT_LIMIT = 4000
_SECRET_PREFIXES = ("ALPACA", "FINNHUB", "OMLX", "EMAIL", "SMTP", "OPENAI", "ANTHROPIC")
_ARTIFACT_EXTS = {".png", ".jpg", ".jpeg", ".svg", ".csv", ".txt", ".html", ".md", ".json"}

_AVAILABLE = "pandas, numpy, yfinance, matplotlib (Agg), scipy + the standard library"

# Common failure -> actionable hint, appended to the ERROR the agent reads back. Turns the
# usual retry-loop ("same error again") into a one-shot self-correction.
_ERROR_HINTS: list[tuple[str, str]] = [
    ("ModuleNotFoundError", f"HINT: only {_AVAILABLE} are installed — rewrite using those."),
    (
        "KeyError",
        "HINT: yf.download() returns MultiIndex columns by default. Flatten first: "
        "df.columns = df.columns.get_level_values(0) — or download with "
        "yf.download(ticker, period=..., multi_level_index=False).",
    ),
    (
        "FileNotFoundError",
        "HINT: each run_analysis call starts in a FRESH empty directory — files from earlier "
        "calls are gone. Make the script self-contained (re-download data in this script).",
    ),
    (
        "NameError",
        "HINT: nothing persists between run_analysis calls — define every variable and import "
        "in THIS script.",
    ),
    (
        "JSONDecodeError",
        "HINT: a data request likely returned empty/error content — print the raw response "
        "first, and add a fallback ticker/period.",
    ),
]


def _hint_for(stderr: str) -> str:
    for needle, hint in _ERROR_HINTS:
        if needle in stderr:
            return "\n" + hint
    if "No data found" in stderr or "no price data" in stderr.lower():
        return "\nHINT: yfinance returned no data — check the ticker symbol and use a simple period like '6mo'."
    return ""


def scrubbed_env(home: str) -> dict[str, str]:
    """A minimal environment with NO secrets — built from scratch, not inherited."""
    env = {
        "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
        "HOME": home,
        "TMPDIR": home,
        "LANG": os.environ.get("LANG", "en_US.UTF-8"),
        "MPLBACKEND": "Agg",  # headless matplotlib
        "MPLCONFIGDIR": home,
    }
    # defensive: never leak anything secret-looking even if the whitelist grows
    return {k: v for k, v in env.items() if not k.upper().startswith(_SECRET_PREFIXES)}


@tool
def run_analysis(code: str) -> str:
    """Run a Python analysis script and return its stdout plus any files it wrote.

    ENVIRONMENT — read once, avoid the classic errors:
    - Installed: pandas, numpy, yfinance (market data), matplotlib (charts, headless Agg),
      scipy, plus the standard library. Nothing else — no requests/sklearn/talib.
    - Each call runs in a FRESH EMPTY temp directory. Files and variables do NOT persist
      between calls: every script must be fully self-contained (imports + data download +
      computation in one script).
    - You MUST print() every value you want to read back; save charts with
      plt.savefig("name.png") (saved files are listed in the reply).
    - ~90s wall-clock limit; no access to account secrets.
    - yfinance pitfalls: yf.download() returns MultiIndex columns — flatten with
      `df.columns = df.columns.get_level_values(0)` (or pass multi_level_index=False);
      it also auto-adjusts prices by default (auto_adjust=True).

    Use it like a real analyst: fetch data, compute indicators/stats/backtests, save charts.

    Example:
        import yfinance as yf, matplotlib.pyplot as plt
        df = yf.download("AAPL", period="6mo", multi_level_index=False)
        df["SMA50"] = df["Close"].rolling(50).mean()
        print(df[["Close","SMA50"]].tail())
        df[["Close","SMA50"]].plot(); plt.savefig("aapl.png")
    """
    s = get_settings()
    s.ensure_dirs()
    workdir = Path(tempfile.mkdtemp(dir=s.sandbox_dir))
    (workdir / "script.py").write_text(code)
    # shared, persistent matplotlib cache so the font cache is built once (not per run) and
    # doesn't land in each run's artifact list
    mpl_cache = s.sandbox_dir / ".mplconfig"
    mpl_cache.mkdir(parents=True, exist_ok=True)
    env = scrubbed_env(str(workdir))
    env["MPLCONFIGDIR"] = str(mpl_cache)
    try:
        proc = subprocess.run(
            [sys.executable, "script.py"],
            cwd=workdir,
            env=env,
            capture_output=True,
            text=True,
            timeout=s.sandbox_timeout_s,
        )
    except subprocess.TimeoutExpired:
        return f"ERROR: analysis timed out ({s.sandbox_timeout_s}s limit)."
    except Exception as e:  # pragma: no cover - defensive
        return f"ERROR: {type(e).__name__}: {e}"

    out = (proc.stdout or "").strip()
    err = (proc.stderr or "").strip()
    artifacts = [
        str(p) for p in sorted(workdir.iterdir())
        if p.is_file() and p.name != "script.py" and p.suffix.lower() in _ARTIFACT_EXTS
    ]

    if proc.returncode != 0:
        return f"ERROR (exit {proc.returncode}): {err[-1500:] or 'no stderr'}{_hint_for(err)}"
    result = out[:_STDOUT_LIMIT] or (
        "(no stdout — you must print() every value you want to read back)"
    )
    if len(out) > _STDOUT_LIMIT:
        result += "\n...[truncated]"
    if err:
        result += f"\n[stderr] {err[-400:]}"
    if artifacts:
        result += "\nArtifacts written: " + ", ".join(artifacts)
    return result