from __future__ import annotations import base64 import hashlib import html import json import os import re import secrets import time from datetime import datetime, timezone from pathlib import Path from typing import Any from urllib.parse import urlencode import httpx from huggingface_hub import HfApi from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from starlette.routing import Mount, Route from starlette.staticfiles import StaticFiles APP_NAME = "Trackio Laboratory" SESSION_COOKIE = "trackio_home_session" SESSION_TTL_SECONDS = 86400 * 30 STATE_TTL_SECONDS = 600 DEFAULT_OAUTH_SCOPES = "openid profile read-repos" CACHE_DIR = Path(os.getenv("TRACKIO_LAB_CACHE_DIR", "/data/trackio-laboratory/cache")) api = HfApi() _sessions: dict[str, dict[str, Any]] = {} _pending_states: dict[str, dict[str, Any]] = {} def _cache_root() -> Path: if CACHE_DIR.exists() or CACHE_DIR.parent.exists(): return CACHE_DIR return Path(".cache/trackio-laboratory") def _cache_key(username: str | None) -> str: raw = (username or "unknown").strip().lower() return hashlib.sha256(raw.encode("utf-8")).hexdigest() def _cache_path(username: str | None) -> Path: return _cache_root() / f"{_cache_key(username)}.json" def _read_cached_spaces(username: str | None) -> dict[str, Any] | None: path = _cache_path(username) try: if not path.is_file(): return None with path.open(encoding="utf-8") as cache_file: payload = json.load(cache_file) if not isinstance(payload, dict) or not isinstance(payload.get("data"), dict): return None data = payload["data"] data["cache"] = { "hit": True, "saved_at": payload.get("saved_at"), } return data except Exception: return None def _write_cached_spaces(username: str | None, data: dict[str, Any]) -> None: root = _cache_root() try: root.mkdir(parents=True, exist_ok=True) path = _cache_path(username) tmp_path = path.with_suffix(".json.tmp") payload = { "saved_at": datetime.now(timezone.utc).isoformat(), "data": data, } with tmp_path.open("w", encoding="utf-8") as cache_file: json.dump(payload, cache_file, separators=(",", ":")) tmp_path.replace(path) except Exception: pass def _current_space_id() -> str: if os.getenv("SPACE_ID"): return os.environ["SPACE_ID"] author = os.getenv("SPACE_AUTHOR_NAME") repo = os.getenv("SPACE_REPO_NAME") if author and repo: return f"{author}/{repo}" return "trackio/laboratory" def _now() -> float: return time.monotonic() def _evict_expired() -> None: now = _now() for state, payload in list(_pending_states.items()): if now - payload["created_at"] > STATE_TTL_SECONDS: del _pending_states[state] for session_id, payload in list(_sessions.items()): if now - payload["created_at"] > SESSION_TTL_SECONDS: del _sessions[session_id] def _safe_next(value: str | None) -> str: if value and value.startswith("/") and not value.startswith("//"): return value return "/" def _oauth_redirect_uri(request: Request) -> str: space_host = os.getenv("SPACE_HOST") if space_host: return f"https://{space_host.split(',')[0]}/login/callback" return str(request.base_url).rstrip("/") + "/login/callback" def _oauth_configured() -> bool: return bool(os.getenv("OAUTH_CLIENT_ID") and os.getenv("OAUTH_CLIENT_SECRET")) def _session_id_from_request(request: Request) -> str | None: session_id = request.cookies.get(SESSION_COOKIE) if session_id and session_id in _sessions: payload = _sessions[session_id] if _now() - payload["created_at"] <= SESSION_TTL_SECONDS: return session_id del _sessions[session_id] return None def _session_from_request(request: Request) -> dict[str, Any] | None: _evict_expired() session_id = _session_id_from_request(request) if not session_id: return None return _sessions.get(session_id) def _set_session_cookie(resp: Response, session_id: str, request: Request) -> None: secure = bool(os.getenv("SPACE_HOST")) or request.url.scheme == "https" same_site = "none" if bool(os.getenv("SPACE_HOST")) else "lax" resp.set_cookie( SESSION_COOKIE, session_id, max_age=SESSION_TTL_SECONDS, httponly=True, samesite=same_site, secure=secure, path="/", ) def _clear_session_cookie(resp: Response, request: Request) -> None: secure = bool(os.getenv("SPACE_HOST")) or request.url.scheme == "https" same_site = "none" if bool(os.getenv("SPACE_HOST")) else "lax" resp.delete_cookie(SESSION_COOKIE, path="/", samesite=same_site, secure=secure) def _display_username(whoami: dict[str, Any] | None) -> str | None: if not whoami: return None return whoami.get("fullname") or whoami.get("name") def _org_names(whoami: dict[str, Any]) -> list[str]: orgs = [] for org in whoami.get("orgs") or []: if isinstance(org, dict) and org.get("name"): orgs.append(org["name"]) elif isinstance(org, str): orgs.append(org) return sorted(set(orgs)) def _namespaces_for_user(token: str) -> tuple[list[str], dict[str, Any]]: whoami = api.whoami(token=token, cache=True) username = whoami.get("name") namespaces = [username] if username else [] namespaces.extend(_org_names(whoami)) return sorted(set(namespaces)), whoami def _has_trackio_tag(space: Any) -> bool: return "trackio" in {str(tag).lower() for tag in (getattr(space, "tags", None) or [])} def _list_trackio_spaces_for_namespace(namespace: str, token: str) -> list[Any]: try: spaces = list( api.list_spaces( author=namespace, filter="trackio", full=True, token=token, ) ) except TypeError: spaces = list(api.list_spaces(author=namespace, full=True, token=token)) return [space for space in spaces if _has_trackio_tag(space)] def _iso(value: Any) -> str | None: if value is None: return None if isinstance(value, datetime): if value.tzinfo is None: value = value.replace(tzinfo=timezone.utc) return value.isoformat() return str(value) def _space_runtime_stage(space: Any) -> str | None: runtime = getattr(space, "runtime", None) stage = getattr(runtime, "stage", None) return str(stage) if stage is not None else None def _space_app_url(space: Any) -> str | None: host = getattr(space, "host", None) if isinstance(host, str) and host: return host if host.startswith("http") else f"https://{host}" subdomain = getattr(space, "subdomain", None) if isinstance(subdomain, str) and subdomain: return f"https://{subdomain}.hf.space" space_id = getattr(space, "id", "") if isinstance(space_id, str) and "/" in space_id: namespace, name = space_id.split("/", 1) subdomain = re.sub(r"[^a-z0-9-]+", "-", f"{namespace}-{name}".lower()).strip( "-" ) if subdomain: return f"https://{subdomain}.hf.space" return None def _serialize_space(space: Any) -> dict[str, Any]: space_id = getattr(space, "id", "") namespace, _, name = space_id.partition("/") last_modified = getattr(space, "last_modified", None) app_url = _space_app_url(space) private = bool(getattr(space, "private", False)) return { "id": space_id, "namespace": namespace, "name": name or space_id, "private": private, "sdk": getattr(space, "sdk", None), "runtime_stage": _space_runtime_stage(space), "last_modified": _iso(last_modified), "repo_url": f"https://huggingface.co/spaces/{space_id}", "app_url": app_url, "embeddable": bool(app_url) and not private, } def discover_trackio_spaces(token: str) -> dict[str, Any]: namespaces, whoami = _namespaces_for_user(token) by_id: dict[str, Any] = {} namespace_errors: dict[str, str] = {} for namespace in namespaces: try: for space in _list_trackio_spaces_for_namespace(namespace, token): space_id = getattr(space, "id", "") if not space_id or space_id == _current_space_id(): continue by_id[space_id] = space except Exception as err: namespace_errors[namespace] = str(err) spaces = [_serialize_space(space) for space in by_id.values()] spaces.sort(key=lambda item: item.get("last_modified") or "", reverse=True) return { "user": { "name": whoami.get("name"), "display_name": _display_username(whoami), "namespaces": namespaces, }, "spaces": spaces, "errors": namespace_errors, } async def homepage(request: Request) -> Response: session = _session_from_request(request) logout_name_source = ( str(session.get("display_name") or session.get("username") or "").strip() if session else "" ) logout_first_name = html.escape( logout_name_source.split()[0] if logout_name_source else "" ) logout_label = f"Logout, {logout_first_name}" if logout_first_name else "Logout" oauth_ready = _oauth_configured() body_class = "is-authed" if session else "is-login" html_doc = f""" {APP_NAME}
Trackio

Open the Trackio dashboards connected to your Hugging Face account and organizations.

Sign in with Hugging Face
No dashboard selected Select a Trackio Space from the sidebar.
""" return HTMLResponse(html_doc) async def login(request: Request) -> Response: client_id = os.getenv("OAUTH_CLIENT_ID") if not client_id: return RedirectResponse("/", status_code=302) _evict_expired() state = secrets.token_urlsafe(32) _pending_states[state] = { "created_at": _now(), "next": _safe_next( request.query_params.get("_target_url") or request.query_params.get("next") ), } redirect_uri = _oauth_redirect_uri(request) scope = os.getenv("OAUTH_SCOPES", DEFAULT_OAUTH_SCOPES).strip() url = "https://huggingface.co/oauth/authorize?" + urlencode( { "client_id": client_id, "redirect_uri": redirect_uri, "response_type": "code", "scope": scope, "state": state, } ) return RedirectResponse(url, status_code=302) async def login_callback(request: Request) -> Response: client_id = os.getenv("OAUTH_CLIENT_ID") client_secret = os.getenv("OAUTH_CLIENT_SECRET") code = request.query_params.get("code") state = request.query_params.get("state") pending = _pending_states.pop(state, None) if state else None if not client_id or not client_secret or not code or not pending: return RedirectResponse("/?oauth_error=1", status_code=302) if _now() - pending["created_at"] > STATE_TTL_SECONDS: return RedirectResponse("/?oauth_error=1", status_code=302) redirect_uri = _oauth_redirect_uri(request) auth_b64 = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() try: async with httpx.AsyncClient(timeout=30) as client: token_resp = await client.post( "https://huggingface.co/oauth/token", headers={"Authorization": f"Basic {auth_b64}"}, data={ "grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri, "client_id": client_id, }, ) token_resp.raise_for_status() access_token = token_resp.json()["access_token"] whoami = api.whoami(token=access_token, cache=True) except Exception: return RedirectResponse("/?oauth_error=1", status_code=302) session_id = secrets.token_urlsafe(32) _sessions[session_id] = { "token": access_token, "created_at": _now(), "username": whoami.get("name"), "display_name": _display_username(whoami), } resp = RedirectResponse(_safe_next(pending.get("next")), status_code=302) _set_session_cookie(resp, session_id, request) return resp async def logout(request: Request) -> Response: session_id = _session_id_from_request(request) if session_id: _sessions.pop(session_id, None) resp = RedirectResponse("/", status_code=302) _clear_session_cookie(resp, request) return resp async def spaces_api(request: Request) -> Response: session = _session_from_request(request) if not session: return JSONResponse({"error": "Not authenticated"}, status_code=401) refresh = request.query_params.get("refresh") in {"1", "true", "yes"} username = session.get("username") if not refresh: cached = _read_cached_spaces(username) if cached is not None: return JSONResponse(cached) try: data = discover_trackio_spaces(session["token"]) data["cache"] = {"hit": False, "saved_at": None} _write_cached_spaces(username, data) return JSONResponse(data) except Exception as err: cached = _read_cached_spaces(username) if cached is not None: cached["cache"]["stale_due_to_error"] = str(err) return JSONResponse(cached) return JSONResponse({"error": str(err)}, status_code=500) routes = [ Route("/", homepage, methods=["GET"]), Route("/login", login, methods=["GET"]), Route("/login/huggingface", login, methods=["GET"]), Route("/login/callback", login_callback, methods=["GET"]), Route("/logout", logout, methods=["GET"]), Route("/api/spaces", spaces_api, methods=["GET"]), Mount("/static", StaticFiles(directory="static"), name="static"), ] app = Starlette(routes=routes) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))