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 

17import time 

18from dataclasses import dataclass, field 

19from enum import Enum 

20from typing import Any, Awaitable, Callable, List, Optional, Type, TypeVar 

21 

22logger = logging.getLogger(__name__) 

23 

24T = TypeVar("T") 

25 

26 

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

28# Config 

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

30 

31class JitterStrategy(str, Enum): 

32 """Jitter algorithms for exponential backoff.""" 

33 

34 NONE = "none" # No jitter — deterministic 

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

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

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

38 

39 

40@dataclass 

41class RetryConfig: 

42 """Configuration for retry behavior.""" 

43 

44 max_retries: int = 3 

45 base_delay: float = 1.0 # seconds 

46 max_delay: float = 60.0 # cap 

47 multiplier: float = 2.0 # exponential factor 

48 jitter: JitterStrategy = JitterStrategy.DECORRELATED 

49 retryable_exceptions: tuple[Type[BaseException], ...] = (Exception,) 

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

51 

52 

53# ============================================================================ 

54# Delay calculators 

55# ============================================================================ 

56 

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

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

59 if strategy == JitterStrategy.NONE: 

60 return delay 

61 if strategy == JitterStrategy.FULL: 

62 return random.uniform(0, delay) 

63 if strategy == JitterStrategy.DECORRELATED: 

64 return random.uniform(0, delay) 

65 # EQUAL 

66 half = delay / 2 

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

68 

69 

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

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

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

73 delay = min(delay, config.max_delay) 

74 return _calc_jitter(delay, config.jitter) 

75 

76 

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

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

79 a, b = 0, 1 

80 for _ in range(attempt): 

81 a, b = b, a + b 

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

83 return _calc_jitter(delay, config.jitter) 

84 

85 

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

87 """Constant delay between retries.""" 

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

89 

90 

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

92 

93 

94# ============================================================================ 

95# Retry executor 

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

97 

98@dataclass 

99class RetryResult: 

100 """Result of a retry operation.""" 

101 

102 attempts: int 

103 success: bool 

104 last_exception: Optional[Exception] = None 

105 total_delay: float = 0.0 

106 attempt_history: List[tuple[int, float, Optional[Exception]]] = field(default_factory=list) 

107 

108 

109class Retry: 

110 """Async retry executor with configurable backoff. 

111 

112 Usage: 

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

114 result = await retry.execute(some_async_fn) 

115 

116 # Decorator-style 

117 @retry.with_retry 

118 async def flaky_call(): ... 

119 """ 

120 

121 def __init__( 

122 self, 

123 config: RetryConfig, 

124 delay_fn: DelayFunc = exponential_delay, 

125 ): 

126 self.config = config 

127 self.delay_fn = delay_fn 

128 

129 async def execute( 

130 self, 

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

132 *args: Any, 

133 **kwargs: Any, 

134 ) -> T: 

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

136 last_exc: Optional[Exception] = None 

137 

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

139 try: 

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

141 except self.config.retryable_exceptions as exc: 

142 last_exc = exc 

143 

144 if attempt > self.config.max_retries: 

145 logger.warning( 

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

147 attempt, type(exc).__name__, exc, 

148 ) 

149 raise 

150 

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

152 

153 if self.config.on_retry: 

154 try: 

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

156 except Exception: 

157 pass # Don't let callback failure break retry 

158 

159 logger.debug( 

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

161 attempt, self.config.max_retries, delay, exc, 

162 ) 

163 await asyncio.sleep(delay) 

164 

165 # Unreachable, but type-safe 

166 assert last_exc is not None 

167 raise last_exc 

168 

169 async def execute_with_result( 

170 self, 

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

172 *args: Any, 

173 **kwargs: Any, 

174 ) -> RetryResult: 

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

176 history: List[tuple[int, float, Optional[Exception]]] = [] 

177 total_delay = 0.0 

178 last_exc: Optional[Exception] = None 

179 

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

181 try: 

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

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

184 return RetryResult( 

185 attempts=attempt, 

186 success=True, 

187 total_delay=total_delay, 

188 attempt_history=history, 

189 ) 

190 except self.config.retryable_exceptions as exc: 

191 last_exc = exc 

192 

193 if attempt > self.config.max_retries: 

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

195 return RetryResult( 

196 attempts=attempt, 

197 success=False, 

198 last_exception=exc, 

199 total_delay=total_delay, 

200 attempt_history=history, 

201 ) 

202 

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

204 total_delay += delay 

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

206 await asyncio.sleep(delay) 

207 

208 return RetryResult( 

209 attempts=self.config.max_retries + 1, 

210 success=False, 

211 last_exception=last_exc, 

212 total_delay=total_delay, 

213 attempt_history=history, 

214 ) 

215 

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

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

218 

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

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

221 

222 wrapper.__name__ = fn.__name__ 

223 wrapper.__doc__ = fn.__doc__ 

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

225 return wrapper 

226 

227 

228# ============================================================================ 

229# Pre-built retry policies 

230# ============================================================================ 

231 

232class RetryPolicies: 

233 """Factory for common retry configurations.""" 

234 

235 @staticmethod 

236 def fast() -> Retry: 

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

238 return Retry(RetryConfig( 

239 max_retries=3, 

240 base_delay=0.1, 

241 max_delay=1.0, 

242 multiplier=2.0, 

243 jitter=JitterStrategy.DECORRELATED, 

244 )) 

245 

246 @staticmethod 

247 def standard() -> Retry: 

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

249 return Retry(RetryConfig( 

250 max_retries=5, 

251 base_delay=0.5, 

252 max_delay=30.0, 

253 multiplier=2.0, 

254 jitter=JitterStrategy.DECORRELATED, 

255 )) 

256 

257 @staticmethod 

258 def persistent() -> Retry: 

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

260 return Retry(RetryConfig( 

261 max_retries=10, 

262 base_delay=1.0, 

263 max_delay=120.0, 

264 multiplier=2.0, 

265 jitter=JitterStrategy.DECORRELATED, 

266 )) 

267 

268 @staticmethod 

269 def immediate() -> Retry: 

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

271 return Retry(RetryConfig( 

272 max_retries=2, 

273 base_delay=0.0, 

274 max_delay=0.0, 

275 jitter=JitterStrategy.NONE, 

276 )) 

277 

278 @staticmethod 

279 def gentle() -> Retry: 

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

281 return Retry( 

282 RetryConfig( 

283 max_retries=5, 

284 base_delay=0.5, 

285 max_delay=60.0, 

286 jitter=JitterStrategy.FULL, 

287 ), 

288 delay_fn=fibonacci_delay, 

289 ) 

290 

291 

292# ============================================================================ 

293# Sync retry (for non-async contexts) 

294# ============================================================================ 

295 

296class SyncRetry: 

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

298 

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

300 self._async_retry = Retry(config, delay_fn) 

301 

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

303 """Synchronous wrapper.""" 

304 import asyncio as _asyncio 

305 

306 async def _wrapper(): 

307 return fn(*args, **kwargs) 

308 

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

310 

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

312 """Decorator for sync functions.""" 

313 

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

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

316 

317 wrapper.__name__ = fn.__name__ 

318 wrapper.__doc__ = fn.__doc__ 

319 return wrapper