tinymind-native-colab-handoff/bundle/model/sandbox_rl.py

"""SandboxRL tool runtime for TinyMind.

The runtime is intentionally embedded as a small audited tool substrate:
Lua-like arithmetic for fast internal scratch work, project/file creation
inside a resolved sandbox root, and command execution without shell expansion.
Every action is written to a CEV ledger.
"""

from __future__ import annotations

import ast
from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
import subprocess
import time
from typing import Any, Callable


ALLOWED_CMD = {
    "echo",
    "python",
    "py",
    "bzip2",
    "bunzip2",
    "bzcat",
    "dig",
    "find",
    "git",
    "gzip",
    "gunzip",
    "host",
    "nslookup",
    "openssl",
    "ping",
    "ps",
    "tar",
    "unzip",
    "which",
    "whois",
    "zstd",
    "unzstd",
    "zstdcat",
}


@dataclass(frozen=True)
class SandboxResult:
    ok: bool
    action: str
    result: Any = None
    error: str | None = None
    stdout: str = ""
    stderr: str = ""
    sandboxed: bool = True

    def to_dict(self) -> dict:
        return {
            "ok": self.ok,
            "action": self.action,
            "result": self.result,
            "error": self.error,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "sandboxed": self.sandboxed,
        }


HttpHandler = Callable[[str, str, str | None], dict]
EvoHandler = Callable[[str, tuple[Any, ...]], Any]


def _sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def _snapshot_workspace(root: Path, max_files: int = 256) -> dict:
    files: list[dict] = []
    for path in sorted(p for p in root.rglob("*") if p.is_file()):
        if len(files) >= max_files:
            break
        try:
            data = path.read_bytes()
        except OSError:
            continue
        files.append(
            {
                "path": path.relative_to(root).as_posix(),
                "bytes": len(data),
                "sha256": _sha256_bytes(data),
            }
        )
    digest = _sha256_bytes(json.dumps(files, sort_keys=True, separators=(",", ":")).encode("utf-8"))
    return {"file_count": len(files), "truncated": len(files) >= max_files, "sha256": digest, "files": files}


class _SafeExpr(ast.NodeVisitor):
    allowed_nodes = (
        ast.Expression,
        ast.BinOp,
        ast.UnaryOp,
        ast.Constant,
        ast.Name,
        ast.Load,
        ast.Add,
        ast.Sub,
        ast.Mult,
        ast.Div,
        ast.FloorDiv,
        ast.Mod,
        ast.Pow,
        ast.USub,
        ast.UAdd,
        ast.Call,
    )

    def __init__(self, env: dict[str, Any], funcs: dict[str, Callable[..., Any]] | None = None):
        self.env = env
        self.funcs = funcs or {}

    def visit(self, node):
        if not isinstance(node, self.allowed_nodes):
            raise ValueError(f"unsupported_expr:{type(node).__name__}")
        return super().visit(node)

    def eval(self, expr: str) -> Any:
        # The Lua subset treats Windows paths as ordinary string bytes. Python's
        # parser would interpret \a, \t, and \U escapes, so escape backslashes
        # before handing the expression to ast.
        parsed = ast.parse(expr.replace("\\", "\\\\"), mode="eval")
        self.visit(parsed)
        scope = dict(self.env)
        scope.update(self.funcs)
        return eval(compile(parsed, "<sandbox-lua>", "eval"), {"__builtins__": {}}, scope)


class SandboxRLRuntime:
    """Audited workspace runtime for TinyMind tool-use training and eval."""

    def __init__(self, root: str | Path, ledger_name: str = "sandbox_ledger.jsonl"):
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self.ledger_path = self.root / ledger_name

    def _resolve(self, relative_path: str | Path) -> Path:
        target = (self.root / relative_path).resolve()
        if self.root != target and self.root not in target.parents:
            raise ValueError("path_escape")
        return target

    def _record(self, result: SandboxResult, claim: str, evidence: str, verification: str) -> dict:
        payload = result.to_dict()
        payload.update(
            {
                "timestamp": time.time(),
                "cev": {
                    "claim": claim,
                    "evidence": evidence,
                    "verification": verification,
                },
            }
        )
        with self.ledger_path.open("a", encoding="utf-8", newline="\n") as f:
            f.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")
        return payload

    def run_lua(self, code: str, http_handler: HttpHandler | None = None, evo_handler: EvoHandler | None = None) -> dict:
        env: dict[str, Any] = {}
        snapshot = _snapshot_workspace(self.root)
        http_calls: list[dict] = []
        evo_calls: list[dict] = []

        def sandbox_http_get(url: str) -> str:
            if http_handler is None:
                raise ValueError("sandbox_proxy_unavailable")
            result = http_handler("GET", str(url), None)
            http_calls.append(
                {
                    "method": "GET",
                    "url": str(url),
                    "ok": result.get("ok"),
                    "status": result.get("status"),
                    "elapsed_ms": result.get("elapsed_ms"),
                }
            )
            if not result.get("ok"):
                raise ValueError(str(result.get("error", "sandbox_proxy_error")))
            return str(result.get("response_text", ""))

        def sandbox_http_post(url: str, body: str = "") -> str:
            if http_handler is None:
                raise ValueError("sandbox_proxy_unavailable")
            result = http_handler("POST", str(url), str(body))
            http_calls.append(
                {
                    "method": "POST",
                    "url": str(url),
                    "ok": result.get("ok"),
                    "status": result.get("status"),
                    "elapsed_ms": result.get("elapsed_ms"),
                }
            )
            if not result.get("ok"):
                raise ValueError(str(result.get("error", "sandbox_proxy_error")))
            return str(result.get("response_text", ""))

        def evo_compact_lora(adapter: str, out: str, target_rank: int) -> dict:
            if evo_handler is None:
                raise ValueError("evo_weight_ops_unavailable")
            result = evo_handler("compact_lora", (str(adapter), str(out), int(target_rank)))
            evo_calls.append(
                {
                    "method": "compact_lora",
                    "adapter": str(adapter),
                    "out": str(out),
                    "target_rank": int(target_rank),
                    "ok": bool(result.get("ok")) if isinstance(result, dict) else True,
                }
            )
            return result

        def evo_choose_lora_rank(adapter: str, min_energy_retained: float = 0.995) -> dict:
            if evo_handler is None:
                raise ValueError("evo_weight_ops_unavailable")
            result = evo_handler("choose_lora_rank", (str(adapter), float(min_energy_retained)))
            evo_calls.append(
                {
                    "method": "choose_lora_rank",
                    "adapter": str(adapter),
                    "min_energy_retained": float(min_energy_retained),
                    "ok": bool(result.get("ok")) if isinstance(result, dict) else True,
                }
            )
            return result

        def evo_compact_lora_auto(adapter: str, out: str, min_energy_retained: float = 0.995) -> dict:
            if evo_handler is None:
                raise ValueError("evo_weight_ops_unavailable")
            result = evo_handler("compact_lora_auto", (str(adapter), str(out), float(min_energy_retained)))
            evo_calls.append(
                {
                    "method": "compact_lora_auto",
                    "adapter": str(adapter),
                    "out": str(out),
                    "min_energy_retained": float(min_energy_retained),
                    "ok": bool(result.get("ok")) if isinstance(result, dict) else True,
                }
            )
            return result

        def evo_patch_gguf(gguf: str) -> dict:
            if evo_handler is None:
                raise ValueError("evo_weight_ops_unavailable")
            result = evo_handler("patch_gguf", (str(gguf),))
            evo_calls.append(
                {
                    "method": "patch_gguf",
                    "gguf": str(gguf),
                    "ok": bool(result.get("ok")) if isinstance(result, dict) else True,
                }
            )
            if isinstance(result, dict) and not result.get("ok", False):
                raise ValueError(str(result.get("error", "evo_patch_gguf_failed")))
            return result

        funcs = {
            "sandbox_http_get": sandbox_http_get,
            "sandbox_http_post": sandbox_http_post,
            "evo_compact_lora": evo_compact_lora,
            "evo_choose_lora_rank": evo_choose_lora_rank,
            "evo_compact_lora_auto": evo_compact_lora_auto,
            "evo_patch_gguf": evo_patch_gguf,
        }
        try:
            result: Any = None
            for raw_line in code.splitlines():
                line = raw_line.strip()
                if not line or line.startswith("--"):
                    continue
                if line.startswith("local "):
                    name, expr = line[6:].split("=", 1)
                    env[name.strip()] = _SafeExpr(env, funcs).eval(expr.strip())
                elif line.startswith("return "):
                    result = _SafeExpr(env, funcs).eval(line[7:].strip())
                else:
                    raise ValueError("unsupported_lua_statement")
            outcome = SandboxResult(True, "lua", result=result)
        except Exception as exc:
            outcome = SandboxResult(False, "lua", error=str(exc))
        payload = self._record(outcome, "lua_subset_execution", "inline_code", "executed_inside_sandbox")
        payload["workspace_snapshot_before"] = snapshot
        payload["http_calls"] = http_calls
        payload["evo_calls"] = evo_calls
        return payload

    def write_file(self, relative_path: str | Path, content: str) -> dict:
        try:
            target = self._resolve(relative_path)
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_text(content, encoding="utf-8", newline="\n")
            outcome = SandboxResult(True, "write_file", result=str(target.relative_to(self.root)))
        except ValueError as exc:
            outcome = SandboxResult(False, "write_file", error=str(exc))
        return self._record(outcome, "workspace_file_created", str(relative_path), "path_resolved_inside_sandbox")

    def read_file(self, relative_path: str | Path) -> dict:
        try:
            target = self._resolve(relative_path)
            outcome = SandboxResult(True, "read_file", result=target.read_text(encoding="utf-8"))
        except ValueError as exc:
            outcome = SandboxResult(False, "read_file", error=str(exc))
        except FileNotFoundError:
            outcome = SandboxResult(False, "read_file", error="not_found")
        return self._record(outcome, "workspace_file_read", str(relative_path), "path_resolved_inside_sandbox")

    def create_project(self, name: str, files: dict[str, str]) -> dict:
        written: list[str] = []
        for rel, content in files.items():
            result = self.write_file(Path(name) / rel, content)
            if not result["ok"]:
                return result
            written.append(result["result"])
        outcome = SandboxResult(True, "create_project", result=written)
        return self._record(outcome, "project_scaffold_created", name, "all_files_inside_sandbox")

    def run_cmd(self, argv: list[str], timeout_s: float = 5.0) -> dict:
        if not argv:
            outcome = SandboxResult(False, "cmd", error="empty_command")
            return self._record(outcome, "command_execution", "empty", "rejected_before_execution")
        executable = Path(argv[0]).name.lower()
        if executable not in ALLOWED_CMD:
            outcome = SandboxResult(False, "cmd", error="command_not_allowlisted")
            return self._record(outcome, "command_execution", " ".join(argv), "rejected_before_execution")
        if executable == "echo":
            outcome = SandboxResult(True, "cmd", result=0, stdout=" ".join(argv[1:]) + "\n")
            return self._record(outcome, "command_execution", " ".join(argv), "executed_inside_sandbox")
        try:
            proc = subprocess.run(
                argv,
                cwd=self.root,
                capture_output=True,
                text=True,
                shell=False,
                timeout=timeout_s,
            )
            outcome = SandboxResult(
                proc.returncode == 0,
                "cmd",
                result=proc.returncode,
                stdout=proc.stdout,
                stderr=proc.stderr,
            )
        except subprocess.TimeoutExpired:
            outcome = SandboxResult(False, "cmd", error="timeout")
        return self._record(outcome, "command_execution", " ".join(argv), "executed_inside_sandbox")