Coverage for src/lilbee/crawler/task.py: 100%
106 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Background crawl task management: start, track, and query crawl operations."""
3import asyncio
4import logging
5import threading
6import uuid
7from dataclasses import dataclass, field
8from datetime import UTC, datetime
9from enum import StrEnum
11from lilbee.core.config.enums import CrawlRenderMode
12from lilbee.crawler import crawl_and_save
13from lilbee.runtime.progress import (
14 CrawlPageEvent,
15 DetailedProgressCallback,
16 EventType,
17 ProgressEvent,
18)
20log = logging.getLogger(__name__)
22# Maximum completed tasks to retain in memory before evicting oldest.
23_MAX_COMPLETED_TASKS = 100
26class TaskStatus(StrEnum):
27 """Lifecycle states for a crawl task."""
29 PENDING = "pending"
30 RUNNING = "running"
31 DONE = "done"
32 FAILED = "failed"
33 CANCELLED = "cancelled"
36@dataclass
37class CrawlTask:
38 """Tracks a single crawl operation.
40 depth / max_pages follow the crawl_and_save three-state convention: None =
41 unbounded, 0 (depth only) = single URL, positive int = explicit cap.
42 """
44 task_id: str
45 url: str
46 depth: int | None
47 max_pages: int | None
48 render_mode: CrawlRenderMode | None = None
49 include_subdomains: bool = False
50 status: TaskStatus = TaskStatus.PENDING
51 pages_crawled: int = 0
52 pages_total: int | None = None
53 error: str | None = None
54 started_at: str = ""
55 finished_at: str = ""
56 # Polled by crawl_and_save between pages and by the follow-on sync between
57 # files, so a stop lands on a boundary with the pages already fetched saved.
58 # Cancelling _async_task instead would abort mid-page and lose them.
59 cancel: threading.Event = field(default_factory=threading.Event, repr=False)
60 _async_task: asyncio.Task[None] | None = field(default=None, repr=False, init=False)
63class TaskRegistry:
64 """In-memory registry of active and completed crawl tasks.
65 A single module-level instance (_registry) is used because task tracking
66 is inherently per-process state (asyncio.Task references, etc.).
67 """
69 def __init__(self) -> None:
70 self.tasks: dict[str, CrawlTask] = {}
72 def clear(self) -> None:
73 """Remove all tasks from the registry."""
74 self.tasks.clear()
77_registry = TaskRegistry()
80def now_iso() -> str:
81 """Current UTC time as ISO 8601 string."""
82 return datetime.now(UTC).isoformat()
85def make_progress_updater(task: CrawlTask) -> DetailedProgressCallback:
86 """Return a progress callback that updates task fields from crawl events."""
88 def _on_progress(event_type: EventType, data: ProgressEvent) -> None:
89 if event_type == EventType.CRAWL_PAGE:
90 if not isinstance(data, CrawlPageEvent):
91 raise TypeError(f"Expected CrawlPageEvent, got {type(data).__name__}")
92 task.pages_crawled = data.current
93 task.pages_total = data.total
95 return _on_progress
98async def run_crawl(task: CrawlTask) -> None:
99 """Execute crawl, save results, and trigger sync."""
100 task.status = TaskStatus.RUNNING
101 task.started_at = now_iso()
102 progress = make_progress_updater(task)
104 try:
105 paths = await crawl_and_save(
106 task.url,
107 depth=task.depth,
108 max_pages=task.max_pages,
109 on_progress=progress,
110 cancel=task.cancel,
111 include_subdomains=task.include_subdomains,
112 render_mode=task.render_mode,
113 )
114 if task.cancel.is_set():
115 task.status = TaskStatus.CANCELLED
116 task.pages_crawled = task.pages_crawled or len(paths)
117 task.finished_at = now_iso()
118 log.info("Crawl cancelled: %s after %d files", task.url, len(paths))
119 return
120 task.status = TaskStatus.DONE
121 task.pages_crawled = task.pages_crawled or len(paths)
122 task.finished_at = now_iso()
123 log.info("Crawl complete: %s → %d files", task.url, len(paths))
124 try:
125 from lilbee.data.ingest import sync
127 await sync(quiet=True, cancel=task.cancel)
128 except Exception:
129 log.warning("Post-crawl sync failed for %s", task.url, exc_info=True)
130 except Exception as exc:
131 task.status = TaskStatus.FAILED
132 task.error = str(exc)
133 task.finished_at = now_iso()
134 log.warning("Crawl failed: %s: %s", task.url, exc)
135 finally:
136 task._async_task = None
139def _evict_completed() -> None:
140 """Remove completed tasks with the earliest ``finished_at`` when over cap.
142 Finish order diverges from start order whenever short tasks complete
143 while a long one is still running, so sort on ``finished_at``
144 rather than dict insertion order.
145 """
146 done_statuses = (TaskStatus.DONE, TaskStatus.FAILED, TaskStatus.CANCELLED)
147 tasks = _registry.tasks
148 completed = [(tid, t) for tid, t in tasks.items() if t.status in done_statuses]
149 excess = len(completed) - _MAX_COMPLETED_TASKS
150 if excess <= 0:
151 return
152 completed.sort(key=lambda pair: pair[1].finished_at)
153 for tid, _ in completed[:excess]:
154 del tasks[tid]
157def start_crawl(
158 url: str,
159 depth: int | None = None,
160 max_pages: int | None = None,
161 render_mode: CrawlRenderMode | None = None,
162 *,
163 include_subdomains: bool = False,
164) -> str:
165 """Create a crawl task and launch it as an asyncio background task.
167 Defaults to whole-site unbounded recursion. Pass depth=0 for single URL.
168 ``render_mode`` of ``None`` defers to ``cfg.crawl_render_mode``.
169 ``include_subdomains`` widens whole-site scope to the host's subdomains.
170 Returns the task_id for status polling.
171 """
172 _evict_completed()
173 task_id = uuid.uuid4().hex[:12]
174 task = CrawlTask(
175 task_id=task_id,
176 url=url,
177 depth=depth,
178 max_pages=max_pages,
179 render_mode=render_mode,
180 include_subdomains=include_subdomains,
181 )
182 _registry.tasks[task_id] = task
183 task._async_task = asyncio.create_task(run_crawl(task))
184 return task_id
187def get_task(task_id: str) -> CrawlTask | None:
188 """Look up a crawl task by ID."""
189 return _registry.tasks.get(task_id)
192def cancel_crawl(task_id: str) -> bool:
193 """Ask a running crawl to stop, returning whether it was still running.
195 Cooperative: the crawl stops at the next page boundary and keeps what it
196 already saved, so the pages fetched so far still reach the corpus. A task
197 that already finished is left alone.
198 """
199 task = _registry.tasks.get(task_id)
200 if task is None or task.status not in (TaskStatus.PENDING, TaskStatus.RUNNING):
201 return False
202 task.cancel.set()
203 return True
206def list_tasks() -> list[CrawlTask]:
207 """Return all tracked crawl tasks (active and completed)."""
208 return list(_registry.tasks.values())
211def clear_tasks() -> None:
212 """Remove all tasks from the registry (for testing)."""
213 _registry.clear()