Make rate limit constants configurable via environment variables (#139)

Add support for configuring rate limits via environment variables:
- MKP_RATE_LIMIT_DEFAULT: default rate limit (default: 60)
- MKP_RATE_LIMIT_READ: read operations limit (default: 120)
- MKP_RATE_LIMIT_WRITE: write operations limit (default: 30)

This resolves the TODO comment in pkg/ratelimit/config.go and allows
operators to tune rate limits based on their deployment needs.

Signed-off-by: majiayu000 <1835304752@qq.com>

Author: majiayu000

README.md

@@ -536,6 +536,17 @@ via the command line flag:
 ./build/mkp-server --enable-rate-limiting=false
 ```
 
+Rate limits can be customized via environment variables:
+
+- `MKP_RATE_LIMIT_DEFAULT`: Default rate limit (default: 60)
+- `MKP_RATE_LIMIT_READ`: Read operations rate limit (default: 120)
+- `MKP_RATE_LIMIT_WRITE`: Write operations rate limit (default: 30)
+
+```bash
+# Run with custom rate limits
+MKP_RATE_LIMIT_READ=200 MKP_RATE_LIMIT_WRITE=50 ./build/mkp-server
+```
+
 ## Development
 
 ### Running tests

pkg/ratelimit/config.go

@@ -1,43 +1,74 @@
 // Package ratelimit provides rate limiting functionality for the MCP server
 package ratelimit
 
-import "github.com/StacklokLabs/mkp/pkg/types"
+import (
+	"os"
+	"strconv"
+
+	"github.com/StacklokLabs/mkp/pkg/types"
+)
 
-// TODO: make these constants configurable
 const (
-	defaultLimit = 60
-	readLimit    = 120 // 120 requests per minute (2 per second)
-	writeLimit   = 30  // 30 requests per minute (0.5 per second)
-	DefaultTool  = "default"
+	// DefaultTool is the key for the default rate limit
+	DefaultTool = "default"
+
+	// Default values for rate limits (requests per minute)
+	defaultDefaultLimit = 60
+	defaultReadLimit    = 120 // 120 requests per minute (2 per second)
+	defaultWriteLimit   = 30  // 30 requests per minute (0.5 per second)
 )
 
-// DefaultConfig defines the default rate limits for different tools
-var DefaultConfig = map[string]int{
-	// Read operations - higher limits
-	types.ListResourcesToolName: readLimit,
-	types.GetResourceToolName:   readLimit,
+// Config holds the rate limiting configuration
+type Config struct {
+	// DefaultLimit is the default rate limit for tools not explicitly configured
+	DefaultLimit int
+	// ReadLimit is the rate limit for read operations (list_resources, get_resource)
+	ReadLimit int
+	// WriteLimit is the rate limit for write operations (apply_resource, delete_resource)
+	WriteLimit int
+}
 
-	// Write operations - lower limits
-	types.ApplyResourceToolName:  writeLimit,
-	types.DeleteResourceToolName: writeLimit,
+// NewDefaultConfig returns a Config with default values, optionally overridden by environment variables:
+//   - MKP_RATE_LIMIT_DEFAULT: default rate limit (default: 60)
+//   - MKP_RATE_LIMIT_READ: read operations rate limit (default: 120)
+//   - MKP_RATE_LIMIT_WRITE: write operations rate limit (default: 30)
+func NewDefaultConfig() *Config {
+	return &Config{
+		DefaultLimit: getEnvInt("MKP_RATE_LIMIT_DEFAULT", defaultDefaultLimit),
+		ReadLimit:    getEnvInt("MKP_RATE_LIMIT_READ", defaultReadLimit),
+		WriteLimit:   getEnvInt("MKP_RATE_LIMIT_WRITE", defaultWriteLimit),
+	}
+}
 
-	// Default for any other tool
-	DefaultTool: defaultLimit,
+// getEnvInt returns the integer value of an environment variable or a default value
+func getEnvInt(key string, defaultValue int) int {
+	if value := os.Getenv(key); value != "" {
+		if intValue, err := strconv.Atoi(value); err == nil && intValue > 0 {
+			return intValue
+		}
+	}
+	return defaultValue
 }
 
 // GetDefaultRateLimiter returns a RateLimiter with default configuration
 func GetDefaultRateLimiter() *RateLimiter {
-	options := []RateLimiterOption{
-		WithDefaultLimit(DefaultConfig[DefaultTool]),
+	return GetRateLimiterWithConfig(nil)
+}
+
+// GetRateLimiterWithConfig returns a RateLimiter with the given configuration.
+// If config is nil, default configuration is used.
+func GetRateLimiterWithConfig(config *Config) *RateLimiter {
+	if config == nil {
+		config = NewDefaultConfig()
 	}
 
-	// Add tool-specific limits
-	for tool, limit := range DefaultConfig {
-		if tool != DefaultTool {
-			options = append(options, WithToolLimit(tool, limit))
-		}
+	options := []RateLimiterOption{
+		WithDefaultLimit(config.DefaultLimit),
+		WithToolLimit(types.ListResourcesToolName, config.ReadLimit),
+		WithToolLimit(types.GetResourceToolName, config.ReadLimit),
+		WithToolLimit(types.ApplyResourceToolName, config.WriteLimit),
+		WithToolLimit(types.DeleteResourceToolName, config.WriteLimit),
 	}
 
-	limiter := NewRateLimiter(options...)
-	return limiter
+	return NewRateLimiter(options...)
 }

pkg/ratelimit/ratelimit.go

@@ -60,7 +60,7 @@ func WithTimeWindow(window time.Duration) RateLimiterOption {
 func NewRateLimiter(opts ...RateLimiterOption) *RateLimiter {
 	rl := &RateLimiter{
 		limits:        make(map[string]int),
-		defaultLimit:  defaultLimit,
+		defaultLimit:  defaultDefaultLimit,
 		requestCounts: make(map[string]map[string]*windowCounter),
 	}
 

pkg/ratelimit/ratelimit_test.go

@@ -2,6 +2,7 @@ package ratelimit
 
 import (
 	"context"
+	"os"
 	"testing"
 	"time"
 
@@ -14,7 +15,7 @@ func TestRateLimiterCreation(t *testing.T) {
 	// Test with default options
 	limiter := NewRateLimiter()
 	assert.NotNil(t, limiter)
-	assert.Equal(t, defaultLimit, limiter.defaultLimit)
+	assert.Equal(t, defaultDefaultLimit, limiter.defaultLimit)
 
 	// Test with custom options
 	customLimiter := NewRateLimiter(
@@ -258,3 +259,71 @@ func TestToolSpecificLimits(t *testing.T) {
 	assert.True(t, result.IsError)
 	assert.Contains(t, result.Content[0].(mcp.TextContent).Text, "Rate limit exceeded")
 }
+
+func TestNewDefaultConfig(t *testing.T) {
+	// Test default values
+	config := NewDefaultConfig()
+	assert.Equal(t, defaultDefaultLimit, config.DefaultLimit)
+	assert.Equal(t, defaultReadLimit, config.ReadLimit)
+	assert.Equal(t, defaultWriteLimit, config.WriteLimit)
+}
+
+func TestConfigFromEnv(t *testing.T) {
+	// Save original env values and restore after test
+	origDefault := os.Getenv("MKP_RATE_LIMIT_DEFAULT")
+	origRead := os.Getenv("MKP_RATE_LIMIT_READ")
+	origWrite := os.Getenv("MKP_RATE_LIMIT_WRITE")
+	defer func() {
+		os.Setenv("MKP_RATE_LIMIT_DEFAULT", origDefault)
+		os.Setenv("MKP_RATE_LIMIT_READ", origRead)
+		os.Setenv("MKP_RATE_LIMIT_WRITE", origWrite)
+	}()
+
+	// Set custom env values
+	os.Setenv("MKP_RATE_LIMIT_DEFAULT", "100")
+	os.Setenv("MKP_RATE_LIMIT_READ", "200")
+	os.Setenv("MKP_RATE_LIMIT_WRITE", "50")
+
+	config := NewDefaultConfig()
+	assert.Equal(t, 100, config.DefaultLimit)
+	assert.Equal(t, 200, config.ReadLimit)
+	assert.Equal(t, 50, config.WriteLimit)
+}
+
+func TestConfigFromEnvInvalidValues(t *testing.T) {
+	// Save original env values and restore after test
+	origDefault := os.Getenv("MKP_RATE_LIMIT_DEFAULT")
+	defer os.Setenv("MKP_RATE_LIMIT_DEFAULT", origDefault)
+
+	// Test with invalid value (non-numeric)
+	os.Setenv("MKP_RATE_LIMIT_DEFAULT", "invalid")
+	config := NewDefaultConfig()
+	assert.Equal(t, defaultDefaultLimit, config.DefaultLimit)
+
+	// Test with invalid value (zero)
+	os.Setenv("MKP_RATE_LIMIT_DEFAULT", "0")
+	config = NewDefaultConfig()
+	assert.Equal(t, defaultDefaultLimit, config.DefaultLimit)
+
+	// Test with invalid value (negative)
+	os.Setenv("MKP_RATE_LIMIT_DEFAULT", "-10")
+	config = NewDefaultConfig()
+	assert.Equal(t, defaultDefaultLimit, config.DefaultLimit)
+}
+
+func TestGetRateLimiterWithConfig(t *testing.T) {
+	// Test with nil config (should use defaults)
+	limiter := GetRateLimiterWithConfig(nil)
+	assert.NotNil(t, limiter)
+	assert.Equal(t, defaultDefaultLimit, limiter.defaultLimit)
+
+	// Test with custom config
+	config := &Config{
+		DefaultLimit: 100,
+		ReadLimit:    200,
+		WriteLimit:   50,
+	}
+	limiter = GetRateLimiterWithConfig(config)
+	assert.NotNil(t, limiter)
+	assert.Equal(t, 100, limiter.defaultLimit)
+}