Configure rate limits on VirtualMCPServer PR A

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

@@ -518,43 +518,43 @@ type SessionStorageConfig struct {
 //
 // +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
+//nolint:lll // kubebuilder marker exceeds line length
 type RateLimitConfig struct {
 	// Shared is a token bucket shared across all users for the entire server.
 	// +optional
-	Shared *RateLimitBucket `json:"shared,omitempty"`
+	Shared *RateLimitBucket `json:"shared,omitempty" yaml:"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"`
+	PerUser *RateLimitBucket `json:"perUser,omitempty" yaml:"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"`
+	Tools []ToolRateLimitConfig `json:"tools,omitempty" yaml:"tools,omitempty"`
 }
 
 // RateLimitBucket defines a token bucket configuration with a maximum capacity
-// and a refill period. Used by both shared (global) and per-user rate limits.
+// and a refill period. Used by both shared 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"`
+	MaxTokens int32 `json:"maxTokens" yaml:"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"`
+	RefillPeriod metav1.Duration `json:"refillPeriod" yaml:"refillPeriod"`
 }
 
 // ToolRateLimitConfig defines rate limits for a specific tool.
@@ -567,15 +567,15 @@ type ToolRateLimitConfig struct {
 	// Name is the MCP tool name this limit applies to.
 	// +kubebuilder:validation:Required
 	// +kubebuilder:validation:MinLength=1
-	Name string `json:"name"`
+	Name string `json:"name" yaml:"name"`
 
 	// Shared token bucket for this specific tool.
 	// +optional
-	Shared *RateLimitBucket `json:"shared,omitempty"`
+	Shared *RateLimitBucket `json:"shared,omitempty" yaml:"shared,omitempty"`
 
 	// PerUser token bucket configuration for this tool.
 	// +optional
-	PerUser *RateLimitBucket `json:"perUser,omitempty"`
+	PerUser *RateLimitBucket `json:"perUser,omitempty" yaml:"perUser,omitempty"`
 }
 
 // Permission profile types

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: &vmcpconfig.RateLimitConfig{
+				Shared: &vmcpconfig.RateLimitBucket{MaxTokens: 10, RefillPeriod: metav1.Duration{Duration: time.Minute}},
+				PerUser: &vmcpconfig.RateLimitBucket{
+					MaxTokens:    2,
+					RefillPeriod: metav1.Duration{Duration: time.Minute},
+				},
+				Tools: []vmcpconfig.ToolRateLimitConfig{
+					{
+						Name: "backend_a_echo",
+						Shared: &vmcpconfig.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/controllers/virtualmcpserver_vmcpconfig_test.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/pkg/vmcpconfig/converter_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: &vmcpconfig.RateLimitConfig{
+					PerUser: &vmcpconfig.RateLimitBucket{
+						MaxTokens:    2,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+					Tools: []vmcpconfig.ToolRateLimitConfig{
+						{
+							Name: "backend_a_echo",
+							Shared: &vmcpconfig.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/test-integration/embedding-server/embeddingserver_update_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())
 			})
 

cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_sessionstorage_cel_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,60 @@ 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.Config.RateLimiting = &vmcpconfig.RateLimitConfig{
+					Shared: &vmcpconfig.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).To(HaveOccurred())
+				Expect(err.Error()).To(ContainSubstring("config.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.Config.RateLimiting = &vmcpconfig.RateLimitConfig{
+					PerUser: &vmcpconfig.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).To(HaveOccurred())
+				Expect(err.Error()).To(ContainSubstring("config.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",
+						Audience: "test-audience",
+					},
+				}
+				vmcp.Spec.Config.RateLimiting = &vmcpconfig.RateLimitConfig{
+					PerUser: &vmcpconfig.RateLimitBucket{
+						MaxTokens:    1,
+						RefillPeriod: metav1.Duration{Duration: time.Minute},
+					},
+				}
+
+				err := k8sClient.Create(ctx, vmcp)
+				Expect(err).NotTo(HaveOccurred())
+			})
+		})
 	})