Coverage for excalidraw_mcp/retry_utils.py: 86%

1"""Retry utilities with exponential backoff and jitter for robust error handling.""" 

2 

3import asyncio 

4import functools 

5import logging 

6import random 

7import time 

8from collections.abc import Awaitable, Callable 

9from typing import Any, TypeVar 

10 

11from .config import config 

12 

13logger = logging.getLogger(__name__) 

14 

15T = TypeVar("T") 

16 

17 

18class RetryConfig: 

19 """Configuration for retry behavior.""" 

20 

21 def __init__( 

22 self, 

23 max_attempts: int = 3, 

24 base_delay: float = 1.0, 

25 max_delay: float = 60.0, 

26 exponential_base: float = 2.0, 

27 jitter: bool = True, 

28 jitter_factor: float = 0.1, 

29 ) -> None: 

30 self.max_attempts = max_attempts 

31 self.base_delay = base_delay 

32 self.max_delay = max_delay 

33 self.exponential_base = exponential_base 

34 self.jitter = jitter 

35 self.jitter_factor = jitter_factor 

36 

37 

38def calculate_delay( 

39 attempt: int, config: RetryConfig, exception: Exception | None = None 

40) -> float: 

41 """Calculate delay with exponential backoff and optional jitter.""" 

42 # Exponential backoff: base_delay * (exponential_base ^ attempt) 

43 delay = config.base_delay * (config.exponential_base**attempt) 

44 

45 # Cap at max_delay 

46 delay = min(delay, config.max_delay) 

47 

48 # Add jitter if enabled 

49 if config.jitter: 

50 jitter_amount = delay * config.jitter_factor 

51 delay += random.uniform(-jitter_amount, jitter_amount) 

52 # Ensure delay doesn't go below base_delay 

53 delay = max(delay, config.base_delay) 

54 

55 return delay 

56 

57 

58async def retry_async[T]( 

59 func: Callable[..., Awaitable[T]], 

60 retry_config: RetryConfig | None = None, 

61 retry_on_exceptions: tuple[type[Exception], ...] | None = None, 

62 on_retry: Callable[[int, Exception], Awaitable[None]] | None = None, 

63) -> T: 

64 """Execute an async function with retry logic. 

65 

66 Args: 

67 func: The async function to execute 

68 retry_config: Configuration for retry behavior 

69 retry_on_exceptions: Tuple of exception types to retry on 

70 on_retry: Optional callback function called on each retry 

71 

72 Returns: 

73 The result of the function call 

74 

75 Raises: 

76 The last exception encountered after all retries are exhausted 

77 """ 

78 if retry_config is None: 

79 retry_config = RetryConfig( 

80 max_attempts=config.server.sync_retry_attempts, 

81 base_delay=config.server.sync_retry_delay_seconds, 

82 ) 

83 

84 if retry_on_exceptions is None: 

85 retry_on_exceptions = (Exception,) 

86 

87 last_exception: Exception | None = None 

88 

89 for attempt in range(retry_config.max_attempts): 

90 try: 

91 return await func() 

92 except retry_on_exceptions as e: 

93 last_exception = e 

94 

95 # If this is the last attempt, don't retry 

96 if attempt == retry_config.max_attempts - 1: 

97 logger.error( 

98 f"Function {func.__name__} failed after {retry_config.max_attempts} attempts. " 

99 f"Last error: {e}" 

100 ) 

101 raise 

102 

103 # Calculate delay 

104 delay = calculate_delay(attempt, retry_config, e) 

105 

106 logger.warning( 

107 f"Attempt {attempt + 1} of {func.__name__} failed: {e}. " 

108 f"Retrying in {delay:.2f} seconds..." 

109 ) 

110 

111 # Call retry callback if provided 

112 if on_retry: 

113 try: 

114 await on_retry(attempt + 1, e) 

115 except Exception as callback_error: 

116 logger.warning(f"Error in retry callback: {callback_error}") 

117 

118 # Wait before retrying 

119 await asyncio.sleep(delay) 

120 

121 # This should never be reached due to the raise in the loop, 

122 # but included for type safety 

123 if last_exception: 

124 raise last_exception 

125 raise RuntimeError("Retry loop ended without an exception") 

126 

127 

128def retry_sync[T]( 

129 func: Callable[..., T], 

130 retry_config: RetryConfig | None = None, 

131 retry_on_exceptions: tuple[type[Exception], ...] | None = None, 

132 on_retry: Callable[[int, Exception], None] | None = None, 

133) -> T: 

134 """Execute a sync function with retry logic. 

135 

136 Args: 

137 func: The sync function to execute 

138 retry_config: Configuration for retry behavior 

139 retry_on_exceptions: Tuple of exception types to retry on 

140 on_retry: Optional callback function called on each retry 

141 

142 Returns: 

143 The result of the function call 

144 

145 Raises: 

146 The last exception encountered after all retries are exhausted 

147 """ 

148 if retry_config is None: 

149 retry_config = RetryConfig( 

150 max_attempts=config.server.sync_retry_attempts, 

151 base_delay=config.server.sync_retry_delay_seconds, 

152 ) 

153 

154 if retry_on_exceptions is None: 

155 retry_on_exceptions = (Exception,) 

156 

157 last_exception: Exception | None = None 

158 

159 for attempt in range(retry_config.max_attempts): 

160 try: 

161 return func() 

162 except retry_on_exceptions as e: 

163 last_exception = e 

164 

165 # If this is the last attempt, don't retry 

166 if attempt == retry_config.max_attempts - 1: 

167 logger.error( 

168 f"Function {func.__name__} failed after {retry_config.max_attempts} attempts. " 

169 f"Last error: {e}" 

170 ) 

171 raise 

172 

173 # Calculate delay 

174 delay = calculate_delay(attempt, retry_config, e) 

175 

176 logger.warning( 

177 f"Attempt {attempt + 1} of {func.__name__} failed: {e}. " 

178 f"Retrying in {delay:.2f} seconds..." 

179 ) 

180 

181 # Call retry callback if provided 

182 if on_retry: 

183 try: 

184 on_retry(attempt + 1, e) 

185 except Exception as callback_error: 

186 logger.warning(f"Error in retry callback: {callback_error}") 

187 

188 # Wait before retrying 

189 time.sleep(delay) 

190 

191 # This should never be reached due to the raise in the loop, 

192 # but included for type safety 

193 if last_exception: 

194 raise last_exception 

195 raise RuntimeError("Retry loop ended without an exception") 

196 

197 

198def retry_decorator( 

199 retry_config: RetryConfig | None = None, 

200 retry_on_exceptions: tuple[type[Exception], ...] | None = None, 

201 on_retry: Callable[[int, Exception], Any] | None = None, 

202) -> Callable[[Callable[..., T]], Callable[..., T]]: 

203 """Decorator for adding retry logic to functions. 

204 

205 Args: 

206 retry_config: Configuration for retry behavior 

207 retry_on_exceptions: Tuple of exception types to retry on 

208 on_retry: Optional callback function called on each retry 

209 

210 Returns: 

211 Decorated function with retry logic 

212 """ 

213 

214 def decorator(func: Callable[..., T]) -> Callable[..., T]: 

215 @functools.wraps(func) 

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

217 def sync_func() -> T: 

218 return func(*args, **kwargs) 

219 

220 return retry_sync( 

221 sync_func, 

222 retry_config=retry_config, 

223 retry_on_exceptions=retry_on_exceptions, 

224 on_retry=on_retry, 

225 ) 

226 

227 @functools.wraps(func) 

228 async def async_wrapper(*args: Any, **kwargs: Any) -> T: 

229 async def async_func() -> T: 

230 result = func(*args, **kwargs) 

231 if asyncio.iscoroutine(result): 

232 coro_result: T = await result 

233 return coro_result 

234 return result 

235 

236 return await retry_async( 

237 async_func, 

238 retry_config=retry_config, 

239 retry_on_exceptions=retry_on_exceptions, 

240 on_retry=on_retry, 

241 ) 

242 

243 # Check if function is async 

244 if asyncio.iscoroutinefunction(func): 

245 return async_wrapper # type: ignore 

246 return wrapper 

247 

248 return decorator