Configure rate limits on VirtualMCPServer

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

Author: Sanskarzz

cmd/thv-operator/api/v1beta1/mcpserver_types.go

@@ -7,6 +7,8 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+
+	rlconfig "github.com/stacklok/toolhive/pkg/ratelimit/config"
 )
 
 // Condition types for MCPServer
@@ -505,11 +507,15 @@ type SessionStorageConfig struct {
 // RateLimitConfig defines rate limiting configuration for an MCP server.
 // At least one of shared, perUser, or tools must be configured.
 //
-// +kubebuilder:validation:XValidation:rule="has(self.shared) || has(self.perUser) || (has(self.tools) && size(self.tools) > 0)",message="at least one of shared, perUser, or tools must be configured"
+// +kubebuilder:validation:XValidation:rule="has(self.global) || has(self.shared) || has(self.perUser) || (has(self.tools) && size(self.tools) > 0)",message="at least one of global, shared, perUser, or tools must be configured"
 //
 //nolint:lll // CEL validation rules exceed line length limit
 type RateLimitConfig struct {
-	// Shared is a token bucket shared across all users for the entire server.
+	// Global is a token bucket shared across all users for the entire server.
+	// +optional
+	Global *RateLimitBucket `json:"global,omitempty"`
+
+	// Shared is a deprecated alias for Global. Use global instead.
 	// +optional
 	Shared *RateLimitBucket `json:"shared,omitempty"`
 
@@ -547,9 +553,9 @@ type RateLimitBucket struct {
 }
 
 // ToolRateLimitConfig defines rate limits for a specific tool.
-// At least one of shared or perUser must be configured.
+// At least one of global, shared, or perUser must be configured.
 //
-// +kubebuilder:validation:XValidation:rule="has(self.shared) || has(self.perUser)",message="at least one of shared or perUser must be configured"
+// +kubebuilder:validation:XValidation:rule="has(self.global) || has(self.shared) || has(self.perUser)",message="at least one of global, shared, or perUser must be configured"
 //
 //nolint:lll // kubebuilder marker exceeds line length
 type ToolRateLimitConfig struct {
@@ -558,7 +564,11 @@ type ToolRateLimitConfig struct {
 	// +kubebuilder:validation:MinLength=1
 	Name string `json:"name"`
 
-	// Shared token bucket for this specific tool.
+	// Global token bucket for this specific tool.
+	// +optional
+	Global *RateLimitBucket `json:"global,omitempty"`
+
+	// Shared is a deprecated alias for Global. Use global instead.
 	// +optional
 	Shared *RateLimitBucket `json:"shared,omitempty"`
 
@@ -567,6 +577,40 @@ type ToolRateLimitConfig struct {
 	PerUser *RateLimitBucket `json:"perUser,omitempty"`
 }
 
+// ToInternal converts the Kubernetes API rate limit config to the runtime-neutral representation.
+func (in *RateLimitConfig) ToInternal() *rlconfig.Config {
+	if in == nil {
+		return nil
+	}
+	out := &rlconfig.Config{
+		Global:  rateLimitBucketToInternal(in.Global),
+		Shared:  rateLimitBucketToInternal(in.Shared),
+		PerUser: rateLimitBucketToInternal(in.PerUser),
+	}
+	if len(in.Tools) > 0 {
+		out.Tools = make([]rlconfig.ToolConfig, 0, len(in.Tools))
+		for _, tool := range in.Tools {
+			out.Tools = append(out.Tools, rlconfig.ToolConfig{
+				Name:    tool.Name,
+				Global:  rateLimitBucketToInternal(tool.Global),
+				Shared:  rateLimitBucketToInternal(tool.Shared),
+				PerUser: rateLimitBucketToInternal(tool.PerUser),
+			})
+		}
+	}
+	return out
+}
+
+func rateLimitBucketToInternal(in *RateLimitBucket) *rlconfig.Bucket {
+	if in == nil {
+		return nil
+	}
+	return &rlconfig.Bucket{
+		MaxTokens:    in.MaxTokens,
+		RefillPeriod: in.RefillPeriod,
+	}
+}
+
 // Permission profile types
 const (
 	// PermissionProfileTypeBuiltin is the type for built-in permission profiles

cmd/thv-operator/api/v1beta1/mcpserver_types_test.go

@@ -75,6 +75,13 @@ func TestRateLimitConfigJSONRoundtrip(t *testing.T) {
 		input    RateLimitConfig
 		wantJSON string
 	}{
+		{
+			name: "global only",
+			input: RateLimitConfig{
+				Global: &RateLimitBucket{MaxTokens: 100, RefillPeriod: metav1.Duration{Duration: time.Minute}},
+			},
+			wantJSON: `{"global":{"maxTokens":100,"refillPeriod":"1m0s"}}`,
+		},
 		{
 			name: "shared only",
 			input: RateLimitConfig{
@@ -116,6 +123,44 @@ func TestRateLimitConfigJSONRoundtrip(t *testing.T) {
 	}
 }
 
+func TestVirtualMCPServerSpecRateLimitingJSONRoundtrip(t *testing.T) {
+	t.Parallel()
+
+	spec := VirtualMCPServerSpec{
+		IncomingAuth: &IncomingAuthConfig{Type: "oidc"},
+		GroupRef:     &MCPGroupRef{Name: "group-a"},
+		SessionStorage: &SessionStorageConfig{
+			Provider: "redis",
+			Address:  "redis.default.svc.cluster.local:6379",
+		},
+		RateLimiting: &RateLimitConfig{
+			Global: &RateLimitBucket{MaxTokens: 10, RefillPeriod: metav1.Duration{Duration: time.Minute}},
+			PerUser: &RateLimitBucket{
+				MaxTokens:    2,
+				RefillPeriod: metav1.Duration{Duration: time.Minute},
+			},
+			Tools: []ToolRateLimitConfig{
+				{
+					Name: "backend_a_echo",
+					Global: &RateLimitBucket{
+						MaxTokens:    5,
+						RefillPeriod: metav1.Duration{Duration: 30 * time.Second},
+					},
+				},
+			},
+		},
+	}
+
+	b, err := json.Marshal(spec)
+	require.NoError(t, err)
+	out := string(b)
+	assert.Contains(t, out, `"rateLimiting"`)
+	assert.Contains(t, out, `"global"`)
+	assert.Contains(t, out, `"perUser"`)
+	assert.Contains(t, out, `"backend_a_echo"`)
+	assert.NotContains(t, out, `"config":{"rateLimiting"`)
+}
+
 func TestMCPServerSpecScalingFieldsJSONRoundtrip(t *testing.T) {
 	t.Parallel()
 

cmd/thv-operator/api/v1beta1/virtualmcpserver_types.go

@@ -15,6 +15,10 @@ import (
 
 // VirtualMCPServerSpec defines the desired state of VirtualMCPServer
 //
+// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || (has(self.sessionStorage) && self.sessionStorage.provider == 'redis')",message="rateLimiting requires sessionStorage with provider 'redis'"
+// +kubebuilder:validation:XValidation:rule="!(has(self.rateLimiting) && has(self.rateLimiting.perUser)) || (has(self.incomingAuth) && self.incomingAuth.type == 'oidc')",message="rateLimiting.perUser requires incomingAuth.type oidc"
+// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || !has(self.rateLimiting.tools) || self.rateLimiting.tools.all(t, !has(t.perUser)) || (has(self.incomingAuth) && self.incomingAuth.type == 'oidc')",message="per-tool perUser rate limiting requires incomingAuth.type oidc"
+//
 //nolint:lll // CEL validation rules exceed line length limit
 type VirtualMCPServerSpec struct {
 	// IncomingAuth configures authentication for clients connecting to the Virtual MCP server.
@@ -107,6 +111,11 @@ type VirtualMCPServerSpec struct {
 	// When nil, no session storage is configured.
 	// +optional
 	SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty"`
+
+	// RateLimiting defines rate limiting configuration for the Virtual MCP server.
+	// Requires Redis session storage to be configured for distributed rate limiting.
+	// +optional
+	RateLimiting *RateLimitConfig `json:"rateLimiting,omitempty"`
 }
 
 // EmbeddingServerRef references an existing EmbeddingServer resource by name.

cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go

@@ -459,6 +459,25 @@ func TestEnsureVmcpConfigConfigMap(t *testing.T) {
 		},
 		Spec: mcpv1beta1.VirtualMCPServerSpec{
 			GroupRef: &mcpv1beta1.MCPGroupRef{Name: "test-group"},
+			SessionStorage: &mcpv1beta1.SessionStorageConfig{
+				Provider: mcpv1beta1.SessionStorageProviderRedis,
+				Address:  "redis.default.svc.cluster.local:6379",
+			},
+			RateLimiting: &mcpv1beta1.RateLimitConfig{
+				PerUser: &mcpv1beta1.RateLimitBucket{
+					MaxTokens:    2,
+					RefillPeriod: metav1.Duration{Duration: time.Minute},
+				},
+				Tools: []mcpv1beta1.ToolRateLimitConfig{
+					{
+						Name: "backend_a_echo",
+						Global: &mcpv1beta1.RateLimitBucket{
+							MaxTokens:    5,
+							RefillPeriod: metav1.Duration{Duration: time.Minute},
+						},
+					},
+				},
+			},
 		},
 	}
 
@@ -507,6 +526,16 @@ func TestEnsureVmcpConfigConfigMap(t *testing.T) {
 	assert.Equal(t, "test-vmcp-vmcp-config", cm.Name)
 	assert.Contains(t, cm.Data, "config.yaml")
 	assert.NotEmpty(t, cm.Annotations["toolhive.stacklok.dev/content-checksum"])
+
+	var cfg vmcpconfig.Config
+	require.NoError(t, yaml.Unmarshal([]byte(cm.Data["config.yaml"]), &cfg))
+	require.NotNil(t, cfg.RateLimiting, "runtime config must include spec.rateLimiting")
+	require.NotNil(t, cfg.RateLimiting.PerUser)
+	assert.Equal(t, int32(2), cfg.RateLimiting.PerUser.MaxTokens)
+	require.Len(t, cfg.RateLimiting.Tools, 1)
+	assert.Equal(t, "backend_a_echo", cfg.RateLimiting.Tools[0].Name)
+	require.NotNil(t, cfg.RateLimiting.Tools[0].Global)
+	assert.Equal(t, int32(5), cfg.RateLimiting.Tools[0].Global.MaxTokens)
 }
 
 // TestSetAuthConfigConditions tests that auth config conditions reflect the current state

cmd/thv-operator/controllers/virtualmcpserver_vmcpconfig_test.go

@@ -140,6 +140,7 @@ func (c *Converter) Convert(
 	}
 
 	config.SessionStorage = convertSessionStorage(vmcp)
+	config.RateLimiting = vmcp.Spec.RateLimiting.ToInternal()
 
 	// Apply operational defaults (fills missing values)
 	config.EnsureOperationalDefaults()

cmd/thv-operator/pkg/vmcpconfig/converter.go

@@ -97,6 +97,42 @@ func newTestConverterWithObjects(t *testing.T, resolver oidc.Resolver, objects .
 	return converter
 }
 
+func TestConverter_RateLimitingFromTopLevelSpec(t *testing.T) {
+	t.Parallel()
+
+	vmcp := &mcpv1beta1.VirtualMCPServer{
+		ObjectMeta: metav1.ObjectMeta{Name: "test-vmcp", Namespace: "default"},
+		Spec: mcpv1beta1.VirtualMCPServerSpec{
+			GroupRef:     &mcpv1beta1.MCPGroupRef{Name: "test-group"},
+			IncomingAuth: &mcpv1beta1.IncomingAuthConfig{Type: "anonymous"},
+			RateLimiting: &mcpv1beta1.RateLimitConfig{
+				Global: &mcpv1beta1.RateLimitBucket{
+					MaxTokens:    10,
+					RefillPeriod: metav1.Duration{Duration: time.Minute},
+				},
+				Tools: []mcpv1beta1.ToolRateLimitConfig{
+					{
+						Name: "backend_a_echo",
+						Global: &mcpv1beta1.RateLimitBucket{
+							MaxTokens:    1,
+							RefillPeriod: metav1.Duration{Duration: time.Minute},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	converter := newTestConverter(t, newNoOpMockResolver(t))
+	cfg, _, err := converter.Convert(context.Background(), vmcp, nil)
+	require.NoError(t, err)
+	require.NotNil(t, cfg.RateLimiting)
+	require.NotNil(t, cfg.RateLimiting.Global)
+	assert.Equal(t, int32(10), cfg.RateLimiting.Global.MaxTokens)
+	require.Len(t, cfg.RateLimiting.Tools, 1)
+	assert.Equal(t, "backend_a_echo", cfg.RateLimiting.Tools[0].Name)
+}
+
 func TestConverter_OIDCResolution(t *testing.T) {
 	t.Parallel()
 

cmd/thv-operator/pkg/vmcpconfig/converter_test.go

@@ -5,6 +5,8 @@
 package controllers
 
 import (
+	"time"
+
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -106,4 +108,57 @@ var _ = Describe("CEL Validation for SessionStorageConfig on VirtualMCPServer",
 				Expect(err).To(HaveOccurred())
 			})
 		})
+
+		Context("rateLimiting", func() {
+			It("should reject rate limiting without redis session storage", func() {
+				vmcp := newVirtualMCPServerWithSessionStorage("vmcp-rl-no-redis", nil)
+				vmcp.Spec.RateLimiting = &mcpv1beta1.RateLimitConfig{
+					Global: &mcpv1beta1.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).To(HaveOccurred())
+				Expect(err.Error()).To(ContainSubstring("rateLimiting requires sessionStorage with provider 'redis'"))
+			})
+
+			It("should reject perUser rate limiting with anonymous auth", func() {
+				vmcp := newVirtualMCPServerWithSessionStorage("vmcp-rl-peruser-anon", &mcpv1beta1.SessionStorageConfig{
+					Provider: "redis",
+					Address:  "redis:6379",
+				})
+				vmcp.Spec.RateLimiting = &mcpv1beta1.RateLimitConfig{
+					PerUser: &mcpv1beta1.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).To(HaveOccurred())
+				Expect(err.Error()).To(ContainSubstring("rateLimiting.perUser requires incomingAuth.type oidc"))
+			})
+
+			It("should accept perUser rate limiting with oidc auth and redis session storage", func() {
+				vmcp := newVirtualMCPServerWithSessionStorage("vmcp-rl-peruser-oidc", &mcpv1beta1.SessionStorageConfig{
+					Provider: "redis",
+					Address:  "redis:6379",
+				})
+				vmcp.Spec.IncomingAuth = &mcpv1beta1.IncomingAuthConfig{
+					Type:          "oidc",
+					OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: "oidc"},
+				}
+				vmcp.Spec.RateLimiting = &mcpv1beta1.RateLimitConfig{
+					PerUser: &mcpv1beta1.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).NotTo(HaveOccurred())
+			})
+		})
 	})