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

1"""AgentOS Retry — production-grade retry with exponential backoff + jitter. 

2 

3Strategies: 

4- ExponentialBackoff: base * multiplier^n with decorrelated jitter 

5- FixedDelay: constant interval 

6- FibonacciBackoff: fib sequence for gentler ramp 

7- CompositeRetry: chain multiple strategies 

8 

9Design: ~350 lines, zero external deps beyond stdlib + asyncio. 

10""" 

11 

12from __future__ import annotations 

13 

14import asyncio 

15import logging 

16import random 

17from collections.abc import Awaitable, Callable 

18from dataclasses import dataclass, field 

19from enum import StrEnum 

20from typing import Any, TypeVar 

21 

22logger = logging.getLogger(__name__) 

23 

24T = TypeVar("T") 

25 

26 

27# ============================================================================ 

28# Config 

29# ============================================================================ 

30 

31 

32class JitterStrategy(StrEnum): 

33 """Jitter algorithms for exponential backoff.""" 

34 

35 NONE = "none" # No jitter — deterministic 

36 FULL = "full" # random(0, delay) 

37 DECORRELATED = "decorrelated" # random(base, delay) — AWS-style 

38 EQUAL = "equal" # delay/2 + random(0, delay/2) 

39 

40 

41@dataclass 

42class RetryConfig: 

43 """Configuration for retry behavior.""" 

44 

45 max_retries: int = 3 

46 base_delay: float = 1.0 # seconds 

47 max_delay: float = 60.0 # cap 

48 multiplier: float = 2.0 # exponential factor 

49 jitter: JitterStrategy = JitterStrategy.DECORRELATED 

50 retryable_exceptions: tuple[type[BaseException], ...] = (Exception,) 

51 on_retry: Callable[[int, Exception, float], None] | None = None # (attempt, exc, delay) → None 

52 

53 

54# ============================================================================ 

55# Delay calculators 

56# ============================================================================ 

57 

58 

59def _calc_jitter(delay: float, strategy: JitterStrategy) -> float: 

60 """Apply jitter to a computed delay.""" 

61 if strategy == JitterStrategy.NONE: 

62 return delay 

63 if strategy == JitterStrategy.FULL: 

64 return random.uniform(0, delay) 

65 if strategy == JitterStrategy.DECORRELATED: 

66 return random.uniform(0, delay) 

67 # EQUAL 

68 half = delay / 2 

69 return half + random.uniform(0, half) 

70 

71 

72def exponential_delay(attempt: int, config: RetryConfig) -> float: 

73 """Exponential backoff: base * multiplier^(attempt-1).""" 

74 delay = config.base_delay * (config.multiplier ** (attempt - 1)) 

75 delay = min(delay, config.max_delay) 

76 return _calc_jitter(delay, config.jitter) 

77 

78 

79def fibonacci_delay(attempt: int, config: RetryConfig) -> float: 

80 """Fibonacci backoff — gentler than exponential.""" 

81 a, b = 0, 1 

82 for _ in range(attempt): 

83 a, b = b, a + b 

84 delay = min(a * config.base_delay, config.max_delay) 

85 return _calc_jitter(delay, config.jitter) 

86 

87 

88def fixed_delay(attempt: int, config: RetryConfig) -> float: 

89 """Constant delay between retries.""" 

90 return min(config.base_delay, config.max_delay) 

91 

92 

93DelayFunc = Callable[[int, RetryConfig], float] 

94 

95 

96# ============================================================================ 

97# Retry executor 

98# ============================================================================ 

99 

100 

101@dataclass 

102class RetryResult: 

103 """Result of a retry operation.""" 

104 

105 attempts: int 

106 success: bool 

107 last_exception: Exception | None = None 

108 total_delay: float = 0.0 

109 attempt_history: list[tuple[int, float, Exception | None]] = field(default_factory=list) 

110 

111 

112class Retry: 

113 """Async retry executor with configurable backoff. 

114 

115 Usage: 

116 retry = Retry(RetryConfig(max_retries=3, base_delay=0.5)) 

117 result = await retry.execute(some_async_fn) 

118 

119 # Decorator-style 

120 @retry.with_retry 

121 async def flaky_call(): ... 

122 """ 

123 

124 def __init__( 

125 self, 

126 config: RetryConfig, 

127 delay_fn: DelayFunc = exponential_delay, 

128 ): 

129 self.config = config 

130 self.delay_fn = delay_fn 

131 

132 async def execute( 

133 self, 

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

135 *args: Any, 

136 **kwargs: Any, 

137 ) -> T: 

138 """Execute fn with retry logic. Raises last exception if all retries exhausted.""" 

139 last_exc: Exception | None = None 

140 

141 for attempt in range(1, self.config.max_retries + 2): # +2 for initial attempt + retries 

142 try: 

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

144 except self.config.retryable_exceptions as exc: 

145 last_exc = exc 

146 

147 if attempt > self.config.max_retries: 

148 logger.warning( 

149 "Retry exhausted after %d attempts — %s: %s", 

150 attempt, 

151 type(exc).__name__, 

152 exc, 

153 ) 

154 raise 

155 

156 delay = self.delay_fn(attempt, self.config) 

157 

158 if self.config.on_retry: 

159 try: 

160 self.config.on_retry(attempt, exc, delay) 

161 except Exception: 

162 pass # Don't let callback failure break retry 

163 

164 logger.debug( 

165 "Retry attempt %d/%d after %.2fs — %s", 

166 attempt, 

167 self.config.max_retries, 

168 delay, 

169 exc, 

170 ) 

171 await asyncio.sleep(delay) 

172 

173 # Unreachable, but type-safe 

174 assert last_exc is not None 

175 raise last_exc 

176 

177 async def execute_with_result( 

178 self, 

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

180 *args: Any, 

181 **kwargs: Any, 

182 ) -> RetryResult: 

183 """Execute with retry and return detailed Result.""" 

184 history: list[tuple[int, float, Exception | None]] = [] 

185 total_delay = 0.0 

186 last_exc: Exception | None = None 

187 

188 for attempt in range(1, self.config.max_retries + 2): 

189 try: 

190 await fn(*args, **kwargs) 

191 history.append((attempt, 0.0, None)) 

192 return RetryResult( 

193 attempts=attempt, 

194 success=True, 

195 total_delay=total_delay, 

196 attempt_history=history, 

197 ) 

198 except self.config.retryable_exceptions as exc: 

199 last_exc = exc 

200 

201 if attempt > self.config.max_retries: 

202 history.append((attempt, 0.0, exc)) 

203 return RetryResult( 

204 attempts=attempt, 

205 success=False, 

206 last_exception=exc, 

207 total_delay=total_delay, 

208 attempt_history=history, 

209 ) 

210 

211 delay = self.delay_fn(attempt, self.config) 

212 total_delay += delay 

213 history.append((attempt, delay, exc)) 

214 await asyncio.sleep(delay) 

215 

216 return RetryResult( 

217 attempts=self.config.max_retries + 1, 

218 success=False, 

219 last_exception=last_exc, 

220 total_delay=total_delay, 

221 attempt_history=history, 

222 ) 

223 

224 def with_retry(self, fn: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]: 

225 """Decorator: apply retry logic to an async function.""" 

226 

227 async def wrapper(*args: Any, **kwargs: Any) -> T: 

228 return await self.execute(fn, *args, **kwargs) 

229 

230 wrapper.__name__ = fn.__name__ 

231 wrapper.__doc__ = fn.__doc__ 

232 wrapper._retry = self # type: ignore[attr-defined] 

233 return wrapper 

234 

235 

236# ============================================================================ 

237# Pre-built retry policies 

238# ============================================================================ 

239 

240 

241class RetryPolicies: 

242 """Factory for common retry configurations.""" 

243 

244 @staticmethod 

245 def fast() -> Retry: 

246 """3 retries, 100ms base, 2x multiplier — for idempotent API calls.""" 

247 return Retry( 

248 RetryConfig( 

249 max_retries=3, 

250 base_delay=0.1, 

251 max_delay=1.0, 

252 multiplier=2.0, 

253 jitter=JitterStrategy.DECORRELATED, 

254 ) 

255 ) 

256 

257 @staticmethod 

258 def standard() -> Retry: 

259 """5 retries, 500ms base, 2x — general purpose.""" 

260 return Retry( 

261 RetryConfig( 

262 max_retries=5, 

263 base_delay=0.5, 

264 max_delay=30.0, 

265 multiplier=2.0, 

266 jitter=JitterStrategy.DECORRELATED, 

267 ) 

268 ) 

269 

270 @staticmethod 

271 def persistent() -> Retry: 

272 """10 retries, 1s base, 2x, max 120s — for critical operations.""" 

273 return Retry( 

274 RetryConfig( 

275 max_retries=10, 

276 base_delay=1.0, 

277 max_delay=120.0, 

278 multiplier=2.0, 

279 jitter=JitterStrategy.DECORRELATED, 

280 ) 

281 ) 

282 

283 @staticmethod 

284 def immediate() -> Retry: 

285 """2 retries, no delay — for infallible operations.""" 

286 return Retry( 

287 RetryConfig( 

288 max_retries=2, 

289 base_delay=0.0, 

290 max_delay=0.0, 

291 jitter=JitterStrategy.NONE, 

292 ) 

293 ) 

294 

295 @staticmethod 

296 def gentle() -> Retry: 

297 """5 retries, Fibonacci backoff — rate-limit friendly.""" 

298 return Retry( 

299 RetryConfig( 

300 max_retries=5, 

301 base_delay=0.5, 

302 max_delay=60.0, 

303 jitter=JitterStrategy.FULL, 

304 ), 

305 delay_fn=fibonacci_delay, 

306 ) 

307 

308 

309# ============================================================================ 

310# Sync retry (for non-async contexts) 

311# ============================================================================ 

312 

313 

314class SyncRetry: 

315 """Synchronous retry executor — calls asyncio.run internally.""" 

316 

317 def __init__(self, config: RetryConfig, delay_fn: DelayFunc = exponential_delay): 

318 self._async_retry = Retry(config, delay_fn) 

319 

320 def execute(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: 

321 """Synchronous wrapper.""" 

322 import asyncio as _asyncio 

323 

324 async def _wrapper(): 

325 return fn(*args, **kwargs) 

326 

327 return _asyncio.run(self._async_retry.execute(_wrapper)) 

328 

329 def with_retry(self, fn: Callable[..., T]) -> Callable[..., T]: 

330 """Decorator for sync functions.""" 

331 

332 def wrapper(*args: Any, **kwargs: Any) -> T: 

333 return self.execute(fn, *args, **kwargs) 

334 

335 wrapper.__name__ = fn.__name__ 

336 wrapper.__doc__ = fn.__doc__ 

337 return wrapper