Coverage for src/lilbee/crawler/task.py: 100%

1"""Background crawl task management: start, track, and query crawl operations.""" 

2 

3import asyncio 

4import logging 

5import threading 

6import uuid 

7from dataclasses import dataclass, field 

8from datetime import UTC, datetime 

9from enum import StrEnum 

10 

11from lilbee.core.config.enums import CrawlRenderMode 

12from lilbee.crawler import crawl_and_save 

13from lilbee.runtime.progress import ( 

14 CrawlPageEvent, 

15 CrawlPageFailedEvent, 

16 DetailedProgressCallback, 

17 EventType, 

18 ProgressEvent, 

19) 

20 

21log = logging.getLogger(__name__) 

22 

23# Maximum completed tasks to retain in memory before evicting oldest. 

24_MAX_COMPLETED_TASKS = 100 

25 

26# Maximum per-page failure reasons stored on a task; pages_failed keeps the 

27# true count so a large all-failing crawl doesn't grow task memory unboundedly. 

28_MAX_FAILURE_REASONS = 20 

29 

30 

31class TaskStatus(StrEnum): 

32 """Lifecycle states for a crawl task.""" 

33 

34 PENDING = "pending" 

35 RUNNING = "running" 

36 DONE = "done" 

37 FAILED = "failed" 

38 CANCELLED = "cancelled" 

39 

40 

41@dataclass 

42class CrawlTask: 

43 """Tracks a single crawl operation. 

44 

45 depth / max_pages follow the crawl_and_save three-state convention: None = 

46 unbounded, 0 (depth only) = single URL, positive int = explicit cap. 

47 """ 

48 

49 task_id: str 

50 url: str 

51 depth: int | None 

52 max_pages: int | None 

53 render_mode: CrawlRenderMode | None = None 

54 include_subdomains: bool = False 

55 status: TaskStatus = TaskStatus.PENDING 

56 pages_crawled: int = 0 

57 pages_total: int | None = None 

58 pages_failed: int = 0 

59 failure_reasons: list[str] = field(default_factory=list) 

60 error: str | None = None 

61 started_at: str = "" 

62 finished_at: str = "" 

63 # Polled by crawl_and_save between pages and by the follow-on sync between 

64 # files, so a stop lands on a boundary with the pages already fetched saved. 

65 # Cancelling _async_task instead would abort mid-page and lose them. 

66 cancel: threading.Event = field(default_factory=threading.Event, repr=False) 

67 _async_task: asyncio.Task[None] | None = field(default=None, repr=False, init=False) 

68 

69 

70class TaskRegistry: 

71 """In-memory registry of active and completed crawl tasks. 

72 A single module-level instance (_registry) is used because task tracking 

73 is inherently per-process state (asyncio.Task references, etc.). 

74 """ 

75 

76 def __init__(self) -> None: 

77 self.tasks: dict[str, CrawlTask] = {} 

78 

79 def clear(self) -> None: 

80 """Remove all tasks from the registry.""" 

81 self.tasks.clear() 

82 

83 

84_registry = TaskRegistry() 

85 

86 

87def now_iso() -> str: 

88 """Current UTC time as ISO 8601 string.""" 

89 return datetime.now(UTC).isoformat() 

90 

91 

92def make_progress_updater(task: CrawlTask) -> DetailedProgressCallback: 

93 """Return a progress callback that updates task fields from crawl events.""" 

94 

95 def _on_progress(event_type: EventType, data: ProgressEvent) -> None: 

96 if event_type == EventType.CRAWL_PAGE: 

97 if not isinstance(data, CrawlPageEvent): 

98 raise TypeError(f"Expected CrawlPageEvent, got {type(data).__name__}") 

99 task.pages_crawled = data.current 

100 task.pages_total = data.total 

101 elif event_type == EventType.CRAWL_PAGE_FAILED: 

102 if not isinstance(data, CrawlPageFailedEvent): 

103 raise TypeError(f"Expected CrawlPageFailedEvent, got {type(data).__name__}") 

104 task.pages_failed += 1 

105 if len(task.failure_reasons) < _MAX_FAILURE_REASONS: 

106 task.failure_reasons.append(f"{data.url}: {data.reason}") 

107 

108 return _on_progress 

109 

110 

111async def run_crawl(task: CrawlTask) -> None: 

112 """Execute crawl, save results, and trigger sync.""" 

113 task.status = TaskStatus.RUNNING 

114 task.started_at = now_iso() 

115 progress = make_progress_updater(task) 

116 

117 try: 

118 paths = await crawl_and_save( 

119 task.url, 

120 depth=task.depth, 

121 max_pages=task.max_pages, 

122 on_progress=progress, 

123 cancel=task.cancel, 

124 include_subdomains=task.include_subdomains, 

125 render_mode=task.render_mode, 

126 ) 

127 if task.cancel.is_set(): 

128 task.status = TaskStatus.CANCELLED 

129 task.pages_crawled = task.pages_crawled or len(paths) 

130 task.finished_at = now_iso() 

131 log.info("Crawl cancelled: %s after %d files", task.url, len(paths)) 

132 return 

133 task.status = TaskStatus.DONE 

134 task.pages_crawled = task.pages_crawled or len(paths) 

135 task.finished_at = now_iso() 

136 log.info("Crawl complete: %s → %d files", task.url, len(paths)) 

137 try: 

138 from lilbee.data.ingest import sync 

139 

140 await sync(quiet=True, cancel=task.cancel) 

141 except Exception: 

142 log.warning("Post-crawl sync failed for %s", task.url, exc_info=True) 

143 except Exception as exc: 

144 task.status = TaskStatus.FAILED 

145 task.error = str(exc) 

146 task.finished_at = now_iso() 

147 log.warning("Crawl failed: %s: %s", task.url, exc) 

148 finally: 

149 task._async_task = None 

150 

151 

152def _evict_completed() -> None: 

153 """Remove completed tasks with the earliest ``finished_at`` when over cap. 

154 

155 Finish order diverges from start order whenever short tasks complete 

156 while a long one is still running, so sort on ``finished_at`` 

157 rather than dict insertion order. 

158 """ 

159 done_statuses = (TaskStatus.DONE, TaskStatus.FAILED, TaskStatus.CANCELLED) 

160 tasks = _registry.tasks 

161 completed = [(tid, t) for tid, t in tasks.items() if t.status in done_statuses] 

162 excess = len(completed) - _MAX_COMPLETED_TASKS 

163 if excess <= 0: 

164 return 

165 completed.sort(key=lambda pair: pair[1].finished_at) 

166 for tid, _ in completed[:excess]: 

167 del tasks[tid] 

168 

169 

170def start_crawl( 

171 url: str, 

172 depth: int | None = None, 

173 max_pages: int | None = None, 

174 render_mode: CrawlRenderMode | None = None, 

175 *, 

176 include_subdomains: bool = False, 

177) -> str: 

178 """Create a crawl task and launch it as an asyncio background task. 

179 

180 Defaults to whole-site unbounded recursion. Pass depth=0 for single URL. 

181 ``render_mode`` of ``None`` defers to ``cfg.crawl_render_mode``. 

182 ``include_subdomains`` widens whole-site scope to the host's subdomains. 

183 Returns the task_id for status polling. 

184 """ 

185 _evict_completed() 

186 task_id = uuid.uuid4().hex[:12] 

187 task = CrawlTask( 

188 task_id=task_id, 

189 url=url, 

190 depth=depth, 

191 max_pages=max_pages, 

192 render_mode=render_mode, 

193 include_subdomains=include_subdomains, 

194 ) 

195 _registry.tasks[task_id] = task 

196 task._async_task = asyncio.create_task(run_crawl(task)) 

197 return task_id 

198 

199 

200def get_task(task_id: str) -> CrawlTask | None: 

201 """Look up a crawl task by ID.""" 

202 return _registry.tasks.get(task_id) 

203 

204 

205def cancel_crawl(task_id: str) -> bool: 

206 """Ask a running crawl to stop, returning whether it was still running. 

207 

208 Cooperative: the crawl stops at the next page boundary and keeps what it 

209 already saved, so the pages fetched so far still reach the corpus. A task 

210 that already finished is left alone. 

211 """ 

212 task = _registry.tasks.get(task_id) 

213 if task is None or task.status not in (TaskStatus.PENDING, TaskStatus.RUNNING): 

214 return False 

215 task.cancel.set() 

216 return True 

217 

218 

219def list_tasks() -> list[CrawlTask]: 

220 """Return all tracked crawl tasks (active and completed).""" 

221 return list(_registry.tasks.values()) 

222 

223 

224def clear_tasks() -> None: 

225 """Remove all tasks from the registry (for testing).""" 

226 _registry.clear()