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 

21 

22# ============================================================================ 

23# Rate Limit Exceeded 

24# ============================================================================ 

25 

26 

27class RateLimitExceeded(Exception): # noqa: N818 

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 

40@dataclass 

41class _BucketState: 

42 tokens: float 

43 last_refill: float 

44 

45 

46class TokenBucket: 

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

48 

49 Usage: 

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

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

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

53 """ 

54 

55 def __init__(self, rate: float, burst: float | None = None): 

56 if rate <= 0: 

57 raise ValueError("rate must be positive") 

58 self._rate = rate 

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

60 self._buckets: dict[str, _BucketState] = {} 

61 self._lock = threading.RLock() 

62 self._total_acquired: int = 0 

63 self._total_rejected: int = 0 

64 

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

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

67 now = time.monotonic() 

68 with self._lock: 

69 bucket = self._buckets.get(key) 

70 if bucket is None: 

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

72 self._buckets[key] = bucket 

73 else: 

74 # Refill 

75 elapsed = now - bucket.last_refill 

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

77 bucket.last_refill = now 

78 

79 if bucket.tokens >= tokens: 

80 bucket.tokens -= tokens 

81 self._total_acquired += 1 

82 return True 

83 else: 

84 self._total_rejected += 1 

85 return False 

86 

87 def acquire_or_wait(self, key: str, timeout: float | None = None, tokens: float = 1.0) -> bool: 

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

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

90 while True: 

91 if self.try_acquire(key, tokens): 

92 return True 

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

94 return False 

95 time.sleep(0.01) 

96 

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

98 with self._lock: 

99 self._buckets.pop(key, None) 

100 

101 def reset_all(self) -> None: 

102 with self._lock: 

103 self._buckets.clear() 

104 

105 def stats(self) -> dict[str, Any]: 

106 with self._lock: 

107 return { 

108 "rate": self._rate, 

109 "burst": self._burst, 

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

111 "total_acquired": self._total_acquired, 

112 "total_rejected": self._total_rejected, 

113 } 

114 

115 @property 

116 def rate(self) -> float: 

117 return self._rate 

118 

119 

120# ============================================================================ 

121# SlidingWindow 

122# ============================================================================ 

123 

124 

125class SlidingWindow: 

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

127 

128 Usage: 

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

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

131 """ 

132 

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

134 if limit <= 0: 

135 raise ValueError("limit must be positive") 

136 self._limit = limit 

137 self._window = window 

138 self._windows: dict[str, list[float]] = {} 

139 self._lock = threading.RLock() 

140 self._total_acquired: int = 0 

141 self._total_rejected: int = 0 

142 

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

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

145 now = time.monotonic() 

146 with self._lock: 

147 timestamps = self._windows.get(key) 

148 if timestamps is None: 

149 timestamps = [] 

150 self._windows[key] = timestamps 

151 

152 # Evict expired entries 

153 cutoff = now - self._window 

154 while timestamps and timestamps[0] < cutoff: 

155 timestamps.pop(0) 

156 

157 if len(timestamps) < self._limit: 

158 timestamps.append(now) 

159 self._total_acquired += 1 

160 return True 

161 else: 

162 self._total_rejected += 1 

163 return False 

164 

165 def acquire_or_wait(self, key: str, timeout: float | None = None) -> bool: 

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

167 while True: 

168 if self.try_acquire(key): 

169 return True 

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

171 return False 

172 time.sleep(0.02) 

173 

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

175 with self._lock: 

176 self._windows.pop(key, None) 

177 

178 def reset_all(self) -> None: 

179 with self._lock: 

180 self._windows.clear() 

181 

182 def stats(self) -> dict[str, Any]: 

183 with self._lock: 

184 return { 

185 "limit": self._limit, 

186 "window": self._window, 

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

188 "total_acquired": self._total_acquired, 

189 "total_rejected": self._total_rejected, 

190 } 

191 

192 @property 

193 def limit(self) -> int: 

194 return self._limit