Gate proxyrunner StatefulSet apply by MCPServer generation (#5024)

Guard proxyrunner StatefulSet apply with a generation gate

During a proxy Deployment rolling update the old and new ReplicaSet
pods both run DeployWorkload, each calling server-side apply on the
backing StatefulSet with the same field manager and Force: true. The
stale pod's apply can land after the fresh pod's and clobber the
image, leaving the StatefulSet pinned to the old digest even after
MCPServer.spec.image has been bumped.

The operator now stamps MCPServer.metadata.generation into the
RunConfig as a monotonic version. Proxyrunner threads it through to
DeployWorkloadOptions. Before apply, the K8s runtime reads the
existing StatefulSet's mcpserver-generation annotation; if it is
strictly greater than ours, the apply is skipped. Otherwise our
generation is stamped on the pod template and apply proceeds.

Zero generation is the backward-compat signal and bypasses both the
gate and the stamp, preserving existing behavior for callers that
don't pass one.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Yolanda Robla Mota <yolanda@stacklok.com>

Author: 3 people

cmd/thv-operator/controllers/mcpserver_runconfig.go

@@ -144,6 +144,7 @@ func (r *MCPServerReconciler) createRunConfigFromMCPServer(m *mcpv1beta1.MCPServ
 	options := []runner.RunConfigBuilderOption{
 		runner.WithName(m.Name),
 		runner.WithImage(m.Spec.Image),
+		runner.WithMCPServerGeneration(m.Generation),
 		runner.WithCmdArgs(m.Spec.Args),
 		runner.WithTransportAndPorts(m.Spec.Transport, int(m.GetProxyPort()), int(m.GetMCPPort())),
 		runner.WithProxyMode(transporttypes.ProxyMode(effectiveProxyMode)),

cmd/thv-operator/controllers/mcpserver_runconfig_test.go

@@ -1782,3 +1782,34 @@ func TestCreateRunConfigFromMCPServer_RateLimiting(t *testing.T) {
 		})
 	}
 }
+
+func TestCreateRunConfigFromMCPServer_SetsMCPServerGeneration(t *testing.T) {
+	t.Parallel()
+
+	m := &mcpv1beta1.MCPServer{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:       "generation-server",
+			Namespace:  "default",
+			Generation: 7,
+		},
+		Spec: mcpv1beta1.MCPServerSpec{
+			Image:     "ghcr.io/example/mcp:v1",
+			Transport: stdioTransport,
+			ProxyPort: 8080,
+		},
+	}
+
+	r := newTestMCPServerReconciler(
+		fake.NewClientBuilder().WithScheme(createRunConfigTestScheme()).WithObjects(m).Build(),
+		createRunConfigTestScheme(),
+		kubernetes.PlatformKubernetes,
+	)
+
+	rc, err := r.createRunConfigFromMCPServer(m)
+
+	require.NoError(t, err)
+	require.NotNil(t, rc)
+
+	assert.Equal(t, int64(7), rc.MCPServerGeneration,
+		"MCPServerGeneration should match MCPServer .metadata.generation")
+}

docs/server/docs.go

@@ -49,6 +49,15 @@ const (
 	defaultNamespace = "default"
 	// serviceFieldManager is the field manager name for server-side apply operations
 	serviceFieldManager = "toolhive-container-manager"
+
+	// RunConfigMCPServerGenerationAnnotation carries the MCPServer .metadata.generation that
+	// produced the RunConfig applied to this StatefulSet. Used as a monotonic version stamp
+	// to prevent stale proxyrunner pods (from an old Deployment ReplicaSet) from clobbering
+	// a newer RunConfig's apply. The gate only becomes effective once proxyrunner is upgraded
+	// to a version that reads this annotation; operator-only upgrades leave the race window
+	// in place until proxyrunner is also rolled. Exported because it forms a wire contract
+	// that external readers (operator, diagnostic tooling) may consume.
+	RunConfigMCPServerGenerationAnnotation = "toolhive.stacklok.dev/mcpserver-generation"
 )
 
 // RuntimeName is the name identifier for the Kubernetes runtime
@@ -397,22 +406,26 @@ func (c *Client) DeployWorkload(ctx context.Context,
 		return 0, err
 	}
 
-	// Create an apply configuration for the statefulset
-	statefulSetApply := appsv1apply.StatefulSet(containerName, namespace).
-		WithLabels(containerLabels).
-		WithSpec(buildStatefulSetSpec(containerName, podTemplateSpec, options))
-
-	// Apply the statefulset using server-side apply
-	createdStatefulSet, err := c.client.AppsV1().StatefulSets(namespace).
-		Apply(ctx, statefulSetApply, metav1.ApplyOptions{
-			FieldManager: serviceFieldManager,
-			Force:        true,
-		})
+	ourGen := runConfigGeneration(options)
+	skip, err := c.shouldSkipStatefulSetApply(ctx, namespace, containerName, ourGen)
 	if err != nil {
-		return 0, fmt.Errorf("failed to apply statefulset: %w", err)
+		return 0, err
+	}
+	if skip {
+		// Intentionally skip ensureBackendServices in the gated path: this pod's RunConfig
+		// is stale, so reconciling services here would clobber port/config fields set by
+		// the newer-generation pod under the same field manager + Force: true — the same
+		// race this gate prevents for the StatefulSet. The newer pod already reconciled
+		// services; if that failed, it returns an error and retries on its own.
+		return 0, nil
 	}
 
-	slog.Debug("applied statefulset", "name", createdStatefulSet.Name)
+	createdStatefulSet, err := c.applyStatefulSet(
+		ctx, namespace, containerName, containerLabels, podTemplateSpec, options, ourGen,
+	)
+	if err != nil {
+		return 0, err
+	}
 
 	err = c.ensureBackendServices(
 		ctx, containerName, namespace, containerLabels, transportType, options, createdStatefulSet)
@@ -435,6 +448,84 @@ func (c *Client) DeployWorkload(ctx context.Context,
 	return 0, nil
 }
 
+// runConfigGeneration extracts the RunConfig MCPServer generation from options,
+// returning 0 when options is nil (backward-compat / non-operator callers).
+func runConfigGeneration(options *runtime.DeployWorkloadOptions) int64 {
+	if options == nil {
+		return 0
+	}
+	return options.RunConfigMCPServerGeneration
+}
+
+// applyStatefulSet stamps the MCPServer generation annotation when non-zero,
+// builds the StatefulSet apply configuration, and performs the server-side apply.
+func (c *Client) applyStatefulSet(
+	ctx context.Context,
+	namespace, containerName string,
+	containerLabels map[string]string,
+	podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
+	options *runtime.DeployWorkloadOptions,
+	ourGen int64,
+) (*appsv1.StatefulSet, error) {
+	if ourGen > 0 {
+		podTemplateSpec = podTemplateSpec.WithAnnotations(map[string]string{
+			RunConfigMCPServerGenerationAnnotation: strconv.FormatInt(ourGen, 10),
+		})
+	}
+	statefulSetApply := appsv1apply.StatefulSet(containerName, namespace).
+		WithLabels(containerLabels).
+		WithSpec(buildStatefulSetSpec(containerName, podTemplateSpec, options))
+	createdStatefulSet, err := c.client.AppsV1().StatefulSets(namespace).
+		Apply(ctx, statefulSetApply, metav1.ApplyOptions{
+			FieldManager: serviceFieldManager,
+			Force:        true,
+		})
+	if err != nil {
+		return nil, fmt.Errorf("failed to apply statefulset: %w", err)
+	}
+	slog.Debug("applied statefulset", "name", createdStatefulSet.Name)
+	return createdStatefulSet, nil
+}
+
+// shouldSkipStatefulSetApply returns true when the existing StatefulSet is already
+// stamped with a strictly greater MCPServer generation than ours, meaning a newer
+// proxyrunner pod has already reconciled the workload and ours would be a regression.
+// Returns false (apply as normal) when ourGen is zero or negative, when the StatefulSet
+// does not yet exist, when the annotation is absent, or when the annotation is unparsable.
+func (c *Client) shouldSkipStatefulSetApply(
+	ctx context.Context, namespace, name string, ourGen int64,
+) (bool, error) {
+	if ourGen <= 0 {
+		return false, nil
+	}
+	existing, err := c.client.AppsV1().StatefulSets(namespace).Get(ctx, name, metav1.GetOptions{})
+	if err != nil {
+		if errors.IsNotFound(err) {
+			return false, nil
+		}
+		return false, fmt.Errorf("failed to get existing statefulset: %w", err)
+	}
+	if existing.Spec.Template.Annotations == nil {
+		return false, nil
+	}
+	theirs := existing.Spec.Template.Annotations[RunConfigMCPServerGenerationAnnotation]
+	if theirs == "" {
+		return false, nil
+	}
+	theirsGen, parseErr := strconv.ParseInt(theirs, 10, 64)
+	if parseErr != nil {
+		slog.Warn("unparsable mcpserver-generation annotation; proceeding with apply",
+			"sts", name, "value", theirs, "err", parseErr)
+		return false, nil
+	}
+	if theirsGen > ourGen {
+		slog.Debug("skipping StatefulSet apply; newer MCPServer generation already applied",
+			"sts", name, "ours", ourGen, "theirs", theirsGen)
+		return true, nil
+	}
+	return false, nil
+}
+
 // buildStatefulSetSpec constructs the StatefulSet spec apply configuration.
 // WithReplicas is only included when BackendReplicas is explicitly set; omitting
 // the field lets the existing field manager (e.g. HPA or kubectl) retain control