Detect ZeroGPU token expiry and prompt refresh

app.py

@@ -1,3 +1,4 @@
+import base64
 import io
 import json
 import math
@@ -60,14 +61,16 @@ DEFAULT_ZERO_GPU_MAX_SHARDS = int(
     os.getenv("CHIZZLER_ZERO_GPU_MAX_SHARDS", "1")
 )
 SPACE_ID = os.getenv("SPACE_ID")
-if SPACE_ID:
-    warnings.filterwarnings(
-        "ignore",
-        message=(
-            "LoginButton created outside of a Blocks context\\. "
-            "May not work unless you call its `activate\\(\\)` method manually\\."
-        ),
-    )
+ZERO_GPU_TOKEN_REFRESH_BUFFER_SEC = int(
+    os.getenv("CHIZZLER_ZERO_GPU_TOKEN_REFRESH_BUFFER_SEC", "120")
+)
+warnings.filterwarnings(
+    "ignore",
+    message=(
+        "LoginButton created outside of a Blocks context\\. "
+        "May not work unless you call its `activate\\(\\)` method manually\\."
+    ),
+)
 
 
 def log_progress(message: str, level: int = 1, enabled: bool = True) -> None:
@@ -80,6 +83,39 @@ def log_progress(message: str, level: int = 1, enabled: bool = True) -> None:
     sys.stdout.flush()
 
 
+def get_zero_gpu_token_payload(request: Optional[gr.Request]) -> Optional[dict]:
+    if request is None:
+        return None
+    headers = getattr(request, "headers", None)
+    if headers is None or not hasattr(headers, "get"):
+        return None
+    token = headers.get("x-ip-token")
+    if not token:
+        return None
+    try:
+        payload = token.split(".")[1]
+        payload += "=" * (-len(payload) % 4)
+        data = json.loads(base64.urlsafe_b64decode(payload.encode()).decode("utf-8"))
+        if isinstance(data, dict):
+            return data
+    except Exception:
+        return None
+    return None
+
+
+def zero_gpu_token_seconds_left(request: Optional[gr.Request]) -> Optional[int]:
+    payload = get_zero_gpu_token_payload(request)
+    if not payload:
+        return None
+    exp = payload.get("exp")
+    if exp is None:
+        return None
+    try:
+        return int(exp) - int(time.time())
+    except Exception:
+        return None
+
+
 # Load environment variables
 load_dotenv()
 
@@ -1175,6 +1211,16 @@ def process_dataset_and_push(
             )
     attempts = 0
     while True:
+        seconds_left = zero_gpu_token_seconds_left(request)
+        if (
+            SPACE_ID
+            and seconds_left is not None
+            and seconds_left <= ZERO_GPU_TOKEN_REFRESH_BUFFER_SEC
+        ):
+            return (
+                "ZeroGPU login token is expiring "
+                f"({seconds_left}s left). Click Process/Resume again to refresh it."
+            )
         try:
             result = _process_dataset_and_push_gpu(
                 dataset_id,
@@ -1199,6 +1245,11 @@ def process_dataset_and_push(
             )
         except Exception as exc:
             message = str(exc)
+            if "ZeroGPU proxy token expired" in message:
+                return (
+                    "ZeroGPU login token expired. Click Process/Resume again "
+                    "to refresh your session."
+                )
             if auto_resume and "GPU task aborted" in message:
                 attempts += 1
                 log_progress(