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

1""" 

2Production health check with dependency liveness probes. 

3 

4Extends the basic health endpoint with: 

5- Database connectivity check (async) 

6- Redis connectivity check (async) 

7- Component-level health status 

8 

9Usage: 

10 from agentos.core.health import HealthChecker 

11 checker = HealthChecker(db_url="...", redis_url="...") 

12 status = await checker.check() 

13""" 

14 

15from __future__ import annotations 

16 

17import asyncio 

18import logging 

19import time 

20from dataclasses import dataclass, field 

21 

22logger = logging.getLogger(__name__) 

23 

24 

25@dataclass 

26class ComponentHealth: 

27 name: str 

28 status: str # "healthy" | "degraded" | "unhealthy" 

29 latency_ms: float 

30 error: str | None = None 

31 

32 

33@dataclass 

34class HealthReport: 

35 status: str # "healthy" | "degraded" | "unhealthy" 

36 uptime_seconds: float 

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

38 timestamp: float = field(default_factory=time.time) 

39 

40 

41class HealthChecker: 

42 """Async health checker with component-level probing.""" 

43 

44 def __init__( 

45 self, 

46 start_time: float, 

47 db_url: str | None = None, 

48 redis_url: str | None = None, 

49 ): 

50 self.start_time = start_time 

51 self.db_url = db_url 

52 self.redis_url = redis_url 

53 

54 async def _probe(self, name: str, check_fn, timeout: float = 3.0) -> ComponentHealth: 

55 """Run a single component health probe with timeout.""" 

56 t0 = time.perf_counter() 

57 try: 

58 await asyncio.wait_for(check_fn(), timeout=timeout) 

59 latency = (time.perf_counter() - t0) * 1000 

60 return ComponentHealth(name=name, status="healthy", latency_ms=latency) 

61 except TimeoutError: 

62 latency = (time.perf_counter() - t0) * 1000 

63 return ComponentHealth( 

64 name=name, status="unhealthy", latency_ms=latency, error=f"Timeout after {timeout}s" 

65 ) 

66 except Exception as e: 

67 latency = (time.perf_counter() - t0) * 1000 

68 return ComponentHealth(name=name, status="unhealthy", latency_ms=latency, error=str(e)) 

69 

70 async def check(self) -> HealthReport: 

71 """Run a full health check across all configured components.""" 

72 probes = [] 

73 

74 # DB probe 

75 if self.db_url: 

76 probes.append(self._probe("database", self._check_db)) 

77 

78 # Redis probe 

79 if self.redis_url: 

80 probes.append(self._probe("redis", self._check_redis)) 

81 

82 # Always probe disk (write test) 

83 probes.append(self._probe("disk", self._check_disk)) 

84 

85 results = await asyncio.gather(*probes, return_exceptions=True) 

86 

87 components: dict[str, ComponentHealth] = {} 

88 overall = "healthy" 

89 

90 for r in results: 

91 if isinstance(r, ComponentHealth): 

92 components[r.name] = r 

93 if r.status == "unhealthy": 

94 if overall == "healthy": 

95 overall = "degraded" 

96 elif r.status == "degraded" and overall == "healthy": 

97 overall = "degraded" 

98 elif isinstance(r, Exception): 

99 # Probe itself crashed 

100 components["internal"] = ComponentHealth( 

101 name="internal", status="unhealthy", latency_ms=0, error=str(r) 

102 ) 

103 overall = "unhealthy" 

104 

105 return HealthReport( 

106 status=overall, 

107 uptime_seconds=time.time() - self.start_time, 

108 components=components, 

109 ) 

110 

111 async def _check_db(self): 

112 """Database connectivity probe.""" 

113 from sqlalchemy import text 

114 from sqlalchemy.ext.asyncio import create_async_engine 

115 

116 engine = create_async_engine(self.db_url, echo=False) 

117 async with engine.connect() as conn: 

118 await conn.execute(text("SELECT 1")) 

119 await engine.dispose() 

120 

121 async def _check_redis(self): 

122 """Redis connectivity probe.""" 

123 import redis.asyncio as redis 

124 

125 r = redis.from_url(self.redis_url) 

126 await r.ping() 

127 await r.close() 

128 

129 async def _check_disk(self): 

130 """Filesystem write test.""" 

131 import os 

132 import tempfile 

133 

134 with tempfile.NamedTemporaryFile(delete=False, prefix="health_", suffix=".tmp") as f: 

135 f.write(b"ok") 

136 

137 try: 

138 os.unlink(f.name) 

139 except OSError: 

140 pass 

141 

142 

143__all__ = ["HealthChecker", "HealthReport", "ComponentHealth"]