Coverage for src/lilbee/server/app.py: 100%
99 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +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
11import anyio.to_thread
12from litestar import Litestar, MediaType, Request, Response
13from litestar.config.cors import CORSConfig
14from litestar.exceptions import HTTPException
15from litestar.exceptions.responses import create_exception_response
16from litestar.middleware.base import DefineMiddleware
17from litestar.openapi import OpenAPIConfig
19from lilbee.app.services import get_services, peek_services
20from lilbee.app.version import get_version
21from lilbee.core.config import cfg
22from lilbee.providers.sdk_llm_provider import inject_provider_keys
23from lilbee.server.anthropic_api.routes import anthropic_router
24from lilbee.server.auth import AuthMiddleware, session_manager
25from lilbee.server.chat_completions_api.routes import completions_router
26from lilbee.server.mcp_mount import build_mcp_mount
27from lilbee.server.routes.agent_config import (
28 agent_config_index_route,
29 agent_config_route,
30)
31from lilbee.server.routes.crawl import crawl_route
32from lilbee.server.routes.documents import (
33 add_route,
34 add_upload_route,
35 documents_list_route,
36 documents_remove_route,
37 export_route,
38 import_route,
39 sync_route,
40)
41from lilbee.server.routes.general import (
42 config_defaults_route,
43 config_route,
44 config_update_route,
45 health_route,
46 shutdown_route,
47 source_content_route,
48 status_route,
49 warm_stream_route,
50)
51from lilbee.server.routes.memory import (
52 memories_list_route,
53 memories_remember_route,
54 memories_remove_route,
55 memories_update_route,
56)
57from lilbee.server.routes.models import (
58 models_catalog_route,
59 models_delete_route,
60 models_external_route,
61 models_installed_route,
62 models_list_route,
63 models_pull_route,
64 models_set_chat_route,
65 models_set_embedding_route,
66 models_set_reranker_route,
67 models_set_vision_route,
68 models_show_route,
69)
70from lilbee.server.routes.placement import (
71 gpu_stats_stream_route,
72 gpus_route,
73 placement_clear_route,
74 placement_preview_route,
75 placement_route,
76 placement_set_route,
77)
78from lilbee.server.routes.search import (
79 ask_route,
80 ask_stream_route,
81 chat_route,
82 chat_stream_route,
83 search_route,
84)
85from lilbee.server.routes.sessions import (
86 session_add_message_route,
87 session_claim_route,
88 session_create_route,
89 session_delete_route,
90 session_get_route,
91 session_rename_route,
92 session_set_summary_route,
93 sessions_list_route,
94)
95from lilbee.server.routes.setup import (
96 setup_crawler_route,
97 setup_crawler_status_route,
98)
99from lilbee.server.wiki import (
100 wiki_build_route,
101 wiki_citations_reverse_route,
102 wiki_draft_accept_route,
103 wiki_draft_diff_route,
104 wiki_draft_reject_route,
105 wiki_drafts_route,
106 wiki_generate_route,
107 wiki_index_route,
108 wiki_lint_route,
109 wiki_list_route,
110 wiki_prune_route,
111 wiki_read_route,
112 wiki_status_route,
113 wiki_stubs_route,
114 wiki_synthesize_route,
115 wiki_update_route,
116 wiki_wipe_route,
117)
119if TYPE_CHECKING:
120 from lilbee.retrieval.embedder import Embedder
122log = logging.getLogger(__name__)
125# Below this soft open-file limit a large agent fleet meets the limit as
126# connection failures before it saturates the machine.
127_FD_SOFT_LIMIT_NUDGE = 4096
130def _warn_if_few_file_descriptors() -> None:
131 """Log an advisory when the open-file limit is low enough to cap the agent fleet."""
132 try:
133 import resource
134 except ImportError: # pragma: no cover - Windows has no resource module
135 return
136 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
137 if soft == resource.RLIM_INFINITY or soft >= _FD_SOFT_LIMIT_NUDGE:
138 return
139 log.info(
140 "Open-file limit is %d. Each connected agent holds a socket, so a large "
141 "fleet will hit this before it saturates the machine; raise it with "
142 "'ulimit -n %d' before starting the server (this shell allows up to %s).",
143 soft,
144 _FD_SOFT_LIMIT_NUDGE,
145 "unlimited" if hard == resource.RLIM_INFINITY else hard,
146 )
149def _raise_thread_pool_ceiling() -> None:
150 """Set anyio's shared thread-pool size to ``mcp_tool_threads``.
152 Resizes anyio's own default limiter, not a private one, so every offload in
153 the process is lifted (Litestar and MCP sync handlers included), not only ours.
154 Must run on the server event loop, where the default limiter lives.
155 """
156 limiter = anyio.to_thread.current_default_thread_limiter()
157 if limiter.total_tokens != cfg.mcp_tool_threads:
158 limiter.total_tokens = cfg.mcp_tool_threads
161class _ServerLoop:
162 """Holds the running server's event loop so a config change off the loop can reach it.
164 A settings write from a worker thread (the MCP handler) needs the loop to
165 resize its thread pool there; nothing is held when no server runs (CLI/TUI).
166 """
168 def __init__(self) -> None:
169 self._loop: asyncio.AbstractEventLoop | None = None
171 def set(self, loop: asyncio.AbstractEventLoop | None) -> None:
172 self._loop = loop
174 def running(self) -> asyncio.AbstractEventLoop | None:
175 loop = self._loop
176 return loop if loop is not None and not loop.is_closed() else None
179_server_loop = _ServerLoop()
182def reapply_thread_pool_ceiling() -> None:
183 """Resize the running server's thread pool after ``mcp_tool_threads`` changes.
185 Marshalled onto the server loop, so it is safe from the HTTP handler (on the
186 loop) and the MCP handler (a worker thread) alike. With no server running the
187 new value simply takes effect at the next start.
188 """
189 loop = _server_loop.running()
190 if loop is not None:
191 loop.call_soon_threadsafe(_raise_thread_pool_ceiling)
194def _log_embedding_model_state(embedder: Embedder) -> None:
195 """Report whether embeddings will work, without letting that check stop startup.
197 A model that reports itself unavailable and one whose check raises both leave
198 the server usable for everything that does not embed, so neither is fatal.
199 """
200 try:
201 validated = embedder.validate_model()
202 except Exception:
203 log.warning("Failed to validate embedding model", exc_info=True)
204 return
205 if validated:
206 log.info("Embedding model validated")
207 elif cfg.embedding_model:
208 log.warning(
209 "Embedding model %s is unavailable; search and chat will run without embeddings",
210 cfg.embedding_model,
211 )
212 # Unconfigured: validate_model already logged the one INFO line.
215@asynccontextmanager
216async def _lifespan(app: Litestar) -> AsyncIterator[None]:
217 """Pre-load LLM provider and embedding model on server startup."""
218 _raise_thread_pool_ceiling()
219 _server_loop.set(asyncio.get_running_loop())
220 _warn_if_few_file_descriptors()
221 session_manager.load_or_generate()
223 inject_provider_keys()
225 try:
226 services = get_services() # pre-load all services (provider, embedder, etc.)
227 log.info("LLM provider pre-loaded")
228 except Exception:
229 log.warning("Failed to pre-load LLM provider", exc_info=True)
230 else:
231 _log_embedding_model_state(services.embedder)
232 try:
233 yield
234 finally:
235 _server_loop.set(None)
236 session_manager.cleanup()
237 # Terminate the provider's worker/fleet subprocesses so they don't
238 # outlive the server (e.g. on parent-death shutdown in managed mode).
239 svc = peek_services()
240 if svc is not None:
241 svc.provider.shutdown()
244def _json_error_response(request: Request, exc: Exception) -> Response:
245 """Answer errors as JSON even on a route whose success media type is SSE.
247 Litestar builds an error response with the route handler's media type, so a
248 streaming route would label a JSON error body text/event-stream.
249 """
250 response = create_exception_response(request, exc)
251 response.media_type = MediaType.JSON
252 return response
255def _unhandled_error_response(request: Request, exc: Exception) -> Response:
256 """Log an unmapped failure, then answer it as the same JSON 500.
258 Litestar's own handler logs nothing outside debug mode.
259 """
260 log.error("Unhandled error serving %s %s", request.method, request.url.path, exc_info=exc)
261 return _json_error_response(request, exc)
264def create_app() -> Litestar:
265 """Create the Litestar application instance."""
266 cors = CORSConfig(
267 allow_origins=cfg.cors_origins,
268 allow_origin_regex=cfg.cors_origin_regex,
269 allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
270 allow_headers=["Content-Type", "Authorization"],
271 )
272 mcp_route, mcp_session_lifespan = build_mcp_mount()
273 return Litestar(
274 lifespan=[_lifespan, mcp_session_lifespan],
275 middleware=[DefineMiddleware(AuthMiddleware)],
276 exception_handlers={
277 HTTPException: _json_error_response,
278 Exception: _unhandled_error_response,
279 },
280 route_handlers=[
281 mcp_route,
282 health_route,
283 warm_stream_route,
284 status_route,
285 shutdown_route,
286 config_route,
287 config_defaults_route,
288 config_update_route,
289 source_content_route,
290 search_route,
291 ask_route,
292 ask_stream_route,
293 chat_route,
294 chat_stream_route,
295 completions_router,
296 anthropic_router,
297 sync_route,
298 add_route,
299 add_upload_route,
300 models_list_route,
301 models_external_route,
302 models_set_chat_route,
303 models_set_embedding_route,
304 models_set_vision_route,
305 models_set_reranker_route,
306 models_catalog_route,
307 models_installed_route,
308 models_pull_route,
309 models_show_route,
310 models_delete_route,
311 documents_list_route,
312 documents_remove_route,
313 memories_list_route,
314 memories_remember_route,
315 memories_update_route,
316 memories_remove_route,
317 sessions_list_route,
318 session_get_route,
319 session_create_route,
320 session_add_message_route,
321 session_claim_route,
322 session_set_summary_route,
323 session_rename_route,
324 session_delete_route,
325 export_route,
326 import_route,
327 placement_route,
328 placement_preview_route,
329 placement_set_route,
330 placement_clear_route,
331 gpus_route,
332 gpu_stats_stream_route,
333 crawl_route,
334 agent_config_index_route,
335 agent_config_route,
336 setup_crawler_route,
337 setup_crawler_status_route,
338 wiki_list_route,
339 wiki_read_route,
340 wiki_drafts_route,
341 wiki_draft_diff_route,
342 wiki_draft_accept_route,
343 wiki_draft_reject_route,
344 wiki_citations_reverse_route,
345 wiki_lint_route,
346 wiki_prune_route,
347 wiki_index_route,
348 wiki_stubs_route,
349 wiki_generate_route,
350 wiki_wipe_route,
351 wiki_build_route,
352 wiki_update_route,
353 wiki_status_route,
354 wiki_synthesize_route,
355 ],
356 request_max_body_size=10 * 1024 * 1024,
357 cors_config=cors,
358 openapi_config=OpenAPIConfig(
359 title="lilbee",
360 description="Local knowledge base REST API",
361 version=get_version(),
362 path="/schema",
363 ),
364 )