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 

8from pathlib import Path 

9from typing import TYPE_CHECKING 

10 

11import typer 

12 

13from lilbee.cli.app import ( 

14 apply_overrides, 

15 console, 

16 data_dir_option, 

17 global_option, 

18) 

19from lilbee.cli.commands.serve_logging import setup_server_log_file, setup_server_logging 

20from lilbee.core.config import cfg 

21 

22if TYPE_CHECKING: 

23 import uvicorn 

24 

25 

26def _port_file() -> Path: 

27 return cfg.data_dir / "server.port" 

28 

29 

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

31 exc = context.get("exception") 

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

33 if isinstance(exc, BaseException): 

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

35 else: 

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

37 

38 

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

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

41 import atexit 

42 

43 from lilbee.parent_monitor import parse_parent_pid, watch_parent_async 

44 

45 loop = asyncio.get_running_loop() 

46 loop.set_exception_handler(_log_loop_exception) 

47 

48 port_path = _port_file() 

49 

50 def _cleanup_port_file() -> None: 

51 port_path.unlink(missing_ok=True) 

52 

53 if not config.loaded: 

54 config.load() 

55 server.lifespan = config.lifespan_class(config) 

56 

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

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

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

60 started = False 

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

62 try: 

63 await server.startup() 

64 started = True 

65 

66 parent_pid = parse_parent_pid() 

67 if parent_pid is not None: 

68 

69 def _on_parent_death() -> None: 

70 server.should_exit = True 

71 

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

73 

74 if server.servers: 

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

76 actual_port = sock.getsockname()[1] 

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

78 port_path.write_text(str(actual_port)) 

79 atexit.register(_cleanup_port_file) 

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

81 await server.main_loop() 

82 finally: 

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

84 parent_watcher.cancel() 

85 port_path.unlink(missing_ok=True) 

86 if started: 

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

88 # original exception from main_loop reaches the caller intact. 

89 with contextlib.suppress(AttributeError): 

90 await server.shutdown() 

91 

92 

93def serve( 

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

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

96 data_dir: Path | None = data_dir_option, 

97 use_global: bool = global_option, 

98) -> None: 

99 """Start the HTTP API server.""" 

100 apply_overrides(data_dir=data_dir, use_global=use_global) 

101 if host is not None: 

102 cfg.server_host = host 

103 if port is not None: 

104 cfg.server_port = port 

105 

106 setup_server_logging() 

107 

108 import uvicorn 

109 

110 from lilbee.server import create_app 

111 

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

113 

114 app = create_app() 

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

116 setup_server_log_file() 

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

118 server = uvicorn.Server(config) 

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

120 

121 

122def mcp_cmd() -> None: 

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

124 from lilbee.mcp_server import main 

125 

126 main()