Coverage for agentos/models/resilience.py: 37%

1""" 

2AgentOS v1.1.5 Resilience — 韧性层。 

3Retry with jitter + Circuit Breaker + Timeout + Fallback chain + Cancellation-aware retry。 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import random 

10import time 

11from dataclasses import dataclass, field 

12from enum import Enum 

13from typing import Callable, Awaitable, TypeVar 

14 

15T = TypeVar("T") 

16 

17 

18class CircuitState(str, Enum): 

19 

20 """熔断器状态枚举。""" 

21 

22 CLOSED = "closed" # 正常 

23 OPEN = "open" # 熔断 

24 HALF_OPEN = "half_open" # 半开（探测） 

25 

26 

27@dataclass 

28class CircuitBreakerConfig: 

29 """熔断器配置。""" 

30 

31 failure_threshold: int = 5 # 连续失败N次后熔断 

32 success_threshold: int = 2 # 半开状态下N次成功后恢复 

33 timeout: float = 60.0 # 熔断持续时间（秒） 

34 half_open_max_requests: int = 1 # 半开状态下最大探测请求 

35 track_duration: float = 300.0 # 统计窗口 

36 

37 

38@dataclass 

39class CircuitBreakerStats: 

40 """熔断器运行统计。""" 

41 

42 state: CircuitState 

43 failure_count: int = 0 

44 success_count: int = 0 

45 last_failure_time: float = 0.0 

46 last_success_time: float = 0.0 

47 opened_at: float = 0.0 

48 total_failures: int = 0 

49 total_successes: int = 0 

50 

51 

52class CircuitBreaker: 

53 """熔断器：检测连续失败，自动熔断/恢复。""" 

54 

55 def __init__(self, name: str, config: CircuitBreakerConfig | None = None): 

56 self.name = name 

57 self.config = config or CircuitBreakerConfig() 

58 self.state = CircuitState.CLOSED 

59 self._failure_count: int = 0 

60 self._success_count: int = 0 

61 self._last_failure_time: float = 0.0 

62 self._opened_at: float = 0.0 

63 self._lock = asyncio.Lock() 

64 self._stats = CircuitBreakerStats(state=CircuitState.CLOSED) 

65 

66 async def call(self, fn: Callable[..., Awaitable[T]], *args, **kwargs) -> T: 

67 """通过熔断器调用函数。""" 

68 async with self._lock: 

69 if not self._allow_request(): 

70 raise CircuitBreakerOpenError(f"Circuit {self.name} is OPEN") 

71 

72 try: 

73 result = await fn(*args, **kwargs) 

74 await self._on_success() 

75 return result 

76 except Exception as e: 

77 await self._on_failure() 

78 raise e 

79 

80 def _allow_request(self) -> bool: 

81 if self.state == CircuitState.CLOSED: 

82 return True 

83 

84 if self.state == CircuitState.OPEN: 

85 if time.time() - self._opened_at >= self.config.timeout: 

86 self.state = CircuitState.HALF_OPEN 

87 self._success_count = 0 

88 return True 

89 return False 

90 

91 if self.state == CircuitState.HALF_OPEN: 

92 return self._success_count < self.config.half_open_max_requests 

93 

94 return True 

95 

96 async def _on_success(self): 

97 async with self._lock: 

98 self._stats.total_successes += 1 

99 self._stats.last_success_time = time.time() 

100 

101 if self.state == CircuitState.HALF_OPEN: 

102 self._success_count += 1 

103 if self._success_count >= self.config.success_threshold: 

104 self.state = CircuitState.CLOSED 

105 self._failure_count = 0 

106 else: 

107 self._failure_count = 0 

108 

109 self._stats.state = self.state 

110 

111 async def _on_failure(self): 

112 async with self._lock: 

113 self._failure_count += 1 

114 self._stats.failure_count = self._failure_count 

115 self._stats.total_failures += 1 

116 self._stats.last_failure_time = time.time() 

117 

118 if self._failure_count >= self.config.failure_threshold: 

119 self.state = CircuitState.OPEN 

120 self._opened_at = time.time() 

121 self._stats.opened_at = self._opened_at 

122 

123 self._stats.state = self.state 

124 

125 @property 

126 def stats(self) -> CircuitBreakerStats: 

127 s = CircuitBreakerStats(state=self.state) 

128 s.failure_count = self._failure_count 

129 s.last_failure_time = self._last_failure_time 

130 s.last_success_time = self._stats.last_success_time 

131 s.opened_at = self._opened_at 

132 s.total_failures = self._stats.total_failures 

133 s.total_successes = self._stats.total_successes 

134 return s 

135 

136 def reset(self): 

137 self.state = CircuitState.CLOSED 

138 self._failure_count = 0 

139 self._success_count = 0 

140 

141 

142class CircuitBreakerOpenError(Exception): 

143 """熔断器打开异常。""" 

144 pass 

145 

146 

147# ── Retry with Jitter ────────────────────────────────────────────────────── 

148 

149@dataclass 

150class RetryConfig: 

151 """重试策略配置。""" 

152 max_retries: int = 3 

153 base_delay: float = 1.0 # 基础延迟（秒） 

154 max_delay: float = 60.0 # 最大延迟 

155 backoff_multiplier: float = 2.0 # 退避乘数 

156 jitter: bool = True # 是否加抖动 

157 jitter_factor: float = 0.1 # 抖动比例 

158 retry_on: tuple[type[Exception], ...] = (Exception,) 

159 

160 

161class CancellationSource(str, Enum): 

162 """取消来源，区分用户主动取消与系统取消。""" 

163 

164 USER = "user" # 用户主动取消 — 不重试 

165 SYSTEM = "system" # 系统级别取消（超时、熔断等）— 按配置重试 

166 

167 

168class CancelledError(Exception): 

169 """带取消来源的取消异常。""" 

170 

171 def __init__(self, message: str, source: CancellationSource = CancellationSource.SYSTEM): 

172 super().__init__(message) 

173 self.source = source 

174 

175 

176class RetryExhaustedError(Exception): 

177 """重试耗尽异常。""" 

178 

179 def __init__(self, attempts: int, last_error: Exception): 

180 super().__init__(f"Retry exhausted after {attempts} attempts. Last error: {last_error}") 

181 self.attempts = attempts 

182 self.last_error = last_error 

183 

184 

185async def retry_with_backoff( 

186 fn: Callable[..., Awaitable[T]], 

187 *args, 

188 config: RetryConfig | None = None, 

189 circuit_breaker: CircuitBreaker | None = None, 

190 on_retry: Callable[[int, Exception, float], None] | None = None, 

191 **kwargs, 

192) -> T: 

193 """带指数退避和抖动的重试函数，区分用户取消（不重试）与系统取消（按配置重试）。""" 

194 cfg = config or RetryConfig() 

195 last_error: Exception | None = None 

196 

197 def _apply_delay(attempt_num: int, err: Exception): 

198 delay = min(cfg.base_delay * (cfg.backoff_multiplier ** attempt_num), cfg.max_delay) 

199 if cfg.jitter: 

200 delay = delay * (1 + random.uniform(-cfg.jitter_factor, cfg.jitter_factor)) 

201 delay = max(0.1, delay) 

202 if on_retry: 

203 on_retry(attempt_num + 1, err, delay) 

204 return delay 

205 

206 for attempt in range(cfg.max_retries + 1): 

207 try: 

208 if circuit_breaker: 

209 return await circuit_breaker.call(fn, *args, **kwargs) 

210 return await fn(*args, **kwargs) 

211 

212 except CircuitBreakerOpenError: 

213 raise # 熔断打开不重试 

214 

215 except CancelledError as e: 

216 if e.source == CancellationSource.USER: 

217 raise # 用户取消不重试，直接上抛 

218 # 系统级取消按重试配置处理 

219 last_error = e 

220 if attempt == cfg.max_retries: 

221 raise RetryExhaustedError(cfg.max_retries + 1, e) 

222 delay = _apply_delay(attempt, e) 

223 await asyncio.sleep(delay) 

224 

225 except cfg.retry_on as e: 

226 last_error = e 

227 if attempt == cfg.max_retries: 

228 raise RetryExhaustedError(cfg.max_retries + 1, e) 

229 delay = _apply_delay(attempt, e) 

230 await asyncio.sleep(delay) 

231 

232 raise RetryExhaustedError(cfg.max_retries, last_error or RuntimeError("unknown")) 

233 

234 

235# ── Timeout ───────────────────────────────────────────────────────────────── 

236 

237class TimeoutError(Exception): 

238 """超时异常。""" 

239 pass 

240 

241 

242async def with_timeout( 

243 fn: Callable[..., Awaitable[T]], 

244 *args, 

245 timeout: float = 120.0, 

246 **kwargs, 

247) -> T: 

248 """为异步函数添加超时保护。""" 

249 try: 

250 return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout) 

251 except asyncio.TimeoutError: 

252 raise TimeoutError(f"Operation timed out after {timeout}s") 

253 

254 

255# ── Fallback Chain ────────────────────────────────────────────────────────── 

256 

257async def with_fallback( 

258 primary: Callable[..., Awaitable[T]], 

259 fallbacks: list[Callable[..., Awaitable[T]]], 

260 *args, **kwargs, 

261) -> T: 

262 """依次尝试主函数和降级函数链。""" 

263 errors: list[Exception] = [] 

264 

265 try: 

266 return await primary(*args, **kwargs) 

267 except Exception as e: 

268 errors.append(e) 

269 

270 for i, fallback in enumerate(fallbacks): 

271 try: 

272 return await fallback(*args, **kwargs) 

273 except Exception as e: 

274 errors.append(e) 

275 

276 raise FallbackExhaustedError(errors) 

277 

278 

279class FallbackExhaustedError(Exception): 

280 

281 """回退耗尽异常。""" 

282 

283 def __init__(self, errors: list[Exception]): 

284 msg = f"All {len(errors)} attempts failed: " + "; ".join(str(e) for e in errors[:3]) 

285 super().__init__(msg) 

286 self.errors = errors 

287 

288 

289# ── Composite Resilience ──────────────────────────────────────────────────── 

290 

291@dataclass 

292class ResilienceConfig: 

293 """弹性总配置。""" 

294 retry: RetryConfig = field(default_factory=RetryConfig) 

295 circuit_breaker: CircuitBreakerConfig | None = None 

296 timeout: float = 120.0 

297 

298 

299class ResilientCall: 

300 """组合韧性调用器：重试 + 熔断 + 超时 + 降级。""" 

301 

302 def __init__(self, config: ResilienceConfig | None = None): 

303 cfg = config or ResilienceConfig() 

304 self.retry_config = cfg.retry 

305 self.timeout = cfg.timeout 

306 self._breaker: CircuitBreaker | None = None 

307 if cfg.circuit_breaker: 

308 self._breaker = CircuitBreaker("default", cfg.circuit_breaker) 

309 

310 async def call(self, fn: Callable[..., Awaitable[T]], *args, **kwargs) -> T: 

311 async def _inner(): 

312 return await with_timeout(fn, *args, timeout=self.timeout, **kwargs) 

313 

314 return await retry_with_backoff( 

315 _inner, 

316 config=self.retry_config, 

317 circuit_breaker=self._breaker, 

318 ) 

319 

320 

321# ── Auto-generated compat stubs ── 

322 

323def retry_with_backoff(*args, **kwargs): pass 

324def with_timeout(*args, **kwargs): pass 

325def with_fallback(*args, **kwargs): pass