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 

21from typing import Any, Dict, List, Optional 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26@dataclass 

27class ComponentHealth: 

28 name: str 

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

30 latency_ms: float 

31 error: Optional[str] = None 

32 

33 

34@dataclass 

35class HealthReport: 

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

37 uptime_seconds: float 

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

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

40 

41 

42class HealthChecker: 

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

44 

45 def __init__( 

46 self, 

47 start_time: float, 

48 db_url: Optional[str] = None, 

49 redis_url: Optional[str] = None, 

50 ): 

51 self.start_time = start_time 

52 self.db_url = db_url 

53 self.redis_url = redis_url 

54 

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

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

57 t0 = time.perf_counter() 

58 try: 

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

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

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

62 except asyncio.TimeoutError: 

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

64 return ComponentHealth( 

65 name=name, status="unhealthy", latency_ms=latency, 

66 error=f"Timeout after {timeout}s" 

67 ) 

68 except Exception as e: 

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

70 return ComponentHealth( 

71 name=name, status="unhealthy", latency_ms=latency, 

72 error=str(e) 

73 ) 

74 

75 async def check(self) -> HealthReport: 

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

77 probes = [] 

78 

79 # DB probe 

80 if self.db_url: 

81 probes.append( 

82 self._probe("database", self._check_db) 

83 ) 

84 

85 # Redis probe 

86 if self.redis_url: 

87 probes.append( 

88 self._probe("redis", self._check_redis) 

89 ) 

90 

91 # Always probe disk (write test) 

92 probes.append( 

93 self._probe("disk", self._check_disk) 

94 ) 

95 

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

97 

98 components: Dict[str, ComponentHealth] = {} 

99 overall = "healthy" 

100 

101 for r in results: 

102 if isinstance(r, ComponentHealth): 

103 components[r.name] = r 

104 if r.status == "unhealthy": 

105 if overall == "healthy": 

106 overall = "degraded" 

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

108 overall = "degraded" 

109 elif isinstance(r, Exception): 

110 # Probe itself crashed 

111 components["internal"] = ComponentHealth( 

112 name="internal", status="unhealthy", 

113 latency_ms=0, error=str(r) 

114 ) 

115 overall = "unhealthy" 

116 

117 return HealthReport( 

118 status=overall, 

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

120 components=components, 

121 ) 

122 

123 async def _check_db(self): 

124 """Database connectivity probe.""" 

125 from sqlalchemy.ext.asyncio import create_async_engine 

126 from sqlalchemy import text 

127 

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

129 async with engine.connect() as conn: 

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

131 await engine.dispose() 

132 

133 async def _check_redis(self): 

134 """Redis connectivity probe.""" 

135 import redis.asyncio as redis 

136 

137 r = redis.from_url(self.redis_url) 

138 await r.ping() 

139 await r.close() 

140 

141 async def _check_disk(self): 

142 """Filesystem write test.""" 

143 import tempfile 

144 import os 

145 

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

147 f.write(b"ok") 

148 

149 try: 

150 os.unlink(f.name) 

151 except OSError: 

152 pass 

153 

154 

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