Coverage for src/lilbee/crawler/bootstrap.py: 100%
183 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Playwright headless Chromium shell detection and install."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import os
8import re
9import shlex
10import sys
11from pathlib import Path
13from filelock import FileLock
14from filelock import Timeout as FileLockTimeout
16from lilbee._frozen import is_frozen
17from lilbee.runtime.progress import (
18 DetailedProgressCallback,
19 EventType,
20 SetupDoneEvent,
21 SetupProgressEvent,
22 SetupStartEvent,
23)
26class CrawlerBrowserError(RuntimeError):
27 """A browser-mode crawl cannot start its browser."""
30class ChromiumMissingError(CrawlerBrowserError):
31 """The headless Chromium shell a browser-mode crawl launches is not on disk."""
34class CrawlEngineRefusedError(RuntimeError):
35 """crawlberg refuses to build the engine for a crawl."""
38class CrawlerBackendError(RuntimeError):
39 """The ``crawler`` extra was never installed."""
42_CHROMIUM_COMPONENT = "chromium"
43# Playwright's name for the headless shell in browsers.json, and its cache directory prefix.
44_HEADLESS_SHELL_BROWSER = "chromium-headless-shell"
45_HEADLESS_SHELL_DIR_PREFIX = "chromium_headless_shell-"
46# The file Playwright writes into a browser directory once its install has finished.
47_INSTALL_COMPLETE_MARKER = "INSTALLATION_COMPLETE"
48# File names Playwright gives the headless shell executable on each platform.
49_HEADLESS_SHELL_NAMES = frozenset(
50 {"chrome-headless-shell", "chrome-headless-shell.exe", "headless_shell"}
51)
53# A Chromium unpack runs into the minutes on a slow link, so a caller queued
54# behind one waits well past any normal request timeout before giving up.
55_BOOTSTRAP_LOCK_TIMEOUT_S = 900.0
56# Rough size estimate for the Chromium download; Playwright bundles vary
57# slightly per platform but this gives the UI a decent denominator before
58# 'Total bytes' parses out of stdout.
59_CHROMIUM_ESTIMATE_MB = 180
60_CHROMIUM_SIZE_ESTIMATE_BYTES = _CHROMIUM_ESTIMATE_MB * 1024 * 1024
62# Unit -> bytes scale for Playwright stdout progress lines.
63_BYTE_UNIT_SCALE: dict[str, int] = {
64 "b": 1,
65 "kb": 1024,
66 "kib": 1024,
67 "mb": 1024 * 1024,
68 "mib": 1024 * 1024,
69}
71# Playwright 1.58 prints lines like
72# ``|■■■■■■■■ | 10% of 162.3 MiB`` during
73# the chromium download. The percent comes first, then "of <total> <unit>".
74_PROGRESS_LINE_RE = re.compile(
75 r"(\d+)\s*%\s*of\s*(\d+(?:\.\d+)?)\s*(MiB|Mb|MB|KiB|KB|B)",
76 re.IGNORECASE,
77)
80def _browsers_cache_path() -> Path:
81 """Return the root path where Playwright stores browser binaries."""
82 override = os.environ.get("PLAYWRIGHT_BROWSERS_PATH")
83 if override:
84 return Path(override).expanduser()
85 if sys.platform == "darwin":
86 return Path.home() / "Library" / "Caches" / "ms-playwright"
87 if sys.platform == "win32":
88 local = os.environ.get("LOCALAPPDATA", str(Path.home() / "AppData" / "Local"))
89 return Path(local) / "ms-playwright"
90 return Path.home() / ".cache" / "ms-playwright"
93def _read_chromium_revision(browsers_json: Path) -> str | None:
94 """Pull the headless shell revision out of a Playwright ``browsers.json``."""
95 import json
97 try:
98 data = json.loads(browsers_json.read_text(encoding="utf-8"))
99 except (OSError, ValueError):
100 return None
101 for browser in data.get("browsers", []):
102 if browser.get("name") == _HEADLESS_SHELL_BROWSER:
103 revision = browser.get("revision")
104 return str(revision) if revision is not None else None
105 return None
108def _expected_chromium_revision() -> str | None:
109 """Revision string Playwright was built against (e.g. ``"1217"``).
111 ``None`` means "unknown" -- treat as "do install" so a missing
112 browsers.json never short-circuits the bootstrap path.
113 """
114 try:
115 import playwright as _pw
116 except ImportError:
117 return None
118 for path in Path(_pw.__file__).parent.rglob("browsers.json"):
119 return _read_chromium_revision(path)
120 return None
123def _shell_in(directory: Path) -> Path | None:
124 """The headless shell path under *directory*, if Playwright finished installing it.
126 A path that cannot run counts, so a crawl reports it and installs nothing over it.
127 """
128 if not (directory / _INSTALL_COMPLETE_MARKER).is_file():
129 return None
130 matches = (path for path in sorted(directory.rglob("*")) if path.name in _HEADLESS_SHELL_NAMES)
131 return next(matches, None)
134def headless_shell_executable() -> Path | None:
135 """The headless shell binary at the revision Playwright expects, or None if absent.
137 With the revision unknown, any installed revision counts.
138 """
139 root = _browsers_cache_path()
140 if not root.is_dir():
141 return None
142 expected = _expected_chromium_revision()
143 if expected is None:
144 candidates = sorted(root.glob(f"{_HEADLESS_SHELL_DIR_PREFIX}*"))
145 else:
146 candidates = [root / f"{_HEADLESS_SHELL_DIR_PREFIX}{expected}"]
147 for directory in candidates:
148 executable = _shell_in(directory)
149 if executable is not None:
150 return executable
151 return None
154def chromium_installed() -> bool:
155 """Return True if the headless shell a browser-mode crawl launches is on disk."""
156 return headless_shell_executable() is not None
159def _bootstrap_lock_path() -> Path:
160 """Lock file serializing Chromium installs into the browsers cache."""
161 root = _browsers_cache_path()
162 root.mkdir(parents=True, exist_ok=True)
163 return root / ".lilbee-bootstrap.lock"
166def crawler_browsers_path() -> Path:
167 """Public accessor for the crawler browser cache root.
169 Used by the HTTP status endpoint to tell plugins where Chromium
170 lives. The underlying resolver stays private because callers should
171 not depend on the Playwright-specific directory layout.
172 """
173 return _browsers_cache_path()
176def _bytes_from_stdout(line: str) -> tuple[int, int] | None:
177 """Extract (downloaded_bytes, total_bytes) from a Playwright stdout line.
179 Matches the ``NN% of N.N MiB`` shape Playwright 1.58+ emits for the
180 Chromium download. Returns None when the line doesn't match. The
181 percent and total both parse out of the same line so callers never
182 have to handle a missing total.
183 """
184 match = _PROGRESS_LINE_RE.search(line)
185 if match is None:
186 return None
187 pct = int(match.group(1))
188 raw_total = float(match.group(2))
189 unit = match.group(3).lower()
190 scale = _BYTE_UNIT_SCALE.get(unit, 1)
191 total = int(raw_total * scale)
192 downloaded = int(total * pct / 100)
193 return downloaded, total
196def _emit_setup_start(on_progress: DetailedProgressCallback | None) -> None:
197 if on_progress is None:
198 return
199 on_progress(
200 EventType.SETUP_START,
201 SetupStartEvent(
202 component=_CHROMIUM_COMPONENT,
203 size_estimate_bytes=_CHROMIUM_SIZE_ESTIMATE_BYTES,
204 ),
205 )
208def _emit_setup_done(
209 on_progress: DetailedProgressCallback | None,
210 *,
211 success: bool,
212 error: str | None,
213) -> None:
214 if on_progress is None:
215 return
216 on_progress(
217 EventType.SETUP_DONE,
218 SetupDoneEvent(component=_CHROMIUM_COMPONENT, success=success, error=error),
219 )
222async def _drain_stdout_to_progress(
223 stream: asyncio.StreamReader,
224 on_progress: DetailedProgressCallback | None,
225) -> None:
226 while True:
227 line_bytes = await stream.readline()
228 if not line_bytes:
229 return
230 line = line_bytes.decode(errors="replace").rstrip()
231 parsed = _bytes_from_stdout(line)
232 if parsed is None or on_progress is None:
233 continue
234 downloaded, total = parsed
235 on_progress(
236 EventType.SETUP_PROGRESS,
237 SetupProgressEvent(
238 component=_CHROMIUM_COMPONENT,
239 downloaded_bytes=downloaded,
240 total_bytes=total,
241 detail=line,
242 ),
243 )
246async def _drain_stderr(stream: asyncio.StreamReader, tail: list[str]) -> None:
247 while True:
248 line_bytes = await stream.readline()
249 if not line_bytes:
250 return
251 tail.append(line_bytes.decode(errors="replace").rstrip())
254_PLAYWRIGHT_MISSING_HINT = (
255 "Chromium bootstrap requires the playwright Python package, which is "
256 "bundled with the release binary and the lilbee[crawler] extra. "
257 "Reinstall with 'pip install lilbee[crawler]' or download a fresh "
258 "release binary."
259)
262def _resolve_playwright_runner() -> tuple[list[str], dict[str, str]]:
263 """Return ``(argv_prefix, env)`` for invoking ``playwright install``.
265 Spawns Playwright's bundled Node driver directly so the call works under a
266 pip install, ``uv tool install``, or a frozen (Nuitka onefile) binary. Falls
267 back to ``[sys.executable, '-m', 'playwright']`` for unfrozen builds when the
268 driver lookup fails; re-raises for frozen builds, where ``sys.executable``
269 is the lilbee exe and ``-m playwright`` would leak into typer.
270 """
271 try:
272 from playwright._impl._driver import compute_driver_executable, get_driver_env
273 except ImportError as exc:
274 raise CrawlerBrowserError(_PLAYWRIGHT_MISSING_HINT) from exc
275 try:
276 driver_exe, driver_cli = compute_driver_executable()
277 except Exception:
278 if not is_frozen():
279 return [sys.executable, "-m", "playwright"], dict(os.environ)
280 raise
281 return [str(driver_exe), str(driver_cli)], dict(get_driver_env())
284def install_deps_command() -> str:
285 """The shell command that installs the system libraries Chromium needs.
287 Built from the same runner as ``playwright install``, so it runs on a frozen
288 build and on a tool install, where no ``playwright`` is on PATH. Raises
289 :class:`CrawlerBrowserError` when no runner is available.
290 """
291 argv, _env = _resolve_playwright_runner()
292 return shlex.join([*argv, "install-deps", "chromium"])
295async def bootstrap_chromium(
296 on_progress: DetailedProgressCallback | None = None,
297) -> None:
298 """Run ``playwright install --only-shell chromium`` as a subprocess, emitting events.
300 Short-circuits when ``chromium_installed()`` is already True. Emits
301 ``setup_start`` before spawning, ``setup_progress`` for each recognizable
302 progress line on stdout, and ``setup_done`` on exit (``success=False`` plus
303 the subprocess stderr tail on failure). Raises :class:`CrawlerBrowserError`
304 with the tail so task workers route to FAILED cleanly.
305 """
306 if chromium_installed():
307 _emit_setup_done(on_progress, success=True, error=None)
308 return
310 # Two callers (a second POST /setup/crawler, or a crawl bootstrapping on
311 # first use) would unpack a browser bundle into the same directory and
312 # corrupt it. On disk, not in memory, because the CLI and MCP are separate
313 # processes sharing this path. thread_local=False because the acquire runs
314 # in a worker thread and the release on the loop thread; with the default
315 # the release would not count and the lock would never be freed.
316 lock = FileLock(str(_bootstrap_lock_path()), thread_local=False)
317 try:
318 await asyncio.to_thread(lock.acquire, timeout=_BOOTSTRAP_LOCK_TIMEOUT_S)
319 except FileLockTimeout as exc:
320 message = (
321 "Timed out waiting for another Chromium install to finish. "
322 "If no other lilbee process is installing, retry."
323 )
324 _emit_setup_done(on_progress, success=False, error=message)
325 raise CrawlerBrowserError(message) from exc
326 try:
327 # The install we were queued behind may have been the one we needed.
328 if chromium_installed():
329 _emit_setup_done(on_progress, success=True, error=None)
330 return
331 await _install_chromium(on_progress)
332 finally:
333 lock.release()
336async def _install_chromium(on_progress: DetailedProgressCallback | None) -> None:
337 """Run the install subprocess. Caller holds the bootstrap lock."""
338 _emit_setup_start(on_progress)
340 try:
341 runner, runner_env = _resolve_playwright_runner()
342 except CrawlerBrowserError as exc:
343 _emit_setup_done(on_progress, success=False, error=str(exc))
344 raise
346 proc = await asyncio.create_subprocess_exec(
347 *runner,
348 "install",
349 "--only-shell",
350 "chromium",
351 stdout=asyncio.subprocess.PIPE,
352 stderr=asyncio.subprocess.PIPE,
353 env=runner_env,
354 )
355 # mypy narrowing: asyncio.create_subprocess_exec with PIPE guarantees
356 # non-None streams at runtime; the asserts only satisfy the type checker.
357 assert proc.stdout is not None # noqa: S101
358 assert proc.stderr is not None # noqa: S101
360 stderr_tail: list[str] = []
361 try:
362 await asyncio.gather(
363 _drain_stdout_to_progress(proc.stdout, on_progress),
364 _drain_stderr(proc.stderr, stderr_tail),
365 )
366 returncode = await proc.wait()
367 finally:
368 # On cancel (SSE client disconnect) or any error, don't leave the
369 # ~180MB chromium install running orphaned; terminate, then kill.
370 await _terminate_process(proc)
372 if returncode != 0:
373 tail = "\n".join(stderr_tail[-10:]) or f"exit code {returncode}"
374 _emit_setup_done(on_progress, success=False, error=tail)
375 raise CrawlerBrowserError(f"Chromium bootstrap failed (exit {returncode}): {tail}")
377 _emit_setup_done(on_progress, success=True, error=None)
380_TERMINATE_TIMEOUT_S = 5.0
383async def _terminate_process(proc: asyncio.subprocess.Process) -> None:
384 """Terminate then kill *proc* if it is still running, and reap it.
386 Best-effort and re-cancel-safe: the waits are shielded so a second
387 cancellation while unwinding cannot leave the child orphaned.
388 """
389 if proc.returncode is not None:
390 return
391 with contextlib.suppress(ProcessLookupError):
392 proc.terminate()
393 with contextlib.suppress(TimeoutError, asyncio.CancelledError):
394 await asyncio.wait_for(asyncio.shield(proc.wait()), timeout=_TERMINATE_TIMEOUT_S)
395 if proc.returncode is None:
396 with contextlib.suppress(ProcessLookupError):
397 proc.kill()
398 # SIGKILL is uncatchable so the reap should be near-instant; cap it anyway
399 # so a wedged reap can't hang the cleanup path indefinitely.
400 with contextlib.suppress(TimeoutError, asyncio.CancelledError):
401 await asyncio.wait_for(asyncio.shield(proc.wait()), timeout=_TERMINATE_TIMEOUT_S)