Coverage for src/lilbee/server/app.py: 100%
147 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Litestar application factory: imports routes from modules, creates app with lifespan."""
3from __future__ import annotations
5import asyncio
6import logging
7from collections.abc import AsyncIterator
8from contextlib import asynccontextmanager
9from typing import TYPE_CHECKING, Any
11import anyio.to_thread
12from litestar import Litestar, MediaType, Request, Response
13from litestar.config.cors import CORSConfig
14from litestar.exceptions import HTTPException, ValidationException
15from litestar.exceptions.responses import create_exception_response
16from litestar.middleware import ASGIMiddleware
17from litestar.middleware.base import DefineMiddleware
18from litestar.openapi import OpenAPIConfig
20from lilbee.app.config_follow import ConfigFileFollower
21from lilbee.app.endpoints import MCP_PATH
22from lilbee.app.services import (
23 get_services,
24 peek_services,
25 set_config_refresh,
26 wait_for_opens,
27)
28from lilbee.app.version import get_version
29from lilbee.core.config import cfg
30from lilbee.runtime.progress import LibraryUpgradeEvent
31from lilbee.server.anthropic_api.errors import anthropic_error_body, anthropic_error_type
32from lilbee.server.anthropic_api.routes import anthropic_router
33from lilbee.server.auth import AuthMiddleware, session_manager
34from lilbee.server.chat_completions_api.errors import CompletionsErrorCode, completions_error_body
35from lilbee.server.chat_completions_api.routes import completions_router
36from lilbee.server.content_disposition import CONTENT_DISPOSITION
37from lilbee.server.mcp_mount import build_mcp_mount
38from lilbee.server.routes.agent_config import (
39 agent_config_index_route,
40 agent_config_route,
41)
42from lilbee.server.routes.analyze import analyze_router
43from lilbee.server.routes.crawl import crawl_route
44from lilbee.server.routes.documents import (
45 add_route,
46 add_upload_route,
47 documents_list_route,
48 documents_remove_route,
49 export_route,
50 import_route,
51 sync_route,
52)
53from lilbee.server.routes.general import (
54 config_defaults_route,
55 config_reset_route,
56 config_route,
57 config_schema_route,
58 config_sources_route,
59 config_update_route,
60 health_route,
61 shutdown_route,
62 source_content_route,
63 status_route,
64 warm_stream_route,
65)
66from lilbee.server.routes.memory import (
67 memories_list_route,
68 memories_remember_route,
69 memories_remove_route,
70 memories_update_route,
71)
72from lilbee.server.routes.models import (
73 models_catalog_route,
74 models_delete_route,
75 models_external_route,
76 models_installed_route,
77 models_list_route,
78 models_pull_route,
79 models_set_chat_route,
80 models_set_embedding_route,
81 models_set_reranker_route,
82 models_set_vision_route,
83 models_show_route,
84)
85from lilbee.server.routes.placement import (
86 gpu_stats_stream_route,
87 gpus_route,
88 placement_clear_route,
89 placement_preview_route,
90 placement_route,
91 placement_set_route,
92)
93from lilbee.server.routes.profiles import profiles_router
94from lilbee.server.routes.search import (
95 ask_route,
96 ask_stream_route,
97 chat_route,
98 chat_stream_route,
99 search_route,
100)
101from lilbee.server.routes.sessions import (
102 session_add_message_route,
103 session_claim_route,
104 session_create_route,
105 session_delete_route,
106 session_fork_route,
107 session_get_route,
108 session_markdown_route,
109 session_rename_route,
110 session_set_summary_route,
111 sessions_list_route,
112)
113from lilbee.server.routes.setup import (
114 setup_crawler_route,
115 setup_crawler_status_route,
116)
117from lilbee.server.validation_format import plain_validation
118from lilbee.server.wiki import (
119 wiki_build_route,
120 wiki_citations_reverse_route,
121 wiki_draft_accept_route,
122 wiki_draft_diff_route,
123 wiki_draft_reject_route,
124 wiki_drafts_route,
125 wiki_generate_route,
126 wiki_index_route,
127 wiki_lint_route,
128 wiki_list_route,
129 wiki_prune_route,
130 wiki_read_route,
131 wiki_status_route,
132 wiki_stubs_route,
133 wiki_synthesize_route,
134 wiki_update_route,
135 wiki_wipe_route,
136)
138if TYPE_CHECKING:
139 from litestar.types import ASGIApp, Receive, Scope, Send
141 from lilbee.retrieval.embedder import Embedder
143log = logging.getLogger(__name__)
146# Below this soft open-file limit a large agent fleet meets the limit as
147# connection failures before it saturates the machine.
148_FD_SOFT_LIMIT_NUDGE = 4096
151def _warn_if_few_file_descriptors() -> None:
152 """Log an advisory when the open-file limit is low enough to cap the agent fleet."""
153 try:
154 import resource
155 except ImportError: # pragma: no cover - Windows has no resource module
156 return
157 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
158 if soft == resource.RLIM_INFINITY or soft >= _FD_SOFT_LIMIT_NUDGE:
159 return
160 log.info(
161 "Open-file limit is %d. Each connected agent holds a socket, so a large "
162 "fleet will hit this before it saturates the machine; raise it with "
163 "'ulimit -n %d' before starting the server (this shell allows up to %s).",
164 soft,
165 _FD_SOFT_LIMIT_NUDGE,
166 "unlimited" if hard == resource.RLIM_INFINITY else hard,
167 )
170def _raise_thread_pool_ceiling() -> None:
171 """Set anyio's shared thread-pool size to ``mcp_tool_threads``.
173 Resizes anyio's own default limiter, not a private one, so every offload in
174 the process is lifted (Litestar and MCP sync handlers included), not only ours.
175 Must run on the server event loop, where the default limiter lives.
176 """
177 limiter = anyio.to_thread.current_default_thread_limiter()
178 if limiter.total_tokens != cfg.mcp_tool_threads:
179 limiter.total_tokens = cfg.mcp_tool_threads
182class _ServerLoop:
183 """Holds the running server's event loop so a config change off the loop can reach it.
185 A settings write from a worker thread (the MCP handler) needs the loop to
186 resize its thread pool there; nothing is held when no server runs (CLI/TUI).
187 """
189 def __init__(self) -> None:
190 self._loop: asyncio.AbstractEventLoop | None = None
192 def set(self, loop: asyncio.AbstractEventLoop | None) -> None:
193 self._loop = loop
195 def running(self) -> asyncio.AbstractEventLoop | None:
196 loop = self._loop
197 return loop if loop is not None and not loop.is_closed() else None
200_server_loop = _ServerLoop()
203def reapply_thread_pool_ceiling() -> None:
204 """Resize the running server's thread pool after ``mcp_tool_threads`` changes.
206 Marshalled onto the server loop, so it is safe from the HTTP handler (on the
207 loop) and the MCP handler (a worker thread) alike. With no server running the
208 new value simply takes effect at the next start.
209 """
210 loop = _server_loop.running()
211 if loop is not None:
212 loop.call_soon_threadsafe(_raise_thread_pool_ceiling)
215def _log_embedding_model_state(embedder: Embedder) -> None:
216 """Report whether embeddings will work, without letting that check stop startup.
218 A model that reports itself unavailable and one whose check raises both leave
219 the server usable for everything that does not embed, so neither is fatal.
220 """
221 try:
222 validated = embedder.validate_model()
223 except Exception:
224 log.warning("Failed to validate embedding model", exc_info=True)
225 return
226 if validated:
227 log.info("Embedding model validated")
228 elif cfg.embedding_model:
229 log.warning(
230 "Embedding model %s is unavailable; search and chat will run without embeddings",
231 cfg.embedding_model,
232 )
233 # Unconfigured: validate_model already logged the one INFO line.
236_UPGRADE_WAIT_SECONDS = 10.0
237# How long a stopping server waits for the library open it started, beside the provider
238# shutdown; an open still running then ends the process at once (exit_if_opens_remain).
239_OPEN_EXIT_WAIT_SECONDS = 5.0
240_UPGRADE_RETRY_AFTER_SECONDS = 10
241_UPGRADING_MESSAGE = "lilbee is upgrading the library. Try again in a few seconds."
244def _upgrading_response(path: str, upgrade: LibraryUpgradeEvent | None) -> Response:
245 """The 503 for a request that waited out the limit, in the error shape of its API.
247 Every shape also carries the upgrade's progress, as health reports it.
248 """
249 code = CompletionsErrorCode.INTERNAL_ERROR
250 body: dict[str, Any]
251 if path.startswith("/v1/messages"):
252 body = anthropic_error_body(anthropic_error_type(code), _UPGRADING_MESSAGE)
253 elif path.startswith("/v1/"):
254 body = completions_error_body(code, _UPGRADING_MESSAGE)
255 else:
256 body = {"status_code": 503, "detail": _UPGRADING_MESSAGE}
257 body["upgrade"] = None if upgrade is None else upgrade.model_dump(mode="json")
258 return Response(
259 body,
260 status_code=503,
261 headers={"Retry-After": str(_UPGRADE_RETRY_AFTER_SECONDS)},
262 media_type=MediaType.JSON,
263 )
266class _LibraryOpenWait(ASGIMiddleware):
267 """Hold a request up to 10 s for the library to open, then answer 503 with Retry-After.
269 The wait runs off the loop, so an upgrade never blocks it. The limit sits under the
270 15 s timeout of the Obsidian plugin, so the plugin shows the reason. Health, the warm
271 stream and shutdown pass at once: health reports the upgrade. MCP tool calls wait
272 for the library where they use it. A process with no services has no open to wait
273 for, and the request passes without building them.
274 """
276 exclude_path_pattern = (
277 "^/api/health$",
278 "^/api/warm/stream$",
279 "^/api/shutdown$",
280 f"^{MCP_PATH}(/.*)?$",
281 )
283 async def handle(self, scope: Scope, receive: Receive, send: Send, next_app: ASGIApp) -> None:
284 services = peek_services()
285 store = None if services is None else services.store
286 if (
287 store is not None
288 and not store.is_open
289 and not await asyncio.to_thread(store.wait_for_open, _UPGRADE_WAIT_SECONDS)
290 ):
291 response = _upgrading_response(scope["path"], store.upgrade_progress())
292 await response.to_asgi_response(None, Request(scope, receive, send))(
293 scope, receive, send
294 )
295 return
296 await next_app(scope, receive, send)
299class _ConfigFileReload(ASGIMiddleware):
300 """Take what another process changed in config.toml before each request runs.
302 The MCP mount is served by this middleware too, so its server does not follow the
303 file itself.
304 """
306 def __init__(self) -> None:
307 self._follower = ConfigFileFollower()
308 set_config_refresh(self._follower.take_now)
310 async def handle(self, scope: Scope, receive: Receive, send: Send, next_app: ASGIApp) -> None:
311 await self._follower.take()
312 await next_app(scope, receive, send)
315@asynccontextmanager
316async def _lifespan(app: Litestar) -> AsyncIterator[None]:
317 """Pre-load LLM provider and embedding model on server startup."""
318 _raise_thread_pool_ceiling()
319 _server_loop.set(asyncio.get_running_loop())
320 _warn_if_few_file_descriptors()
321 session_manager.load_or_generate()
323 try:
324 services = get_services() # pre-load all services (provider, embedder, etc.)
325 log.info("LLM provider pre-loaded")
326 except Exception:
327 log.warning("Failed to pre-load LLM provider", exc_info=True)
328 else:
329 _log_embedding_model_state(services.embedder)
330 services.store.open_in_background()
331 try:
332 yield
333 finally:
334 _server_loop.set(None)
335 session_manager.cleanup()
336 # Terminate the provider's worker/fleet subprocesses so they don't
337 # outlive the server (e.g. on parent-death shutdown in managed mode).
338 # The open wait runs beside it, so its bound counts from the same start.
339 svc = peek_services()
340 stops = [_wait_for_the_opens()]
341 if svc is not None:
342 stops.append(asyncio.to_thread(svc.provider.shutdown))
343 # A failed shutdown must not leave the open wait unawaited.
344 for outcome in await asyncio.gather(*stops, return_exceptions=True):
345 if isinstance(outcome, BaseException):
346 log.warning("A server shutdown step failed", exc_info=outcome)
349async def _wait_for_the_opens() -> None:
350 """Wait, within one bound, for every open; an open in LanceDB at exit aborts 3.12 on Linux."""
351 if not await asyncio.to_thread(wait_for_opens, _OPEN_EXIT_WAIT_SECONDS):
352 log.warning(
353 "The library is still opening after %.0f s; the server exits anyway.",
354 _OPEN_EXIT_WAIT_SECONDS,
355 )
358def _json_error_response(request: Request, exc: Exception) -> Response:
359 """Answer errors as JSON even on a route whose success media type is SSE.
361 Litestar builds an error response with the route handler's media type, so a
362 streaming route would label a JSON error body text/event-stream.
363 """
364 response = create_exception_response(request, exc)
365 response.media_type = MediaType.JSON
366 return response
369def _validation_error_response(request: Request, exc: ValidationException) -> Response:
370 """Answer a body that fails validation as the JSON 400, with no internal class name."""
371 return _json_error_response(request, plain_validation(exc))
374def _unhandled_error_response(request: Request, exc: Exception) -> Response:
375 """Log an unmapped failure, then answer it as the same JSON 500.
377 Litestar's own handler logs nothing outside debug mode.
378 """
379 log.error("Unhandled error serving %s %s", request.method, request.url.path, exc_info=exc)
380 return _json_error_response(request, exc)
383def create_app() -> Litestar:
384 """Create the Litestar application instance."""
385 cors = CORSConfig(
386 allow_origins=cfg.cors_origins,
387 allow_origin_regex=cfg.cors_origin_regex,
388 allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
389 allow_headers=["Content-Type", "Authorization"],
390 expose_headers=[CONTENT_DISPOSITION],
391 )
392 mcp_route, mcp_session_lifespan = build_mcp_mount()
393 return Litestar(
394 lifespan=[_lifespan, mcp_session_lifespan],
395 logging_config=None,
396 middleware=[DefineMiddleware(AuthMiddleware), _LibraryOpenWait(), _ConfigFileReload()],
397 exception_handlers={
398 HTTPException: _json_error_response,
399 ValidationException: _validation_error_response,
400 Exception: _unhandled_error_response,
401 },
402 route_handlers=[
403 mcp_route,
404 health_route,
405 warm_stream_route,
406 status_route,
407 shutdown_route,
408 config_route,
409 config_defaults_route,
410 config_reset_route,
411 config_schema_route,
412 config_sources_route,
413 config_update_route,
414 source_content_route,
415 search_route,
416 ask_route,
417 ask_stream_route,
418 chat_route,
419 chat_stream_route,
420 completions_router,
421 anthropic_router,
422 sync_route,
423 add_route,
424 add_upload_route,
425 models_list_route,
426 models_external_route,
427 models_set_chat_route,
428 models_set_embedding_route,
429 models_set_vision_route,
430 models_set_reranker_route,
431 models_catalog_route,
432 models_installed_route,
433 models_pull_route,
434 models_show_route,
435 models_delete_route,
436 documents_list_route,
437 documents_remove_route,
438 memories_list_route,
439 memories_remember_route,
440 memories_update_route,
441 memories_remove_route,
442 sessions_list_route,
443 session_get_route,
444 session_markdown_route,
445 session_create_route,
446 session_add_message_route,
447 session_fork_route,
448 session_claim_route,
449 session_set_summary_route,
450 session_rename_route,
451 session_delete_route,
452 export_route,
453 import_route,
454 profiles_router,
455 analyze_router,
456 placement_route,
457 placement_preview_route,
458 placement_set_route,
459 placement_clear_route,
460 gpus_route,
461 gpu_stats_stream_route,
462 crawl_route,
463 agent_config_index_route,
464 agent_config_route,
465 setup_crawler_route,
466 setup_crawler_status_route,
467 wiki_list_route,
468 wiki_read_route,
469 wiki_drafts_route,
470 wiki_draft_diff_route,
471 wiki_draft_accept_route,
472 wiki_draft_reject_route,
473 wiki_citations_reverse_route,
474 wiki_lint_route,
475 wiki_prune_route,
476 wiki_index_route,
477 wiki_stubs_route,
478 wiki_generate_route,
479 wiki_wipe_route,
480 wiki_build_route,
481 wiki_update_route,
482 wiki_status_route,
483 wiki_synthesize_route,
484 ],
485 request_max_body_size=10 * 1024 * 1024,
486 cors_config=cors,
487 openapi_config=OpenAPIConfig(
488 title="lilbee",
489 description="Local knowledge base REST API",
490 version=get_version(),
491 path="/schema",
492 ),
493 )