releaseops_arena/server.py

import csv
import io
import json
import os
import subprocess
import sys
import tarfile
import threading
import time
import uuid
from pathlib import Path

from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.responses import FileResponse, HTMLResponse, PlainTextResponse, Response, StreamingResponse

from releaseops_arena.tool_env import ReleaseOpsToolEnv
from releaseops_arena.space_paths import get_outputs_root

SUPPORTS_CONCURRENT_SESSIONS: bool = True
MAX_CONCURRENT_ENVS = int(os.getenv("MAX_CONCURRENT_ENVS", "64"))

app = FastAPI(title="ReleaseOps Arena Env")
from releaseops_arena.eval_api import router as eval_router
app.include_router(eval_router, prefix="/api")


_env_sessions: dict[str, ReleaseOpsToolEnv] = {}
_env_lock = threading.Lock()
_training_lock = threading.Lock()
_training_process: subprocess.Popen | None = None
_training_started_at: float | None = None
_training_command: list[str] | None = None
_training_mode: str | None = None

OUTPUTS_ROOT = get_outputs_root()
_DEMO_DIR = Path(__file__).resolve().parent.parent / "demo"
_FLOW_HTML = _DEMO_DIR / "releaseops_episode_flow.html"
TRAINING_LOG_PATH = OUTPUTS_ROOT / "space_training.log"
SMOKE_METRICS_PATH = OUTPUTS_ROOT / "grpo_env_smoke_metrics.json"
PILOT_METRICS_PATH = OUTPUTS_ROOT / "grpo_env_metrics_pilot.json"


@app.get("/")
def root():
    return {
        "name": "ReleaseOps Arena Env",
        "ok": True,
        "message": "Use /health for status, /reset and /step for env rollouts, or /train/smoke to start a GRPO smoke run.",
        "docs": "/docs",
        "ui": {
            "episode_flow": "/ui/flow",
        },
        "notebook": {
            "on_hub": "https://huggingface.co/spaces/hiitsesh/New_gpu_space/blob/main/notebooks/ReleaseOps_final_walkthrough.ipynb",
            "raw": "https://huggingface.co/spaces/hiitsesh/New_gpu_space/resolve/main/notebooks/ReleaseOps_final_walkthrough.ipynb",
        },
        "pitch_video": "https://www.youtube.com/shorts/OxfBH7jDOwg",
        "training": {
            "start_smoke": "/train/smoke",
            "start_pilot": "/train/pilot?max_steps=100",
            "status": "/train/status",
            "kill": "/train/kill",
            "logs": "/train/logs",
            "metrics": "/train/metrics",
            "summary": "/train/summary",
            "live_metrics": "/train/live",
            "live_plot_png": "/train/plot.png",
            "live_plot_html": "/train/plot",
            "push_to_hub": "/train/push_to_hub?repo_id=your-user/releaseops-grpo-artifacts",
        },
        "artifacts": {
            "list": "/outputs/ls",
            "download_file": "/outputs/file?path=grpo_env_metrics_pilot.json",
            "download_archive": "/outputs/archive?path=releaseops-grpo-pilot",
            "write_test_csv": "/outputs/write_test",
        },
        "persistence": {
            "outputs_root": str(OUTPUTS_ROOT),
            "warning": (
                "Without Space Storage, local disk is ephemeral on restart. If you mount `/data` "
                f"(Storage Buckets), this server writes under `{OUTPUTS_ROOT}` so logs and "
                "checkpoints survive sleep. Still recommend Hub push for backup."
            ),
            "push_to_hub": "/train/push_to_hub?repo_id=YOUR_USER/dataset-repo (set HF_TOKEN secret)",
            "sleep": "Space Settings → Sleep time: longer idle window reduces surprise restarts.",
        },
    }


@app.get("/ui/flow")
def episode_flow_explanation():
    """
    Mermaid flowchart (episode + n8n-style CI). Needs outbound HTTPS for cdn.jsdelivr.net / fonts.
    """
    if not _FLOW_HTML.is_file():
        raise HTTPException(status_code=404, detail="demo/releaseops_episode_flow.html not in image")
    return FileResponse(
        str(_FLOW_HTML),
        media_type="text/html; charset=utf-8",
    )


@app.get("/health")
def health():
    return {
        "ok": True,
        "active_sessions": len(_env_sessions),
        "max_concurrent_envs": MAX_CONCURRENT_ENVS,
    }


def _tail_text(path: Path, max_lines: int = 80) -> str:
    if not path.exists():
        return ""
    with open(path, "r", encoding="utf-8", errors="replace") as handle:
        lines = handle.readlines()
    return "".join(lines[-max_lines:])


def _read_metrics(path: Path):
    if not path.exists():
        return None
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _summarize_metrics(rows):
    if not rows:
        return None
    reward_rows = [row for row in rows if "reward" in row]
    if not reward_rows:
        return {"reward_points": 0}
    first = reward_rows[0]
    last = reward_rows[-1]
    return {
        "reward_points": len(reward_rows),
        "first_step": first.get("step"),
        "last_step": last.get("step"),
        "reward_first": first.get("reward"),
        "reward_last": last.get("reward"),
        "reward_delta": (
            last.get("reward") - first.get("reward")
            if isinstance(first.get("reward"), (int, float))
            and isinstance(last.get("reward"), (int, float))
            else None
        ),
        "tool_call_frequency_first": first.get("tools/call_frequency"),
        "tool_call_frequency_last": last.get("tools/call_frequency"),
        "tool_failure_frequency_first": first.get("tools/failure_frequency"),
        "tool_failure_frequency_last": last.get("tools/failure_frequency"),
        "loss_last": last.get("loss"),
        "train_runtime": rows[-1].get("train_runtime"),
    }


def _training_state():
    global _training_process
    running = _training_process is not None and _training_process.poll() is None
    exit_code = None if _training_process is None else _training_process.poll()
    elapsed_seconds = None
    if _training_started_at is not None:
        elapsed_seconds = round(time.time() - _training_started_at, 1)

    return {
        "running": running,
        "exit_code": exit_code,
        "mode": _training_mode,
        "command": _training_command,
        "elapsed_seconds": elapsed_seconds,
        "log_path": str(TRAINING_LOG_PATH),
        "recent_log": _tail_text(TRAINING_LOG_PATH),
    }


def _start_training(command: list[str], mode: str):
    global _training_process, _training_started_at, _training_command, _training_mode
    with _training_lock:
        if _training_process is not None and _training_process.poll() is None:
            return {
                "started": False,
                "message": "Training is already running.",
                **_training_state(),
            }

        TRAINING_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
        log_handle = open(TRAINING_LOG_PATH, "w", encoding="utf-8", buffering=1)
        log_handle.write(f"Starting {mode} training at {time.ctime()}\n")
        log_handle.write("Command: " + " ".join(command) + "\n\n")
        log_handle.flush()

        # Line-buffer log file + unbuffered Python child so /train/logs updates live (not stuck until process exit).
        train_env = os.environ.copy()
        train_env["PYTHONUNBUFFERED"] = "1"
        # Align with train_grpo: avoid torch._dynamo mega-cache duplicate registration
        # on some Space images, and keep compile off for GRPO.
        train_env.setdefault("TORCH_COMPILE_DISABLE", "1")
        # Inductor calls getpass.getuser() for default cache path; uids in Space
        # often have no /etc/passwd entry — fixed dirs under OUTPUTS avoid KeyError.
        train_env.setdefault(
            "TORCHINDUCTOR_CACHE_DIR",
            str(OUTPUTS_ROOT / ".torch_inductor_cache"),
        )
        train_env.setdefault("TRITON_CACHE_DIR", str(OUTPUTS_ROOT / "triton_cache"))
        # HuggingFace datasets / hub: default ~/.cache can become /.cache when HOME is "/".
        _hf = str(OUTPUTS_ROOT / ".hf")
        train_env.setdefault("HF_HOME", _hf)
        train_env.setdefault("HUGGINGFACE_HUB_CACHE", f"{_hf}/hub")
        train_env.setdefault("HF_DATASETS_CACHE", f"{_hf}/datasets")
        train_env.setdefault("TRANSFORMERS_CACHE", f"{_hf}/transformers")
        train_env.setdefault("XDG_CACHE_HOME", f"{_hf}/xdg")
        if (train_env.get("HOME") or "").strip() in ("", "/"):
            train_env["HOME"] = f"{_hf}/user_home"

        _training_process = subprocess.Popen(
            command,
            stdout=log_handle,
            stderr=subprocess.STDOUT,
            text=True,
            env=train_env,
        )
        _training_started_at = time.time()
        _training_command = command
        _training_mode = mode

    return {
        "started": True,
        "message": "Training started. Poll /train/status or /train/logs for progress.",
        **_training_state(),
    }


@app.api_route("/train/smoke", methods=["GET", "POST"])
def train_smoke():
    command = [
        sys.executable,
        "training/train_grpo.py",
        "--smoke",
        "--max-steps",
        "8",
        "--num-generations",
        "2",
        "--output-dir",
        str(OUTPUTS_ROOT / "releaseops-grpo"),
        "--metrics-json",
        str(SMOKE_METRICS_PATH),
    ]
    return _start_training(command, "smoke")


@app.api_route("/train/pilot", methods=["GET", "POST"])
def train_pilot(
    max_steps: int = 100,
    num_generations: int = 4,
    gradient_accumulation_steps: int = 4,
    max_completion_length: int = 512,
    learning_rate: float = 1e-5,
    logging_steps: int = 5,
    model_name: str = "Qwen/Qwen3-0.6B",
    bf16: bool = False,
    best_loss_dir: str = "",
    hub_model_repo: str = "",
    hub_upload_include: str = "best",
):
    bld = (best_loss_dir or "").strip() or str(OUTPUTS_ROOT / "best_by_loss")
    max_steps = max(1, min(max_steps, 500))
    num_generations = max(2, min(num_generations, 16))
    gradient_accumulation_steps = max(1, min(gradient_accumulation_steps, 32))
    max_completion_length = max(128, min(max_completion_length, 2048))
    logging_steps = max(1, min(logging_steps, 50))
    command = [
        sys.executable,
        "training/train_grpo.py",
        "--model-name",
        model_name,
        "--max-steps",
        str(max_steps),
        "--num-generations",
        str(num_generations),
        "--gradient-accumulation-steps",
        str(gradient_accumulation_steps),
        "--max-completion-length",
        str(max_completion_length),
        "--learning-rate",
        str(learning_rate),
        "--logging-steps",
        str(logging_steps),
        "--output-dir",
        str(OUTPUTS_ROOT / "releaseops-grpo-pilot"),
        "--metrics-json",
        str(PILOT_METRICS_PATH),
        "--best-loss-dir",
        bld,
    ]
    if bf16:
        command.append("--bf16")
    if hub_model_repo.strip():
        command.extend(["--hub-model-repo", hub_model_repo.strip()])
        if hub_upload_include in ("best", "final", "both"):
            command.extend(["--hub-upload-include", hub_upload_include])
    return _start_training(command, "pilot")


@app.get("/train/status")
def train_status():
    return _training_state()


@app.api_route("/train/kill", methods=["GET", "POST"])
def train_kill():
    """Terminate the currently running training process, if any."""
    global _training_process
    with _training_lock:
        if _training_process is None or _training_process.poll() is not None:
            return {"killed": False, "message": "No active training process.", **_training_state()}
        try:
            _training_process.terminate()
            try:
                _training_process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                _training_process.kill()
                _training_process.wait(timeout=5)
        except Exception as exc:
            return {"killed": False, "message": f"Failed to kill: {exc}", **_training_state()}

    return {"killed": True, "message": "Training process terminated.", **_training_state()}


@app.get("/train/logs", response_class=PlainTextResponse)
def train_logs(lines: int = 200):
    lines = max(20, min(lines, 2000))
    body = _tail_text(TRAINING_LOG_PATH, max_lines=lines) or "No training log yet. Start /train/smoke first."
    return PlainTextResponse(content=body, media_type="text/plain; charset=utf-8")


@app.get("/train/metrics")
def train_metrics():
    return {
        "smoke": _read_metrics(SMOKE_METRICS_PATH),
        "pilot": _read_metrics(PILOT_METRICS_PATH),
    }


@app.get("/train/summary")
def train_summary():
    smoke = _read_metrics(SMOKE_METRICS_PATH)
    pilot = _read_metrics(PILOT_METRICS_PATH)
    return {
        "status": _training_state(),
        "smoke": _summarize_metrics(smoke),
        "pilot": _summarize_metrics(pilot),
    }


def _latest_metrics_path() -> Path | None:
    """Pick whichever metrics file was most recently updated, for live plotting."""
    candidates = [p for p in (PILOT_METRICS_PATH, SMOKE_METRICS_PATH) if p.exists()]
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def _extract_series(rows):
    """Extract reward, loss, tool-call, tool-failure series from a TRL log_history."""
    steps, rewards, losses, tcalls, tfails, completions = [], [], [], [], [], []
    for row in rows or []:
        step = row.get("step")
        if step is None:
            continue
        if "reward" in row:
            steps.append(step)
            rewards.append(row.get("reward"))
            losses.append(row.get("loss"))
            tcalls.append(row.get("tools/call_frequency"))
            tfails.append(row.get("tools/failure_frequency"))
            completions.append(row.get("completions/mean_length"))
    return {
        "steps": steps,
        "reward": rewards,
        "loss": losses,
        "tool_call_frequency": tcalls,
        "tool_failure_frequency": tfails,
        "completion_mean_length": completions,
    }


@app.get("/train/live")
def train_live():
    """Return the current live log history series (available after every logging step)."""
    latest = _latest_metrics_path()
    if latest is None:
        return {"available": False, "message": "No metrics yet. Start /train/smoke or /train/pilot."}
    rows = _read_metrics(latest) or []
    series = _extract_series(rows)
    return {
        "available": True,
        "metrics_source": str(latest),
        "num_points": len(series["steps"]),
        "series": series,
        "status": _training_state(),
    }


@app.get("/train/plot.png")
def train_plot_png():
    """Render the current training curves as a PNG image."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise HTTPException(status_code=500, detail="matplotlib not installed in Space.") from exc

    latest = _latest_metrics_path()
    if latest is None:
        raise HTTPException(status_code=404, detail="No metrics yet.")

    rows = _read_metrics(latest) or []
    series = _extract_series(rows)
    if not series["steps"]:
        raise HTTPException(status_code=404, detail="No reward points logged yet.")

    fig, axes = plt.subplots(3, 1, figsize=(9, 9), sharex=True)
    steps = series["steps"]

    axes[0].plot(steps, series["reward"], marker="o", color="#2563eb", label="reward")
    axes[0].axhline(0, color="#999", linewidth=0.5, linestyle="--")
    axes[0].set_ylabel("reward (mean)")
    axes[0].legend(loc="upper left")
    axes[0].grid(alpha=0.3)

    loss_vals = [v for v in series["loss"] if v is not None]
    if loss_vals:
        axes[1].plot(steps, series["loss"], marker="o", color="#dc2626", label="loss")
        axes[1].set_ylabel("loss")
        axes[1].legend(loc="upper left")
        axes[1].grid(alpha=0.3)
    else:
        axes[1].text(0.5, 0.5, "loss not yet logged", ha="center", va="center", transform=axes[1].transAxes)

    axes[2].plot(steps, series["tool_call_frequency"], marker="o", color="#16a34a", label="tool_call_freq")
    axes[2].plot(steps, series["tool_failure_frequency"], marker="o", color="#ea580c", label="tool_failure_freq")
    axes[2].set_ylabel("tool usage")
    axes[2].set_xlabel("step")
    axes[2].legend(loc="upper left")
    axes[2].grid(alpha=0.3)

    mode = _training_mode or "idle"
    elapsed = _training_state().get("elapsed_seconds")
    fig.suptitle(
        f"ReleaseOps GRPO training  |  mode={mode}  |  elapsed={elapsed}s  |  points={len(steps)}",
        fontsize=11,
    )
    fig.tight_layout(rect=[0, 0, 1, 0.97])

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110)
    plt.close(fig)
    buf.seek(0)
    return Response(content=buf.getvalue(), media_type="image/png")


@app.get("/train/plot", response_class=HTMLResponse)
def train_plot_html(refresh: int = 20):
    """Auto-refreshing HTML page that shows /train/plot.png and current summary."""
    refresh = max(5, min(refresh, 300))
    html = f"""<!doctype html>
<html><head>
<meta charset="utf-8" />
<meta http-equiv="refresh" content="{refresh}" />
<title>ReleaseOps GRPO Live Metrics</title>
<style>
  body {{ font-family: system-ui, -apple-system, sans-serif; background:#0b1020; color:#e2e8f0; margin:0; padding:20px; }}
  h1 {{ font-size:18px; margin:0 0 12px 0; color:#93c5fd; }}
  img {{ max-width:100%; height:auto; border:1px solid #1e293b; border-radius:6px; background:white; }}
  .note {{ font-size:12px; color:#94a3b8; margin-top:8px; }}
  a {{ color:#60a5fa; }}
  code {{ background:#1e293b; padding:1px 6px; border-radius:4px; font-size:12px; }}
</style>
</head>
<body>
  <h1>ReleaseOps GRPO — live training metrics</h1>
  <div class="note">
    auto-refresh every {refresh}s &middot;
    <a href="/train/summary">/train/summary</a> &middot;
    <a href="/train/live">/train/live</a> &middot;
    <a href="/train/logs">/train/logs</a> &middot;
    <a href="/outputs/ls">/outputs/ls</a>
  </div>
  <p><img src="/train/plot.png?t={int(time.time())}" alt="training plot" /></p>
  <div class="note">
    Source JSON: <code>outputs/grpo_env_metrics_pilot.json</code> (updated every logging step).
  </div>
</body></html>
"""
    return HTMLResponse(content=html)


def _safe_outputs_path(relative_path: str) -> Path:
    """Resolve a user-provided relative path, but only allow targets under outputs/."""
    if not relative_path:
        raise HTTPException(status_code=400, detail="Missing path parameter.")
    candidate = (OUTPUTS_ROOT / relative_path).resolve()
    try:
        candidate.relative_to(OUTPUTS_ROOT)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail="Path must stay under outputs/.") from exc
    return candidate


def _walk_outputs():
    if not OUTPUTS_ROOT.exists():
        return []
    entries = []
    for path in sorted(OUTPUTS_ROOT.rglob("*")):
        try:
            stat = path.stat()
        except FileNotFoundError:
            continue
        rel = path.relative_to(OUTPUTS_ROOT).as_posix()
        entries.append(
            {
                "path": rel,
                "is_dir": path.is_dir(),
                "size_bytes": stat.st_size if path.is_file() else None,
                "mtime": stat.st_mtime,
            }
        )
    return entries


@app.get("/outputs/ls")
def outputs_ls():
    """List all files and folders under outputs/ that exist in the running container."""
    return {
        "outputs_root": str(OUTPUTS_ROOT),
        "entries": _walk_outputs(),
    }


WRITE_PROBE_FILENAME = "simpllll.csv"


@app.api_route("/outputs/write_test", methods=["GET", "POST"])
def outputs_write_test(request: Request):
    """
    Create a tiny CSV under outputs/ to verify the Space can write to disk (same volume as
    training checkpoints and push_to_hub sources). Does not use the token; only reports whether
    any Hub token env var is set (boolean, never the value).
    """
    try:
        OUTPUTS_ROOT.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        raise HTTPException(
            status_code=500, detail=f"Could not create outputs root: {exc}"
        ) from exc

    target = OUTPUTS_ROOT / WRITE_PROBE_FILENAME
    ts = time.time()
    try:
        with open(target, "w", newline="", encoding="utf-8") as handle:
            writer = csv.writer(handle)
            writer.writerow(
                [
                    "status",
                    "message",
                    "unix_time",
                    "outputs_root",
                ]
            )
            writer.writerow(
                [
                    "ok",
                    "disk_write_probe",
                    f"{ts:.3f}",
                    str(OUTPUTS_ROOT),
                ]
            )
    except OSError as exc:
        raise HTTPException(
            status_code=500, detail=f"Write failed: {exc}"
        ) from exc

    try:
        stat = target.stat()
    except OSError as exc:
        raise HTTPException(
            status_code=500, detail=f"Stat after write failed: {exc}"
        ) from exc

    token_present = bool(
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    )
    base = str(request.base_url).rstrip("/")
    download_url = f"{base}/outputs/file?path={WRITE_PROBE_FILENAME}"
    return {
        "wrote": True,
        "path_relative": WRITE_PROBE_FILENAME,
        "path_absolute": str(target),
        "size_bytes": stat.st_size,
        "mtime": stat.st_mtime,
        "hf_token_env_set": token_present,
        "fetch_csv": f"/outputs/file?path={WRITE_PROBE_FILENAME}",
        "download_url": download_url,
        "list_files": f"{base}/outputs/ls",
        "not_in_repo_file_browser": (
            "huggingface.co/spaces/.../tree shows GIT files only. This file lives on the live "
            "container disk. Open download_url in a browser or curl it; use list_files to see "
            "runtime outputs."
        ),
    }


@app.get("/outputs/file")
def outputs_file(path: str = Query(..., description="Path relative to outputs/ root.")):
    """Download a single file from the outputs/ directory."""
    target = _safe_outputs_path(path)
    if not target.exists() or not target.is_file():
        raise HTTPException(status_code=404, detail=f"File not found: {path}")
    return FileResponse(
        str(target),
        media_type="application/octet-stream",
        filename=target.name,
    )


@app.get("/outputs/archive")
def outputs_archive(path: str = Query(..., description="Directory path relative to outputs/ root.")):
    """Stream a tar.gz of a directory inside outputs/ (useful for downloading a checkpoint)."""
    target = _safe_outputs_path(path)
    if not target.exists() or not target.is_dir():
        raise HTTPException(status_code=404, detail=f"Directory not found: {path}")

    def _iter_archive():
        buffer = io.BytesIO()
        with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
            tar.add(str(target), arcname=target.name)
        buffer.seek(0)
        while True:
            chunk = buffer.read(65536)
            if not chunk:
                break
            yield chunk

    filename = f"{target.name}.tar.gz"
    return StreamingResponse(
        _iter_archive(),
        media_type="application/gzip",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@app.api_route("/train/push_to_hub", methods=["GET", "POST"])
def push_artifacts_to_hub(
    repo_id: str = Query(..., description="Target HF repo, e.g. your-user/releaseops-grpo-artifacts"),
    repo_type: str = Query("dataset", description="Repo type: dataset (recommended) or model."),
    path: str = Query("", description="Optional sub-path under outputs/ to push. Empty pushes all of outputs/."),
):
    """
    Upload outputs/ (or a subpath) to a Hugging Face Hub repo for durable storage.

    Requires HF_TOKEN env var to be set on the Space with write permission to repo_id.
    """
    token = (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    )
    if not token:
        raise HTTPException(
            status_code=400,
            detail=(
                "No HF token found. In the Space: Settings → Repository secrets → add "
                "HF_TOKEN (or HUGGINGFACE_HUB_TOKEN) with write access to the target repo, "
                "then restart the Space and retry."
            ),
        )

    if repo_type not in {"dataset", "model", "space"}:
        raise HTTPException(status_code=400, detail="repo_type must be dataset, model, or space.")

    source = _safe_outputs_path(path) if path else OUTPUTS_ROOT
    if not source.exists():
        raise HTTPException(status_code=404, detail=f"Source path does not exist: {path or 'outputs/'}")

    try:
        from huggingface_hub import HfApi, create_repo
    except ImportError as exc:
        raise HTTPException(
            status_code=500,
            detail="huggingface_hub not installed in the Space image.",
        ) from exc

    api = HfApi(token=token)
    try:
        create_repo(repo_id, token=token, repo_type=repo_type, exist_ok=True, private=False)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"create_repo failed: {exc}") from exc

    commit_message = f"Upload artifacts from Space at {time.ctime()}"
    try:
        if source.is_file():
            api.upload_file(
                path_or_fileobj=str(source),
                path_in_repo=source.name,
                repo_id=repo_id,
                repo_type=repo_type,
                commit_message=commit_message,
            )
            uploaded = {"files": 1, "source": str(source)}
        else:
            api.upload_folder(
                folder_path=str(source),
                repo_id=repo_id,
                repo_type=repo_type,
                path_in_repo=path or "",
                commit_message=commit_message,
                ignore_patterns=["**/.git/**", "**/__pycache__/**"],
            )
            uploaded = {"folder": str(source)}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"upload failed: {exc}") from exc

    return {
        "pushed": True,
        "repo_id": repo_id,
        "repo_type": repo_type,
        "source": str(source),
        "commit_message": commit_message,
        "repo_url": f"https://huggingface.co/{'datasets/' if repo_type == 'dataset' else ''}{repo_id}",
        **uploaded,
    }


@app.post("/reset")
def reset(params: dict):
    with _env_lock:
        if len(_env_sessions) >= MAX_CONCURRENT_ENVS:
            raise HTTPException(
                status_code=429,
                detail=(
                    f"Maximum concurrent environments reached: {MAX_CONCURRENT_ENVS}. "
                    "Close an environment before creating a new one."
                ),
            )

        env = ReleaseOpsToolEnv()
        observation = env.reset(**params)
        env_id = str(uuid.uuid4())
        _env_sessions[env_id] = env

    return {
        "env_id": env_id,
        "observation": json.loads(observation),
        "reward": env.reward,
        "done": env.done,
    }


@app.post("/step")
def step(action: dict):
    env_id = action.get("env_id")
    tool = action.get("tool")
    arguments = action.get("arguments", {})

    if not env_id:
        raise HTTPException(status_code=400, detail="Missing required field: env_id")
    if not tool:
        raise HTTPException(status_code=400, detail="Missing required field: tool")
    if not isinstance(arguments, dict):
        raise HTTPException(status_code=400, detail="Field 'arguments' must be an object")

    with _env_lock:
        env = _env_sessions.get(env_id)

    if env is None:
        raise HTTPException(status_code=404, detail=f"Unknown env_id: {env_id}")

    if tool.startswith("_") or not hasattr(env, tool):
        raise HTTPException(status_code=400, detail=f"Unknown tool: {tool}")

    method = getattr(env, tool)
    if not callable(method):
        raise HTTPException(status_code=400, detail=f"Tool is not callable: {tool}")

    try:
        result = method(**arguments)
    except TypeError as exc:
        raise HTTPException(status_code=400, detail=f"Invalid arguments for {tool}: {exc}") from exc
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    parsed_result = result
    if isinstance(result, str):
        try:
            parsed_result = json.loads(result)
        except json.JSONDecodeError:
            parsed_result = result

    response = {
        "env_id": env_id,
        "result": parsed_result,
        "observation": json.loads(env.render_observation()),
        "reward": env.reward,
        "done": env.done,
        "terminal_reason": env.state.get("terminal_reason") if env.state else None,
    }

    if env.done:
        with _env_lock:
            _env_sessions.pop(env_id, None)

    return response


@app.post("/close")
def close(payload: dict):
    env_id = payload.get("env_id")
    if not env_id:
        raise HTTPException(status_code=400, detail="Missing required field: env_id")

    with _env_lock:
        removed = _env_sessions.pop(env_id, None)

    if removed is None:
        raise HTTPException(status_code=404, detail=f"Unknown env_id: {env_id}")

    return {"env_id": env_id, "closed": True}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("releaseops_arena.server:app", host="0.0.0.0", port=8000)