Coverage for src/lilbee/server/handlers/crawl.py: 100%

1"""Crawl streaming handler.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from collections.abc import AsyncGenerator 

7from pathlib import Path 

8 

9from lilbee.core.config.enums import CrawlRenderMode 

10from lilbee.crawler.models import NO_LOGIN, CrawlLogin 

11from lilbee.server.handlers.sse import SseStream 

12 

13 

14async def crawl_stream( 

15 url: str, 

16 depth: int | None = None, 

17 max_pages: int | None = None, 

18 render_mode: CrawlRenderMode | None = None, 

19 include_subdomains: bool = False, 

20 login: CrawlLogin = NO_LOGIN, 

21) -> AsyncGenerator[str, None]: 

22 """Stream crawl progress as SSE events. 

23 

24 Emits crawl_start, crawl_page, crawl_done events, then a final done event 

25 with the list of files written. On error emits crawl_error. 

26 Sets a cancel event on client disconnect so the crawl stops between pages. 

27 

28 A browser crawl that finds no Chromium installed inlines 

29 setup_start/progress/done events before the crawl begins, so a consumer can 

30 render a matching 'setup' progress indicator. These are not part of every 

31 stream: an http crawl never launches a browser, and that is the default 

32 render mode, so a client must not block waiting for a setup phase. 

33 """ 

34 sse = SseStream() 

35 

36 async def _run_crawl() -> list[Path]: 

37 from lilbee.crawler import crawl_and_save 

38 

39 # crawl_and_save runs the Chromium bootstrap itself on first use, 

40 # relaying setup_* events through the same on_progress callback 

41 # so the SSE stream carries them before any crawl_* events. 

42 try: 

43 return await crawl_and_save( 

44 url, 

45 depth=depth, 

46 max_pages=max_pages, 

47 on_progress=sse.callback, 

48 cancel=sse.cancel, 

49 include_subdomains=include_subdomains, 

50 render_mode=render_mode, 

51 login=login, 

52 ) 

53 finally: 

54 sse.queue.put_nowait(None) 

55 

56 task = asyncio.create_task(_run_crawl()) 

57 async for event in sse.drain(task, "Crawl stream"): 

58 yield event 

59 frame = sse.terminal_frame(task, lambda paths: {"files_written": [str(p) for p in paths]}) 

60 if frame is not None: 

61 yield frame