tsikol 0.1.0

data/docs/cookbook/rate-limiting.md

@@ -0,0 +1,717 @@
+# Rate Limiting Recipe
+
+This recipe shows various strategies for implementing rate limiting in your MCP server to prevent abuse and ensure fair usage.
+
+## Basic Rate Limiting
+
+### Simple Request Counter
+
+```ruby
+class BasicRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @max_requests = options[:max_requests] || 60
+    @window_seconds = options[:window_seconds] || 60
+    @requests = Hash.new { |h, k| h[k] = [] }
+    @mutex = Mutex.new
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    
+    if rate_limited?(client_id)
+      return rate_limit_error(request["id"])
+    end
+    
+    track_request(client_id)
+    @app.call(request)
+  end
+  
+  private
+  
+  def identify_client(request)
+    # Identify by various methods
+    request.dig("params", "_client_id") ||
+      request.dig("meta", "client_id") ||
+      request.dig("authenticated_user", "id") ||
+      "anonymous"
+  end
+  
+  def rate_limited?(client_id)
+    @mutex.synchronize do
+      clean_old_requests(client_id)
+      @requests[client_id].size >= @max_requests
+    end
+  end
+  
+  def track_request(client_id)
+    @mutex.synchronize do
+      @requests[client_id] << Time.now
+    end
+  end
+  
+  def clean_old_requests(client_id)
+    cutoff = Time.now - @window_seconds
+    @requests[client_id].reject! { |time| time < cutoff }
+  end
+  
+  def rate_limit_error(id)
+    {
+      jsonrpc: "2.0",
+      id: id,
+      error: {
+        code: -32005,
+        message: "Rate limit exceeded",
+        data: {
+          retry_after: @window_seconds,
+          limit: @max_requests
+        }
+      }
+    }
+  end
+end
+```
+
+## Token Bucket Algorithm
+
+### Advanced Rate Limiting with Bursts
+
+```ruby
+class TokenBucketRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @capacity = options[:capacity] || 10        # Max tokens
+    @refill_rate = options[:refill_rate] || 1  # Tokens per second
+    @buckets = Hash.new { |h, k| h[k] = create_bucket }
+    @mutex = Mutex.new
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    cost = request_cost(request)
+    
+    if consume_tokens(client_id, cost)
+      @app.call(request)
+    else
+      rate_limit_error(request["id"], client_id)
+    end
+  end
+  
+  private
+  
+  def create_bucket
+    {
+      tokens: @capacity.to_f,
+      last_refill: Time.now
+    }
+  end
+  
+  def consume_tokens(client_id, cost)
+    @mutex.synchronize do
+      bucket = @buckets[client_id]
+      refill_bucket(bucket)
+      
+      if bucket[:tokens] >= cost
+        bucket[:tokens] -= cost
+        true
+      else
+        false
+      end
+    end
+  end
+  
+  def refill_bucket(bucket)
+    now = Time.now
+    elapsed = now - bucket[:last_refill]
+    
+    # Add tokens based on elapsed time
+    tokens_to_add = elapsed * @refill_rate
+    bucket[:tokens] = [bucket[:tokens] + tokens_to_add, @capacity].min
+    bucket[:last_refill] = now
+  end
+  
+  def request_cost(request)
+    # Different costs for different operations
+    case request["method"]
+    when "tools/call"
+      case request.dig("params", "name")
+      when "expensive_operation" then 5
+      when "ai_generation" then 3
+      else 1
+      end
+    when "sampling/sample"
+      10  # AI sampling is expensive
+    else
+      1
+    end
+  end
+  
+  def rate_limit_error(id, client_id)
+    bucket = @buckets[client_id]
+    wait_time = (1 - bucket[:tokens]) / @refill_rate
+    
+    {
+      jsonrpc: "2.0",
+      id: id,
+      error: {
+        code: -32005,
+        message: "Rate limit exceeded",
+        data: {
+          retry_after: wait_time.ceil,
+          tokens_remaining: bucket[:tokens].floor,
+          refill_rate: @refill_rate
+        }
+      }
+    }
+  end
+end
+```
+
+## Sliding Window Rate Limiter
+
+### Precise Rate Limiting
+
+```ruby
+class SlidingWindowRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @max_requests = options[:max_requests] || 100
+    @window_seconds = options[:window_seconds] || 60
+    @precision = options[:precision] || 10  # Sub-windows
+    @counters = Hash.new { |h, k| h[k] = Array.new(@precision, 0) }
+    @mutex = Mutex.new
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    
+    if exceeds_limit?(client_id)
+      return rate_limit_error(request["id"])
+    end
+    
+    increment_counter(client_id)
+    @app.call(request)
+  end
+  
+  private
+  
+  def exceeds_limit?(client_id)
+    @mutex.synchronize do
+      current_count = calculate_current_count(client_id)
+      current_count >= @max_requests
+    end
+  end
+  
+  def calculate_current_count(client_id)
+    now = Time.now.to_f
+    window_size = @window_seconds.to_f / @precision
+    current_window = (now / window_size).floor % @precision
+    
+    # Age out old windows
+    @counters[client_id].each_with_index do |count, i|
+      window_age = ((current_window - i) % @precision) * window_size
+      if window_age >= @window_seconds
+        @counters[client_id][i] = 0
+      end
+    end
+    
+    @counters[client_id].sum
+  end
+  
+  def increment_counter(client_id)
+    @mutex.synchronize do
+      window_size = @window_seconds.to_f / @precision
+      current_window = (Time.now.to_f / window_size).floor % @precision
+      @counters[client_id][current_window] += 1
+    end
+  end
+end
+```
+
+## Distributed Rate Limiting
+
+### Redis-Based Rate Limiter
+
+```ruby
+require 'redis'
+
+class RedisRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @redis = options[:redis] || Redis.new
+    @max_requests = options[:max_requests] || 100
+    @window_seconds = options[:window_seconds] || 60
+    @key_prefix = options[:key_prefix] || "mcp:rate_limit"
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    
+    if rate_limited?(client_id)
+      return rate_limit_error(request["id"], client_id)
+    end
+    
+    @app.call(request)
+  end
+  
+  private
+  
+  def rate_limited?(client_id)
+    key = "#{@key_prefix}:#{client_id}"
+    
+    # Use Redis pipeline for atomic operations
+    result = @redis.pipelined do |pipeline|
+      pipeline.incr(key)
+      pipeline.expire(key, @window_seconds)
+    end
+    
+    current_count = result[0]
+    current_count > @max_requests
+  end
+  
+  def rate_limit_error(id, client_id)
+    key = "#{@key_prefix}:#{client_id}"
+    ttl = @redis.ttl(key)
+    
+    {
+      jsonrpc: "2.0",
+      id: id,
+      error: {
+        code: -32005,
+        message: "Rate limit exceeded",
+        data: {
+          retry_after: ttl > 0 ? ttl : @window_seconds,
+          limit: @max_requests,
+          window: @window_seconds
+        }
+      }
+    }
+  end
+end
+```
+
+### Redis Sliding Window
+
+```ruby
+class RedisSlidingWindowLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @redis = options[:redis] || Redis.new
+    @max_requests = options[:max_requests] || 100
+    @window_seconds = options[:window_seconds] || 60
+    @key_prefix = options[:key_prefix] || "mcp:sliding"
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    
+    if rate_limited?(client_id)
+      return rate_limit_error(request["id"])
+    end
+    
+    @app.call(request)
+  end
+  
+  private
+  
+  def rate_limited?(client_id)
+    key = "#{@key_prefix}:#{client_id}"
+    now = Time.now.to_f
+    window_start = now - @window_seconds
+    
+    # Redis Lua script for atomic sliding window
+    lua_script = <<~LUA
+      local key = KEYS[1]
+      local now = tonumber(ARGV[1])
+      local window_start = tonumber(ARGV[2])
+      local max_requests = tonumber(ARGV[3])
+      
+      -- Remove old entries
+      redis.call('ZREMRANGEBYSCORE', key, 0, window_start)
+      
+      -- Count current entries
+      local current = redis.call('ZCARD', key)
+      
+      if current < max_requests then
+        -- Add new entry
+        redis.call('ZADD', key, now, now)
+        redis.call('EXPIRE', key, ARGV[4])
+        return 0
+      else
+        return 1
+      end
+    LUA
+    
+    result = @redis.eval(lua_script, 1, key, now, window_start, @max_requests, @window_seconds + 1)
+    result == 1
+  end
+end
+```
+
+## Tiered Rate Limiting
+
+### Different Limits for Different Users
+
+```ruby
+class TieredRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @tiers = options[:tiers] || default_tiers
+    @default_tier = options[:default_tier] || "free"
+    @limiters = {}
+    
+    # Create limiter for each tier
+    @tiers.each do |tier, config|
+      @limiters[tier] = TokenBucketRateLimiter.new(nil, config)
+    end
+  end
+  
+  def call(request)
+    client_tier = determine_tier(request)
+    limiter = @limiters[client_tier] || @limiters[@default_tier]
+    
+    # Use tier-specific limiter
+    limiter_response = limiter.call(request)
+    
+    if limiter_response[:error]
+      # Add tier information to error
+      limiter_response[:error][:data][:tier] = client_tier
+      limiter_response[:error][:data][:upgrade_available] = can_upgrade?(client_tier)
+      return limiter_response
+    end
+    
+    @app.call(request)
+  end
+  
+  private
+  
+  def default_tiers
+    {
+      "free" => {
+        capacity: 10,
+        refill_rate: 0.1  # 0.1 tokens/second = 6/minute
+      },
+      "basic" => {
+        capacity: 50,
+        refill_rate: 1    # 1 token/second = 60/minute
+      },
+      "premium" => {
+        capacity: 200,
+        refill_rate: 5    # 5 tokens/second = 300/minute
+      },
+      "enterprise" => {
+        capacity: 1000,
+        refill_rate: 20   # 20 tokens/second = 1200/minute
+      }
+    }
+  end
+  
+  def determine_tier(request)
+    # Check authenticated user's subscription
+    user = request["authenticated_user"]
+    return @default_tier unless user
+    
+    user["subscription_tier"] || user["tier"] || @default_tier
+  end
+  
+  def can_upgrade?(tier)
+    tier_order = ["free", "basic", "premium", "enterprise"]
+    current_index = tier_order.index(tier) || 0
+    current_index < tier_order.length - 1
+  end
+end
+```
+
+## Adaptive Rate Limiting
+
+### Dynamic Limits Based on Load
+
+```ruby
+class AdaptiveRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @base_limit = options[:base_limit] || 100
+    @min_limit = options[:min_limit] || 10
+    @max_limit = options[:max_limit] || 1000
+    @load_threshold = options[:load_threshold] || 0.8
+    @limiters = {}
+    @mutex = Mutex.new
+    
+    # Start monitoring thread
+    start_monitoring
+  end
+  
+  def call(request)
+    client_id = identify_client(request)
+    current_limit = calculate_current_limit
+    
+    if exceeds_limit?(client_id, current_limit)
+      return adaptive_rate_limit_error(request["id"], current_limit)
+    end
+    
+    track_request(client_id)
+    @app.call(request)
+  end
+  
+  private
+  
+  def calculate_current_limit
+    # Adjust limit based on system load
+    load_average = get_system_load
+    
+    if load_average > @load_threshold
+      # Reduce limits under high load
+      reduction_factor = (load_average - @load_threshold) / (1 - @load_threshold)
+      reduced_limit = @base_limit * (1 - reduction_factor * 0.5)
+      [@min_limit, reduced_limit].max.to_i
+    else
+      # Increase limits under low load
+      increase_factor = 1 + (1 - load_average) * 0.5
+      increased_limit = @base_limit * increase_factor
+      [@max_limit, increased_limit].min.to_i
+    end
+  end
+  
+  def get_system_load
+    # Get 1-minute load average normalized by CPU count
+    if File.exist?("/proc/loadavg")
+      load_avg = File.read("/proc/loadavg").split.first.to_f
+      cpu_count = `nproc`.to_i rescue 1
+      load_avg / cpu_count
+    else
+      # Fallback for non-Linux systems
+      0.5
+    end
+  end
+  
+  def start_monitoring
+    Thread.new do
+      loop do
+        sleep 10
+        cleanup_old_requests
+        log_metrics if @app.respond_to?(:log)
+      end
+    end
+  end
+  
+  def cleanup_old_requests
+    @mutex.synchronize do
+      cutoff = Time.now - 300  # 5 minutes
+      @limiters.each do |client_id, requests|
+        requests.reject! { |time| time < cutoff }
+        @limiters.delete(client_id) if requests.empty?
+      end
+    end
+  end
+end
+```
+
+## Method-Specific Rate Limiting
+
+### Different Limits for Different Operations
+
+```ruby
+class MethodRateLimiter < Tsikol::Middleware
+  def initialize(app, options = {})
+    @app = app
+    @method_limits = options[:method_limits] || default_method_limits
+    @default_limit = options[:default_limit] || 60
+    @window_seconds = options[:window_seconds] || 60
+    @counters = Hash.new { |h, k| h[k] = Hash.new { |h2, k2| h2[k2] = [] } }
+    @mutex = Mutex.new
+  end
+  
+  def call(request)
+    method = request["method"]
+    client_id = identify_client(request)
+    limit = @method_limits[method] || @default_limit
+    
+    if exceeds_method_limit?(client_id, method, limit)
+      return method_rate_limit_error(request["id"], method, limit)
+    end
+    
+    track_method_request(client_id, method)
+    @app.call(request)
+  end
+  
+  private
+  
+  def default_method_limits
+    {
+      # Expensive operations
+      "sampling/sample" => 10,
+      "tools/call" => 60,
+      
+      # Moderate operations
+      "completion/complete" => 100,
+      "resources/read" => 100,
+      
+      # Cheap operations
+      "tools/list" => 300,
+      "resources/list" => 300,
+      "prompts/list" => 300,
+      
+      # Very cheap
+      "ping" => 600
+    }
+  end
+  
+  def exceeds_method_limit?(client_id, method, limit)
+    @mutex.synchronize do
+      clean_old_requests(client_id, method)
+      @counters[client_id][method].size >= limit
+    end
+  end
+  
+  def track_method_request(client_id, method)
+    @mutex.synchronize do
+      @counters[client_id][method] << Time.now
+    end
+  end
+  
+  def clean_old_requests(client_id, method)
+    cutoff = Time.now - @window_seconds
+    @counters[client_id][method].reject! { |time| time < cutoff }
+  end
+  
+  def method_rate_limit_error(id, method, limit)
+    {
+      jsonrpc: "2.0",
+      id: id,
+      error: {
+        code: -32005,
+        message: "Rate limit exceeded for #{method}",
+        data: {
+          method: method,
+          limit: limit,
+          window: @window_seconds,
+          retry_after: @window_seconds
+        }
+      }
+    }
+  end
+end
+```
+
+## Testing Rate Limiting
+
+```ruby
+require 'minitest/autorun'
+
+class RateLimitTest < Minitest::Test
+  def setup
+    @server = create_rate_limited_server
+    @client = Tsikol::TestHelpers::TestClient.new(@server)
+  end
+  
+  def test_allows_requests_under_limit
+    3.times do |i|
+      response = @client.call_tool("echo", { "message" => "Test #{i}" })
+      assert_successful_response(response)
+    end
+  end
+  
+  def test_blocks_requests_over_limit
+    # Make requests up to limit
+    5.times do
+      @client.call_tool("echo", { "message" => "Test" })
+    end
+    
+    # Next request should be rate limited
+    response = @client.call_tool("echo", { "message" => "Over limit" })
+    assert_error_response(response, -32005)
+    assert_match /rate limit/i, response[:error][:message]
+  end
+  
+  def test_resets_after_window
+    # Exhaust limit
+    5.times { @client.call_tool("echo", { "message" => "Test" }) }
+    
+    # Wait for window to reset
+    sleep(2)  # Assuming 1-second window in test
+    
+    # Should allow new requests
+    response = @client.call_tool("echo", { "message" => "After reset" })
+    assert_successful_response(response)
+  end
+  
+  def test_token_bucket_allows_bursts
+    server = Tsikol::Server.new(name: "test")
+    server.use TokenBucketRateLimiter, capacity: 10, refill_rate: 1
+    server.tool "echo" do |message:| message end
+    
+    client = Tsikol::TestHelpers::TestClient.new(server)
+    
+    # Burst of 10 requests should succeed
+    10.times do
+      response = client.call_tool("echo", { "message" => "Burst" })
+      assert_successful_response(response)
+    end
+    
+    # 11th request should fail
+    response = client.call_tool("echo", { "message" => "Over capacity" })
+    assert_error_response(response, -32005)
+  end
+  
+  private
+  
+  def create_rate_limited_server
+    server = Tsikol::Server.new(name: "test")
+    server.use BasicRateLimiter, max_requests: 5, window_seconds: 1
+    server.tool "echo" do |message:| message end
+    server
+  end
+end
+```
+
+## Best Practices
+
+1. **Choose the right algorithm** for your use case
+2. **Consider distributed systems** for multi-server deployments
+3. **Provide clear error messages** with retry information
+4. **Monitor rate limiting metrics** to tune limits
+5. **Different limits for different operations** based on cost
+6. **Graceful degradation** under high load
+7. **Consider user tiers** for SaaS applications
+8. **Test rate limiting** thoroughly including edge cases
+
+## Integration Example
+
+```ruby
+# Complete server with rate limiting
+Tsikol.start(name: "production-server") do
+  # Authentication first
+  use JWTAuthMiddleware, secret: ENV['JWT_SECRET']
+  
+  # Then rate limiting based on authenticated user
+  use TieredRateLimiter, tiers: {
+    "free" => { capacity: 10, refill_rate: 0.1 },
+    "pro" => { capacity: 100, refill_rate: 2 },
+    "enterprise" => { capacity: 1000, refill_rate: 20 }
+  }
+  
+  # Method-specific limits
+  use MethodRateLimiter, method_limits: {
+    "sampling/sample" => 5,    # AI operations
+    "tools/call" => 60,        # General tools
+    "resources/read" => 200    # Read operations
+  }
+  
+  # Adaptive limiting under load
+  use AdaptiveRateLimiter, 
+      base_limit: 100,
+      load_threshold: 0.7
+  
+  # Your application
+  tool AiTool
+  resource DataResource
+  prompt Assistant
+end
+```
+
+## Next Steps
+
+- Implement [Authentication](authentication.md) to identify users
+- Add [Logging](logging.md) for monitoring
+- Set up [Caching](caching.md) to reduce load
+- Review [Error Handling](error-handling.md) for rate limit errors