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 TypeVar 

19 

20logger = logging.getLogger(__name__) 

21 

22T = TypeVar("T") 

23 

24 

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

26# Core limiters 

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

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(self, tokens: int = 1, timeout: float | None = None) -> bool: 

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

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

71 

72 while True: 

73 if await self.acquire(tokens): 

74 return True 

75 

76 async with self._lock: 

77 self._refill() 

78 if self._tokens >= tokens: 

79 self._tokens -= tokens 

80 return True 

81 

82 # Calculate wait time 

83 needed = tokens - self._tokens 

84 wait_time = needed / self.rate 

85 

86 if deadline is not None: 

87 remaining = deadline - time.monotonic() 

88 if remaining <= 0: 

89 return False 

90 wait_time = min(wait_time, remaining) 

91 

92 await asyncio.sleep(wait_time) 

93 

94 @property 

95 def available_tokens(self) -> float: 

96 return self._tokens 

97 

98 @property 

99 def fill_level(self) -> float: 

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

101 return self._tokens / self.capacity 

102 

103 

104class SlidingWindow: 

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

106 

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

108 """ 

109 

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

111 if max_requests <= 0: 

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

113 if window_seconds <= 0: 

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

115 

116 self.max_requests = max_requests 

117 self.window_seconds = window_seconds 

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

119 self._lock = asyncio.Lock() 

120 

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

122 cutoff = now - self.window_seconds 

123 # Find first timestamp within window 

124 idx = 0 

125 for ts in self._timestamps: 

126 if ts >= cutoff: 

127 break 

128 idx += 1 

129 if idx > 0: 

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

131 

132 async def acquire(self) -> bool: 

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

134 async with self._lock: 

135 now = time.monotonic() 

136 self._cleanup(now) 

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

138 self._timestamps.append(now) 

139 return True 

140 return False 

141 

142 @property 

143 def current_count(self) -> int: 

144 return len(self._timestamps) 

145 

146 @property 

147 def remaining(self) -> int: 

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

149 

150 

151class ConcurrentLimiter: 

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

153 

154 def __init__(self, max_concurrent: int): 

155 if max_concurrent <= 0: 

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

157 self._semaphore = asyncio.Semaphore(max_concurrent) 

158 

159 async def acquire(self) -> bool: 

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

161 try: 

162 await self._semaphore.acquire() 

163 return True 

164 except asyncio.CancelledError: 

165 return False 

166 

167 def release(self) -> None: 

168 """Release a slot.""" 

169 self._semaphore.release() 

170 

171 @property 

172 def available(self) -> int: 

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

174 

175 

176class CompositeLimiter: 

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

178 

179 def __init__(self, limiters: list): 

180 self.limiters = limiters 

181 

182 async def acquire(self) -> bool: 

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

184 results = await asyncio.gather( 

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

186 return_exceptions=True, 

187 ) 

188 for r in results: 

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

190 return False 

191 return all(results) 

192 

193 

194# ============================================================================ 

195# Rate limit decorator / context manager 

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

197 

198 

199class RateLimitError(Exception): 

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

201 

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(cls, name: str, max_requests: int, window_seconds: float) -> RateLimiter: 

231 return cls( 

232 name=name, 

233 strategy=SlidingWindow(max_requests=max_requests, window_seconds=window_seconds), 

234 ) 

235 

236 @classmethod 

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

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

239 

240 

241# ============================================================================ 

242# Endpoint-level registry 

243# ============================================================================ 

244 

245 

246@dataclass 

247class EndpointRateLimit: 

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

249 

250 endpoint: str 

251 requests_per_second: float | None = None 

252 requests_per_minute: int | None = None 

253 concurrent: int | None = None 

254 burst: int = 1 

255 

256 

257class RateLimitRegistry: 

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

259 

260 def __init__(self): 

261 self._limiters: dict[str, RateLimiter] = {} 

262 self._lock = asyncio.Lock() 

263 

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

265 async with self._lock: 

266 limiters = [] 

267 

268 if spec.requests_per_second is not None: 

269 limiters.append( 

270 TokenBucket( 

271 rate=spec.requests_per_second, 

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

273 ) 

274 ) 

275 

276 if spec.requests_per_minute is not None: 

277 limiters.append( 

278 SlidingWindow( 

279 max_requests=spec.requests_per_minute, 

280 window_seconds=60.0, 

281 ) 

282 ) 

283 

284 if spec.concurrent is not None: 

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

286 

287 strategy = ( 

288 CompositeLimiter(limiters) 

289 if len(limiters) > 1 

290 else limiters[0] if limiters else None 

291 ) 

292 

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

294 self._limiters[spec.endpoint] = limiter 

295 return limiter 

296 

297 async def get(self, endpoint: str) -> RateLimiter | None: 

298 return self._limiters.get(endpoint) 

299 

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

301 limiter = self._limiters.get(endpoint) 

302 if limiter is None: 

303 return True 

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