feat(selenium-devtools-py): register the collector at document start via BiDi preload

examples/selenium/python-test/login.py

@@ -47,7 +47,7 @@
     # wait.until(EC.visibility_of_element_located((By.ID, "flash")))
     assert "/secure" in driver.current_url, driver.current_url
     flash = driver.find_element(By.ID, "flash").text
-    assert "You logged into a secure area1" in flash, flash
+    assert "You logged into a secure area" in flash, flash
 
     # wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "a.button")))
     driver.find_element(By.CSS_SELECTOR, "a.button").click()

packages/selenium-devtools-py/src/selenium_devtools/bidi_preload.py

@@ -0,0 +1,68 @@
+"""Document-start registration of the collector via a BiDi preload script.
+
+The ``<script>``-append path in ``snapshot.py`` can only instrument the document
+that exists when it runs, and a ``<script>`` dies with its document. So every
+navigation produces a document we learn about afterwards, and everything built on
+that — when to re-inject, when to drain, which action owns the new DOM — is
+reconstruction, and races. A preload registered for the session runs in EVERY
+document before any of that document's own script, so each one instruments itself
+and anchors its own DOM at its own ``performance.timeOrigin``. Nothing to detect,
+nothing to attribute.
+
+Mirrors ``core/bidi-preload.ts``, which the JS adapters use for the same reason.
+
+Requires a session created with ``webSocketUrl: true``; ``driver.script`` raises
+without it, so this reports False and the caller keeps the ``<script>`` path.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from .constants import BIDI_CAPABILITY, LOGGER_NAME
+
+_log = logging.getLogger(f"{LOGGER_NAME}.preload")
+
+
+def as_function_declaration(source: str) -> str:
+    """Wrap raw collector source as the function declaration BiDi expects.
+
+    Not the ``wrap_injectable`` IIFE: a preload script IS a function, so the
+    bundle's top-level await works directly in its body.
+    """
+    return f"async () => {{ {source} }}"
+
+
+def register_collector_preload(driver: Any, source: str) -> bool:
+    """Register the collector to run at document-start in every document of this
+    session. False means the caller must fall back to ``<script>`` injection.
+
+    Registered with NO browsing-context id, which is what makes it global —
+    contexts created later are covered, which is exactly the set this exists to
+    catch. `driver.script.pin` is the public API for it; its docstring says
+    "current browsing context", but it forwards no `contexts` argument, and BiDi
+    treats that as every context. That reliance is pinned by a test.
+    """
+    caps = getattr(driver, "caps", None)
+    if not (isinstance(caps, dict) and caps.get(BIDI_CAPABILITY)):
+        # Not a warning: the caller already reports the missing capability once,
+        # and the `<script>` path still captures DOM — just with the races this
+        # module exists to remove.
+        _log.debug(
+            "%s not set on the session — no document-start preload, using "
+            "per-document injection",
+            BIDI_CAPABILITY,
+        )
+        return False
+    try:
+        script = driver.script
+        script.pin(as_function_declaration(source))
+    except Exception as exc:  # noqa: BLE001 — any selenium/BiDi failure degrades
+        _log.warning(
+            "BiDi preload unavailable, falling back to per-document injection: %s",
+            exc,
+        )
+        return False
+    _log.info("collector registered at document-start (BiDi preload)")
+    return True

packages/selenium-devtools-py/src/selenium_devtools/instrumentation.py

@@ -19,7 +19,7 @@
 import weakref
 from typing import Any, Optional
 
-from . import bidi, frames
+from . import bidi, bidi_preload, frames
 from .capturer import SessionCapturer
 from .collector_source import reset_cache as reset_collector_cache
 from .constants import (
@@ -32,7 +32,11 @@
 )
 from .output_dir import resolve_adapter_output_dir
 from .screencast import ScreencastRecorder
-from .snapshot import SnapshotCapturer, start_snapshot_capture
+from .snapshot import (
+    SnapshotCapturer,
+    collector_source_text,
+    start_snapshot_capture,
+)
 from .sources import read_source
 from .utils import call_source, now_ms
 
@@ -325,7 +329,10 @@ def _ensure_session_setup(driver: Any, capturer: SessionCapturer) -> Optional[di
         if entry["session_id"] == session_id:
             return entry
         _close_entry(capturer, entry)  # same driver, new session
-    entry = {"session_id": session_id, "screencast": None, "snapshot": None}
+    entry = {
+        "session_id": session_id, "screencast": None, "snapshot": None,
+        "preloaded": False,
+    }
     try:
         sessions[driver] = entry
     except TypeError:  # not weak-referenceable
@@ -347,16 +354,30 @@ def _ensure_session_setup(driver: Any, capturer: SessionCapturer) -> Optional[di
     except Exception as exc:  # noqa: BLE001
         _log.warning("screencast start threw: %s", exc)
     try:
+        # Register the collector at document-start FIRST, while this runs before
+        # the session's first command executes — a preload registered after a
+        # navigation has already missed the document it needed to instrument.
+        # Falls back to per-document `<script>` injection when BiDi is absent.
+        origin = _backend_origin(capturer)
+        source = collector_source_text(backend=origin)
+        preloaded = bool(
+            source and bidi_preload.register_collector_preload(driver, source)
+        )
+        entry["preloaded"] = preloaded
         # Inject the packages/script DOM observer so the snapshot iframe fills.
         # Use a capture-bypassing execute_script so injection/readback scripts
-        # don't pollute the Actions timeline.
+        # don't pollute the Actions timeline. With the preload registered the
+        # capturer injects nothing and exists only to DRAIN the buffer.
         entry["snapshot"] = start_snapshot_capture(
             driver,
             execute_fn=_guarded_execute_script(driver),
-            backend=_backend_origin(capturer),
+            backend=origin,
+            preloaded=preloaded,
         )
         snapshot_cap = entry["snapshot"]
-        if snapshot_cap is not None and snapshot_cap.injected:
+        if preloaded:
+            pass  # already reported by register_collector_preload
+        elif snapshot_cap is not None and snapshot_cap.injected:
             _log.info("DOM snapshot collector injected")
         elif snapshot_cap is not None:
             # Kept rather than dropped, so a later command retries. Saying so

packages/selenium-devtools-py/src/selenium_devtools/snapshot.py

@@ -50,6 +50,10 @@
 #: Cheap readiness probe used after injection.
 _READY_SCRIPT = 'return typeof window.wdioTraceCollector !== "undefined";'
 
+#: Distinguishes "the read itself failed" from "the collector is absent here",
+#: which the page returns as a plain null and which is the recovery signal.
+_UNREADABLE = object()
+
 #: Session-lost signatures — expected during teardown, so silenced rather than
 #: logged (matches the JS adapter's error filter).
 _QUIET_ERRORS = ("ECONNREFUSED", "no such session", "invalid session id")
@@ -93,12 +97,17 @@ def wrap_injectable(script_content: str) -> str:
     return f"(async function() {{ {script_content} }})()"
 
 
-def load_injectable_script(
+def collector_source_text(
     path: Optional[str] = None,
     *,
     backend: Optional[tuple] = None,
 ) -> Optional[str]:
-    """The IIFE-wrapped collector source, or None if it cannot be obtained.
+    """The collector's RAW source, or None if it cannot be obtained.
+
+    Raw rather than wrapped, because the two ways of installing it need
+    different envelopes: a ``<script>`` body needs the async IIFE
+    (`wrap_injectable`), while a BiDi preload script is itself a function
+    declaration and takes the source in its body.
 
     ``backend`` is the ``(host, port)`` of the connected backend, which serves
     the collector out of the package it depends on — the only route that works
@@ -110,7 +119,7 @@ def load_injectable_script(
     if path is None and backend is not None:
         source = fetch_collector_source(*backend)
         if source is not None:
-            return wrap_injectable(source)
+            return source
     resolved = path or resolve_script_path()
     if not resolved:
         _warn(
@@ -120,11 +129,20 @@ def load_injectable_script(
         return None
     try:
         with open(resolved, "r", encoding="utf-8") as handle:
-            content = handle.read()
+            return handle.read()
     except OSError as exc:
         _warn(f"could not read injected script ({resolved}): {exc}")
         return None
-    return wrap_injectable(content)
+
+
+def load_injectable_script(
+    path: Optional[str] = None,
+    *,
+    backend: Optional[tuple] = None,
+) -> Optional[str]:
+    """The IIFE-wrapped collector source for the ``<script>`` path, or None."""
+    source = collector_source_text(path, backend=backend)
+    return None if source is None else wrap_injectable(source)
 
 
 def normalize_mutations(trace_data: Any) -> List[Any]:
@@ -155,11 +173,16 @@ def __init__(
         *,
         script_path: Optional[str] = None,
         backend: Optional[tuple] = None,
+        preloaded: bool = False,
     ) -> None:
         self._execute = execute_fn
         self._script_path = script_path
         self._backend = backend
-        self._injected = False
+        # A document-start preload already installed the collector in every
+        # document of this session, so there is nothing to inject — but the
+        # capturer is still what DRAINS the buffer, so it is still needed.
+        self._preloaded = preloaded
+        self._injected = preloaded
 
     @property
     def injected(self) -> bool:
@@ -171,17 +194,32 @@ def inject(self) -> bool:
         Navigation wipes the injected collector, so we probe the live page each
         call and re-install if it's gone (matching the JS adapter, which injects
         per navigation) rather than trusting a one-time flag. Failures are logged
-        no-ops."""
-        wrapped = load_injectable_script(self._script_path, backend=self._backend)
-        if wrapped is None:
-            return False
+        no-ops. A no-op when a document-start preload is registered: the
+        collector is present in every document before any of its script runs,
+        which is the whole point of the preload. A document where the preload
+        nonetheless did not take is recovered by `pull_mutations`, whose own null
+        is the only signal that says so."""
+        if self._preloaded:
+            return True
         try:
             if self._execute(_READY_SCRIPT) is True:
                 self._injected = True
                 return True
         except BaseException as exc:  # noqa: BLE001 — probe failure → try install
             if not _is_quiet_error(exc):
                 _warn(f"readiness probe failed: {exc}")
+        self._injected = self._install_now()
+        if not self._injected:
+            _warn("collector not detected immediately after injection")
+        return self._injected
+
+    def _install_now(self) -> bool:
+        """Install the collector into the CURRENT document via ``<script>``.
+        Used both by `inject` and by the drain's recovery, which already knows
+        the collector is absent and so skips the readiness probe."""
+        wrapped = load_injectable_script(self._script_path, backend=self._backend)
+        if wrapped is None:
+            return False
         try:
             self._execute(_INJECT_SCRIPT, wrapped)
             ready = self._execute(_READY_SCRIPT)
@@ -190,20 +228,36 @@ def inject(self) -> bool:
                 _warn(f"injection failed: {exc}")
             return False
         self._injected = ready is True
-        if ready is not True:
-            _warn("collector not detected immediately after injection")
         return self._injected
 
-    def pull_mutations(self) -> List[Any]:
-        """Read and drain the buffered mutations (``getTraceData()`` resets the
-        page-side buffer). Returns [] on any failure or when nothing's buffered."""
+    def _read_trace(self) -> Any:
+        """The page's trace payload, None when the collector is absent from this
+        document, or `_UNREADABLE` when the read itself failed."""
         try:
-            trace_data = self._execute(_READ_TRACE_SCRIPT)
+            return self._execute(_READ_TRACE_SCRIPT)
         except BaseException as exc:  # noqa: BLE001
             if not _is_quiet_error(exc):
                 _warn(f"trace read failed: {exc}")
+            return _UNREADABLE
+
+    def pull_mutations(self) -> List[Any]:
+        """Read and drain the buffered mutations (``getTraceData()`` resets the
+        page-side buffer). Returns [] on any failure or when nothing's buffered.
+
+        A null payload means the collector is not in THIS document, and that is
+        the only signal which says so — a preload can still miss one (its script
+        can throw, or a document can predate registration), and with the preload
+        registered nothing else probes. Recovered once here, mirroring core's
+        `drainCollectorWithRecovery`, which costs nothing on the happy path
+        because the drain happens either way."""
+        data = self._read_trace()
+        if data is _UNREADABLE:
             return []
-        return normalize_mutations(trace_data)
+        if data is None and self._install_now():
+            data = self._read_trace()
+            if data is _UNREADABLE:
+                return []
+        return normalize_mutations(data)
 
 
 def start_snapshot_capture(
@@ -212,6 +266,7 @@ def start_snapshot_capture(
     script_path: Optional[str] = None,
     execute_fn: Optional[ExecuteFn] = None,
     backend: Optional[tuple] = None,
+    preloaded: bool = False,
 ) -> Optional[SnapshotCapturer]:
     """Build a ``SnapshotCapturer`` and attempt the first injection. Returns the
     capturer — including when that injection fails, so a later command can retry
@@ -223,7 +278,9 @@ def start_snapshot_capture(
     if not callable(run):
         _warn("driver has no execute_script — snapshot capture skipped")
         return None
-    capturer = SnapshotCapturer(run, script_path=script_path, backend=backend)
+    capturer = SnapshotCapturer(
+        run, script_path=script_path, backend=backend, preloaded=preloaded
+    )
     # The capturer is returned even when this first injection fails. It used to
     # return None, which made the failure terminal: the caller stores None, its
     # post-command refresh skips a missing capturer, and `inject()` is never