Coverage for src/lilbee/cli/tui/widgets/task_bar_controller.py: 100%
212 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""TaskBarController and the per-task ProgressReporter."""
3from __future__ import annotations
5import logging
6import threading
7from collections.abc import Callable
8from enum import StrEnum
9from typing import TYPE_CHECKING, Any
11from textual.app import App
13from lilbee.catalog.formatting import download_task_name
14from lilbee.cli.tui import messages as msg
15from lilbee.cli.tui.task_queue import Task, TaskQueue, TaskStatus, TaskType
16from lilbee.cli.tui.thread_safe import call_from_thread
17from lilbee.crawler import bootstrap_chromium, chromium_installed
18from lilbee.runtime import asyncio_loop
19from lilbee.runtime.cancellation import TaskCancelledError
20from lilbee.runtime.progress import EventType, SetupProgressEvent
22if TYPE_CHECKING:
23 from lilbee.catalog import CatalogModel
25log = logging.getLogger(__name__)
27# Each download runs in its own child process (catalog.download_process), so
28# transfers neither share a xet session nor die with a cancelled sibling.
29_DOWNLOAD_CONCURRENCY = 4
30_BYTES_PER_MB = 1024 * 1024
33class TaskOutcome(StrEnum):
34 """How a task terminated. Passed from worker thread to finalizer."""
36 DONE = "done"
37 FAILED = "failed"
38 CANCELLED = "cancelled"
41class ProgressReporter:
42 """Thread-safe handle a worker uses to report progress and check cancellation.
44 The worker only sees this object; it never touches ``self.app``,
45 ``call_from_thread``, or any screen. Writes to the lock-protected
46 ``TaskQueue`` so updates survive any UI navigation.
47 """
49 def __init__(self, controller: TaskBarController, task_id: str) -> None:
50 self._controller = controller
51 self._task_id = task_id
53 @property
54 def task_id(self) -> str:
55 return self._task_id
57 def is_set(self) -> bool:
58 """True once the UI cancelled this task; the ``CancelSignal`` a download polls."""
59 task = self._controller.queue.get_task(self._task_id)
60 return task is not None and task.status is TaskStatus.CANCELLED
62 def check_cancelled(self) -> None:
63 """Raise ``TaskCancelledError`` if the task was cancelled from the UI."""
64 if self.is_set():
65 raise TaskCancelledError
67 def update(
68 self, progress: float, detail: str = "", *, indeterminate: bool | None = None
69 ) -> None:
70 """Write a progress snapshot to the shared queue.
72 Raises ``TaskCancelledError`` first if the UI cancelled the task, so
73 callers can use ``update`` as both a progress write and a cancel
74 checkpoint.
75 """
76 self.check_cancelled()
77 self._controller.queue.update_task(
78 self._task_id, progress, detail, indeterminate=indeterminate
79 )
82TaskTarget = Callable[[ProgressReporter], None]
85def _chromium_bootstrap_target(reporter: ProgressReporter) -> None:
86 """Worker target for the SETUP task: run bootstrap_chromium with progress forwarding.
88 Module-level so ``TaskBarController.ensure_chromium`` stays short and
89 tests can stub the target in isolation.
90 """
92 def _forward(event_type: EventType, data: Any) -> None:
93 if event_type != EventType.SETUP_PROGRESS:
94 return
95 if not isinstance(data, SetupProgressEvent):
96 return
97 total = data.total_bytes or 0
98 pct = int(data.downloaded_bytes * 100 / total) if total > 0 else 0
99 mb = data.downloaded_bytes // _BYTES_PER_MB
100 if total > 0:
101 detail = msg.SETUP_CHROMIUM_DETAIL.format(done=mb, total=total // _BYTES_PER_MB)
102 else:
103 detail = msg.SETUP_CHROMIUM_DETAIL_UNKNOWN.format(done=mb)
104 reporter.update(pct, detail)
106 asyncio_loop.run(bootstrap_chromium(on_progress=_forward))
109class TaskBarController:
110 """App-level owner of the shared TaskQueue + all long-running work.
112 The controller is attached as ``app.task_bar`` during
113 ``LilbeeApp.__init__``. All task lifecycle methods
114 (add/update/complete/fail/cancel) go through here so every ``TaskBar``
115 widget sees the same state, and every long-running op is spawned by
116 this controller: never by a screen that may dismiss mid-flight.
117 """
119 def __init__(self, app: App[Any]) -> None:
120 self.app = app
121 self.queue = TaskQueue(capacity={TaskType.DOWNLOAD.value: _DOWNLOAD_CONCURRENCY})
122 # task_id -> (target, on_success). Worker looks up its target here
123 # so we don't capture in a closure that outlives the task.
124 self._task_targets: dict[str, tuple[TaskTarget, Callable[[], None] | None]] = {}
125 # Number of files in documents/ that are out of date with the store.
126 # Set by start_detect_pending; read by TaskBar to render the
127 # "N docs to sync · S to sync" hint when no live tasks are running.
128 # Atomic int writes are safe under the GIL; the bar polls at 10 Hz.
129 self.pending_sync_count: int = 0
130 self._detect_thread: threading.Thread | None = None
131 # Roles whose worker is currently in the spawn window (1-3 s cold
132 # start). Surfaced as a single TaskBar hint instead of one toast
133 # per role so the chat screen isn't drowned in implementation
134 # detail on first prompt.
135 self.spawning_roles: set[str] = set()
137 def add_task(
138 self,
139 name: str,
140 task_type: str,
141 fn: Callable[[], None] | None = None,
142 *,
143 indeterminate: bool = False,
144 ) -> str:
145 """Enqueue a task. Returns the new task_id."""
146 return self.queue.enqueue(
147 fn or (lambda: None), name, task_type, indeterminate=indeterminate
148 )
150 def update_task(
151 self,
152 task_id: str,
153 progress: float,
154 detail: str = "",
155 *,
156 indeterminate: bool | None = None,
157 ) -> None:
158 """Update progress and detail text for a task."""
159 self.queue.update_task(task_id, progress, detail, indeterminate=indeterminate)
161 def complete_task(self, task_id: str) -> None:
162 """Mark a task done. Row lingers in history until the user clears it."""
163 task_type = self._task_type_of(task_id)
164 self.queue.complete_task(task_id)
165 self._after_done_hooks(task_type)
166 self._advance_all(task_type)
168 def fail_task(self, task_id: str, detail: str = "") -> None:
169 """Mark a task failed. Row lingers in history until the user clears it."""
170 self.queue.fail_task(task_id, detail)
171 self._advance_all(self._task_type_of(task_id))
173 def cancel_task(self, task_id: str) -> None:
174 """Mark a task cancelled; its worker notices at the next cancel checkpoint.
176 A download worker polls the row's status and terminates its child
177 process, which stops the bytes; other workers raise out of their next
178 progress update.
179 """
180 started = task_id in self._task_targets
181 task_type = self._task_type_of(task_id)
182 self.queue.cancel(task_id)
183 if not started:
184 # Rows put straight on the queue have no worker whose exit advances
185 # it, so the cancel must. A started row either has a worker that
186 # finalizes and advances when it exits, or is queued and holds no
187 # slot worth freeing.
188 self._advance_all(task_type)
190 def _after_done_hooks(self, task_type: str | None) -> None:
191 """Side effects triggered by a DONE completion.
193 Callable from both the direct ``complete_task`` convenience and
194 the worker-thread ``_finalize_task`` path so every success route
195 stays in sync. Does NOT advance the queue; each caller picks the
196 advance strategy that fits its context (``_advance_all`` vs
197 ``_try_start_next``).
198 """
199 if task_type == TaskType.DOWNLOAD.value:
200 self._notify_model_installed()
201 elif task_type in (TaskType.SYNC.value, TaskType.WIKI.value):
202 self.reload_wiki_screens()
204 def _task_type_of(self, task_id: str) -> str | None:
205 task = self.queue.get_task(task_id)
206 return task.task_type if task else None
208 def _advance_all(self, task_type: str | None) -> None:
209 """Start the freed type's next task, then any other idle type's.
211 Promotion has to spawn the worker. A task advanced without one sits
212 ACTIVE with no thread behind it, which renders as a live row that never
213 progresses.
214 """
215 if task_type:
216 self._try_start_next(task_type)
217 while (task := self.queue.advance()) is not None:
218 self._spawn_task_worker(task.task_id)
220 def downloading_label_for(self, ref: str) -> str | None:
221 """Return the task name if *ref*'s download is queued or active, else None.
223 ``ref`` is a model reference (catalog repo id or native GGUF
224 ref); the helper maps it to the canonical
225 :attr:`CatalogModel.display_name` and matches against
226 in-flight DOWNLOAD tasks. The returned label is suitable for
227 embedding in a user-facing toast.
228 """
229 label = download_task_name(ref)
230 if not label:
231 return None
232 for task in self.queue.active_tasks + self.queue.queued_tasks:
233 if task.task_type == TaskType.DOWNLOAD.value and task.name == label:
234 return task.name
235 return None
237 def set_pending_sync(self, count: int) -> None:
238 """Update the pending-sync count surfaced in the TaskBar hint."""
239 self.pending_sync_count = max(count, 0)
241 def clear_pending_sync(self) -> None:
242 """Drop the pending hint. Called when sync starts so the bar shows live progress instead."""
243 self.pending_sync_count = 0
245 def mark_role_spawning(self, role: str) -> None:
246 """Add *role* to the set of workers whose pool process is starting."""
247 self.spawning_roles.add(role)
249 def mark_role_spawned(self, role: str) -> None:
250 """Drop *role* from the spawn-in-progress set; harmless if already absent."""
251 self.spawning_roles.discard(role)
253 def start_detect_pending(self) -> None:
254 """Run the cheap sync-detection (filesystem walk + hash compare) on a daemon thread.
256 Writes the result via ``set_pending_sync``. No-op if a detect job
257 is already running. Errors are logged and silently swallowed: a
258 failed detect just leaves the previous count in place rather
259 than blocking the UI.
260 """
261 if self._detect_thread is not None and self._detect_thread.is_alive():
262 return
263 thread = threading.Thread(
264 target=self._run_detect_pending, daemon=True, name="detect-pending"
265 )
266 self._detect_thread = thread
267 thread.start()
269 def _run_detect_pending(self) -> None:
270 # Local import: lilbee.data.ingest pulls in lancedb + the embedder
271 # transitively; the TUI shouldn't pay for that just to import the
272 # task bar widget.
273 from lilbee.data.ingest import detect_pending
275 try:
276 count = detect_pending()
277 except Exception:
278 log.warning("detect_pending failed", exc_info=True)
279 return
280 self.set_pending_sync(count)
282 def ensure_chromium(self, on_ready: Callable[[], None]) -> None:
283 """Kick off a Chromium bootstrap if missing, then call ``on_ready``.
285 If Chromium is already installed, ``on_ready`` runs immediately on
286 the caller's thread. Otherwise a single SETUP task is enqueued
287 that runs ``bootstrap_chromium``; on success the controller
288 invokes ``on_ready`` on the worker thread via the task's
289 ``on_success`` hook. On failure the SETUP task surfaces as FAILED
290 and ``on_ready`` is NOT called (the follow-up work shouldn't
291 proceed against a missing browser).
293 bb-wq8g: the on_ready hook is how callers like ``_do_crawl`` chain
294 their real work behind the one-time bootstrap.
295 """
296 if chromium_installed():
297 on_ready()
298 return
300 self.start_task(
301 msg.SETUP_CHROMIUM_NAME,
302 TaskType.SETUP,
303 _chromium_bootstrap_target,
304 indeterminate=False,
305 on_success=on_ready,
306 )
308 def start_task(
309 self,
310 name: str,
311 task_type: TaskType,
312 target: TaskTarget,
313 *,
314 indeterminate: bool = False,
315 on_success: Callable[[], None] | None = None,
316 dedupe_key: str | None = None,
317 ) -> str:
318 """Enqueue a task, spawn its worker, return task_id.
320 The *target* receives a ``ProgressReporter`` as its only argument.
321 It should periodically call ``reporter.update(percent, detail)`` and
322 may call ``reporter.check_cancelled()`` to cooperatively abort.
324 On success (target returns normally) the queue marks the task DONE
325 and ``on_success`` (if provided) runs after on the same worker
326 thread. On ``TaskCancelledError`` the task is marked CANCELLED. On any
327 other exception the task is marked FAILED with ``str(exc)`` as
328 detail. Rows linger in the Task Center under their final status
329 until the user presses capital ``C`` to clear; the bottom bar
330 flashes the outcome once and then hides when idle.
332 Per-type capacity in ``TaskQueue`` (1 for every type) controls
333 concurrency: a second task of the same type queues behind the first.
334 """
335 task_id = self.queue.enqueue(
336 lambda: None,
337 name,
338 task_type.value,
339 indeterminate=indeterminate,
340 dedupe_key=dedupe_key,
341 )
342 if task_id in self._task_targets:
343 # Deduplicated: the live task keeps its original target.
344 return task_id
345 self._task_targets[task_id] = (target, on_success)
346 self._try_start_next(task_type.value)
347 return task_id
349 def _try_start_next(self, task_type: str) -> None:
350 """Promote queued tasks of this type into any free capacity slots."""
351 while (task := self.queue.advance(task_type)) is not None:
352 self._spawn_task_worker(task.task_id)
354 def _spawn_task_worker(self, task_id: str) -> None:
355 """Start a daemon thread for the task. Safe to call from any thread."""
356 if task_id not in self._task_targets:
357 return
358 thread = threading.Thread(
359 target=self._run_task_worker,
360 args=(task_id,),
361 daemon=True,
362 name=f"task-{task_id}",
363 )
364 thread.start()
366 def _run_task_worker(self, task_id: str) -> None:
367 """Body of the daemon worker thread."""
368 entry = self._task_targets.get(task_id)
369 if entry is None:
370 return
371 target, on_success = entry
372 task = self.queue.get_task(task_id)
373 task_type = task.task_type if task is not None else None
374 reporter = ProgressReporter(self, task_id)
375 try:
376 target(reporter)
377 except TaskCancelledError:
378 log.info("Task %s cancelled", task_id)
379 self._post_finalize(task_id, TaskOutcome.CANCELLED, "", task_type)
380 except Exception as exc:
381 log.warning("Task %s failed: %s", task_id, exc)
382 self._post_finalize(task_id, TaskOutcome.FAILED, str(exc), task_type)
383 else:
384 self._post_finalize(task_id, TaskOutcome.DONE, "", task_type)
385 if on_success is not None:
386 try:
387 on_success()
388 except Exception:
389 log.warning("on_success for %s raised", task_id, exc_info=True)
390 finally:
391 self._task_targets.pop(task_id, None)
393 def _post_finalize(
394 self, task_id: str, outcome: TaskOutcome, detail: str, task_type: str | None
395 ) -> None:
396 """Marshal finalization back to the main thread.
398 Main-thread execution matters because ``set_timer`` (used for the
399 flash-then-remove cycle) isn't safe from workers. ``call_from_thread``
400 targets ``self.app``: the App is long-lived; screens are not.
401 """
402 call_from_thread(self.app, self._finalize_task, task_id, outcome, detail, task_type)
404 def _finalize_task(
405 self, task_id: str, outcome: TaskOutcome, detail: str, task_type: str | None
406 ) -> None:
407 """Mark the queue state, refresh dependents, promote next queued task.
409 Runs on the main thread. Atomically: free the active slot, notify
410 anything downstream that needs a repaint (e.g. model dropdowns
411 after a download lands), and advance the queue. Rows stay in
412 history; the bottom bar flash expires on its own. Users clear
413 finished rows from the Task Center manually.
414 """
415 if outcome is TaskOutcome.DONE:
416 self.queue.complete_task(task_id)
417 self._after_done_hooks(task_type)
418 elif outcome is TaskOutcome.FAILED:
419 self.queue.fail_task(task_id, detail)
420 elif outcome is TaskOutcome.CANCELLED:
421 self.queue.cancel(task_id)
422 if task_type:
423 self._try_start_next(task_type)
425 def _notify_model_installed(self) -> None:
426 """Refresh any ChatScreen's ModelBar so the new model is selectable.
428 The dropdowns are built once on mount from the registry; without
429 this nudge, a freshly-downloaded model only appears after the
430 user reopens the screen. A bar that is not mounted yet is no longer
431 this caller's problem: refresh_model_bar is a no-op until the bar
432 exists, and the bar scans on its own mount.
433 """
434 # Late import to avoid a circular (ChatScreen imports this module).
435 from lilbee.cli.tui.screens.chat import ChatScreen
437 for screen in self.app.screen_stack:
438 # screen_stack is typed Screen[Any]; narrow at runtime to
439 # locate the one screen that owns the ModelBar.
440 if isinstance(screen, ChatScreen):
441 screen.refresh_model_bar()
442 break
444 def reload_wiki_screens(self) -> None:
445 """Rescan open wiki and drafts screens after a task rewrote pages or drafts on disk."""
446 from textual.css.query import QueryError
448 from lilbee.cli.tui.screens.wiki import WikiScreen
449 from lilbee.cli.tui.screens.wiki_drafts import WikiDraftsScreen
451 for screen in self.app.screen_stack:
452 # screen_stack is typed Screen[Any]; narrow at runtime to the
453 # wiki views that expose a reload.
454 if isinstance(screen, WikiScreen | WikiDraftsScreen):
455 try:
456 screen.reload()
457 except QueryError:
458 log.debug("Wiki screen not mounted yet; skipping reload", exc_info=True)
460 def start_download(
461 self,
462 model: CatalogModel,
463 *,
464 allow_unsupported: bool = False,
465 on_success: Callable[[], None] | None = None,
466 ) -> str:
467 """Enqueue a download; ``on_success`` runs on the worker thread once the file is on disk.
469 Re-requesting a download that is already queued or running returns that
470 task's id and starts nothing.
471 """
472 return self.start_task(
473 model.display_name,
474 TaskType.DOWNLOAD,
475 lambda reporter: _download_target(reporter, model, allow_unsupported=allow_unsupported),
476 on_success=on_success,
477 dedupe_key=download_key(model),
478 )
480 def pending_download(self, model: CatalogModel) -> Task | None:
481 """Return the queued-or-active download for *model*, if there is one."""
482 return self.queue.find_pending(TaskType.DOWNLOAD.value, download_key(model))
485def download_key(model: CatalogModel) -> str:
486 """Identity of a pull: repo plus filename, since one repo ships several quants."""
487 return f"{model.hf_repo}::{model.gguf_filename}"
490def _download_target(
491 reporter: ProgressReporter, model: CatalogModel, *, allow_unsupported: bool = False
492) -> None:
493 """``start_task`` target for a HuggingFace model download.
495 Translates ``PermissionError`` into the gated-repo friendly message and
496 ``UnsupportedArchError`` into a user-facing arch-mismatch message so
497 every call site gets consistent error UX.
498 """
499 from lilbee.app.models import pull_model_data
500 from lilbee.catalog import DownloadProgress
501 from lilbee.catalog.compat import UnsupportedArchError
502 from lilbee.catalog.types import ModelSource
504 def _on_progress(p: DownloadProgress) -> None:
505 reporter.update(p.percent, f"{model.display_name}: {p.detail}")
507 try:
508 pull_model_data(
509 model.ref,
510 ModelSource.NATIVE,
511 on_update=_on_progress,
512 allow_unsupported=allow_unsupported,
513 cancel=reporter,
514 )
515 except PermissionError as exc:
516 raise RuntimeError(msg.CATALOG_GATED_REPO.format(name=model.display_name)) from exc
517 except UnsupportedArchError as exc:
518 raise RuntimeError(
519 f"Architecture {exc.architecture!r} not supported by this lilbee build."
520 ) from exc