Update server.py

server.py

@@ -8,10 +8,13 @@ all original behaviour unchanged.
 
 Behaviour:
  - If ZD_HEADLESS is True, zendriver will run headless as before.
- - If ZD_HEADLESS is False and a DISPLAY is not present, we attempt to
-   start a pyvirtualdisplay.Display (Xvfb) automatically before launching
-   browsers. If pyvirtualdisplay is not installed or starting Xvfb fails,
-   we log a warning and continue.
+ - If ZD_HEADLESS is False and USE_VIRTUAL_DISPLAY is True and a DISPLAY
+   is not present, we attempt to start a pyvirtualdisplay.Display (Xvfb)
+   automatically before launching browsers. If pyvirtualdisplay is not
+   installed or starting Xvfb fails, we log a warning and continue.
+ - If USE_VIRTUAL_DISPLAY is False we will NOT attempt to start a virtual
+   display — you must provide a DISPLAY yourself (or set ZD_HEADLESS=True)
+   if running on a headless host.
 
 To run on Hugging Face Spaces, add `pyvirtualdisplay` to requirements.txt
 and ensure `xvfb` is available in the runtime (HF Spaces typically provide it).
@@ -71,6 +74,13 @@ CLICK_INTERVAL   = 0.35
 CLICK_JITTER     = 8.0
 KEY_PREFIX       = "userKey"
 
+# --- Virtual display toggle (new) ---
+# If True the server will attempt to auto-start a pyvirtualdisplay (Xvfb)
+# when no DISPLAY is present and ZD_HEADLESS is False. If False the server
+# will not try to start a virtual display and you must provide a DISPLAY
+# or set ZD_HEADLESS=True.
+USE_VIRTUAL_DISPLAY = True
+
 # --- HTTP / generation ---
 HTTP_TIMEOUT      = 30
 MAX_DOWNLOAD_WAIT = 180
@@ -157,7 +167,7 @@ def _stamp() -> str:
 def _start_virtual_display_if_needed(headless: bool):
     """
     Start pyvirtualdisplay.Display() if we're running non-headless in an
-    environment without DISPLAY.
+    environment without DISPLAY and USE_VIRTUAL_DISPLAY is True.
     This function is synchronous and safe to be run in a thread executor.
     """
     global VDISPLAY
@@ -166,6 +176,10 @@ def _start_virtual_display_if_needed(headless: bool):
         log.info("ZD_HEADLESS=True → not starting virtual display")
         return
 
+    if not USE_VIRTUAL_DISPLAY:
+        log.info("USE_VIRTUAL_DISPLAY=False → not starting virtual display; expecting manual DISPLAY or headless mode.")
+        return
+
     if os.environ.get("DISPLAY"):
         log.info("DISPLAY already set: %s", os.environ.get("DISPLAY"))
         return
@@ -1034,9 +1048,7 @@ async def lifespan(app: FastAPI):
     _key_refresh_lock = asyncio.Lock()
     JOB_QUEUE         = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
 
-    # If running non-headless and no DISPLAY is present, attempt to start
-    # a virtual display (pyvirtualdisplay/Xvfb). This runs in a thread
-    # executor because pyvirtualdisplay is synchronous.
+    # If configured to start virtual display, attempt to do so when needed.
     loop = asyncio.get_running_loop()
     try:
         await loop.run_in_executor(None, partial(_start_virtual_display_if_needed, ZD_HEADLESS))
@@ -1325,6 +1337,7 @@ if __name__ == "__main__":
 
     # Allow overriding via environment variables
     ZD_HEADLESS = os.environ.get("ZD_HEADLESS", str(ZD_HEADLESS)) in ("1", "true", "True")
+    USE_VIRTUAL_DISPLAY = os.environ.get("USE_VIRTUAL_DISPLAY", str(USE_VIRTUAL_DISPLAY)) in ("1", "true", "True")
 
     # If running locally and not headless, start virtual display if needed
     try: