Coverage for src/lilbee/cli/commands/servers.py: 100%

1"""Serve (HTTP API) and mcp (stdio) server-boot commands.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import contextlib 

7import logging 

8import os 

9from pathlib import Path 

10from typing import TYPE_CHECKING, NoReturn 

11 

12import typer 

13 

14from lilbee.app.services import set_server_exit_hook, wait_for_hard_exit_teardown 

15from lilbee.cli.app import ( 

16 apply_overrides, 

17 console, 

18 data_dir_option, 

19 global_option, 

20) 

21from lilbee.cli.commands.serve_logging import ( 

22 install_health_access_filter, 

23 setup_server_log_file, 

24 setup_server_logging, 

25) 

26from lilbee.core.config import cfg 

27from lilbee.runtime.lock import ( 

28 SERVER_LOCK_TIMEOUT, 

29 acquire_scope_lock, 

30 acquire_server_lock, 

31 read_scope_owner, 

32) 

33 

34if TYPE_CHECKING: 

35 import uvicorn 

36 

37 

38SCOPE_ENV = "LILBEE_EXCLUSIVE_SCOPE" 

39"""A directory that at most one server may serve at a time (a plugin's shared root).""" 

40 

41LOCK_REFUSAL_EXIT_CODE = 3 

42"""Exit code for a start refused because another live server holds a lock. 

43 

44Distinct from a generic failure so a supervisor can tell "another server owns 

45this" apart from a crash without parsing output. 

46""" 

47 

48 

49def port_file() -> Path: 

50 """Path to the running server's port file under ``cfg.data_dir``.""" 

51 return cfg.data_dir / "server.port" 

52 

53 

54def _log_loop_exception(_loop: asyncio.AbstractEventLoop, context: dict[str, object]) -> None: 

55 exc = context.get("exception") 

56 # isinstance: asyncio's context dict is untyped; "exception" may be absent 

57 if isinstance(exc, BaseException): 

58 logging.getLogger(__name__).error("asyncio task error", exc_info=exc) 

59 else: 

60 logging.getLogger(__name__).error("asyncio task error: %s", context.get("message")) 

61 

62 

63async def _run_server(server: uvicorn.Server, config: uvicorn.Config, host: str) -> None: 

64 """Start uvicorn, write port file, and clean up on shutdown.""" 

65 import atexit 

66 

67 from lilbee.parent_monitor import parse_parent_pid, watch_parent_async 

68 

69 loop = asyncio.get_running_loop() 

70 loop.set_exception_handler(_log_loop_exception) 

71 

72 port_path = port_file() 

73 

74 def _cleanup_port_file() -> None: 

75 port_path.unlink(missing_ok=True) 

76 

77 def _request_exit() -> None: 

78 server.should_exit = True 

79 

80 if not config.loaded: 

81 config.load() 

82 server.lifespan = config.lifespan_class(config) 

83 

84 # `server.servers` is set inside `startup()`. The finally below must skip 

85 # `shutdown()` when startup never ran: uvicorn dereferences `self.servers` 

86 # there and the resulting AttributeError would mask the original failure. 

87 started = False 

88 parent_watcher: asyncio.Task[None] | None = None 

89 try: 

90 set_server_exit_hook(_request_exit) 

91 await server.startup() 

92 started = True 

93 

94 parent_pid = parse_parent_pid() 

95 if parent_pid is not None: 

96 

97 def _on_parent_death() -> None: 

98 server.should_exit = True 

99 

100 parent_watcher = asyncio.create_task(watch_parent_async(parent_pid, _on_parent_death)) 

101 

102 if server.servers: 

103 sock = server.servers[0].sockets[0] 

104 actual_port = sock.getsockname()[1] 

105 port_path.parent.mkdir(parents=True, exist_ok=True) 

106 port_path.write_text(str(actual_port), encoding="utf-8") 

107 atexit.register(_cleanup_port_file) 

108 console.print(f"Listening on http://{host}:{actual_port}") 

109 await server.main_loop() 

110 finally: 

111 set_server_exit_hook(None) 

112 if parent_watcher is not None and not parent_watcher.done(): 

113 parent_watcher.cancel() 

114 port_path.unlink(missing_ok=True) 

115 if started: 

116 # Suppress AttributeError from a partial uvicorn bring-up so any 

117 # original exception from main_loop reaches the caller intact. 

118 with contextlib.suppress(AttributeError): 

119 await server.shutdown() 

120 

121 

122def _refuse_to_start(message: str) -> NoReturn: 

123 """Report why the server will not start, to the log and the terminal, then exit.""" 

124 logging.getLogger(__name__).error(message) 

125 console.print(message) 

126 raise typer.Exit(LOCK_REFUSAL_EXIT_CODE) 

127 

128 

129def serve( 

130 host: str = typer.Option(None, "--host", "-H", help="Bind address (default: 127.0.0.1)"), 

131 port: int = typer.Option(None, "--port", "-p", help="Port (default: 0/random)"), 

132 data_dir: Path | None = data_dir_option, 

133 use_global: bool = global_option, 

134) -> None: 

135 """Start the HTTP API server.""" 

136 apply_overrides(data_dir=data_dir, use_global=use_global) 

137 if host is not None: 

138 cfg.server_host = host 

139 if port is not None: 

140 cfg.server_port = port 

141 

142 setup_server_logging() 

143 

144 # One managed server per scope: the plugin passes its shared root here so a 

145 # second vault's server cannot start while another vault's is serving it. 

146 scope_hold = None 

147 scope_env = os.environ.get(SCOPE_ENV) 

148 if scope_env: 

149 scope_dir = Path(scope_env) 

150 scope_hold = acquire_scope_lock(scope_dir, cfg.data_dir, timeout=SERVER_LOCK_TIMEOUT) 

151 if scope_hold is None: 

152 owner = read_scope_owner(scope_dir) 

153 serving = f" It is serving {owner.data_dir}." if owner else "" 

154 message = ( 

155 f"Another lilbee server is already running for this installation.{serving}" 

156 " Stop it or wait for it to exit, then retry." 

157 ) 

158 _refuse_to_start(message) 

159 

160 # One server per data dir: a second instance would overwrite server.port 

161 # and spawn a second engine fleet against the same models and vector store. 

162 server_lock = acquire_server_lock(cfg.data_dir, timeout=SERVER_LOCK_TIMEOUT) 

163 if server_lock is None: 

164 if scope_hold is not None: 

165 scope_hold.release() 

166 message = ( 

167 "Another lilbee server is already running for this data directory. " 

168 "Stop it or wait for it to exit, then retry." 

169 ) 

170 _refuse_to_start(message) 

171 

172 import uvicorn 

173 

174 from lilbee.server import create_app 

175 

176 logging.getLogger("asyncio").setLevel(logging.ERROR) 

177 

178 try: 

179 app = create_app() 

180 # Litestar's app construction reconfigures root logging; re-install the file handler. 

181 setup_server_log_file() 

182 # Install after app construction so Litestar's logging reconfig cannot drop the filter. 

183 install_health_access_filter() 

184 config = uvicorn.Config(app, host=cfg.server_host, port=cfg.server_port) 

185 server = uvicorn.Server(config) 

186 asyncio.run(_run_server(server, config, cfg.server_host)) 

187 finally: 

188 # A signal-driven shutdown stops the fleet on its own thread; hold the 

189 # locks until it finishes so a successor cannot start while this 

190 # server's models still occupy memory. 

191 wait_for_hard_exit_teardown() 

192 server_lock.release() 

193 if scope_hold is not None: 

194 scope_hold.release() 

195 

196 

197def mcp_cmd( 

198 data_dir: Path | None = data_dir_option, 

199 use_global: bool = global_option, 

200) -> None: 

201 """Start the MCP server (stdio transport) for agent integration.""" 

202 apply_overrides(data_dir=data_dir, use_global=use_global) 

203 from lilbee.mcp_server import main 

204 

205 main()