Update server.py

server.py

@@ -6,8 +6,8 @@ Perchance Image Generation API v3.0
 Lightweight, self-healing, scalable image generation server.
 Handles thousands of concurrent users on Hugging Face Spaces.
 
-Main Endpoint:  POST /v1/generation/Image/Create
-SSE Stream:     GET  /v1/generation/stream/{task_id}
+Key fix: Server starts accepting connections IMMEDIATELY.
+         Browser key fetch runs in BACKGROUND after startup.
 """
 
 # ═══════════════════════════════════════════════════════════════
@@ -263,10 +263,6 @@ class Broadcaster:
                 except asyncio.QueueFull:
                     pass
 
-    async def drop_channel(self, task_id: str) -> None:
-        async with self._lock:
-            self._subs.pop(task_id, None)
-
     @property
     def total_subscribers(self) -> int:
         return sum(len(v) for v in self._subs.values())
@@ -347,6 +343,18 @@ class KeyManager:
         self._refresh_gate       = asyncio.Lock()
         self._last_ok: float     = 0.0
         self._consec_fails: int  = 0
+        self._ready              = asyncio.Event()   # set once first key obtained
+
+    @property
+    def is_ready(self) -> bool:
+        return self._ready.is_set()
+
+    async def wait_ready(self, timeout: float = None):
+        """Wait until at least one key has been obtained."""
+        if timeout:
+            await asyncio.wait_for(self._ready.wait(), timeout=timeout)
+        else:
+            await self._ready.wait()
 
     async def get(self) -> Optional[str]:
         await self._usable.wait()
@@ -359,6 +367,7 @@ class KeyManager:
         self._last_ok      = time.time()
         self._consec_fails = 0
         self._usable.set()
+        self._ready.set()
         log.info("userKey set OK (len=%d)", len(new_key))
 
     def reset_failures(self) -> None:
@@ -448,7 +457,8 @@ class PerchanceClient:
         except Exception:
             pass
 
-    def generate(self, *, prompt, negative, seed, resolution, guidance, style, user_key, ad_code) -> dict:
+    def generate(self, *, prompt, negative, seed, resolution,
+                 guidance, style, user_key, ad_code) -> dict:
         rid = _rid()
         params = {
             "userKey": user_key, "requestId": rid,
@@ -722,6 +732,9 @@ _tasks    = TaskStore()
 _sse      = Broadcaster()
 _queue: Optional[asyncio.Queue] = None
 
+# Track startup readiness (server is up but key may still be loading)
+_server_start_time: float = 0.0
+
 
 # ═══════════════════════════════════════════════════════════════
 #  GENERATION ENGINE
@@ -842,6 +855,26 @@ async def _run_task(task: dict) -> None:
     loop = asyncio.get_running_loop()
     tid  = task["id"]
 
+    # Wait for key to be available before starting
+    if not _keys.is_ready:
+        log.info("Task %s waiting for initial key …", tid[:8])
+        await _sse.emit(tid, {
+            "event": "waiting", "task_id": tid,
+            "message": "Waiting for server key initialization …",
+            "time": _now(),
+        })
+        try:
+            await _keys.wait_ready(timeout=120)
+        except asyncio.TimeoutError:
+            task["status"] = "failed"
+            task["error"]  = "Key initialization timeout"
+            task["finished_at"] = _now()
+            await _sse.emit(tid, {
+                "event": "failed", "task_id": tid,
+                "error": "Key initialization timeout", "time": _now(),
+            })
+            return
+
     task["status"]     = "running"
     task["started_at"] = _now()
 
@@ -935,52 +968,88 @@ async def _janitor():
             log.exception("Janitor error")
 
 
+async def _background_key_fetch():
+    """
+    Fetch userKey via browser IN THE BACKGROUND.
+    Server is already accepting connections while this runs.
+    """
+    skip = os.environ.get("NO_INITIAL_FETCH", "0") in ("1", "true", "True")
+    if skip:
+        log.info("NO_INITIAL_FETCH → skipping browser key fetch")
+        return
+
+    log.info("Background key fetch starting …")
+    try:
+        key = await extract_key_via_browser(BROWSER_TIMEOUT, BROWSER_HEADLESS)
+        if key:
+            await _keys.set(key)
+            log.info("Background key fetch OK")
+        else:
+            log.warning("Background key fetch failed — use /v1/keys/set or /v1/keys/refresh")
+    except Exception as exc:
+        log.exception("Background key fetch error: %s", exc)
+
+
 # ═══════════════════════════════════════════════════════════════
 #  FASTAPI — LIFESPAN
+#
+#  KEY CHANGE: yield IMMEDIATELY so HF Spaces sees port 7860
+#              key fetch runs as a background task AFTER yield
 # ═══════════════════════════════════════════════��═══════════════
 
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
-    global _queue
+    global _queue, _server_start_time
 
+    _server_start_time = time.time()
     _queue = asyncio.Queue(maxsize=QUEUE_CAPACITY)
     loop   = asyncio.get_running_loop()
 
+    # Start Xvfb (fast, synchronous)
     await loop.run_in_executor(None, partial(_display.ensure, BROWSER_HEADLESS))
 
-    skip = os.environ.get("NO_INITIAL_FETCH", "0") in ("1", "true", "True")
-    if not skip:
-        try:
-            key = await extract_key_via_browser(BROWSER_TIMEOUT, BROWSER_HEADLESS)
-            if key:
-                await _keys.set(key)
-                log.info("Startup key OK")
-            else:
-                log.warning("Startup key failed — use /v1/keys/set or /v1/keys/refresh")
-        except Exception as exc:
-            log.exception("Startup key error: %s", exc)
-    else:
-        log.info("NO_INITIAL_FETCH → skipping browser key fetch")
-
+    # Start workers (instant)
     workers = [asyncio.create_task(_worker(i + 1)) for i in range(WORKER_COUNT)]
+
+    # Start janitor (instant)
     janitor = asyncio.create_task(_janitor())
 
+    # Start key fetch IN BACKGROUND (non-blocking!)
+    key_task = asyncio.create_task(_background_key_fetch())
+
     log.info(
-        "═══ SERVER READY ═══  workers=%d  queue_cap=%d  headless=%s",
+        "═══ SERVER READY ═══  workers=%d  queue_cap=%d  headless=%s  (key fetching in background)",
         WORKER_COUNT, QUEUE_CAPACITY, BROWSER_HEADLESS,
     )
 
+    # ════════════════════════════════════════
+    #  yield IMMEDIATELY — port 7860 opens NOW
+    #  HF Spaces sees 200 OK on GET /
+    #  No more "stuck on Starting"!
+    # ════════════════════════════════════════
     yield
 
+    # ──────── shutdown ────────
     log.info("Shutting down …")
+
+    # Cancel background key fetch if still running
+    if not key_task.done():
+        key_task.cancel()
+        try:
+            await key_task
+        except (asyncio.CancelledError, Exception):
+            pass
+
     for _ in range(WORKER_COUNT):
         await _queue.put(None)
     await asyncio.gather(*workers, return_exceptions=True)
+
     janitor.cancel()
     try:
         await janitor
     except asyncio.CancelledError:
         pass
+
     _client.close()
     _executor.shutdown(wait=False)
     await loop.run_in_executor(None, _display.shutdown)
@@ -1001,21 +1070,18 @@ async def _global_err(request: Request, exc: Exception):
     return JSONResponse(status_code=500, content={"error": "Internal server error"})
 
 
-# ══════════════════════════════════════════════
-#  ROOT  ← THIS IS THE FIX FOR "STUCK STARTING"
-# ══════════════════════════════════════════════
+# ══════════════════════════════════════
+#  ROOT  — HF Spaces readiness check
+# ══════════════════════════════════════
 
 @app.get("/")
 async def root():
-    """
-    Root endpoint — required by Hugging Face Spaces.
-    HF checks GET / to confirm the app is alive.
-    Without this, the Space shows 'Starting...' forever.
-    """
     return {
         "name":    "Perchance Image Generation API",
         "version": "3.0",
         "status":  "running",
+        "key_ready": _keys.is_ready,
+        "uptime":  round(time.time() - _server_start_time, 1),
         "docs":    "/docs",
         "endpoints": {
             "generate":     "POST /v1/generation/Image/Create",
@@ -1026,7 +1092,6 @@ async def root():
             "set_key":      "POST /v1/keys/set",
             "refresh_key":  "POST /v1/keys/refresh",
         },
-        "time": _now(),
     }
 
 
@@ -1035,11 +1100,13 @@ async def health():
     return {
         "status":          "healthy",
         "has_key":          _keys.key is not None,
+        "key_ready":        _keys.is_ready,
         "queue_size":       _queue.qsize() if _queue else 0,
         "queue_capacity":   QUEUE_CAPACITY,
         "active_tasks":     _tasks.active_count(),
         "stored_tasks":     _tasks.size,
         "sse_subscribers":  _sse.total_subscribers,
+        "uptime":           round(time.time() - _server_start_time, 1),
         "time":             _now(),
     }
 
@@ -1128,7 +1195,6 @@ async def stream_task(request: Request, task_id: str):
                         "data":  json.dumps({"time": _now()}),
                     }
                     if await request.is_disconnected():
-                        log.info("SSE client disconnected  task=%s", task_id[:8])
                         break
                     if task["status"] in ("completed", "failed"):
                         yield {
@@ -1143,13 +1209,9 @@ async def stream_task(request: Request, task_id: str):
                     continue
 
                 etype = ev.get("event", "message")
-                yield {
-                    "event": etype,
-                    "data":  json.dumps(ev),
-                }
+                yield {"event": etype, "data": json.dumps(ev)}
                 if etype in ("completed", "failed"):
                     break
-
         finally:
             await _sse.unsubscribe(task_id, sub_q)