"""Serves a static frontend alongside the Trackio HTTP API.""" import hashlib import json import logging import re from pathlib import Path from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import FileResponse, HTMLResponse, JSONResponse from starlette.routing import Mount, Route, get_route_path from starlette.staticfiles import StaticFiles ASSETS_DIR = Path(__file__).parent / "assets" _logger = logging.getLogger(__name__) _ABSOLUTE_REF_RE = re.compile(r'(href|src)="(/(?:assets|static/trackio)/[^"]+)"') _LIVE_RELOAD_SCRIPT = """ """.strip() def _frontend_version(frontend_root: Path) -> str: digest = hashlib.sha256() for path in sorted(frontend_root.rglob("*")): if not path.is_file(): continue relative = path.relative_to(frontend_root) digest.update(str(relative).encode("utf-8")) stat = path.stat() digest.update(str(stat.st_mtime_ns).encode("utf-8")) digest.update(str(stat.st_size).encode("utf-8")) return digest.hexdigest() def _inject_live_reload(html: str, root_path: str) -> str: if "/__trackio/frontend_version" in html: return html script = _LIVE_RELOAD_SCRIPT.replace("__TRACKIO_ROOT_PATH__", root_path) if "" in html: return html.replace("", f"{script}\n ") return html + script def _prefix_absolute_refs(html: str, root_path: str) -> str: if not root_path: return html return _ABSOLUTE_REF_RE.sub( lambda m: f'{m.group(1)}="{root_path}{m.group(2)}"', html ) def _inject_base_script(html: str, root_path: str) -> str: if not root_path: return html script = f"" if "" in html: return html.replace("", f"\n {script}", 1) return script + html def _render_index_html(path: Path, root_path: str = "") -> str: html = path.read_text(encoding="utf-8") html = _prefix_absolute_refs(html, root_path) html = _inject_base_script(html, root_path) return _inject_live_reload(html, root_path) def _html_response(path: Path, root_path: str = "") -> HTMLResponse: return HTMLResponse(_render_index_html(path, root_path)) class FrontendMiddleware(BaseHTTPMiddleware): def __init__(self, app, frontend_root: Path, index_html_path: Path): super().__init__(app) self.frontend_root = frontend_root self.index_html_path = index_html_path self.reserved_prefixes = ( "/api/", "/file", "/version", "/artifact_blob/", "/static/trackio", "/__trackio/frontend_version", "/oauth/", "/login/", "/mcp", ) self.reserved_exact = { "/api", "/oauth", "/login", } async def dispatch(self, request, call_next): if request.method not in {"GET", "HEAD"}: return await call_next(request) path = get_route_path(request.scope) if path in self.reserved_exact or path.startswith(self.reserved_prefixes): return await call_next(request) root_path = request.scope.get("root_path", "") relative_path = path.lstrip("/") requested_path = (self.frontend_root / relative_path).resolve() if ( relative_path and requested_path.is_file() and requested_path.is_relative_to(self.frontend_root) ): if requested_path.suffix.lower() == ".html": return _html_response(requested_path, root_path) return FileResponse(requested_path) return _html_response(self.index_html_path, root_path) def mount_frontend(app, frontend_dir: str | Path): frontend_root = Path(frontend_dir).resolve() if not frontend_root.exists(): _logger.warning( "Trackio dashboard UI was not mounted: %s is missing. " "Build the frontend or provide a custom frontend directory.", frontend_root, ) return index_html_path = frontend_root / "index.html" if not index_html_path.exists(): _logger.warning( "Trackio dashboard UI was not mounted: %s is missing.", index_html_path, ) return async def frontend_version(_request): return JSONResponse({"version": _frontend_version(frontend_root)}) static_assets = StaticFiles(directory=str(ASSETS_DIR)) app.add_middleware( FrontendMiddleware, frontend_root=frontend_root, index_html_path=index_html_path, ) app.routes.append(Mount("/static/trackio", app=static_assets)) app.routes.append( Route("/__trackio/frontend_version", frontend_version, methods=["GET"]) )