Update loader.py

loader.py

@@ -1,5 +1,6 @@
-# loader.py — Public proxy loader for private static Hugging Face Space (redirect-safe)
+# loader.py — Public proxy for PRIVATE Static Space using huggingface.co/resolve
 import os
+import mimetypes
 from urllib.parse import urljoin
 
 from fastapi import FastAPI, Request
@@ -8,73 +9,78 @@ import httpx
 
 app = FastAPI()
 
-# ========= CONFIG (from Secrets) =========
-PRIVATE_STATIC_URL = os.getenv("PRIVATE_STATIC_URL")
-HF_TOKEN = os.getenv("HF_TOKEN")
+# ==== Config from Secrets ====
+SPACE_ID = os.getenv("PRIVATE_SPACE_ID")  # e.g. "cheekeong2025/iip-grading"
+HF_TOKEN = os.getenv("HF_TOKEN")          # must have read access
+REVISION = os.getenv("REVISION", "main")  # optional: pin a branch/tag/commit
+# =============================
 
-if not PRIVATE_STATIC_URL:
+if not SPACE_ID:
     raise RuntimeError(
-        "❌ PRIVATE_STATIC_URL is not set.\n"
-        "Add it in Settings → Repository secrets, e.g.\n"
-        "PRIVATE_STATIC_URL=https://cheekeong2025-iip-grading.static.hf.space"
+        "PRIVATE_SPACE_ID is not set. Go to Settings → Repository secrets and add "
+        "PRIVATE_SPACE_ID like 'cheekeong2025/iip-grading'."
     )
-
 if not HF_TOKEN:
-    print("[loader] ⚠️  Warning: HF_TOKEN not set — attempting anonymous access.")
+    raise RuntimeError(
+        "HF_TOKEN is not set. Add a token with READ access to the private Space."
+    )
+
+BASE_RESOLVE = f"https://huggingface.co/spaces/{SPACE_ID}/resolve/{REVISION}/"
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
-# ========================================
+def build_resolve_url(path: str) -> str:
+    # Ensure trailing slash join is correct
+    return urljoin(BASE_RESOLVE, path.lstrip("/"))
 
-def _with_query(url: str, request: Request) -> str:
-    """Append the query string (if any) from the incoming request."""
+def with_query(url: str, request: Request) -> str:
     q = str(request.url.query or "")
     return f"{url}?{q}" if q else url
 
+def guess_mime(path: str, fallback: str = "application/octet-stream") -> str:
+    if path.endswith("/") or path == "":
+        return "text/html; charset=utf-8"
+    mime, _ = mimetypes.guess_type(path)
+    return mime or fallback
+
 @app.get("/health")
 async def health():
-    return {"status": "ok", "source": PRIVATE_STATIC_URL}
+    return {"status": "ok", "space": SPACE_ID, "revision": REVISION, "source": "huggingface.co/resolve"}
 
 @app.get("/")
 async def root(request: Request):
-    """
-    Serve index.html. Many static sites redirect '/' -> '/index.html' (302).
-    We request index.html directly to avoid the 302, and we also enable
-    follow_redirects for extra safety.
-    """
-    index_url = urljoin(PRIVATE_STATIC_URL.rstrip("/") + "/", "index.html")
-    index_url = _with_query(index_url, request)
-
+    # Serve index.html explicitly (avoid upstream redirects)
+    target = with_query(build_resolve_url("index.html"), request)
     try:
         async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
-            r = await client.get(index_url, headers=HEADERS)
-            # If upstream doesn't have index.html, fall back to root (may redirect)
-            if r.status_code == 404:
-                root_url = _with_query(PRIVATE_STATIC_URL, request)
-                r = await client.get(root_url, headers=HEADERS)
+            r = await client.get(target, headers=HEADERS)
     except Exception as e:
-        return PlainTextResponse(f"Error fetching root: {e}", status_code=500)
+        return PlainTextResponse(f"Error fetching index.html: {e}", status_code=500)
+
+    # If index.html is missing, try repo root as a fallback (rare)
+    if r.status_code == 404:
+        try:
+            fallback = with_query(build_resolve_url(""), request)
+            async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
+                r = await client.get(fallback, headers=HEADERS)
+        except Exception as e:
+            return PlainTextResponse(f"Error fetching root: {e}", status_code=500)
 
-    content_type = r.headers.get("content-type", "text/html; charset=utf-8")
-    # Use HTMLResponse for better default headers on HTML; Response otherwise
-    if "text/html" in content_type:
-        return HTMLResponse(r.text, status_code=r.status_code, media_type=content_type)
-    return Response(content=r.content, status_code=r.status_code, media_type=content_type)
+    ctype = r.headers.get("content-type", "text/html; charset=utf-8")
+    return HTMLResponse(r.text, status_code=r.status_code, media_type=ctype)
 
 @app.get("/{path:path}")
 async def proxy(path: str, request: Request):
-    """
-    Proxy all other static assets (CSS, JS, images, fonts, etc.).
-    We honor redirects and forward the query string.
-    """
-    base = PRIVATE_STATIC_URL.rstrip("/") + "/"
-    target = urljoin(base, path)  # robust path join
-    target = _with_query(target, request)
+    # Proxy any other static asset (CSS, JS, images, fonts, etc.)
+    # Map directories to index.html for SPA-style paths
+    resolved_path = path if "." in path.split("/")[-1] else (path.rstrip("/") + "/index.html")
+    target = with_query(build_resolve_url(resolved_path), request)
 
     try:
         async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
             r = await client.get(target, headers=HEADERS)
     except Exception as e:
-        return PlainTextResponse(f"Error fetching {target}: {e}", status_code=500)
+        return PlainTextResponse(f"Error fetching {resolved_path}: {e}", status_code=500)
 
-    content_type = r.headers.get("content-type", "application/octet-stream")
-    return Response(content=r.content, media_type=content_type, status_code=r.status_code)
+    # Prefer upstream header; otherwise guess from file extension
+    ctype = r.headers.get("content-type") or guess_mime(resolved_path)
+    return Response(content=r.content, media_type=ctype, status_code=r.status_code)