From 66ce2c50b0449eac183722fb6fe880308aa1bee6 Mon Sep 17 00:00:00 2001 From: "Ankit Kr. Chowdhury" Date: Fri, 7 Aug 2026 09:04:29 +0000 Subject: [PATCH 1/4] feat: implement AgentDeployment admission webhooks and resource quota enforcement --- .agents/AGENTS.md | 2 +- .agents/skills/agentrax-context/SKILL.md | 1 + api/v1alpha1/error_rate.go | 39 ++ api/v1alpha1/error_rate_test.go | 62 +++ cmd/main.go | 17 +- config/rbac/role.yaml | 10 +- config/webhook/kustomization.yaml | 3 + config/webhook/manifests.yaml | 52 +++ config/webhook/service.yaml | 13 + go.mod | 5 +- .../agentdeployment_controller_test.go | 20 + internal/controller/enqueue_handlers.go | 49 ++ internal/controller/suite_test.go | 30 ++ internal/controller/tenantquota_controller.go | 123 ++++- .../controller/tenantquota_controller_test.go | 362 +++++++++++++-- internal/controller/test_helpers_test.go | 50 +++ internal/quota/enforcer.go | 302 +++++++++++++ internal/quota/enforcer_test.go | 380 ++++++++++++++++ .../webhook/agentdeployment_validator_test.go | 421 ++++++++++++++++++ internal/webhook/agentdeployment_webhook.go | 326 ++++++++++++++ .../webhook/agentdeployment_webhook_test.go | 136 ++++++ 21 files changed, 2333 insertions(+), 70 deletions(-) create mode 100644 api/v1alpha1/error_rate.go create mode 100644 api/v1alpha1/error_rate_test.go create mode 100644 config/webhook/kustomization.yaml create mode 100644 config/webhook/manifests.yaml create mode 100644 config/webhook/service.yaml create mode 100644 internal/controller/enqueue_handlers.go create mode 100644 internal/controller/test_helpers_test.go create mode 100644 internal/quota/enforcer.go create mode 100644 internal/quota/enforcer_test.go create mode 100644 internal/webhook/agentdeployment_validator_test.go create mode 100644 internal/webhook/agentdeployment_webhook.go create mode 100644 internal/webhook/agentdeployment_webhook_test.go diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md index d385f60..0d1be80 100644 --- a/.agents/AGENTS.md +++ b/.agents/AGENTS.md @@ -44,7 +44,7 @@ ### Serena (MCP — Semantic Code Intelligence) This project supports **Serena** configured as an MCP server. -Serena exposes `gopls`-backed semantic tools via the Model Context Protocol. Note that Go and `gopls` must be installed and available on `PATH` before use. +Serena exposes `gopls`-backed semantic tools via the Model Context Protocol. The project is indexed and the config lives at `.serena/project.yml`. **Configuration & Setup:** diff --git a/.agents/skills/agentrax-context/SKILL.md b/.agents/skills/agentrax-context/SKILL.md index d1aef4a..4105f90 100644 --- a/.agents/skills/agentrax-context/SKILL.md +++ b/.agents/skills/agentrax-context/SKILL.md @@ -31,6 +31,7 @@ description: Project context and settled architecture decisions for the Agentrax | `internal/scaling/` | HPA generation and quota-capped scaling logic. | | `internal/registry/` | MCP registrar, registry HTTP handler, TTL sweep. | | `internal/quota/` | Quota arithmetic and in-flight reservation. Shared by webhook and TenantQuota reconciler. | +| `internal/webhook/` | Validating and mutating admission webhooks. Lives here (not `api/`) to import `internal/quota` without creating an import cycle. | | `internal/metrics/` | Shared Prometheus client plumbing used by rollout and scaling. | ## Where the hard logic lives diff --git a/api/v1alpha1/error_rate.go b/api/v1alpha1/error_rate.go new file mode 100644 index 0000000..21c6514 --- /dev/null +++ b/api/v1alpha1/error_rate.go @@ -0,0 +1,39 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package v1alpha1 + +import "fmt" + +// ParseErrorRate parses a percentage string like "2%" and returns the float64 +// value (e.g. 0.02 for "2%"). Returns an error if the format is invalid. +// Used by the validating webhook and the rollout threshold evaluator. +func ParseErrorRate(s string) (float64, error) { + if len(s) == 0 { + return 0, fmt.Errorf("empty error rate string") + } + if s[len(s)-1] != '%' { + return 0, fmt.Errorf("error rate must end with '%%': got %q", s) + } + var pct float64 + if _, err := fmt.Sscanf(s[:len(s)-1], "%f", &pct); err != nil { + return 0, fmt.Errorf("parsing error rate %q: %w", s, err) + } + if pct < 0 || pct > 100 { + return 0, fmt.Errorf("error rate %q out of range [0, 100]", s) + } + return pct / 100.0, nil +} diff --git a/api/v1alpha1/error_rate_test.go b/api/v1alpha1/error_rate_test.go new file mode 100644 index 0000000..96dde68 --- /dev/null +++ b/api/v1alpha1/error_rate_test.go @@ -0,0 +1,62 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package v1alpha1_test + +import ( + "testing" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +func TestParseErrorRate(t *testing.T) { + t.Parallel() + tests := []struct { + input string + want float64 + wantErr bool + }{ + {"2%", 0.02, false}, + {"100%", 1.0, false}, + {"0%", 0.0, false}, + {"0.5%", 0.005, false}, + {"", 0, true}, + {"5", 0, true}, + {"-1%", 0, true}, + {"101%", 0, true}, + {"abc%", 0, true}, + } + for _, tc := range tests { + tc := tc + t.Run(tc.input, func(t *testing.T) { + t.Parallel() + got, err := agentraxv1alpha1.ParseErrorRate(tc.input) + if (err != nil) != tc.wantErr { + t.Errorf("ParseErrorRate(%q) error=%v, wantErr=%v", tc.input, err, tc.wantErr) + } + if err == nil && abs(got-tc.want) > 1e-9 { + t.Errorf("ParseErrorRate(%q) = %v, want %v", tc.input, got, tc.want) + } + }) + } +} + +func abs(f float64) float64 { + if f < 0 { + return -f + } + return f +} diff --git a/cmd/main.go b/cmd/main.go index af8196b..fbadaa4 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -38,6 +38,8 @@ import ( agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" "github.com/gitcommitankit/agentrax/internal/controller" + "github.com/gitcommitankit/agentrax/internal/quota" + agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" // +kubebuilder:scaffold:imports ) @@ -59,6 +61,7 @@ func main() { var probeAddr string var secureMetrics bool var enableHTTP2 bool + var gpuResourceName string var tlsOpts []func(*tls.Config) flag.StringVar(&metricsAddr, "metrics-bind-address", "0", "The address the metrics endpoint binds to. "+ "Use :8443 for HTTPS or :8080 for HTTP, or leave as 0 to disable the metrics service.") @@ -70,6 +73,8 @@ func main() { "If set, the metrics endpoint is served securely via HTTPS. Use --metrics-secure=false to use HTTP instead.") flag.BoolVar(&enableHTTP2, "enable-http2", false, "If set, HTTP/2 will be enabled for the metrics and webhook servers") + flag.StringVar(&gpuResourceName, "gpu-resource-name", quota.DefaultGPUResourceName, + "Kubernetes resource name used to count GPU units in AgentDeployment resource limits.") opts := zap.Options{ Development: true, } @@ -145,6 +150,9 @@ func main() { os.Exit(1) } + // Shared quota enforcer used by both the webhook validator and TenantQuota reconciler. + quotaEnforcer := quota.NewEnforcer(gpuResourceName) + if err = (&controller.AgentDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), @@ -153,12 +161,17 @@ func main() { os.Exit(1) } if err = (&controller.TenantQuotaReconciler{ - Client: mgr.GetClient(), - Scheme: mgr.GetScheme(), + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + Enforcer: quotaEnforcer, }).SetupWithManager(mgr); err != nil { setupLog.Error(err, "unable to create controller", "controller", "TenantQuota") os.Exit(1) } + if err = agentraxwebhook.SetupAgentDeploymentWebhookWithManager(mgr, quotaEnforcer); err != nil { + setupLog.Error(err, "unable to register webhook", "webhook", "AgentDeployment") + os.Exit(1) + } // +kubebuilder:scaffold:builder if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil { diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index dab9cf2..bc3d84c 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -8,7 +8,6 @@ rules: - agentrax.io resources: - agentdeployments - - tenantquotas verbs: - create - delete @@ -21,7 +20,6 @@ rules: - agentrax.io resources: - agentdeployments/finalizers - - tenantquotas/finalizers verbs: - update - apiGroups: @@ -33,6 +31,14 @@ rules: - get - patch - update +- apiGroups: + - agentrax.io + resources: + - tenantquotas + verbs: + - get + - list + - watch - apiGroups: - apiextensions.k8s.io resources: diff --git a/config/webhook/kustomization.yaml b/config/webhook/kustomization.yaml new file mode 100644 index 0000000..93e025a --- /dev/null +++ b/config/webhook/kustomization.yaml @@ -0,0 +1,3 @@ +resources: + - manifests.yaml + - service.yaml diff --git a/config/webhook/manifests.yaml b/config/webhook/manifests.yaml new file mode 100644 index 0000000..fa32a1a --- /dev/null +++ b/config/webhook/manifests.yaml @@ -0,0 +1,52 @@ +--- +apiVersion: admissionregistration.k8s.io/v1 +kind: MutatingWebhookConfiguration +metadata: + name: mutating-webhook-configuration +webhooks: + - admissionReviewVersions: + - v1 + clientConfig: + service: + name: webhook-service + namespace: system + path: /mutate-agentrax-io-v1alpha1-agentdeployment + failurePolicy: Fail + name: magentdeployment.kb.io + rules: + - apiGroups: + - agentrax.io + apiVersions: + - v1alpha1 + operations: + - CREATE + - UPDATE + resources: + - agentdeployments + sideEffects: None +--- +apiVersion: admissionregistration.k8s.io/v1 +kind: ValidatingWebhookConfiguration +metadata: + name: validating-webhook-configuration +webhooks: + - admissionReviewVersions: + - v1 + clientConfig: + service: + name: webhook-service + namespace: system + path: /validate-agentrax-io-v1alpha1-agentdeployment + failurePolicy: Fail + name: vagentdeployment.kb.io + rules: + - apiGroups: + - agentrax.io + apiVersions: + - v1alpha1 + operations: + - CREATE + - UPDATE + resources: + - agentdeployments + sideEffects: None diff --git a/config/webhook/service.yaml b/config/webhook/service.yaml new file mode 100644 index 0000000..efaf524 --- /dev/null +++ b/config/webhook/service.yaml @@ -0,0 +1,13 @@ +--- +apiVersion: v1 +kind: Service +metadata: + name: webhook-service + namespace: system +spec: + ports: + - port: 443 + protocol: TCP + targetPort: 9443 + selector: + control-plane: controller-manager diff --git a/go.mod b/go.mod index 19b71e9..3ee41c1 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/gitcommitankit/agentrax go 1.22.0 require ( + github.com/go-logr/logr v1.4.2 github.com/onsi/ginkgo/v2 v2.19.0 github.com/onsi/gomega v1.33.1 github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring v0.75.0 @@ -26,7 +27,6 @@ require ( github.com/felixge/httpsnoop v1.0.4 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/fxamacker/cbor/v2 v2.7.0 // indirect - github.com/go-logr/logr v1.4.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/go-logr/zapr v1.3.0 // indirect github.com/go-openapi/jsonpointer v0.19.6 // indirect @@ -84,6 +84,7 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect google.golang.org/grpc v1.65.0 // indirect google.golang.org/protobuf v1.34.2 // indirect + gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect @@ -96,4 +97,4 @@ require ( sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect sigs.k8s.io/yaml v1.4.0 // indirect -) \ No newline at end of file +) diff --git a/internal/controller/agentdeployment_controller_test.go b/internal/controller/agentdeployment_controller_test.go index 6f65ae5..c63f493 100644 --- a/internal/controller/agentdeployment_controller_test.go +++ b/internal/controller/agentdeployment_controller_test.go @@ -44,9 +44,28 @@ const ( testNginxImage = "nginx:latest" ) +// ensureTenantQuota creates the TenantQuota "team-test" in the given namespace +// if it does not already exist. Phase 2 webhooks require spec.tenantRef to +// resolve to a real TenantQuota before admitting an AgentDeployment. +func ensureTenantQuota(namespace string) { + tq := &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{Name: "team-test", Namespace: namespace}, + Spec: agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 100, + MaxGPUs: 0, + MaxTotalReplicas: 300, + MaxReplicasPerAgent: 10, + }, + } + err := k8sClient.Create(ctx, tq) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue(), + "ensuring TenantQuota team-test in %s: %v", namespace, err) +} + // createAgentDeployment is a test helper that creates a minimal AgentDeployment // and returns its NamespacedName. func createAgentDeployment(name, namespace, image string, port, minReplicas int32) types.NamespacedName { + ensureTenantQuota(namespace) ad := &agentraxv1alpha1.AgentDeployment{ ObjectMeta: metav1.ObjectMeta{ Name: name, @@ -647,6 +666,7 @@ var _ = Describe("AgentDeployment Controller", func() { ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-env-args"}} err := k8sClient.Create(ctx, ns) Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue(), "creating namespace test-env-args: %v", err) + ensureTenantQuota("test-env-args") ad := &agentraxv1alpha1.AgentDeployment{ ObjectMeta: metav1.ObjectMeta{ diff --git a/internal/controller/enqueue_handlers.go b/internal/controller/enqueue_handlers.go new file mode 100644 index 0000000..87e64f3 --- /dev/null +++ b/internal/controller/enqueue_handlers.go @@ -0,0 +1,49 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package controller + +import ( + "context" + + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +// enqueueTenantQuota returns an EventHandler that maps every AgentDeployment +// event to a reconcile request for the TenantQuota named by spec.tenantRef in +// the same namespace. This causes the TenantQuota reconciler to recompute usage +// whenever any AD in the namespace is created, updated, or deleted. +func enqueueTenantQuota() handler.EventHandler { + return handler.EnqueueRequestsFromMapFunc(func(ctx context.Context, obj client.Object) []reconcile.Request { + ad, ok := obj.(*agentraxv1alpha1.AgentDeployment) + if !ok || ad.Spec.TenantRef == "" { + return nil + } + return []reconcile.Request{ + { + NamespacedName: types.NamespacedName{ + Namespace: ad.Namespace, + Name: ad.Spec.TenantRef, + }, + }, + } + }) +} diff --git a/internal/controller/suite_test.go b/internal/controller/suite_test.go index 5049429..cc0c4f8 100644 --- a/internal/controller/suite_test.go +++ b/internal/controller/suite_test.go @@ -37,10 +37,13 @@ import ( logf "sigs.k8s.io/controller-runtime/pkg/log" "sigs.k8s.io/controller-runtime/pkg/log/zap" metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server" + "sigs.k8s.io/controller-runtime/pkg/webhook" monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/quota" + agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" // +kubebuilder:scaffold:imports ) @@ -59,6 +62,10 @@ var mgrDone chan struct{} // and clear it in AfterEach to avoid contaminating other tests. var testReconciler *AgentDeploymentReconciler +// testEnforcer is the shared quota Enforcer used by the TenantQuota reconciler +// and the validating webhook in integration tests. +var testEnforcer *quota.Enforcer + func TestControllers(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Controller Suite") @@ -79,6 +86,12 @@ var _ = BeforeSuite(func() { }, ErrorIfCRDPathMissing: true, + // Configure envtest to install and run the webhooks during integration + // tests. envtest generates self-signed TLS certs automatically. + WebhookInstallOptions: envtest.WebhookInstallOptions{ + Paths: []string{filepath.Join("..", "..", "config", "webhook")}, + }, + // The BinaryAssetsDirectory is only required if you want to run the tests directly // without calling the makefile target test. If not informed it will look for the // default path defined in controller-runtime which is /usr/local/kubebuilder/. @@ -110,19 +123,36 @@ var _ = BeforeSuite(func() { Expect(k8sClient).NotTo(BeNil()) // Start the controller manager so the reconciler runs during integration tests. + // Use envtest's webhook host/port so the manager's webhook server binds to the + // same address the webhook install options configured the API server to call. mgr, err := ctrl.NewManager(cfg, ctrl.Options{ Scheme: scheme.Scheme, // Disable the metrics server in tests to avoid port conflicts. Metrics: metricsserver.Options{BindAddress: "0"}, + WebhookServer: webhook.NewServer(webhook.Options{ + Host: testEnv.WebhookInstallOptions.LocalServingHost, + Port: testEnv.WebhookInstallOptions.LocalServingPort, + CertDir: testEnv.WebhookInstallOptions.LocalServingCertDir, + }), }) Expect(err).NotTo(HaveOccurred()) + testEnforcer = quota.NewEnforcer(quota.DefaultGPUResourceName) + testReconciler = &AgentDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), } Expect(testReconciler.SetupWithManager(mgr)).To(Succeed()) + Expect((&TenantQuotaReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + Enforcer: testEnforcer, + }).SetupWithManager(mgr)).To(Succeed()) + + Expect(agentraxwebhook.SetupAgentDeploymentWebhookWithManager(mgr, testEnforcer)).To(Succeed()) + mgrDone = make(chan struct{}) go func() { defer GinkgoRecover() diff --git a/internal/controller/tenantquota_controller.go b/internal/controller/tenantquota_controller.go index 81d7760..4b0ce7e 100644 --- a/internal/controller/tenantquota_controller.go +++ b/internal/controller/tenantquota_controller.go @@ -19,45 +19,134 @@ package controller import ( "context" + "fmt" + "time" + "k8s.io/apimachinery/pkg/api/equality" + apierrors "k8s.io/apimachinery/pkg/api/errors" + apimeta "k8s.io/apimachinery/pkg/api/meta" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/log" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/quota" ) // TenantQuotaReconciler reconciles a TenantQuota object. type TenantQuotaReconciler struct { client.Client - Scheme *runtime.Scheme + Scheme *runtime.Scheme + Enforcer *quota.Enforcer } -// +kubebuilder:rbac:groups=agentrax.io,resources=tenantquotas,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups=agentrax.io,resources=tenantquotas,verbs=get;list;watch // +kubebuilder:rbac:groups=agentrax.io,resources=tenantquotas/status,verbs=get;update;patch -// +kubebuilder:rbac:groups=agentrax.io,resources=tenantquotas/finalizers,verbs=update - -// Reconcile is part of the main kubernetes reconciliation loop which aims to -// move the current state of the cluster closer to the desired state. -// TODO(user): Modify the Reconcile function to compare the state specified by -// the TenantQuota object against the actual cluster state, and then -// perform operations to make the cluster state reflect the state specified by -// the user. -// -// For more details, check Reconcile and its Result here: -// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.19.0/pkg/reconcile +// +kubebuilder:rbac:groups=agentrax.io,resources=agentdeployments,verbs=get;list;watch + +// Reconcile computes actual AgentDeployment usage within the TenantQuota's +// namespace and writes accurate usage counters to TenantQuota.status. +// If usage exceeds any quota ceiling (e.g. because maxAgents was lowered), +// it sets the OverQuota condition without forcibly deleting any resources. func (r *TenantQuotaReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { - _ = log.FromContext(ctx) + logger := log.FromContext(ctx) + + // 1. Fetch the TenantQuota; return immediately if it has been deleted. + tq := &agentraxv1alpha1.TenantQuota{} + if err := r.Get(ctx, req.NamespacedName, tq); err != nil { + if apierrors.IsNotFound(err) { + return ctrl.Result{}, nil + } + return ctrl.Result{}, fmt.Errorf("fetching TenantQuota: %w", err) + } + + // 2. List all AgentDeployment objects in the same namespace that reference + // this TenantQuota via spec.tenantRef. + adList := &agentraxv1alpha1.AgentDeploymentList{} + if err := r.List(ctx, adList, client.InNamespace(tq.Namespace)); err != nil { + return ctrl.Result{}, fmt.Errorf("listing AgentDeployments: %w", err) + } + + // Filter to only ADs referencing this TenantQuota. + var specs []agentraxv1alpha1.AgentDeploymentSpec + for i := range adList.Items { + if adList.Items[i].Spec.TenantRef == tq.Name { + specs = append(specs, adList.Items[i].Spec) + // Release the in-flight reservation for this AD now that it is + // committed to etcd. This prevents the 5s TTL window from + // blocking rapid sequential creates of the same or sibling ADs. + r.Enforcer.Release(fmt.Sprintf("%s/%s", adList.Items[i].Namespace, adList.Items[i].Name)) + } + } + + // 3. Compute accurate usage from the live AD list. + usage := r.Enforcer.ComputeUsage(specs) + + // 4. Re-fetch with the latest resourceVersion before writing status to avoid + // optimistic concurrency conflicts. + latest := &agentraxv1alpha1.TenantQuota{} + if err := r.Get(ctx, req.NamespacedName, latest); err != nil { + if apierrors.IsNotFound(err) { + return ctrl.Result{}, nil + } + return ctrl.Result{}, fmt.Errorf("re-fetching TenantQuota for status update: %w", err) + } + + prevStatus := latest.Status.DeepCopy() + + latest.Status.UsedAgents = usage.UsedAgents + latest.Status.UsedGPUs = usage.UsedGPUs + latest.Status.UsedTotalReplicas = usage.UsedTotalReplicas + + // 5. Set or clear the OverQuota condition based on whether usage exceeds spec. + // Use latest.Spec (re-fetched) rather than tq.Spec (first fetch) to avoid + // evaluating against a ceiling that may have changed between the two Gets. + over, overMsg := r.Enforcer.IsOverQuota(latest.Spec, usage) + if over { + apimeta.SetStatusCondition(&latest.Status.Conditions, metav1.Condition{ + Type: agentraxv1alpha1.ConditionOverQuota, + Status: metav1.ConditionTrue, + Reason: "UsageExceedsQuota", + Message: overMsg, + ObservedGeneration: latest.Generation, + }) + } else { + apimeta.RemoveStatusCondition(&latest.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) + } - // TODO(user): your logic here + // 6. Only write status when something actually changed. + if !equality.Semantic.DeepEqual(prevStatus, &latest.Status) { + if err := r.Status().Update(ctx, latest); err != nil { + return ctrl.Result{}, fmt.Errorf("updating TenantQuota status: %w", err) + } + logger.Info("updated TenantQuota status", + "name", latest.Name, + "namespace", latest.Namespace, + "usedAgents", latest.Status.UsedAgents, + "usedGPUs", latest.Status.UsedGPUs, + "usedTotalReplicas", latest.Status.UsedTotalReplicas, + "overQuota", over, + ) + } - return ctrl.Result{}, nil + // Requeue periodically as a safety net: if an AgentDeployment is deleted + // without firing a watch event (e.g. namespace force-deletion), the usage + // counters will self-heal on the next reconcile rather than staying stale. + return ctrl.Result{RequeueAfter: 5 * time.Minute}, nil } -// SetupWithManager sets up the controller with the Manager. +// SetupWithManager sets up the TenantQuota controller with the Manager. +// It watches both TenantQuota objects and AgentDeployment objects (to trigger +// reconciliation when ADs are created, updated, or deleted in the namespace). func (r *TenantQuotaReconciler) SetupWithManager(mgr ctrl.Manager) error { return ctrl.NewControllerManagedBy(mgr). For(&agentraxv1alpha1.TenantQuota{}). + // Enqueue the owning TenantQuota whenever an AgentDeployment changes. + Watches( + &agentraxv1alpha1.AgentDeployment{}, + enqueueTenantQuota(), + ). Complete(r) } diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index e0631ec..6b3520f 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -18,72 +18,342 @@ package controller import ( "context" + "fmt" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "k8s.io/apimachinery/pkg/api/errors" - "k8s.io/apimachinery/pkg/types" - "sigs.k8s.io/controller-runtime/pkg/reconcile" + apierrors "k8s.io/apimachinery/pkg/api/errors" + apimeta "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" ) var _ = Describe("TenantQuota Controller", func() { - Context("When reconciling a resource", func() { - const resourceName = "test-resource" + // tqNS holds the namespace for this test group. Each Describe block uses a + // unique namespace suffix to avoid cross-test interference. + const tqNS = "tq-ctrl-test" - ctx := context.Background() + ctx := context.Background() - typeNamespacedName := types.NamespacedName{ - Name: resourceName, - Namespace: "default", // TODO(user):Modify as needed + // ── Setup / teardown ────────────────────────────────────────────────────── + + BeforeEach(func() { + By("ensuring the test namespace exists") + ns := namespaceObject(tqNS) + err := k8sClient.Create(ctx, ns) + if err != nil && !apierrors.IsAlreadyExists(err) { + Expect(err).NotTo(HaveOccurred()) } - tenantquota := &agentraxv1alpha1.TenantQuota{} - - BeforeEach(func() { - By("creating the custom resource for the Kind TenantQuota") - err := k8sClient.Get(ctx, typeNamespacedName, tenantquota) - if err != nil && errors.IsNotFound(err) { - resource := &agentraxv1alpha1.TenantQuota{ - ObjectMeta: metav1.ObjectMeta{ - Name: resourceName, - Namespace: "default", - }, - Spec: agentraxv1alpha1.TenantQuotaSpec{ - MaxAgents: 6, - MaxGPUs: 4, - MaxTotalReplicas: 12, - MaxReplicasPerAgent: 6, - }, - } - Expect(k8sClient.Create(ctx, resource)).To(Succeed()) + }) + + AfterEach(func() { + By("deleting all AgentDeployments in the test namespace") + adList := &agentraxv1alpha1.AgentDeploymentList{} + Expect(k8sClient.List(ctx, adList, inNamespace(tqNS))).To(Succeed()) + for i := range adList.Items { + // Remove finalizer so deletion doesn't block. + ad := &adList.Items[i] + ad.Finalizers = nil + _ = k8sClient.Update(ctx, ad) + _ = k8sClient.Delete(ctx, ad) + } + + By("deleting all TenantQuotas in the test namespace") + tqList := &agentraxv1alpha1.TenantQuotaList{} + Expect(k8sClient.List(ctx, tqList, inNamespace(tqNS))).To(Succeed()) + for i := range tqList.Items { + tq := &tqList.Items[i] + _ = k8sClient.Delete(ctx, tq) + } + }) + + // ── Status accuracy ─────────────────────────────────────────────────────── + + Describe("status accuracy", func() { + It("reflects zero usage when no AgentDeployments exist", func() { + tq := makeTQ("tq-empty", tqNS, 6, 4, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-empty", tqNS), fetched)).To(Succeed()) + // Status may not have been written yet if reconcile hasn't run. + // usedAgents == 0 is the expected steady state. + g.Expect(fetched.Status.UsedAgents).To(BeNumerically("==", 0)) + g.Expect(fetched.Status.UsedTotalReplicas).To(BeNumerically("==", 0)) + }, timeout, interval).Should(Succeed()) + }) + + It("increments usedAgents when an AgentDeployment is created (bypassing webhook)", func() { + // We bypass the webhook by directly patching status / using the k8sClient + // with pre-created fixtures. The webhook integration test below covers + // admission-path increments. + tq := makeTQ("tq-count", tqNS, 6, 4, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + // Create an AD via the webhook path (webhook is active in this suite). + ad := makeBasicAD("ad-count-1", tqNS, "tq-count", 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-count", tqNS), fetched)).To(Succeed()) + g.Expect(fetched.Status.UsedAgents).To(BeNumerically("==", 1)) + g.Expect(fetched.Status.UsedTotalReplicas).To(BeNumerically("==", 2)) + }, timeout, interval).Should(Succeed()) + }) + + It("decrements usedAgents when an AgentDeployment is deleted", func() { + tq := makeTQ("tq-del", tqNS, 6, 4, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + ad := makeBasicAD("ad-del-1", tqNS, "tq-del", 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Wait for status to show 1 agent. + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-del", tqNS), fetched)).To(Succeed()) + g.Expect(fetched.Status.UsedAgents).To(BeNumerically("==", 1)) + }, timeout, interval).Should(Succeed()) + + // Remove finalizer so we can delete immediately. + fetched := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, namespacedName("ad-del-1", tqNS), fetched)).To(Succeed()) + fetched.Finalizers = nil + Expect(k8sClient.Update(ctx, fetched)).To(Succeed()) + Expect(k8sClient.Delete(ctx, fetched)).To(Succeed()) + + Eventually(func(g Gomega) { + tqFetched := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-del", tqNS), tqFetched)).To(Succeed()) + g.Expect(tqFetched.Status.UsedAgents).To(BeNumerically("==", 0)) + }, timeout, interval).Should(Succeed()) + }) + }) + + // ── OverQuota condition ─────────────────────────────────────────────────── + + Describe("OverQuota condition", func() { + It("sets OverQuota condition when quota is lowered below current usage", func() { + tq := makeTQ("tq-overq", tqNS, 6, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + // Create 2 ADs within the original quota of 6. + for i := 1; i <= 2; i++ { + ad := makeBasicAD(fmt.Sprintf("ad-overq-%d", i), tqNS, "tq-overq", 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) } + + // Wait for status to reflect 2 agents. + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-overq", tqNS), f)).To(Succeed()) + g.Expect(f.Status.UsedAgents).To(BeNumerically("==", 2)) + }, timeout, interval).Should(Succeed()) + + // Lower maxAgents to 1 while 2 exist. + tqFetched := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName("tq-overq", tqNS), tqFetched)).To(Succeed()) + tqFetched.Spec.MaxAgents = 1 + Expect(k8sClient.Update(ctx, tqFetched)).To(Succeed()) + + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-overq", tqNS), f)).To(Succeed()) + cond := apimeta.FindStatusCondition(f.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) + g.Expect(cond).NotTo(BeNil()) + g.Expect(cond.Status).To(Equal(metav1.ConditionTrue)) + }, timeout, interval).Should(Succeed()) }) - AfterEach(func() { - // TODO(user): Cleanup logic after each test, like removing the resource instance. - resource := &agentraxv1alpha1.TenantQuota{} - err := k8sClient.Get(ctx, typeNamespacedName, resource) - Expect(err).NotTo(HaveOccurred()) + It("clears OverQuota condition when usage returns to within quota", func() { + tq := makeTQ("tq-clearoq", tqNS, 2, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + // Create 2 ADs (at the limit). + for i := 1; i <= 2; i++ { + ad := makeBasicAD(fmt.Sprintf("ad-clearoq-%d", i), tqNS, "tq-clearoq", 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + } + + // Lower maxAgents to 1 → OverQuota. + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-clearoq", tqNS), f)).To(Succeed()) + g.Expect(f.Status.UsedAgents).To(BeNumerically("==", 2)) + }, timeout, interval).Should(Succeed()) - By("Cleanup the specific resource instance TenantQuota") - Expect(k8sClient.Delete(ctx, resource)).To(Succeed()) + tqFetched := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName("tq-clearoq", tqNS), tqFetched)).To(Succeed()) + tqFetched.Spec.MaxAgents = 1 + Expect(k8sClient.Update(ctx, tqFetched)).To(Succeed()) + + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-clearoq", tqNS), f)).To(Succeed()) + cond := apimeta.FindStatusCondition(f.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) + g.Expect(cond).NotTo(BeNil()) + g.Expect(cond.Status).To(Equal(metav1.ConditionTrue)) + }, timeout, interval).Should(Succeed()) + + // Delete one AD → usage falls back to 1 == maxAgents; OverQuota clears. + ad1 := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, namespacedName("ad-clearoq-1", tqNS), ad1)).To(Succeed()) + ad1.Finalizers = nil + Expect(k8sClient.Update(ctx, ad1)).To(Succeed()) + Expect(k8sClient.Delete(ctx, ad1)).To(Succeed()) + + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-clearoq", tqNS), f)).To(Succeed()) + cond := apimeta.FindStatusCondition(f.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) + // Condition should be absent or False once usage normalises. + if cond != nil { + g.Expect(cond.Status).To(Equal(metav1.ConditionFalse)) + } + }, timeout, interval).Should(Succeed()) }) - It("should successfully reconcile the resource", func() { - By("Reconciling the created resource") - controllerReconciler := &TenantQuotaReconciler{ - Client: k8sClient, - Scheme: k8sClient.Scheme(), + + It("never forcibly deletes existing ADs when quota is lowered", func() { + tq := makeTQ("tq-nodel", tqNS, 4, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + for i := 1; i <= 3; i++ { + ad := makeBasicAD(fmt.Sprintf("ad-nodel-%d", i), tqNS, "tq-nodel", 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) } - _, err := controllerReconciler.Reconcile(ctx, reconcile.Request{ - NamespacedName: typeNamespacedName, - }) - Expect(err).NotTo(HaveOccurred()) - // TODO(user): Add more specific assertions depending on your controller's reconciliation logic. - // Example: If you expect a certain status condition after reconciliation, verify it here. + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-nodel", tqNS), f)).To(Succeed()) + g.Expect(f.Status.UsedAgents).To(BeNumerically("==", 3)) + }, timeout, interval).Should(Succeed()) + + // Lower quota to 1. + tqFetched := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName("tq-nodel", tqNS), tqFetched)).To(Succeed()) + tqFetched.Spec.MaxAgents = 1 + Expect(k8sClient.Update(ctx, tqFetched)).To(Succeed()) + + // Wait for OverQuota to be set. + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-nodel", tqNS), f)).To(Succeed()) + cond := apimeta.FindStatusCondition(f.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) + g.Expect(cond).NotTo(BeNil()) + }, timeout, interval).Should(Succeed()) + + // All 3 ADs must still exist — no forced deletion. + adList := &agentraxv1alpha1.AgentDeploymentList{} + Expect(k8sClient.List(ctx, adList, inNamespace(tqNS))).To(Succeed()) + Expect(adList.Items).To(HaveLen(3)) + }) + }) + + // ── Webhook quota rejection ─────────────────────────────────────────────── + + Describe("webhook quota rejection", func() { + It("rejects an AgentDeployment that would exceed maxAgents", func() { + tq := makeTQ("tq-reject", tqNS, 1, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + // First AD is within quota → admitted. + ad1 := makeBasicAD("ad-reject-1", tqNS, "tq-reject", 2) + Expect(k8sClient.Create(ctx, ad1)).To(Succeed()) + + // Second AD would push usedAgents to 2 > maxAgents=1 → rejected. + ad2 := makeBasicAD("ad-reject-2", tqNS, "tq-reject", 2) + err := k8sClient.Create(ctx, ad2) + Expect(err).To(HaveOccurred()) + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue(), + "expected Invalid or Forbidden error, got: %v", err) + }) + + It("rejects an AgentDeployment where replicas.max > maxReplicasPerAgent", func() { + tq := makeTQ("tq-reject-rpa", tqNS, 6, 0, 12, 3) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + // replicas.max=4 > maxReplicasPerAgent=3 → rejected. + ad := makeBasicAD("ad-reject-rpa", tqNS, "tq-reject-rpa", 4) + err := k8sClient.Create(ctx, ad) + Expect(err).To(HaveOccurred()) + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue(), + "expected Invalid or Forbidden error, got: %v", err) + }) + + It("rejects an AgentDeployment that references a non-existent TenantQuota", func() { + ad := makeBasicAD("ad-no-tq", tqNS, "no-such-tq", 2) + err := k8sClient.Create(ctx, ad) + Expect(err).To(HaveOccurred()) + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue(), + "expected Invalid or Forbidden error, got: %v", err) + }) + }) + + // ── Mutating webhook defaults ───────────────────────────────────────────── + + Describe("mutating webhook defaults", func() { + It("defaults port to 8080 when omitted", func() { + tq := makeTQ("tq-defaults", tqNS, 6, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + ad := makeBasicAD("ad-defaults", tqNS, "tq-defaults", 2) + ad.Spec.Port = 0 // explicitly unset + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + fetched := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, namespacedName("ad-defaults", tqNS), fetched)).To(Succeed()) + Expect(fetched.Spec.Port).To(BeNumerically("==", 8080)) + }) + + It("defaults rollout.strategy to Recreate when omitted", func() { + tq := makeTQ("tq-strat-def", tqNS, 6, 0, 12, 6) + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + + ad := makeBasicAD("ad-strat-def", tqNS, "tq-strat-def", 2) + ad.Spec.Rollout.Strategy = "" + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + fetched := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, namespacedName("ad-strat-def", tqNS), fetched)).To(Succeed()) + Expect(fetched.Spec.Rollout.Strategy).To(Equal("Recreate")) }) }) }) + +// ── Test helpers ────────────────────────────────────────────────────────────── + +func makeTQ(name, ns string, maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) *agentraxv1alpha1.TenantQuota { + return &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns}, + Spec: agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: maxAgents, + MaxGPUs: maxGPUs, + MaxTotalReplicas: maxTotalReplicas, + MaxReplicasPerAgent: maxReplicasPerAgent, + }, + } +} + +// makeBasicAD builds a minimal AgentDeployment suitable for TQ reconciler tests. +// It uses a real image name that won't pull (but pod scheduling isn't needed here). +func makeBasicAD(name, ns, tenantRef string, maxReplicas int32) *agentraxv1alpha1.AgentDeployment { + return &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "test-image:v1", + TenantRef: tenantRef, + Port: 8080, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: maxReplicas, + Metric: "queueDepth", + Target: 50, + }, + Rollout: agentraxv1alpha1.RolloutPolicy{Strategy: "Recreate"}, + }, + } +} diff --git a/internal/controller/test_helpers_test.go b/internal/controller/test_helpers_test.go new file mode 100644 index 0000000..ead3df4 --- /dev/null +++ b/internal/controller/test_helpers_test.go @@ -0,0 +1,50 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package controller + +import ( + "time" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" +) + +const ( + // timeout is the maximum time any Eventually assertion waits in this package. + timeout = 30 * time.Second + // interval is how often Eventually polls. + interval = 250 * time.Millisecond +) + +// namespacedName is a convenience wrapper for building types.NamespacedName. +func namespacedName(name, namespace string) types.NamespacedName { + return types.NamespacedName{Name: name, Namespace: namespace} +} + +// namespaceObject creates a Namespace object for use in test setup. +func namespaceObject(name string) *corev1.Namespace { + return &corev1.Namespace{ + ObjectMeta: metav1.ObjectMeta{Name: name}, + } +} + +// inNamespace returns a ListOption that filters by namespace. +func inNamespace(ns string) client.ListOption { + return client.InNamespace(ns) +} diff --git a/internal/quota/enforcer.go b/internal/quota/enforcer.go new file mode 100644 index 0000000..8e0c769 --- /dev/null +++ b/internal/quota/enforcer.go @@ -0,0 +1,302 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package quota implements tenant quota arithmetic and admission-time in-flight +// reservation used by the validating webhook and TenantQuota reconciler. +package quota + +import ( + "fmt" + "sync" + "time" + + corev1 "k8s.io/api/core/v1" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +// DefaultGPUResourceName is the Kubernetes resource name used to count GPU +// units when --gpu-resource-name is not overridden by the operator flag. +const DefaultGPUResourceName = "nvidia.com/gpu" + +// reservationEntry holds in-flight resource counts that have been pre-reserved +// during admission but not yet committed to etcd. +type reservationEntry struct { + agents int32 + gpus int32 + replicas int32 + expiry time.Time +} + +// Enforcer provides quota arithmetic and manages an in-flight reservation map +// to prevent concurrent near-limit creates from both slipping past the quota +// ceiling. The zero value is not usable; use NewEnforcer. +type Enforcer struct { + gpuResourceName string + + // mu guards reservations. + mu sync.Mutex + reservations map[string]*reservationEntry // keyed by "namespace/adName" + + // done is closed by Stop() to terminate the background sweep goroutine. + done chan struct{} + + // nowFn is overridden in tests to control time. + nowFn func() time.Time +} + +// NewEnforcer creates a new Enforcer. gpuResourceName is the Kubernetes resource +// name used to count GPU units (e.g. "nvidia.com/gpu"). Pass DefaultGPUResourceName +// when the operator --gpu-resource-name flag is not overridden. +func NewEnforcer(gpuResourceName string) *Enforcer { + if gpuResourceName == "" { + gpuResourceName = DefaultGPUResourceName + } + e := &Enforcer{ + gpuResourceName: gpuResourceName, + reservations: make(map[string]*reservationEntry), + done: make(chan struct{}), + nowFn: time.Now, + } + // Start the background sweep goroutine to purge expired reservations. + go e.sweepLoop() + return e +} + +// Stop terminates the background sweep goroutine. Call this when the Enforcer +// is no longer needed (e.g. in test teardown) to avoid goroutine leaks. +func (e *Enforcer) Stop() { + close(e.done) +} + +// sweepLoop removes expired in-flight reservations every second. +// It exits cleanly when Stop() is called. +func (e *Enforcer) sweepLoop() { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + e.sweepExpired() + case <-e.done: + return + } + } +} + +// sweepExpired removes all entries whose expiry has passed. +func (e *Enforcer) sweepExpired() { + now := e.nowFn() + e.mu.Lock() + defer e.mu.Unlock() + for k, v := range e.reservations { + if now.After(v.expiry) { + delete(e.reservations, k) + } + } +} + +// extractGPUs returns the GPU count declared in the resource limits for one +// AgentDeployment replica, using the configured GPU resource name. +// Returns 0 if no GPU limit is set. +func (e *Enforcer) extractGPUs(resources corev1.ResourceRequirements) int64 { + if resources.Limits == nil { + return 0 + } + qty, ok := resources.Limits[corev1.ResourceName(e.gpuResourceName)] + if !ok { + return 0 + } + return qty.Value() +} + +// gpusForAD returns the total GPU units for one AgentDeployment: +// gpuPerReplica × spec.replicas.max. +func (e *Enforcer) gpusForAD(ad agentraxv1alpha1.AgentDeploymentSpec) int32 { + perReplica := e.extractGPUs(ad.Resources) + return int32(perReplica) * ad.Replicas.Max +} + +// ComputeUsage aggregates resource usage across a slice of AgentDeployment specs +// belonging to the same tenant. It is called by the TenantQuota reconciler to +// compute the accurate observed state. +func (e *Enforcer) ComputeUsage(ads []agentraxv1alpha1.AgentDeploymentSpec) agentraxv1alpha1.TenantQuotaStatus { + var status agentraxv1alpha1.TenantQuotaStatus + for _, ad := range ads { + status.UsedAgents++ + status.UsedTotalReplicas += ad.Replicas.Max + status.UsedGPUs += e.gpusForAD(ad) + } + return status +} + +// CanAdmit checks whether admitting a new or updated AgentDeployment with the +// given spec would stay within the quota ceilings. It accounts for both the +// already-committed usage (from status) and any in-flight reservations held +// by concurrent admission requests. +// +// admissionKey must be unique per AgentDeployment — use "namespace/adName". +// It identifies the reservation slot for this specific AD so that a retried +// webhook call replaces (not duplicates) any prior reservation, and so that +// sibling ADs in the same TenantQuota each hold independent slots. +// oldSpec may be nil for CREATE requests; for UPDATE requests it is the +// previous spec so we can compute the delta rather than a full new addition. +// +// Returns (true, "") if admission is allowed, or (false, reason) if not. +func (e *Enforcer) CanAdmit( + admissionKey string, + quota agentraxv1alpha1.TenantQuotaSpec, + committedUsage agentraxv1alpha1.TenantQuotaStatus, + requested agentraxv1alpha1.AgentDeploymentSpec, + oldSpec *agentraxv1alpha1.AgentDeploymentSpec, +) (bool, string) { + // Compute the delta this request adds on top of committed usage. + // For CREATE: delta = full requested resources. + // For UPDATE: delta = requested - old (can be negative if scaling down). + var deltaAgents, deltaGPUs, deltaReplicas int32 + if oldSpec == nil { + // CREATE + deltaAgents = 1 + deltaGPUs = e.gpusForAD(requested) + deltaReplicas = requested.Replicas.Max + } else { + // UPDATE — agent count doesn't change, only resources may shift. + deltaAgents = 0 + deltaGPUs = e.gpusForAD(requested) - e.gpusForAD(*oldSpec) + deltaReplicas = requested.Replicas.Max - oldSpec.Replicas.Max + } + + // Add in-flight reservations (excluding this AD's own slot, which will + // be replaced when we Reserve below). + inFlight := e.sumInflight(admissionKey) + + projectedAgents := committedUsage.UsedAgents + inFlight.agents + deltaAgents + projectedGPUs := committedUsage.UsedGPUs + inFlight.gpus + deltaGPUs + projectedReplicas := committedUsage.UsedTotalReplicas + inFlight.replicas + deltaReplicas + + // For UPDATE requests, only reject when the delta increases a dimension that + // is already at or over quota. If quota was lowered below current usage, the + // existing ADs are already OverQuota (indicated by the TQ condition) — we + // must not block updates that don't make things worse, otherwise finalizer + // removal and spec corrections are deadlocked. + isUpdate := oldSpec != nil + + if projectedAgents > quota.MaxAgents && (!isUpdate || deltaAgents > 0) { + return false, fmt.Sprintf( + "would exceed maxAgents (%d): current=%d in-flight=%d delta=%d", + quota.MaxAgents, committedUsage.UsedAgents, inFlight.agents, deltaAgents, + ) + } + if quota.MaxGPUs > 0 && projectedGPUs > quota.MaxGPUs && (!isUpdate || deltaGPUs > 0) { + return false, fmt.Sprintf( + "would exceed maxGPUs (%d): current=%d in-flight=%d delta=%d", + quota.MaxGPUs, committedUsage.UsedGPUs, inFlight.gpus, deltaGPUs, + ) + } + if projectedReplicas > quota.MaxTotalReplicas && (!isUpdate || deltaReplicas > 0) { + return false, fmt.Sprintf( + "would exceed maxTotalReplicas (%d): current=%d in-flight=%d delta=%d", + quota.MaxTotalReplicas, committedUsage.UsedTotalReplicas, inFlight.replicas, deltaReplicas, + ) + } + // MaxReplicasPerAgent is only enforced when the request would increase the + // per-agent replica ceiling. If the quota ceiling was lowered below an + // existing AD's replicas.max, updates that don't raise replicas.max further + // must still be allowed — blocking them would deadlock spec corrections. + if requested.Replicas.Max > quota.MaxReplicasPerAgent && (!isUpdate || requested.Replicas.Max > oldSpec.Replicas.Max) { + return false, fmt.Sprintf( + "spec.replicas.max (%d) exceeds maxReplicasPerAgent (%d)", + requested.Replicas.Max, quota.MaxReplicasPerAgent, + ) + } + + return true, "" +} + +// sumInflight returns the total in-flight resource counts excluding the entry +// for excludeKey (so this AD's existing slot is not double-counted when the +// same AD retries admission). +func (e *Enforcer) sumInflight(excludeKey string) reservationEntry { + e.mu.Lock() + defer e.mu.Unlock() + now := e.nowFn() + var total reservationEntry + for k, v := range e.reservations { + if k == excludeKey { + continue + } + if now.After(v.expiry) { + continue + } + total.agents += v.agents + total.gpus += v.gpus + total.replicas += v.replicas + } + return total +} + +// Reserve creates or replaces an in-flight reservation for admissionKey (a +// per-request unique string, e.g. UID of the AdmissionRequest) lasting ttl. +// The reservation is automatically swept when it expires. Call Release when +// the webhook handler returns (succeeded or failed), or let it expire on its +// own if the process crashes. +func (e *Enforcer) Reserve(admissionKey string, spec agentraxv1alpha1.AgentDeploymentSpec, oldSpec *agentraxv1alpha1.AgentDeploymentSpec, ttl time.Duration) { + var deltaAgents, deltaGPUs, deltaReplicas int32 + if oldSpec == nil { + deltaAgents = 1 + deltaGPUs = e.gpusForAD(spec) + deltaReplicas = spec.Replicas.Max + } else { + deltaGPUs = e.gpusForAD(spec) - e.gpusForAD(*oldSpec) + deltaReplicas = spec.Replicas.Max - oldSpec.Replicas.Max + } + + e.mu.Lock() + defer e.mu.Unlock() + e.reservations[admissionKey] = &reservationEntry{ + agents: deltaAgents, + gpus: deltaGPUs, + replicas: deltaReplicas, + expiry: e.nowFn().Add(ttl), + } +} + +// Release removes the in-flight reservation for admissionKey. +// It is safe to call even if the key does not exist. +func (e *Enforcer) Release(admissionKey string) { + e.mu.Lock() + defer e.mu.Unlock() + delete(e.reservations, admissionKey) +} + +// IsOverQuota returns true and a condition message when the committed usage +// exceeds any of the quota ceilings. Used by the TenantQuota reconciler to +// set the OverQuota condition without forcing deletions. +func (e *Enforcer) IsOverQuota( + quota agentraxv1alpha1.TenantQuotaSpec, + usage agentraxv1alpha1.TenantQuotaStatus, +) (bool, string) { + if usage.UsedAgents > quota.MaxAgents { + return true, fmt.Sprintf("usedAgents (%d) exceeds maxAgents (%d)", usage.UsedAgents, quota.MaxAgents) + } + if quota.MaxGPUs > 0 && usage.UsedGPUs > quota.MaxGPUs { + return true, fmt.Sprintf("usedGPUs (%d) exceeds maxGPUs (%d)", usage.UsedGPUs, quota.MaxGPUs) + } + if usage.UsedTotalReplicas > quota.MaxTotalReplicas { + return true, fmt.Sprintf("usedTotalReplicas (%d) exceeds maxTotalReplicas (%d)", usage.UsedTotalReplicas, quota.MaxTotalReplicas) + } + return false, "" +} diff --git a/internal/quota/enforcer_test.go b/internal/quota/enforcer_test.go new file mode 100644 index 0000000..61ad48e --- /dev/null +++ b/internal/quota/enforcer_test.go @@ -0,0 +1,380 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package quota_test + +import ( + "testing" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/quota" +) + +// ── helpers ────────────────────────────────────────────────────────────────── + +func makeSpec(maxReplicas int32, gpuLimit string) agentraxv1alpha1.AgentDeploymentSpec { + spec := agentraxv1alpha1.AgentDeploymentSpec{ + Image: "test-image:v1", + TenantRef: "test-tenant", + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: maxReplicas, + Metric: "queueDepth", + Target: 50, + }, + } + if gpuLimit != "" { + spec.Resources = corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse(gpuLimit), + }, + } + } + return spec +} + +func makeQuota(maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) agentraxv1alpha1.TenantQuotaSpec { + return agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: maxAgents, + MaxGPUs: maxGPUs, + MaxTotalReplicas: maxTotalReplicas, + MaxReplicasPerAgent: maxReplicasPerAgent, + } +} + +func makeUsage(agents, gpus, replicas int32) agentraxv1alpha1.TenantQuotaStatus { + return agentraxv1alpha1.TenantQuotaStatus{ + UsedAgents: agents, + UsedGPUs: gpus, + UsedTotalReplicas: replicas, + } +} + +// newTestEnforcer creates an Enforcer for tests and registers Stop() as a cleanup +// function so the background sweep goroutine is terminated when the test ends. +func newTestEnforcer(t *testing.T) *quota.Enforcer { + t.Helper() + e := quota.NewEnforcer(quota.DefaultGPUResourceName) + t.Cleanup(e.Stop) + return e +} + +// ── CanAdmit tests ──────────────────────────────────────────────────────────── + +func TestCanAdmit_Create(t *testing.T) { + t.Parallel() + tests := []struct { + name string + quota agentraxv1alpha1.TenantQuotaSpec + usage agentraxv1alpha1.TenantQuotaStatus + spec agentraxv1alpha1.AgentDeploymentSpec + wantAdmit bool + wantContain string // substring that must appear in denial reason + }{ + { + name: "within all limits", + quota: makeQuota(6, 4, 12, 6), + usage: makeUsage(2, 0, 4), + spec: makeSpec(3, ""), + wantAdmit: true, + }, + { + name: "at agent limit → rejected", + quota: makeQuota(3, 4, 12, 6), + usage: makeUsage(3, 0, 6), + spec: makeSpec(2, ""), + wantAdmit: false, + wantContain: "maxAgents", + }, + { + name: "would exceed maxTotalReplicas", + quota: makeQuota(6, 4, 10, 6), + usage: makeUsage(2, 0, 9), + spec: makeSpec(2, ""), // 9+2 = 11 > 10 + wantAdmit: false, + wantContain: "maxTotalReplicas", + }, + { + name: "would exceed maxReplicasPerAgent", + quota: makeQuota(6, 4, 12, 4), + usage: makeUsage(0, 0, 0), + spec: makeSpec(5, ""), // max=5 > maxReplicasPerAgent=4 + wantAdmit: false, + wantContain: "maxReplicasPerAgent", + }, + { + name: "GPU within limit", + quota: makeQuota(6, 4, 12, 6), + usage: makeUsage(1, 2, 3), + spec: makeSpec(3, "1"), // 1 GPU × 3 replicas = 3; 2+3=5 > 4? no, quota is 4 wait: 2+3=5 > 4 = reject + wantAdmit: false, + // actually 1GPU×3replicas=3 + used=2 = 5 > 4 → rejected + wantContain: "maxGPUs", + }, + { + name: "GPU exactly at limit", + quota: makeQuota(6, 4, 12, 6), + usage: makeUsage(1, 2, 3), + spec: makeSpec(2, "1"), // 1 GPU × 2 replicas = 2; 2+2=4 == 4 → ok + wantAdmit: true, + }, + { + name: "no GPU limit set on spec → GPU check skipped", + quota: makeQuota(6, 4, 12, 6), + usage: makeUsage(1, 4, 3), // already at GPU limit + spec: makeSpec(2, ""), // no GPU requested → no increase + wantAdmit: true, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + got, reason := e.CanAdmit("ns/ad-test", tc.quota, tc.usage, tc.spec, nil) + if got != tc.wantAdmit { + t.Errorf("CanAdmit() = %v, want %v; reason: %q", got, tc.wantAdmit, reason) + } + if !tc.wantAdmit && tc.wantContain != "" { + if reason == "" { + t.Errorf("CanAdmit() denied but returned empty reason") + } else if !containsSubstring(reason, tc.wantContain) { + t.Errorf("CanAdmit() reason %q does not contain %q", reason, tc.wantContain) + } + } + }) + } +} + +func TestCanAdmit_Update(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + q := makeQuota(6, 4, 10, 6) + usage := makeUsage(2, 0, 8) + oldSpec := makeSpec(4, "") + newSpec := makeSpec(5, "") // delta replicas = +1; 8+1=9 ≤ 10 → ok + ok, reason := e.CanAdmit("ns/ad-A", q, usage, newSpec, &oldSpec) + if !ok { + t.Errorf("expected update to be admitted but got reason: %q", reason) + } + + // Delta that hits the ceiling exactly → ok. + newSpec2 := makeSpec(6, "") // delta replicas = +2; 8+2=10 ≤ 10 → ok + ok2, _ := e.CanAdmit("ns/ad-A", q, usage, newSpec2, &oldSpec) + if !ok2 { + t.Errorf("expected exact-limit update to be admitted") + } + + // Delta that exceeds ceiling → rejected. + newSpec3 := makeSpec(7, "") // delta replicas = +3; 8+3=11 > 10 → rejected + ok3, reason3 := e.CanAdmit("ns/ad-A", q, usage, newSpec3, &oldSpec) + if ok3 { + t.Errorf("expected over-limit update to be rejected; got reason %q", reason3) + } +} + +func TestCanAdmit_Update_MaxReplicasPerAgent_Downgrade(t *testing.T) { + // When maxReplicasPerAgent is lowered below an existing AD's replicas.max, + // updates that do NOT further increase replicas.max must still be allowed. + // Blocking them would deadlock spec corrections on over-quota ADs. + t.Parallel() + e := newTestEnforcer(t) + q := makeQuota(6, 0, 20, 3) // maxReplicasPerAgent lowered to 3 + usage := makeUsage(1, 0, 5) + + oldSpec := makeSpec(5, "") // existing AD already has max=5 > new ceiling of 3 + + // UPDATE that keeps replicas.max unchanged → must be admitted (no increase). + sameSpec := makeSpec(5, "") + ok, reason := e.CanAdmit("ns/ad-existing", q, usage, sameSpec, &oldSpec) + if !ok { + t.Errorf("update keeping replicas.max unchanged should be allowed after quota downgrade; got: %q", reason) + } + + // UPDATE that reduces replicas.max → must also be admitted. + smallerSpec := makeSpec(4, "") + ok2, reason2 := e.CanAdmit("ns/ad-existing", q, usage, smallerSpec, &oldSpec) + if !ok2 { + t.Errorf("update reducing replicas.max should be allowed; got: %q", reason2) + } + + // UPDATE that further increases replicas.max → must be rejected. + largerSpec := makeSpec(6, "") + ok3, reason3 := e.CanAdmit("ns/ad-existing", q, usage, largerSpec, &oldSpec) + if ok3 { + t.Errorf("update increasing replicas.max beyond maxReplicasPerAgent should be rejected; got reason: %q", reason3) + } + if !containsSubstring(reason3, "maxReplicasPerAgent") { + t.Errorf("denial reason %q should mention maxReplicasPerAgent", reason3) + } +} + +// ── In-flight reservation tests ─────────────────────────────────────────────── + +func TestReservation_BlocksConcurrentCreate(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + q := makeQuota(2, 0, 4, 2) + usage := makeUsage(1, 0, 2) // 1 of 2 agent slots used + + spec := makeSpec(2, "") + + // First admission check passes; then we reserve. + ok1, _ := e.CanAdmit("ns/ad-A", q, usage, spec, nil) + if !ok1 { + t.Fatal("first CanAdmit should have passed") + } + e.Reserve("ns/ad-A", spec, nil, 5*time.Second) + + // Second concurrent request for the same remaining slot should now be blocked + // because ad-A's reservation already claimed it. + ok2, reason2 := e.CanAdmit("ns/ad-B", q, usage, spec, nil) + if ok2 { + t.Errorf("second CanAdmit should have been blocked by in-flight reservation; reason=%q", reason2) + } + + // After releasing ad-A's reservation, the second request passes again. + e.Release("ns/ad-A") + ok3, _ := e.CanAdmit("ns/ad-B", q, usage, spec, nil) + if !ok3 { + t.Error("after Release, CanAdmit should pass again") + } +} + +func TestReservation_DoesNotDoubleCount(t *testing.T) { + // A re-admission for the same AD key should exclude its own prior reservation + // so it isn't double-counted. + t.Parallel() + e := newTestEnforcer(t) + q := makeQuota(3, 0, 6, 3) + usage := makeUsage(1, 0, 2) + spec := makeSpec(2, "") + + e.Reserve("ns/ad-X", spec, nil, 5*time.Second) + + // Calling CanAdmit with the same admissionKey should exclude its own + // reservation from the in-flight sum (no double-count). + ok, _ := e.CanAdmit("ns/ad-X", q, usage, spec, nil) + // usage.agents=1, in-flight from ad-X is excluded, delta=1 → projected=2 ≤ 3 → ok + if !ok { + t.Error("CanAdmit for the same AD key should not be blocked by its own reservation") + } +} + +// ── ComputeUsage tests ──────────────────────────────────────────────────────── + +func TestComputeUsage(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + ads := []agentraxv1alpha1.AgentDeploymentSpec{ + makeSpec(3, "1"), // 1 GPU × 3 = 3 GPU units + makeSpec(2, "2"), // 2 GPU × 2 = 4 GPU units + makeSpec(4, ""), // 0 GPUs + } + got := e.ComputeUsage(ads) + if got.UsedAgents != 3 { + t.Errorf("UsedAgents = %d, want 3", got.UsedAgents) + } + if got.UsedGPUs != 7 { // 3 + 4 + t.Errorf("UsedGPUs = %d, want 7", got.UsedGPUs) + } + if got.UsedTotalReplicas != 9 { // 3+2+4 + t.Errorf("UsedTotalReplicas = %d, want 9", got.UsedTotalReplicas) + } +} + +func TestComputeUsage_Empty(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + got := e.ComputeUsage(nil) + if got.UsedAgents != 0 || got.UsedGPUs != 0 || got.UsedTotalReplicas != 0 { + t.Errorf("expected all-zero for empty input, got %+v", got) + } +} + +// ── IsOverQuota tests ───────────────────────────────────────────────────────── + +func TestIsOverQuota(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + q := makeQuota(3, 4, 10, 3) + + tests := []struct { + name string + usage agentraxv1alpha1.TenantQuotaStatus + wantOQ bool + wantMsg string + }{ + {"within limits", makeUsage(2, 3, 8), false, ""}, + {"agents over", makeUsage(4, 3, 8), true, "maxAgents"}, + {"GPUs over", makeUsage(2, 5, 8), true, "maxGPUs"}, + {"replicas over", makeUsage(2, 3, 11), true, "maxTotalReplicas"}, + } + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + over, msg := e.IsOverQuota(q, tc.usage) + if over != tc.wantOQ { + t.Errorf("IsOverQuota() = %v, want %v; msg=%q", over, tc.wantOQ, msg) + } + if tc.wantOQ && !containsSubstring(msg, tc.wantMsg) { + t.Errorf("IsOverQuota() msg %q does not contain %q", msg, tc.wantMsg) + } + }) + } +} + +// ── ParseErrorRate tests ────────────────────────────────────────────────────── +// ParseErrorRate lives in api/v1alpha1 to avoid an import cycle. +// Its tests are in api/v1alpha1/webhook_test.go. + +func TestParseErrorRate_ViaV1alpha1(t *testing.T) { + // Smoke-test that ParseErrorRate is accessible from outside api/v1alpha1. + t.Parallel() + got, err := agentraxv1alpha1.ParseErrorRate("5%") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if absFloat(got-0.05) > 1e-9 { + t.Errorf("ParseErrorRate(5%%) = %v, want 0.05", got) + } +} + +// ── helpers ─────────────────────────────────────────────────────────────────── + +func containsSubstring(s, sub string) bool { + return len(sub) == 0 || (len(s) >= len(sub) && func() bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false + }()) +} + +func absFloat(f float64) float64 { + if f < 0 { + return -f + } + return f +} diff --git a/internal/webhook/agentdeployment_validator_test.go b/internal/webhook/agentdeployment_validator_test.go new file mode 100644 index 0000000..8f177e6 --- /dev/null +++ b/internal/webhook/agentdeployment_validator_test.go @@ -0,0 +1,421 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package webhook_test + +import ( + "context" + "testing" + "time" + + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/quota" + agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +// newValidatorWithTQ builds a validator backed by a fake client that has one +// pre-existing TenantQuota with the given spec and status in namespace "ns". +func newValidatorWithTQ(t *testing.T, tqSpec agentraxv1alpha1.TenantQuotaSpec, tqStatus agentraxv1alpha1.TenantQuotaStatus) *agentraxwebhook.AgentDeploymentCustomValidator { + t.Helper() + tq := &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{Name: "tq", Namespace: "ns"}, + Spec: tqSpec, + Status: tqStatus, + } + scheme := runtime.NewScheme() + if err := agentraxv1alpha1.AddToScheme(scheme); err != nil { + t.Fatalf("AddToScheme: %v", err) + } + fakeClient := fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(tq). + WithStatusSubresource(tq). + Build() + e := quota.NewEnforcer(quota.DefaultGPUResourceName) + t.Cleanup(e.Stop) + return &agentraxwebhook.AgentDeploymentCustomValidator{Client: fakeClient, Enforcer: e} +} + +// newValidatorNoTQ builds a validator backed by a fake client with no TenantQuota. +func newValidatorNoTQ(t *testing.T) *agentraxwebhook.AgentDeploymentCustomValidator { + t.Helper() + scheme := runtime.NewScheme() + if err := agentraxv1alpha1.AddToScheme(scheme); err != nil { + t.Fatalf("AddToScheme: %v", err) + } + fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build() + e := quota.NewEnforcer(quota.DefaultGPUResourceName) + t.Cleanup(e.Stop) + return &agentraxwebhook.AgentDeploymentCustomValidator{Client: fakeClient, Enforcer: e} +} + +// permissiveTQSpec returns a TenantQuotaSpec with generous headroom. +func permissiveTQSpec() agentraxv1alpha1.TenantQuotaSpec { + return agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 10, MaxGPUs: 0, MaxTotalReplicas: 50, MaxReplicasPerAgent: 10, + } +} + +// validCanaryAD returns an AD with a fully-valid Canary rollout spec. +func validCanaryAD() *agentraxv1alpha1.AgentDeployment { + w10, w100 := int32(10), int32(100) + pause := metav1.Duration{Duration: 5 * time.Minute} + return &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v2", + TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 3, Metric: "queueDepth", Target: 50}, + Rollout: agentraxv1alpha1.RolloutPolicy{ + Strategy: "Canary", + Steps: []agentraxv1alpha1.CanaryStep{ + {SetWeight: &w10}, + {Pause: &pause}, + {SetWeight: &w100}, + }, + Rollback: agentraxv1alpha1.RollbackPolicy{ + MaxErrorRate: "2%", + MaxP99LatencyMs: 3000, + MinRequestSample: 200, + }, + }, + }, + } +} + +// ── ValidateCreate tests ────────────────────────────────────────────────────── + +func TestValidateCreate_HappyPath(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", + TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 3, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err != nil { + t.Errorf("ValidateCreate() unexpected error: %v", err) + } +} + +func TestValidateCreate_TenantRefNotFound(t *testing.T) { + t.Parallel() + v := newValidatorNoTQ(t) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "nonexistent", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 2, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing TenantQuota, got nil") + } +} + +func TestValidateCreate_ReplicasMinGtMax(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 5, Max: 2, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for min > max, got nil") + } +} + +func TestValidateCreate_OverMaxAgents(t *testing.T) { + t.Parallel() + tqSpec := agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 2, MaxGPUs: 0, MaxTotalReplicas: 20, MaxReplicasPerAgent: 10, + } + // Status already shows 2 agents used — a new one would exceed maxAgents. + tqStatus := agentraxv1alpha1.TenantQuotaStatus{UsedAgents: 2} + v := newValidatorWithTQ(t, tqSpec, tqStatus) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad-new", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 2, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected quota rejection for over-maxAgents, got nil") + } +} + +func TestValidateCreate_OverMaxReplicasPerAgent(t *testing.T) { + t.Parallel() + tqSpec := agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 10, MaxGPUs: 0, MaxTotalReplicas: 50, MaxReplicasPerAgent: 3, + } + v := newValidatorWithTQ(t, tqSpec, agentraxv1alpha1.TenantQuotaStatus{}) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + // replicas.max=5 exceeds maxReplicasPerAgent=3 + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 5, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected rejection for maxReplicasPerAgent violation, got nil") + } +} + +func TestValidateCreate_WrongType(t *testing.T) { + t.Parallel() + v := newValidatorNoTQ(t) + _, err := v.ValidateCreate(context.Background(), &agentraxv1alpha1.TenantQuota{}) + if err == nil { + t.Error("expected error for wrong object type, got nil") + } +} + +// ── validateCanarySpec tests ────────────────────────────────────────────────── + +func TestValidateCreate_Canary_HappyPath(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + _, err := v.ValidateCreate(context.Background(), validCanaryAD()) + if err != nil { + t.Errorf("ValidateCreate() for valid canary AD unexpected error: %v", err) + } +} + +func TestValidateCreate_Canary_NoSteps(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + ad.Spec.Rollout.Steps = nil + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing canary steps, got nil") + } +} + +func TestValidateCreate_Canary_StepBothFields(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + w := int32(50) + pause := metav1.Duration{Duration: 5 * time.Minute} + // Set both setWeight and pause on the same step — must be rejected. + ad.Spec.Rollout.Steps[0] = agentraxv1alpha1.CanaryStep{SetWeight: &w, Pause: &pause} + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for step with both setWeight and pause, got nil") + } +} + +func TestValidateCreate_Canary_NoFullWeight(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + w50 := int32(50) + // Replace all steps with one that never reaches 100. + ad.Spec.Rollout.Steps = []agentraxv1alpha1.CanaryStep{{SetWeight: &w50}} + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing setWeight: 100 step, got nil") + } +} + +func TestValidateCreate_Canary_MissingMaxErrorRate(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + ad.Spec.Rollout.Rollback.MaxErrorRate = "" + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing maxErrorRate, got nil") + } +} + +func TestValidateCreate_Canary_InvalidMaxErrorRate(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + ad.Spec.Rollout.Rollback.MaxErrorRate = "not-a-percent" + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for invalid maxErrorRate format, got nil") + } +} + +func TestValidateCreate_Canary_MissingMaxP99(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + ad.Spec.Rollout.Rollback.MaxP99LatencyMs = 0 + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing maxP99LatencyMs, got nil") + } +} + +func TestValidateCreate_Canary_MissingMinRequestSample(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := validCanaryAD() + ad.Spec.Rollout.Rollback.MinRequestSample = 0 + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for missing minRequestSample, got nil") + } +} + +// ── validateMCPTools tests ──────────────────────────────────────────────────── + +func TestValidateCreate_MCPTools_EmptyName(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 2, Metric: "queueDepth", Target: 50}, + MCP: agentraxv1alpha1.MCPConfig{Expose: true, Tools: []string{"search", ""}}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for empty MCP tool name, got nil") + } +} + +func TestValidateCreate_MCPTools_Duplicate(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 2, Metric: "queueDepth", Target: 50}, + MCP: agentraxv1alpha1.MCPConfig{Expose: true, Tools: []string{"search", "search"}}, + }, + } + _, err := v.ValidateCreate(context.Background(), ad) + if err == nil { + t.Error("expected error for duplicate MCP tool name, got nil") + } +} + +// ── ValidateUpdate tests ────────────────────────────────────────────────────── + +func TestValidateUpdate_DeletionTimestampBypass(t *testing.T) { + t.Parallel() + // Even with a TQ that would reject quota, updates with DeletionTimestamp + // must always pass — otherwise finalizer removal is deadlocked. + tqSpec := agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 0, MaxGPUs: 0, MaxTotalReplicas: 0, MaxReplicasPerAgent: 0, + } + v := newValidatorWithTQ(t, tqSpec, agentraxv1alpha1.TenantQuotaStatus{}) + now := metav1.Now() + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "ad", Namespace: "ns", DeletionTimestamp: &now, + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 5, Metric: "queueDepth", Target: 50}, + }, + } + _, err := v.ValidateUpdate(context.Background(), ad, ad) + if err != nil { + t.Errorf("ValidateUpdate() with DeletionTimestamp should bypass all checks, got: %v", err) + } +} + +func TestValidateUpdate_RolloutInProgress_BlocksImageChange(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + oldAD := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 3, Metric: "queueDepth", Target: 50}, + }, + Status: agentraxv1alpha1.AgentDeploymentStatus{Phase: agentraxv1alpha1.PhaseRolloutInProgress}, + } + newAD := oldAD.DeepCopy() + newAD.Spec.Image = "img:v2" // image change while rollout is active + _, err := v.ValidateUpdate(context.Background(), oldAD, newAD) + if err == nil { + t.Error("expected rejection of image change during RolloutInProgress, got nil") + } +} + +func TestValidateUpdate_RolloutInProgress_AllowsNonImageChange(t *testing.T) { + t.Parallel() + v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) + oldAD := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad", Namespace: "ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "img:v1", TenantRef: "tq", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 3, Metric: "queueDepth", Target: 50}, + }, + Status: agentraxv1alpha1.AgentDeploymentStatus{Phase: agentraxv1alpha1.PhaseRolloutInProgress}, + } + newAD := oldAD.DeepCopy() + // Non-image change (e.g. label): must not be blocked by the rollout guard. + newAD.Labels = map[string]string{"env": "staging"} + _, err := v.ValidateUpdate(context.Background(), oldAD, newAD) + if err != nil { + t.Errorf("non-image update during RolloutInProgress should be allowed, got: %v", err) + } +} + +func TestValidateUpdate_WrongType(t *testing.T) { + t.Parallel() + v := newValidatorNoTQ(t) + _, err := v.ValidateUpdate(context.Background(), + &agentraxv1alpha1.TenantQuota{}, + &agentraxv1alpha1.TenantQuota{}, + ) + if err == nil { + t.Error("expected error for wrong object type in ValidateUpdate, got nil") + } +} + +// ── ValidateDelete test ─────────────────────────────────────────────────────── + +func TestValidateDelete_AlwaysNil(t *testing.T) { + t.Parallel() + v := newValidatorNoTQ(t) + _, err := v.ValidateDelete(context.Background(), &agentraxv1alpha1.AgentDeployment{}) + if err != nil { + t.Errorf("ValidateDelete() should always return nil, got: %v", err) + } +} diff --git a/internal/webhook/agentdeployment_webhook.go b/internal/webhook/agentdeployment_webhook.go new file mode 100644 index 0000000..8913faf --- /dev/null +++ b/internal/webhook/agentdeployment_webhook.go @@ -0,0 +1,326 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package webhook implements the validating and mutating admission webhooks for +// Agentrax CRDs. It imports internal/quota to enforce tenant resource ceilings +// at admission time, which is why it lives in internal/ rather than api/ — +// keeping the import graph acyclic (api/v1alpha1 ← internal/webhook, never the reverse). +package webhook + +import ( + "context" + "fmt" + "strings" + "time" + + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/resource" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/util/validation/field" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + logf "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/webhook" + "sigs.k8s.io/controller-runtime/pkg/webhook/admission" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/quota" +) + +var webhookLog = logf.Log.WithName("agentdeployment-webhook") + +// reservationTTL is how long an in-flight reservation is kept while the +// webhook response is in transit. Generous to tolerate slow API server calls. +const reservationTTL = 5 * time.Second + +// ── SetupWebhookWithManager ─────────────────────────────────────────────────── + +// SetupAgentDeploymentWebhookWithManager registers the mutating and validating +// webhook handlers for AgentDeployment with the controller-runtime Manager. +// enforcer is the shared Enforcer instance (must be non-nil). +func SetupAgentDeploymentWebhookWithManager(mgr ctrl.Manager, enforcer *quota.Enforcer) error { + return ctrl.NewWebhookManagedBy(mgr). + For(&agentraxv1alpha1.AgentDeployment{}). + WithDefaulter(&AgentDeploymentCustomDefaulter{}). + WithValidator(&AgentDeploymentCustomValidator{ + Client: mgr.GetClient(), + Enforcer: enforcer, + }). + Complete() +} + +// ── Mutating webhook (defaulter) ───────────────────────────────────────────── + +// AgentDeploymentCustomDefaulter applies defaults to AgentDeployment specs +// before they are persisted. It implements admission.CustomDefaulter. +type AgentDeploymentCustomDefaulter struct{} + +var _ webhook.CustomDefaulter = &AgentDeploymentCustomDefaulter{} + +// Default fills in missing optional fields with sensible defaults. +func (d *AgentDeploymentCustomDefaulter) Default(_ context.Context, obj runtime.Object) error { + ad, ok := obj.(*agentraxv1alpha1.AgentDeployment) + if !ok { + return fmt.Errorf("expected AgentDeployment, got %T", obj) + } + webhookLog.Info("applying defaults", "name", ad.Name, "namespace", ad.Namespace) + + // Default port to 8080 when omitted. + if ad.Spec.Port == 0 { + ad.Spec.Port = 8080 + } + + // Default rollout strategy to Recreate when omitted. + if ad.Spec.Rollout.Strategy == "" { + ad.Spec.Rollout.Strategy = "Recreate" + } + + // Default resources to a conservative baseline when the entire Resources + // field is zero-value (no requests AND no limits). We do not overwrite + // partially-specified resource requirements. + if isZeroResources(ad.Spec.Resources) { + ad.Spec.Resources = corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("100m"), + corev1.ResourceMemory: resource.MustParse("128Mi"), + }, + Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("500m"), + corev1.ResourceMemory: resource.MustParse("512Mi"), + }, + } + } + + return nil +} + +// isZeroResources returns true when both Requests and Limits are nil or empty. +func isZeroResources(r corev1.ResourceRequirements) bool { + return len(r.Requests) == 0 && len(r.Limits) == 0 +} + +// ── Validating webhook ──────────────────────────────────────────────────────── + +// AgentDeploymentCustomValidator validates AgentDeployment specs at admission +// time. It implements admission.CustomValidator. +type AgentDeploymentCustomValidator struct { + Client client.Client + Enforcer *quota.Enforcer +} + +var _ webhook.CustomValidator = &AgentDeploymentCustomValidator{} + +// ValidateCreate validates a new AgentDeployment against quota and spec rules. +func (v *AgentDeploymentCustomValidator) ValidateCreate(ctx context.Context, obj runtime.Object) (admission.Warnings, error) { + ad, ok := obj.(*agentraxv1alpha1.AgentDeployment) + if !ok { + return nil, fmt.Errorf("expected AgentDeployment, got %T", obj) + } + webhookLog.Info("validating create", "name", ad.Name, "namespace", ad.Namespace) + + allErrs := v.validateSpec(ctx, ad, nil) + if len(allErrs) > 0 { + return nil, apierrors.NewInvalid( + agentraxv1alpha1.GroupVersion.WithKind("AgentDeployment").GroupKind(), + ad.Name, allErrs) + } + return nil, nil +} + +// ValidateUpdate validates an updated AgentDeployment against quota and spec rules. +func (v *AgentDeploymentCustomValidator) ValidateUpdate(ctx context.Context, oldObj, newObj runtime.Object) (admission.Warnings, error) { + ad, ok := newObj.(*agentraxv1alpha1.AgentDeployment) + if !ok { + return nil, fmt.Errorf("expected AgentDeployment, got %T", newObj) + } + oldAD, ok := oldObj.(*agentraxv1alpha1.AgentDeployment) + if !ok { + return nil, fmt.Errorf("expected old AgentDeployment, got %T", oldObj) + } + webhookLog.Info("validating update", "name", ad.Name, "namespace", ad.Namespace) + + // When the object is being deleted (DeletionTimestamp set), the reconciler is + // updating it only to strip the finalizer. Blocking that with quota or TQ + // checks would deadlock deletion — particularly if the TenantQuota was already + // deleted before the AD's finalizer could be removed. + if ad.DeletionTimestamp != nil { + return nil, nil + } + + // Block image changes while a rollout is already in progress. + if oldAD.Spec.Image != ad.Spec.Image && oldAD.Status.Phase == agentraxv1alpha1.PhaseRolloutInProgress { + return nil, apierrors.NewForbidden( + agentraxv1alpha1.GroupVersion.WithResource("agentdeployments").GroupResource(), + ad.Name, + fmt.Errorf("image update rejected: a rollout is already in progress (status.phase=%s)", + agentraxv1alpha1.PhaseRolloutInProgress), + ) + } + + allErrs := v.validateSpec(ctx, ad, &oldAD.Spec) + if len(allErrs) > 0 { + return nil, apierrors.NewInvalid( + agentraxv1alpha1.GroupVersion.WithKind("AgentDeployment").GroupKind(), + ad.Name, allErrs) + } + return nil, nil +} + +// ValidateDelete is a no-op; deletion is controlled by the finalizer. +func (v *AgentDeploymentCustomValidator) ValidateDelete(_ context.Context, _ runtime.Object) (admission.Warnings, error) { + return nil, nil +} + +// validateSpec runs all spec-level validation rules and quota checks. +// oldSpec is nil for CREATE; non-nil for UPDATE (used for delta quota arithmetic). +func (v *AgentDeploymentCustomValidator) validateSpec( + ctx context.Context, + ad *agentraxv1alpha1.AgentDeployment, + oldSpec *agentraxv1alpha1.AgentDeploymentSpec, +) field.ErrorList { + var allErrs field.ErrorList + specPath := field.NewPath("spec") + + // ── 1. tenantRef must reference an existing TenantQuota in the same namespace ── + tq := &agentraxv1alpha1.TenantQuota{} + tqKey := client.ObjectKey{Namespace: ad.Namespace, Name: ad.Spec.TenantRef} + if err := v.Client.Get(ctx, tqKey, tq); err != nil { + if apierrors.IsNotFound(err) { + allErrs = append(allErrs, field.Invalid( + specPath.Child("tenantRef"), ad.Spec.TenantRef, + fmt.Sprintf("TenantQuota %q not found in namespace %q", ad.Spec.TenantRef, ad.Namespace), + )) + } else { + allErrs = append(allErrs, field.InternalError(specPath.Child("tenantRef"), err)) + } + // Cannot proceed with quota checks without a valid TQ. + return allErrs + } + + // ── 2. replicas.min ≤ replicas.max ── + replicasPath := specPath.Child("replicas") + if ad.Spec.Replicas.Min > ad.Spec.Replicas.Max { + allErrs = append(allErrs, field.Invalid( + replicasPath.Child("min"), ad.Spec.Replicas.Min, + fmt.Sprintf("must be ≤ spec.replicas.max (%d)", ad.Spec.Replicas.Max), + )) + } + + // ── 3. Canary-specific spec rules ── + if ad.Spec.Rollout.Strategy == "Canary" { + allErrs = append(allErrs, v.validateCanarySpec(ad)...) + } + + // ── 4. MCP tools uniqueness ── + if len(ad.Spec.MCP.Tools) > 0 { + allErrs = append(allErrs, validateMCPTools(specPath.Child("mcp", "tools"), ad.Spec.MCP.Tools)...) + } + + // ── 5. Quota admission check ── + // Compute current usage from the TenantQuota status. The status is kept + // accurate by the TenantQuota reconciler; we add in-flight reservations to + // handle concurrent near-limit creates. + admissionKey := fmt.Sprintf("%s/%s", ad.Namespace, ad.Name) + ok, reason := v.Enforcer.CanAdmit(admissionKey, tq.Spec, tq.Status, ad.Spec, oldSpec) + if !ok { + allErrs = append(allErrs, field.Forbidden(specPath, fmt.Sprintf("quota exceeded: %s", reason))) + } else { + // Reserve in-flight slot for the duration the webhook response is in transit. + // The reservation is automatically swept after reservationTTL. + v.Enforcer.Reserve(admissionKey, ad.Spec, oldSpec, reservationTTL) + } + + return allErrs +} + +// validateCanarySpec checks constraints that only apply when strategy == Canary. +func (v *AgentDeploymentCustomValidator) validateCanarySpec(ad *agentraxv1alpha1.AgentDeployment) field.ErrorList { + var errs field.ErrorList + rolloutPath := field.NewPath("spec", "rollout") + + // steps must be non-empty. + if len(ad.Spec.Rollout.Steps) == 0 { + errs = append(errs, field.Required( + rolloutPath.Child("steps"), + "at least one step is required when strategy is Canary", + )) + } else { + // Each step must set exactly one of setWeight or pause. + for i, step := range ad.Spec.Rollout.Steps { + stepPath := rolloutPath.Child("steps").Index(i) + setBoth := step.SetWeight != nil && step.Pause != nil + setNeither := step.SetWeight == nil && step.Pause == nil + if setBoth || setNeither { + errs = append(errs, field.Invalid( + stepPath, step, + "exactly one of setWeight or pause must be set per step", + )) + } + } + + // At least one setWeight step must reach 100 (full promotion). + hasFullWeight := false + for _, step := range ad.Spec.Rollout.Steps { + if step.SetWeight != nil && *step.SetWeight == 100 { + hasFullWeight = true + break + } + } + if !hasFullWeight { + errs = append(errs, field.Invalid( + rolloutPath.Child("steps"), ad.Spec.Rollout.Steps, + "canary steps must include at least one setWeight: 100 for full promotion", + )) + } + } + + // All rollback fields are required when strategy is Canary. + rb := ad.Spec.Rollout.Rollback + rollbackPath := rolloutPath.Child("rollback") + if rb.MaxErrorRate == "" { + errs = append(errs, field.Required(rollbackPath.Child("maxErrorRate"), + "required when strategy is Canary")) + } else if _, err := agentraxv1alpha1.ParseErrorRate(rb.MaxErrorRate); err != nil { + errs = append(errs, field.Invalid(rollbackPath.Child("maxErrorRate"), rb.MaxErrorRate, err.Error())) + } + if rb.MaxP99LatencyMs == 0 { + errs = append(errs, field.Required(rollbackPath.Child("maxP99LatencyMs"), + "required when strategy is Canary")) + } + if rb.MinRequestSample == 0 { + errs = append(errs, field.Required(rollbackPath.Child("minRequestSample"), + "required when strategy is Canary; must be > 0 to prevent false positives at low traffic")) + } + + return errs +} + +// validateMCPTools checks that tool names are unique and non-empty. +func validateMCPTools(fldPath *field.Path, tools []string) field.ErrorList { + var errs field.ErrorList + seen := make(map[string]bool, len(tools)) + for i, t := range tools { + if strings.TrimSpace(t) == "" { + errs = append(errs, field.Invalid(fldPath.Index(i), t, "tool name must not be empty")) + } + if seen[t] { + errs = append(errs, field.Invalid(fldPath.Index(i), t, fmt.Sprintf("duplicate tool name %q", t))) + } + seen[t] = true + } + return errs +} diff --git a/internal/webhook/agentdeployment_webhook_test.go b/internal/webhook/agentdeployment_webhook_test.go new file mode 100644 index 0000000..910ce05 --- /dev/null +++ b/internal/webhook/agentdeployment_webhook_test.go @@ -0,0 +1,136 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package webhook_test + +import ( + "context" + "testing" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" +) + +// newDefaulter constructs the defaulter under test. +func newDefaulter() *agentraxwebhook.AgentDeploymentCustomDefaulter { + return &agentraxwebhook.AgentDeploymentCustomDefaulter{} +} + +func baseAD() *agentraxv1alpha1.AgentDeployment { + return &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "test-ad", Namespace: "test-ns"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "test:v1", + TenantRef: "tenant", + Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 2, Metric: "queueDepth", Target: 50}, + }, + } +} + +func TestDefaulter_PortDefaulted(t *testing.T) { + t.Parallel() + ad := baseAD() + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + if ad.Spec.Port != 8080 { + t.Errorf("Port = %d, want 8080", ad.Spec.Port) + } +} + +func TestDefaulter_PortNotOverwritten(t *testing.T) { + t.Parallel() + ad := baseAD() + ad.Spec.Port = 9090 + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + if ad.Spec.Port != 9090 { + t.Errorf("Port was overwritten: got %d, want 9090", ad.Spec.Port) + } +} + +func TestDefaulter_StrategyDefaulted(t *testing.T) { + t.Parallel() + ad := baseAD() + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + if ad.Spec.Rollout.Strategy != "Recreate" { + t.Errorf("Strategy = %q, want Recreate", ad.Spec.Rollout.Strategy) + } +} + +func TestDefaulter_StrategyNotOverwritten(t *testing.T) { + t.Parallel() + ad := baseAD() + ad.Spec.Rollout.Strategy = "Canary" + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + if ad.Spec.Rollout.Strategy != "Canary" { + t.Errorf("Strategy was overwritten: got %q, want Canary", ad.Spec.Rollout.Strategy) + } +} + +func TestDefaulter_ResourcesDefaulted(t *testing.T) { + t.Parallel() + ad := baseAD() + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + if len(ad.Spec.Resources.Requests) == 0 { + t.Error("Resources.Requests not defaulted") + } + if len(ad.Spec.Resources.Limits) == 0 { + t.Error("Resources.Limits not defaulted") + } + // Verify the defaults are sane values. + cpuReq := ad.Spec.Resources.Requests[corev1.ResourceCPU] + if cpuReq.Cmp(resource.MustParse("100m")) != 0 { + t.Errorf("default CPU request = %s, want 100m", cpuReq.String()) + } +} + +func TestDefaulter_ResourcesNotOverwritten(t *testing.T) { + t.Parallel() + ad := baseAD() + ad.Spec.Resources = corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("200m"), + }, + } + if err := newDefaulter().Default(context.Background(), ad); err != nil { + t.Fatalf("Default() error: %v", err) + } + cpu := ad.Spec.Resources.Requests[corev1.ResourceCPU] + if cpu.Cmp(resource.MustParse("200m")) != 0 { + t.Errorf("Resources.Requests[cpu] was changed: got %s, want 200m", cpu.String()) + } +} + +func TestDefaulter_WrongType(t *testing.T) { + t.Parallel() + d := newDefaulter() + // Passing a non-AgentDeployment object should return an error. + if err := d.Default(context.Background(), &agentraxv1alpha1.TenantQuota{}); err == nil { + t.Error("expected error when passing wrong type, got nil") + } +} From 6c5c2da2413adad016dacd762fb7192a4ff104f4 Mon Sep 17 00:00:00 2001 From: "Ankit Kr. Chowdhury" Date: Fri, 7 Aug 2026 09:23:00 +0000 Subject: [PATCH 2/4] refactor: rename test resources for clarity and add unparam lint suppressions to test helpers --- internal/controller/tenantquota_controller_test.go | 10 +++++----- internal/controller/test_helpers_test.go | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index 6b3520f..fa65fea 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -310,15 +310,15 @@ var _ = Describe("TenantQuota Controller", func() { }) It("defaults rollout.strategy to Recreate when omitted", func() { - tq := makeTQ("tq-strat-def", tqNS, 6, 0, 12, 6) + tq := makeTQ("tq-strategy-def", tqNS, 6, 0, 12, 6) Expect(k8sClient.Create(ctx, tq)).To(Succeed()) - ad := makeBasicAD("ad-strat-def", tqNS, "tq-strat-def", 2) + ad := makeBasicAD("ad-strategy-def", tqNS, "tq-strategy-def", 2) ad.Spec.Rollout.Strategy = "" Expect(k8sClient.Create(ctx, ad)).To(Succeed()) fetched := &agentraxv1alpha1.AgentDeployment{} - Expect(k8sClient.Get(ctx, namespacedName("ad-strat-def", tqNS), fetched)).To(Succeed()) + Expect(k8sClient.Get(ctx, namespacedName("ad-strategy-def", tqNS), fetched)).To(Succeed()) Expect(fetched.Spec.Rollout.Strategy).To(Equal("Recreate")) }) }) @@ -326,7 +326,7 @@ var _ = Describe("TenantQuota Controller", func() { // ── Test helpers ────────────────────────────────────────────────────────────── -func makeTQ(name, ns string, maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) *agentraxv1alpha1.TenantQuota { +func makeTQ(name, ns string, maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) *agentraxv1alpha1.TenantQuota { //nolint:unparam return &agentraxv1alpha1.TenantQuota{ ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns}, Spec: agentraxv1alpha1.TenantQuotaSpec{ @@ -340,7 +340,7 @@ func makeTQ(name, ns string, maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPe // makeBasicAD builds a minimal AgentDeployment suitable for TQ reconciler tests. // It uses a real image name that won't pull (but pod scheduling isn't needed here). -func makeBasicAD(name, ns, tenantRef string, maxReplicas int32) *agentraxv1alpha1.AgentDeployment { +func makeBasicAD(name, ns, tenantRef string, maxReplicas int32) *agentraxv1alpha1.AgentDeployment { //nolint:unparam return &agentraxv1alpha1.AgentDeployment{ ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns}, Spec: agentraxv1alpha1.AgentDeploymentSpec{ diff --git a/internal/controller/test_helpers_test.go b/internal/controller/test_helpers_test.go index ead3df4..9eea77c 100644 --- a/internal/controller/test_helpers_test.go +++ b/internal/controller/test_helpers_test.go @@ -33,7 +33,7 @@ const ( ) // namespacedName is a convenience wrapper for building types.NamespacedName. -func namespacedName(name, namespace string) types.NamespacedName { +func namespacedName(name, namespace string) types.NamespacedName { //nolint:unparam return types.NamespacedName{Name: name, Namespace: namespace} } From 14cfebeae6905c2636c0ecc6b10b0d521da00207 Mon Sep 17 00:00:00 2001 From: "Ankit Kr. Chowdhury" Date: Fri, 7 Aug 2026 09:40:36 +0000 Subject: [PATCH 3/4] fix: improve arithmetic safety, parsing validation, and test reliability in quota enforcer and controller logic --- api/v1alpha1/error_rate.go | 17 ++++++++++++--- api/v1alpha1/error_rate_test.go | 6 ++++++ internal/controller/suite_test.go | 2 ++ .../controller/tenantquota_controller_test.go | 18 ++++++++++++---- internal/quota/enforcer.go | 21 +++++++++++++------ internal/quota/enforcer_test.go | 18 ++++------------ 6 files changed, 55 insertions(+), 27 deletions(-) diff --git a/api/v1alpha1/error_rate.go b/api/v1alpha1/error_rate.go index 21c6514..6c4bc02 100644 --- a/api/v1alpha1/error_rate.go +++ b/api/v1alpha1/error_rate.go @@ -16,7 +16,11 @@ limitations under the License. package v1alpha1 -import "fmt" +import ( + "fmt" + "math" + "strconv" +) // ParseErrorRate parses a percentage string like "2%" and returns the float64 // value (e.g. 0.02 for "2%"). Returns an error if the format is invalid. @@ -28,10 +32,17 @@ func ParseErrorRate(s string) (float64, error) { if s[len(s)-1] != '%' { return 0, fmt.Errorf("error rate must end with '%%': got %q", s) } - var pct float64 - if _, err := fmt.Sscanf(s[:len(s)-1], "%f", &pct); err != nil { + // strconv.ParseFloat rejects trailing garbage (e.g. "5x") and leading + // whitespace (e.g. " 5"), unlike fmt.Sscanf which silently ignores them. + pct, err := strconv.ParseFloat(s[:len(s)-1], 64) + if err != nil { return 0, fmt.Errorf("parsing error rate %q: %w", s, err) } + // Reject non-finite values (NaN, ±Inf) that ParseFloat may return for + // inputs like "NaN" or "Inf". + if math.IsNaN(pct) || math.IsInf(pct, 0) { + return 0, fmt.Errorf("error rate %q is not a finite number", s) + } if pct < 0 || pct > 100 { return 0, fmt.Errorf("error rate %q out of range [0, 100]", s) } diff --git a/api/v1alpha1/error_rate_test.go b/api/v1alpha1/error_rate_test.go index 96dde68..5d94d26 100644 --- a/api/v1alpha1/error_rate_test.go +++ b/api/v1alpha1/error_rate_test.go @@ -38,6 +38,12 @@ func TestParseErrorRate(t *testing.T) { {"-1%", 0, true}, {"101%", 0, true}, {"abc%", 0, true}, + // trailing garbage — strconv.ParseFloat must reject these + {"5x%", 0, true}, + // non-finite numeric input + {"NaN%", 0, true}, + // leading whitespace — strconv.ParseFloat must reject " 5" + {" 5%", 0, true}, } for _, tc := range tests { tc := tc diff --git a/internal/controller/suite_test.go b/internal/controller/suite_test.go index cc0c4f8..41eca0b 100644 --- a/internal/controller/suite_test.go +++ b/internal/controller/suite_test.go @@ -166,5 +166,7 @@ var _ = AfterSuite(func() { cancel() // Wait for the manager goroutine to finish before stopping envtest. <-mgrDone + // Stop the shared enforcer's background sweep goroutine. + testEnforcer.Stop() Expect(testEnv.Stop()).To(Succeed()) }) diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index fa65fea..6da930e 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -210,10 +210,9 @@ var _ = Describe("TenantQuota Controller", func() { f := &agentraxv1alpha1.TenantQuota{} g.Expect(k8sClient.Get(ctx, namespacedName("tq-clearoq", tqNS), f)).To(Succeed()) cond := apimeta.FindStatusCondition(f.Status.Conditions, agentraxv1alpha1.ConditionOverQuota) - // Condition should be absent or False once usage normalises. - if cond != nil { - g.Expect(cond.Status).To(Equal(metav1.ConditionFalse)) - } + // The reconciler calls RemoveStatusCondition, so the condition + // must be fully absent once usage normalises. + g.Expect(cond).To(BeNil()) }, timeout, interval).Should(Succeed()) }) @@ -264,6 +263,17 @@ var _ = Describe("TenantQuota Controller", func() { ad1 := makeBasicAD("ad-reject-1", tqNS, "tq-reject", 2) Expect(k8sClient.Create(ctx, ad1)).To(Succeed()) + // Wait for the reconciler to commit ad1's usage into TQ status before + // testing ad2. Without this, ad2 may be tried while the in-flight + // reservation is still live and the quota arithmetic has not yet been + // persisted, causing the rejection to rely solely on the reservation + // TTL which may have already expired. + Eventually(func(g Gomega) { + f := &agentraxv1alpha1.TenantQuota{} + g.Expect(k8sClient.Get(ctx, namespacedName("tq-reject", tqNS), f)).To(Succeed()) + g.Expect(f.Status.UsedAgents).To(BeNumerically("==", 1)) + }, timeout, interval).Should(Succeed()) + // Second AD would push usedAgents to 2 > maxAgents=1 → rejected. ad2 := makeBasicAD("ad-reject-2", tqNS, "tq-reject", 2) err := k8sClient.Create(ctx, ad2) diff --git a/internal/quota/enforcer.go b/internal/quota/enforcer.go index 8e0c769..e2e225e 100644 --- a/internal/quota/enforcer.go +++ b/internal/quota/enforcer.go @@ -20,6 +20,7 @@ package quota import ( "fmt" + "math" "sync" "time" @@ -52,7 +53,9 @@ type Enforcer struct { reservations map[string]*reservationEntry // keyed by "namespace/adName" // done is closed by Stop() to terminate the background sweep goroutine. - done chan struct{} + // stopOnce ensures close(done) is called exactly once. + done chan struct{} + stopOnce sync.Once // nowFn is overridden in tests to control time. nowFn func() time.Time @@ -76,10 +79,10 @@ func NewEnforcer(gpuResourceName string) *Enforcer { return e } -// Stop terminates the background sweep goroutine. Call this when the Enforcer -// is no longer needed (e.g. in test teardown) to avoid goroutine leaks. +// Stop terminates the background sweep goroutine. It is idempotent; repeated +// calls are safe and will not panic. func (e *Enforcer) Stop() { - close(e.done) + e.stopOnce.Do(func() { close(e.done) }) } // sweepLoop removes expired in-flight reservations every second. @@ -124,10 +127,16 @@ func (e *Enforcer) extractGPUs(resources corev1.ResourceRequirements) int64 { } // gpusForAD returns the total GPU units for one AgentDeployment: -// gpuPerReplica × spec.replicas.max. +// gpuPerReplica × spec.replicas.max. The multiplication is performed in +// int64 to prevent overflow, then clamped to the int32 range. func (e *Enforcer) gpusForAD(ad agentraxv1alpha1.AgentDeploymentSpec) int32 { perReplica := e.extractGPUs(ad.Resources) - return int32(perReplica) * ad.Replicas.Max + total := perReplica * int64(ad.Replicas.Max) + const maxInt32 = int64(math.MaxInt32) + if total > maxInt32 { + return math.MaxInt32 + } + return int32(total) } // ComputeUsage aggregates resource usage across a slice of AgentDeployment specs diff --git a/internal/quota/enforcer_test.go b/internal/quota/enforcer_test.go index 61ad48e..018c0b9 100644 --- a/internal/quota/enforcer_test.go +++ b/internal/quota/enforcer_test.go @@ -17,6 +17,7 @@ limitations under the License. package quota_test import ( + "strings" "testing" "time" @@ -156,7 +157,7 @@ func TestCanAdmit_Create(t *testing.T) { if !tc.wantAdmit && tc.wantContain != "" { if reason == "" { t.Errorf("CanAdmit() denied but returned empty reason") - } else if !containsSubstring(reason, tc.wantContain) { + } else if !strings.Contains(reason, tc.wantContain) { t.Errorf("CanAdmit() reason %q does not contain %q", reason, tc.wantContain) } } @@ -222,7 +223,7 @@ func TestCanAdmit_Update_MaxReplicasPerAgent_Downgrade(t *testing.T) { if ok3 { t.Errorf("update increasing replicas.max beyond maxReplicasPerAgent should be rejected; got reason: %q", reason3) } - if !containsSubstring(reason3, "maxReplicasPerAgent") { + if !strings.Contains(reason3, "maxReplicasPerAgent") { t.Errorf("denial reason %q should mention maxReplicasPerAgent", reason3) } } @@ -336,7 +337,7 @@ func TestIsOverQuota(t *testing.T) { if over != tc.wantOQ { t.Errorf("IsOverQuota() = %v, want %v; msg=%q", over, tc.wantOQ, msg) } - if tc.wantOQ && !containsSubstring(msg, tc.wantMsg) { + if tc.wantOQ && !strings.Contains(msg, tc.wantMsg) { t.Errorf("IsOverQuota() msg %q does not contain %q", msg, tc.wantMsg) } }) @@ -361,17 +362,6 @@ func TestParseErrorRate_ViaV1alpha1(t *testing.T) { // ── helpers ─────────────────────────────────────────────────────────────────── -func containsSubstring(s, sub string) bool { - return len(sub) == 0 || (len(s) >= len(sub) && func() bool { - for i := 0; i <= len(s)-len(sub); i++ { - if s[i:i+len(sub)] == sub { - return true - } - } - return false - }()) -} - func absFloat(f float64) float64 { if f < 0 { return -f From ddaa9c31c7c3c3bd90af805ea2881985af19aa01 Mon Sep 17 00:00:00 2001 From: "Ankit Kr. Chowdhury" Date: Fri, 7 Aug 2026 16:49:39 +0000 Subject: [PATCH 4/4] fix: improve deletion reliability in tests, harden GPU quota calculation against overflow, and expand error rate parsing test cases --- api/v1alpha1/error_rate_test.go | 4 +++- .../controller/tenantquota_controller_test.go | 22 +++++++++++++++---- internal/quota/enforcer.go | 13 +++++++++++ 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/api/v1alpha1/error_rate_test.go b/api/v1alpha1/error_rate_test.go index 5d94d26..933b58a 100644 --- a/api/v1alpha1/error_rate_test.go +++ b/api/v1alpha1/error_rate_test.go @@ -40,8 +40,10 @@ func TestParseErrorRate(t *testing.T) { {"abc%", 0, true}, // trailing garbage — strconv.ParseFloat must reject these {"5x%", 0, true}, - // non-finite numeric input + // non-finite numeric input — caught by math.IsNaN / math.IsInf guard {"NaN%", 0, true}, + {"Inf%", 0, true}, + {"-Inf%", 0, true}, // leading whitespace — strconv.ParseFloat must reject " 5" {" 5%", 0, true}, } diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index 6da930e..08ce49d 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -53,11 +53,20 @@ var _ = Describe("TenantQuota Controller", func() { adList := &agentraxv1alpha1.AgentDeploymentList{} Expect(k8sClient.List(ctx, adList, inNamespace(tqNS))).To(Succeed()) for i := range adList.Items { - // Remove finalizer so deletion doesn't block. ad := &adList.Items[i] + // Remove finalizer so deletion is not blocked by the controller. ad.Finalizers = nil - _ = k8sClient.Update(ctx, ad) - _ = k8sClient.Delete(ctx, ad) + if err := k8sClient.Update(ctx, ad); err != nil && !apierrors.IsNotFound(err) { + Expect(err).NotTo(HaveOccurred(), "removing finalizer from AD %s", ad.Name) + } + if err := k8sClient.Delete(ctx, ad); err != nil && !apierrors.IsNotFound(err) { + Expect(err).NotTo(HaveOccurred(), "deleting AD %s", ad.Name) + } + // Wait until the API server confirms the object is gone. + Eventually(func() bool { + err := k8sClient.Get(ctx, namespacedName(ad.Name, tqNS), &agentraxv1alpha1.AgentDeployment{}) + return apierrors.IsNotFound(err) + }, timeout, interval).Should(BeTrue(), "AD %s should be fully deleted", ad.Name) } By("deleting all TenantQuotas in the test namespace") @@ -65,7 +74,9 @@ var _ = Describe("TenantQuota Controller", func() { Expect(k8sClient.List(ctx, tqList, inNamespace(tqNS))).To(Succeed()) for i := range tqList.Items { tq := &tqList.Items[i] - _ = k8sClient.Delete(ctx, tq) + if err := k8sClient.Delete(ctx, tq); err != nil && !apierrors.IsNotFound(err) { + Expect(err).NotTo(HaveOccurred(), "deleting TQ %s", tq.Name) + } } }) @@ -213,6 +224,9 @@ var _ = Describe("TenantQuota Controller", func() { // The reconciler calls RemoveStatusCondition, so the condition // must be fully absent once usage normalises. g.Expect(cond).To(BeNil()) + // Usage counters must reflect the one remaining AD (maxReplicas=2). + g.Expect(f.Status.UsedAgents).To(BeNumerically("==", 1)) + g.Expect(f.Status.UsedTotalReplicas).To(BeNumerically("==", 2)) }, timeout, interval).Should(Succeed()) }) diff --git a/internal/quota/enforcer.go b/internal/quota/enforcer.go index e2e225e..c3e3601 100644 --- a/internal/quota/enforcer.go +++ b/internal/quota/enforcer.go @@ -129,8 +129,21 @@ func (e *Enforcer) extractGPUs(resources corev1.ResourceRequirements) int64 { // gpusForAD returns the total GPU units for one AgentDeployment: // gpuPerReplica × spec.replicas.max. The multiplication is performed in // int64 to prevent overflow, then clamped to the int32 range. +// Fails closed (returns MaxInt32) on unexpected negative inputs or int64 +// overflow so the quota check rejects rather than under-counts. func (e *Enforcer) gpusForAD(ad agentraxv1alpha1.AgentDeploymentSpec) int32 { perReplica := e.extractGPUs(ad.Resources) + // GPU quantities and replica counts must be non-negative. If either is + // negative (should never happen given CRD validation), fail closed. + if perReplica < 0 || ad.Replicas.Max < 0 { + return math.MaxInt32 + } + // Detect int64 multiplication overflow before computing total. + // perReplica and Replicas.Max are both non-negative at this point, so + // overflow can only occur in the positive direction. + if perReplica > 0 && int64(ad.Replicas.Max) > math.MaxInt64/perReplica { + return math.MaxInt32 + } total := perReplica * int64(ad.Replicas.Max) const maxInt32 = int64(math.MaxInt32) if total > maxInt32 {