Class: A2A::Server::Middleware::RateLimitMiddleware

Inherits:
Object
  • Object
show all
Defined in:
lib/a2a/server/middleware/rate_limit_middleware.rb

Overview

Rate limiting middleware for A2A requests

Implements rate limiting using various strategies including in-memory, Redis-backed, and sliding window algorithms.

Examples:

Basic usage

middleware = RateLimitMiddleware.new(
  limit: 100,
  window: 3600, # 1 hour
  strategy: :sliding_window
)

Constant Summary collapse

STRATEGIES =

Rate limiting strategies

i[fixed_window sliding_window token_bucket].freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(limit: 100, window: 3600, strategy: :sliding_window, store: nil, key_generator: nil) ⇒ RateLimitMiddleware

Initialize rate limiting middleware

Parameters:

  • (defaults to: 100)

    Maximum number of requests per window

  • (defaults to: 3600)

    Time window in seconds

  • (defaults to: :sliding_window)

    Rate limiting strategy

  • (defaults to: nil)

    Storage backend (defaults to in-memory)

  • (defaults to: nil)

    Custom key generator for rate limiting



35
36
37
38
39
40
41
42
43
44
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 35

def initialize(limit: 100, window: 3600, strategy: :sliding_window,
               store: nil, key_generator: nil)
  @limit = limit
  @window = window
  @strategy = strategy
  @store = store || InMemoryStore.new
  @key_generator = key_generator || method(:default_key_generator)

  validate_strategy!
end

Instance Attribute Details

#limit ⇒ Object (readonly)

Returns the value of attribute limit.



22
23
24
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 22

def limit
  @limit
end

#store ⇒ Object (readonly)

Returns the value of attribute store.



22
23
24
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 22

def store
  @store
end

#strategy ⇒ Object (readonly)

Returns the value of attribute strategy.



22
23
24
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 22

def strategy
  @strategy
end

#window ⇒ Object (readonly)

Returns the value of attribute window.



22
23
24
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 22

def window
  @window
end

Instance Method Details

#call(request, context) { ... } ⇒ Object

Process rate limiting for a request

Parameters:

  • The JSON-RPC request

  • The request context

Yields:

  • Block to continue the middleware chain

Returns:

  • The result from the next middleware or handler

Raises:

  • If rate limit is exceeded



54
55
56
57
58
59
60
61
62
63
64
65
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 54

def call(request, context)
  # Generate rate limiting key
  key = @key_generator.call(request, context)

  # Check rate limit
  unless check_rate_limit(key)
    raise A2A::Errors::RateLimitExceeded, "Rate limit exceeded: #{@limit} requests per #{@window} seconds"
  end

  # Continue to next middleware
  yield
end

#check_fixed_window(key) ⇒ Boolean (private)

Fixed window rate limiting

Parameters:

  • The rate limiting key

Returns:

  • True if within limit



139
140
141
142
143
144
145
146
147
148
149
150
151
152
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 139

def check_fixed_window(key)
  now = Time.now.to_i
  window_start = (now / @window) * @window
  window_key = "#{key}:#{window_start}"

  current_count = @store.get(window_key) || 0

  if current_count >= @limit
    false
  else
    @store.increment(window_key, expires_at: window_start + @window)
    true
  end
end

#check_rate_limit(key) ⇒ Boolean

Check if request is within rate limit

Parameters:

  • The rate limiting key

Returns:

  • True if within limit, false otherwise



72
73
74
75
76
77
78
79
80
81
82
83
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 72

def check_rate_limit(key)
  case @strategy
  when :fixed_window
    check_fixed_window(key)
  when :sliding_window
    check_sliding_window(key)
  when :token_bucket
    check_token_bucket(key)
  else
    true # Fallback to allow request
  end
end

#check_sliding_window(key) ⇒ Boolean (private)

Sliding window rate limiting

Parameters:

  • The rate limiting key

Returns:

  • True if within limit



159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 159

def check_sliding_window(key)
  now = Time.now.to_f
  window_start = now - @window

  # Get timestamps of requests in the current window
  timestamps = @store.get_list("#{key}:timestamps") || []

  # Remove old timestamps
  timestamps = timestamps.select { |ts| ts > window_start }

  if timestamps.length >= @limit
    false
  else
    # Add current timestamp
    timestamps << now
    @store.set_list("#{key}:timestamps", timestamps, expires_at: now + @window)
    true
  end
end

#check_token_bucket(key) ⇒ Boolean (private)

Token bucket rate limiting

Parameters:

  • The rate limiting key

Returns:

  • True if within limit



184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 184

def check_token_bucket(key)
  now = Time.now.to_f
  bucket_key = "#{key}:bucket"

  # Get current bucket state
  bucket = @store.get(bucket_key) || { tokens: @limit, last_refill: now }

  # Calculate tokens to add based on time elapsed
  time_elapsed = now - bucket[:last_refill]
  tokens_to_add = (time_elapsed / @window) * @limit

  # Refill bucket
  bucket[:tokens] = [@limit, bucket[:tokens] + tokens_to_add].min
  bucket[:last_refill] = now

  if bucket[:tokens] >= 1
    bucket[:tokens] -= 1
    @store.set(bucket_key, bucket, expires_at: now + (@window * 2))
    true
  else
    @store.set(bucket_key, bucket, expires_at: now + (@window * 2))
    false
  end
end

#default_key_generator(_request, context) ⇒ String (private)

Default key generator based on authentication or IP

Parameters:

  • The request

  • The context

Returns:

  • The rate limiting key



119
120
121
122
123
124
125
126
127
128
129
130
131
132
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 119

def default_key_generator(_request, context)
  # Try to use authenticated user/API key
  if context.authenticated?
    auth_data = context.instance_variable_get(:@auth_schemes)&.values&.first
    return "user:#{auth_data[:username] || auth_data[:api_key] || auth_data[:token]}" if auth_data.is_a?(Hash)
  end

  # Fall back to IP address if available
  ip = context.get_metadata(:remote_ip) || context.get_metadata("REMOTE_ADDR")
  return "ip:#{ip}" if ip

  # Default fallback
  "anonymous"
end

#fixed_window_status(key) ⇒ Object (private)

Get fixed window status



211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 211

def fixed_window_status(key)
  now = Time.now.to_i
  window_start = (now / @window) * @window
  window_key = "#{key}:#{window_start}"

  current_count = @store.get(window_key) || 0
  reset_time = window_start + @window

  {
    limit: @limit,
    remaining: [@limit - current_count, 0].max,
    reset_time: Time.zone.at(reset_time),
    window_start: Time.zone.at(window_start)
  }
end

#sliding_window_status(key) ⇒ Object (private)

Get sliding window status



229
230
231
232
233
234
235
236
237
238
239
240
241
242
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 229

def sliding_window_status(key)
  now = Time.now.to_f
  window_start = now - @window

  timestamps = @store.get_list("#{key}:timestamps") || []
  current_count = timestamps.count { |ts| ts > window_start }

  {
    limit: @limit,
    remaining: [@limit - current_count, 0].max,
    reset_time: nil, # No fixed reset time for sliding window
    window_start: Time.zone.at(window_start)
  }
end

#status(key) ⇒ Hash

Get current rate limit status for a key

Parameters:

  • The rate limiting key

Returns:

  • Status information



90
91
92
93
94
95
96
97
98
99
100
101
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 90

def status(key)
  case @strategy
  when :fixed_window
    fixed_window_status(key)
  when :sliding_window
    sliding_window_status(key)
  when :token_bucket
    token_bucket_status(key)
  else
    { limit: @limit, remaining: @limit, reset_time: nil }
  end
end

#token_bucket_status(key) ⇒ Object (private)

Get token bucket status



246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 246

def token_bucket_status(key)
  now = Time.now.to_f
  bucket_key = "#{key}:bucket"

  bucket = @store.get(bucket_key) || { tokens: @limit, last_refill: now }

  # Calculate current tokens
  time_elapsed = now - bucket[:last_refill]
  tokens_to_add = (time_elapsed / @window) * @limit
  current_tokens = [@limit, bucket[:tokens] + tokens_to_add].min

  {
    limit: @limit,
    remaining: current_tokens.floor,
    reset_time: nil, # Continuous refill
    tokens: current_tokens
  }
end

#validate_strategy! ⇒ Object (private)

Validate the rate limiting strategy

Raises:



107
108
109
110
111
 
# File 'lib/a2a/server/middleware/rate_limit_middleware.rb', line 107

def validate_strategy!
  return if STRATEGIES.include?(@strategy)

  raise ArgumentError, "Invalid strategy: #{@strategy}. Must be one of: #{STRATEGIES.join(', ')}"
end