Update app.py

app.py

@@ -1,6 +1,12 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-import threading, queue, time, json
+import threading, queue, time, os
+
+try:
+    from pyvirtualdisplay import Display
+except Exception:
+    Display = None
+
 from seleniumbase import SB
 
 app = FastAPI()
@@ -9,26 +15,50 @@ class URLRequest(BaseModel):
     url: str
 
 url_queue = queue.Queue()
-sb_ref = {"driver": None} 
+sb_ref = {"driver": None}
 
 def browser_worker():
-    with SB(uc=True, headless=False) as sb:
-        sb_ref["driver"] = sb
-        sb.get("https://www.google.com") 
-        while True:
-            url = url_queue.get()
+    display = None
+    sb = None
+    try:
+        try:
+            sb = SB(uc=True, headless=False, xvfb=True)
+        except TypeError:
+            if Display is None:
+                raise RuntimeError("pyvirtualdisplay not available; cannot start Xvfb")
+            display = Display(visible=0, size=(1920, 1080))
+            display.start()
+            sb = SB(uc=True, headless=False)
+
+        with sb:
+            sb_ref["driver"] = sb
+            sb.get("https://www.google.com") 
+            while True:
+                url = url_queue.get()
+                try:
+                    print(f"[open] {url}")
+                    sb.get(url)
+                    time.sleep(0.5)
+                except Exception as e:
+                    print("[error]", e)
+                finally:
+                    url_queue.task_done()
+    finally:
+        try:
+            if sb:
+                sb.driver.quit()
+        except Exception:
+            pass
+        if display:
             try:
-                print(f"[open] {url}")
-                sb.get(url)
-                time.sleep(0.5)
-            except Exception as e:
-                print("[error]", e)
-            url_queue.task_done()
+                display.stop()
+            except Exception:
+                pass
 
 @app.get("/cookies")
 def get_cookies():
     sb = sb_ref.get("driver")
-    return sb.get_cookies() if sb else {}
+    return sb.get_cookies() if sb else []
 
 @app.post("/open_url")
 async def open_url(request: URLRequest):
@@ -36,4 +66,12 @@ async def open_url(request: URLRequest):
     url_queue.put(request.url)
     return {"status": "queued", "url": request.url}
 
-threading.Thread(target=browser_worker, daemon=True).start()
+def _start_worker_once():
+    if not getattr(_start_worker_once, "started", False):
+        t = threading.Thread(target=browser_worker, daemon=True)
+        t.start()
+        _start_worker_once.started = True
+
+@app.on_event("startup")
+def on_startup():
+    _start_worker_once()