feat(selenium-devtools-py): export a trace archive through the backend

packages/selenium-devtools-py/scripts/gen_contract.py

@@ -71,6 +71,22 @@ def _ws_scopes(routes_ts: str) -> dict[str, str]:
     return dict(re.findall(r"(\w+):\s*'([^']+)'", m.group(1)))
 
 
+def _trace_export_scopes(trace_export_ts: str) -> dict[str, str]:
+    """`TRACE_EXPORT_SCOPE` — the worker↔backend frames that ask the backend to
+    build a trace and answer with where it landed. Python cannot run the
+    transforms itself, so these two strings are the whole route to a trace."""
+    m = re.search(
+        r"export const TRACE_EXPORT_SCOPE = \{(.*?)\n\} as const",
+        trace_export_ts,
+        re.DOTALL,
+    )
+    if not m:
+        raise SystemExit(
+            "could not find `TRACE_EXPORT_SCOPE` in shared/trace-export.ts"
+        )
+    return dict(re.findall(r"(\w+):\s*'([^']+)'", m.group(1)))
+
+
 def _collector_path(collector_ts: str) -> str:
     """The route the backend serves the page-side collector from."""
     m = re.search(r"export const COLLECTOR_API = \{(.*?)\} as const", collector_ts, re.DOTALL)
@@ -161,6 +177,9 @@ def main() -> int:
     data_keys = _trace_log_keys(types_ts)
     runner_ids = _test_runner_ids(types_ts)
     collector_path = _collector_path((shared / "src" / "collector.ts").read_text())
+    trace_export = _trace_export_scopes(
+        (shared / "src" / "trace-export.ts").read_text()
+    )
     routes_ts = (shared / "src" / "routes.ts").read_text()
     control = _ws_scopes(routes_ts)
     worker_query = _worker_query(routes_ts)
@@ -209,6 +228,14 @@ def main() -> int:
             "three to report into the dashboard that launched it."
         )
 
+    missing_export = [k for k in ("request", "result") if k not in trace_export]
+    if missing_export:
+        raise SystemExit(
+            f"contract drift: TRACE_EXPORT_SCOPE key(s) {missing_export} no "
+            f"longer in shared (present: {sorted(trace_export)}). Python has no "
+            "other route to a trace — it cannot run the transforms itself."
+        )
+
     if REQUIRED_RUNNER_ID not in runner_ids:
         raise SystemExit(
             f"contract drift: runner id {REQUIRED_RUNNER_ID!r} is no longer in "
@@ -241,6 +268,9 @@ def main() -> int:
         f'RERUN_SLOT_TEST_ID = "{rerun_slot["testId"]}"',
         f'ENV_RUNNER_CWD = "{runner_cwd_env}"',
         "",
+        f'SCOPE_TRACE_EXPORT = "{trace_export["request"]}"',
+        f'SCOPE_TRACE_EXPORTED = "{trace_export["result"]}"',
+        "",
         f'ENV_REUSE = "{reuse_env["REUSE"]}"',
         f'ENV_REUSE_HOST = "{reuse_env["HOST"]}"',
         f'ENV_REUSE_PORT = "{reuse_env["PORT"]}"',

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

@@ -18,13 +18,22 @@
 import os
 import subprocess
 import sys
+import threading
 from typing import Optional
 
-from . import backend, instrumentation, lifecycle, rerun
+from . import backend, instrumentation, lifecycle, rerun, trace_export
 from ._contract import CONTRACT_VERSION
 from .capturer import SessionCapturer
-from .run_id import reset_run_id
-from .constants import DEFAULT_HOST, DEFAULT_PORT, ENV_HOST, ENV_PORT, LOGGER_NAME
+from .output_dir import resolve_adapter_output_dir
+from .run_id import reset_run_id, resolve_run_id
+from .constants import (
+    DEFAULT_HOST,
+    DEFAULT_PORT,
+    ENV_HOST,
+    ENV_PORT,
+    ENV_TRACE,
+    LOGGER_NAME,
+)
 from .logcapture import LogCapturer
 from .terminal import TerminalCapturer
 from .transport import WSClient
@@ -66,14 +75,95 @@ def _restore_excepthook() -> None:
 _active: dict = {
     "capturer": None, "transport": None, "process": None, "url": None,
     "handle": None, "terminal": None, "logs": None, "excepthook": None,
+    "trace": False, "traced": False,
 }
 
 
+def _trace_enabled(trace: Optional[bool]) -> bool:
+    """Whether this run writes a trace archive. The argument wins over the
+    environment so a script can opt out of an exported default."""
+    if trace is not None:
+        return trace
+    return os.environ.get(ENV_TRACE, "").lower() in ("1", "true", "yes")
+
+
+#: Serializes exports. Teardown can run on the WS reader thread while a caller
+#: is mid-export on the main one, and `trace_export` holds ONE pending slot: a
+#: second request replaces it, so the first caller's reply is dropped and it
+#: waits out the full timeout while the backend writes the same archive twice.
+#: Holding it across the wait is also what stops a MAIN-THREAD teardown closing
+#: the transport out from under an export still listening on it. Off-thread
+#: callers must not wait on it at all — see export_trace.
+_export_lock = threading.Lock()
+
+
+def export_trace(output_dir: Optional[str] = None) -> Optional[str]:
+    """Write this run's trace archive now. No-op unless trace mode is on.
+
+    Called when the RUN finishes rather than when the process tears down. An
+    interactive run blocks on the dashboard window in between, and CI has no
+    window at all; an artifact that depends on either is an artifact that is
+    missing exactly when it is wanted.
+
+    Only a SUCCESSFUL export closes the door on the teardown fallback. This is
+    public, so a caller may run it early, get nothing, and still expect an
+    archive at the end — latching on the attempt would spend that one chance on
+    a transport that was not ready. The cost is that an unresponsive backend is
+    waited on twice, once here and once at teardown; losing the artifact
+    outright is the worse of the two, and by then the run is already broken.
+    """
+    # Off the main thread this never waits. Teardown can arrive on the WS
+    # reader thread — `_trigger_shutdown` runs it there when nobody is parked
+    # in wait_for_shutdown — and that thread is the ONLY one that can deliver
+    # the reply an in-flight export is blocked on. Waiting for that export from
+    # here deadlocks both until the timeout, and the shutdown's `os._exit`
+    # timer may kill the process first. An interrupted run losing its archive
+    # is the better failure; by construction it was interrupted.
+    on_main = threading.current_thread() is threading.main_thread()
+    if not _export_lock.acquire(blocking=on_main):
+        _log.debug("a trace export is already in flight; not starting another")
+        return None
+    try:
+        if not _active["trace"] or _active["traced"]:
+            return None
+        path = _export_trace(
+            _active["capturer"],
+            output_dir
+            if output_dir is not None
+            else instrumentation.resolved_output_dir(),
+        )
+        if path is not None:
+            _active["traced"] = True
+        return path
+    finally:
+        _export_lock.release()
+
+
+def _export_trace(
+    capturer: Optional[SessionCapturer], output_dir: Optional[str]
+) -> Optional[str]:
+    """Ask the backend for this run's archive. Never raises — a run that
+    captured everything and failed to write a file still passed."""
+    try:
+        session_id = (
+            getattr(capturer, "session_id", None) or resolve_run_id()
+        )
+        return trace_export.export(
+            _active["transport"],
+            output_dir=output_dir or resolve_adapter_output_dir(),
+            session_id=session_id,
+        )
+    except Exception as exc:  # noqa: BLE001
+        _log.warning("trace export skipped (%s)", exc)
+    return None
+
+
 def enable(
     host: Optional[str] = None,
     port: Optional[int] = None,
     *,
     webdriver_cls: Optional[type] = None,
+    trace: Optional[bool] = None,
 ) -> Optional[SessionCapturer]:
     """Connect to the backend and instrument Selenium. Idempotent.
 
@@ -85,6 +175,10 @@ def enable(
     if _active["capturer"] is not None:
         return _active["capturer"]
 
+    # Decided before anything reads it: the screencast recorder, the dashboard
+    # window and the teardown export all branch on this.
+    trace_mode = _trace_enabled(trace)
+
     # Before the backend is launched: the directory a rerun spawns in travels
     # through the environment the backend process inherits. A framework plugin
     # has already published richer commands by now and this leaves those alone.
@@ -118,7 +212,7 @@ def enable(
         return None
 
     capturer = SessionCapturer(transport)
-    instrumentation.install(capturer, webdriver_cls)
+    instrumentation.install(capturer, webdriver_cls, trace=trace_mode)
     # Plain scripts only: a framework plugin calls
     # `set_external_suites`, which turns this back off.
     instrumentation.start_assertion_tracing(capturer)
@@ -132,12 +226,16 @@ def enable(
     url = f"http://{host}:{port}"
     _active.update(
         capturer=capturer, transport=transport, process=process, url=url,
-        terminal=term, logs=logs,
+        terminal=term, logs=logs, trace=trace_mode, traced=False,
     )
 
     # Open the dashboard window and wire exit/signal + control-frame teardown so
     # closing the window (clientDisconnected) or ending the process both tidy up.
-    handle = lifecycle.open_dashboard(url) if lifecycle.auto_open_enabled() else None
+    handle = (
+        lifecycle.open_dashboard(url)
+        if lifecycle.auto_open_enabled(trace=trace_mode)
+        else None
+    )
     _active["handle"] = handle
     lifecycle.register_exit_handlers(disable, handle)
     return capturer
@@ -156,13 +254,22 @@ def disable() -> None:
         # learns afterwards could still be sent. `finalize_run` therefore reads
         # the live exception itself. Must precede transport.close() either way.
         instrumentation.finalize_run(capturer)
+    # Read before uninstall clears it: the trace belongs beside this run's
+    # video, and the fallback is the cwd — the repo root, for a runner invoked
+    # from one.
+    output_dir = instrumentation.resolved_output_dir()
     instrumentation.uninstall()
     term = _active["terminal"]
     if term is not None:  # restore stdout/stderr before tearing the transport down
         term.stop()
     logs = _active["logs"]
     if logs is not None:  # detach the logging handler + restore logger levels
         logs.stop()
+    # Fallback for a plain script that never called export_trace() itself.
+    # Before the transport closes: the answer comes back on this same socket.
+    # Through export_trace, not around it: one lock and one latch, so a public
+    # call still in flight is waited for rather than raced.
+    export_trace(output_dir)
     transport = _active["transport"]
     if transport is not None:
         transport.close()
@@ -179,9 +286,10 @@ def disable() -> None:
             process.wait(timeout=5)
         except subprocess.TimeoutExpired:
             process.kill()  # backend ignored SIGTERM — force it
+    trace_export.reset()
     _active.update(
         capturer=None, transport=None, process=None, url=None, handle=None,
-        terminal=None, logs=None, excepthook=None,
+        terminal=None, logs=None, excepthook=None, trace=False, traced=False,
     )
 
 

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

@@ -25,6 +25,9 @@
 RERUN_SLOT_TEST_ID = "{{testId}}"
 ENV_RUNNER_CWD = "DEVTOOLS_RUNNER_CWD"
 
+SCOPE_TRACE_EXPORT = "traceExport"
+SCOPE_TRACE_EXPORTED = "traceExported"
+
 ENV_REUSE = "DEVTOOLS_APP_REUSE"
 ENV_REUSE_HOST = "DEVTOOLS_APP_HOST"
 ENV_REUSE_PORT = "DEVTOOLS_APP_PORT"

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

@@ -76,6 +76,15 @@
 # floor its dependencies require; below it the process starts and then dies on
 # syntax it cannot parse, which surfaces here only as "exited before reporting
 # a port". Checked up front so the message names the real problem.
+# How long to wait for the backend to answer a trace export. The archive is
+# assembled from a whole run's frames, so it is not instant; but a run that
+# captured everything and then hung waiting for a file is worse than one that
+# reports the wait timed out.
+TRACE_EXPORT_TIMEOUT_S = 60.0
+
+#: Opt in to writing a trace archive at the end of the run.
+ENV_TRACE = "DEVTOOLS_TRACE"
+
 MIN_NODE_MAJOR = 18
 NODE_VERSION_TIMEOUT_S = 5.0
 

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

@@ -57,6 +57,11 @@
 _skip_frames_cache: Optional[tuple] = None
 
 
+class _SkipScreencast(Exception):
+    """Control-flow marker: no recorder for this session. Named rather than a
+    branch so the reason lands in one place with the other bring-up failures."""
+
+
 def _skip_frames() -> tuple:
     """Call-source skip fragments: the adapter package + the REAL selenium
     library dir (resolved from selenium.__file__), cached. Resolving the actual
@@ -241,6 +246,10 @@ def _capture_source(capturer: SessionCapturer, call_src: Optional[str]) -> None:
     # Set by enable()'s excepthook when an exception reaches top level. The
     # synthetic suite's final state reads this rather than assuming success.
     "run_failed": False,
+    # Trace mode. The archive carries per-command screenshots, not the
+    # screencast — `screencastFrames` does not cross the wire yet — so
+    # recording one writes a .webm nothing reads. Set by install().
+    "trace": False,
 }
 
 
@@ -316,6 +325,18 @@ def _attach_performance(
         _log.debug("could not replace the navigation row: %s", exc)
 
 
+def resolved_output_dir() -> Optional[str]:
+    """The ``test-results`` dir this run resolved from its first test file, or
+    None if no command carried a user call source. Screencast videos already
+    write here; a trace belongs beside them rather than in the cwd, which for a
+    runner invoked from a repo root is the repo root.
+
+    Cleared by ``uninstall``, so a caller tearing a run down must read it before
+    that rather than after.
+    """
+    return _state.get("output_dir")
+
+
 def _begin_screencast_run(entry: Optional[dict], shot: Optional[str] = None) -> None:
     """Let a pushed stream start keeping frames. Idempotent, never raises.
 
@@ -496,6 +517,10 @@ def _ensure_session_setup(driver: Any, capturer: SessionCapturer) -> Optional[di
     except Exception as exc:  # noqa: BLE001 — capture must never break the test
         _log.warning("BiDi attach threw: %s", exc)
     try:
+        if _state["trace"]:
+            # The archive's frames are the per-command screenshots; the video is
+            # a live-dashboard artifact, and trace mode opens no dashboard.
+            raise _SkipScreencast
         recorder = ScreencastRecorder()
         recorder.start(driver)
         entry["screencast"] = recorder
@@ -520,6 +545,8 @@ def _ensure_session_setup(driver: Any, capturer: SessionCapturer) -> Optional[di
         entry["screencast_push"] = push
         if push is None:
             _log.info("screencast recording started (one frame per command)")
+    except _SkipScreencast:
+        _log.info("trace mode — skipping the screencast recording")
     except Exception as exc:  # noqa: BLE001
         _log.warning("screencast start threw: %s", exc)
     try:
@@ -695,7 +722,12 @@ def finalize_run(capturer: SessionCapturer) -> None:
     _send_default_suite(capturer, _live_run_state())
 
 
-def install(capturer: SessionCapturer, webdriver_cls: Optional[type] = None) -> None:
+def install(
+    capturer: SessionCapturer,
+    webdriver_cls: Optional[type] = None,
+    *,
+    trace: bool = False,
+) -> None:
     if _state["installed"]:
         return
     if webdriver_cls is None:
@@ -775,6 +807,7 @@ def patched_execute(self, driver_command: str, params: Any = None):  # noqa: ANN
     _state.update(
         installed=True, cls=webdriver_cls, orig=orig_execute,
         sessions=weakref.WeakKeyDictionary(), output_dir=None, default_suite=None,
+        trace=trace,
     )