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.server.handlers.sse import SseStream 

11 

12 

13async def crawl_stream( 

14 url: str, 

15 depth: int | None = None, 

16 max_pages: int | None = None, 

17 render_mode: CrawlRenderMode | None = None, 

18 include_subdomains: bool = False, 

19) -> AsyncGenerator[str, None]: 

20 """Stream crawl progress as SSE events. 

21 

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

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

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

25 

26 A browser crawl that finds no Chromium installed inlines 

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

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

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

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

31 """ 

32 sse = SseStream() 

33 

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

35 from lilbee.crawler import crawl_and_save 

36 

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

38 # relaying setup_* events through the same on_progress callback 

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

40 try: 

41 return await crawl_and_save( 

42 url, 

43 depth=depth, 

44 max_pages=max_pages, 

45 on_progress=sse.callback, 

46 cancel=sse.cancel, 

47 include_subdomains=include_subdomains, 

48 render_mode=render_mode, 

49 ) 

50 finally: 

51 sse.queue.put_nowait(None) 

52 

53 task = asyncio.create_task(_run_crawl()) 

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

55 yield event 

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

57 if frame is not None: 

58 yield frame