Update loader.py

loader.py

@@ -1,7 +1,8 @@
-# loader.py — Public proxy for PRIVATE Static Space using huggingface.co/resolve
+# loader.py — robust public proxy for PRIVATE Static Space using huggingface.co/resolve
 import os
 import mimetypes
 from urllib.parse import urljoin
+from typing import Optional
 
 from fastapi import FastAPI, Request
 from fastapi.responses import Response, HTMLResponse, PlainTextResponse
@@ -10,15 +11,14 @@ import httpx
 app = FastAPI()
 
 # ==== 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
+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 branch/tag/commit
 # =============================
 
 if not SPACE_ID:
     raise RuntimeError(
-        "PRIVATE_SPACE_ID is not set. Go to Settings → Repository secrets and add "
-        "PRIVATE_SPACE_ID like 'cheekeong2025/iip-grading'."
+        "PRIVATE_SPACE_ID is not set. Add it in Settings → Repository secrets (e.g. 'cheekeong2025/iip-grading')."
     )
 if not HF_TOKEN:
     raise RuntimeError(
@@ -28,19 +28,59 @@ if not HF_TOKEN:
 BASE_RESOLVE = f"https://huggingface.co/spaces/{SPACE_ID}/resolve/{REVISION}/"
 HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-def build_resolve_url(path: str) -> str:
-    # Ensure trailing slash join is correct
-    return urljoin(BASE_RESOLVE, path.lstrip("/"))
+def build(url_base: str, path: str) -> str:
+    return urljoin(url_base, path.lstrip("/"))
 
 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:
+def looks_like_file(path: str) -> bool:
+    # crude but effective: treat last segment with a dot as a file
+    return "." in path.split("/")[-1]
+
+def guess_mime(path: str, default: 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
+    return mime or default
+
+async def fetch(client: httpx.AsyncClient, path: str, request: Request) -> httpx.Response:
+    """
+    Try multiple resolve locations:
+      1) path
+      2) static/path
+      3) if directory: path/index.html
+      4) if directory: static/path/index.html
+    """
+    # 1) as-is
+    url = with_query(build(BASE_RESOLVE, path), request)
+    r = await client.get(url, headers=HEADERS)
+    if r.status_code == 200:
+        return r
+
+    # 2) static/<path>
+    url_static = with_query(build(BASE_RESOLVE, f"static/{path.lstrip('/')}"), request)
+    r2 = await client.get(url_static, headers=HEADERS)
+    if r2.status_code == 200:
+        return r2
+
+    # If it doesn't look like a file, try directory index.html in both places
+    if not looks_like_file(path):
+        # 3) path/index.html
+        url_dir = with_query(build(BASE_RESOLVE, path.rstrip("/") + "/index.html"), request)
+        r3 = await client.get(url_dir, headers=HEADERS)
+        if r3.status_code == 200:
+            return r3
+
+        # 4) static/path/index.html
+        url_dir_static = with_query(build(BASE_RESOLVE, f"static/{path.rstrip('/')}/index.html"), request)
+        r4 = await client.get(url_dir_static, headers=HEADERS)
+        if r4.status_code == 200:
+            return r4
+
+    # Return the "best" (most recent) response even if 404, so caller can show message
+    return r2 if r2.status_code != 404 else r
 
 @app.get("/health")
 async def health():
@@ -48,39 +88,37 @@ async def health():
 
 @app.get("/")
 async def root(request: Request):
-    # Serve index.html explicitly (avoid upstream redirects)
-    target = with_query(build_resolve_url("index.html"), request)
+    """
+    Serve index.html robustly:
+      - Try /index.html
+      - Try /static/index.html
+      - As a last resort, try repo root (may 404)
+    """
     try:
         async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
-            r = await client.get(target, headers=HEADERS)
+            r = await fetch(client, "index.html", request)
+            if r.status_code == 404:
+                r = await fetch(client, "", request)  # fallback
     except Exception as e:
-        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)
+        return PlainTextResponse(f"Error fetching root: {e}", status_code=500)
 
     ctype = r.headers.get("content-type", "text/html; charset=utf-8")
-    return HTMLResponse(r.text, status_code=r.status_code, media_type=ctype)
+    # If HTML, return text; else raw bytes
+    if "text/html" in ctype:
+        return HTMLResponse(r.text, status_code=r.status_code, media_type=ctype)
+    return Response(content=r.content, status_code=r.status_code, media_type=ctype)
 
 @app.get("/{path:path}")
 async def proxy(path: str, request: 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)
-
+    """
+    Proxy all assets (CSS, JS, images, fonts, SPA routes).
+    Tries root, then static/, and directory index.html fallbacks.
+    """
     try:
         async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client:
-            r = await client.get(target, headers=HEADERS)
+            r = await fetch(client, path, request)
     except Exception as e:
-        return PlainTextResponse(f"Error fetching {resolved_path}: {e}", status_code=500)
+        return PlainTextResponse(f"Error fetching {path}: {e}", status_code=500)
 
-    # Prefer upstream header; otherwise guess from file extension
-    ctype = r.headers.get("content-type") or guess_mime(resolved_path)
+    ctype = r.headers.get("content-type") or guess_mime(path)
     return Response(content=r.content, media_type=ctype, status_code=r.status_code)