Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions cmd/thv-operator/api/v1beta1/mcpserver_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -507,43 +507,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.
Expand All @@ -556,15 +556,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
Expand Down
38 changes: 38 additions & 0 deletions cmd/thv-operator/api/v1beta1/mcpserver_types_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,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{
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.NotContains(t, out, `"config":{"rateLimiting"`)
}

func TestMCPServerSpecScalingFieldsJSONRoundtrip(t *testing.T) {
t.Parallel()

Expand Down
9 changes: 9 additions & 0 deletions cmd/thv-operator/api/v1beta1/virtualmcpserver_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,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.
Expand Down Expand Up @@ -143,6 +147,11 @@ type VirtualMCPServerSpec struct {
// +listType=atomic
// +optional
ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,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.
Expand Down
5 changes: 5 additions & 0 deletions cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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())
})

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.RateLimiting = &mcpv1beta1.RateLimitConfig{
Shared: &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",
Audience: "test-audience",
},
}
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())
})
})
})
Loading
Loading