Coverage for agentos/tools/rate_limiter.py: 27%

1""" 

2RateLimiter — token bucket and sliding window rate limiters. 

3 

4Supports two algorithms: 

5 - TokenBucket: supports burst (tokens accumulate up to burst size), smooth refill 

6 - SlidingWindow: strict per-window limit, no burst 

7 

8Common interface: 

9 - try_acquire(key) → bool 

10 - acquire_or_wait(key, timeout) → bool (blocking with timeout) 

11 - reset(key) 

12 - stats() → dict 

13""" 

14 

15from __future__ import annotations 

16 

17import threading 

18import time 

19from dataclasses import dataclass 

20from typing import Any, Dict, List, Optional 

21 

22 

23# ============================================================================ 

24# Rate Limit Exceeded 

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

26 

27class RateLimitExceeded(Exception): 

28 def __init__(self, key: str, limit: float, window: float): 

29 self.key = key 

30 self.limit = limit 

31 self.window = window 

32 super().__init__(f"Rate limit exceeded for '{key}': {limit}/{window}s") 

33 

34 

35# ============================================================================ 

36# TokenBucket 

37# ============================================================================ 

38 

39@dataclass 

40class _BucketState: 

41 tokens: float 

42 last_refill: float 

43 

44 

45class TokenBucket: 

46 """Token bucket rate limiter with burst support. 

47 

48 Usage: 

49 limiter = TokenBucket(rate=10.0, burst=20.0) # 10 tokens/sec, burst up to 20 

50 limiter.try_acquire("api:user:42") # → True/False 

51 limiter.try_acquire("api:user:42", tokens=5) # consume 5 tokens 

52 """ 

53 

54 def __init__(self, rate: float, burst: Optional[float] = None): 

55 if rate <= 0: 

56 raise ValueError("rate must be positive") 

57 self._rate = rate 

58 self._burst = burst if burst is not None else rate 

59 self._buckets: Dict[str, _BucketState] = {} 

60 self._lock = threading.RLock() 

61 self._total_acquired: int = 0 

62 self._total_rejected: int = 0 

63 

64 def try_acquire(self, key: str, tokens: float = 1.0) -> bool: 

65 """Try to acquire tokens. Returns True if allowed.""" 

66 now = time.monotonic() 

67 with self._lock: 

68 bucket = self._buckets.get(key) 

69 if bucket is None: 

70 bucket = _BucketState(tokens=self._burst, last_refill=now) 

71 self._buckets[key] = bucket 

72 else: 

73 # Refill 

74 elapsed = now - bucket.last_refill 

75 bucket.tokens = min(self._burst, bucket.tokens + elapsed * self._rate) 

76 bucket.last_refill = now 

77 

78 if bucket.tokens >= tokens: 

79 bucket.tokens -= tokens 

80 self._total_acquired += 1 

81 return True 

82 else: 

83 self._total_rejected += 1 

84 return False 

85 

86 def acquire_or_wait(self, key: str, timeout: Optional[float] = None, tokens: float = 1.0) -> bool: 

87 """Block until tokens available or timeout.""" 

88 deadline = time.monotonic() + timeout if timeout else None 

89 while True: 

90 if self.try_acquire(key, tokens): 

91 return True 

92 if deadline and time.monotonic() >= deadline: 

93 return False 

94 time.sleep(0.01) 

95 

96 def reset(self, key: str) -> None: 

97 with self._lock: 

98 self._buckets.pop(key, None) 

99 

100 def reset_all(self) -> None: 

101 with self._lock: 

102 self._buckets.clear() 

103 

104 def stats(self) -> Dict[str, Any]: 

105 with self._lock: 

106 return { 

107 "rate": self._rate, 

108 "burst": self._burst, 

109 "active_keys": len(self._buckets), 

110 "total_acquired": self._total_acquired, 

111 "total_rejected": self._total_rejected, 

112 } 

113 

114 @property 

115 def rate(self) -> float: 

116 return self._rate 

117 

118 

119# ============================================================================ 

120# SlidingWindow 

121# ============================================================================ 

122 

123class SlidingWindow: 

124 """Sliding window rate limiter — strict per-window limit, no burst. 

125 

126 Usage: 

127 limiter = SlidingWindow(limit=100, window=60.0) # 100 req per 60s 

128 limiter.try_acquire("api:endpoint") # → True/False 

129 """ 

130 

131 def __init__(self, limit: int, window: float = 60.0): 

132 if limit <= 0: 

133 raise ValueError("limit must be positive") 

134 self._limit = limit 

135 self._window = window 

136 self._windows: Dict[str, List[float]] = {} 

137 self._lock = threading.RLock() 

138 self._total_acquired: int = 0 

139 self._total_rejected: int = 0 

140 

141 def try_acquire(self, key: str) -> bool: 

142 """Try to acquire a slot. Returns True if within limit.""" 

143 now = time.monotonic() 

144 with self._lock: 

145 timestamps = self._windows.get(key) 

146 if timestamps is None: 

147 timestamps = [] 

148 self._windows[key] = timestamps 

149 

150 # Evict expired entries 

151 cutoff = now - self._window 

152 while timestamps and timestamps[0] < cutoff: 

153 timestamps.pop(0) 

154 

155 if len(timestamps) < self._limit: 

156 timestamps.append(now) 

157 self._total_acquired += 1 

158 return True 

159 else: 

160 self._total_rejected += 1 

161 return False 

162 

163 def acquire_or_wait(self, key: str, timeout: Optional[float] = None) -> bool: 

164 deadline = time.monotonic() + timeout if timeout else None 

165 while True: 

166 if self.try_acquire(key): 

167 return True 

168 if deadline and time.monotonic() >= deadline: 

169 return False 

170 time.sleep(0.02) 

171 

172 def reset(self, key: str) -> None: 

173 with self._lock: 

174 self._windows.pop(key, None) 

175 

176 def reset_all(self) -> None: 

177 with self._lock: 

178 self._windows.clear() 

179 

180 def stats(self) -> Dict[str, Any]: 

181 with self._lock: 

182 return { 

183 "limit": self._limit, 

184 "window": self._window, 

185 "active_keys": len(self._windows), 

186 "total_acquired": self._total_acquired, 

187 "total_rejected": self._total_rejected, 

188 } 

189 

190 @property 

191 def limit(self) -> int: 

192 return self._limit