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

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 dataclasses import dataclass, field 

18from enum import Enum 

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

20 

21logger = logging.getLogger(__name__) 

22 

23T = TypeVar("T") 

24 

25 

26# ============================================================================ 

27# Config 

28# ============================================================================ 

29 

30class JitterStrategy(str, Enum): 

31 """Jitter algorithms for exponential backoff.""" 

32 

33 NONE = "none" # No jitter — deterministic 

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

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

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

37 

38 

39@dataclass 

40class RetryConfig: 

41 """Configuration for retry behavior.""" 

42 

43 max_retries: int = 3 

44 base_delay: float = 1.0 # seconds 

45 max_delay: float = 60.0 # cap 

46 multiplier: float = 2.0 # exponential factor 

47 jitter: JitterStrategy = JitterStrategy.DECORRELATED 

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

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

50 

51 

52# ============================================================================ 

53# Delay calculators 

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

55 

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

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

58 if strategy == JitterStrategy.NONE: 

59 return delay 

60 if strategy == JitterStrategy.FULL: 

61 return random.uniform(0, delay) 

62 if strategy == JitterStrategy.DECORRELATED: 

63 return random.uniform(0, delay) 

64 # EQUAL 

65 half = delay / 2 

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

67 

68 

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

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

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

72 delay = min(delay, config.max_delay) 

73 return _calc_jitter(delay, config.jitter) 

74 

75 

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

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

78 a, b = 0, 1 

79 for _ in range(attempt): 

80 a, b = b, a + b 

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

82 return _calc_jitter(delay, config.jitter) 

83 

84 

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

86 """Constant delay between retries.""" 

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

88 

89 

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

91 

92 

93# ============================================================================ 

94# Retry executor 

95# ============================================================================ 

96 

97@dataclass 

98class RetryResult: 

99 """Result of a retry operation.""" 

100 

101 attempts: int 

102 success: bool 

103 last_exception: Optional[Exception] = None 

104 total_delay: float = 0.0 

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

106 

107 

108class Retry: 

109 """Async retry executor with configurable backoff. 

110 

111 Usage: 

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

113 result = await retry.execute(some_async_fn) 

114 

115 # Decorator-style 

116 @retry.with_retry 

117 async def flaky_call(): ... 

118 """ 

119 

120 def __init__( 

121 self, 

122 config: RetryConfig, 

123 delay_fn: DelayFunc = exponential_delay, 

124 ): 

125 self.config = config 

126 self.delay_fn = delay_fn 

127 

128 async def execute( 

129 self, 

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

131 *args: Any, 

132 **kwargs: Any, 

133 ) -> T: 

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

135 last_exc: Optional[Exception] = None 

136 

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

138 try: 

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

140 except self.config.retryable_exceptions as exc: 

141 last_exc = exc 

142 

143 if attempt > self.config.max_retries: 

144 logger.warning( 

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

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

147 ) 

148 raise 

149 

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

151 

152 if self.config.on_retry: 

153 try: 

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

155 except Exception: 

156 pass # Don't let callback failure break retry 

157 

158 logger.debug( 

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

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

161 ) 

162 await asyncio.sleep(delay) 

163 

164 # Unreachable, but type-safe 

165 assert last_exc is not None 

166 raise last_exc 

167 

168 async def execute_with_result( 

169 self, 

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

171 *args: Any, 

172 **kwargs: Any, 

173 ) -> RetryResult: 

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

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

176 total_delay = 0.0 

177 last_exc: Optional[Exception] = None 

178 

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

180 try: 

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

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

183 return RetryResult( 

184 attempts=attempt, 

185 success=True, 

186 total_delay=total_delay, 

187 attempt_history=history, 

188 ) 

189 except self.config.retryable_exceptions as exc: 

190 last_exc = exc 

191 

192 if attempt > self.config.max_retries: 

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

194 return RetryResult( 

195 attempts=attempt, 

196 success=False, 

197 last_exception=exc, 

198 total_delay=total_delay, 

199 attempt_history=history, 

200 ) 

201 

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

203 total_delay += delay 

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

205 await asyncio.sleep(delay) 

206 

207 return RetryResult( 

208 attempts=self.config.max_retries + 1, 

209 success=False, 

210 last_exception=last_exc, 

211 total_delay=total_delay, 

212 attempt_history=history, 

213 ) 

214 

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

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

217 

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

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

220 

221 wrapper.__name__ = fn.__name__ 

222 wrapper.__doc__ = fn.__doc__ 

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

224 return wrapper 

225 

226 

227# ============================================================================ 

228# Pre-built retry policies 

229# ============================================================================ 

230 

231class RetryPolicies: 

232 """Factory for common retry configurations.""" 

233 

234 @staticmethod 

235 def fast() -> Retry: 

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

237 return Retry(RetryConfig( 

238 max_retries=3, 

239 base_delay=0.1, 

240 max_delay=1.0, 

241 multiplier=2.0, 

242 jitter=JitterStrategy.DECORRELATED, 

243 )) 

244 

245 @staticmethod 

246 def standard() -> Retry: 

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

248 return Retry(RetryConfig( 

249 max_retries=5, 

250 base_delay=0.5, 

251 max_delay=30.0, 

252 multiplier=2.0, 

253 jitter=JitterStrategy.DECORRELATED, 

254 )) 

255 

256 @staticmethod 

257 def persistent() -> Retry: 

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

259 return Retry(RetryConfig( 

260 max_retries=10, 

261 base_delay=1.0, 

262 max_delay=120.0, 

263 multiplier=2.0, 

264 jitter=JitterStrategy.DECORRELATED, 

265 )) 

266 

267 @staticmethod 

268 def immediate() -> Retry: 

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

270 return Retry(RetryConfig( 

271 max_retries=2, 

272 base_delay=0.0, 

273 max_delay=0.0, 

274 jitter=JitterStrategy.NONE, 

275 )) 

276 

277 @staticmethod 

278 def gentle() -> Retry: 

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

280 return Retry( 

281 RetryConfig( 

282 max_retries=5, 

283 base_delay=0.5, 

284 max_delay=60.0, 

285 jitter=JitterStrategy.FULL, 

286 ), 

287 delay_fn=fibonacci_delay, 

288 ) 

289 

290 

291# ============================================================================ 

292# Sync retry (for non-async contexts) 

293# ============================================================================ 

294 

295class SyncRetry: 

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

297 

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

299 self._async_retry = Retry(config, delay_fn) 

300 

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

302 """Synchronous wrapper.""" 

303 import asyncio as _asyncio 

304 

305 async def _wrapper(): 

306 return fn(*args, **kwargs) 

307 

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

309 

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

311 """Decorator for sync functions.""" 

312 

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

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

315 

316 wrapper.__name__ = fn.__name__ 

317 wrapper.__doc__ = fn.__doc__ 

318 return wrapper