app/mimo/runner/core.py

"""One rollout, start to finish: sandbox up, environment set, agent runs, task's own verifier grades.

Everything a rollout does becomes an *event*, kept in memory for live viewers and appended to
`events.jsonl` in batches. Event kinds the UI understands:

    phase    {name, status: "start"|"done"|"error", detail}
    text     {text}                                 an assistant message
    tool     {tool, title, input, output, status, ms}
    step     {tokens: {...}, cost}                  one model call finished
    log      {text}                                 setup / verifier output worth showing
    checks   {checks: [{id, passed, score, method, tier, weight, message}], reward, summary}
    error    {text}

The user's token is held here for the rollout's lifetime only. It never reaches an event or a file.
"""

from __future__ import annotations

import json
import re
import secrets
import threading
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor

from .. import catalog, config, models, store, version

_pool = ThreadPoolExecutor(max_workers=config.MAX_ACTIVE_ROLLOUTS, thread_name_prefix="rollout")
_live: dict[str, "Rollout"] = {}
_caps: dict[str, "Rollout"] = {}   # model-proxy capability -> its rollout, while it runs
_live_lock = threading.Lock()


# key-shaped: a known prefix, then a run with an uppercase letter or a digit (code names like hf_raise_for_status don't match)
SECRET_RE = re.compile(r"\b(hf_|sk-|sk_|api_key=|Bearer )(?=[A-Za-z0-9\-]*[A-Z0-9])[A-Za-z0-9\-]{4,}(?![a-z_])")


ACTIVE_STATES = {"queued", "starting", "setup", "running", "verifying"}


class Cancelled(Exception):
    pass


class Rollout:
    def __init__(self, run: dict, token: str, agent_key: str | None = None):
        self.run = run
        self.id = run["id"]
        self.token = token
        self.agent_key = agent_key   # a bring-your-own endpoint's key: memory only, like the HF token
        # what the sandbox gets instead of credentials (see config.PUBLIC_URL): valid while this rollout is live
        self.cap = secrets.token_urlsafe(32)
        self.stream: dict | None = None   # the model reply streaming through the proxy right now, counted as it goes
        self.events: list[dict] = store.read_events(self.id)
        self._pending: list[dict] = []
        self._last_flush = time.time()
        self._lock = threading.Lock()
        self.t0 = time.time()
        self.sandbox = None
        self.cancelled = threading.Event()
        self.tokens = {"input": 0, "output": 0, "reasoning": 0, "cache_read": 0, "cache_write": 0}
        self.model_cost = 0.0

    # ── events ───────────────────────────────────────────────────────────────
    def _redact(self, v):
        """Nothing secret reaches a trace: the user's HF token, an endpoint key, or anything shaped like one
        (the General verifier, for one, prints the first characters of its judge key)."""
        if isinstance(v, str):
            for secret in (self.token, self.agent_key, getattr(self, "cap", None)):
                if secret:
                    v = v.replace(secret, "[redacted]")
            return SECRET_RE.sub(lambda m: m.group(1) + "[redacted]", v)
        if isinstance(v, dict):
            return {k: self._redact(x) for k, x in v.items()}
        if isinstance(v, list):
            return [self._redact(x) for x in v]
        return v

    def emit(self, kind: str, **data) -> None:
        ev = {"i": len(self.events), "t": round(time.time() - self.t0, 2), "kind": kind, **self._redact(data)}
        with self._lock:
            self.events.append(ev)
            self._pending.append(ev)
            if time.time() - self._last_flush > 1.0 or kind in ("phase", "checks", "error", "prompt"):
                self._flush()

    def _flush(self) -> None:
        pending, self._pending = self._pending, []
        self._last_flush = time.time()
        store.append_events(self.id, pending)

    def phase(self, name: str, status: str = "start", detail: str = "") -> None:
        self.emit("phase", name=name, status=status, detail=detail)
        if status == "start":
            self.update(status={"sandbox": "starting", "setup": "setup", "agent": "running",
                                "verify": "verifying"}.get(name, self.run.get("status")), phase=name)

    def log(self, text: str) -> None:
        if text and text.strip():
            self.emit("log", text=text.strip()[-6000:])

    def add_tokens(self, t: dict) -> None:
        for k in self.tokens:
            self.tokens[k] += int(t.get(k) or 0)
        ep = self.run.get("endpoint")
        if ep:   # your own endpoint: metered only if you told us its prices
            t = self.tokens
            self.model_cost = ((t["input"] + t["cache_read"] + t["cache_write"]) * (ep.get("price_in") or 0)
                               + (t["output"] + t["reasoning"]) * (ep.get("price_out") or 0)) / 1e6
        else:
            self.model_cost = models.cost(self.run["model"], self.run.get("provider"), self.tokens)
        self.update(tokens=self.tokens, cost=self.cost_now())

    def cost_now(self) -> dict:
        hours = (time.time() - self.t0) / 3600
        sandbox = hours * config.FLAVOR_PRICE_PER_HOUR.get(self.run.get("flavor") or "", 0.0)
        judge = self.run.get("judge_cost") or 0.0
        return {"model": round(self.model_cost, 5), "judge": round(judge, 5), "sandbox": round(sandbox, 5),
                "total": round(self.model_cost + judge + sandbox, 5)}

    def update(self, **fields) -> None:
        fields = self._redact(fields)
        self.run.update(fields)
        store.update(self.id, **fields)

    def check_cancel(self) -> None:
        if self.cancelled.is_set():
            raise Cancelled()

    # ── lifecycle ────────────────────────────────────────────────────────────
    def start_sandbox(self, image: str, flavor: str) -> "object":
        from huggingface_hub import Sandbox

        self.phase("sandbox", detail=f"{image.rsplit(':', 1)[-1]} on {flavor}")
        self.update(flavor=flavor)
        t = time.time()
        self.sandbox = Sandbox.create(image=image, flavor=flavor, idle_timeout=config.SANDBOX_IDLE_TIMEOUT,
                                      forward_hf_token=False, token=self.token, start_timeout=900,
                                      labels={"app": "rl-explorer", "runner": "mimo", "run": self.id})
        self.update(sandbox_id=self.sandbox.id)
        self.phase("sandbox", "done", f"ready in {time.time() - t:.0f}s")
        return self.sandbox

    def agent_api(self) -> tuple[str, str | None, str]:
        """(base URL, key, model string) the agent's model calls go to."""
        ep = self.run.get("endpoint")
        if ep:
            return ep["base_url"], self.agent_key, self.run["model"]
        mid = f"{self.run['model']}:{self.run['provider']}" if self.run.get("provider") else self.run["model"]
        return config.ROUTER, self.token, mid

    def llm_base(self) -> str | None:
        """Only a per-run capability reaches the sandbox, including local previews."""
        if config.PUBLIC_URL:
            return f"{config.PUBLIC_URL}/api/llm/{self.cap}/v1"
        from ...runner import _proxy_url
        return f"{_proxy_url().rstrip('/')}/rlx/llm/{self.cap}/v1"

    def check_model_network(self):
        """Check the exact relay path from the sandbox, without a model call."""
        url = self.llm_base().removesuffix("/v1") + "/health"
        script = ("import urllib.request,urllib.error\n"
                  f"url={url!r}\n"
                  "try:\n"
                  " r=urllib.request.urlopen(url,timeout=15)\n"
                  " print('relay HTTP',r.status)\n"
                  "except urllib.error.HTTPError as e:\n"
                  " print('relay HTTP',e.code);raise SystemExit(1)\n"
                  "except Exception as e:\n"
                  " print(type(e).__name__);raise SystemExit(1)\n")
        import shlex
        result = self.sh("python3 -c " + shlex.quote(script), timeout=25)
        if result.exit_code:
            raise RuntimeError("The sandbox cannot reach the model relay. Check HTTPS egress from the sandbox and ingress to the explorer. "
                               "No account token was forwarded. " + (result.stdout or "")[-120:])
        self.log("Network check passed: sandbox → model relay. HF account token stays on the explorer.")

    def upstream(self, model: str) -> tuple[str, str | None] | None:
        """For the model proxy: (base URL, key) for a model this rollout may call, else None."""
        base, key, mid = self.agent_api()
        if model == mid:
            return base, key
        if self.run.get("judge") and model == self.run["judge"]:
            return config.ROUTER, self.token
        return None

    def sh(self, cmd: str, timeout: float = 600, **kw):
        """Run a shell command in the sandbox; never raises on a non-zero exit."""
        self.check_cancel()
        return self.sandbox.run(cmd, shell=True, check=False, timeout=timeout, **kw)

    def settle_model(self) -> None:
        """Check, with tiny requests, that the model takes this rollout's settings, and adjust rather than fail.

        Some models take no thinking level at all (Kimi-K2-Instruct answers HTTP 400 to reasoning_effort=low): run
        those with thinking off. A provider that caps replies below the agent's output limit gets its own cap."""
        p = self.run.get("params") or {}
        base, key, mid = self.agent_api()
        byo = bool(self.run.get("endpoint"))
        eff = p.get("thinking")
        if eff in ("low", "medium", "high") and _refused(base, key, mid, byo, {"reasoning_effort": eff}, r"reason|think|effort"):
            p = {**p, "thinking": "none", "thinking_requested": eff}
            self.update(params=p)
            self.log(f"{self.run['model']} doesn't take a thinking level (it refused reasoning_effort={eff}), "
                     "so this rollout runs with thinking off.")
        if self.run["domain"] != "music" and not p.get("max_tokens"):
            from .opencode import OUTPUT_LIMIT
            if _refused(base, key, mid, byo, {"max_tokens": OUTPUT_LIMIT}, r"token|length|context|max"):
                self.update(params={**p, "output_limit": 32000})
                self.log(f"The provider doesn't allow {OUTPUT_LIMIT:,}-token replies for this model; "
                         "the agent's replies are capped at 32,000 tokens.")

    def execute(self) -> None:
        from .domains import ADAPTERS

        if self.run.get("runner") == "nemo-gym":
            from ... import nemo_runner as adapter
        else:
            adapter = ADAPTERS[self.run["domain"]]
        try:
            self.update(status="starting", started_at=time.time())
            self.settle_model()
            result = adapter.run(self)
            self.update(status="done", phase="done", reward=result.get("reward"), reward_error=result.get("error"),
                        graded=result.get("reward") is not None, wall_s=round(time.time() - self.t0, 2),
                        finished_at=time.time(), cost=self.cost_now())
            self.phase("done", "done", f"reward {result.get('reward')}" if result.get("reward") is not None
                       else f"not scored: {result.get('error')}")
        except Cancelled:
            self.update(status="cancelled", finished_at=time.time(), cost=self.cost_now())
            self.phase("done", "error", "cancelled")
        except Exception as e:  # noqa: BLE001
            if self.cancelled.is_set():   # killing the sandbox makes the blocked call fail: that is the cancel, not a crash
                self.update(status="cancelled", finished_at=time.time(), cost=self.cost_now())
            else:
                msg = friendly(e) or f"{type(e).__name__}: {str(e)[:300]}"
                self.emit("error", text=msg, trace=traceback.format_exc()[-3000:])
                self.update(status="failed", error=msg, finished_at=time.time(), cost=self.cost_now())
        finally:
            with self._lock:
                self._flush()
            self.stop_sandbox()
            self.agent_key = None
            self.native_task = None
            with _live_lock:
                _caps.pop(self.cap, None)
            with _live_lock:
                _live.pop(self.id, None)
            self.token = ""

    def stop_sandbox(self) -> None:
        if self.sandbox is None:
            return
        try:
            from huggingface_hub import Sandbox

            Sandbox.kill(self.sandbox.id, token=self.token)
        except Exception:  # it idles out on its own anyway
            pass
        self.sandbox = None


def friendly(e: Exception) -> str | None:
    """The Hub errors people actually hit, said so they know what to do."""
    status = getattr(getattr(e, "response", None), "status_code", None)
    text = str(e)
    if status == 402 or "402 Payment Required" in text:
        return ("Your Hugging Face account has no prepaid credit for sandboxes, so none could start. Add credit at "
                "https://huggingface.co/settings/billing and run again. Nothing was charged.")
    if status in (401, 403) and "/api/jobs" in text:
        return ("Your sign-in isn't allowed to start sandboxes. Sign in again and allow Jobs, or use a write token.")
    if status == 429:
        return "Hugging Face is rate-limiting sandbox starts for your account. Wait a minute and run again."
    return None


# ── public API ───────────────────────────────────────────────────────────────
def active_for(user: str) -> int:
    with _live_lock:
        return sum(1 for r in _live.values() if r.run["user"] == user and not r.cancelled.is_set())


def submit(user: str, token: str, task: dict, model: str, provider: str | None, judge: str | None,
           endpoint: dict | None = None, agent_key: str | None = None, params: dict | None = None,
           visibility: str = "public", extra: dict | None = None, reserved: bool = False) -> dict:
    """`extra`: fields the explorer adds to the record (its dataset, ref, environment and runner). `reserved`: the
    explorer already holds this rollout's place (app/runner._slot, with the admin's limits), so no check of its own."""
    if not reserved and active_for(user) >= config.MAX_ACTIVE_PER_USER:
        raise RuntimeError(f"You already have {config.MAX_ACTIVE_PER_USER} rollouts running. Wait for one to finish.")
    with _live_lock:
        if not reserved and len(_live) >= config.MAX_ACTIVE_ROLLOUTS:
            raise RuntimeError("The explorer is at capacity right now. Try again in a few minutes.")
    run = store.create({
        "id": time.strftime("%Y%m%d-%H%M%S-") + uuid.uuid4().hex[:6],
        "user": user, "task_id": task["id"], "domain": task["domain"], "title": task["title"],
        "facets": task.get("facets"), "model": model, "provider": provider, "judge": judge,
        "status": "queued", "harness": "opencode", "endpoint": endpoint, "params": params or {},
        "visibility": "private" if visibility == "private" else "public",
        "image": catalog.image_for(task["id"]) if task["domain"] not in ("music", "nemo-gym") else None,
        "provenance": version.provenance(),
        **(extra or {}),
    })
    r = Rollout(run, token, agent_key)
    r.native_task = task.get("native_task")  # reference action stays in memory; never in the run record
    with _live_lock:
        _live[r.id] = r
        _caps[r.cap] = r
    _pool.submit(r.execute)
    return run


_refusals: dict[tuple, bool] = {}


def _refused(base: str, key: str | None, mid: str, byo: bool, extra: dict, about: str) -> bool:
    """Does the model answer HTTP 400/422, about `about`, to a two-word request carrying `extra`? Cached."""
    k = (base, mid, json.dumps(extra, sort_keys=True))
    if k not in _refusals:
        import httpx

        from .. import endpoints
        body = {"model": mid, "max_tokens": 16, **extra, "messages": [{"role": "user", "content": "Say OK."}]}
        try:
            if byo:   # the user's endpoint: pinned connection, capped reply
                r = endpoints.request("POST", f"{base}/chat/completions", key, timeout=60, json=body)
            else:
                r = httpx.post(f"{base}/chat/completions", json=body, timeout=60,
                               headers={"Authorization": f"Bearer {key}"} if key else {})
        except Exception:
            return False   # can't tell; let the rollout itself report a real connection problem
        refused = r.status_code in (400, 422) and bool(re.search(about, r.text, re.I))
        if r.status_code in (200, 400, 422):   # only cache real answers, not rate limits or outages
            _refusals[k] = refused
        return refused
    return _refusals[k]


def by_cap(cap: str) -> Rollout | None:
    with _live_lock:
        r = _caps.get(cap)
    return r if r is not None and not r.cancelled.is_set() else None


def live(run_id: str) -> Rollout | None:
    with _live_lock:
        return _live.get(run_id)


def events(run_id: str, after: int = 0) -> list[dict]:
    r = live(run_id)
    if r is not None:
        with r._lock:
            return r.events[after:]
    return store.read_events(run_id, after)


def cancel(run_id: str) -> bool:
    r = live(run_id)
    if r is None:
        return False
    r.cancelled.set()
    with _live_lock:
        _caps.pop(r.cap, None)
    # Say so at once: the worker may sit in a blocking sandbox call for a while after the kill.
    r.update(status="cancelled", finished_at=time.time(), cost=r.cost_now())
    r.phase("done", "error", "stopped by you")
    threading.Thread(target=r.stop_sandbox, daemon=True).start()
    return True