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

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 time 

19from collections.abc import Callable 

20from dataclasses import dataclass, field 

21from enum import StrEnum 

22from typing import Any 

23 

24logger = logging.getLogger(__name__) 

25 

26 

27# ============================================================================ 

28# Types 

29# ============================================================================ 

30 

31 

32class LifecyclePhase(StrEnum): 

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(StrEnum): 

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: LifecyclePhase | None = None 

78 message: str = "" 

79 error: str | None = None 

80 started_at: float | None = 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 (ComponentStatus.HEALTHY, ComponentStatus.DEGRADED) 

101 

102 

103# ============================================================================ 

104# Lifecycle Manager 

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

106 

107 

108class LifecycleManager: 

109 """Orchestrates ordered startup and graceful shutdown. 

110 

111 Usage: 

112 lm = LifecycleManager() 

113 

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

115 async def init_db(): 

116 ... 

117 

118 @lm.on_shutdown 

119 async def close_db(): 

120 ... 

121 

122 async with lm: 

123 # Application runs here 

124 ... 

125 

126 Signal handling (SIGTERM, SIGINT) integrated automatically. 

127 """ 

128 

129 PHASE_ORDER: list[LifecyclePhase] = [ 

130 LifecyclePhase.CONFIG, 

131 LifecyclePhase.INFRA, 

132 LifecyclePhase.SECURITY, 

133 LifecyclePhase.SERVICES, 

134 LifecyclePhase.MIDDLEWARE, 

135 LifecyclePhase.API, 

136 LifecyclePhase.READY, 

137 ] 

138 

139 def __init__( 

140 self, 

141 grace_period: float = 30.0, 

142 startup_timeout: float = 120.0, 

143 ): 

144 self._startup_hooks: list[LifecycleHook] = [] 

145 self._shutdown_hooks: list[LifecycleHook] = [] 

146 self._health: dict[str, ComponentHealth] = {} 

147 self._status = ComponentStatus.UNINITIALIZED 

148 self._start_time: float | None = None 

149 self._grace_period = grace_period 

150 self._startup_timeout = startup_timeout 

151 self._shutdown_event = asyncio.Event() 

152 self._ready_event = asyncio.Event() 

153 

154 # ── Registration ────────────────────────────────────────────────────── 

155 

156 def on_startup( 

157 self, 

158 name: str | None = None, 

159 *, 

160 phase: LifecyclePhase = LifecyclePhase.SERVICES, 

161 critical: bool = True, 

162 timeout_seconds: float = 30.0, 

163 weight: int = 50, 

164 retries: int = 0, 

165 retry_delay: float = 1.0, 

166 ) -> Callable: 

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

168 

169 def decorator(fn): 

170 hook_name = name or fn.__name__ 

171 is_async = asyncio.iscoroutinefunction(fn) 

172 self._startup_hooks.append( 

173 LifecycleHook( 

174 name=hook_name, 

175 phase=phase, 

176 fn=fn, 

177 timeout_seconds=timeout_seconds, 

178 is_async=is_async, 

179 critical=critical, 

180 weight=weight, 

181 retries=retries, 

182 retry_delay=retry_delay, 

183 ) 

184 ) 

185 self._health[hook_name] = ComponentHealth(name=hook_name, phase=phase) 

186 return fn 

187 

188 return decorator 

189 

190 def on_shutdown( 

191 self, 

192 name: str | None = None, 

193 *, 

194 timeout_seconds: float = 10.0, 

195 weight: int = 50, 

196 ) -> Callable: 

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

198 

199 def decorator(fn): 

200 hook_name = name or fn.__name__ 

201 is_async = asyncio.iscoroutinefunction(fn) 

202 self._shutdown_hooks.append( 

203 LifecycleHook( 

204 name=hook_name, 

205 phase=LifecyclePhase.READY, # irrelevant for shutdown 

206 fn=fn, 

207 timeout_seconds=timeout_seconds, 

208 is_async=is_async, 

209 critical=False, 

210 weight=weight, 

211 ) 

212 ) 

213 return fn 

214 

215 return decorator 

216 

217 # ── Startup ─────────────────────────────────────────────────────────── 

218 

219 async def start(self) -> LifecycleReport: 

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

221 self._start_time = time.perf_counter() 

222 self._status = ComponentStatus.STARTING 

223 

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

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

226 sorted_hooks = sorted( 

227 self._startup_hooks, 

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

229 ) 

230 

231 for hook in sorted_hooks: 

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

233 if not ok and hook.critical: 

234 self._status = ComponentStatus.UNHEALTHY 

235 return self.report() 

236 

237 self._status = ComponentStatus.HEALTHY 

238 self._ready_event.set() 

239 

240 return self.report() 

241 

242 async def _execute_hook(self, hook: LifecycleHook, is_startup: bool = True) -> bool: 

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

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

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

246 health.started_at = time.time() 

247 

248 attempt = 0 

249 last_error = None 

250 

251 while attempt <= hook.retries: 

252 t0 = time.perf_counter() 

253 try: 

254 if hook.is_async: 

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

256 else: 

257 loop = asyncio.get_event_loop() 

258 await asyncio.wait_for( 

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

260 timeout=hook.timeout_seconds, 

261 ) 

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

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

264 health.message = "OK" 

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

266 return True 

267 

268 except TimeoutError: 

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

270 health.error = last_error 

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

272 except Exception as e: 

273 last_error = str(e) 

274 health.error = last_error 

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

276 

277 attempt += 1 

278 if attempt <= hook.retries: 

279 await asyncio.sleep(hook.retry_delay) 

280 

281 health.status = ComponentStatus.UNHEALTHY 

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

283 return False 

284 

285 # ── Shutdown ────────────────────────────────────────────────────────── 

286 

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

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

289 if self._status == ComponentStatus.STOPPED: 

290 return self.report() 

291 

292 self._status = ComponentStatus.SHUTTING_DOWN 

293 self._shutdown_event.set() 

294 logger.info( 

295 f"[lifecycle] Shutting down gracefully{f' ({signal_name})' if signal_name else ''}" 

296 ) 

297 

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

299 for hook in reversed(self._shutdown_hooks): 

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

301 

302 self._status = ComponentStatus.STOPPED 

303 return self.report() 

304 

305 # ── Signal integration ──────────────────────────────────────────────── 

306 

307 def setup_signal_handlers(self, loop: asyncio.AbstractEventLoop | None = None): 

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

309 if loop is None: 

310 loop = asyncio.get_event_loop() 

311 

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

313 try: 

314 loop.add_signal_handler( 

315 sig, 

316 lambda s=sig: asyncio.ensure_future(self.shutdown(signal.Signals(s).name)), 

317 ) 

318 except (NotImplementedError, RuntimeError): 

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

320 signal.signal( 

321 sig, lambda s, f: asyncio.ensure_future(self.shutdown(signal.Signals(s).name)) 

322 ) 

323 

324 # ── Probes ──────────────────────────────────────────────────────────── 

325 

326 def is_ready(self) -> bool: 

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

328 return self._ready_event.is_set() 

329 

330 def is_live(self) -> bool: 

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

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

333 

334 # ── Report ──────────────────────────────────────────────────────────── 

335 

336 def report(self) -> LifecycleReport: 

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

338 startup_ms = 0.0 

339 if self._start_time: 

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

341 

342 return LifecycleReport( 

343 overall_status=self._status, 

344 phase=self._status.value, 

345 components=dict(self._health), 

346 startup_duration_ms=startup_ms, 

347 shutdown_remaining_hooks=len( 

348 [ 

349 h 

350 for h in self._shutdown_hooks 

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

352 not in (ComponentStatus.STOPPED,) 

353 ] 

354 ), 

355 ) 

356 

357 # ── Context manager ─────────────────────────────────────────────────── 

358 

359 async def __aenter__(self): 

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

361 self.setup_signal_handlers() 

362 await self.start() 

363 return self 

364 

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

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

367 await self.shutdown() 

368 return False # Don't suppress exceptions 

369 

370 

371# ============================================================================ 

372# Singleton helper 

373# ============================================================================ 

374 

375_default_lifecycle: LifecycleManager | None = None 

376 

377 

378def get_lifecycle( 

379 grace_period: float = 30.0, 

380 startup_timeout: float = 120.0, 

381) -> LifecycleManager: 

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

383 global _default_lifecycle 

384 if _default_lifecycle is None: 

385 _default_lifecycle = LifecycleManager( 

386 grace_period=grace_period, 

387 startup_timeout=startup_timeout, 

388 ) 

389 return _default_lifecycle