Configure rate limits on VirtualMCPServer PR A (#5079)

* Configure rate limits on VirtualMCPServer

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

* Only schema changes, no runtime behavior

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

* Address vMCP rate limit review feedback

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

* removed wrapper and reverted path marshall directly

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

* moved onto config.Config as spec.config.rateLimiting

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

* removed the duplicate rate-limit structs from pkg/vmcp/config

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

---------

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

Author: Sanskarzz

cmd/thv-operator/Taskfile.yml

@@ -168,7 +168,7 @@ tasks:
         platforms: [windows]
         ignore_error: true   # Windows has no mkdir -p, so just ignore error if it exists
       - go install sigs.k8s.io/controller-tools/cmd/controller-gen@v0.17.3
-      - $(go env GOPATH)/bin/controller-gen object:headerFile="hack/boilerplate.go.txt" paths="./cmd/thv-operator/..." paths="./pkg/json/..." paths="./pkg/vmcp/config/..." paths="./pkg/vmcp/auth/types/..." paths="./pkg/telemetry/..." paths="./pkg/audit/..."
+      - $(go env GOPATH)/bin/controller-gen object:headerFile="hack/boilerplate.go.txt" paths="./cmd/thv-operator/..." paths="./pkg/json/..." paths="./pkg/ratelimit/types/..." paths="./pkg/vmcp/config/..." paths="./pkg/vmcp/auth/types/..." paths="./pkg/telemetry/..." paths="./pkg/audit/..."
 
   operator-manifests:
     desc: Generate WebhookConfiguration, ClusterRole and CustomResourceDefinition objects
@@ -286,6 +286,7 @@ tasks:
     sources:
       - '{{.ROOT_DIR}}/cmd/thv-operator/config/crd/bases/**/*.yaml'
       - '{{.ROOT_DIR}}/cmd/thv-operator/api/**/*.go'
+      - '{{.ROOT_DIR}}/pkg/ratelimit/types/*.go'
       - '{{.ROOT_DIR}}/pkg/vmcp/config/*.go'
       - '{{.ROOT_DIR}}/pkg/vmcp/auth/types/*.go'
       - '{{.ROOT_DIR}}/pkg/telemetry/*.go'

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"
+
+	ratelimittypes "github.com/stacklok/toolhive/pkg/ratelimit/types"
 )
 
 // Condition types for MCPServer
@@ -373,7 +375,7 @@ type MCPServerSpec struct {
 	// RateLimiting defines rate limiting configuration for the MCP server.
 	// Requires Redis session storage to be configured for distributed rate limiting.
 	// +optional
-	RateLimiting *RateLimitConfig `json:"rateLimiting,omitempty"`
+	RateLimiting *ratelimittypes.RateLimitConfig `json:"rateLimiting,omitempty"`
 }
 
 // ResourceOverrides defines overrides for annotations and labels on created resources
@@ -525,69 +527,17 @@ 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"
-//
-//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.
-	// +optional
-	Shared *RateLimitBucket `json:"shared,omitempty"`
-
-	// PerUser is a token bucket applied independently to each authenticated user
-	// at the server level. Requires authentication to be enabled.
-	// Each unique userID creates Redis keys that expire after 2x refillPeriod.
-	// Memory formula: unique_users_per_TTL_window * (1 + num_tools_with_per_user_limits) keys.
-	// +optional
-	PerUser *RateLimitBucket `json:"perUser,omitempty"`
-
-	// Tools defines per-tool rate limit overrides.
-	// Each entry applies additional rate limits to calls targeting a specific tool name.
-	// A request must pass both the server-level limit and the per-tool limit.
-	// +listType=map
-	// +listMapKey=name
-	// +optional
-	Tools []ToolRateLimitConfig `json:"tools,omitempty"`
-}
+// +gendoc
+type RateLimitConfig = ratelimittypes.RateLimitConfig
 
 // RateLimitBucket defines a token bucket configuration with a maximum capacity
-// and a refill period. Used by both shared (global) and per-user rate limits.
-type RateLimitBucket struct {
-	// MaxTokens is the maximum number of tokens (bucket capacity).
-	// This is also the burst size: the maximum number of requests that can be served
-	// instantaneously before the bucket is depleted.
-	// +kubebuilder:validation:Required
-	// +kubebuilder:validation:Minimum=1
-	MaxTokens int32 `json:"maxTokens"`
-
-	// RefillPeriod is the duration to fully refill the bucket from zero to maxTokens.
-	// The effective refill rate is maxTokens / refillPeriod tokens per second.
-	// Format: Go duration string (e.g., "1m0s", "30s", "1h0m0s").
-	// +kubebuilder:validation:Required
-	RefillPeriod metav1.Duration `json:"refillPeriod"`
-}
+// and a refill period. Used by both shared and per-user rate limits.
+// +gendoc
+type RateLimitBucket = ratelimittypes.RateLimitBucket
 
 // ToolRateLimitConfig defines rate limits for a specific tool.
-// At least one of 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"
-//
-//nolint:lll // kubebuilder marker exceeds line length
-type ToolRateLimitConfig struct {
-	// Name is the MCP tool name this limit applies to.
-	// +kubebuilder:validation:Required
-	// +kubebuilder:validation:MinLength=1
-	Name string `json:"name"`
-
-	// Shared token bucket for this specific tool.
-	// +optional
-	Shared *RateLimitBucket `json:"shared,omitempty"`
-
-	// PerUser token bucket configuration for this tool.
-	// +optional
-	PerUser *RateLimitBucket `json:"perUser,omitempty"`
-}
+// +gendoc
+type ToolRateLimitConfig = ratelimittypes.ToolRateLimitConfig
 
 // Permission profile types
 const (

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

@@ -11,6 +11,8 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
 )
 
 func TestSessionStorageConfigJSONRoundtrip(t *testing.T) {
@@ -116,6 +118,46 @@ 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",
+		},
+		Config: vmcpconfig.Config{
+			RateLimiting: &RateLimitConfig{
+				Shared: &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",
+						Shared: &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, `"shared"`)
+	assert.Contains(t, out, `"perUser"`)
+	assert.Contains(t, out, `"backend_a_echo"`)
+	assert.Contains(t, out, `"config":{"rateLimiting"`)
+}
+
 func TestMCPServerSpecScalingFieldsJSONRoundtrip(t *testing.T) {
 	t.Parallel()
 

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

@@ -16,6 +16,10 @@ import (
 
 // VirtualMCPServerSpec defines the desired state of VirtualMCPServer
 //
+// +kubebuilder:validation:XValidation:rule="!has(self.config) || !has(self.config.rateLimiting) || (has(self.sessionStorage) && self.sessionStorage.provider == 'redis')",message="config.rateLimiting requires sessionStorage with provider 'redis'"
+// +kubebuilder:validation:XValidation:rule="!(has(self.config) && has(self.config.rateLimiting) && has(self.config.rateLimiting.perUser)) || (has(self.incomingAuth) && self.incomingAuth.type == 'oidc')",message="config.rateLimiting.perUser requires incomingAuth.type oidc"
+// +kubebuilder:validation:XValidation:rule="!has(self.config) || !has(self.config.rateLimiting) || !has(self.config.rateLimiting.tools) || self.config.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.

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

@@ -507,6 +507,11 @@ 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))
+	assert.Equal(t, "test-vmcp", cfg.Name)
+	assert.Equal(t, "test-group", cfg.Group)
 }
 
 // TestSetAuthConfigConditions tests that auth config conditions reflect the current state

cmd/thv-operator/controllers/virtualmcpserver_vmcpconfig_test.go

@@ -1601,6 +1601,51 @@ func TestConverter_SessionStorage(t *testing.T) {
 	}
 }
 
+func TestConverter_RateLimitingPassThrough(t *testing.T) {
+	t.Parallel()
+
+	vmcpServer := &mcpv1beta1.VirtualMCPServer{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-vmcp",
+			Namespace: "default",
+		},
+		Spec: mcpv1beta1.VirtualMCPServerSpec{
+			GroupRef: &mcpv1beta1.MCPGroupRef{Name: "test-group"},
+			Config: vmcpconfig.Config{
+				RateLimiting: &mcpv1beta1.RateLimitConfig{
+					PerUser: &mcpv1beta1.RateLimitBucket{
+						MaxTokens:    2,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+					Tools: []mcpv1beta1.ToolRateLimitConfig{
+						{
+							Name: "backend_a_echo",
+							Shared: &mcpv1beta1.RateLimitBucket{
+								MaxTokens:    5,
+								RefillPeriod: metav1.Duration{Duration: 30 * time.Second},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	converter := newTestConverter(t, newNoOpMockResolver(t))
+	ctx := log.IntoContext(context.Background(), logr.Discard())
+
+	config, _, err := converter.Convert(ctx, vmcpServer, nil)
+	require.NoError(t, err)
+	require.NotNil(t, config)
+	require.NotNil(t, config.RateLimiting)
+
+	assert.EqualValues(t, 2, config.RateLimiting.PerUser.MaxTokens)
+	require.Len(t, config.RateLimiting.Tools, 1)
+	assert.Equal(t, "backend_a_echo", config.RateLimiting.Tools[0].Name)
+	require.NotNil(t, config.RateLimiting.Tools[0].Shared)
+	assert.EqualValues(t, 5, config.RateLimiting.Tools[0].Shared.MaxTokens)
+}
+
 func TestDeriveAllowedAudiences(t *testing.T) {
 	t.Parallel()
 

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

@@ -466,6 +466,7 @@ var _ = Describe("EmbeddingServer Controller Update Tests", func() {
 				Expect(k8sClient.Create(ctx, embeddingServer)).To(Succeed())
 				Eventually(func(g Gomega) {
 					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(embeddingServer), &appsv1.StatefulSet{})).To(Succeed())
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(embeddingServer), &corev1.Service{})).To(Succeed())
 				}, timeout, interval).Should(Succeed())
 			})