Update app.py

app.py

@@ -16,7 +16,7 @@ from starlette.responses import JSONResponse
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.chrome.service import Service
-from selenium.common.exceptions import WebDriverException
+from selenium.common.exceptions import WebDriverException, SessionNotCreatedException
 from webdriver_manager.chrome import ChromeDriverManager
 from bs4 import BeautifulSoup
 from selenium.webdriver.common.by import By
@@ -24,7 +24,10 @@ from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.common.exceptions import TimeoutException
 
-# Configure logging
+# virtual display
+from pyvirtualdisplay import Display
+
+# Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("fast_fetcher")
 
@@ -50,25 +53,24 @@ class BrowserManager:
         ]
         self._driver_lock = threading.Lock()
         self._driver: Optional[webdriver.Chrome] = None
+        self._display: Optional[Display] = None
         self._start_driver_with_retries()
 
     def _build_options(self) -> Options:
         opts = Options()
-
-        # set binary location from env, fallback to default
+        # If CHROME_BIN is present, point to it
         chrome_bin = os.environ.get("CHROME_BIN", "/usr/bin/google-chrome-stable")
         if os.path.exists(chrome_bin):
             opts.binary_location = chrome_bin
             logger.debug("Using chrome binary: %s", chrome_bin)
         else:
-            logger.warning("Chrome binary not found at %s; Selenium Manager may still locate a browser.", chrome_bin)
+            logger.warning("Chrome binary not found at %s (will rely on system/browser manager).", chrome_bin)
 
         if self.headless:
-            # keep compatibility for versions: add both
             opts.add_argument("--headless=new")
             opts.add_argument("--headless")
 
-        # container-friendly and stability flags
+        # container-friendly flags (and stable fallback)
         opts.add_argument("--no-sandbox")
         opts.add_argument("--disable-setuid-sandbox")
         opts.add_argument("--disable-dev-shm-usage")
@@ -76,7 +78,6 @@ class BrowserManager:
         opts.add_argument("--disable-extensions")
         opts.add_argument("--disable-blink-features=AutomationControlled")
         opts.add_argument("--disable-software-rasterizer")
-        opts.add_argument("--disable-accelerated-2d-canvas")
         opts.add_argument(f"--window-size={self.window_size}")
         opts.add_argument("--remote-debugging-port=0")
 
@@ -103,75 +104,106 @@ class BrowserManager:
                 self._start_driver()
                 logger.info("Chrome driver started successfully.")
                 return
-            except Exception as e:
-                logger.exception("Failed to start driver on attempt %d: %s", attempt, e)
-                last_exc = e
+            except Exception as exc:
+                logger.exception("Failed to start driver on attempt %d: %s", attempt, exc)
+                last_exc = exc
                 time.sleep(delay_seconds)
         raise RuntimeError(f"Unable to start Chrome driver after {attempts} attempts: {last_exc}") from last_exc
 
+    def _start_xvfb_if_needed(self):
+        # If headless=False AND no DISPLAY, start Xvfb via pyvirtualdisplay
+        if not self.headless and os.environ.get("DISPLAY", "") == "":
+            try:
+                logger.info("No DISPLAY found and headless=False — starting virtual X display (Xvfb).")
+                self._display = Display(visible=0, size=(int(self.window_size.split(",")[0]), int(self.window_size.split(",")[1])))
+                self._display.start()
+                logger.info("Virtual X display started (DISPLAY=%s).", os.environ.get("DISPLAY"))
+            except Exception as e:
+                logger.exception("Failed to start virtual display: %s", e)
+                raise
+
+    def _stop_xvfb_if_started(self):
+        if self._display:
+            try:
+                self._display.stop()
+                logger.info("Virtual X display stopped.")
+            except Exception:
+                pass
+            self._display = None
+
     def _start_driver(self):
+        # start virtual display if required BEFORE launching Chrome
+        self._start_xvfb_if_needed()
+
         opts = self._build_options()
 
-        # First: try letting Selenium Manager handle driver (selenium >=4.14)
+        # 1) Try Selenium Manager (webdriver.Chrome(options=opts)). Selenium >=4.14 may use driver manager itself.
+        primary_exc = None
+        fallback_exc = None
         try:
-            logger.debug("Trying webdriver.Chrome(options=opts) to allow Selenium Manager to find driver.")
+            logger.debug("Attempting to start Chrome via Selenium Manager (webdriver.Chrome(options=opts))")
             self._driver = webdriver.Chrome(options=opts)
-            # quick smoke test: get version via execute_script -> may throw if browser crashed
+            # quick smoke test: ensure browser is responsive (may throw)
             try:
                 self._driver.execute_script("return navigator.userAgent")
-            except Exception:
-                # driver started but browser died; fallthrough to fallback
-                raise RuntimeError("Selenium Manager started a session but the browser crashed immediately.")
-            # if success, configure network blocking best-effort
+            except Exception as e:
+                # browser started but died quickly
+                raise RuntimeError("Browser started by Selenium Manager but crashed immediately.") from e
+
             self._post_start_setup()
             return
         except Exception as e_primary:
+            primary_exc = e_primary
             logger.warning("Selenium Manager attempt failed: %s", e_primary)
 
-        # Fallback: use webdriver-manager to download compatible chromedriver, ensure executable, and use it
+        # 2) Fallback: use webdriver-manager to download driver and start with the explicit Service
         try:
             driver_path = ChromeDriverManager().install()
             logger.info("webdriver-manager installed chromedriver: %s", driver_path)
-            # ensure executable bit set
             try:
                 os.chmod(driver_path, 0o755)
-            except Exception as ex_perm:
-                logger.warning("Could not chmod chromedriver: %s", ex_perm)
+            except Exception:
+                logger.debug("chmod on chromedriver failed or unnecessary.")
+
             service = Service(driver_path)
             self._driver = webdriver.Chrome(service=service, options=opts)
             self._post_start_setup()
             return
         except Exception as e_fallback:
+            fallback_exc = e_fallback
             logger.exception("webdriver-manager fallback failed: %s", e_fallback)
-            # try one more variation: if /usr/bin/chromedriver exists, use it
-            try:
-                system_path = "/usr/bin/chromedriver"
-                if os.path.exists(system_path):
-                    logger.info("Trying system chromedriver at %s", system_path)
-                    os.chmod(system_path, 0o755)
-                    service = Service(system_path)
-                    self._driver = webdriver.Chrome(service=service, options=opts)
-                    self._post_start_setup()
-                    return
-                else:
-                    logger.debug("No system chromedriver at %s", system_path)
-            except Exception as e_sys:
-                logger.exception("System chromedriver attempt failed: %s", e_sys)
 
-            # raise combined error
-            raise RuntimeError(f"Failed to start Chrome driver. primary_error={e_primary}, fallback_error={e_fallback}")
+        # 3) Final fallback: attempt system /usr/bin/chromedriver if available
+        try:
+            sys_path = "/usr/bin/chromedriver"
+            if os.path.exists(sys_path):
+                logger.info("Trying system chromedriver at %s", sys_path)
+                try:
+                    os.chmod(sys_path, 0o755)
+                except Exception:
+                    pass
+                service = Service(sys_path)
+                self._driver = webdriver.Chrome(service=service, options=opts)
+                self._post_start_setup()
+                return
+        except Exception as e_sys:
+            logger.exception("System chromedriver attempt failed: %s", e_sys)
+
+        # If all failed, stop virtual display (if started) and raise a helpful error
+        self._stop_xvfb_if_started()
+        # Include both primary and fallback messages in the raised exception
+        raise RuntimeError(f"Failed to start Chrome driver. primary_error={primary_exc}, fallback_error={fallback_exc}")
 
     def _post_start_setup(self):
-        # set reasonable timeout and (best-effort) block some urls using CDP
         try:
             self._driver.set_page_load_timeout(60)
-            self._driver.execute_cdp_cmd("Network.enable", {})
-            if self.block_resource_urls:
-                try:
+            # best-effort CDP network blocking
+            try:
+                self._driver.execute_cdp_cmd("Network.enable", {})
+                if self.block_resource_urls:
                     self._driver.execute_cdp_cmd("Network.setBlockedURLs", {"urls": self.block_resource_urls})
-                except Exception:
-                    # some Chrome builds may not support this; ignore
-                    pass
+            except Exception:
+                pass
         except Exception:
             pass
 
@@ -223,12 +255,14 @@ class BrowserManager:
             except Exception:
                 pass
         self._driver = None
+        # stop display if we started one
+        self._stop_xvfb_if_started()
 
     def close(self):
         self._safe_quit_driver()
 
 
-# ---------------- EXTRACT_DATA ----------------
+# ---------------- EXTRACT_DATA (same as your earlier implementation) ----------------
 def EXTRACT_DATA(html: str) -> Dict[str, Any]:
     soup = BeautifulSoup(html, "html.parser")
     BASE_URL = "https://www.google.com"
@@ -473,8 +507,8 @@ CACHE = SimpleTTLCache(ttl_seconds=25)
 @app.on_event("startup")
 async def startup_event():
     global POOL, EXECUTOR
-    # in Docker use headless True
-    POOL = BrowserPool(pool_size=1, headless=True)
+    # Switch headless here to False as you asked. The BrowserManager will start an Xvfb display automatically.
+    POOL = BrowserPool(pool_size=1, headless=False)
     EXECUTOR = ThreadPoolExecutor(max_workers=2)
     app.state.executor = EXECUTOR
     app.state.pool = POOL