diff --git a/.golangci.yml b/.golangci.yml index e324620..6e02878 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -21,7 +21,7 @@ linters: enable: - dupl - errcheck - - exportloopref + - copyloopvar - ginkgolinter - goconst - gocyclo diff --git a/Makefile b/Makefile index 2c027f5..9d98e18 100644 --- a/Makefile +++ b/Makefile @@ -76,6 +76,10 @@ lint: golangci-lint ## Run golangci-lint linter lint-fix: golangci-lint ## Run golangci-lint linter and perform fixes $(GOLANGCI_LINT) run --fix +.PHONY: setup-git-hooks +setup-git-hooks: ## Install pre-commit and pre-push Git hooks. + ./hack/setup-git-hooks.sh + ##@ Build .PHONY: build diff --git a/api/v1alpha1/error_rate_test.go b/api/v1alpha1/error_rate_test.go index bcc3793..a799a68 100644 --- a/api/v1alpha1/error_rate_test.go +++ b/api/v1alpha1/error_rate_test.go @@ -49,7 +49,6 @@ func TestParseErrorRate(t *testing.T) { {" 5%", 0, true}, } for _, tc := range tests { - tc := tc t.Run(tc.input, func(t *testing.T) { t.Parallel() got, err := agentraxv1alpha1.ParseErrorRate(tc.input) diff --git a/cmd/main.go b/cmd/main.go index 7a6957a..376fc5b 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -18,8 +18,11 @@ limitations under the License. package main import ( + "context" "crypto/tls" + "errors" "flag" + "net/http" "os" "time" @@ -28,11 +31,14 @@ import ( _ "k8s.io/client-go/plugin/pkg/client/auth" autoscalingv2 "k8s.io/api/autoscaling/v2" + corev1 "k8s.io/api/core/v1" apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" clientgoscheme "k8s.io/client-go/kubernetes/scheme" ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/cache" + "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/healthz" "sigs.k8s.io/controller-runtime/pkg/log/zap" "sigs.k8s.io/controller-runtime/pkg/metrics/filters" @@ -46,6 +52,7 @@ import ( "github.com/gitcommitankit/agentrax/internal/controller" "github.com/gitcommitankit/agentrax/internal/metrics" "github.com/gitcommitankit/agentrax/internal/quota" + "github.com/gitcommitankit/agentrax/internal/registry" "github.com/gitcommitankit/agentrax/internal/rollout" agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" // +kubebuilder:scaffold:imports @@ -56,6 +63,41 @@ var ( setupLog = ctrl.Log.WithName("setup") ) +// registryServerRunnable runs the MCP registry HTTP server on every manager replica. +type registryServerRunnable struct { + registryAddr string + mcpRegistry *registry.Registry +} + +// Start starts the MCP discovery HTTP server and listens until context cancellation. +func (r *registryServerRunnable) Start(ctx context.Context) error { + r.mcpRegistry.Start(ctx) + srv := &http.Server{ + Addr: r.registryAddr, + Handler: r.mcpRegistry.Handler(), + ReadHeaderTimeout: 10 * time.Second, + ReadTimeout: 30 * time.Second, + WriteTimeout: 30 * time.Second, + IdleTimeout: 120 * time.Second, + } + go func() { + <-ctx.Done() + shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _ = srv.Shutdown(shutdownCtx) + }() + setupLog.Info("starting MCP discovery registry server", "addr", r.registryAddr) + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + return err + } + return nil +} + +// NeedLeaderElection returns false so the registry server runs across all manager replicas. +func (r *registryServerRunnable) NeedLeaderElection() bool { + return false +} + // init registers all Kubernetes core, CRD, and monitoring schemes. func init() { utilruntime.Must(clientgoscheme.AddToScheme(scheme)) @@ -68,6 +110,8 @@ func init() { // +kubebuilder:scaffold:scheme } +// +kubebuilder:rbac:groups="",resources=configmaps,verbs=get;list;watch;create;update;patch + // main is the entrypoint for the Agentrax controller manager binary. func main() { var metricsAddr string @@ -79,6 +123,7 @@ func main() { var prometheusURL string var gatewayName string var gatewayNamespace string + var registryAddr 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.") @@ -99,6 +144,8 @@ func main() { "Name of the Gateway API Gateway object used for canary traffic splitting.") flag.StringVar(&gatewayNamespace, "gateway-namespace", "agentrax-system", "Namespace of the Gateway API Gateway object used for canary traffic splitting.") + flag.StringVar(®istryAddr, "registry-bind-address", ":9090", + "The address the MCP discovery registry HTTP endpoint binds to.") opts := zap.Options{ Development: true, } @@ -159,6 +206,12 @@ func main() { metricsServerOptions.FilterProvider = filters.WithAuthenticationAndAuthorization } + // Determine registry namespace once for both manager cache and registry construction. + registryNamespace := os.Getenv("POD_NAMESPACE") + if registryNamespace == "" { + registryNamespace = "agentrax-system" + } + mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{ Scheme: scheme, Metrics: metricsServerOptions, @@ -166,6 +219,15 @@ func main() { HealthProbeBindAddress: probeAddr, LeaderElection: enableLeaderElection, LeaderElectionID: "ddf1aac5.agentrax.io", + Cache: cache.Options{ + ByObject: map[client.Object]cache.ByObject{ + &corev1.ConfigMap{}: { + Namespaces: map[string]cache.Config{ + registryNamespace: {}, + }, + }, + }, + }, // LeaderElectionReleaseOnCancel defines if the leader should step down voluntarily // when the Manager ends. This requires the binary to immediately end when the // Manager is stopped, otherwise, this setting is unsafe. Setting this significantly @@ -204,12 +266,33 @@ func main() { setupLog.Info("canary rollout disabled (no --prometheus-url)") } - if err = (&controller.AgentDeploymentReconciler{ + // Initialize MCP discovery registry and registrar. + mcpRegistry := registry.NewRegistry(mgr.GetClient(), registryNamespace, registry.DefaultTTL) + mcpRegistrar := registry.NewRegistrar(mcpRegistry, registry.NewHTTPMCPClient()) + + if registryAddr != "" && registryAddr != "0" { + registryRunnable := ®istryServerRunnable{ + registryAddr: registryAddr, + mcpRegistry: mcpRegistry, + } + if err := mgr.Add(registryRunnable); err != nil { + setupLog.Error(err, "unable to add registry server to manager") + os.Exit(1) + } + } + + agentDeploymentReconciler := &controller.AgentDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), GPUResourceName: gpuResourceName, CanaryController: canaryController, - }).SetupWithManager(mgr); err != nil { + Registrar: mcpRegistrar, + } + if canaryController != nil { + canaryController.Registrar = mcpRegistrar + } + + if err = agentDeploymentReconciler.SetupWithManager(mgr); err != nil { setupLog.Error(err, "unable to create controller", "controller", "AgentDeployment") os.Exit(1) } diff --git a/config/manager/kustomization.yaml b/config/manager/kustomization.yaml index 5c5f0b8..5eb2c3f 100644 --- a/config/manager/kustomization.yaml +++ b/config/manager/kustomization.yaml @@ -1,2 +1,9 @@ resources: - manager.yaml +- registry_service.yaml +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization +images: +- name: controller + newName: controller + newTag: v0.1.0 diff --git a/config/manager/registry_service.yaml b/config/manager/registry_service.yaml new file mode 100644 index 0000000..1a57013 --- /dev/null +++ b/config/manager/registry_service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: agentrax-registry + namespace: system + labels: + app.kubernetes.io/name: agentrax + app.kubernetes.io/component: registry +spec: + selector: + control-plane: controller-manager + ports: + - name: registry + port: 9090 + targetPort: 9090 + protocol: TCP + type: ClusterIP diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index c2b186e..4c8c90a 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -4,6 +4,17 @@ kind: ClusterRole metadata: name: manager-role rules: +- apiGroups: + - "" + resources: + - configmaps + verbs: + - create + - get + - list + - patch + - update + - watch - apiGroups: - agentrax.io resources: diff --git a/dist/install.yaml b/dist/install.yaml new file mode 100644 index 0000000..8c86c1e --- /dev/null +++ b/dist/install.yaml @@ -0,0 +1,1124 @@ +apiVersion: v1 +kind: Namespace +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + control-plane: controller-manager + name: agentrax-system +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.16.1 + name: agentdeployments.agentrax.io +spec: + group: agentrax.io + names: + kind: AgentDeployment + listKind: AgentDeploymentList + plural: agentdeployments + singular: agentdeployment + scope: Namespaced + versions: + - additionalPrinterColumns: + - jsonPath: .status.phase + name: Phase + type: string + - jsonPath: .status.currentReplicas + name: Replicas + type: integer + - jsonPath: .status.stableVersion + name: Stable + type: string + - jsonPath: .status.registered + name: Registered + type: boolean + - jsonPath: .metadata.creationTimestamp + name: Age + type: date + name: v1alpha1 + schema: + openAPIV3Schema: + description: |- + AgentDeployment is the Schema for the agentdeployments API. + It manages the full lifecycle of a model or autonomous agent workload on Kubernetes, + including autoscaling, canary rollout, and MCP-based service discovery. + properties: + apiVersion: + description: |- + APIVersion defines the versioned schema of this representation of an object. + Servers should convert recognized schemas to the latest internal value, and + may reject unrecognized values. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources + type: string + kind: + description: |- + Kind is a string value representing the REST resource this object represents. + Servers may infer this from the endpoint the client submits requests to. + Cannot be updated. + In CamelCase. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds + type: string + metadata: + type: object + spec: + description: AgentDeploymentSpec defines the desired state of an AgentDeployment. + properties: + args: + description: Args overrides the container entrypoint arguments. + items: + type: string + type: array + env: + description: Env allows passing environment variables to the agent + container. + items: + description: EnvVar represents an environment variable present in + a Container. + properties: + name: + description: Name of the environment variable. Must be a C_IDENTIFIER. + type: string + value: + description: |- + Variable references $(VAR_NAME) are expanded + using the previously defined environment variables in the container and + any service environment variables. If a variable cannot be resolved, + the reference in the input string will be unchanged. Double $$ are reduced + to a single $, which allows for escaping the $(VAR_NAME) syntax: i.e. + "$$(VAR_NAME)" will produce the string literal "$(VAR_NAME)". + Escaped references will never be expanded, regardless of whether the variable + exists or not. + Defaults to "". + type: string + valueFrom: + description: Source for the environment variable's value. Cannot + be used if value is not empty. + properties: + configMapKeyRef: + description: Selects a key of a ConfigMap. + properties: + key: + description: The key to select. + type: string + name: + default: "" + description: |- + Name of the referent. + This field is effectively required, but due to backwards compatibility is + allowed to be empty. Instances of this type with an empty value here are + almost certainly wrong. + More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names + type: string + optional: + description: Specify whether the ConfigMap or its key + must be defined + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + fieldRef: + description: |- + Selects a field of the pod: supports metadata.name, metadata.namespace, `metadata.labels['']`, `metadata.annotations['']`, + spec.nodeName, spec.serviceAccountName, status.hostIP, status.podIP, status.podIPs. + properties: + apiVersion: + description: Version of the schema the FieldPath is + written in terms of, defaults to "v1". + type: string + fieldPath: + description: Path of the field to select in the specified + API version. + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + description: |- + Selects a resource of the container: only resources limits and requests + (limits.cpu, limits.memory, limits.ephemeral-storage, requests.cpu, requests.memory and requests.ephemeral-storage) are currently supported. + properties: + containerName: + description: 'Container name: required for volumes, + optional for env vars' + type: string + divisor: + anyOf: + - type: integer + - type: string + description: Specifies the output format of the exposed + resources, defaults to "1" + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + description: 'Required: resource to select' + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + description: Selects a key of a secret in the pod's namespace + properties: + key: + description: The key of the secret to select from. Must + be a valid secret key. + type: string + name: + default: "" + description: |- + Name of the referent. + This field is effectively required, but due to backwards compatibility is + allowed to be empty. Instances of this type with an empty value here are + almost certainly wrong. + More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names + type: string + optional: + description: Specify whether the Secret or its key must + be defined + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + type: object + required: + - name + type: object + type: array + image: + description: Image is the container image serving the model or agent. + minLength: 1 + type: string + mcp: + description: MCP controls whether this agent registers itself for + MCP-based discovery. + properties: + expose: + description: Expose, when true, registers this agent in the MCP + registry upon reaching a stable state. + type: boolean + tools: + description: Tools is the list of tool names advertised by this + agent in the MCP registry. + items: + type: string + type: array + type: object + port: + default: 8080 + description: Port is the container port the model/agent listens on. + format: int32 + maximum: 65535 + minimum: 1 + type: integer + replicas: + description: Replicas defines the autoscaling policy. + properties: + max: + description: Max is the maximum number of replicas. Must be at + least 1. + format: int32 + minimum: 1 + type: integer + metric: + description: Metric selects the custom metric used for autoscaling + decisions. + enum: + - queueDepth + - gpuUtilization + type: string + min: + description: Min is the minimum number of replicas. Must be at + least 1. + format: int32 + minimum: 1 + type: integer + target: + description: Target is the desired value of the chosen metric + per replica. + format: int32 + minimum: 1 + type: integer + required: + - max + - metric + - min + - target + type: object + resources: + description: Resources are passed through to the underlying pod template. + properties: + claims: + description: |- + Claims lists the names of resources, defined in spec.resourceClaims, + that are used by this container. + + This is an alpha field and requires enabling the + DynamicResourceAllocation feature gate. + + This field is immutable. It can only be set for containers. + items: + description: ResourceClaim references one entry in PodSpec.ResourceClaims. + properties: + name: + description: |- + Name must match the name of one entry in pod.spec.resourceClaims of + the Pod where this field is used. It makes that resource available + inside a container. + type: string + request: + description: |- + Request is the name chosen for a request in the referenced claim. + If empty, everything from the claim is made available, otherwise + only the result of this request. + type: string + required: + - name + type: object + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + limits: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + description: |- + Limits describes the maximum amount of compute resources allowed. + More info: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/ + type: object + requests: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + description: |- + Requests describes the minimum amount of compute resources required. + If Requests is omitted for a container, it defaults to Limits if that is explicitly specified, + otherwise to an implementation-defined value. Requests cannot exceed Limits. + More info: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/ + type: object + type: object + rollout: + description: Rollout defines how new versions are shipped. Optional; + defaults to Recreate. + properties: + abort: + description: Abort, when set to true, triggers an immediate rollback + of any in-progress canary. + type: boolean + rollback: + description: |- + Rollback defines the automatic rollback thresholds. + Required when strategy is Canary. + properties: + maxErrorRate: + description: MaxErrorRate is the maximum acceptable error + rate, expressed as a percentage string (e.g. "2%"). + type: string + maxP99LatencyMs: + description: MaxP99LatencyMs is the maximum acceptable p99 + latency in milliseconds. + format: int32 + minimum: 1 + type: integer + minRequestSample: + description: |- + MinRequestSample is the minimum number of requests that must be observed before + threshold evaluation occurs. Guards against false positives at low traffic weight. + format: int32 + minimum: 1 + type: integer + type: object + steps: + description: |- + Steps lists the ordered canary traffic-shift and pause steps. + Required when strategy is Canary. + items: + description: |- + CanaryStep represents a single step in a canary rollout. + Exactly one of SetWeight or Pause must be set per step. + properties: + pause: + description: |- + Pause is the duration to wait before evaluating rollback thresholds. + Mutually exclusive with SetWeight. + type: string + setWeight: + description: SetWeight is the percentage of traffic to send + to the canary (0–100). + format: int32 + maximum: 100 + minimum: 0 + type: integer + type: object + type: array + strategy: + default: Recreate + description: Strategy determines the rollout approach. Defaults + to Recreate. + enum: + - Recreate + - Canary + type: string + type: object + tenantRef: + description: TenantRef names the owning TenantQuota object in the + same namespace. + minLength: 1 + type: string + required: + - image + - replicas + - tenantRef + type: object + status: + description: AgentDeploymentStatus defines the observed state of an AgentDeployment. + properties: + canaryStepIndex: + description: |- + CanaryStepIndex is the index of the currently executing rollout step. + Persisted so the state machine survives operator restarts. + type: integer + canaryVersion: + description: CanaryVersion is the container image tag of the canary + deployment, if one is in progress. + type: string + canaryWeight: + description: CanaryWeight is the current percentage of traffic routed + to the canary (0–100). + format: int32 + type: integer + conditions: + description: Conditions holds the conditions for the AgentDeployment. + items: + description: Condition contains details for one aspect of the current + state of this API Resource. + properties: + lastTransitionTime: + description: |- + lastTransitionTime is the last time the condition transitioned from one status to another. + This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. + format: date-time + type: string + message: + description: |- + message is a human readable message indicating details about the transition. + This may be an empty string. + maxLength: 32768 + type: string + observedGeneration: + description: |- + observedGeneration represents the .metadata.generation that the condition was set based upon. + For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date + with respect to the current state of the instance. + format: int64 + minimum: 0 + type: integer + reason: + description: |- + reason contains a programmatic identifier indicating the reason for the condition's last transition. + Producers of specific condition types may define expected values and meanings for this field, + and whether the values are considered a guaranteed API. + The value should be a CamelCase string. + This field may not be empty. + maxLength: 1024 + minLength: 1 + pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ + type: string + status: + description: status of the condition, one of True, False, Unknown. + enum: + - "True" + - "False" + - Unknown + type: string + type: + description: type of condition in CamelCase or in foo.example.com/CamelCase. + maxLength: 316 + pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ + type: string + required: + - lastTransitionTime + - message + - reason + - status + - type + type: object + type: array + currentReplicas: + description: CurrentReplicas is the number of replicas currently running. + format: int32 + type: integer + pauseStartedAt: + description: |- + PauseStartedAt records when the current pause step began. + Used to enforce maximum pause extensions and the fail-safe rollback timeout. + format: date-time + type: string + phase: + description: Phase is the high-level lifecycle phase of this deployment. + enum: + - Pending + - Running + - RolloutInProgress + - RolloutFailed + - Degraded + type: string + promUnreachableSince: + description: |- + PromUnreachableSince records when Prometheus last became unreachable. + When non-nil and age exceeds FailSafeTimeout, a fail-safe rollback fires. + format: date-time + type: string + registered: + description: Registered is true when this agent is currently registered + in the MCP registry. + type: boolean + stableVersion: + description: StableVersion is the container image tag of the currently + stable deployment. + type: string + type: object + type: object + served: true + storage: true + subresources: + status: {} +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.16.1 + name: tenantquotas.agentrax.io +spec: + group: agentrax.io + names: + kind: TenantQuota + listKind: TenantQuotaList + plural: tenantquotas + singular: tenantquota + scope: Namespaced + versions: + - additionalPrinterColumns: + - jsonPath: .spec.maxAgents + name: MaxAgents + type: integer + - jsonPath: .status.usedAgents + name: UsedAgents + type: integer + - jsonPath: .spec.maxGPUs + name: MaxGPUs + type: integer + - jsonPath: .status.usedGPUs + name: UsedGPUs + type: integer + - jsonPath: .metadata.creationTimestamp + name: Age + type: date + name: v1alpha1 + schema: + openAPIV3Schema: + description: |- + TenantQuota is the Schema for the tenantquotas API. + It declares and enforces per-tenant resource ceilings for AgentDeployment objects. + properties: + apiVersion: + description: |- + APIVersion defines the versioned schema of this representation of an object. + Servers should convert recognized schemas to the latest internal value, and + may reject unrecognized values. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources + type: string + kind: + description: |- + Kind is a string value representing the REST resource this object represents. + Servers may infer this from the endpoint the client submits requests to. + Cannot be updated. + In CamelCase. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds + type: string + metadata: + type: object + spec: + description: TenantQuotaSpec defines resource ceilings for a tenant. + properties: + maxAgents: + description: MaxAgents is the maximum number of AgentDeployment objects + allowed in this tenant. + format: int32 + minimum: 1 + type: integer + maxGPUs: + description: |- + MaxGPUs is the total number of GPU units that can be allocated across all agents. + GPU count is derived from spec.resources.limits["nvidia.com/gpu"] × spec.replicas.max. + format: int32 + minimum: 0 + type: integer + maxReplicasPerAgent: + description: MaxReplicasPerAgent is the maximum spec.replicas.max + value any single agent may request. + format: int32 + minimum: 1 + type: integer + maxTotalReplicas: + description: MaxTotalReplicas is the ceiling on the sum of spec.replicas.max + across all agents in this tenant. + format: int32 + minimum: 1 + type: integer + required: + - maxAgents + - maxGPUs + - maxReplicasPerAgent + - maxTotalReplicas + type: object + status: + description: TenantQuotaStatus reports current usage against the spec + ceilings. + properties: + conditions: + description: Conditions holds the conditions for the TenantQuota (e.g. + OverQuota). + items: + description: Condition contains details for one aspect of the current + state of this API Resource. + properties: + lastTransitionTime: + description: |- + lastTransitionTime is the last time the condition transitioned from one status to another. + This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. + format: date-time + type: string + message: + description: |- + message is a human readable message indicating details about the transition. + This may be an empty string. + maxLength: 32768 + type: string + observedGeneration: + description: |- + observedGeneration represents the .metadata.generation that the condition was set based upon. + For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date + with respect to the current state of the instance. + format: int64 + minimum: 0 + type: integer + reason: + description: |- + reason contains a programmatic identifier indicating the reason for the condition's last transition. + Producers of specific condition types may define expected values and meanings for this field, + and whether the values are considered a guaranteed API. + The value should be a CamelCase string. + This field may not be empty. + maxLength: 1024 + minLength: 1 + pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ + type: string + status: + description: status of the condition, one of True, False, Unknown. + enum: + - "True" + - "False" + - Unknown + type: string + type: + description: type of condition in CamelCase or in foo.example.com/CamelCase. + maxLength: 316 + pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ + type: string + required: + - lastTransitionTime + - message + - reason + - status + - type + type: object + type: array + usedAgents: + description: UsedAgents is the number of AgentDeployment objects currently + in this tenant. + format: int32 + type: integer + usedGPUs: + description: UsedGPUs is the total GPU units currently allocated across + all agents in this tenant. + format: int32 + type: integer + usedTotalReplicas: + description: UsedTotalReplicas is the sum of spec.replicas.max across + all agents currently in this tenant. + format: int32 + type: integer + type: object + type: object + served: true + storage: true + subresources: + status: {} +--- +apiVersion: v1 +kind: ServiceAccount +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-controller-manager + namespace: agentrax-system +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-leader-election-role + namespace: agentrax-system +rules: +- apiGroups: + - "" + resources: + - configmaps + verbs: + - get + - list + - watch + - create + - update + - patch + - delete +- apiGroups: + - coordination.k8s.io + resources: + - leases + verbs: + - get + - list + - watch + - create + - update + - patch + - delete +- apiGroups: + - "" + resources: + - events + verbs: + - create + - patch +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-agentdeployment-editor-role +rules: +- apiGroups: + - agentrax.agentrax.io + resources: + - agentdeployments + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - agentrax.agentrax.io + resources: + - agentdeployments/status + verbs: + - get +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-agentdeployment-viewer-role +rules: +- apiGroups: + - agentrax.agentrax.io + resources: + - agentdeployments + verbs: + - get + - list + - watch +- apiGroups: + - agentrax.agentrax.io + resources: + - agentdeployments/status + verbs: + - get +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: agentrax-manager-role +rules: +- apiGroups: + - "" + resources: + - configmaps + verbs: + - create + - get + - list + - patch + - update + - watch +- apiGroups: + - agentrax.io + resources: + - agentdeployments + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - agentrax.io + resources: + - agentdeployments/finalizers + verbs: + - update +- apiGroups: + - agentrax.io + resources: + - agentdeployments/status + - tenantquotas/status + verbs: + - get + - patch + - update +- apiGroups: + - agentrax.io + resources: + - tenantquotas + verbs: + - get + - list + - watch +- apiGroups: + - apiextensions.k8s.io + resources: + - customresourcedefinitions + verbs: + - get + - list + - watch +- apiGroups: + - apps + resources: + - deployments + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - autoscaling + resources: + - horizontalpodautoscalers + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - "" + resources: + - events + verbs: + - create + - patch +- apiGroups: + - "" + resources: + - pods + verbs: + - get + - list + - watch +- apiGroups: + - "" + resources: + - services + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - gateway.networking.k8s.io + resources: + - httproutes + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - monitoring.coreos.com + resources: + - servicemonitors + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: agentrax-metrics-auth-role +rules: +- apiGroups: + - authentication.k8s.io + resources: + - tokenreviews + verbs: + - create +- apiGroups: + - authorization.k8s.io + resources: + - subjectaccessreviews + verbs: + - create +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: agentrax-metrics-reader +rules: +- nonResourceURLs: + - /metrics + verbs: + - get +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-tenantquota-editor-role +rules: +- apiGroups: + - agentrax.agentrax.io + resources: + - tenantquotas + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - agentrax.agentrax.io + resources: + - tenantquotas/status + verbs: + - get +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-tenantquota-viewer-role +rules: +- apiGroups: + - agentrax.agentrax.io + resources: + - tenantquotas + verbs: + - get + - list + - watch +- apiGroups: + - agentrax.agentrax.io + resources: + - tenantquotas/status + verbs: + - get +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-leader-election-rolebinding + namespace: agentrax-system +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: agentrax-leader-election-role +subjects: +- kind: ServiceAccount + name: agentrax-controller-manager + namespace: agentrax-system +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + name: agentrax-manager-rolebinding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: agentrax-manager-role +subjects: +- kind: ServiceAccount + name: agentrax-controller-manager + namespace: agentrax-system +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: agentrax-metrics-auth-rolebinding +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: agentrax-metrics-auth-role +subjects: +- kind: ServiceAccount + name: agentrax-controller-manager + namespace: agentrax-system +--- +apiVersion: v1 +kind: Service +metadata: + labels: + app.kubernetes.io/component: registry + app.kubernetes.io/name: agentrax + name: agentrax-agentrax-registry + namespace: agentrax-system +spec: + ports: + - name: registry + port: 9090 + protocol: TCP + targetPort: 9090 + selector: + control-plane: controller-manager + type: ClusterIP +--- +apiVersion: v1 +kind: Service +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + control-plane: controller-manager + name: agentrax-controller-manager-metrics-service + namespace: agentrax-system +spec: + ports: + - name: https + port: 8443 + protocol: TCP + targetPort: 8443 + selector: + control-plane: controller-manager +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + labels: + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/name: agentrax + control-plane: controller-manager + name: agentrax-controller-manager + namespace: agentrax-system +spec: + replicas: 1 + selector: + matchLabels: + control-plane: controller-manager + template: + metadata: + annotations: + kubectl.kubernetes.io/default-container: manager + labels: + control-plane: controller-manager + spec: + containers: + - args: + - --metrics-bind-address=:8443 + - --leader-elect + - --health-probe-bind-address=:8081 + command: + - /manager + image: controller:latest + livenessProbe: + httpGet: + path: /healthz + port: 8081 + initialDelaySeconds: 15 + periodSeconds: 20 + name: manager + readinessProbe: + httpGet: + path: /readyz + port: 8081 + initialDelaySeconds: 5 + periodSeconds: 10 + resources: + limits: + cpu: 500m + memory: 128Mi + requests: + cpu: 10m + memory: 64Mi + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + securityContext: + runAsNonRoot: true + serviceAccountName: agentrax-controller-manager + terminationGracePeriodSeconds: 10 diff --git a/docs/agentrax.md b/docs/agentrax.md index e70f062..94ee94e 100644 --- a/docs/agentrax.md +++ b/docs/agentrax.md @@ -39,7 +39,7 @@ The canonical list of top-level end-to-end test scenarios. Do not add new top-le | 2 | Multi-tenancy | **Done** | `TenantQuota`, validating + mutating webhook, quota enforcement | | 3 | Autoscaling | **Done** | Prometheus Adapter integration, managed HPA | | 4 | Canary rollout | **Done** | Rollout state machine, Gateway API traffic shifting, PromQL threshold evaluation | -| 5 | MCP registry | Pending | Registrar, registry HTTP handler, discovery API | +| 5 | MCP registry | **Done** | Registrar, registry HTTP handler, discovery API, TTL/heartbeat, ConfigMap persistence | | 6 | Hardening & demo | Pending | E2e tests in CI, Helm chart, README, recorded demo | Phases 2, 3, and 5 are independent of each other and can run in parallel once Phase 1 is complete. Phase 4 depends on both 2 and 3. @@ -53,3 +53,26 @@ The canary controller (`internal/rollout.Controller`) is a pure helper driven by **New operator flags**: `--prometheus-url` (required for Canary), `--gateway-name`, `--gateway-namespace`. **New status fields**: `canaryStepIndex`, `pauseStartedAt`, `promUnreachableSince` — all persisted so the state machine is re-entrant across operator restarts. + +### Phase 5 — MCP Registration & Discovery (§7) + +The MCP registry is an embedded HTTP service inside the operator manager process, served on `--registry-bind-address` (default `:9090`) and exposed to the cluster as the `agentrax-registry` ClusterIP Service in `agentrax-system`. + +**Registration flow**: When `spec.mcp.expose: true` and the agent's underlying Deployment reports `rolloutComplete` (`phase=Running`), the reconciler calls `registry.Registrar.Register()`, which performs an MCP-level `initialize` JSON-RPC 2.0 handshake by posting to the `/initialize` sub-path at `http://{name}.{namespace}.svc:{port}/initialize`. Discovered tools are merged with `spec.mcp.tools` and persisted to the in-memory registry map and `agentrax-registry` ConfigMap. On failure, `MCPHandshakeFailed` condition is set and `status.registered` remains `false`. + +**TTL & Heartbeat**: Every registry entry carries a `TTL` (default 90s) and `HeartbeatAt` timestamp. When an agent is `status.registered: true` and healthy in `Running` phase, the reconciler triggers `Registrar.Heartbeat()`. A background sweeper running every 30s removes expired entries if heartbeats lapse (e.g. following an ungraceful pod termination). After 3 consecutive heartbeat probe failures, the agent is automatically deregistered. + +**Deregistration & Finalizers**: Clean deregistration occurs when: (1) an `AgentDeployment` is deleted (via the `agentrax.io/mcp-deregister` finalizer before child services are garbage-collected), (2) `spec.mcp.expose` is toggled to `false`, or (3) 3 consecutive heartbeat probes fail. + +**State Recovery**: On operator startup/restart, existing registrations are reloaded from the `agentrax-registry` ConfigMap. + +**REST API Endpoints**: +- `GET /agents`: List all active, non-expired registered agents. +- `POST /agents`: Register / update an agent. This endpoint requires a successful MCP initialize handshake and currently requires no authentication. It bypasses the handshake verification when called directly. +- `GET /agents/{namespace}/{name}`: Get metadata and tools for a specific agent. +- `DELETE /agents/{namespace}/{name}`: Deregister an agent. This endpoint currently requires no authentication. +- *(Legacy aliases `POST /register` and `DELETE /deregister` are supported for backward compatibility)*. + +**New operator flags**: `--registry-bind-address` (default `:9090`). + + diff --git a/hack/setup-git-hooks.sh b/hack/setup-git-hooks.sh new file mode 100755 index 0000000..fc3bf7e --- /dev/null +++ b/hack/setup-git-hooks.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# 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. + +set -o errexit +set -o nounset +set -o pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +HOOKS_DIR="${REPO_ROOT}/.git/hooks" + +if [ ! -d "${HOOKS_DIR}" ]; then + echo "Error: .git/hooks directory not found. Are you in a git repository?" + exit 1 +fi + +echo "Setting up Agentrax Git hooks in ${HOOKS_DIR}..." + +# 1. Create pre-commit hook +cat << 'EOF' > "${HOOKS_DIR}/pre-commit" +#!/usr/bin/env bash +set -e + +echo "🔍 [pre-commit] Checking formatting, vet, and lint..." + +# Run go fmt +echo " -> go fmt ./..." +go fmt ./... + +# Run go vet +echo " -> go vet ./..." +go vet ./... + +# Run golangci-lint +echo " -> golangci-lint run..." +if [ -f "./bin/golangci-lint" ]; then + ./bin/golangci-lint run +else + make lint +fi + +echo "✅ [pre-commit] All pre-commit checks passed!" +EOF + +chmod +x "${HOOKS_DIR}/pre-commit" + +# 2. Create pre-push hook +cat << 'EOF' > "${HOOKS_DIR}/pre-push" +#!/usr/bin/env bash +set -e + +echo "🧪 [pre-push] Running full test suite and manifest verification..." + +# Verify manifests and code generation +echo " -> make manifests generate..." +make manifests generate + +# Run test suite +echo " -> make test..." +make test + +# Verify kustomize render +echo " -> Verifying config/default kustomization..." +if [ -f "./bin/kustomize" ]; then + ./bin/kustomize build config/default > /dev/null +fi + +echo "✅ [pre-push] All pre-push checks passed!" +EOF + +chmod +x "${HOOKS_DIR}/pre-push" + +echo "✨ Git hooks installed successfully! (pre-commit & pre-push)" diff --git a/internal/controller/agentdeployment_controller.go b/internal/controller/agentdeployment_controller.go index 3bf982c..1827406 100644 --- a/internal/controller/agentdeployment_controller.go +++ b/internal/controller/agentdeployment_controller.go @@ -19,6 +19,7 @@ package controller import ( "context" + "errors" "fmt" "sync" "time" @@ -43,6 +44,7 @@ import ( apimeta "k8s.io/apimachinery/pkg/api/meta" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/registry" "github.com/gitcommitankit/agentrax/internal/rollout" "github.com/gitcommitankit/agentrax/internal/scaling" ) @@ -83,6 +85,10 @@ type AgentDeploymentReconciler struct { // AgentDeployments that request it are treated as Recreate. CanaryController *rollout.Controller + // Registrar manages registration and discovery of this agent in the MCP registry. + // When nil, MCP registration features are disabled. + Registrar *registry.Registrar + // hasServiceMonitorCRD is set once during SetupWithManager and determines // whether ServiceMonitor reconciliation is attempted at all. hasServiceMonitorCRD bool @@ -249,6 +255,10 @@ func (r *AgentDeploymentReconciler) runDeletionCleanup(ctx context.Context, ad * if err := deregister(ctx, ad); err != nil { return fmt.Errorf("deregistering agent: %w", err) } + } else if r.Registrar != nil { + if err := r.Registrar.Deregister(ctx, ad); err != nil { + return fmt.Errorf("deregistering agent: %w", err) + } } logger.Info("deletion cleanup complete", "name", ad.Name, "namespace", ad.Namespace) return nil @@ -549,6 +559,12 @@ func (r *AgentDeploymentReconciler) updateStatus(ctx context.Context, ad *agentr } } + // Synchronize MCP registration status. + var mcpRequeue time.Duration + if r.Registrar != nil { + mcpRequeue = r.reconcileMCPRegistration(ctx, latest) + } + // Only write status when something actually changed to avoid spurious API // calls and watch events on every reconcile. if !equality.Semantic.DeepEqual(prevStatus, &latest.Status) { @@ -562,6 +578,11 @@ func (r *AgentDeploymentReconciler) updateStatus(ctx context.Context, ad *agentr return ctrl.Result{RequeueAfter: 5 * time.Second}, nil } + // If MCP registration needs periodic heartbeats, requeue accordingly. + if mcpRequeue > 0 { + return ctrl.Result{RequeueAfter: mcpRequeue}, nil + } + logger.Info("reconciled AgentDeployment", "phase", latest.Status.Phase, "readyReplicas", latest.Status.CurrentReplicas) return ctrl.Result{}, nil } @@ -591,6 +612,62 @@ func (r *AgentDeploymentReconciler) detectImagePullFailure(ctx context.Context, return false, "", nil } +// reconcileMCPRegistration synchronizes the MCP registry state with the AgentDeployment. +// It registers when running and expose is true, heartbeats when already registered, +// or deregisters when expose is false. +// Returns a requeue duration if periodic heartbeats are needed. +func (r *AgentDeploymentReconciler) reconcileMCPRegistration(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) time.Duration { + if r.Registrar == nil { + return 0 + } + + // If MCP exposure is disabled, ensure agent is deregistered. + if !ad.Spec.MCP.Expose { + if ad.Status.Registered { + if err := r.Registrar.Deregister(ctx, ad); err != nil { + SetCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed, metav1.ConditionTrue, "DeregisterFailed", err.Error()) + return 5 * time.Second + } + ad.Status.Registered = false + } + RemoveCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed) + return 0 + } + + // Only register or heartbeat if the deployment is in PhaseRunning. + if ad.Status.Phase != agentraxv1alpha1.PhaseRunning { + return 0 + } + + // Requeue interval shorter than TTL (90s) to ensure heartbeats occur before expiry. + requeueInterval := 60 * time.Second + + // If already registered, perform heartbeat probe. + if ad.Status.Registered { + if err := r.Registrar.Heartbeat(ctx, ad); err != nil { + // Only mark unregistered if Heartbeat confirms deregistration (after 3 consecutive failures). + if errors.Is(err, registry.ErrHeartbeatDeregistered) { + ad.Status.Registered = false + } + SetCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed, metav1.ConditionTrue, "HeartbeatFailed", err.Error()) + return requeueInterval + } + RemoveCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed) + return requeueInterval + } + + // Not registered yet — attempt initial registration handshake. + if err := r.Registrar.Register(ctx, ad); err != nil { + ad.Status.Registered = false + SetCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed, metav1.ConditionTrue, "HandshakeFailed", err.Error()) + return requeueInterval + } + + ad.Status.Registered = true + RemoveCondition(ad, agentraxv1alpha1.ConditionMCPHandshakeFailed) + return requeueInterval +} + // ── Desired-state builders ──────────────────────────────────────────────────── // agentLabels returns the canonical label set applied to all resources owned by ad. diff --git a/internal/controller/mcp_registration_test.go b/internal/controller/mcp_registration_test.go new file mode 100644 index 0000000..893ed40 --- /dev/null +++ b/internal/controller/mcp_registration_test.go @@ -0,0 +1,346 @@ +/* +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" + "errors" + "fmt" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +var _ = Describe("MCP Registration Integration", func() { + const ( + timeout = 10 * time.Second + interval = 250 * time.Millisecond + ) + + var ( + ns string + tq *agentraxv1alpha1.TenantQuota + ) + + BeforeEach(func() { + testMockMCP.SetError(nil) + testMockMCP.SetTools([]string{"search", "calculator"}) + + ns = fmt.Sprintf("tenant-mcp-%d", time.Now().UnixNano()) + namespaceObj := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: ns}} + Expect(k8sClient.Create(ctx, namespaceObj)).To(Succeed()) + + tq = &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{ + Name: "quota-mcp", + Namespace: ns, + }, + Spec: agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 10, + MaxTotalReplicas: 20, + MaxReplicasPerAgent: 5, + MaxGPUs: 8, + }, + } + Expect(k8sClient.Create(ctx, tq)).To(Succeed()) + }) + + AfterEach(func() { + testMockMCP.SetError(nil) + }) + + It("registers an AgentDeployment with mcp.expose: true upon reaching Running phase", func() { + adName := "mcp-agent-basic" + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: adName, + Namespace: ns, + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "nginx:latest", + Port: 8080, + TenantRef: tq.Name, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: 3, + Metric: "queueDepth", + Target: 10, + }, + MCP: agentraxv1alpha1.MCPConfig{ + Expose: true, + Tools: []string{"customTool"}, + }, + }, + } + + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Simulate Deployment pods becoming ready + depKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func() error { + dep := &appsv1.Deployment{} + if err := k8sClient.Get(ctx, depKey, dep); err != nil { + return err + } + dep.Status.ObservedGeneration = dep.Generation + dep.Status.Replicas = 1 + dep.Status.UpdatedReplicas = 1 + dep.Status.AvailableReplicas = 1 + dep.Status.ReadyReplicas = 1 + return k8sClient.Status().Update(ctx, dep) + }, timeout, interval).Should(Succeed()) + + // Verify AgentDeployment reaches Running phase and Registered=true + adKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + g.Expect(fetched.Status.Phase).To(Equal(agentraxv1alpha1.PhaseRunning)) + g.Expect(fetched.Status.Registered).To(BeTrue()) + }, timeout, interval).Should(Succeed()) + + // Verify entry in testRegistry + entry, ok := testRegistry.Get(ns, adName) + Expect(ok).To(BeTrue()) + Expect(entry.Endpoint).To(Equal(fmt.Sprintf("http://%s.%s.svc:8080", adName, ns))) + Expect(entry.Tools).To(ContainElements("search", "calculator", "customTool")) + }) + + It("sets MCPHandshakeFailed condition when MCP initialize fails", func() { + testMockMCP.SetError(errors.New("connection refused")) + + adName := "mcp-agent-fail" + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: adName, + Namespace: ns, + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "nginx:latest", + Port: 8080, + TenantRef: tq.Name, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: 3, + Metric: "queueDepth", + Target: 10, + }, + MCP: agentraxv1alpha1.MCPConfig{ + Expose: true, + }, + }, + } + + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Simulate Deployment pods becoming ready + depKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func() error { + dep := &appsv1.Deployment{} + if err := k8sClient.Get(ctx, depKey, dep); err != nil { + return err + } + dep.Status.ObservedGeneration = dep.Generation + dep.Status.Replicas = 1 + dep.Status.UpdatedReplicas = 1 + dep.Status.AvailableReplicas = 1 + dep.Status.ReadyReplicas = 1 + return k8sClient.Status().Update(ctx, dep) + }, timeout, interval).Should(Succeed()) + + // Verify MCPHandshakeFailed condition is set and Registered=false + adKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + g.Expect(fetched.Status.Phase).To(Equal(agentraxv1alpha1.PhaseRunning)) + g.Expect(fetched.Status.Registered).To(BeFalse()) + + cond := GetCondition(fetched, agentraxv1alpha1.ConditionMCPHandshakeFailed) + g.Expect(cond).NotTo(BeNil()) + g.Expect(cond.Status).To(Equal(metav1.ConditionTrue)) + g.Expect(cond.Reason).To(Equal("HandshakeFailed")) + }, timeout, interval).Should(Succeed()) + + // Verify NOT in registry + _, ok := testRegistry.Get(ns, adName) + Expect(ok).To(BeFalse()) + }) + + It("deregisters when spec.mcp.expose is toggled to false", func() { + adName := "mcp-agent-toggle" + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: adName, + Namespace: ns, + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "nginx:latest", + Port: 8080, + TenantRef: tq.Name, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: 3, + Metric: "queueDepth", + Target: 10, + }, + MCP: agentraxv1alpha1.MCPConfig{ + Expose: true, + }, + }, + } + + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Simulate Deployment ready + depKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func() error { + dep := &appsv1.Deployment{} + if err := k8sClient.Get(ctx, depKey, dep); err != nil { + return err + } + dep.Status.ObservedGeneration = dep.Generation + dep.Status.Replicas = 1 + dep.Status.UpdatedReplicas = 1 + dep.Status.AvailableReplicas = 1 + dep.Status.ReadyReplicas = 1 + return k8sClient.Status().Update(ctx, dep) + }, timeout, interval).Should(Succeed()) + + adKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + g.Expect(fetched.Status.Registered).To(BeTrue()) + }, timeout, interval).Should(Succeed()) + + // Toggle expose to false + Eventually(func() error { + fetched := &agentraxv1alpha1.AgentDeployment{} + if err := k8sClient.Get(ctx, adKey, fetched); err != nil { + return err + } + fetched.Spec.MCP.Expose = false + return k8sClient.Update(ctx, fetched) + }, timeout, interval).Should(Succeed()) + + // Verify Registered becomes false and removed from registry + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + g.Expect(fetched.Status.Registered).To(BeFalse()) + }, timeout, interval).Should(Succeed()) + + _, ok := testRegistry.Get(ns, adName) + Expect(ok).To(BeFalse()) + }) + + It("deregisters before finalizer is removed on deletion", func() { + adName := "mcp-agent-delete" + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: adName, + Namespace: ns, + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "nginx:latest", + Port: 8080, + TenantRef: tq.Name, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: 3, + Metric: "queueDepth", + Target: 10, + }, + MCP: agentraxv1alpha1.MCPConfig{ + Expose: true, + }, + }, + } + + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Simulate Deployment ready + depKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func() error { + dep := &appsv1.Deployment{} + if err := k8sClient.Get(ctx, depKey, dep); err != nil { + return err + } + dep.Status.ObservedGeneration = dep.Generation + dep.Status.Replicas = 1 + dep.Status.UpdatedReplicas = 1 + dep.Status.AvailableReplicas = 1 + dep.Status.ReadyReplicas = 1 + return k8sClient.Status().Update(ctx, dep) + }, timeout, interval).Should(Succeed()) + + adKey := types.NamespacedName{Name: adName, Namespace: ns} + Eventually(func(g Gomega) { + fetched := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + g.Expect(fetched.Status.Registered).To(BeTrue()) + }, timeout, interval).Should(Succeed()) + + // Install a Deregister hook that verifies Service still exists. + var serviceExistedDuringDeregister atomic.Bool + testReconciler.SetDeregister(func(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) error { + svc := &corev1.Service{} + svcKey := types.NamespacedName{Name: adName, Namespace: ns} + err := k8sClient.Get(ctx, svcKey, svc) + serviceExistedDuringDeregister.Store(err == nil) + return testRegistrar.Deregister(ctx, ad) + }) + defer testReconciler.SetDeregister(nil) + + // Delete the AgentDeployment + fetched := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, adKey, fetched)).To(Succeed()) + + // Verify child Service has controlling owner reference matching AgentDeployment UID + svc := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: adName, Namespace: ns}, svc)).To(Succeed()) + ownerRef := metav1.GetControllerOf(svc) + Expect(ownerRef).NotTo(BeNil()) + Expect(ownerRef.UID).To(Equal(fetched.UID)) + + Expect(k8sClient.Delete(ctx, fetched)).To(Succeed()) + + // Wait for object to be completely removed + Eventually(func() bool { + err := k8sClient.Get(ctx, adKey, &agentraxv1alpha1.AgentDeployment{}) + return apierrors.IsNotFound(err) + }, timeout, interval).Should(BeTrue()) + + // Verify deregistration occurred while Service was present + Expect(serviceExistedDuringDeregister.Load()).To(BeTrue(), "Service should exist during deregistration") + + // Entry must be gone from registry + _, ok := testRegistry.Get(ns, adName) + Expect(ok).To(BeFalse()) + }) +}) diff --git a/internal/controller/suite_test.go b/internal/controller/suite_test.go index adc0d32..9280829 100644 --- a/internal/controller/suite_test.go +++ b/internal/controller/suite_test.go @@ -21,6 +21,7 @@ import ( "fmt" "path/filepath" "runtime" + "sync" "testing" . "github.com/onsi/ginkgo/v2" @@ -30,6 +31,8 @@ import ( autoscalingv2 "k8s.io/api/autoscaling/v2" corev1 "k8s.io/api/core/v1" apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes/scheme" "k8s.io/client-go/rest" ctrl "sigs.k8s.io/controller-runtime" @@ -45,6 +48,7 @@ import ( agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" "github.com/gitcommitankit/agentrax/internal/quota" + "github.com/gitcommitankit/agentrax/internal/registry" agentraxwebhook "github.com/gitcommitankit/agentrax/internal/webhook" // +kubebuilder:scaffold:imports ) @@ -68,6 +72,41 @@ var testReconciler *AgentDeploymentReconciler // and the validating webhook in integration tests. var testEnforcer *quota.Enforcer +// testRegistry and testRegistrar are the shared MCP registry fixtures used in tests. +var testRegistry *registry.Registry +var testRegistrar *registry.Registrar +var testMockMCP *testMockMCPClient + +type testMockMCPClient struct { + mu sync.Mutex + tools []string + err error +} + +func (m *testMockMCPClient) Initialize(ctx context.Context, endpoint string) ([]string, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.err != nil { + return nil, m.err + } + if len(m.tools) == 0 { + return []string{"default_tool"}, nil + } + return m.tools, nil +} + +func (m *testMockMCPClient) SetError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.err = err +} + +func (m *testMockMCPClient) SetTools(tools []string) { + m.mu.Lock() + defer m.mu.Unlock() + m.tools = tools +} + // TestControllers is the Ginkgo test suite runner for controller integration tests. func TestControllers(t *testing.T) { RegisterFailHandler(Fail) @@ -129,6 +168,18 @@ var _ = BeforeSuite(func() { Expect(err).NotTo(HaveOccurred()) Expect(k8sClient).NotTo(BeNil()) + // Ensure system namespace exists for registry ConfigMap + systemNamespace := &corev1.Namespace{ + ObjectMeta: metav1.ObjectMeta{ + Name: "agentrax-system", + }, + } + if err = k8sClient.Create(ctx, systemNamespace); err != nil { + if !apierrors.IsAlreadyExists(err) { + Expect(err).NotTo(HaveOccurred(), "failed to create agentrax-system namespace") + } + } + // 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. @@ -146,10 +197,15 @@ var _ = BeforeSuite(func() { testEnforcer = quota.NewEnforcer(quota.DefaultGPUResourceName) + testRegistry = registry.NewRegistry(k8sClient, "agentrax-system", registry.DefaultTTL) + testMockMCP = &testMockMCPClient{} + testRegistrar = registry.NewRegistrar(testRegistry, testMockMCP) + testReconciler = &AgentDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), GPUResourceName: quota.DefaultGPUResourceName, + Registrar: testRegistrar, } Expect(testReconciler.SetupWithManager(mgr)).To(Succeed()) diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index 808c4ff..0455ab7 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -26,6 +26,8 @@ import ( apierrors "k8s.io/apimachinery/pkg/api/errors" apimeta "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/util/retry" + "sigs.k8s.io/controller-runtime/pkg/client" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" ) @@ -55,10 +57,14 @@ var _ = Describe("TenantQuota Controller", func() { for i := range adList.Items { ad := &adList.Items[i] // Remove finalizer so deletion is not blocked by the controller. - ad.Finalizers = nil - if err := k8sClient.Update(ctx, ad); err != nil && !apierrors.IsNotFound(err) { - Expect(err).NotTo(HaveOccurred(), "removing finalizer from AD %s", ad.Name) - } + _ = retry.RetryOnConflict(retry.DefaultRetry, func() error { + latest := &agentraxv1alpha1.AgentDeployment{} + if err := k8sClient.Get(ctx, namespacedName(ad.Name, tqNS), latest); err != nil { + return client.IgnoreNotFound(err) + } + latest.Finalizers = nil + return k8sClient.Update(ctx, latest) + }) if err := k8sClient.Delete(ctx, ad); err != nil && !apierrors.IsNotFound(err) { Expect(err).NotTo(HaveOccurred(), "deleting AD %s", ad.Name) } diff --git a/internal/controller/webhook_integration_test.go b/internal/controller/webhook_integration_test.go index 33187d2..fd49ad2 100644 --- a/internal/controller/webhook_integration_test.go +++ b/internal/controller/webhook_integration_test.go @@ -307,7 +307,6 @@ var _ = Describe("Admission Webhook (integration)", func() { start := make(chan struct{}) errs := make([]error, 2) for i, name := range []string{"ad-race-a", "ad-race-b"} { - i, name := i, name wg.Add(1) go func() { defer wg.Done() diff --git a/internal/metrics/prometheus_test.go b/internal/metrics/prometheus_test.go index aaec1a2..b53c121 100644 --- a/internal/metrics/prometheus_test.go +++ b/internal/metrics/prometheus_test.go @@ -93,7 +93,6 @@ func TestParseScalarFromQueryResponse(t *testing.T) { } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() got, err := parseScalarFromQueryResponse([]byte(tc.body)) @@ -151,7 +150,6 @@ func TestParseLastValueFromRangeResponse(t *testing.T) { } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() got, err := parseLastValueFromRangeResponse([]byte(tc.body)) diff --git a/internal/quota/enforcer_test.go b/internal/quota/enforcer_test.go index fa72fb3..aa7dedf 100644 --- a/internal/quota/enforcer_test.go +++ b/internal/quota/enforcer_test.go @@ -170,7 +170,6 @@ func TestCanAdmit_Create(t *testing.T) { } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -297,7 +296,6 @@ func TestCanAdmit_Update_GPUCeiling(t *testing.T) { } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -421,7 +419,6 @@ func TestCanAdmit_Update_CrossTenantMove(t *testing.T) { } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -505,7 +502,6 @@ func TestRelease_Concurrent(t *testing.T) { {"10 concurrent releases", 10}, } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -529,7 +525,6 @@ func TestRelease_Concurrent(t *testing.T) { start := make(chan struct{}) var wg sync.WaitGroup for _, k := range keys { - k := k wg.Add(1) go func() { defer wg.Done() @@ -581,7 +576,6 @@ func TestAdmitAndReserve_AtomicRaceProtection(t *testing.T) { // AdmitAndReserve, maximising the chance of a real concurrent execution. start := make(chan struct{}) for _, key := range []string{"ns/ad-A", "ns/ad-B"} { - key := key wg.Add(1) go func() { defer wg.Done() @@ -687,7 +681,6 @@ func TestIsOverQuota(t *testing.T) { {"zero GPU quota with 0 used GPUs ok", makeQuota(3, 0, 10, 3), makeUsage(2, 0, 8), false, ""}, } for _, tc := range tests { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() over, msg := e.IsOverQuota(tc.quota, tc.usage) diff --git a/internal/registry/mcp_client.go b/internal/registry/mcp_client.go new file mode 100644 index 0000000..7d1d934 --- /dev/null +++ b/internal/registry/mcp_client.go @@ -0,0 +1,189 @@ +/* +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 registry + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +const ( + // DefaultHandshakeTimeout is the HTTP timeout for MCP initialization requests. + DefaultHandshakeTimeout = 10 * time.Second + + // DefaultMCPProtocolVersion is the Model Context Protocol specification version advertised. + DefaultMCPProtocolVersion = "2024-11-05" +) + +// MCPClient defines the interface for communicating with an agent's MCP endpoint. +type MCPClient interface { + // Initialize performs the MCP protocol handshake against the given agent endpoint URL + // (e.g. "http://agent-svc.namespace.svc:8080") and returns the list of advertised tool names. + Initialize(ctx context.Context, endpoint string) ([]string, error) +} + +// httpMCPClient is the default HTTP-based implementation of MCPClient. +type httpMCPClient struct { + httpClient *http.Client +} + +// NewHTTPMCPClient creates a new HTTP-based MCPClient with default timeout settings. +func NewHTTPMCPClient() MCPClient { + return &httpMCPClient{ + httpClient: &http.Client{ + Timeout: DefaultHandshakeTimeout, + }, + } +} + +// mcpInitializeRequest represents the JSON-RPC 2.0 initialize request payload. +type mcpInitializeRequest struct { + JSONRPC string `json:"jsonrpc"` + ID int `json:"id"` + Method string `json:"method"` + Params mcpInitializeParams `json:"params"` +} + +// mcpInitializeParams encapsulates the parameters passed to the MCP initialize RPC. +type mcpInitializeParams struct { + ProtocolVersion string `json:"protocolVersion"` + ClientInfo mcpClientInfo `json:"clientInfo"` + Capabilities map[string]interface{} `json:"capabilities"` +} + +// mcpClientInfo identifies the client implementation and version to the MCP server. +type mcpClientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// mcpInitializeResponse represents the JSON-RPC 2.0 initialize response payload. +type mcpInitializeResponse struct { + JSONRPC string `json:"jsonrpc"` + ID int `json:"id"` + Result *mcpInitializeResult `json:"result,omitempty"` + Error *mcpError `json:"error,omitempty"` +} + +// mcpError represents a JSON-RPC 2.0 error object. +type mcpError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// mcpInitializeResult encapsulates the capability and tool declarations returned by an MCP server. +type mcpInitializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities mcpCapabilities `json:"capabilities"` + ServerInfo *mcpClientInfo `json:"serverInfo,omitempty"` + Tools []mcpToolDefinition `json:"tools,omitempty"` +} + +// mcpCapabilities declares server-supported features such as tools or resources. +type mcpCapabilities struct { + Tools *mcpToolsCapability `json:"tools,omitempty"` +} + +// mcpToolsCapability specifies tool availability and listing support. +type mcpToolsCapability struct { + Available []string `json:"available,omitempty"` + List bool `json:"list,omitempty"` +} + +// mcpToolDefinition describes an individual tool declared in the initialize result. +type mcpToolDefinition struct { + Name string `json:"name"` +} + +// Initialize performs an HTTP POST handshake to the agent's MCP endpoint and extracts tool names. +func (c *httpMCPClient) Initialize(ctx context.Context, endpoint string) ([]string, error) { + url := strings.TrimRight(endpoint, "/") + // If endpoint doesn't end with a specific path, target /initialize or root + if !strings.HasSuffix(url, "/initialize") { + url += "/initialize" + } + + reqBody := mcpInitializeRequest{ + JSONRPC: "2.0", + ID: 1, + Method: "initialize", + Params: mcpInitializeParams{ + ProtocolVersion: DefaultMCPProtocolVersion, + ClientInfo: mcpClientInfo{ + Name: "agentrax", + Version: "1.0", + }, + Capabilities: map[string]interface{}{}, + }, + } + + bodyBytes, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("marshaling MCP initialize request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyBytes)) + if err != nil { + return nil, fmt.Errorf("creating MCP initialize request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("sending MCP initialize to %s: %w", url, err) + } + defer resp.Body.Close() //nolint:errcheck + + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + return nil, fmt.Errorf("MCP initialize returned HTTP %d: %s", resp.StatusCode, string(respBody)) + } + + var initResp mcpInitializeResponse + if err := json.NewDecoder(resp.Body).Decode(&initResp); err != nil { + return nil, fmt.Errorf("decoding MCP initialize response: %w", err) + } + + if initResp.Error != nil { + return nil, fmt.Errorf("MCP initialize RPC error (code %d): %s", initResp.Error.Code, initResp.Error.Message) + } + + if initResp.Result == nil { + return nil, fmt.Errorf("MCP initialize response missing result object") + } + + // Extract tools: check capabilities.tools.available first, then result.tools list + var tools []string + if initResp.Result.Capabilities.Tools != nil && len(initResp.Result.Capabilities.Tools.Available) > 0 { + tools = append(tools, initResp.Result.Capabilities.Tools.Available...) + } else if len(initResp.Result.Tools) > 0 { + for _, t := range initResp.Result.Tools { + if t.Name != "" { + tools = append(tools, t.Name) + } + } + } + + return tools, nil +} diff --git a/internal/registry/mcp_registrar.go b/internal/registry/mcp_registrar.go new file mode 100644 index 0000000..b7e3635 --- /dev/null +++ b/internal/registry/mcp_registrar.go @@ -0,0 +1,174 @@ +/* +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 registry + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +const ( + // MaxConsecutiveHeartbeatFailures is the number of failed heartbeat probes before an agent is automatically deregistered. + MaxConsecutiveHeartbeatFailures = 3 +) + +// ErrHeartbeatDeregistered is returned when an agent is deregistered due to reaching MaxConsecutiveHeartbeatFailures. +var ErrHeartbeatDeregistered = errors.New("deregistered after consecutive heartbeat failures") + +// Registrar coordinates MCP initialize handshakes and updates the discovery Registry. +type Registrar struct { + // Registry is the underlying registry store and HTTP server handler. + Registry *Registry + + // MCPClient is the protocol client used for MCP initialize handshakes. + MCPClient MCPClient + + failuresMu sync.Mutex + failures map[string]int +} + +// NewRegistrar creates a new Registrar instance with the provided registry and MCP client. +func NewRegistrar(reg *Registry, client MCPClient) *Registrar { + if client == nil { + client = NewHTTPMCPClient() + } + return &Registrar{ + Registry: reg, + MCPClient: client, + failures: make(map[string]int), + } +} + +// EndpointForAgent computes the standard Kubernetes Service URL for an AgentDeployment. +func EndpointForAgent(ad *agentraxv1alpha1.AgentDeployment) string { + port := ad.Spec.Port + if port == 0 { + port = 8080 + } + return fmt.Sprintf("http://%s.%s.svc:%d", ad.Name, ad.Namespace, port) +} + +// Register performs the MCP initialize handshake against the agent's endpoint and records the agent in the registry. +func (r *Registrar) Register(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) error { + if r.Registry == nil { + return fmt.Errorf("registry not initialized") + } + + endpoint := EndpointForAgent(ad) + + // Perform initialize handshake to discover advertised tools and verify MCP readiness. + tools, err := r.MCPClient.Initialize(ctx, endpoint) + if err != nil { + return fmt.Errorf("MCP handshake with %s failed: %w", endpoint, err) + } + + // If spec declared tools explicitly, combine them with discovered tools. + if len(ad.Spec.MCP.Tools) > 0 { + seen := make(map[string]bool) + for _, t := range tools { + seen[t] = true + } + for _, t := range ad.Spec.MCP.Tools { + if !seen[t] { + tools = append(tools, t) + seen[t] = true + } + } + } + + entry := Entry{ + Namespace: ad.Namespace, + Name: ad.Name, + Endpoint: endpoint, + Tools: tools, + } + + if err := r.Registry.Register(ctx, entry); err != nil { + return fmt.Errorf("registering agent in store: %w", err) + } + + r.resetFailures(ad.Namespace, ad.Name) + return nil +} + +// Deregister removes the agent from the registry store. +func (r *Registrar) Deregister(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) error { + if r.Registry == nil { + return nil + } + + r.resetFailures(ad.Namespace, ad.Name) + return r.Registry.Deregister(ctx, ad.Namespace, ad.Name) +} + +// Heartbeat performs an initialize probe against the agent and refreshes the registry TTL if healthy. +// After 3 consecutive probe failures, the agent is automatically deregistered. +func (r *Registrar) Heartbeat(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) error { + if r.Registry == nil { + return nil + } + + endpoint := EndpointForAgent(ad) + _, err := r.MCPClient.Initialize(ctx, endpoint) + if err != nil { + failCount := r.incrementFailures(ad.Namespace, ad.Name) + if failCount >= MaxConsecutiveHeartbeatFailures { + deregErr := r.Registry.Deregister(ctx, ad.Namespace, ad.Name) + if deregErr != nil { + return fmt.Errorf("failed to deregister after %d consecutive heartbeat failures: %w (latest probe error: %v)", + failCount, deregErr, err) + } + return fmt.Errorf("%w: deregistered after %d consecutive heartbeat failures (latest error: %w)", + ErrHeartbeatDeregistered, failCount, err) + } + return fmt.Errorf("heartbeat probe failed (%d/%d): %w", failCount, MaxConsecutiveHeartbeatFailures, err) + } + + r.resetFailures(ad.Namespace, ad.Name) + hbErr := r.Registry.Heartbeat(ctx, ad.Namespace, ad.Name) + if hbErr != nil && strings.Contains(hbErr.Error(), "not found in registry") { + // Agent not found - attempt re-registration + return r.Register(ctx, ad) + } + return hbErr +} + +// incrementFailures records a probe failure and returns the new consecutive failure count. +func (r *Registrar) incrementFailures(namespace, name string) int { + r.failuresMu.Lock() + defer r.failuresMu.Unlock() + if r.failures == nil { + r.failures = make(map[string]int) + } + k := namespace + "/" + name + r.failures[k]++ + return r.failures[k] +} + +// resetFailures clears the consecutive failure counter for the specified agent. +func (r *Registrar) resetFailures(namespace, name string) { + r.failuresMu.Lock() + defer r.failuresMu.Unlock() + if r.failures != nil { + delete(r.failures, namespace+"/"+name) + } +} diff --git a/internal/registry/registry.go b/internal/registry/registry.go new file mode 100644 index 0000000..7e4888a --- /dev/null +++ b/internal/registry/registry.go @@ -0,0 +1,472 @@ +/* +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 registry provides an in-process MCP (Model Context Protocol) service registry +// backed by a Kubernetes ConfigMap with automated TTL expiration and health tracking. +package registry + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "sync" + "time" + + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/util/retry" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/log" +) + +const ( + // DefaultTTL is the default time-to-live for a registry entry without a heartbeat. + DefaultTTL = 90 * time.Second + + // DefaultSweepInterval is the frequency at which the background TTL sweeper runs. + DefaultSweepInterval = 30 * time.Second + + // DefaultConfigMapName is the name of the ConfigMap storing registry state. + DefaultConfigMapName = "agentrax-registry" + + // configMapKey is the data key under which the serialized registry entries are stored. + configMapKey = "state" +) + +// Entry represents a registered agent in the MCP discovery registry. +type Entry struct { + // Namespace is the Kubernetes namespace of the agent. + Namespace string `json:"namespace"` + + // Name is the name of the AgentDeployment. + Name string `json:"name"` + + // Endpoint is the full HTTP address for the agent's MCP service (e.g. "http://my-agent.tenant-a.svc:8080"). + Endpoint string `json:"endpoint"` + + // Tools is the list of tools advertised by this agent via MCP initialize. + Tools []string `json:"tools,omitempty"` + + // RegisteredAt is the timestamp when the agent was first registered. + RegisteredAt time.Time `json:"registeredAt"` + + // HeartbeatAt is the timestamp of the latest successful heartbeat or registration update. + HeartbeatAt time.Time `json:"heartbeatAt"` + + // TTL is the time-to-live duration for this entry. + TTL time.Duration `json:"ttl"` +} + +// Registry manages in-memory agent registration entries, exposes HTTP discovery endpoints, +// persists state to a ConfigMap, and sweeps expired entries. +type Registry struct { + mu sync.RWMutex + persistMu sync.Mutex + entries map[string]*Entry + client client.Client + namespace string + configMapName string + defaultTTL time.Duration + sweepInterval time.Duration +} + +// NewRegistry creates a new in-memory MCP Registry. +// If k8sClient is non-nil, registry state is persisted to and recovered from a ConfigMap. +func NewRegistry(k8sClient client.Client, namespace string, defaultTTL time.Duration) *Registry { + if defaultTTL <= 0 { + defaultTTL = DefaultTTL + } + if namespace == "" { + namespace = "agentrax-system" + } + return &Registry{ + entries: make(map[string]*Entry), + client: k8sClient, + namespace: namespace, + configMapName: DefaultConfigMapName, + defaultTTL: defaultTTL, + sweepInterval: DefaultSweepInterval, + } +} + +// SetSweepInterval overrides the default background sweep frequency (useful for testing). +func (r *Registry) SetSweepInterval(interval time.Duration) { + r.mu.Lock() + defer r.mu.Unlock() + r.sweepInterval = interval +} + +// Start launches the background TTL sweeper and loads any existing state from ConfigMap. +func (r *Registry) Start(ctx context.Context) { + logger := log.FromContext(ctx).WithName("mcp-registry") + if r.client != nil { + if err := r.loadFromConfigMap(ctx); err != nil { + logger.Error(err, "failed to load initial registry state from ConfigMap") + } + } + + go r.runSweeper(ctx) +} + +// runSweeper periodically deletes entries whose heartbeats have exceeded their TTL. +func (r *Registry) runSweeper(ctx context.Context) { + logger := log.FromContext(ctx).WithName("mcp-registry-sweeper") + + r.mu.RLock() + interval := r.sweepInterval + r.mu.RUnlock() + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + removed := r.sweepExpired() + if removed > 0 { + logger.Info("swept expired MCP registry entries", "removedCount", removed) + if r.client != nil { + if err := r.persistToConfigMap(ctx); err != nil { + logger.Error(err, "failed to persist registry state after sweep") + } + } + } + } + } +} + +// sweepExpired cleans up entries that have exceeded their TTL and returns the count of removed entries. +func (r *Registry) sweepExpired() int { + r.mu.Lock() + defer r.mu.Unlock() + + now := time.Now() + removed := 0 + for key, entry := range r.entries { + ttl := entry.TTL + if ttl <= 0 { + ttl = r.defaultTTL + } + if now.Sub(entry.HeartbeatAt) > ttl { + delete(r.entries, key) + removed++ + } + } + return removed +} + +// key returns the map key for a namespace/name pair. +func key(namespace, name string) string { + return namespace + "/" + name +} + +// Register registers or updates an agent in the registry. +func (r *Registry) Register(ctx context.Context, e Entry) error { + if e.Namespace == "" || e.Name == "" || e.Endpoint == "" { + return errors.New("namespace, name, and endpoint are required for registration") + } + + now := time.Now() + k := key(e.Namespace, e.Name) + + r.mu.Lock() + existing, found := r.entries[k] + if !found { + e.RegisteredAt = now + } else { + e.RegisteredAt = existing.RegisteredAt + } + e.HeartbeatAt = now + if e.TTL <= 0 { + e.TTL = r.defaultTTL + } + r.entries[k] = &e + r.mu.Unlock() + + if r.client != nil { + return r.persistToConfigMap(ctx) + } + return nil +} + +// Deregister removes an agent from the registry by namespace and name. +func (r *Registry) Deregister(ctx context.Context, namespace, name string) error { + k := key(namespace, name) + + r.mu.Lock() + _, found := r.entries[k] + if !found { + r.mu.Unlock() + return nil + } + delete(r.entries, k) + r.mu.Unlock() + + if r.client != nil { + return r.persistToConfigMap(ctx) + } + return nil +} + +// Heartbeat refreshes the heartbeat timestamp of a registered agent. +func (r *Registry) Heartbeat(ctx context.Context, namespace, name string) error { + k := key(namespace, name) + + r.mu.Lock() + entry, found := r.entries[k] + if !found { + r.mu.Unlock() + return fmt.Errorf("agent %s not found in registry", k) + } + entry.HeartbeatAt = time.Now() + r.mu.Unlock() + + // Heartbeats only update timestamps in-memory; persisting on every heartbeat + // would generate unnecessary ConfigMap write load. + return nil +} + +// Get looks up an agent entry by namespace and name. Returns false if not found or expired. +func (r *Registry) Get(namespace, name string) (*Entry, bool) { + k := key(namespace, name) + + r.mu.RLock() + defer r.mu.RUnlock() + + entry, found := r.entries[k] + if !found { + return nil, false + } + + ttl := entry.TTL + if ttl <= 0 { + ttl = r.defaultTTL + } + if time.Since(entry.HeartbeatAt) > ttl { + return nil, false + } + + entryCopy := *entry + return &entryCopy, true +} + +// List returns all non-expired agent entries in the registry. +func (r *Registry) List() []*Entry { + r.mu.RLock() + defer r.mu.RUnlock() + + now := time.Now() + result := make([]*Entry, 0, len(r.entries)) + for _, entry := range r.entries { + ttl := entry.TTL + if ttl <= 0 { + ttl = r.defaultTTL + } + if now.Sub(entry.HeartbeatAt) <= ttl { + entryCopy := *entry + result = append(result, &entryCopy) + } + } + return result +} + +// ── ConfigMap Persistence ───────────────────────────────────────────────────── + +// persistToConfigMap writes the current entries to the agentrax-registry ConfigMap. +func (r *Registry) persistToConfigMap(ctx context.Context) error { + r.persistMu.Lock() + defer r.persistMu.Unlock() + + return retry.RetryOnConflict(retry.DefaultRetry, func() error { + r.mu.RLock() + data, err := json.Marshal(r.entries) + r.mu.RUnlock() + if err != nil { + return fmt.Errorf("marshaling registry state: %w", err) + } + + cm := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: r.configMapName, + Namespace: r.namespace, + }, + } + + _, err = controllerutil.CreateOrUpdate(ctx, r.client, cm, func() error { + if cm.Data == nil { + cm.Data = make(map[string]string) + } + cm.Data[configMapKey] = string(data) + return nil + }) + if err != nil { + return fmt.Errorf("persisting registry ConfigMap: %w", err) + } + return nil + }) +} + +// loadFromConfigMap reads and restores entries from the agentrax-registry ConfigMap. +func (r *Registry) loadFromConfigMap(ctx context.Context) error { + cm := &corev1.ConfigMap{} + err := r.client.Get(ctx, types.NamespacedName{Name: r.configMapName, Namespace: r.namespace}, cm) + if err != nil { + if apierrors.IsNotFound(err) { + return nil + } + return fmt.Errorf("fetching registry ConfigMap: %w", err) + } + + raw, ok := cm.Data[configMapKey] + if !ok || raw == "" { + return nil + } + + var loaded map[string]*Entry + if err := json.Unmarshal([]byte(raw), &loaded); err != nil { + return fmt.Errorf("unmarshaling registry state: %w", err) + } + + r.mu.Lock() + defer r.mu.Unlock() + for k, v := range loaded { + r.entries[k] = v + } + return nil +} + +// ── HTTP Handler ────────────────────────────────────────────────────────────── + +// Handler returns an http.Handler implementing the registry REST API. +// It exposes standard RESTful endpoints under /agents as well as legacy aliases (/register, /deregister). +func (r *Registry) Handler() http.Handler { + mux := http.NewServeMux() + + // Consolidated RESTful /agents collection endpoints + mux.HandleFunc("GET /agents", r.handleListAgents) + mux.HandleFunc("POST /agents", r.handleRegister) + mux.HandleFunc("GET /agents/{namespace}/{name}", r.handleGetAgent) + mux.HandleFunc("DELETE /agents/{namespace}/{name}", r.handleDeregisterAgentPath) + + // Backward-compatible RPC action aliases + mux.HandleFunc("POST /register", r.handleRegister) + mux.HandleFunc("DELETE /deregister", r.handleDeregister) + + return mux +} + +// handleRegister handles HTTP POST requests to register or update an agent entry. +func (r *Registry) handleRegister(w http.ResponseWriter, req *http.Request) { + var entry Entry + if err := json.NewDecoder(req.Body).Decode(&entry); err != nil { + http.Error(w, fmt.Sprintf("invalid request body: %v", err), http.StatusBadRequest) + return + } + + if err := r.Register(req.Context(), entry); err != nil { + http.Error(w, fmt.Sprintf("registration failed: %v", err), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]string{"status": "registered"}) +} + +// handleDeregisterAgentPath handles HTTP DELETE requests targeting /agents/{namespace}/{name}. +func (r *Registry) handleDeregisterAgentPath(w http.ResponseWriter, req *http.Request) { + namespace := req.PathValue("namespace") + name := req.PathValue("name") + + if namespace == "" || name == "" { + http.Error(w, "namespace and name are required in URL path", http.StatusBadRequest) + return + } + + if err := r.Deregister(req.Context(), namespace, name); err != nil { + http.Error(w, fmt.Sprintf("deregistration failed: %v", err), http.StatusInternalServerError) + return + } + + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]string{"status": "deregistered"}) +} + +// handleDeregister handles legacy HTTP DELETE requests targeting /deregister. +func (r *Registry) handleDeregister(w http.ResponseWriter, req *http.Request) { + namespace := req.URL.Query().Get("namespace") + name := req.URL.Query().Get("name") + + if namespace == "" || name == "" { + // Also support JSON body if query params are missing + var body struct { + Namespace string `json:"namespace"` + Name string `json:"name"` + } + if err := json.NewDecoder(req.Body).Decode(&body); err == nil { + namespace = body.Namespace + name = body.Name + } + } + + if namespace == "" || name == "" { + http.Error(w, "namespace and name are required", http.StatusBadRequest) + return + } + + if err := r.Deregister(req.Context(), namespace, name); err != nil { + http.Error(w, fmt.Sprintf("deregistration failed: %v", err), http.StatusInternalServerError) + return + } + + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]string{"status": "deregistered"}) +} + +// handleListAgents handles HTTP GET requests to list all active, non-expired registered agents. +func (r *Registry) handleListAgents(w http.ResponseWriter, req *http.Request) { + agents := r.List() + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(agents) +} + +// handleGetAgent handles HTTP GET requests to fetch details of a specific agent. +func (r *Registry) handleGetAgent(w http.ResponseWriter, req *http.Request) { + namespace := req.PathValue("namespace") + name := req.PathValue("name") + + if namespace == "" || name == "" { + http.Error(w, "namespace and name are required", http.StatusBadRequest) + return + } + + entry, found := r.Get(namespace, name) + if !found { + http.Error(w, fmt.Sprintf("agent %s/%s not found", namespace, name), http.StatusNotFound) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(entry) +} diff --git a/internal/registry/registry_test.go b/internal/registry/registry_test.go new file mode 100644 index 0000000..26d80ba --- /dev/null +++ b/internal/registry/registry_test.go @@ -0,0 +1,462 @@ +/* +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 registry + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + corev1 "k8s.io/api/core/v1" + 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" +) + +func TestRegistry_CRUD(t *testing.T) { + ctx := context.Background() + reg := NewRegistry(nil, "agentrax-system", 10*time.Second) + + // 1. Initial lookup should be empty + if _, ok := reg.Get("tenant-a", "agent-1"); ok { + t.Fatalf("expected agent-1 to not be found") + } + + // 2. Register agent + err := reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-1", + Endpoint: "http://agent-1.tenant-a.svc:8080", + Tools: []string{"search", "calculator"}, + }) + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + // 3. Get agent + entry, ok := reg.Get("tenant-a", "agent-1") + if !ok || entry == nil { + t.Fatalf("expected agent-1 to be found") + } + if entry.Endpoint != "http://agent-1.tenant-a.svc:8080" { + t.Errorf("got endpoint %q, want http://agent-1.tenant-a.svc:8080", entry.Endpoint) + } + if len(entry.Tools) != 2 { + t.Errorf("got %d tools, want 2", len(entry.Tools)) + } + + // 4. List agents + list := reg.List() + if len(list) != 1 { + t.Fatalf("expected 1 agent in list, got %d", len(list)) + } + + // 5. Deregister agent + err = reg.Deregister(ctx, "tenant-a", "agent-1") + if err != nil { + t.Fatalf("Deregister failed: %v", err) + } + + if _, ok := reg.Get("tenant-a", "agent-1"); ok { + t.Fatalf("expected agent-1 to be removed") + } + if len(reg.List()) != 0 { + t.Fatalf("expected 0 agents in list after deregister") + } +} + +func TestRegistry_RegisterIsIdempotent(t *testing.T) { + ctx := context.Background() + reg := NewRegistry(nil, "agentrax-system", 10*time.Second) + + // 1. First registration + err := reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-idem", + Endpoint: "http://agent-idem.tenant-a.svc:8080", + Tools: []string{"search"}, + }) + if err != nil { + t.Fatalf("first Register failed: %v", err) + } + + entry1, ok := reg.Get("tenant-a", "agent-idem") + if !ok || entry1 == nil { + t.Fatalf("expected agent-idem to be found after first registration") + } + firstRegisteredAt := entry1.RegisteredAt + firstHeartbeatAt := entry1.HeartbeatAt + + // Wait a moment to ensure timestamps would differ + time.Sleep(10 * time.Millisecond) + + // 2. Second registration (idempotent) + err = reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-idem", + Endpoint: "http://agent-idem.tenant-a.svc:8080", + Tools: []string{"search", "calculator"}, + }) + if err != nil { + t.Fatalf("second Register failed: %v", err) + } + + // 3. Verify only one entry exists + list := reg.List() + if len(list) != 1 { + t.Fatalf("expected 1 agent in list after idempotent register, got %d", len(list)) + } + + // 4. Verify RegisteredAt is preserved, HeartbeatAt is refreshed + entry2, ok := reg.Get("tenant-a", "agent-idem") + if !ok || entry2 == nil { + t.Fatalf("expected agent-idem to be found after second registration") + } + if !entry2.RegisteredAt.Equal(firstRegisteredAt) { + t.Errorf("RegisteredAt changed: was %v, now %v", firstRegisteredAt, entry2.RegisteredAt) + } + if !entry2.HeartbeatAt.After(firstHeartbeatAt) { + t.Errorf("HeartbeatAt not refreshed: was %v, now %v", firstHeartbeatAt, entry2.HeartbeatAt) + } + + // 5. Verify tools updated + if len(entry2.Tools) != 2 { + t.Errorf("expected 2 tools after second registration, got %d", len(entry2.Tools)) + } +} + +func TestRegistry_TTLSweep(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + reg := NewRegistry(nil, "agentrax-system", 100*time.Millisecond) + reg.SetSweepInterval(50 * time.Millisecond) + + // Register entry with 100ms TTL + err := reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-short", + Endpoint: "http://agent-short.tenant-a.svc:8080", + TTL: 80 * time.Millisecond, + }) + if err != nil { + t.Fatalf("register failed: %v", err) + } + + // Register second entry with long TTL + err = reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-long", + Endpoint: "http://agent-long.tenant-a.svc:8080", + TTL: 5 * time.Second, + }) + if err != nil { + t.Fatalf("register failed: %v", err) + } + + reg.Start(ctx) + + // Wait 150ms for short entry to expire and sweep + time.Sleep(200 * time.Millisecond) + + if _, ok := reg.Get("tenant-a", "agent-short"); ok { + t.Errorf("expected agent-short to have expired and been swept") + } + if _, ok := reg.Get("tenant-a", "agent-long"); !ok { + t.Errorf("expected agent-long to remain registered") + } +} + +func TestRegistry_Heartbeat(t *testing.T) { + ctx := context.Background() + reg := NewRegistry(nil, "agentrax-system", 150*time.Millisecond) + + err := reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-hb", + Endpoint: "http://agent-hb.tenant-a.svc:8080", + }) + if err != nil { + t.Fatalf("register failed: %v", err) + } + + // Sleep 100ms + time.Sleep(100 * time.Millisecond) + + // Send heartbeat + err = reg.Heartbeat(ctx, "tenant-a", "agent-hb") + if err != nil { + t.Fatalf("heartbeat failed: %v", err) + } + + // Sleep another 100ms (total 200ms elapsed since register, but only 100ms since heartbeat) + time.Sleep(100 * time.Millisecond) + + if _, ok := reg.Get("tenant-a", "agent-hb"); !ok { + t.Errorf("expected agent-hb to still be alive due to heartbeat") + } +} + +func TestRegistry_ConfigMapPersistence(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + scheme := runtime.NewScheme() + _ = corev1.AddToScheme(scheme) + + fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build() + reg := NewRegistry(fakeClient, "agentrax-system", 60*time.Second) + + // 1. Register agent + err := reg.Register(ctx, Entry{ + Namespace: "tenant-a", + Name: "agent-cm", + Endpoint: "http://agent-cm.tenant-a.svc:8080", + Tools: []string{"search"}, + }) + if err != nil { + t.Fatalf("register failed: %v", err) + } + + // 2. Create new registry instance to test recovery from ConfigMap + recoveredReg := NewRegistry(fakeClient, "agentrax-system", 60*time.Second) + recoveredReg.Start(ctx) + + entry, ok := recoveredReg.Get("tenant-a", "agent-cm") + if !ok || entry == nil { + t.Fatalf("expected agent-cm to be recovered from ConfigMap") + } + if entry.Endpoint != "http://agent-cm.tenant-a.svc:8080" { + t.Errorf("got endpoint %q, want http://agent-cm.tenant-a.svc:8080", entry.Endpoint) + } +} + +func TestRegistry_HTTPHandler(t *testing.T) { + reg := NewRegistry(nil, "agentrax-system", 60*time.Second) + handler := reg.Handler() + + // 1. POST /agents (RESTful create/register) + regPayload := `{"namespace":"tenant-1","name":"agent-http","endpoint":"http://agent-http:8080","tools":["translate"]}` + req := httptest.NewRequest(http.MethodPost, "/agents", bytes.NewBufferString(regPayload)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("POST /agents returned %d: %s", w.Code, w.Body.String()) + } + + // 2. GET /agents (List) + req = httptest.NewRequest(http.MethodGet, "/agents", nil) + w = httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("GET /agents returned %d", w.Code) + } + var list []*Entry + if err := json.NewDecoder(w.Body).Decode(&list); err != nil { + t.Fatalf("decoding GET /agents response: %v", err) + } + if len(list) != 1 || list[0].Name != "agent-http" { + t.Fatalf("unexpected GET /agents response: %+v", list) + } + + // 3. GET /agents/{namespace}/{name} (Get) + req = httptest.NewRequest(http.MethodGet, "/agents/tenant-1/agent-http", nil) + w = httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("GET /agents/tenant-1/agent-http returned %d", w.Code) + } + + // 4. DELETE /agents/{namespace}/{name} (RESTful delete) + req = httptest.NewRequest(http.MethodDelete, "/agents/tenant-1/agent-http", nil) + w = httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("DELETE /agents/tenant-1/agent-http returned %d: %s", w.Code, w.Body.String()) + } + + // Verify not found after delete + if _, ok := reg.Get("tenant-1", "agent-http"); ok { + t.Fatalf("expected agent-http to be deregistered") + } + + // 5. Test legacy aliases (POST /register and DELETE /deregister) + req = httptest.NewRequest(http.MethodPost, "/register", bytes.NewBufferString(regPayload)) + req.Header.Set("Content-Type", "application/json") + w = httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("POST /register legacy alias returned %d", w.Code) + } + + req = httptest.NewRequest(http.MethodDelete, "/deregister?namespace=tenant-1&name=agent-http", nil) + w = httptest.NewRecorder() + handler.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("DELETE /deregister legacy alias returned %d", w.Code) + } +} + +func TestHTTPMCPClient_Initialize(t *testing.T) { + // Mock MCP Server + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/initialize" { + http.NotFound(w, r) + return + } + var req mcpInitializeRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, "bad request", http.StatusBadRequest) + return + } + + resp := mcpInitializeResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: &mcpInitializeResult{ + ProtocolVersion: "2024-11-05", + Capabilities: mcpCapabilities{ + Tools: &mcpToolsCapability{ + Available: []string{"web_search", "code_eval"}, + }, + }, + ServerInfo: &mcpClientInfo{ + Name: "mock-agent", + Version: "0.1.0", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + defer ts.Close() + + client := NewHTTPMCPClient() + tools, err := client.Initialize(context.Background(), ts.URL) + if err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + if len(tools) != 2 || tools[0] != "web_search" || tools[1] != "code_eval" { + t.Errorf("unexpected tools extracted: %+v", tools) + } +} + +// mockMCPClient is a mock implementation of MCPClient for testing. +type mockMCPClient struct { + tools []string + err error +} + +func (m *mockMCPClient) Initialize(ctx context.Context, endpoint string) ([]string, error) { + if m.err != nil { + return nil, m.err + } + return m.tools, nil +} + +func TestRegistrar_RegisterAndHeartbeat(t *testing.T) { + ctx := context.Background() + reg := NewRegistry(nil, "agentrax-system", 10*time.Second) + mockClient := &mockMCPClient{ + tools: []string{"toolA", "toolB"}, + } + + registrar := NewRegistrar(reg, mockClient) + + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-agent", + Namespace: "tenant-x", + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Port: 9000, + MCP: agentraxv1alpha1.MCPConfig{ + Expose: true, + Tools: []string{"customTool"}, + }, + }, + } + + // 1. Register + err := registrar.Register(ctx, ad) + if err != nil { + t.Fatalf("Registrar.Register failed: %v", err) + } + + entry, ok := reg.Get("tenant-x", "test-agent") + if !ok || entry == nil { + t.Fatalf("expected test-agent in registry") + } + if entry.Endpoint != "http://test-agent.tenant-x.svc:9000" { + t.Errorf("unexpected endpoint: %s", entry.Endpoint) + } + if len(entry.Tools) != 3 { // toolA, toolB, customTool + t.Errorf("got %d tools, want 3: %+v", len(entry.Tools), entry.Tools) + } + + // 2. Successful Heartbeat + err = registrar.Heartbeat(ctx, ad) + if err != nil { + t.Fatalf("Heartbeat failed: %v", err) + } + + // 3. Failing Heartbeat -> 3 strikes deregisters + mockClient.err = errors.New("connection refused") + // Strike 1 + err = registrar.Heartbeat(ctx, ad) + if err == nil { + t.Errorf("expected error on strike 1") + } + if errors.Is(err, ErrHeartbeatDeregistered) { + t.Errorf("strike 1 should not return ErrHeartbeatDeregistered") + } + if _, ok := reg.Get("tenant-x", "test-agent"); !ok { + t.Errorf("entry should still exist after 1 strike") + } + + // Strike 2 + err = registrar.Heartbeat(ctx, ad) + if err == nil { + t.Errorf("expected error on strike 2") + } + if errors.Is(err, ErrHeartbeatDeregistered) { + t.Errorf("strike 2 should not return ErrHeartbeatDeregistered") + } + if _, ok := reg.Get("tenant-x", "test-agent"); !ok { + t.Errorf("entry should still exist after 2 strikes") + } + + // Strike 3 -> deregistration + err = registrar.Heartbeat(ctx, ad) + if err == nil { + t.Errorf("expected error on strike 3") + } + if !errors.Is(err, ErrHeartbeatDeregistered) { + t.Errorf("strike 3 should return ErrHeartbeatDeregistered, got %v", err) + } + if _, ok := reg.Get("tenant-x", "test-agent"); ok { + t.Errorf("entry should be removed after 3 consecutive failures") + } +} diff --git a/internal/rollout/canary.go b/internal/rollout/canary.go index 698cfd8..5ae0e9f 100644 --- a/internal/rollout/canary.go +++ b/internal/rollout/canary.go @@ -40,6 +40,7 @@ import ( agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" "github.com/gitcommitankit/agentrax/internal/metrics" + "github.com/gitcommitankit/agentrax/internal/registry" "github.com/gitcommitankit/agentrax/internal/scaling" ) @@ -70,6 +71,9 @@ type Controller struct { Scheme *runtime.Scheme // PromClient is the Prometheus query client used for threshold evaluation. PromClient *metrics.Client + // Registrar manages registration in the MCP discovery registry. + // When non-nil, the canary controller triggers MCP re-registration after promotion. + Registrar *registry.Registrar // GatewayName is the name of the Gateway API Gateway object. GatewayName string // GatewayNamespace is the namespace of the Gateway API Gateway object. @@ -424,6 +428,19 @@ func (c *Controller) promote(ctx context.Context, ad *agentraxv1alpha1.AgentDepl return fmt.Errorf("updating status after promotion: %w", err) } + // Trigger MCP re-registration with updated endpoint / image after promotion. + if c.Registrar != nil && latest.Spec.MCP.Expose { + if err := c.Registrar.Register(ctx, latest); err != nil { + logger.Error(err, "MCP re-registration after promotion failed; reconciler will retry on next cycle") + } else if !latest.Status.Registered { + // Registration succeeded but status not yet set - update it. + latest.Status.Registered = true + if err := c.Client.Status().Update(ctx, latest); err != nil { + logger.Error(err, "failed to update registration status after promotion") + } + } + } + logger.Info("canary promoted successfully", "stableVersion", promoted) return nil }