import os import logging from typing import Any from fastapi import FastAPI, Header, HTTPException, Query, Request from huggingface_hub import HfApi logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI() HF_TOKEN = os.environ["HF_TOKEN"] WEBHOOK_SECRET = os.environ["WEBHOOK_SECRET"] api = HfApi(token=HF_TOKEN) def _normalize_prefix(value: str) -> str: value = value.strip().strip("/") return f"{value}/" if value else "" def _load_allowed_prefixes() -> list[str]: raw = os.getenv("ALLOWED_PREFIXES", "").strip() if raw: prefixes = sorted( { _normalize_prefix(item) for item in raw.split(",") if _normalize_prefix(item) } ) logger.info("Using ALLOWED_PREFIXES override: %s", prefixes) return prefixes me = api.whoami() prefixes = {_normalize_prefix(me["name"])} for org in me.get("orgs", []): if isinstance(org, dict): name = org.get("name", "") else: name = str(org) normalized = _normalize_prefix(name) if normalized: prefixes.add(normalized) prefixes = sorted(prefixes) logger.info("Derived allowed prefixes from whoami(): %s", prefixes) return prefixes ALLOWED_PREFIXES = _load_allowed_prefixes() def _namespace_of_repo_id(repo_id: str) -> str: parts = repo_id.strip().split("/", 1) if len(parts) != 2 or not parts[0] or not parts[1]: raise ValueError("repo_id must be in 'namespace/name' format") return parts[0] def _is_allowed_space(repo_id: str) -> bool: try: namespace = _namespace_of_repo_id(repo_id) except ValueError: return False return _normalize_prefix(namespace) in ALLOWED_PREFIXES def _get_webhook_secret( x_webhook_secret: str | None, secret_query: str | None, ) -> str | None: return x_webhook_secret or secret_query @app.get("/healthz") async def healthz() -> dict[str, Any]: return { "ok": True, "allowed_prefixes": ALLOWED_PREFIXES, } @app.post("/webhook") async def handle_webhook( request: Request, space: str = Query(..., description="Target Space repo_id, e.g. user/my-space"), expected_repo: str | None = Query( default=None, description="Optional source repo_id to require for this webhook", ), secret: str | None = Query( default=None, description="Optional webhook secret in query params if header is unavailable", ), x_webhook_secret: str | None = Header(default=None), ) -> dict[str, Any]: supplied_secret = _get_webhook_secret(x_webhook_secret, secret) if supplied_secret != WEBHOOK_SECRET: raise HTTPException(status_code=403, detail="Bad secret") if not _is_allowed_space(space): raise HTTPException( status_code=403, detail=f"Target Space namespace is not allowed: {space}", ) try: payload = await request.json() except Exception as exc: raise HTTPException(status_code=400, detail="Invalid JSON payload") from exc repo = payload.get("repo", {}) event = payload.get("event", {}) repo_name = repo.get("name") repo_type = repo.get("type") event_action = event.get("action") event_scope = event.get("scope") if repo_type != "model": return { "status": "ignored", "reason": "repo is not a model", "source_repo": repo_name, "space": space, } if event_action != "update": return { "status": "ignored", "reason": "event action is not update", "source_repo": repo_name, "space": space, } if event_scope not in (None, "repo.content"): return { "status": "ignored", "reason": "event scope is not repo.content", "source_repo": repo_name, "space": space, } if expected_repo and repo_name != expected_repo: return { "status": "ignored", "reason": "source repo does not match expected_repo", "source_repo": repo_name, "expected_repo": expected_repo, "space": space, } logger.info( "Restarting Space '%s' due to update on source repo '%s'", space, repo_name, ) api.restart_space(repo_id=space) return { "status": "restarted", "space": space, "source_repo": repo_name, "allowed_prefixes": ALLOWED_PREFIXES, }