Coverage for src/lexigram/web/websocket/rate_limiter.py: 43%

1"""WebSocket message rate limiter. 

2 

3Implements a per-connection token-bucket rate limiter for WebSocket messages. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from dataclasses import dataclass 

10 

11from lexigram.primitives import clock as ambient_clock 

12 

13 

14@dataclass 

15class _BucketState: 

16 """Token-bucket state for a single connection.""" 

17 

18 tokens: float 

19 last_refill: float = 0.0 

20 

21 

22class WebSocketRateLimiter: 

23 """Per-connection token-bucket rate limiter for WebSocket messages. 

24 

25 Each connection gets its own token bucket that refills at ``max_messages_per_second`` 

26 tokens per second. An incoming message consumes one token. If the bucket is empty 

27 the message is rejected and ``check()`` returns ``False``. 

28 

29 State for connections that have been cleaned up via ``remove()`` is discarded 

30 automatically. 

31 

32 Attributes: 

33 max_messages_per_second: Maximum sustained message rate per connection. 

34 

35 Example: 

36 ```python 

37 limiter = WebSocketRateLimiter(max_messages_per_second=10.0) 

38 connection_id_manager = get_connection_id_manager() 

39 

40 async def on_message(self, websocket, message): 

41 conn_id = connection_id_manager.get_or_create(websocket) 

42 if not await limiter.check(conn_id): 

43 await websocket.send_json({"type": "error", "error": "rate_limit_exceeded"}) 

44 return 

45 ... 

46 ``` 

47 """ 

48 

49 def __init__( 

50 self, 

51 max_messages_per_second: float = 10.0, 

52 ) -> None: 

53 """Initialise the rate limiter. 

54 

55 Args: 

56 max_messages_per_second: Maximum number of messages per second per connection. 

57 Must be a positive number. 

58 """ 

59 if max_messages_per_second <= 0: 

60 raise ValueError("max_messages_per_second must be positive") 

61 self._max_rate = max_messages_per_second 

62 self._connections: dict[str, _BucketState] = {} 

63 self._lock = asyncio.Lock() 

64 

65 async def check(self, connection_id: str) -> bool: 

66 """Check whether a message from ``connection_id`` is within the rate limit. 

67 

68 Refills the token bucket proportionally to elapsed time, then attempts 

69 to consume one token. Returns ``True`` if the message is allowed, 

70 ``False`` if the rate limit is exceeded. 

71 

72 Args: 

73 connection_id: A unique identifier for the connection (UUID string). 

74 

75 Returns: 

76 True if the message is allowed; False if rate-limited. 

77 """ 

78 now = ambient_clock.monotonic() 

79 

80 async with self._lock: 

81 if connection_id not in self._connections: 

82 # New connection: start with a full bucket 

83 self._connections[connection_id] = _BucketState( 

84 tokens=self._max_rate, 

85 last_refill=now, 

86 ) 

87 

88 bucket = self._connections[connection_id] 

89 

90 # Refill tokens based on elapsed time 

91 elapsed = now - bucket.last_refill 

92 bucket.tokens = min( 

93 self._max_rate, 

94 bucket.tokens + elapsed * self._max_rate, 

95 ) 

96 bucket.last_refill = now 

97 

98 if bucket.tokens < 1.0: 

99 return False 

100 

101 bucket.tokens -= 1.0 

102 return True 

103 

104 async def remove(self, connection_id: str) -> None: 

105 """Remove tracking state for a connection. 

106 

107 Call this when a WebSocket connection closes to free memory. 

108 

109 Args: 

110 connection_id: The same identifier passed to ``check()``. 

111 """ 

112 async with self._lock: 

113 self._connections.pop(connection_id, None) 

114 

115 @property 

116 def active_connection_count(self) -> int: 

117 """Number of connections currently being tracked.""" 

118 return len(self._connections) 

119 

120 

121__all__ = ["WebSocketRateLimiter"]