Update main.py

main.py

@@ -61,10 +61,6 @@ class BrowserWorker:
     def tag(self, msg: str) -> str:
         return f"[W{self.id}] {msg}"
 
-    # ──────────────────────────────────────────────
-    # Init
-    # ──────────────────────────────────────────────
-
     async def init(self, playwright):
         self._playwright = playwright
         self._browser = await playwright.chromium.launch(
@@ -109,10 +105,6 @@ class BrowserWorker:
         finally:
             await page.close()
 
-    # ──────────────────────────────────────────────
-    # Context Rotation
-    # ──────────────────────────────────────────────
-
     async def _rotate_context(self):
         print(self.tag(f"Rotating context after {self._request_count} requests..."))
         try:
@@ -124,10 +116,6 @@ class BrowserWorker:
         self._request_count = 0
         print(self.tag("Context rotated ✓"))
 
-    # ──────────────────────────────────────────────
-    # Chat
-    # ──────────────────────────────────────────────
-
     async def chat(self, model_label: str, prompt: str) -> str:
         async with self._lock:
             self.busy = True
@@ -172,10 +160,6 @@ class BrowserWorker:
         finally:
             await page.close()
 
-    # ──────────────────────────────────────────────
-    # Helpers
-    # ──────────────────────────────────────────────
-
     async def _select_model(self, page, model_label: str):
         try:
             btn  = page.locator('button[data-testid="model-select-button"]')
@@ -455,6 +439,15 @@ def _auth(request: Request) -> bool:
     )
 
 
+def _auth_token(request: Request, token: str = "") -> bool:
+    """يقبل المصادقة من Header أو من ?token= في الـ URL"""
+    header_key = (
+        request.headers.get("authorization", "")
+        .replace("Bearer ", "").strip()
+    )
+    return header_key == API_SECRET_KEY or token == API_SECRET_KEY
+
+
 def _get_model_label(model: str) -> str:
     return DUCK_MODELS.get(model, DUCK_MODELS[DEFAULT_MODEL])
 
@@ -496,6 +489,7 @@ async def chat_completions(request: Request):
         return JSONResponse(status_code=400, content={"error": {"message": "Invalid JSON"}})
     if not _auth(request):
         return JSONResponse(status_code=401, content={"error": {"message": "Invalid API Key"}})
+
     messages = data.get("messages", [])
     if not messages:
         return JSONResponse(status_code=400, content={"error": {"message": "messages required"}})
@@ -589,8 +583,11 @@ async def list_models(request: Request):
 
 
 @app.get("/health")
-@app.get("/")
-async def health():
+async def health(request: Request, token: str = ""):
+    # ── /health يقبل Header أو ?token= ──────────────────────
+    if not _auth_token(request, token):
+        return JSONResponse(status_code=401, content={"error": {"message": "Invalid API Key"}})
+
     busy  = sum(1 for w in pool.workers if w.busy)
     stats = [
         {
@@ -622,10 +619,25 @@ async def health():
     }
 
 
+@app.get("/")
+async def root():
+    return {
+        "status":  "running",
+        "message": "Duck.ai API Pool Server is active!",
+        "docs":    "/docs",
+        "health":  "/health",
+    }
+
+
 @app.get("/dashboard", response_class=HTMLResponse)
-async def dashboard(request: Request):
-    if not _auth(request):
-        return HTMLResponse("<h1>401 Unauthorized</h1>", status_code=401)
+async def dashboard(request: Request, token: str = ""):
+    # ── يقبل Header أو ?token= في الـ URL ───────────────────
+    if not _auth_token(request, token):
+        return HTMLResponse(
+            "<h1 style='font-family:sans-serif;text-align:center;margin-top:20%;"
+            "color:#ef5350'>401 — Unauthorized</h1>",
+            status_code=401
+        )
     with open("dashboard.html", "r", encoding="utf-8") as f:
         return HTMLResponse(f.read())