Coverage for src/lilbee/cli/commands/servers.py: 100%
144 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Serve (HTTP API) and mcp (stdio) server-boot commands."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8import os
9from pathlib import Path
10from typing import TYPE_CHECKING, NoReturn
12import typer
14from lilbee.app.services import (
15 exit_after_signal,
16 exit_if_opens_remain,
17 report_library_upgrade,
18 set_server_exit_hook,
19 wait_for_hard_exit_teardown,
20)
21from lilbee.cli.app import (
22 apply_overrides,
23 console,
24 data_dir_option,
25 global_option,
26)
27from lilbee.cli.commands.serve_logging import (
28 install_health_access_filter,
29 setup_server_logging,
30)
31from lilbee.cli.upgrade_line import UpgradeLine
32from lilbee.core.config import cfg
33from lilbee.runtime.console import PlainConsole
34from lilbee.runtime.lock import (
35 SERVER_LOCK_TIMEOUT,
36 LockingUnsupportedError,
37 acquire_scope_lock,
38 acquire_server_lock,
39 read_scope_owner,
40)
42if TYPE_CHECKING:
43 import uvicorn
44 from filelock import FileLock
46 from lilbee.runtime.lock import ScopeHold
49SCOPE_ENV = "LILBEE_EXCLUSIVE_SCOPE"
50"""A directory that at most one server may serve at a time (a plugin's shared root)."""
52LOCK_REFUSAL_EXIT_CODE = 3
53"""Exit code for a start refused because another live server holds a lock.
55Distinct from a generic failure so a supervisor can tell "another server owns
56this" apart from a crash without parsing output.
57"""
60def port_file() -> Path:
61 """Path to the running server's port file under ``cfg.data_dir``."""
62 return cfg.data_dir / "server.port"
65def _log_loop_exception(_loop: asyncio.AbstractEventLoop, context: dict[str, object]) -> None:
66 exc = context.get("exception")
67 # isinstance: asyncio's context dict is untyped; "exception" may be absent
68 if isinstance(exc, BaseException):
69 logging.getLogger(__name__).error("asyncio task error", exc_info=exc)
70 else:
71 logging.getLogger(__name__).error("asyncio task error: %s", context.get("message"))
74class _ExitRequest:
75 """The exit a serving loop was asked for; a nonzero status ends the responses in flight."""
77 def __init__(self, server: uvicorn.Server, loop: asyncio.AbstractEventLoop) -> None:
78 self.status = 0
79 self._server = server
80 self._loop = loop
82 def _cancel_responses(self) -> None:
83 for task in self._server.server_state.tasks:
84 task.cancel()
86 def __call__(self, status: int) -> None:
87 self.status = status
88 self._server.should_exit = True
89 if status:
90 # A signal handler calls this, so the loop runs the cancel itself.
91 self._loop.call_soon_threadsafe(self._cancel_responses)
94async def _run_server(server: uvicorn.Server, config: uvicorn.Config, host: str) -> int:
95 """Start uvicorn, write port file, clean up on shutdown, and return the exit status."""
96 import atexit
98 from lilbee.parent_monitor import parse_parent_pid, watch_parent_async
100 loop = asyncio.get_running_loop()
101 loop.set_exception_handler(_log_loop_exception)
103 port_path = port_file()
105 def _cleanup_port_file() -> None:
106 port_path.unlink(missing_ok=True)
108 exit_request = _ExitRequest(server, loop)
110 if not config.loaded:
111 config.load()
112 server.lifespan = config.lifespan_class(config)
114 # `server.servers` is set inside `startup()`. The finally below must skip
115 # `shutdown()` when startup never ran: uvicorn dereferences `self.servers`
116 # there and the resulting AttributeError would mask the original failure.
117 started = False
118 parent_watcher: asyncio.Task[None] | None = None
119 try:
120 set_server_exit_hook(exit_request)
121 await server.startup()
122 started = True
124 parent_pid = parse_parent_pid()
125 if parent_pid is not None:
127 def _on_parent_death() -> None:
128 server.should_exit = True
130 parent_watcher = asyncio.create_task(watch_parent_async(parent_pid, _on_parent_death))
132 if server.servers:
133 sock = server.servers[0].sockets[0]
134 actual_port = sock.getsockname()[1]
135 port_path.parent.mkdir(parents=True, exist_ok=True)
136 port_path.write_text(str(actual_port), encoding="utf-8")
137 atexit.register(_cleanup_port_file)
138 console.print(f"Listening on http://{host}:{actual_port}")
139 await server.main_loop()
140 finally:
141 set_server_exit_hook(None)
142 if parent_watcher is not None and not parent_watcher.done():
143 parent_watcher.cancel()
144 port_path.unlink(missing_ok=True)
145 if started:
146 # Suppress AttributeError from a partial uvicorn bring-up so any
147 # original exception from main_loop reaches the caller intact.
148 with contextlib.suppress(AttributeError):
149 await server.shutdown()
150 return exit_request.status
153def _report_upgrade_as_log_lines() -> None:
154 """Show a library upgrade of this server as plain stderr lines, beside its log."""
155 report_library_upgrade(UpgradeLine(PlainConsole(stderr=True), live=False))
158def _refuse_to_start(message: str) -> NoReturn:
159 """Report why the server will not start, to the log and the terminal, then exit."""
160 logging.getLogger(__name__).error(message)
161 console.print(message, soft_wrap=True)
162 raise typer.Exit(LOCK_REFUSAL_EXIT_CODE)
165def serve(
166 host: str = typer.Option(None, "--host", "-H", help="Bind address (default: 127.0.0.1)"),
167 port: int = typer.Option(None, "--port", "-p", help="Port (default: 0/random)"),
168 data_dir: Path | None = data_dir_option,
169 use_global: bool = global_option,
170) -> None:
171 """Start the HTTP API server."""
172 apply_overrides(data_dir=data_dir, use_global=use_global)
173 _report_upgrade_as_log_lines()
174 if host is not None:
175 cfg.server_host = host
176 if port is not None:
177 cfg.server_port = port
179 setup_server_logging()
181 # One managed server per scope: the plugin passes its shared root here so a
182 # second vault's server cannot start while another vault's is serving it.
183 scope_hold = None
184 scope_env = os.environ.get(SCOPE_ENV)
185 if scope_env:
186 scope_dir = Path(scope_env)
187 try:
188 scope_hold = acquire_scope_lock(scope_dir, cfg.data_dir, timeout=SERVER_LOCK_TIMEOUT)
189 except LockingUnsupportedError as exc:
190 _refuse_to_start(str(exc))
191 if scope_hold is None:
192 owner = read_scope_owner(scope_dir)
193 serving = f" It is serving {owner.data_dir}." if owner else ""
194 message = (
195 f"Another lilbee server is already running for this installation.{serving}"
196 " Stop it or wait for it to exit, then retry."
197 )
198 _refuse_to_start(message)
200 # One server per data dir: a second instance would overwrite server.port
201 # and spawn a second engine fleet against the same models and vector store.
202 message = (
203 "Another lilbee server is already running for this data directory. "
204 "Stop it or wait for it to exit, then retry."
205 )
206 try:
207 server_lock = acquire_server_lock(cfg.data_dir, timeout=SERVER_LOCK_TIMEOUT)
208 except LockingUnsupportedError as exc:
209 server_lock, message = None, str(exc)
210 if server_lock is None:
211 if scope_hold is not None:
212 scope_hold.release()
213 _refuse_to_start(message)
215 import uvicorn
217 from lilbee.server import create_app
219 logging.getLogger("asyncio").setLevel(logging.ERROR)
221 status = 1
222 try:
223 app = create_app()
224 install_health_access_filter()
225 config = uvicorn.Config(app, host=cfg.server_host, port=cfg.server_port)
226 server = uvicorn.Server(config)
227 status = asyncio.run(_run_server(server, config, cfg.server_host))
228 except Exception:
229 # The exit in the finally can end the process before the traceback prints.
230 logging.getLogger(__name__).error("The server run failed.", exc_info=True)
231 raise
232 finally:
233 _stop_server(status, server_lock, scope_hold)
234 raise typer.Exit(status)
237def _stop_server(status: int, server_lock: FileLock, scope_hold: ScopeHold | None) -> None:
238 """Free the locks after the teardown, then end the process if a library open still runs."""
239 # A signal-driven shutdown stops the fleet on its own thread; hold the
240 # locks until it finishes so a successor cannot start while this
241 # server's models still occupy memory.
242 wait_for_hard_exit_teardown()
243 server_lock.release()
244 if scope_hold is not None:
245 scope_hold.release()
246 exit_if_opens_remain(status)
249def mcp_cmd(
250 data_dir: Path | None = data_dir_option,
251 use_global: bool = global_option,
252) -> None:
253 """Start the MCP server (stdio transport) for agent integration."""
254 apply_overrides(data_dir=data_dir, use_global=use_global)
255 _report_upgrade_as_log_lines()
256 from lilbee.mcp_server import main
258 try:
259 main()
260 except SystemExit as exc:
261 exit_after_signal(exc.code if isinstance(exc.code, int) else 1)
262 raise