Coverage for src/lilbee/app/services.py: 100%
214 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Typed service container: single point of access for all singletons.
3All runtime dependencies (provider, store, embedder, reranker, concepts,
4clusterer, searcher, worker pool) are created lazily on first call to
5``get_services()`` and cached for the process lifetime. Tests call
6``reset_services()`` between runs.
8``build_services(config)`` is the construction seam: it builds a full container
9against an arbitrary Config without touching the process-global singleton. The
10library API (:class:`lilbee.Lilbee`) builds one per instance and installs it for
11the duration of each call via :func:`services_scope`, so ingest code reaching for
12``get_services()`` resolves the caller's container. The override is a ContextVar,
13so ``reset_services`` / ``set_services`` / ``peek_services`` (which operate only
14on the global singleton) never see it.
15"""
17from __future__ import annotations
19import asyncio
20import atexit
21import logging
22import os
23import signal
24import sys
25import threading
26import time
27from collections.abc import Callable
28from contextlib import contextmanager
29from contextvars import ContextVar
30from dataclasses import dataclass, field
31from typing import TYPE_CHECKING
33if TYPE_CHECKING:
34 from collections.abc import Iterator
36 from lilbee.catalog.hf_client import HfClient
37 from lilbee.core.config import Config
38 from lilbee.data.store import Store
39 from lilbee.modelhub.model_manager import ModelManager
40 from lilbee.modelhub.model_manager.discovery import KnownModelCache
41 from lilbee.modelhub.registry import ModelRegistry
42 from lilbee.providers.base import LLMProvider
43 from lilbee.providers.roles import WorkerRole
44 from lilbee.retrieval.clustering import Clusterer
45 from lilbee.retrieval.concepts import ConceptGraph
46 from lilbee.retrieval.embedder import Embedder
47 from lilbee.retrieval.query import Searcher
48 from lilbee.retrieval.reranker import Reranker
49 from lilbee.runtime.ingest_lock import IngestLockRegistry
50 from lilbee.sessions import SessionStore
53log = logging.getLogger(__name__)
55_SIGNAL_EXIT_BASE = 128
57_HARD_EXIT_THREAD_NAME = "hard-exit-teardown"
60def _default_session_store() -> SessionStore:
61 """Build the file-backed session store, importing it lazily.
63 ``lilbee.sessions`` pulls in the config/catalog import chain, so importing it
64 at this module's top would form a cycle during CLI config load.
65 """
66 from lilbee.sessions import SessionStore
68 return SessionStore()
71@dataclass
72class CrawlerSyncState:
73 """Process-wide sync coordination state (lock + last-run timestamp)."""
75 lock: threading.Lock = field(default_factory=threading.Lock)
76 last_run: float = 0.0
79@dataclass(frozen=True)
80class Services:
81 """Holds all runtime service instances.
83 Inference lifecycle (cancel, per-role reload, spawn notifications) is owned
84 by the provider, which manages the llama-server fleet. Services exposes thin
85 pass-throughs so callers (Ctrl+C, the chat-stream cancel action, the settings
86 and model-bar pickers, the TUI task bar) need not reach into the provider's
87 API. ``cancel_inference()`` is the canonical cancel entry point.
88 """
90 provider: LLMProvider
91 store: Store
92 embedder: Embedder
93 reranker: Reranker
94 concepts: ConceptGraph
95 clusterer: Clusterer
96 searcher: Searcher
97 registry: ModelRegistry
98 hf_client: HfClient
99 ingest_lock_registry: IngestLockRegistry
100 model_manager: ModelManager
101 crawler_semaphore: asyncio.Semaphore | None
102 crawler_sync_state: CrawlerSyncState
103 known_models: KnownModelCache
104 session_store: SessionStore = field(default_factory=_default_session_store)
106 def cancel_inference(self) -> None:
107 """Interrupt any in-flight generation. Idempotent.
109 The fleet engine severs its live chat streams (llama-server stops
110 generating when the connection drops); providers with nothing in
111 flight treat this as a no-op.
112 """
113 self.provider.cancel_inference()
115 def reload_role(self, role_name: WorkerRole, *, wait: bool = False) -> None:
116 """Respawn only *role_name*'s model server so it picks up changed cfg.
118 Other roles' servers and any in-flight stream they own are untouched. Use
119 when one role-bound model setting changed (e.g. embedding_model). The
120 respawn runs off the caller's thread, so this returns immediately, unless
121 ``wait=True`` (the caller is already off the event loop and wants to block
122 until the new model has loaded).
123 """
124 self.provider.reload_role(role_name, wait=wait)
126 def add_pool_listener(
127 self,
128 *,
129 on_spawning: Callable[[WorkerRole], None] | None = None,
130 on_spawned: Callable[[WorkerRole], None] | None = None,
131 ) -> None:
132 """Subscribe to server spawn lifecycle events.
134 Forwards to :meth:`LLMProvider.add_spawn_listener`. The TUI uses this to
135 surface "Starting <role>..." / "<role> ready" notifications when a role's
136 server (re)spawns (cold start after a non-eager boot, or a reload).
137 """
138 self.provider.add_spawn_listener(on_spawning=on_spawning, on_spawned=on_spawned)
141class _ServicesState:
142 """The cached process-global singleton plus the per-task scoped override.
144 ``singleton`` is set on first ``get_services()`` call. Concurrency
145 contract: creation is serialized by ``_singleton_create_lock`` (several
146 worker threads can first-touch services at once, and a duplicate build
147 would collide in xberg's process-global backend registry), and the
148 Services dataclass is logically immutable post-construction, so concurrent
149 reads are safe without a lock. Tests that need a custom container call
150 ``set_services(make_mock_services(...))``; ``peek_services()`` is the
151 read-only inspector for cleanup fixtures.
153 ``override`` shadows the singleton for the entering task: set by
154 :func:`services_scope` (the library API's per-call binding), read by
155 :func:`get_services`, and invisible to ``reset_services`` /
156 ``set_services`` / ``peek_services``, which only touch the singleton.
157 """
159 def __init__(self) -> None:
160 self.singleton: Services | None = None
161 self.override: ContextVar[Services | None] = ContextVar(
162 "lilbee_services_override", default=None
163 )
164 # Whether this process is an interactive session (the TUI). Recorded by
165 # the interactive entry point before anything builds the container, and
166 # read once at build so the provider it creates holds its fleet resident
167 # for the session. Build-time intent only; the state that matters after
168 # that lives on the provider itself.
169 self.interactive: bool = False
172_state = _ServicesState()
175def build_services(
176 config: Config,
177 *,
178 provider: LLMProvider | None = None,
179 registry: ModelRegistry | None = None,
180 interactive: bool = False,
181) -> Services:
182 """Build a full Services container bound to *config*, without caching it.
184 ``get_services()`` calls this with the process-global cfg; the library API
185 calls it per instance. Service modules are imported inside the function to
186 keep CLI startup fast (they transitively pull in lancedb / xberg). Pass
187 *provider* to reuse a caller-supplied one; otherwise it is built from
188 *config* via the provider factory. Pass *registry* to reuse one already built
189 (get_services builds it for embedding-dim reconciliation). Embedding-dim
190 reconciliation is a global-cfg concern owned by :func:`get_services`, not
191 done here.
193 Side effect: binds *provider* into xberg's process-global OCR/embedding/
194 tokenizer backends. The registry is a single global slot, so every container
195 (singleton or per-instance library) binds its own, not only get_services.
196 """
197 from lilbee.catalog.hf_client import HfClient
198 from lilbee.data.store import Store
199 from lilbee.modelhub.model_manager import ModelManager
200 from lilbee.modelhub.model_manager.discovery import KnownModelCache
201 from lilbee.modelhub.registry import ModelRegistry
202 from lilbee.providers.factory import create_provider
203 from lilbee.retrieval.clustering import Clusterer
204 from lilbee.retrieval.concepts import ConceptGraph
205 from lilbee.retrieval.embedder import Embedder
206 from lilbee.retrieval.query import Searcher
207 from lilbee.retrieval.reranker import Reranker
208 from lilbee.runtime.ingest_lock import IngestLockRegistry
210 provider = provider or create_provider(config, hold_warm=interactive)
211 from lilbee.data.extract.backends import sync_xberg_backends
213 sync_xberg_backends(provider)
214 registry = registry or ModelRegistry(config.models_dir)
215 store = Store(config)
216 embedder = Embedder(config, provider)
217 reranker = Reranker(config)
218 concepts = ConceptGraph(config, store)
219 clusterer = Clusterer(config, store)
220 searcher = Searcher(config, provider, store, embedder, reranker, concepts)
221 hf_client = HfClient()
222 ingest_lock_registry = IngestLockRegistry()
223 model_manager = ModelManager(config.models_dir)
224 crawler_semaphore = (
225 asyncio.Semaphore(config.crawl_max_concurrent) if config.crawl_max_concurrent > 0 else None
226 )
227 crawler_sync_state = CrawlerSyncState()
228 known_models = KnownModelCache()
229 return Services(
230 provider=provider,
231 store=store,
232 embedder=embedder,
233 reranker=reranker,
234 concepts=concepts,
235 clusterer=clusterer,
236 searcher=searcher,
237 registry=registry,
238 hf_client=hf_client,
239 ingest_lock_registry=ingest_lock_registry,
240 model_manager=model_manager,
241 crawler_semaphore=crawler_semaphore,
242 crawler_sync_state=crawler_sync_state,
243 known_models=known_models,
244 )
247# Serializes first-touch singleton creation: several worker threads (e.g.
248# concurrent downloads) can call get_services() before the singleton exists,
249# and a duplicate build re-registers xberg's process-global backends mid-flight,
250# raising 'already registered' in the losing thread.
251_singleton_create_lock = threading.Lock()
254def get_services() -> Services:
255 """Return the active container: a scoped override if set, else the cached singleton.
257 Creates the singleton on first call (against the process-global cfg). A
258 config-file embedding_model with no embedding_dim would otherwise build the
259 store at the stale 768 default, so the width is pinned to the embedder before
260 the store is built.
261 """
262 override = _state.override.get()
263 if override is not None:
264 return override
265 if _state.singleton is not None:
266 return _state.singleton
268 with _singleton_create_lock:
269 if _state.singleton is not None:
270 return _state.singleton
272 from lilbee.app.settings import reconcile_embedding_dim
273 from lilbee.core.config import cfg
274 from lilbee.modelhub.registry import ModelRegistry
276 registry = ModelRegistry(cfg.models_dir)
277 # Pin the store width to the embedder before Store(); pass the registry so
278 # resolution doesn't re-enter this half-built get_services.
279 reconcile_embedding_dim(registry)
280 _state.singleton = build_services(cfg, registry=registry, interactive=_state.interactive)
281 # Eager start is the default: pay the spawn cost per role server at TUI mount
282 # so the first user action lands on a warm fleet. Roles whose model is unset
283 # are skipped, so a setup with only chat + embed never spawns rerank or
284 # vision. Set ``cfg.worker_pool_eager_start = false`` for headless scripts
285 # where mount time matters more than first-call latency.
286 if cfg.worker_pool_eager_start:
287 from contextlib import suppress
289 with suppress(Exception):
290 _state.singleton.provider.warm_up_pool()
291 return _state.singleton
294@contextmanager
295def services_scope(services: Services) -> Iterator[None]:
296 """Bind *services* as the container ``get_services()`` returns for this block.
298 Isolated to the entering task via a ContextVar (it propagates into
299 ``to_ingest_thread`` workers), and never affects the global singleton, so
300 ``reset_services`` is unnecessary and unused around a scoped call.
301 """
302 token = _state.override.set(services)
303 try:
304 yield
305 finally:
306 _state.override.reset(token)
309def mark_interactive_session() -> None:
310 """Record that this process is an interactive session before services build.
312 The TUI owns the process for its whole lifetime, so an engine bound to it
313 keeps its weights resident rather than idle-unloading under a user who is
314 still in the app; closing lilbee releases it. A keep-warm engine outlives the
315 session and keeps its idle window instead. Called by the interactive entry point
316 before anything touches ``get_services``, so the provider is constructed with
317 that intent; a one-shot CLI or the MCP server never calls it.
318 """
319 _state.interactive = True
322def set_services(services: Services | None) -> None:
323 """Replace the cached Services singleton (for testing)."""
324 _state.singleton = services
327def peek_services() -> Services | None:
328 """Return the cached Services container, or None if not yet initialized.
330 Public read-only accessor for test cleanup helpers that need to
331 inspect the singleton without forcing initialization.
332 """
333 return _state.singleton
336# Serializes the singleton swap in reset_services: a signal's teardown thread
337# and an exiting caller's can race, and both tearing down the same container
338# would double-close the store.
339_reset_swap_lock = threading.Lock()
342def reset_services() -> None:
343 """Shut down and discard all cached instances.
345 Swap the module reference to ``None`` *before* tearing the old instances
346 down, so a new caller never observes a half-closed container. The swap is
347 locked so concurrent callers (a signal's teardown thread plus an exiting
348 one) tear the container down exactly once. On the shared HTTP daemon every
349 entry point that would call this mid-flight is refused, so it only ever
350 runs single-client (CLI, TUI, stdio MCP).
351 """
352 with _reset_swap_lock:
353 old = _state.singleton
354 _state.singleton = None
355 if old is not None:
356 old.provider.shutdown()
357 old.store.close()
360def reset_store() -> None:
361 """Close and rebuild only the Store and its dependents; keep providers loaded.
363 Used after a data-dir wipe (``/reset``) where the LanceDB handle is invalid
364 but the running provider/embedder/reranker are still good. Avoids the
365 multi-second reload cost of ``reset_services()``.
366 """
367 svc = _state.singleton
368 if svc is None:
369 return
370 from dataclasses import replace
372 from lilbee.core.config import cfg
373 from lilbee.data.store import Store
374 from lilbee.retrieval.clustering import Clusterer
375 from lilbee.retrieval.concepts import ConceptGraph
376 from lilbee.retrieval.query import Searcher
378 # Build the replacement, swap the reference, then close the old store last so
379 # a new caller never observes a closed handle mid-swap.
380 old_store = svc.store
381 store = Store(cfg)
382 concepts = ConceptGraph(cfg, store)
383 clusterer = Clusterer(cfg, store)
384 searcher = Searcher(cfg, svc.provider, store, svc.embedder, svc.reranker, concepts)
385 _state.singleton = replace(
386 svc,
387 store=store,
388 concepts=concepts,
389 clusterer=clusterer,
390 searcher=searcher,
391 )
392 old_store.close()
395class _EngineLifecycle:
396 """Owns the hard-exit hooks that stop the engine fleet."""
398 def __init__(self) -> None:
399 self._installed = False
401 @staticmethod
402 def _hard_exit_signals() -> tuple[signal.Signals, ...]: # pragma: no cover - platform split
403 """Signals whose default disposition kills us without running atexit."""
404 if sys.platform == "win32": # Windows has no SIGHUP
405 return (signal.SIGTERM,)
406 return (signal.SIGTERM, signal.SIGHUP)
408 def install(self) -> None:
409 """Route hard-exit signals through teardown. Idempotent; no-op off the main thread."""
410 if self._installed:
411 return
412 try:
413 for sig in self._hard_exit_signals():
414 signal.signal(sig, self._on_hard_exit)
415 except ValueError:
416 return
417 self._installed = True
419 def reset(self) -> None:
420 """Forget that handlers were installed."""
421 self._installed = False
423 def _on_hard_exit(self, signum: int, frame: object) -> None:
424 """Stop the fleet on its own thread, then exit with the signal status.
426 Signal handlers all run on the main thread, and a second signal can
427 interrupt this one mid-teardown: the kernel pairs SIGCONT with SIGHUP
428 for an orphaned process group, and Textual's SIGCONT handler raises
429 once the event loop is gone, which aborted the reap half-done and
430 orphaned a loaded fleet. A dedicated non-daemon thread cannot be
431 interrupted by signals, and the interpreter waits for it even as the
432 SystemExit unwinds the main thread.
433 """
434 del frame
435 threading.Thread(
436 target=_teardown_for_signal, args=(signum,), name=_HARD_EXIT_THREAD_NAME
437 ).start()
438 raise SystemExit(_SIGNAL_EXIT_BASE + signum)
441_FORCE_QUIT_EXIT_CODE = 130
443_STOPPING_NOTE = "lilbee: stopping the engine. Press Ctrl-C again to force quit.\n"
445_STOPPING_NOTE_GRACE_S = 1.0
446"""Seconds a teardown may run silently; a fast exit stays as quiet as ever."""
448_FORCE_QUIT_NOTE = "lilbee: force quit. The engine stops with it.\n"
450_STRAGGLER_BUDGET_S = 5.0
451"""Seconds every leftover non-daemon thread gets, together, before the exit."""
453_STRAGGLER_NOTE = "lilbee: a background task would not stop; exiting anyway.\n"
456def _write_exit_note(note: str) -> None:
457 """Print an exit-path status line; stderr may already be gone at exit."""
458 try:
459 sys.stderr.write(note)
460 sys.stderr.flush()
461 except (OSError, ValueError):
462 pass
465def _all_threads() -> list[threading.Thread]:
466 """Seam over threading.enumerate, so tests never patch the stdlib global."""
467 return threading.enumerate()
470def wait_for_hard_exit_teardown() -> None:
471 """Block until any teardown thread (signal-driven or exit-driven) finishes.
473 Lets ``serve`` hold its OS locks through the fleet stop, so a successor
474 cannot acquire them while this server's models still occupy memory.
476 An interrupt while waiting force-quits the process: waiting is the only
477 thing left, the interpreter would otherwise still join the non-daemon
478 teardown thread, and every platform's child guard reaps the engine when
479 the process dies (kill-on-close job object, pdeathsig, death pipe).
480 """
481 try:
482 for thread in _all_threads():
483 if thread.name != _HARD_EXIT_THREAD_NAME:
484 continue
485 thread.join(_STOPPING_NOTE_GRACE_S)
486 if thread.is_alive():
487 _write_exit_note(_STOPPING_NOTE)
488 thread.join()
489 except KeyboardInterrupt:
490 _write_exit_note(_FORCE_QUIT_NOTE)
491 os._exit(_FORCE_QUIT_EXIT_CODE)
494def reset_services_on_exit() -> None:
495 """Tear the container down on a thread no signal reaches, and wait for it.
497 Engine release waits on the fleet build lock before it releases anything, so
498 a Ctrl-C on the main thread skips the release, and atexit cannot retry: the
499 singleton is already cleared. The teardown thread is non-daemon and takes
500 no signals; a wait that outlives the grace names what is happening.
501 """
502 if peek_services() is None:
503 return
504 threading.Thread(target=reset_services, name=_HARD_EXIT_THREAD_NAME).start()
505 wait_for_hard_exit_teardown()
508def _straggler_threads() -> list[threading.Thread]:
509 """Non-daemon threads the interpreter would join at shutdown, besides us."""
510 return [
511 thread
512 for thread in _all_threads()
513 if thread.is_alive()
514 and not thread.daemon
515 and thread is not threading.main_thread()
516 and thread is not threading.current_thread()
517 ]
520def exit_when_stragglers_would_hang() -> None:
521 """Exit rather than let interpreter shutdown join a wedged thread forever.
523 threading joins every non-daemon thread after atexit, so one worker
524 blocked in an unbounded network read hangs the process after all real
525 work is done. Stragglers share a short budget; whatever remains cannot
526 finish, and the child guard reaps the engine when the process dies.
527 Called from the TUI's own exit path and inert elsewhere: a test process
528 or embedding host owns its threads, and an exit would take them with it.
529 """
530 if not _state.interactive:
531 return
532 try:
533 deadline = time.monotonic() + _STRAGGLER_BUDGET_S
534 for thread in _straggler_threads():
535 thread.join(max(0.0, deadline - time.monotonic()))
536 if _straggler_threads():
537 _write_exit_note(_STRAGGLER_NOTE)
538 os._exit(0)
539 except KeyboardInterrupt:
540 _write_exit_note(_FORCE_QUIT_NOTE)
541 os._exit(_FORCE_QUIT_EXIT_CODE)
544def _teardown_for_signal(signum: int) -> None:
545 """Log the fatal signal, then stop services; runs off the signal handler's thread."""
546 log.info(
547 "Received signal %s; stopping the engine fleet before exit", signal.Signals(signum).name
548 )
549 reset_services()
552_lifecycle = _EngineLifecycle()
555def install_engine_lifecycle_hooks() -> None:
556 """Make a terminal close or ``kill`` stop the engine fleet instead of orphaning it."""
557 _lifecycle.install()
560atexit.register(reset_services_on_exit)