scripts/huggingface/entrypoint.py

#!/usr/bin/env python3
"""Runtime bootstrap for Hugging Face Spaces."""

from __future__ import annotations

import base64
import json
import os
import re
import sys
from pathlib import Path


APP_DIR = Path(__file__).resolve().parents[2]
AUTH_DIR = APP_DIR / "auth_profiles" / "active"
KEY_FILE = APP_DIR / "auth_profiles" / "key.txt"
HF_AUTH_PATH = AUTH_DIR / "hf-auth.json"


def log(message: str) -> None:
    print(f"[hf-entrypoint] {message}", flush=True)


def _decode_auth_secret() -> str | None:
    raw_json = os.environ.get("AISTUDIO_AUTH_JSON", "").strip()
    if raw_json:
        return raw_json

    encoded_json = os.environ.get("AISTUDIO_AUTH_JSON_B64", "").strip()
    if not encoded_json:
        return None

    compact = "".join(encoded_json.split())
    try:
        return base64.b64decode(compact, validate=True).decode("utf-8")
    except Exception as exc:
        raise RuntimeError("AISTUDIO_AUTH_JSON_B64 is not valid base64 JSON") from exc


def _write_auth_from_secret() -> str | None:
    auth_secret = _decode_auth_secret()
    if not auth_secret:
        return None

    try:
        auth_data = json.loads(auth_secret)
    except json.JSONDecodeError as exc:
        raise RuntimeError("AISTUDIO auth secret is not valid JSON") from exc

    if not isinstance(auth_data, dict) or not isinstance(auth_data.get("cookies"), list):
        raise RuntimeError("AISTUDIO auth JSON must be a Playwright storage_state file")

    AUTH_DIR.mkdir(parents=True, exist_ok=True)
    HF_AUTH_PATH.write_text(
        json.dumps(auth_data, ensure_ascii=False),
        encoding="utf-8",
    )
    HF_AUTH_PATH.chmod(0o600)
    log(f"Wrote auth profile to {HF_AUTH_PATH}")
    return str(HF_AUTH_PATH)


def resolve_auth_path() -> str:
    configured_path = os.environ.get("ACTIVE_AUTH_JSON_PATH", "").strip()
    if configured_path:
        path = Path(configured_path)
        if path.exists():
            log(f"Using ACTIVE_AUTH_JSON_PATH={path}")
            return str(path)
        raise RuntimeError(f"ACTIVE_AUTH_JSON_PATH does not exist: {path}")

    generated_path = _write_auth_from_secret()
    if generated_path:
        return generated_path

    existing_profiles = sorted(AUTH_DIR.glob("*.json"))
    if existing_profiles:
        log(f"Using existing auth profile {existing_profiles[0]}")
        return str(existing_profiles[0])

    raise RuntimeError(
        "No auth profile found. Add AISTUDIO_AUTH_JSON_B64 or AISTUDIO_AUTH_JSON "
        "as a Hugging Face Space Secret."
    )


def write_api_keys() -> None:
    raw_keys = (
        os.environ.get("AISTUDIO_PROXY_API_KEYS")
        or os.environ.get("AISTUDIO_PROXY_API_KEY")
        or os.environ.get("API_KEYS")
        or os.environ.get("API_KEY")
        or ""
    )
    keys = [key.strip() for key in re.split(r"[\n,;]+", raw_keys) if key.strip()]

    KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
    if keys:
        KEY_FILE.write_text("\n".join(keys) + "\n", encoding="utf-8")
        KEY_FILE.chmod(0o600)
        log(f"Configured {len(keys)} API key(s)")
    else:
        KEY_FILE.touch(exist_ok=True)
        log("No API key secret configured; /v1 endpoints will be public")


def main() -> None:
    os.chdir(APP_DIR)
    active_auth_path = resolve_auth_path()
    write_api_keys()

    server_port = os.environ.get("PORT") or os.environ.get("SERVER_PORT") or "7860"
    stream_port = os.environ.get("STREAM_PORT", "0")

    os.environ["PORT"] = server_port
    os.environ["SERVER_PORT"] = server_port
    os.environ.setdefault("DEFAULT_FASTAPI_PORT", server_port)
    os.environ["ACTIVE_AUTH_JSON_PATH"] = active_auth_path

    command = [
        sys.executable,
        "launch_camoufox.py",
        "--headless",
        "--server-port",
        server_port,
        "--stream-port",
        stream_port,
        "--active-auth-json",
        active_auth_path,
        "--helper",
        "",
    ]

    internal_proxy = os.environ.get("INTERNAL_CAMOUFOX_PROXY")
    if internal_proxy is not None:
        command.extend(["--internal-camoufox-proxy", internal_proxy])

    log("Starting AI Studio Proxy API")
    os.execvp(command[0], command)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        log(f"Startup failed: {exc}")
        sys.exit(1)