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

1"""AgentOS Rate Limiter — multi-strategy throttling for production APIs. 

2 

3Strategies: 

4- TokenBucket: classic token refill, burst-tolerant 

5- SlidingWindow: precise time-window counting 

6- ConcurrentLimiter: limit in-flight requests (semaphore-based) 

7- CompositeLimiter: chain multiple limiters together 

8 

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

10""" 

11 

12from __future__ import annotations 

13 

14import asyncio 

15import logging 

16import time 

17from dataclasses import dataclass 

18from enum import Enum 

19from typing import Awaitable, Callable, Dict, Optional, TypeVar 

20 

21logger = logging.getLogger(__name__) 

22 

23T = TypeVar("T") 

24 

25 

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

27# Core limiters 

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

29 

30class TokenBucket: 

31 """Classic token bucket for rate limiting with burst support. 

32 

33 Tokens refill at a fixed rate up to capacity. Each request costs 1 token. 

34 """ 

35 

36 def __init__(self, rate: float, capacity: int): 

37 """ 

38 Args: 

39 rate: Tokens per second refill rate 

40 capacity: Maximum tokens (burst size) 

41 """ 

42 if rate <= 0: 

43 raise ValueError("rate must be > 0") 

44 if capacity <= 0: 

45 raise ValueError("capacity must be > 0") 

46 

47 self.rate = rate 

48 self.capacity = capacity 

49 self._tokens = float(capacity) 

50 self._last_refill = time.monotonic() 

51 self._lock = asyncio.Lock() 

52 

53 def _refill(self) -> None: 

54 now = time.monotonic() 

55 elapsed = now - self._last_refill 

56 self._tokens = min(self.capacity, self._tokens + elapsed * self.rate) 

57 self._last_refill = now 

58 

59 async def acquire(self, tokens: int = 1) -> bool: 

60 """Try to acquire tokens. Returns True if successful, False otherwise.""" 

61 async with self._lock: 

62 self._refill() 

63 if self._tokens >= tokens: 

64 self._tokens -= tokens 

65 return True 

66 return False 

67 

68 async def wait_and_acquire( 

69 self, tokens: int = 1, timeout: Optional[float] = None 

70 ) -> bool: 

71 """Wait until tokens are available or timeout expires.""" 

72 deadline = (time.monotonic() + timeout) if timeout is not None else None 

73 

74 while True: 

75 if await self.acquire(tokens): 

76 return True 

77 

78 async with self._lock: 

79 self._refill() 

80 if self._tokens >= tokens: 

81 self._tokens -= tokens 

82 return True 

83 

84 # Calculate wait time 

85 needed = tokens - self._tokens 

86 wait_time = needed / self.rate 

87 

88 if deadline is not None: 

89 remaining = deadline - time.monotonic() 

90 if remaining <= 0: 

91 return False 

92 wait_time = min(wait_time, remaining) 

93 

94 await asyncio.sleep(wait_time) 

95 

96 @property 

97 def available_tokens(self) -> float: 

98 return self._tokens 

99 

100 @property 

101 def fill_level(self) -> float: 

102 """0.0 (empty) to 1.0 (full).""" 

103 return self._tokens / self.capacity 

104 

105 

106class SlidingWindow: 

107 """Sliding window rate limiter using precise timestamps. 

108 

109 Tracks request timestamps in a deque for O(1) amortized cleanup. 

110 """ 

111 

112 def __init__(self, max_requests: int, window_seconds: float): 

113 if max_requests <= 0: 

114 raise ValueError("max_requests must be > 0") 

115 if window_seconds <= 0: 

116 raise ValueError("window_seconds must be > 0") 

117 

118 self.max_requests = max_requests 

119 self.window_seconds = window_seconds 

120 self._timestamps: list[float] = [] 

121 self._lock = asyncio.Lock() 

122 

123 def _cleanup(self, now: float) -> None: 

124 cutoff = now - self.window_seconds 

125 # Find first timestamp within window 

126 idx = 0 

127 for ts in self._timestamps: 

128 if ts >= cutoff: 

129 break 

130 idx += 1 

131 if idx > 0: 

132 self._timestamps = self._timestamps[idx:] 

133 

134 async def acquire(self) -> bool: 

135 """Try to add a request to the window. Returns True if within limit.""" 

136 async with self._lock: 

137 now = time.monotonic() 

138 self._cleanup(now) 

139 if len(self._timestamps) < self.max_requests: 

140 self._timestamps.append(now) 

141 return True 

142 return False 

143 

144 @property 

145 def current_count(self) -> int: 

146 return len(self._timestamps) 

147 

148 @property 

149 def remaining(self) -> int: 

150 return max(0, self.max_requests - len(self._timestamps)) 

151 

152 

153class ConcurrentLimiter: 

154 """Limit the number of in-flight concurrent operations.""" 

155 

156 def __init__(self, max_concurrent: int): 

157 if max_concurrent <= 0: 

158 raise ValueError("max_concurrent must be > 0") 

159 self._semaphore = asyncio.Semaphore(max_concurrent) 

160 

161 async def acquire(self) -> bool: 

162 """Acquire a slot. Returns False if cancelled.""" 

163 try: 

164 await self._semaphore.acquire() 

165 return True 

166 except asyncio.CancelledError: 

167 return False 

168 

169 def release(self) -> None: 

170 """Release a slot.""" 

171 self._semaphore.release() 

172 

173 @property 

174 def available(self) -> int: 

175 return self._semaphore._value # pyright: ignore[reportPrivateUsage] 

176 

177 

178class CompositeLimiter: 

179 """Chain multiple limiters — all must pass for a request to proceed.""" 

180 

181 def __init__(self, limiters: list): 

182 self.limiters = limiters 

183 

184 async def acquire(self) -> bool: 

185 """Try to acquire from all limiters simultaneously.""" 

186 results = await asyncio.gather( 

187 *[limiter.acquire() for limiter in self.limiters], 

188 return_exceptions=True, 

189 ) 

190 for r in results: 

191 if r is False or isinstance(r, Exception): 

192 return False 

193 return all(results) 

194 

195 

196# ============================================================================ 

197# Rate limit decorator / context manager 

198# ============================================================================ 

199 

200class RateLimitError(Exception): 

201 """Raised when rate limit is exceeded.""" 

202 pass 

203 

204 

205class RateLimiter: 

206 """Unified rate limiter wrapping any strategy.""" 

207 

208 def __init__( 

209 self, 

210 name: str = "default", 

211 strategy=None, 

212 ): 

213 self.name = name 

214 self.strategy = strategy 

215 

216 async def __aenter__(self): 

217 ok = await self.strategy.acquire() if self.strategy else True 

218 if not ok: 

219 raise RateLimitError(f"Rate limit exceeded: {self.name}") 

220 return self 

221 

222 async def __aexit__(self, *args): 

223 if isinstance(self.strategy, ConcurrentLimiter): 

224 self.strategy.release() 

225 

226 @classmethod 

227 def token_bucket(cls, name: str, rate: float, capacity: int) -> "RateLimiter": 

228 return cls(name=name, strategy=TokenBucket(rate=rate, capacity=capacity)) 

229 

230 @classmethod 

231 def sliding_window( 

232 cls, name: str, max_requests: int, window_seconds: float 

233 ) -> "RateLimiter": 

234 return cls( 

235 name=name, 

236 strategy=SlidingWindow( 

237 max_requests=max_requests, window_seconds=window_seconds 

238 ), 

239 ) 

240 

241 @classmethod 

242 def concurrent(cls, name: str, max_concurrent: int) -> "RateLimiter": 

243 return cls(name=name, strategy=ConcurrentLimiter(max_concurrent=max_concurrent)) 

244 

245 

246# ============================================================================ 

247# Endpoint-level registry 

248# ============================================================================ 

249 

250@dataclass 

251class EndpointRateLimit: 

252 """Per-endpoint rate limit configuration.""" 

253 

254 endpoint: str 

255 requests_per_second: Optional[float] = None 

256 requests_per_minute: Optional[int] = None 

257 concurrent: Optional[int] = None 

258 burst: int = 1 

259 

260 

261class RateLimitRegistry: 

262 """Manage per-endpoint rate limiters.""" 

263 

264 def __init__(self): 

265 self._limiters: Dict[str, RateLimiter] = {} 

266 self._lock = asyncio.Lock() 

267 

268 async def configure(self, spec: EndpointRateLimit) -> RateLimiter: 

269 async with self._lock: 

270 limiters = [] 

271 

272 if spec.requests_per_second is not None: 

273 limiters.append( 

274 TokenBucket( 

275 rate=spec.requests_per_second, 

276 capacity=max(spec.burst, 1), 

277 ) 

278 ) 

279 

280 if spec.requests_per_minute is not None: 

281 limiters.append( 

282 SlidingWindow( 

283 max_requests=spec.requests_per_minute, 

284 window_seconds=60.0, 

285 ) 

286 ) 

287 

288 if spec.concurrent is not None: 

289 limiters.append(ConcurrentLimiter(max_concurrent=spec.concurrent)) 

290 

291 strategy = ( 

292 CompositeLimiter(limiters) if len(limiters) > 1 

293 else limiters[0] if limiters 

294 else None 

295 ) 

296 

297 limiter = RateLimiter(name=spec.endpoint, strategy=strategy) 

298 self._limiters[spec.endpoint] = limiter 

299 return limiter 

300 

301 async def get(self, endpoint: str) -> Optional[RateLimiter]: 

302 return self._limiters.get(endpoint) 

303 

304 async def acquire(self, endpoint: str) -> bool: 

305 limiter = self._limiters.get(endpoint) 

306 if limiter is None: 

307 return True 

308 return await limiter.strategy.acquire() if limiter.strategy else True