Skip to content

Commit 65a78f4

Browse files
jerm-droclaude
andauthored
Add per-user rate limit types and limiter support (#4692)
* Add per-user rate limit CRD types and CEL validation Add PerUser field to RateLimitConfig and ToolRateLimitConfig so administrators can configure per-user token bucket rate limits on MCPServer. Make ToolRateLimitConfig.Shared optional since a tool entry may now have only a perUser limit. CEL admission validation enforces that perUser rate limiting requires authentication (oidcConfig, oidcConfigRef, or externalAuthConfigRef) at both server-level and per-tool level. The existing "at least one scope" rule is updated to include perUser alongside shared and tools. Add RateLimitConfigValid condition type and reason constants for use in the operator reconciler (wired in a following commit). Part of #4550 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Add RateLimitConfigValid status condition to reconciler Validate that per-user rate limiting has authentication enabled at reconciliation time (defense-in-depth alongside CEL admission). Set RateLimitConfigValid condition with appropriate reason: - RateLimitConfigValid when configuration is valid - PerUserRequiresAuth when perUser is set without auth - RateLimitNotApplicable when rate limiting is not configured Part of #4550 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Support per-user buckets in rate limiter Extend the limiter to create per-user token buckets keyed by userID. Per-user buckets are stored as deferred specs (bucketSpec) at construction time and materialized into TokenBucket structs at Allow() time since the userID is request-scoped. bucket.New() only allocates a struct (no I/O), so per-request creation is cheap. All applicable buckets (shared server, shared per-tool, per-user server, per-user per-tool) are checked atomically via ConsumeAll. The Lua script's two-phase check-then-consume ensures a per-user rejection does not drain the shared bucket. Redis keys follow the RFC format: - Server per-user: thv:rl:{ns:name}:user:{userID} - Tool per-user: thv:rl:{ns:name}:user:{userID}:tool:{toolName} Part of #4550 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix struct field alignment from linter Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Address review feedback on limiter - Use pointer for optional perUserSpec (clearer than bool+value) - Use distinct key prefix "user-tool:" for per-tool per-user buckets to prevent key collisions when userID contains delimiter characters - Extract shared validateBucketCRD helper to deduplicate validation between newBucket and newBucketSpec Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Regenerate swagger docs for perUser rate limit fields Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Regenerate CRD API reference docs for perUser fields Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Address review feedback from JAORMX Must fix: - Replace context.Background() with t.Context() in all limiter tests (new and pre-existing) - Fix RateLimitBucket swagger description by trimming field-level comments so the shared type gets a neutral description - Add comment in Allow() documenting RFC key format deviation and why "user-tool:" prefix prevents cross-type key collisions Should fix: - Change condition to ConditionTrue with NotApplicable when rate limiting is not configured (matches ImageValidated/Skipped pattern) - Add defense-in-depth comment explaining reconciliation continues intentionally (CEL is the primary gate) - Add Redis memory sizing note on PerUser CRD field Nits: - Fix test object names to use K8s-valid slugs (no spaces) - Keep nolint:lll on ToolRateLimitConfig (kubebuilder marker is 146 chars, exceeds the 130 limit) - Improve bucketSpec comment noting state lives in Redis Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 967257e commit 65a78f4

12 files changed

Lines changed: 683 additions & 59 deletions

File tree

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

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,18 @@ const (
172172
ConditionReasonSessionStorageNotApplicable = "SessionStorageWarningNotApplicable"
173173
)
174174

175+
// ConditionRateLimitConfigValid indicates whether the rate limit configuration is valid.
176+
const ConditionRateLimitConfigValid = "RateLimitConfigValid"
177+
178+
const (
179+
// ConditionReasonRateLimitConfigValid indicates the rate limit configuration is valid.
180+
ConditionReasonRateLimitConfigValid = "RateLimitConfigValid"
181+
// ConditionReasonRateLimitPerUserRequiresAuth indicates perUser rate limiting requires authentication.
182+
ConditionReasonRateLimitPerUserRequiresAuth = "PerUserRequiresAuth"
183+
// ConditionReasonRateLimitNotApplicable indicates rate limiting is not configured.
184+
ConditionReasonRateLimitNotApplicable = "RateLimitNotApplicable"
185+
)
186+
175187
// SessionStorageProviderRedis is the provider name for Redis-backed session storage.
176188
const SessionStorageProviderRedis = "redis"
177189

@@ -180,6 +192,8 @@ const SessionStorageProviderRedis = "redis"
180192
// +kubebuilder:validation:XValidation:rule="!(has(self.oidcConfig) && has(self.oidcConfigRef))",message="oidcConfig and oidcConfigRef are mutually exclusive; use oidcConfigRef to reference a shared MCPOIDCConfig"
181193
// +kubebuilder:validation:XValidation:rule="!(has(self.telemetry) && has(self.telemetryConfigRef))",message="telemetry and telemetryConfigRef are mutually exclusive; migrate to telemetryConfigRef"
182194
// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || (has(self.sessionStorage) && self.sessionStorage.provider == 'redis')",message="rateLimiting requires sessionStorage with provider 'redis'"
195+
// +kubebuilder:validation:XValidation:rule="!(has(self.rateLimiting) && has(self.rateLimiting.perUser)) || has(self.oidcConfig) || has(self.oidcConfigRef) || has(self.externalAuthConfigRef)",message="rateLimiting.perUser requires authentication (oidcConfig, oidcConfigRef, or externalAuthConfigRef)"
196+
// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || !has(self.rateLimiting.tools) || self.rateLimiting.tools.all(t, !has(t.perUser)) || has(self.oidcConfig) || has(self.oidcConfigRef) || has(self.externalAuthConfigRef)",message="per-tool perUser rate limiting requires authentication (oidcConfig, oidcConfigRef, or externalAuthConfigRef)"
183197
//
184198
//nolint:lll // CEL validation rules exceed line length limit
185199
type MCPServerSpec struct {
@@ -519,16 +533,23 @@ type SessionStorageConfig struct {
519533
}
520534

521535
// RateLimitConfig defines rate limiting configuration for an MCP server.
522-
// At least one of shared or tools must be configured.
536+
// At least one of shared, perUser, or tools must be configured.
523537
//
524-
// +kubebuilder:validation:XValidation:rule="has(self.shared) || (has(self.tools) && size(self.tools) > 0)",message="at least one of shared or tools must be configured"
538+
// +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"
525539
//
526540
//nolint:lll // CEL validation rules exceed line length limit
527541
type RateLimitConfig struct {
528-
// Shared defines a token bucket shared across all users for the entire server.
542+
// Shared is a token bucket shared across all users for the entire server.
529543
// +optional
530544
Shared *RateLimitBucket `json:"shared,omitempty"`
531545

546+
// PerUser is a token bucket applied independently to each authenticated user
547+
// at the server level. Requires authentication to be enabled.
548+
// Each unique userID creates Redis keys that expire after 2x refillPeriod.
549+
// Memory formula: unique_users_per_TTL_window * (1 + num_tools_with_per_user_limits) keys.
550+
// +optional
551+
PerUser *RateLimitBucket `json:"perUser,omitempty"`
552+
532553
// Tools defines per-tool rate limit overrides.
533554
// Each entry applies additional rate limits to calls targeting a specific tool name.
534555
// A request must pass both the server-level limit and the per-tool limit.
@@ -538,7 +559,8 @@ type RateLimitConfig struct {
538559
Tools []ToolRateLimitConfig `json:"tools,omitempty"`
539560
}
540561

541-
// RateLimitBucket defines a token bucket configuration.
562+
// RateLimitBucket defines a token bucket configuration with a maximum capacity
563+
// and a refill period. Used by both shared (global) and per-user rate limits.
542564
type RateLimitBucket struct {
543565
// MaxTokens is the maximum number of tokens (bucket capacity).
544566
// This is also the burst size: the maximum number of requests that can be served
@@ -555,15 +577,24 @@ type RateLimitBucket struct {
555577
}
556578

557579
// ToolRateLimitConfig defines rate limits for a specific tool.
580+
// At least one of shared or perUser must be configured.
581+
//
582+
// +kubebuilder:validation:XValidation:rule="has(self.shared) || has(self.perUser)",message="at least one of shared or perUser must be configured"
583+
//
584+
//nolint:lll // kubebuilder marker exceeds line length
558585
type ToolRateLimitConfig struct {
559586
// Name is the MCP tool name this limit applies to.
560587
// +kubebuilder:validation:Required
561588
// +kubebuilder:validation:MinLength=1
562589
Name string `json:"name"`
563590

564-
// Shared defines a token bucket shared across all users for this specific tool.
565-
// +kubebuilder:validation:Required
566-
Shared *RateLimitBucket `json:"shared"`
591+
// Shared token bucket for this specific tool.
592+
// +optional
593+
Shared *RateLimitBucket `json:"shared,omitempty"`
594+
595+
// PerUser token bucket configuration for this tool.
596+
// +optional
597+
PerUser *RateLimitBucket `json:"perUser,omitempty"`
567598
}
568599

569600
// Permission profile types

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

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

cmd/thv-operator/controllers/mcpserver_controller.go

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,9 +222,10 @@ func (r *MCPServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (
222222
// Validate CABundleRef if specified
223223
r.validateCABundleRef(ctx, mcpServer)
224224

225-
// Validate stdio replica cap and session storage requirements
225+
// Validate stdio replica cap, session storage, and rate limit config
226226
r.validateStdioReplicaCap(ctx, mcpServer)
227227
r.validateSessionStorageForReplicas(ctx, mcpServer)
228+
r.validateRateLimitConfig(ctx, mcpServer)
228229

229230
// Validate PodTemplateSpec early - before other validations
230231
// This ensures we fail fast if the spec is invalid
@@ -2407,6 +2408,61 @@ func (r *MCPServerReconciler) validateSessionStorageForReplicas(ctx context.Cont
24072408
}
24082409
}
24092410

2411+
// setRateLimitConfigCondition sets the RateLimitConfigValid status condition.
2412+
func setRateLimitConfigCondition(mcpServer *mcpv1alpha1.MCPServer, status metav1.ConditionStatus, reason, message string) {
2413+
meta.SetStatusCondition(&mcpServer.Status.Conditions, metav1.Condition{
2414+
Type: mcpv1alpha1.ConditionRateLimitConfigValid,
2415+
Status: status,
2416+
Reason: reason,
2417+
Message: message,
2418+
ObservedGeneration: mcpServer.Generation,
2419+
})
2420+
}
2421+
2422+
// validateRateLimitConfig validates that per-user rate limiting has authentication enabled.
2423+
// Sets the RateLimitConfigValid condition. This is defense-in-depth only; CEL admission
2424+
// validation is the primary gate. Reconciliation continues even when the condition is False
2425+
// because per-user buckets are silently skipped when userID is empty (graceful degradation).
2426+
func (r *MCPServerReconciler) validateRateLimitConfig(ctx context.Context, mcpServer *mcpv1alpha1.MCPServer) {
2427+
rl := mcpServer.Spec.RateLimiting
2428+
if rl == nil {
2429+
setRateLimitConfigCondition(mcpServer, metav1.ConditionTrue,
2430+
mcpv1alpha1.ConditionReasonRateLimitNotApplicable,
2431+
"rate limiting is not configured")
2432+
if err := r.Status().Update(ctx, mcpServer); err != nil {
2433+
log.FromContext(ctx).Error(err, "Failed to update MCPServer status after rate limit validation")
2434+
}
2435+
return
2436+
}
2437+
2438+
authEnabled := mcpServer.Spec.OIDCConfig != nil ||
2439+
mcpServer.Spec.OIDCConfigRef != nil ||
2440+
mcpServer.Spec.ExternalAuthConfigRef != nil
2441+
2442+
hasPerUser := rl.PerUser != nil
2443+
if !hasPerUser {
2444+
for _, t := range rl.Tools {
2445+
if t.PerUser != nil {
2446+
hasPerUser = true
2447+
break
2448+
}
2449+
}
2450+
}
2451+
2452+
if hasPerUser && !authEnabled {
2453+
setRateLimitConfigCondition(mcpServer, metav1.ConditionFalse,
2454+
mcpv1alpha1.ConditionReasonRateLimitPerUserRequiresAuth,
2455+
"perUser rate limiting requires authentication to be enabled (oidcConfig, oidcConfigRef, or externalAuthConfigRef)")
2456+
} else {
2457+
setRateLimitConfigCondition(mcpServer, metav1.ConditionTrue,
2458+
mcpv1alpha1.ConditionReasonRateLimitConfigValid,
2459+
"rate limit configuration is valid")
2460+
}
2461+
if err := r.Status().Update(ctx, mcpServer); err != nil {
2462+
log.FromContext(ctx).Error(err, "Failed to update MCPServer status after rate limit validation")
2463+
}
2464+
}
2465+
24102466
// SetupWithManager sets up the controller with the Manager.
24112467
func (r *MCPServerReconciler) SetupWithManager(mgr ctrl.Manager) error {
24122468
// Create a handler that maps MCPExternalAuthConfig changes to MCPServer reconciliation requests

cmd/thv-operator/controllers/mcpserver_replicas_test.go

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -979,3 +979,156 @@ func TestUpdateMCPServerStatusExcludesTerminatingPods(t *testing.T) {
979979
assert.Equal(t, int32(2), updatedMCPServer.Status.ReadyReplicas,
980980
"ReadyReplicas should exclude terminating pods")
981981
}
982+
983+
func TestRateLimitConfigValidation(t *testing.T) {
984+
t.Parallel()
985+
986+
tests := []struct {
987+
name string
988+
spec mcpv1alpha1.MCPServerSpec
989+
expectStatus metav1.ConditionStatus
990+
expectReason string
991+
}{
992+
{
993+
name: "no-rate-limiting",
994+
spec: mcpv1alpha1.MCPServerSpec{
995+
Image: "test-image:latest",
996+
Transport: "sse",
997+
ProxyPort: 8080,
998+
},
999+
expectStatus: metav1.ConditionTrue,
1000+
expectReason: mcpv1alpha1.ConditionReasonRateLimitNotApplicable,
1001+
},
1002+
{
1003+
name: "peruser-with-auth",
1004+
spec: mcpv1alpha1.MCPServerSpec{
1005+
Image: "test-image:latest",
1006+
Transport: "sse",
1007+
ProxyPort: 8080,
1008+
SessionStorage: &mcpv1alpha1.SessionStorageConfig{
1009+
Provider: mcpv1alpha1.SessionStorageProviderRedis,
1010+
Address: "redis:6379",
1011+
},
1012+
OIDCConfig: &mcpv1alpha1.OIDCConfigRef{Type: "kubernetes"},
1013+
RateLimiting: &mcpv1alpha1.RateLimitConfig{
1014+
PerUser: &mcpv1alpha1.RateLimitBucket{
1015+
MaxTokens: 100,
1016+
RefillPeriod: metav1.Duration{Duration: time.Minute},
1017+
},
1018+
},
1019+
},
1020+
expectStatus: metav1.ConditionTrue,
1021+
expectReason: mcpv1alpha1.ConditionReasonRateLimitConfigValid,
1022+
},
1023+
{
1024+
name: "peruser-without-auth",
1025+
spec: mcpv1alpha1.MCPServerSpec{
1026+
Image: "test-image:latest",
1027+
Transport: "sse",
1028+
ProxyPort: 8080,
1029+
SessionStorage: &mcpv1alpha1.SessionStorageConfig{
1030+
Provider: mcpv1alpha1.SessionStorageProviderRedis,
1031+
Address: "redis:6379",
1032+
},
1033+
RateLimiting: &mcpv1alpha1.RateLimitConfig{
1034+
PerUser: &mcpv1alpha1.RateLimitBucket{
1035+
MaxTokens: 100,
1036+
RefillPeriod: metav1.Duration{Duration: time.Minute},
1037+
},
1038+
},
1039+
},
1040+
expectStatus: metav1.ConditionFalse,
1041+
expectReason: mcpv1alpha1.ConditionReasonRateLimitPerUserRequiresAuth,
1042+
},
1043+
{
1044+
name: "per-tool-peruser-without-auth",
1045+
spec: mcpv1alpha1.MCPServerSpec{
1046+
Image: "test-image:latest",
1047+
Transport: "sse",
1048+
ProxyPort: 8080,
1049+
SessionStorage: &mcpv1alpha1.SessionStorageConfig{
1050+
Provider: mcpv1alpha1.SessionStorageProviderRedis,
1051+
Address: "redis:6379",
1052+
},
1053+
RateLimiting: &mcpv1alpha1.RateLimitConfig{
1054+
Tools: []mcpv1alpha1.ToolRateLimitConfig{
1055+
{
1056+
Name: "search",
1057+
PerUser: &mcpv1alpha1.RateLimitBucket{
1058+
MaxTokens: 10,
1059+
RefillPeriod: metav1.Duration{Duration: time.Minute},
1060+
},
1061+
},
1062+
},
1063+
},
1064+
},
1065+
expectStatus: metav1.ConditionFalse,
1066+
expectReason: mcpv1alpha1.ConditionReasonRateLimitPerUserRequiresAuth,
1067+
},
1068+
{
1069+
name: "shared-only-no-auth",
1070+
spec: mcpv1alpha1.MCPServerSpec{
1071+
Image: "test-image:latest",
1072+
Transport: "sse",
1073+
ProxyPort: 8080,
1074+
SessionStorage: &mcpv1alpha1.SessionStorageConfig{
1075+
Provider: mcpv1alpha1.SessionStorageProviderRedis,
1076+
Address: "redis:6379",
1077+
},
1078+
RateLimiting: &mcpv1alpha1.RateLimitConfig{
1079+
Shared: &mcpv1alpha1.RateLimitBucket{
1080+
MaxTokens: 1000,
1081+
RefillPeriod: metav1.Duration{Duration: time.Minute},
1082+
},
1083+
},
1084+
},
1085+
expectStatus: metav1.ConditionTrue,
1086+
expectReason: mcpv1alpha1.ConditionReasonRateLimitConfigValid,
1087+
},
1088+
}
1089+
1090+
for _, tt := range tests {
1091+
t.Run(tt.name, func(t *testing.T) {
1092+
t.Parallel()
1093+
1094+
name := "rl-" + tt.name
1095+
namespace := testNamespaceDefault
1096+
1097+
mcpServer := &mcpv1alpha1.MCPServer{
1098+
ObjectMeta: metav1.ObjectMeta{
1099+
Name: name,
1100+
Namespace: namespace,
1101+
},
1102+
Spec: tt.spec,
1103+
}
1104+
1105+
testScheme := createTestScheme()
1106+
fakeClient := fake.NewClientBuilder().
1107+
WithScheme(testScheme).
1108+
WithObjects(mcpServer).
1109+
WithStatusSubresource(&mcpv1alpha1.MCPServer{}).
1110+
Build()
1111+
1112+
reconciler := newTestMCPServerReconciler(fakeClient, testScheme, kubernetes.PlatformKubernetes)
1113+
1114+
_, err := reconciler.Reconcile(t.Context(), ctrl.Request{
1115+
NamespacedName: types.NamespacedName{Name: name, Namespace: namespace},
1116+
})
1117+
require.NoError(t, err)
1118+
1119+
updated := &mcpv1alpha1.MCPServer{}
1120+
err = fakeClient.Get(t.Context(), types.NamespacedName{Name: name, Namespace: namespace}, updated)
1121+
require.NoError(t, err)
1122+
1123+
var found bool
1124+
for _, cond := range updated.Status.Conditions {
1125+
if cond.Type == mcpv1alpha1.ConditionRateLimitConfigValid {
1126+
found = true
1127+
assert.Equal(t, tt.expectStatus, cond.Status)
1128+
assert.Equal(t, tt.expectReason, cond.Reason)
1129+
}
1130+
}
1131+
assert.True(t, found, "ConditionRateLimitConfigValid condition should be set")
1132+
})
1133+
}
1134+
}

0 commit comments

Comments
 (0)