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 builtins 

10import random 

11import time 

12from collections.abc import Awaitable, Callable 

13from dataclasses import dataclass, field 

14from enum import StrEnum 

15from typing import TypeVar 

16 

17T = TypeVar("T") 

18 

19 

20class CircuitState(StrEnum): 

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

22 

23 CLOSED = "closed" # 正常 

24 OPEN = "open" # 熔断 

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

26 

27 

28@dataclass 

29class CircuitBreakerConfig: 

30 """熔断器配置。""" 

31 

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

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

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

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

36 track_duration: float = 300.0 # 统计窗口 

37 

38 

39@dataclass 

40class CircuitBreakerStats: 

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

42 

43 state: CircuitState 

44 failure_count: int = 0 

45 success_count: int = 0 

46 last_failure_time: float = 0.0 

47 last_success_time: float = 0.0 

48 opened_at: float = 0.0 

49 total_failures: int = 0 

50 total_successes: int = 0 

51 

52 

53class CircuitBreaker: 

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

55 

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

57 self.name = name 

58 self.config = config or CircuitBreakerConfig() 

59 self.state = CircuitState.CLOSED 

60 self._failure_count: int = 0 

61 self._success_count: int = 0 

62 self._last_failure_time: float = 0.0 

63 self._opened_at: float = 0.0 

64 self._lock = asyncio.Lock() 

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

66 

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

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

69 async with self._lock: 

70 if not self._allow_request(): 

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

72 

73 try: 

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

75 await self._on_success() 

76 return result 

77 except Exception as e: 

78 await self._on_failure() 

79 raise e 

80 

81 def _allow_request(self) -> bool: 

82 if self.state == CircuitState.CLOSED: 

83 return True 

84 

85 if self.state == CircuitState.OPEN: 

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

87 self.state = CircuitState.HALF_OPEN 

88 self._success_count = 0 

89 return True 

90 return False 

91 

92 if self.state == CircuitState.HALF_OPEN: 

93 return self._success_count < self.config.half_open_max_requests 

94 

95 return True 

96 

97 async def _on_success(self): 

98 async with self._lock: 

99 self._stats.total_successes += 1 

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

101 

102 if self.state == CircuitState.HALF_OPEN: 

103 self._success_count += 1 

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

105 self.state = CircuitState.CLOSED 

106 self._failure_count = 0 

107 else: 

108 self._failure_count = 0 

109 

110 self._stats.state = self.state 

111 

112 async def _on_failure(self): 

113 async with self._lock: 

114 self._failure_count += 1 

115 self._stats.failure_count = self._failure_count 

116 self._stats.total_failures += 1 

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

118 

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

120 self.state = CircuitState.OPEN 

121 self._opened_at = time.time() 

122 self._stats.opened_at = self._opened_at 

123 

124 self._stats.state = self.state 

125 

126 @property 

127 def stats(self) -> CircuitBreakerStats: 

128 s = CircuitBreakerStats(state=self.state) 

129 s.failure_count = self._failure_count 

130 s.last_failure_time = self._last_failure_time 

131 s.last_success_time = self._stats.last_success_time 

132 s.opened_at = self._opened_at 

133 s.total_failures = self._stats.total_failures 

134 s.total_successes = self._stats.total_successes 

135 return s 

136 

137 def reset(self): 

138 self.state = CircuitState.CLOSED 

139 self._failure_count = 0 

140 self._success_count = 0 

141 

142 

143class CircuitBreakerOpenError(Exception): 

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

145 

146 

147 

148# ── Retry with Jitter ────────────────────────────────────────────────────── 

149 

150 

151@dataclass 

152class RetryConfig: 

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

154 

155 max_retries: int = 3 

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

157 max_delay: float = 60.0 # 最大延迟 

158 backoff_multiplier: float = 2.0 # 退避乘数 

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

160 jitter_factor: float = 0.1 # 抖动比例 

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

162 

163 

164class CancellationSource(StrEnum): 

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

166 

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

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

169 

170 

171class CancelledError(Exception): 

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

173 

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

175 super().__init__(message) 

176 self.source = source 

177 

178 

179class RetryExhaustedError(Exception): 

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

181 

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

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

184 self.attempts = attempts 

185 self.last_error = last_error 

186 

187 

188async def retry_with_backoff( 

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

190 *args, 

191 config: RetryConfig | None = None, 

192 circuit_breaker: CircuitBreaker | None = None, 

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

194 **kwargs, 

195) -> T: 

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

197 cfg = config or RetryConfig() 

198 last_error: Exception | None = None 

199 

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

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

202 if cfg.jitter: 

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

204 delay = max(0.1, delay) 

205 if on_retry: 

206 on_retry(attempt_num + 1, err, delay) 

207 return delay 

208 

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

210 try: 

211 if circuit_breaker: 

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

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

214 

215 except CircuitBreakerOpenError: 

216 raise # 熔断打开不重试 

217 

218 except CancelledError as e: 

219 if e.source == CancellationSource.USER: 

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

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

222 last_error = e 

223 if attempt == cfg.max_retries: 

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

225 delay = _apply_delay(attempt, e) 

226 await asyncio.sleep(delay) 

227 

228 except cfg.retry_on as e: 

229 last_error = e 

230 if attempt == cfg.max_retries: 

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

232 delay = _apply_delay(attempt, e) 

233 await asyncio.sleep(delay) 

234 

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

236 

237 

238# ── Timeout ───────────────────────────────────────────────────────────────── 

239 

240 

241class TimeoutError(Exception): 

242 """超时异常。""" 

243 

244 

245 

246async def with_timeout( 

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

248 *args, 

249 timeout: float = 120.0, 

250 **kwargs, 

251) -> T: 

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

253 try: 

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

255 except builtins.TimeoutError: 

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

257 

258 

259# ── Fallback Chain ────────────────────────────────────────────────────────── 

260 

261 

262async def with_fallback( 

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

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

265 *args, 

266 **kwargs, 

267) -> T: 

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

269 errors: list[Exception] = [] 

270 

271 try: 

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

273 except Exception as e: 

274 errors.append(e) 

275 

276 for i, fallback in enumerate(fallbacks): 

277 try: 

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

279 except Exception as e: 

280 errors.append(e) 

281 

282 raise FallbackExhaustedError(errors) 

283 

284 

285class FallbackExhaustedError(Exception): 

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

287 

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

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

290 super().__init__(msg) 

291 self.errors = errors 

292 

293 

294# ── Composite Resilience ──────────────────────────────────────────────────── 

295 

296 

297@dataclass 

298class ResilienceConfig: 

299 """弹性总配置。""" 

300 

301 retry: RetryConfig = field(default_factory=RetryConfig) 

302 circuit_breaker: CircuitBreakerConfig | None = None 

303 timeout: float = 120.0 

304 

305 

306class ResilientCall: 

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

308 

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

310 cfg = config or ResilienceConfig() 

311 self.retry_config = cfg.retry 

312 self.timeout = cfg.timeout 

313 self._breaker: CircuitBreaker | None = None 

314 if cfg.circuit_breaker: 

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

316 

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

318 async def _inner(): 

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

320 

321 return await retry_with_backoff( 

322 _inner, 

323 config=self.retry_config, 

324 circuit_breaker=self._breaker, 

325 ) 

326 

327 

328# ── Auto-generated compat stubs ── 

329 

330 

331def retry_with_backoff(*args, **kwargs): # noqa: F811 

332 pass 

333 

334 

335def with_timeout(*args, **kwargs): # noqa: F811 

336 pass 

337 

338 

339def with_fallback(*args, **kwargs): # noqa: F811 

340 pass