hermes_cli/hf_space_update.py

from __future__ import annotations

import argparse
import json
import os
import subprocess
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Iterable

OFFICIAL_UPSTREAM_REMOTE = "https://github.com/NousResearch/hermes-agent.git"
PROJECT_ROOT = Path(__file__).resolve().parent.parent
UPSTREAM_REF_FILE = PROJECT_ROOT / "docker" / "HERMES_UPSTREAM_REF"
DEFAULT_WAIT_TIMEOUT = 1800
POLL_INTERVAL_SECONDS = 5


def run(cmd, **kwargs):
    kwargs.setdefault("check", True)
    return subprocess.run(cmd, **kwargs)


def is_hf_space_runtime(env: dict[str, str] | None = None) -> bool:
    env = env or os.environ
    return bool(env.get("SPACE_ID") or env.get("SPACE_HOST"))


def get_hf_token(env: dict[str, str] | None = None) -> str | None:
    env = env or os.environ
    for key in ("HF_TOKEN", "HUGGINGFACE_TOKEN", "HF_API_TOKEN"):
        value = env.get(key)
        if value:
            return value
    return None


def list_remote_tags(remote: str = OFFICIAL_UPSTREAM_REMOTE) -> list[str]:
    result = run(
        ["git", "ls-remote", "--tags", "--refs", "--sort=-v:refname", remote],
        capture_output=True,
        text=True,
    )
    tags: list[str] = []
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) != 2:
            continue
        ref = parts[1]
        prefix = "refs/tags/"
        if ref.startswith(prefix):
            tags.append(ref[len(prefix) :])
    return tags


def pick_target_ref(requested_ref: str | None) -> str:
    if requested_ref:
        return requested_ref
    tags = list_remote_tags()
    if not tags:
        raise RuntimeError("Could not determine the latest upstream Hermes tag.")
    return tags[0]


def build_space_clone_url(space_id: str, token: str) -> str:
    return f"https://oauth2:{token}@huggingface.co/spaces/{space_id}"


def clone_space_repo(space_id: str, token: str, workdir: Path | None = None) -> Path:
    root = Path(workdir) if workdir else Path(tempfile.mkdtemp(prefix="hermes-space-update-"))
    repo_dir = root / "repo"
    run(["git", "clone", "--depth", "1", build_space_clone_url(space_id, token), str(repo_dir)])
    return repo_dir


def read_current_ref(path: Path = UPSTREAM_REF_FILE) -> str | None:
    if not path.exists():
        return None
    value = path.read_text(encoding="utf-8").strip()
    return value or None


def apply_space_repo_update(
    space_id: str,
    token: str,
    target_ref: str,
    workdir: Path | None = None,
) -> str | None:
    repo_dir = clone_space_repo(space_id, token, workdir=workdir)
    ref_file = repo_dir / "docker" / "HERMES_UPSTREAM_REF"
    current_ref = read_current_ref(ref_file)
    if current_ref == target_ref:
        return None

    ref_file.parent.mkdir(parents=True, exist_ok=True)
    ref_file.write_text(f"{target_ref}\n", encoding="utf-8")

    run(["git", "config", "user.email", "hermes-space-updater@local"], cwd=repo_dir)
    run(["git", "config", "user.name", "Hermes Space Updater"], cwd=repo_dir)
    run(["git", "add", "docker/HERMES_UPSTREAM_REF"], cwd=repo_dir)
    run(["git", "commit", "-m", f"Update Hermes upstream ref to {target_ref}"], cwd=repo_dir)
    run(["git", "push", "origin", "HEAD:main"], cwd=repo_dir)
    result = run(["git", "rev-parse", "HEAD"], cwd=repo_dir, capture_output=True, text=True)
    return result.stdout.strip()


def _runtime_request(space_id: str, token: str | None = None):
    request = urllib.request.Request(f"https://huggingface.co/api/spaces/{space_id}/runtime")
    if token:
        request.add_header("Authorization", f"Bearer {token}")
    with urllib.request.urlopen(request, timeout=30) as response:
        return json.loads(response.read().decode("utf-8"))


def wait_for_space_runtime(
    space_id: str,
    expected_sha: str,
    token: str | None = None,
    timeout: int = DEFAULT_WAIT_TIMEOUT,
    poll_interval: int = POLL_INTERVAL_SECONDS,
) -> dict:
    deadline = time.time() + timeout
    last_stage = None
    while time.time() < deadline:
        payload = _runtime_request(space_id, token)
        stage = payload.get("stage")
        sha = payload.get("sha")
        if stage != last_stage:
            print(f"• Space stage: {stage} (sha={sha})")
            last_stage = stage
        if sha == expected_sha and stage == "RUNNING":
            return payload
        if stage in {"BUILD_ERROR", "RUNTIME_ERROR", "CONFIG_ERROR"}:
            raise RuntimeError(f"Space rebuild failed with stage {stage} (sha={sha})")
        time.sleep(poll_interval)
    raise TimeoutError(f"Timed out waiting for Space runtime to reach {expected_sha}")


def _gateway_state_paths(env: dict[str, str] | None = None) -> tuple[Path, Path]:
    env = env or os.environ
    hermes_home = Path(env.get("HERMES_HOME", str(Path.home() / ".hermes")))
    return hermes_home / ".update_output.txt", hermes_home / ".update_exit_code"


def write_gateway_exit_code(exit_code: int, env: dict[str, str] | None = None) -> None:
    _output_path, exit_code_path = _gateway_state_paths(env)
    exit_code_path.write_text(str(exit_code), encoding="utf-8")


def run_update_cli(argv: Iterable[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="hermes update",
        description="Update a Hugging Face Space by bumping its upstream Hermes ref and triggering a rebuild.",
    )
    parser.add_argument("--gateway", action="store_true", default=False)
    parser.add_argument("--ref", help="Update to a specific upstream tag or commit instead of the latest tag.")
    parser.add_argument("--no-wait", action="store_true", default=False, help="Push the Space update and exit without waiting for rebuild completion.")
    parser.add_argument("--timeout", type=int, default=DEFAULT_WAIT_TIMEOUT)
    args = parser.parse_args(list(argv) if argv is not None else None)

    try:
        if not is_hf_space_runtime():
            print("✗ HF Space update flow is only available inside a Hugging Face Space runtime.")
            if args.gateway:
                write_gateway_exit_code(1)
            return 1

        space_id = os.environ.get("SPACE_ID")
        if not space_id:
            print("✗ SPACE_ID is not set; cannot locate the current Hugging Face Space.")
            if args.gateway:
                write_gateway_exit_code(1)
            return 1

        token = get_hf_token()
        if not token:
            print("✗ HF_TOKEN (or HUGGINGFACE_TOKEN) is required to update the Space repository.")
            if args.gateway:
                write_gateway_exit_code(1)
            return 1

        target_ref = pick_target_ref(args.ref)
        current_ref = read_current_ref()
        if current_ref == target_ref:
            print(f"Already up to date: {target_ref}")
            if args.gateway:
                write_gateway_exit_code(0)
            return 0

        print(f"⚕ Updating Hugging Face Space {space_id} to upstream ref {target_ref}...")
        commit_sha = apply_space_repo_update(space_id=space_id, token=token, target_ref=target_ref)
        if commit_sha is None:
            print(f"Already up to date: {target_ref}")
            if args.gateway:
                write_gateway_exit_code(0)
            return 0

        print(f"✓ Pushed Space commit {commit_sha}")
        if args.no_wait:
            print("✓ Space rebuild submitted. Current runtime will keep serving until the new build is ready.")
            if args.gateway:
                write_gateway_exit_code(0)
            return 0

        print("⏳ Waiting for Hugging Face Space rebuild to finish...")
        payload = wait_for_space_runtime(space_id, commit_sha, token=token, timeout=args.timeout)
        if args.gateway:
            write_gateway_exit_code(0)
        print(f"✓ Space is RUNNING at {payload.get('sha')}")
        return 0
    except (subprocess.CalledProcessError, RuntimeError, TimeoutError, urllib.error.URLError) as exc:
        print(f"✗ {exc}")
        if args.gateway:
            write_gateway_exit_code(1)
        return 1


if __name__ == "__main__":
    raise SystemExit(run_update_cli())