vMCP rate-limit middleware wiring

Signed-off-by: Sanskarzz <sanskar.gur@gmail.com>

Author: Sanskarzz

pkg/ratelimit/middleware.go

@@ -52,6 +52,17 @@ type rateLimitMiddleware struct {
 	client  redis.UniversalClient
 }
 
+// ToolNameResolver resolves the rate-limit tool name from a parsed MCP request.
+type ToolNameResolver func(*mcp.ParsedMCPRequest) string
+
+// DefaultToolNameResolver uses the parsed MCP resource ID as the rate-limit tool name.
+func DefaultToolNameResolver(parsed *mcp.ParsedMCPRequest) string {
+	if parsed == nil {
+		return ""
+	}
+	return parsed.ResourceID
+}
+
 // Handler returns the middleware function used by the proxy.
 func (m *rateLimitMiddleware) Handler() types.MiddlewareFunction {
 	return m.handler
@@ -99,16 +110,19 @@ func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRun
 	}
 
 	mw := &rateLimitMiddleware{
-		handler: rateLimitHandler(limiter),
+		handler: NewMiddleware(limiter, nil),
 		client:  client,
 	}
 	runner.AddMiddleware(MiddlewareType, mw)
 	return nil
 }
 
-// rateLimitHandler returns a middleware function that enforces rate limits
+// NewMiddleware returns a middleware function that enforces rate limits
 // on tools/call requests.
-func rateLimitHandler(limiter Limiter) types.MiddlewareFunction {
+func NewMiddleware(limiter Limiter, resolveToolName ToolNameResolver) types.MiddlewareFunction {
+	if resolveToolName == nil {
+		resolveToolName = DefaultToolNameResolver
+	}
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			// Rate limits only apply to parsed tools/call requests.
@@ -127,7 +141,7 @@ func rateLimitHandler(limiter Limiter) types.MiddlewareFunction {
 			if identity, ok := auth.IdentityFromContext(r.Context()); ok {
 				userID = identity.Subject
 			}
-			decision, err := limiter.Allow(r.Context(), parsed.ResourceID, userID)
+			decision, err := limiter.Allow(r.Context(), resolveToolName(parsed), userID)
 			if err != nil {
 				slog.Warn("rate limit check failed, allowing request", "error", err)
 				next.ServeHTTP(w, r)
@@ -142,6 +156,11 @@ func rateLimitHandler(limiter Limiter) types.MiddlewareFunction {
 	}
 }
 
+// rateLimitHandler returns the default rate-limit middleware used by tests and legacy callers.
+func rateLimitHandler(limiter Limiter) types.MiddlewareFunction {
+	return NewMiddleware(limiter, nil)
+}
+
 // writeRateLimited writes an HTTP 429 response with a JSON-RPC error body.
 func writeRateLimited(w http.ResponseWriter, requestID any, retryAfter time.Duration) {
 	retrySeconds := int(math.Ceil(retryAfter.Seconds()))

pkg/vmcp/cli/serve.go

@@ -376,6 +376,7 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 
 	serverCfg := &vmcpserver.Config{
 		Name:                    vmcpCfg.Name,
+		Namespace:               vmcpNamespace(),
 		Version:                 versions.Version,
 		GroupRef:                vmcpCfg.Group,
 		Host:                    cfg.Host,
@@ -394,6 +395,7 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 		OptimizerConfig:         optCfg,
 		SessionFactory:          sessionFactory,
 		SessionStorage:          vmcpCfg.SessionStorage,
+		RateLimiting:            vmcpCfg.RateLimiting,
 	}
 
 	// Assign Watcher only when backendWatcher is non-nil. A typed nil
@@ -529,6 +531,14 @@ func generateQuickModeConfig(groupRef string) (*config.Config, error) {
 	return cfg, nil
 }
 
+func vmcpNamespace() string {
+	namespace := os.Getenv("VMCP_NAMESPACE")
+	if namespace == "" {
+		return "local"
+	}
+	return namespace
+}
+
 // loadAuthServerConfig loads the auth server RunConfig from a sibling file
 // alongside the main config. The operator serializes authserver.RunConfig as a
 // separate ConfigMap key (authserver-config.yaml).

pkg/vmcp/server/ratelimit.go

@@ -0,0 +1,81 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package server
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+
+	mcpparser "github.com/stacklok/toolhive/pkg/mcp"
+	"github.com/stacklok/toolhive/pkg/ratelimit"
+	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
+)
+
+const rateLimitRedisPingTimeout = 5 * time.Second
+
+func (s *Server) buildRateLimitMiddleware(
+	ctx context.Context,
+) (func(http.Handler) http.Handler, func(context.Context) error, error) {
+	if s.config.RateLimiting == nil {
+		return nil, nil, nil
+	}
+	if s.config.SessionStorage == nil || s.config.SessionStorage.Provider != "redis" {
+		return nil, nil, fmt.Errorf("rate limiting requires Redis session storage")
+	}
+	if s.config.SessionStorage.Address == "" {
+		return nil, nil, fmt.Errorf("rate limiting requires Redis session storage address")
+	}
+
+	client := redis.NewClient(&redis.Options{
+		Addr:     s.config.SessionStorage.Address,
+		DB:       int(s.config.SessionStorage.DB),
+		Password: os.Getenv(vmcpconfig.RedisPasswordEnvVar),
+	})
+
+	pingCtx, cancel := context.WithTimeout(ctx, rateLimitRedisPingTimeout)
+	defer cancel()
+	if err := client.Ping(pingCtx).Err(); err != nil {
+		_ = client.Close()
+		return nil, nil, fmt.Errorf("rate limit middleware: failed to connect to Redis at %s: %w",
+			s.config.SessionStorage.Address, err)
+	}
+
+	limiter, err := ratelimit.NewLimiter(client, s.config.Namespace, s.config.Name, s.config.RateLimiting)
+	if err != nil {
+		_ = client.Close()
+		return nil, nil, fmt.Errorf("failed to create rate limiter: %w", err)
+	}
+
+	cleanup := func(context.Context) error {
+		return client.Close()
+	}
+	return ratelimit.NewMiddleware(limiter, s.rateLimitToolName), cleanup, nil
+}
+
+func (s *Server) rateLimitToolName(parsed *mcpparser.ParsedMCPRequest) string {
+	if parsed == nil {
+		return ""
+	}
+	toolName := parsed.ResourceID
+	if !s.optimizerEnabled() || toolName != "call_tool" {
+		return toolName
+	}
+	if parsed.Arguments == nil {
+		return toolName
+	}
+	innerToolName, ok := parsed.Arguments["tool_name"].(string)
+	if !ok || innerToolName == "" {
+		return toolName
+	}
+	return innerToolName
+}
+
+func (s *Server) optimizerEnabled() bool {
+	return s.config.OptimizerConfig != nil || s.config.OptimizerFactory != nil
+}