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 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 setup_server_log_file, setup_server_logging 

22from lilbee.core.config import cfg 

23from lilbee.runtime.lock import ( 

24 SERVER_LOCK_TIMEOUT, 

25 acquire_scope_lock, 

26 acquire_server_lock, 

27 read_scope_owner, 

28) 

29 

30if TYPE_CHECKING: 

31 import uvicorn 

32 

33 

34SCOPE_ENV = "LILBEE_EXCLUSIVE_SCOPE" 

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

36 

37LOCK_REFUSAL_EXIT_CODE = 3 

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

39 

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

41this" apart from a crash without parsing output. 

42""" 

43 

44 

45def port_file() -> Path: 

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

47 return cfg.data_dir / "server.port" 

48 

49 

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

51 exc = context.get("exception") 

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

53 if isinstance(exc, BaseException): 

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

55 else: 

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

57 

58 

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

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

61 import atexit 

62 

63 from lilbee.parent_monitor import parse_parent_pid, watch_parent_async 

64 

65 loop = asyncio.get_running_loop() 

66 loop.set_exception_handler(_log_loop_exception) 

67 

68 port_path = port_file() 

69 

70 def _cleanup_port_file() -> None: 

71 port_path.unlink(missing_ok=True) 

72 

73 if not config.loaded: 

74 config.load() 

75 server.lifespan = config.lifespan_class(config) 

76 

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

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

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

80 started = False 

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

82 try: 

83 await server.startup() 

84 started = True 

85 

86 parent_pid = parse_parent_pid() 

87 if parent_pid is not None: 

88 

89 def _on_parent_death() -> None: 

90 server.should_exit = True 

91 

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

93 

94 if server.servers: 

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

96 actual_port = sock.getsockname()[1] 

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

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

99 atexit.register(_cleanup_port_file) 

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

101 await server.main_loop() 

102 finally: 

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

104 parent_watcher.cancel() 

105 port_path.unlink(missing_ok=True) 

106 if started: 

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

108 # original exception from main_loop reaches the caller intact. 

109 with contextlib.suppress(AttributeError): 

110 await server.shutdown() 

111 

112 

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

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

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

116 console.print(message) 

117 raise typer.Exit(LOCK_REFUSAL_EXIT_CODE) 

118 

119 

120def serve( 

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

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

123 data_dir: Path | None = data_dir_option, 

124 use_global: bool = global_option, 

125) -> None: 

126 """Start the HTTP API server.""" 

127 apply_overrides(data_dir=data_dir, use_global=use_global) 

128 if host is not None: 

129 cfg.server_host = host 

130 if port is not None: 

131 cfg.server_port = port 

132 

133 setup_server_logging() 

134 

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

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

137 scope_hold = None 

138 scope_env = os.environ.get(SCOPE_ENV) 

139 if scope_env: 

140 scope_dir = Path(scope_env) 

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

142 if scope_hold is None: 

143 owner = read_scope_owner(scope_dir) 

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

145 message = ( 

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

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

148 ) 

149 _refuse_to_start(message) 

150 

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

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

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

154 if server_lock is None: 

155 if scope_hold is not None: 

156 scope_hold.release() 

157 message = ( 

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

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

160 ) 

161 _refuse_to_start(message) 

162 

163 import uvicorn 

164 

165 from lilbee.server import create_app 

166 

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

168 

169 try: 

170 app = create_app() 

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

172 setup_server_log_file() 

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

174 server = uvicorn.Server(config) 

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

176 finally: 

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

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

179 # server's models still occupy memory. 

180 wait_for_hard_exit_teardown() 

181 server_lock.release() 

182 if scope_hold is not None: 

183 scope_hold.release() 

184 

185 

186def mcp_cmd( 

187 data_dir: Path | None = data_dir_option, 

188 use_global: bool = global_option, 

189) -> None: 

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

191 apply_overrides(data_dir=data_dir, use_global=use_global) 

192 from lilbee.mcp_server import main 

193 

194 main()