| 1 | from collections import deque |
| 2 | import asyncio |
| 3 | from contextlib import asynccontextmanager, contextmanager |
| 4 | from dataclasses import dataclass |
| 5 | import faulthandler |
| 6 | import os |
| 7 | import sys |
| 8 | import threading |
| 9 | import time |
| 10 | from typing import Callable, Iterator |
| 11 | import urllib.request |
| 12 | |
| 13 | import uvicorn |
| 14 | |
| 15 | from helpers import process |
| 16 | from helpers.print_style import PrintStyle |
| 17 | |
| 18 | |
| 19 | def _env_int(name: str, default: int, minimum: int = 0) -> int: |
| 20 | try: |
| 21 | return max(minimum, int(os.getenv(name, str(default)))) |
| 22 | except (TypeError, ValueError): |
| 23 | return default |
| 24 | |
| 25 | |
| 26 | def _env_float(name: str, default: float, minimum: float = 0.0) -> float: |
| 27 | try: |
| 28 | return max(minimum, float(os.getenv(name, str(default)))) |
| 29 | except (TypeError, ValueError): |
| 30 | return default |
| 31 | |
| 32 | |
| 33 | @dataclass(frozen=True) |
| 34 | class StartupConfig: |
| 35 | timeout_seconds: int |
| 36 | max_attempts: int |
| 37 | retry_delay_seconds: float |
| 38 | |
| 39 | @classmethod |
| 40 | def from_env(cls) -> "StartupConfig": |
| 41 | return cls( |
| 42 | timeout_seconds=_env_int("A0_STARTUP_TIMEOUT_SECONDS", 90, minimum=15), |
| 43 | max_attempts=_env_int("A0_STARTUP_MAX_ATTEMPTS", 2, minimum=1), |
| 44 | retry_delay_seconds=_env_float( |
| 45 | "A0_STARTUP_RETRY_DELAY_SECONDS", 2.0, minimum=0.0 |
| 46 | ), |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | @dataclass |
| 51 | class StartupStageRecord: |
| 52 | name: str |
| 53 | timestamp: float |
| 54 | detail: str | None = None |
| 55 | |
| 56 | |
| 57 | class StartupMonitor: |
| 58 | def __init__( |
| 59 | self, |
| 60 | bind_host: str, |
| 61 | probe_host: str, |
| 62 | port: int, |
| 63 | attempt: int, |
| 64 | max_attempts: int, |
| 65 | timeout_seconds: int, |
| 66 | ) -> None: |
| 67 | self.bind_host = bind_host |
| 68 | self.probe_host = probe_host |
| 69 | self.port = port |
| 70 | self.attempt = attempt |
| 71 | self.max_attempts = max_attempts |
| 72 | self.timeout_seconds = timeout_seconds |
| 73 | self.start_time = time.monotonic() |
| 74 | self._stage = "created" |
| 75 | self._stage_detail: str | None = None |
| 76 | self._stage_started_at = self.start_time |
| 77 | self._history: deque[StartupStageRecord] = deque(maxlen=30) |
| 78 | self._history.append(StartupStageRecord(self._stage, self.start_time)) |
| 79 | self._ready = threading.Event() |
| 80 | self._stop = threading.Event() |
| 81 | self._lock = threading.RLock() |
| 82 | self._server: uvicorn.Server | None = None |
| 83 | self._watchdog_thread: threading.Thread | None = None |
| 84 | |
| 85 | def _prefix(self) -> str: |
| 86 | return f"[startup attempt {self.attempt}/{self.max_attempts}]" |
| 87 | |
| 88 | def mark(self, stage: str, detail: str | None = None) -> None: |
| 89 | now = time.monotonic() |
| 90 | with self._lock: |
| 91 | self._stage = stage |
| 92 | self._stage_detail = detail |
| 93 | self._stage_started_at = now |
| 94 | self._history.append(StartupStageRecord(stage, now, detail)) |
| 95 | elapsed = now - self.start_time |
| 96 | |
| 97 | suffix = f" ({detail})" if detail else "" |
| 98 | PrintStyle.debug(f"{self._prefix()} {stage}{suffix} at +{elapsed:.1f}s") |
| 99 | |
| 100 | @contextmanager |
| 101 | def stage(self, stage: str, detail: str | None = None) -> Iterator[None]: |
| 102 | self.mark(f"{stage}.start", detail) |
| 103 | try: |
| 104 | yield |
| 105 | except BaseException as e: |
| 106 | message = f"{type(e).__name__}: {e}" |
| 107 | self.mark(f"{stage}.error", message[:200]) |
| 108 | raise |
| 109 | else: |
| 110 | self.mark(f"{stage}.done", detail) |
| 111 | |
| 112 | def lifespan(self): |
| 113 | @asynccontextmanager |
| 114 | async def _lifespan(_app): |
| 115 | self.mark("starlette.lifespan.startup") |
| 116 | try: |
| 117 | yield |
| 118 | finally: |
| 119 | self.mark("starlette.lifespan.shutdown") |
| 120 | |
| 121 | return _lifespan |
| 122 | |
| 123 | def attach_server(self, server: uvicorn.Server) -> None: |
| 124 | with self._lock: |
| 125 | self._server = server |
| 126 | |
| 127 | def start_watchdog(self) -> None: |
| 128 | if self._watchdog_thread and self._watchdog_thread.is_alive(): |
| 129 | return |
| 130 | self._watchdog_thread = threading.Thread( |
| 131 | target=self._watchdog_loop, |
| 132 | daemon=True, |
| 133 | name=f"StartupWatchdog-{self.attempt}", |
| 134 | ) |
| 135 | self._watchdog_thread.start() |
| 136 | |
| 137 | def mark_ready(self, source: str = "health_check") -> None: |
| 138 | if self._ready.is_set(): |
| 139 | return |
| 140 | self.mark("ready", source) |
| 141 | self._ready.set() |
| 142 | self._stop.set() |
| 143 | |
| 144 | def is_ready(self) -> bool: |
| 145 | return self._ready.is_set() |
| 146 | |
| 147 | def close(self) -> None: |
| 148 | self._stop.set() |
| 149 | |
| 150 | def stop_event(self) -> threading.Event: |
| 151 | return self._stop |
| 152 | |
| 153 | def _watchdog_loop(self) -> None: |
| 154 | next_progress_log = self.start_time + 10 |
| 155 | while not self._stop.wait(timeout=1): |
| 156 | if self._ready.is_set(): |
| 157 | return |
| 158 | |
| 159 | now = time.monotonic() |
| 160 | if now >= next_progress_log: |
| 161 | stage, detail, stage_elapsed, total_elapsed, _history = self.snapshot() |
| 162 | detail_text = f" ({detail})" if detail else "" |
| 163 | PrintStyle.warning( |
| 164 | f"{self._prefix()} still waiting for readiness after " |
| 165 | f"{total_elapsed:.1f}s; current stage '{stage}' has been active " |
| 166 | f"for {stage_elapsed:.1f}s{detail_text}" |
| 167 | ) |
| 168 | next_progress_log = now + 10 |
| 169 | |
| 170 | if now - self.start_time >= self.timeout_seconds: |
| 171 | self._handle_timeout() |
| 172 | return |
| 173 | |
| 174 | def snapshot( |
| 175 | self, |
| 176 | ) -> tuple[str, str | None, float, float, list[StartupStageRecord]]: |
| 177 | now = time.monotonic() |
| 178 | with self._lock: |
| 179 | return ( |
| 180 | self._stage, |
| 181 | self._stage_detail, |
| 182 | now - self._stage_started_at, |
| 183 | now - self.start_time, |
| 184 | list(self._history), |
| 185 | ) |
| 186 | |
| 187 | def _handle_timeout(self) -> None: |
| 188 | stage, detail, stage_elapsed, total_elapsed, history = self.snapshot() |
| 189 | detail_text = f" ({detail})" if detail else "" |
| 190 | PrintStyle.error( |
| 191 | f"{self._prefix()} startup timed out after {total_elapsed:.1f}s while " |
| 192 | f"waiting for bind={self.bind_host}:{self.port} " |
| 193 | f"probe=http://{self.probe_host}:{self.port}/api/health; current stage " |
| 194 | f"'{stage}' has been active for {stage_elapsed:.1f}s{detail_text}" |
| 195 | ) |
| 196 | |
| 197 | PrintStyle.error(f"{self._prefix()} recent stage history follows:") |
| 198 | for record in history: |
| 199 | relative = record.timestamp - self.start_time |
| 200 | suffix = f" ({record.detail})" if record.detail else "" |
| 201 | PrintStyle.standard(f" +{relative:5.1f}s {record.name}{suffix}") |
| 202 | |
| 203 | active_threads = ", ".join( |
| 204 | f"{thread.name}(alive={thread.is_alive()}, daemon={thread.daemon})" |
| 205 | for thread in threading.enumerate() |
| 206 | ) |
| 207 | PrintStyle.error(f"{self._prefix()} active threads: {active_threads}") |
| 208 | PrintStyle.error( |
| 209 | f"{self._prefix()} dumping all thread stack traces for startup diagnosis" |
| 210 | ) |
| 211 | try: |
| 212 | faulthandler.dump_traceback(file=sys.stderr, all_threads=True) |
| 213 | except Exception as e: |
| 214 | PrintStyle.error(f"{self._prefix()} failed to dump thread traces: {e}") |
| 215 | |
| 216 | with self._lock: |
| 217 | server = self._server |
| 218 | |
| 219 | if server is not None: |
| 220 | PrintStyle.warning( |
| 221 | f"{self._prefix()} requesting uvicorn shutdown after startup timeout" |
| 222 | ) |
| 223 | server.should_exit = True |
| 224 | |
| 225 | if not self._stop.wait(timeout=3): |
| 226 | PrintStyle.error( |
| 227 | f"{self._prefix()} forcing process exit so the supervisor can restart it" |
| 228 | ) |
| 229 | os._exit(1) |
| 230 | |
| 231 | |
| 232 | def get_health_probe_host(bind_host: str) -> str: |
| 233 | if bind_host in {"0.0.0.0", "::", "[::]", ""}: |
| 234 | return "127.0.0.1" |
| 235 | return bind_host |
| 236 | |
| 237 | |
| 238 | def run_uvicorn_with_retries( |
| 239 | *, |
| 240 | host: str, |
| 241 | port: int, |
| 242 | build_asgi_app: Callable[[StartupMonitor], object], |
| 243 | flush_callback: Callable[[str], None], |
| 244 | access_log: bool = False, |
| 245 | log_level: str = "info", |
| 246 | ws: str = "wsproto", |
| 247 | startup_config: StartupConfig | None = None, |
| 248 | ) -> None: |
| 249 | startup_config = startup_config or StartupConfig.from_env() |
| 250 | health_host = get_health_probe_host(host) |
| 251 | PrintStyle.debug( |
| 252 | f"[startup] bind={host}:{port} probe=http://{health_host}:{port}/api/health " |
| 253 | f"timeout={startup_config.timeout_seconds}s attempts={startup_config.max_attempts}" |
| 254 | ) |
| 255 | |
| 256 | for attempt in range(1, startup_config.max_attempts + 1): |
| 257 | startup_monitor = StartupMonitor( |
| 258 | bind_host=host, |
| 259 | probe_host=health_host, |
| 260 | port=port, |
| 261 | attempt=attempt, |
| 262 | max_attempts=startup_config.max_attempts, |
| 263 | timeout_seconds=startup_config.timeout_seconds, |
| 264 | ) |
| 265 | try: |
| 266 | if _run_server_attempt( |
| 267 | host=host, |
| 268 | health_host=health_host, |
| 269 | port=port, |
| 270 | startup_monitor=startup_monitor, |
| 271 | build_asgi_app=build_asgi_app, |
| 272 | flush_callback=flush_callback, |
| 273 | access_log=access_log, |
| 274 | log_level=log_level, |
| 275 | ws=ws, |
| 276 | ): |
| 277 | return |
| 278 | except BaseException as e: |
| 279 | if isinstance(e, SystemExit) and startup_monitor.is_ready(): |
| 280 | raise |
| 281 | |
| 282 | PrintStyle.error( |
| 283 | f"[startup attempt {attempt}/{startup_config.max_attempts}] " |
| 284 | f"server startup failed before readiness with " |
| 285 | f"{type(e).__name__}: {e}" |
| 286 | ) |
| 287 | if attempt >= startup_config.max_attempts: |
| 288 | raise |
| 289 | else: |
| 290 | if attempt >= startup_config.max_attempts: |
| 291 | raise RuntimeError( |
| 292 | "Uvicorn exited before readiness on the final startup attempt." |
| 293 | ) |
| 294 | |
| 295 | PrintStyle.warning( |
| 296 | f"[startup attempt {attempt}/{startup_config.max_attempts}] " |
| 297 | "server exited before readiness; retrying" |
| 298 | ) |
| 299 | |
| 300 | if startup_config.retry_delay_seconds > 0: |
| 301 | PrintStyle.warning( |
| 302 | f"[startup attempt {attempt}/{startup_config.max_attempts}] " |
| 303 | f"sleeping {startup_config.retry_delay_seconds:.1f}s before retry" |
| 304 | ) |
| 305 | time.sleep(startup_config.retry_delay_seconds) |
| 306 | |
| 307 | raise RuntimeError("Server failed to reach readiness after all startup attempts.") |
| 308 | |
| 309 | |
| 310 | def _run_server_attempt( |
| 311 | *, |
| 312 | host: str, |
| 313 | health_host: str, |
| 314 | port: int, |
| 315 | startup_monitor: StartupMonitor, |
| 316 | build_asgi_app: Callable[[StartupMonitor], object], |
| 317 | flush_callback: Callable[[str], None], |
| 318 | access_log: bool, |
| 319 | log_level: str, |
| 320 | ws: str, |
| 321 | ) -> bool: |
| 322 | startup_monitor.start_watchdog() |
| 323 | try: |
| 324 | asgi_app = build_asgi_app(startup_monitor) |
| 325 | |
| 326 | with startup_monitor.stage("uvicorn.config.create"): |
| 327 | config = uvicorn.Config( |
| 328 | asgi_app, |
| 329 | host=host, |
| 330 | port=port, |
| 331 | log_level=log_level, |
| 332 | access_log=access_log, |
| 333 | ws=ws, |
| 334 | ) |
| 335 | |
| 336 | with startup_monitor.stage("uvicorn.server.create"): |
| 337 | server = uvicorn.Server(config) |
| 338 | |
| 339 | startup_monitor.attach_server(server) |
| 340 | process.set_server(_UvicornServerWrapper(server, flush_callback)) |
| 341 | |
| 342 | startup_monitor.mark("health.thread.start") |
| 343 | threading.Thread( |
| 344 | target=wait_for_health, |
| 345 | args=(health_host, port, startup_monitor), |
| 346 | daemon=True, |
| 347 | name=f"StartupHealth-{startup_monitor.attempt}", |
| 348 | ).start() |
| 349 | |
| 350 | PrintStyle().debug(f"Starting server at http://{host}:{port} ...") |
| 351 | startup_monitor.mark("uvicorn.run.enter") |
| 352 | _serve_uvicorn(server) |
| 353 | |
| 354 | if startup_monitor.is_ready(): |
| 355 | return True |
| 356 | |
| 357 | PrintStyle.warning( |
| 358 | f"[startup attempt {startup_monitor.attempt}/{startup_monitor.max_attempts}] " |
| 359 | "uvicorn exited before the health probe observed readiness" |
| 360 | ) |
| 361 | return False |
| 362 | finally: |
| 363 | startup_monitor.close() |
| 364 | process.set_server(None) |
| 365 | flush_callback("server_exit") |
| 366 | |
| 367 | |
| 368 | def wait_for_health(host: str, port: int, startup_monitor: StartupMonitor) -> None: |
| 369 | url = f"http://{host}:{port}/api/health" |
| 370 | while not startup_monitor.stop_event().is_set(): |
| 371 | try: |
| 372 | with urllib.request.urlopen(url, timeout=2) as resp: |
| 373 | if resp.status == 200: |
| 374 | startup_monitor.mark_ready("health_probe") |
| 375 | PrintStyle().print("Agent Zero is running.") |
| 376 | return |
| 377 | except Exception: |
| 378 | pass |
| 379 | startup_monitor.stop_event().wait(1) |
| 380 | |
| 381 | |
| 382 | class _UvicornServerWrapper: |
| 383 | def __init__( |
| 384 | self, server: uvicorn.Server, flush_callback: Callable[[str], None] |
| 385 | ) -> None: |
| 386 | self._server = server |
| 387 | self._flush_callback = flush_callback |
| 388 | |
| 389 | def shutdown(self) -> None: |
| 390 | self._flush_callback("shutdown") |
| 391 | self._server.should_exit = True |
| 392 | |
| 393 | |
| 394 | def _serve_uvicorn(server: uvicorn.Server) -> None: |
| 395 | # Avoid uvicorn.Server.run(), which delegates to asyncio.run(...) and can |
| 396 | # conflict with the global nest_asyncio patch used by the runtime. |
| 397 | # The project requires uvicorn>=0.38.0, where loop setup is exposed via |
| 398 | # Config.get_loop_factory(). |
| 399 | loop_factory = server.config.get_loop_factory() |
| 400 | with asyncio.Runner(loop_factory=loop_factory) as runner: |
| 401 | runner.run(server.serve()) |