group_xinghuo_jinrong/scripts/eval/sandbox.py

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


class SandboxSafetyError(RuntimeError):
    pass


@dataclass(frozen=True)
class SandboxConfig:
    run_id: str
    root: Path
    live: bool = False
    keep: bool = False
    env_name: str = "offline"
    resources: dict[str, str] = field(default_factory=dict)

    @property
    def artifact_dir(self) -> Path:
        return self.root / "artifacts" / "eval" / self.run_id

    @property
    def ledger_path(self) -> Path:
        return self.artifact_dir / "ledger.json"

    def safe_environment(self) -> dict[str, Any]:
        return {
            "env_name": self.env_name,
            "live": self.live,
            "run_id": self.run_id,
            "resources": {
                key: _redact_resource(value) for key, value in self.resources.items()
            },
        }


@dataclass
class ResourceLedger:
    config: SandboxConfig
    resources: list[dict[str, Any]] = field(default_factory=list)

    def add(self, kind: str, identifier: str, cleanup: str) -> None:
        self.resources.append(
            {"kind": kind, "identifier": identifier, "cleanup": cleanup, "cleaned": False}
        )
        self.save()

    def mark_cleaned(self, identifier: str) -> None:
        for resource in self.resources:
            if resource["identifier"] == identifier:
                resource["cleaned"] = True
        self.save()

    def save(self) -> None:
        self.config.artifact_dir.mkdir(parents=True, exist_ok=True)
        self.config.ledger_path.write_text(
            json.dumps(
                {
                    "run_id": self.config.run_id,
                    "environment": self.config.safe_environment(),
                    "resources": self.resources,
                },
                ensure_ascii=False,
                indent=2,
            ),
            encoding="utf-8",
        )

    def cleanup(self) -> dict[str, Any]:
        if self.config.keep:
            self.save()
            return {"status": "kept", "remaining": self.resources}

        failures: list[dict[str, Any]] = []
        for resource in reversed(self.resources):
            try:
                _cleanup_resource(resource)
                resource["cleaned"] = True
            except Exception as exc:  # noqa: BLE001
                failures.append(
                    {
                        "identifier": resource["identifier"],
                        "error": repr(exc),
                    }
                )
        self.save()
        return {
            "status": "failed" if failures else "cleaned",
            "remaining": [r for r in self.resources if not r["cleaned"]],
            "failures": failures,
        }


def new_run_id(prefix: str = "eval") -> str:
    return f"{prefix}-{uuid.uuid4().hex[:12]}"


def build_config(
    root: Path,
    *,
    run_id: str | None = None,
    live: bool = False,
    sandbox: bool = False,
    keep: bool = False,
    resources: dict[str, str] | None = None,
) -> SandboxConfig:
    resolved_root = root.resolve()
    rid = run_id or new_run_id()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]{2,63}", rid):
        raise SandboxSafetyError("run_id must contain only letters, digits, '_' or '-'")
    if live and not sandbox:
        raise SandboxSafetyError("live evaluation requires --sandbox")
    if live and os.getenv("JINRONG_EVAL") != "1":
        raise SandboxSafetyError("live evaluation requires JINRONG_EVAL=1")
    if live:
        validate_live_resources(resources or {})
    return SandboxConfig(
        run_id=rid,
        root=resolved_root,
        live=live,
        keep=keep,
        env_name="sandbox-live" if live else "offline",
        resources=dict(resources or {}),
    )


def validate_live_resources(resources: dict[str, str]) -> None:
    required = ("mysql_database", "mysql_core_database", "redis_url", "milvus_uri")
    missing = [key for key in required if not resources.get(key)]
    if missing:
        raise SandboxSafetyError(f"missing explicit sandbox resources: {', '.join(missing)}")

    forbidden_names = {"jinrong_agent", "jinrong_core", "production", "prod", "development", "dev"}
    for key in ("mysql_database", "mysql_core_database"):
        value = resources[key].lower()
        if value in forbidden_names or not re.search(r"(?:eval|sandbox)", value):
            raise SandboxSafetyError(f"{key} must contain eval or sandbox and not be a default database")

    redis_url = resources["redis_url"].lower()
    if "/0" in redis_url and "eval" not in redis_url and "sandbox" not in redis_url:
        raise SandboxSafetyError("redis_url must identify a dedicated eval/sandbox database")

    milvus_uri = Path(resources["milvus_uri"]).resolve()
    if "eval" not in str(milvus_uri).lower() and "sandbox" not in str(milvus_uri).lower():
        raise SandboxSafetyError("milvus_uri must point to an eval/sandbox path")


def write_json(config: SandboxConfig, name: str, payload: Any) -> Path:
    if Path(name).name != name or not name.endswith(".json"):
        raise ValueError("artifact name must be a flat .json filename")
    config.artifact_dir.mkdir(parents=True, exist_ok=True)
    path = config.artifact_dir / name
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
    return path


def write_jsonl(config: SandboxConfig, name: str, rows: list[dict[str, Any]]) -> Path:
    if Path(name).name != name or not name.endswith(".jsonl"):
        raise ValueError("artifact name must be a flat .jsonl filename")
    config.artifact_dir.mkdir(parents=True, exist_ok=True)
    path = config.artifact_dir / name
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
    return path


def hash_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16]


def redact_input(value: str, limit: int = 160) -> dict[str, Any]:
    text = value or ""
    return {"sha256_16": hash_text(text), "length": len(text), "preview": text[:limit]}


def _redact_resource(value: str) -> str:
    text = str(value)
    if "@" in text and "://" in text:
        prefix, suffix = text.split("@", 1)
        if "://" in prefix:
            scheme, auth = prefix.split("://", 1)
            if ":" in auth:
                user = auth.split(":", 1)[0]
                return f"{scheme}://{user}:***@{suffix}"
    return text


def _cleanup_resource(resource: dict[str, Any]) -> None:
    kind = resource.get("kind")
    identifier = resource.get("identifier", "")
    if kind == "directory":
        path = Path(identifier).resolve()
        if path.name.startswith("eval-") and path.exists():
            shutil.rmtree(path)
        return
    if kind == "file":
        path = Path(identifier).resolve()
        if path.name.startswith("eval-") and path.exists():
            path.unlink()
        return
    if kind in {"mysql_database", "redis_key", "milvus_collection"}:
        raise SandboxSafetyError(f"resource cleanup adapter not configured for {kind}: {identifier}")
    raise SandboxSafetyError(f"unknown resource kind: {kind}")


def cleanup_artifact_dir(config: SandboxConfig) -> None:
    if config.keep or not config.artifact_dir.exists():
        return
    if config.artifact_dir.name != config.run_id:
        raise SandboxSafetyError("refusing to remove an unexpected artifact directory")
    shutil.rmtree(config.artifact_dir)