NF Firefox headless WebGL support; FIX Vertex in Firefox WebGL

.github/workflows/run_tests.yml

@@ -60,6 +60,8 @@ jobs:
         restore-keys: |
           ${{ runner.os }}-playwright-
 
+    # Firefox is left out on purpose, to keep CI runs short: its tests
+    # skip here and are run locally instead (see "Tests" in AGENTS.md).
     - name: Install Playwright Chromium
       run: playwright install --with-deps --only-shell chromium
 

AGENTS.md

@@ -13,6 +13,7 @@ The package has two Cython extensions (`cortex/formats.pyx` for surface mesh I/O
 ```bash
 uv sync --all-extras                              # headless extra (playwright) + dev group (test/type-checking deps) by default; builds the Cython extensions
 uv run playwright install --only-shell chromium   # browser for headless WebGL tests
+uv run playwright install firefox                 # optional: also run the Firefox tests (needs Xvfb, see Tests)
 ```
 
 Versioning is by `setuptools-scm` from git tags (written to `cortex/_version.py`); never edit version files by hand. A new git worktree gets its own `.venv` — rerun `uv sync` there to (re)build the compiled extensions.
@@ -28,9 +29,10 @@ uv run pytest -p no:cacheprovider --no-cov -x cortex/tests/test_formats.py  # qu
 
 - `pytest.ini` sets a 240 s per-test timeout (via `pytest-timeout`) because headless browser sessions can hang; override per-test with `@pytest.mark.timeout(N)`.
 - Tests self-skip based on available tools: Inkscape (quickflat/dataset overlay tests), Playwright Chromium (`cortex/tests/testing_utils.py:has_playwright` — WebGL headless tests), and FreeSurfer's `mri_surf2surf`.
+- Firefox tests (every test ID containing `firefox`; `pytest -k firefox` runs just those) self-skip unless Playwright's own Firefox build is installed (`testing_utils.py:has_playwright_firefox`). Headless Firefox needs an X display for WebGL, so they also need `Xvfb` (`apt-get install xvfb`) unless `DISPLAY` is set. These tests catch driver differences Chromium hides. For example, Mesa (Firefox) counts declared-but-unused vertex attributes against the 16-attribute limit, which ANGLE (Chromium) does not. Run them after touching `shaderlib.js` or the attributes `dataset.js` dispatches.
 - `pytest.ini` deselects `-m slow` by default. The one slow test, `cortex/tests/test_autoflatten.py:test_autoflatten_subject_end_to_end`, really runs autoflatten (15-30 min); it needs FreeSurfer sourced with `fsaverage` in `$SUBJECTS_DIR`, autoflatten installed, and `$PYCORTEX_TEST_FS_SUBJECT` naming a recon-all'd subject — the bundled S1 is a pycortex filestore entry, not a FreeSurfer subject directory, so it cannot be flattened. Run it with `pytest -m slow`.
 - Nearly all tests use the stub subject `S1` bundled in `filestore/db/S1` (transform `fullhead`, 304380 vertices, volume shape `(31, 100, 100)`).
-- CI (`.github/workflows/run_tests.yml`) runs `pytest --cov=./` on a matrix of Python versions, with Inkscape and Playwright Chromium installed. The only enforced lint is codespell (config in `pyproject.toml`). mypy is configured in `pyproject.toml` and installed with the dev group but not run in CI — new Python code should still carry type annotations.
+- CI (`.github/workflows/run_tests.yml`) runs `pytest --cov=./` on a matrix of Python versions, with Inkscape and Playwright Chromium installed. It deliberately does not install Firefox, to keep CI runs short, so the Firefox tests only run locally. The only enforced lint is codespell (config in `pyproject.toml`). mypy is configured in `pyproject.toml` and installed with the dev group but not run in CI — new Python code should still carry type annotations.
 
 ### Docs
 
@@ -67,7 +69,7 @@ Flow: dataviews → JSON + PNG mosaics (`data.py:Package`) + compressed CTM surf
 - The bundled Three.js is **r69** — very old; the shader pipeline depends on its conventions, do not casually upgrade.
 - All GLSL lives in `resources/js/shaderlib.js` as arrays of string lines assembled per-configuration — there are no `.glsl` files.
 - CTM packs reorder vertices: `cortex.utils.get_ctmmap` / `get_ctm2webgl_map` translate between CTM/WebGL ordering and the original surface ordering. Indexing viewer data with original vertex indices without remapping is a classic bug.
-- `cortex/export/headless.py` (`headless_viewer`) runs the viewer in headless Chromium via Playwright for screenshots/tests.
+- `cortex/export/headless.py` (`headless_viewer`) runs the viewer in headless Chromium (default) or Firefox (`browser="firefox"`) via Playwright for screenshots/tests.
 - Some JS resources use CRLF line endings (e.g. `dataset.js`); keep the existing endings when editing.
 - `setup.py` explicitly enumerates `cortex.webgl` `package_data` patterns — new resource subdirectories must be added there or they won't ship in wheels.
 

cortex/export/headless.py

@@ -1,6 +1,6 @@
 """headless.py - context manager that runs the pycortex WebGL viewer inside a
-headless Chromium browser via Playwright, allowing ``save_3d_views`` to
-produce screenshots without any manual browser interaction.
+headless Chromium or Firefox browser via Playwright, allowing ``save_3d_views``
+to produce screenshots without any manual browser interaction.
 
 The approach:
 1. Start the Tornado WebGL server with ``open_browser=False`` so that it does
@@ -34,20 +34,31 @@
 
 Requirements
 ------------
+Playwright, plus the one browser you render with; you do not need both::
+
     pip install playwright
-    playwright install chromium
+    playwright install chromium   # for browser="chromium" (the default)
+    playwright install firefox    # only for browser="firefox"
+
+``browser="firefox"`` additionally requires the ``Xvfb`` binary to be on
+``PATH`` (e.g. ``apt-get install xvfb``) unless a ``DISPLAY`` is already set,
+since Firefox (unlike Chromium's bundled SwiftShader) needs a real or virtual
+X server to create a software WebGL context.
 """
 
 import concurrent.futures
 import contextlib
 import logging
 import os
 import queue
+import select
 import shutil
+import subprocess
 import tempfile
 import threading
 import time
-from typing import Any, Callable, Mapping, Optional, TypeVar
+from collections.abc import Generator
+from typing import Any, Callable, Literal, Mapping, Optional, TypeVar, get_args
 
 T = TypeVar("T")
 
@@ -65,6 +76,102 @@
     "--disable-dev-shm-usage",
 ]
 
+BrowserName = Literal["chromium", "firefox"]
+
+
+@contextlib.contextmanager
+def _xvfb_display() -> Generator[str]:
+    """Run a private Xvfb server for the duration of the context.
+
+    Yields its ``DISPLAY`` string (e.g. ``":99"``). Xvfb picks a free display
+    number itself and reports it via ``-displayfd`` once it accepts
+    connections.
+
+    Raises
+    ------
+    RuntimeError
+        If the ``Xvfb`` binary is not installed, or it fails to start.
+    """
+    read_fd, write_fd = os.pipe()
+    try:
+        try:
+            proc = subprocess.Popen(
+                ["Xvfb", "-displayfd", str(write_fd),
+                 "-screen", "0", "1280x1024x24", "-nolisten", "tcp"],
+                pass_fds=(write_fd,),
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        finally:
+            os.close(write_fd)
+    except FileNotFoundError as exc:
+        os.close(read_fd)
+        raise RuntimeError(
+            "Firefox headless WebGL rendering requires Xvfb, which was not "
+            "found on PATH. Install it with e.g. `apt-get install xvfb`, or "
+            "set the DISPLAY environment variable to an existing X server."
+        ) from exc
+
+    try:
+        with os.fdopen(read_fd) as displayfd:
+            if not select.select([displayfd], [], [], 10.0)[0]:
+                raise RuntimeError("Timed out waiting for Xvfb to start.")
+            display_num = displayfd.readline().strip()
+        if not display_num:
+            raise RuntimeError("Xvfb exited unexpectedly during startup.")
+        yield f":{display_num}"
+    finally:
+        try:
+            proc.terminate()
+            proc.wait(timeout=10)
+        except Exception:
+            logger.warning("Failed to terminate Xvfb process", exc_info=True)
+
+
+@contextlib.contextmanager
+def webgl_browser(
+    playwright: Any, browser: BrowserName = "chromium"
+) -> Generator[Any]:
+    """Launch a headless Playwright browser with software WebGL enabled.
+
+    Yields the browser, and closes it (and the private Xvfb server started for
+    Firefox when no ``DISPLAY`` is set) on exit.
+
+    Raises
+    ------
+    ValueError
+        If ``browser`` is not one of the supported browsers.
+    """
+    with contextlib.ExitStack() as stack:
+        if browser == "firefox":
+            # Unlike Chromium (which bundles SwiftShader and can create a
+            # software WebGL context with no display server at all),
+            # Firefox creates its GL context via GLX and needs a real (or
+            # virtual) X display to do so, even when launched "headless".
+            # If the caller hasn't already provided one, start a private
+            # Xvfb so WebGL still works with no GPU and no visible display.
+            launch_env = None
+            if not os.environ.get("DISPLAY"):
+                display = stack.enter_context(_xvfb_display())
+                launch_env = {**os.environ, "DISPLAY": display}
+            instance = playwright.firefox.launch(
+                headless=True,
+                env=launch_env,
+                # Bypasses Firefox's GPU driver blocklist, which could refuse WebGL when
+                # DISPLAY is backed by a blocklisted driver (not applicable for Xvfb).
+                firefox_user_prefs={"webgl.force-enabled": True},
+            )
+        elif browser == "chromium":
+            instance = playwright.chromium.launch(
+                headless=True, args=SWIFTSHADER_CHROMIUM_ARGS
+            )
+        else:
+            raise ValueError(
+                f"Unknown browser {browser!r}; expected one of {get_args(BrowserName)}"
+            )
+        stack.callback(instance.close)
+        yield instance
+
 
 def _wait_for_viewer_loaded(handle, timeout: float = 60.0) -> None:
     """Block until ``window.viewer.loaded`` resolves in the browser.
@@ -182,7 +289,7 @@ def __init__(self, download_dir: Optional[str] = None) -> None:
         self._thread: Optional[threading.Thread] = None
         # Set during worker startup; only touched by the worker thread.
         self._playwright: Any = None
-        self._browser: Any = None
+        self._browser_stack = contextlib.ExitStack()
         self._page: Any = None
         # Browser-side errors collected by Playwright event listeners.
         # Written by the worker thread, read by the main thread via
@@ -192,11 +299,19 @@ def __init__(self, download_dir: Optional[str] = None) -> None:
 
     # -- public API -------------------------------------------------------- #
 
-    def start(self, url: str, *, timeout: float = 60.0) -> None:
-        """Launch the worker thread, open Chromium, and navigate to *url*.
+    def start(
+        self, url: str, *, timeout: float = 60.0, browser: BrowserName = "chromium"
+    ) -> None:
+        """Launch the worker thread, open the browser, and navigate to *url*.
 
         Blocks until the page has finished loading (or raises on failure).
 
+        Parameters
+        ----------
+        browser : str
+            Which Playwright browser type to launch: ``"chromium"`` (default)
+            or ``"firefox"``.
+
         Raises
         ------
         RuntimeError
@@ -206,6 +321,7 @@ def start(self, url: str, *, timeout: float = 60.0) -> None:
         """
         self._url = url
         self._nav_timeout = timeout
+        self._browser_type = browser
         self._thread = threading.Thread(
             target=self._worker, name="PlaywrightThread", daemon=True
         )
@@ -296,18 +412,18 @@ def _worker(self) -> None:
         except ImportError as exc:
             self._error = ImportError(
                 "playwright is required for headless rendering. Install it with:\n"
-                "    pip install playwright && playwright install chromium"
+                f"    pip install playwright && playwright install {self._browser_type}"
             )
             self._error.__cause__ = exc
             self._ready_future.set_result(None)  # unblock caller
             return
 
         try:
             self._playwright = sync_playwright().start()
-            self._browser = self._playwright.chromium.launch(
-                headless=True, args=SWIFTSHADER_CHROMIUM_ARGS
+            browser = self._browser_stack.enter_context(
+                webgl_browser(self._playwright, self._browser_type)
             )
-            self._page = self._browser.new_page()
+            self._page = browser.new_page()
 
             # Register listeners *before* navigation so we capture
             # errors that fire during page load (e.g. WebGL failures).
@@ -410,7 +526,7 @@ def _cleanup(self) -> None:
         individually guarded so a failure in one does not prevent the rest."""
         for obj, method in [
             (self._page, "close"),
-            (self._browser, "close"),
+            (self._browser_stack, "close"),
             (self._playwright, "stop"),
         ]:
             if obj is not None:
@@ -437,9 +553,10 @@ def headless_viewer(
     *,
     timeout: float = 60.0,
     download_dir: Optional[str] = None,
+    browser: BrowserName = "chromium",
 ):
     """Context manager that yields a connected ``JSMixer`` handle rendered in a
-    headless Chromium browser.
+    headless browser.
 
     Parameters
     ----------
@@ -458,6 +575,11 @@ def headless_viewer(
         ``handle._pw_thread.downloads``; ``handle._pw_thread.wait_for_download()``
         waits for one. Default None, meaning a temporary directory removed when
         the viewer closes.
+    browser : str
+        Which Playwright browser to use: ``"chromium"`` (default) or
+        ``"firefox"``. The ``"firefox"`` browser requires the ``Xvfb``
+        binary to be installed unless ``DISPLAY`` is already set (see
+        module docstring).
 
     Yields
     ------
@@ -543,7 +665,7 @@ def _await_client() -> None:
             # Launch the browser and navigate.  python_interface.js runs on
             # load and sends "connect" over WebSocket, which unblocks
             # server.get_client().
-            pw_thread.start(url, timeout=timeout)
+            pw_thread.start(url, timeout=timeout, browser=browser)
 
             # Wait (bounded) for the getter to return; the timeout guard
             # surfaces hung state clearly instead of blocking indefinitely.

cortex/export/panels.py

@@ -17,6 +17,7 @@
 from matplotlib.figure import Figure
 
 from ..dataset import Dataview
+from .headless import BrowserName
 from .save_views import save_3d_views, ViewParams
 from ._default_params import (
     params_inflatedless_lateral_medial_ventral,
@@ -52,6 +53,7 @@ def plot_panels(
     interpolation: str = "nearest",
     layers: int = 1,
     headless: bool = False,
+    browser: BrowserName = "chromium",
 ) -> Figure:
     """Plot on the same figure a number of views, as defined by a list of panel
     specifications.
@@ -102,13 +104,22 @@ def plot_panels(
         plotting the data. (Default: 1).
 
     headless: bool
-        If True, render using a headless Chromium browser via Playwright instead
+        If True, render using a headless browser via Playwright instead
         of requiring the user to manually open a browser window. This allows
         the function to run fully autonomously without any user interaction.
         Requires ``playwright`` to be installed (``pip install playwright``) and
-        Chromium to be available (``playwright install chromium``).
-        Software WebGL (SwiftShader) is used, so no GPU or display server is
-        needed. (Default: False)
+        the chosen ``browser`` to be installed: ``playwright install
+        chromium`` for the default, or ``playwright install firefox`` for
+        ``browser="firefox"`` (only that one is needed).
+        Software WebGL rendering is used, so no GPU is needed, and with the
+        default Chromium no display server either (``browser="firefox"``
+        needs one, see below). (Default: False)
+
+    browser: str
+        Which Playwright browser to use when ``headless=True``: ``"chromium"``
+        (default) or ``"firefox"``. Ignored when ``headless=False``.
+        ``"firefox"`` additionally requires ``Xvfb`` to be installed unless
+        ``DISPLAY`` is already set (see ``cortex.export.headless``).
 
     Returns
     -------
@@ -150,6 +161,7 @@ def plot_panels(
         interpolation=interpolation,
         layers=layers,
         headless=headless,
+        browser=browser,
     )
 
     fig = plt.figure(figsize=figsize)

cortex/export/save_views.py

@@ -11,6 +11,7 @@
 import cortex
 
 from ..dataset import Dataview
+from .headless import BrowserName
 
 file_pattern = "{base}_{view}_{surface}.png"
 
@@ -46,12 +47,13 @@ def save_3d_views(
     trim: bool = True,
     sleep: float = 10,
     headless: bool = False,
+    browser: BrowserName = "chromium",
 ) -> list[str]:
     """Saves 3D views of `volume` under multiple specifications.
 
     By default (``headless=False``), a webgl viewer is launched and a display
-    server is required.  With ``headless=True``, a headless Chromium browser
-    is used instead, so no display server or GPU is needed.
+    server is required.  With ``headless=True``, a headless browser (Chromium
+    or Firefox) is used instead, so no display server or GPU is needed.
 
     Parameters
     ----------
@@ -96,13 +98,22 @@ def save_3d_views(
         Time in seconds, to let the viewer open.
 
     headless: bool
-        If True, render using a headless Chromium browser via Playwright instead
+        If True, render using a headless browser via Playwright instead
         of requiring the user to manually open a browser window.  This allows
         the function to run fully autonomously without any user interaction.
         Requires ``playwright`` to be installed (``pip install playwright``) and
-        Chromium to be available (``playwright install chromium``).
-        Software WebGL (SwiftShader) is used, so no GPU or display server is
-        needed.  (Default: False)
+        the chosen ``browser`` to be installed: ``playwright install
+        chromium`` for the default, or ``playwright install firefox`` for
+        ``browser="firefox"`` (only that one is needed).
+        Software WebGL rendering is used, so no GPU is needed, and with the
+        default Chromium no display server either (``browser="firefox"``
+        needs one, see below).  (Default: False)
+
+    browser: str
+        Which Playwright browser to use when ``headless=True``: ``"chromium"``
+        (default) or ``"firefox"``. Ignored when ``headless=False``.
+        ``"firefox"`` additionally requires ``Xvfb`` to be installed unless
+        ``DISPLAY`` is already set (see ``cortex.export.headless``).
 
     Returns
     -------
@@ -117,7 +128,7 @@ def save_3d_views(
     if headless:
         from cortex.export.headless import headless_viewer as _headless_viewer
 
-        cm = _headless_viewer(volume, viewer_params)
+        cm = _headless_viewer(volume, viewer_params, browser=browser)
     else:
         cm = contextlib.nullcontext(cortex.webshow(volume, **viewer_params))