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 typing import Dict, Optional, TypeVar 

19 

20logger = logging.getLogger(__name__) 

21 

22T = TypeVar("T") 

23 

24 

25# ============================================================================ 

26# Core limiters 

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

28 

29class TokenBucket: 

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

31 

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

33 """ 

34 

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

36 """ 

37 Args: 

38 rate: Tokens per second refill rate 

39 capacity: Maximum tokens (burst size) 

40 """ 

41 if rate <= 0: 

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

43 if capacity <= 0: 

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

45 

46 self.rate = rate 

47 self.capacity = capacity 

48 self._tokens = float(capacity) 

49 self._last_refill = time.monotonic() 

50 self._lock = asyncio.Lock() 

51 

52 def _refill(self) -> None: 

53 now = time.monotonic() 

54 elapsed = now - self._last_refill 

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

56 self._last_refill = now 

57 

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

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

60 async with self._lock: 

61 self._refill() 

62 if self._tokens >= tokens: 

63 self._tokens -= tokens 

64 return True 

65 return False 

66 

67 async def wait_and_acquire( 

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

69 ) -> bool: 

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

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

72 

73 while True: 

74 if await self.acquire(tokens): 

75 return True 

76 

77 async with self._lock: 

78 self._refill() 

79 if self._tokens >= tokens: 

80 self._tokens -= tokens 

81 return True 

82 

83 # Calculate wait time 

84 needed = tokens - self._tokens 

85 wait_time = needed / self.rate 

86 

87 if deadline is not None: 

88 remaining = deadline - time.monotonic() 

89 if remaining <= 0: 

90 return False 

91 wait_time = min(wait_time, remaining) 

92 

93 await asyncio.sleep(wait_time) 

94 

95 @property 

96 def available_tokens(self) -> float: 

97 return self._tokens 

98 

99 @property 

100 def fill_level(self) -> float: 

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

102 return self._tokens / self.capacity 

103 

104 

105class SlidingWindow: 

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

107 

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

109 """ 

110 

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

112 if max_requests <= 0: 

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

114 if window_seconds <= 0: 

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

116 

117 self.max_requests = max_requests 

118 self.window_seconds = window_seconds 

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

120 self._lock = asyncio.Lock() 

121 

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

123 cutoff = now - self.window_seconds 

124 # Find first timestamp within window 

125 idx = 0 

126 for ts in self._timestamps: 

127 if ts >= cutoff: 

128 break 

129 idx += 1 

130 if idx > 0: 

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

132 

133 async def acquire(self) -> bool: 

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

135 async with self._lock: 

136 now = time.monotonic() 

137 self._cleanup(now) 

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

139 self._timestamps.append(now) 

140 return True 

141 return False 

142 

143 @property 

144 def current_count(self) -> int: 

145 return len(self._timestamps) 

146 

147 @property 

148 def remaining(self) -> int: 

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

150 

151 

152class ConcurrentLimiter: 

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

154 

155 def __init__(self, max_concurrent: int): 

156 if max_concurrent <= 0: 

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

158 self._semaphore = asyncio.Semaphore(max_concurrent) 

159 

160 async def acquire(self) -> bool: 

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

162 try: 

163 await self._semaphore.acquire() 

164 return True 

165 except asyncio.CancelledError: 

166 return False 

167 

168 def release(self) -> None: 

169 """Release a slot.""" 

170 self._semaphore.release() 

171 

172 @property 

173 def available(self) -> int: 

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

175 

176 

177class CompositeLimiter: 

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

179 

180 def __init__(self, limiters: list): 

181 self.limiters = limiters 

182 

183 async def acquire(self) -> bool: 

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

185 results = await asyncio.gather( 

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

187 return_exceptions=True, 

188 ) 

189 for r in results: 

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

191 return False 

192 return all(results) 

193 

194 

195# ============================================================================ 

196# Rate limit decorator / context manager 

197# ============================================================================ 

198 

199class RateLimitError(Exception): 

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

201 pass 

202 

203 

204class RateLimiter: 

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

206 

207 def __init__( 

208 self, 

209 name: str = "default", 

210 strategy=None, 

211 ): 

212 self.name = name 

213 self.strategy = strategy 

214 

215 async def __aenter__(self): 

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

217 if not ok: 

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

219 return self 

220 

221 async def __aexit__(self, *args): 

222 if isinstance(self.strategy, ConcurrentLimiter): 

223 self.strategy.release() 

224 

225 @classmethod 

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

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

228 

229 @classmethod 

230 def sliding_window( 

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

232 ) -> "RateLimiter": 

233 return cls( 

234 name=name, 

235 strategy=SlidingWindow( 

236 max_requests=max_requests, window_seconds=window_seconds 

237 ), 

238 ) 

239 

240 @classmethod 

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

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

243 

244 

245# ============================================================================ 

246# Endpoint-level registry 

247# ============================================================================ 

248 

249@dataclass 

250class EndpointRateLimit: 

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

252 

253 endpoint: str 

254 requests_per_second: Optional[float] = None 

255 requests_per_minute: Optional[int] = None 

256 concurrent: Optional[int] = None 

257 burst: int = 1 

258 

259 

260class RateLimitRegistry: 

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

262 

263 def __init__(self): 

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

265 self._lock = asyncio.Lock() 

266 

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

268 async with self._lock: 

269 limiters = [] 

270 

271 if spec.requests_per_second is not None: 

272 limiters.append( 

273 TokenBucket( 

274 rate=spec.requests_per_second, 

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

276 ) 

277 ) 

278 

279 if spec.requests_per_minute is not None: 

280 limiters.append( 

281 SlidingWindow( 

282 max_requests=spec.requests_per_minute, 

283 window_seconds=60.0, 

284 ) 

285 ) 

286 

287 if spec.concurrent is not None: 

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

289 

290 strategy = ( 

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

292 else limiters[0] if limiters 

293 else None 

294 ) 

295 

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

297 self._limiters[spec.endpoint] = limiter 

298 return limiter 

299 

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

301 return self._limiters.get(endpoint) 

302 

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

304 limiter = self._limiters.get(endpoint) 

305 if limiter is None: 

306 return True 

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