Coverage for agentos/core/lifecycle.py: 99%

1"""AgentOS Lifecycle — graceful startup/shutdown with ordered hooks. 

2 

3Provides enterprise-grade process lifecycle management: 

4- Ordered startup hooks with timeout and health gate 

5- Ordered shutdown hooks with grace period 

6- SIGTERM/SIGINT graceful shutdown integration 

7- Component-level health registration 

8- Liveness/readiness probe support 

9 

10Design: ~370 lines, zero external deps beyond stdlib + asyncio. 

11""" 

12 

13from __future__ import annotations 

14 

15import asyncio 

16import logging 

17import signal 

18import sys 

19import time 

20from collections import OrderedDict 

21from dataclasses import dataclass, field 

22from enum import Enum 

23from typing import Any, Awaitable, Callable, Dict, List, Optional 

24 

25logger = logging.getLogger(__name__) 

26 

27 

28# ============================================================================ 

29# Types 

30# ============================================================================ 

31 

32class LifecyclePhase(str, Enum): 

33 """Ordered lifecycle phases during startup.""" 

34 

35 CONFIG = "config" # Configuration loading 

36 INFRA = "infra" # DB, Redis, message queues 

37 SECURITY = "security" # Auth, encryption, certs 

38 SERVICES = "services" # Internal services 

39 MIDDLEWARE = "middleware" # Middleware pipeline 

40 API = "api" # HTTP/gRPC server 

41 READY = "ready" # Final readiness signal 

42 

43 

44class ComponentStatus(str, Enum): 

45 """Component health status.""" 

46 

47 UNINITIALIZED = "uninitialized" 

48 STARTING = "starting" 

49 HEALTHY = "healthy" 

50 DEGRADED = "degraded" 

51 UNHEALTHY = "unhealthy" 

52 SHUTTING_DOWN = "shutting_down" 

53 STOPPED = "stopped" 

54 

55 

56@dataclass 

57class LifecycleHook: 

58 """A single startup or shutdown hook with metadata.""" 

59 

60 name: str 

61 phase: LifecyclePhase 

62 fn: Callable[[], Any] # or async callable 

63 timeout_seconds: float = 30.0 

64 is_async: bool = False 

65 critical: bool = True # Fail startup if critical hook fails 

66 weight: int = 50 # Ordering within same phase (lower = first) 

67 retries: int = 0 

68 retry_delay: float = 1.0 

69 

70 

71@dataclass 

72class ComponentHealth: 

73 """Health status of a single component.""" 

74 

75 name: str 

76 status: ComponentStatus = ComponentStatus.UNINITIALIZED 

77 phase: Optional[LifecyclePhase] = None 

78 message: str = "" 

79 error: Optional[str] = None 

80 started_at: Optional[float] = None 

81 duration_ms: float = 0.0 

82 

83 

84@dataclass 

85class LifecycleReport: 

86 """Full lifecycle status report.""" 

87 

88 overall_status: ComponentStatus = ComponentStatus.UNINITIALIZED 

89 phase: str = "" 

90 components: Dict[str, ComponentHealth] = field(default_factory=dict) 

91 startup_duration_ms: float = 0.0 

92 shutdown_remaining_hooks: int = 0 

93 

94 @property 

95 def is_healthy(self) -> bool: 

96 return self.overall_status == ComponentStatus.HEALTHY 

97 

98 @property 

99 def is_ready(self) -> bool: 

100 return self.overall_status in ( 

101 ComponentStatus.HEALTHY, ComponentStatus.DEGRADED 

102 ) 

103 

104 

105# ============================================================================ 

106# Lifecycle Manager 

107# ============================================================================ 

108 

109class LifecycleManager: 

110 """Orchestrates ordered startup and graceful shutdown. 

111 

112 Usage: 

113 lm = LifecycleManager() 

114 

115 @lm.on_startup(phase=LifecyclePhase.INFRA) 

116 async def init_db(): 

117 ... 

118 

119 @lm.on_shutdown 

120 async def close_db(): 

121 ... 

122 

123 async with lm: 

124 # Application runs here 

125 ... 

126 

127 Signal handling (SIGTERM, SIGINT) integrated automatically. 

128 """ 

129 

130 PHASE_ORDER: List[LifecyclePhase] = [ 

131 LifecyclePhase.CONFIG, 

132 LifecyclePhase.INFRA, 

133 LifecyclePhase.SECURITY, 

134 LifecyclePhase.SERVICES, 

135 LifecyclePhase.MIDDLEWARE, 

136 LifecyclePhase.API, 

137 LifecyclePhase.READY, 

138 ] 

139 

140 def __init__( 

141 self, 

142 grace_period: float = 30.0, 

143 startup_timeout: float = 120.0, 

144 ): 

145 self._startup_hooks: List[LifecycleHook] = [] 

146 self._shutdown_hooks: List[LifecycleHook] = [] 

147 self._health: Dict[str, ComponentHealth] = {} 

148 self._status = ComponentStatus.UNINITIALIZED 

149 self._start_time: Optional[float] = None 

150 self._grace_period = grace_period 

151 self._startup_timeout = startup_timeout 

152 self._shutdown_event = asyncio.Event() 

153 self._ready_event = asyncio.Event() 

154 

155 # ── Registration ────────────────────────────────────────────────────── 

156 

157 def on_startup( 

158 self, 

159 name: Optional[str] = None, 

160 *, 

161 phase: LifecyclePhase = LifecyclePhase.SERVICES, 

162 critical: bool = True, 

163 timeout_seconds: float = 30.0, 

164 weight: int = 50, 

165 retries: int = 0, 

166 retry_delay: float = 1.0, 

167 ) -> Callable: 

168 """Decorator: register a startup hook.""" 

169 def decorator(fn): 

170 hook_name = name or fn.__name__ 

171 is_async = asyncio.iscoroutinefunction(fn) 

172 self._startup_hooks.append(LifecycleHook( 

173 name=hook_name, 

174 phase=phase, 

175 fn=fn, 

176 timeout_seconds=timeout_seconds, 

177 is_async=is_async, 

178 critical=critical, 

179 weight=weight, 

180 retries=retries, 

181 retry_delay=retry_delay, 

182 )) 

183 self._health[hook_name] = ComponentHealth( 

184 name=hook_name, phase=phase 

185 ) 

186 return fn 

187 return decorator 

188 

189 def on_shutdown( 

190 self, 

191 name: Optional[str] = None, 

192 *, 

193 timeout_seconds: float = 10.0, 

194 weight: int = 50, 

195 ) -> Callable: 

196 """Decorator: register a shutdown hook (reverse order).""" 

197 def decorator(fn): 

198 hook_name = name or fn.__name__ 

199 is_async = asyncio.iscoroutinefunction(fn) 

200 self._shutdown_hooks.append(LifecycleHook( 

201 name=hook_name, 

202 phase=LifecyclePhase.READY, # irrelevant for shutdown 

203 fn=fn, 

204 timeout_seconds=timeout_seconds, 

205 is_async=is_async, 

206 critical=False, 

207 weight=weight, 

208 )) 

209 return fn 

210 return decorator 

211 

212 # ── Startup ─────────────────────────────────────────────────────────── 

213 

214 async def start(self) -> LifecycleReport: 

215 """Execute all startup hooks in phase/weight order.""" 

216 self._start_time = time.perf_counter() 

217 self._status = ComponentStatus.STARTING 

218 

219 # Sort: phase order first, then weight within phase 

220 phase_idx = {p: i for i, p in enumerate(self.PHASE_ORDER)} 

221 sorted_hooks = sorted( 

222 self._startup_hooks, 

223 key=lambda h: (phase_idx.get(h.phase, 99), h.weight), 

224 ) 

225 

226 for hook in sorted_hooks: 

227 ok = await self._execute_hook(hook, is_startup=True) 

228 if not ok and hook.critical: 

229 self._status = ComponentStatus.UNHEALTHY 

230 return self.report() 

231 

232 self._status = ComponentStatus.HEALTHY 

233 self._ready_event.set() 

234 

235 return self.report() 

236 

237 async def _execute_hook( 

238 self, hook: LifecycleHook, is_startup: bool = True 

239 ) -> bool: 

240 """Execute a single hook with timeout, retries, and health tracking.""" 

241 health = self._health.get(hook.name) or ComponentHealth(name=hook.name) 

242 health.status = ComponentStatus.STARTING if is_startup else ComponentStatus.SHUTTING_DOWN 

243 health.started_at = time.time() 

244 

245 attempt = 0 

246 last_error = None 

247 

248 while attempt <= hook.retries: 

249 t0 = time.perf_counter() 

250 try: 

251 if hook.is_async: 

252 await asyncio.wait_for(hook.fn(), timeout=hook.timeout_seconds) 

253 else: 

254 loop = asyncio.get_event_loop() 

255 await asyncio.wait_for( 

256 loop.run_in_executor(None, hook.fn), 

257 timeout=hook.timeout_seconds, 

258 ) 

259 health.duration_ms = (time.perf_counter() - t0) * 1000 

260 health.status = ComponentStatus.HEALTHY if is_startup else ComponentStatus.STOPPED 

261 health.message = "OK" 

262 logger.info(f"[lifecycle] {hook.name}: OK ({health.duration_ms:.0f}ms)") 

263 return True 

264 

265 except asyncio.TimeoutError: 

266 last_error = f"Timeout after {hook.timeout_seconds}s" 

267 health.error = last_error 

268 logger.warning(f"[lifecycle] {hook.name}: {last_error}") 

269 except Exception as e: 

270 last_error = str(e) 

271 health.error = last_error 

272 logger.warning(f"[lifecycle] {hook.name}: {last_error}") 

273 

274 attempt += 1 

275 if attempt <= hook.retries: 

276 await asyncio.sleep(hook.retry_delay) 

277 

278 health.status = ComponentStatus.UNHEALTHY 

279 health.duration_ms = (time.perf_counter() - t0) * 1000 

280 return False 

281 

282 # ── Shutdown ────────────────────────────────────────────────────────── 

283 

284 async def shutdown(self, signal_name: str = "") -> LifecycleReport: 

285 """Execute all shutdown hooks in reverse registration order.""" 

286 if self._status == ComponentStatus.STOPPED: 

287 return self.report() 

288 

289 self._status = ComponentStatus.SHUTTING_DOWN 

290 self._shutdown_event.set() 

291 logger.info(f"[lifecycle] Shutting down gracefully{f' ({signal_name})' if signal_name else ''}") 

292 

293 # Reverse order for shutdown (LIFO — last started, first stopped) 

294 for hook in reversed(self._shutdown_hooks): 

295 await self._execute_hook(hook, is_startup=False) 

296 

297 self._status = ComponentStatus.STOPPED 

298 return self.report() 

299 

300 # ── Signal integration ──────────────────────────────────────────────── 

301 

302 def setup_signal_handlers(self, loop: Optional[asyncio.AbstractEventLoop] = None): 

303 """Register SIGTERM/SIGINT handlers on the event loop.""" 

304 if loop is None: 

305 loop = asyncio.get_event_loop() 

306 

307 for sig in (signal.SIGTERM, signal.SIGINT): 

308 try: 

309 loop.add_signal_handler( 

310 sig, 

311 lambda s=sig: asyncio.ensure_future( 

312 self.shutdown(signal.Signals(s).name) 

313 ), 

314 ) 

315 except (NotImplementedError, RuntimeError): 

316 # Windows or non-main-thread — fallback to signal.signal 

317 signal.signal(sig, lambda s, f: asyncio.ensure_future( 

318 self.shutdown(signal.Signals(s).name) 

319 )) 

320 

321 # ── Probes ──────────────────────────────────────────────────────────── 

322 

323 def is_ready(self) -> bool: 

324 """Readiness probe: is the service ready to accept requests?""" 

325 return self._ready_event.is_set() 

326 

327 def is_live(self) -> bool: 

328 """Liveness probe: is the process alive (not hung)?""" 

329 return self._status not in (ComponentStatus.STOPPED, ComponentStatus.UNHEALTHY) 

330 

331 # ── Report ──────────────────────────────────────────────────────────── 

332 

333 def report(self) -> LifecycleReport: 

334 """Generate a full lifecycle status report.""" 

335 startup_ms = 0.0 

336 if self._start_time: 

337 startup_ms = (time.perf_counter() - self._start_time) * 1000 

338 

339 return LifecycleReport( 

340 overall_status=self._status, 

341 phase=self._status.value, 

342 components=dict(self._health), 

343 startup_duration_ms=startup_ms, 

344 shutdown_remaining_hooks=len([ 

345 h for h in self._shutdown_hooks 

346 if self._health.get(h.name, ComponentHealth(name=h.name)).status 

347 not in (ComponentStatus.STOPPED,) 

348 ]), 

349 ) 

350 

351 # ── Context manager ─────────────────────────────────────────────────── 

352 

353 async def __aenter__(self): 

354 """Async context manager: start lifecycle.""" 

355 self.setup_signal_handlers() 

356 await self.start() 

357 return self 

358 

359 async def __aexit__(self, exc_type, exc_val, exc_tb): 

360 """Async context manager: graceful shutdown.""" 

361 await self.shutdown() 

362 return False # Don't suppress exceptions 

363 

364 

365# ============================================================================ 

366# Singleton helper 

367# ============================================================================ 

368 

369_default_lifecycle: Optional[LifecycleManager] = None 

370 

371 

372def get_lifecycle( 

373 grace_period: float = 30.0, 

374 startup_timeout: float = 120.0, 

375) -> LifecycleManager: 

376 """Get or create the global LifecycleManager singleton.""" 

377 global _default_lifecycle 

378 if _default_lifecycle is None: 

379 _default_lifecycle = LifecycleManager( 

380 grace_period=grace_period, 

381 startup_timeout=startup_timeout, 

382 ) 

383 return _default_lifecycle