fix: preserve pod template annotations during warm adoption

Author: noeljackson

controllers/sandbox_controller.go

@@ -424,6 +424,9 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		if pod.Labels == nil {
 			pod.Labels = make(map[string]string)
 		}
+		if pod.Annotations == nil {
+			pod.Annotations = make(map[string]string)
+		}
 		changed := false
 		if pod.Labels[sandboxLabel] != nameHash {
 			pod.Labels[sandboxLabel] = nameHash
@@ -436,6 +439,12 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 				changed = true
 			}
 		}
+		for k, v := range sandbox.Spec.PodTemplate.ObjectMeta.Annotations {
+			if pod.Annotations[k] != v {
+				pod.Annotations[k] = v
+				changed = true
+			}
+		}
 
 		// Set controller reference if the pod is not controlled by anything.
 		if controllerRef := metav1.GetControllerOf(pod); controllerRef == nil {

controllers/sandbox_controller_test.go

@@ -613,6 +613,9 @@ func TestReconcilePod(t *testing.T) {
 						"agents.x-k8s.io/sandbox-name-hash": nameHash,
 						"custom-label":                      "label-val",
 					},
+					Annotations: map[string]string{
+						"custom-annotation": "anno-val",
+					},
 					OwnerReferences: []metav1.OwnerReference{sandboxControllerRef(sandboxName)},
 				},
 				Spec: corev1.PodSpec{
@@ -685,7 +688,7 @@ func TestReconcilePod(t *testing.T) {
 			wantPod: nil,
 		},
 		{
-			name: "adopts existing pod via annotation - pod gets label and owner reference",
+			name: "adopts existing pod via annotation - pod gets metadata and owner reference",
 			initialObjs: []runtime.Object{
 				&corev1.Pod{
 					ObjectMeta: metav1.ObjectMeta{
@@ -713,6 +716,11 @@ func TestReconcilePod(t *testing.T) {
 				Spec: sandboxv1alpha1.SandboxSpec{
 					Replicas: ptr.To(int32(1)),
 					PodTemplate: sandboxv1alpha1.PodTemplate{
+						ObjectMeta: sandboxv1alpha1.PodMetadata{
+							Annotations: map[string]string{
+								"io.codewire.sh/workspace": "true",
+							},
+						},
 						Spec: corev1.PodSpec{
 							Containers: []corev1.Container{
 								{
@@ -731,6 +739,9 @@ func TestReconcilePod(t *testing.T) {
 					Labels: map[string]string{
 						sandboxLabel: nameHash,
 					},
+					Annotations: map[string]string{
+						"io.codewire.sh/workspace": "true",
+					},
 					OwnerReferences: []metav1.OwnerReference{sandboxControllerRef(sandboxName)},
 				},
 				Spec: corev1.PodSpec{
@@ -783,6 +794,9 @@ func TestReconcilePod(t *testing.T) {
 						"agents.x-k8s.io/sandbox-name-hash": nameHash,
 						"custom-label":                      "label-val",
 					},
+					Annotations: map[string]string{
+						"custom-annotation": "anno-val",
+					},
 					// Should still have the original controller reference
 					OwnerReferences: []metav1.OwnerReference{
 						{

extensions/controllers/sandboxclaim_controller.go

@@ -379,8 +379,32 @@ func applyClaimWorkspaceResourcesToPodSpec(spec *corev1.PodSpec, claim *extensio
 		applyWorkspaceResourceOverrides(container, claim.Spec.WorkspaceResources)
 	}
 }
+
+func mergeTemplatePodMetadata(target *v1alpha1.PodMetadata, template v1alpha1.PodMetadata) {
+	if len(template.Labels) > 0 {
+		if target.Labels == nil {
+			target.Labels = make(map[string]string, len(template.Labels))
+		}
+		for k, v := range template.Labels {
+			if _, exists := target.Labels[k]; !exists {
+				target.Labels[k] = v
+			}
+		}
+	}
+	if len(template.Annotations) > 0 {
+		if target.Annotations == nil {
+			target.Annotations = make(map[string]string, len(template.Annotations))
+		}
+		for k, v := range template.Annotations {
+			if _, exists := target.Annotations[k]; !exists {
+				target.Annotations[k] = v
+			}
+		}
+	}
+}
+
 // adoptSandboxFromCandidates picks the best candidate and transfers ownership to the claim.
-func (r *SandboxClaimReconciler) adoptSandboxFromCandidates(ctx context.Context, claim *extensionsv1alpha1.SandboxClaim, candidates []*v1alpha1.Sandbox) (*v1alpha1.Sandbox, error) {
+func (r *SandboxClaimReconciler) adoptSandboxFromCandidates(ctx context.Context, claim *extensionsv1alpha1.SandboxClaim, template *extensionsv1alpha1.SandboxTemplate, candidates []*v1alpha1.Sandbox) (*v1alpha1.Sandbox, error) {
 	log := log.FromContext(ctx)
 
 	// Sort: ready sandboxes first, then by creation time (oldest first)
@@ -441,6 +465,10 @@ func (r *SandboxClaimReconciler) adoptSandboxFromCandidates(ctx context.Context,
 			adopted.Annotations[asmetrics.TraceContextAnnotation] = tc
 		}
 
+		if template != nil {
+			mergeTemplatePodMetadata(&adopted.Spec.PodTemplate.ObjectMeta, template.Spec.PodTemplate.ObjectMeta)
+		}
+
 		// Add sandbox ID label to pod template for NetworkPolicy targeting
 		if adopted.Spec.PodTemplate.ObjectMeta.Labels == nil {
 			adopted.Spec.PodTemplate.ObjectMeta.Labels = make(map[string]string)
@@ -647,7 +675,12 @@ func (r *SandboxClaimReconciler) getOrCreateSandbox(ctx context.Context, claim *
 
 	// Try to adopt from warm pool
 	if len(adoptionCandidates) > 0 {
-		adopted, err := r.adoptSandboxFromCandidates(ctx, claim, adoptionCandidates)
+		var template *extensionsv1alpha1.SandboxTemplate
+		template, err := r.getTemplate(ctx, claim)
+		if err != nil && !k8errors.IsNotFound(err) {
+			return nil, err
+		}
+		adopted, err := r.adoptSandboxFromCandidates(ctx, claim, template, adoptionCandidates)
 		if err != nil {
 			return nil, err
 		}

extensions/controllers/sandboxclaim_controller_test.go

@@ -1277,6 +1277,12 @@ func TestSandboxClaimAdoptionAppliesWorkspaceResources(t *testing.T) {
 		ObjectMeta: metav1.ObjectMeta{Name: "test-template", Namespace: "default"},
 		Spec: extensionsv1alpha1.SandboxTemplateSpec{
 			PodTemplate: sandboxv1alpha1.PodTemplate{
+				ObjectMeta: sandboxv1alpha1.PodMetadata{
+					Annotations: map[string]string{
+						"io.codewire.sh/workspace": "true",
+						"test-annotation":          "template",
+					},
+				},
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
 						{Name: "workspace", Image: "workspace:latest"},
@@ -1320,6 +1326,11 @@ func TestSandboxClaimAdoptionAppliesWorkspaceResources(t *testing.T) {
 		Spec: sandboxv1alpha1.SandboxSpec{
 			Replicas: ptr.To(int32(1)),
 			PodTemplate: sandboxv1alpha1.PodTemplate{
+				ObjectMeta: sandboxv1alpha1.PodMetadata{
+					Annotations: map[string]string{
+						"io.codewire.sh/workspace": "true",
+					},
+				},
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
 						{Name: "workspace", Image: "workspace:latest"},
@@ -1389,6 +1400,12 @@ func TestSandboxClaimAdoptionAppliesWorkspaceResources(t *testing.T) {
 	if len(sidecar.Resources.Requests) != 0 || len(sidecar.Resources.Limits) != 0 {
 		t.Fatalf("expected adopted sidecar resources to remain untouched, got requests=%v limits=%v", sidecar.Resources.Requests, sidecar.Resources.Limits)
 	}
+	if got := adopted.Spec.PodTemplate.ObjectMeta.Annotations["io.codewire.sh/workspace"]; got != "true" {
+		t.Fatalf("expected adopted workspace annotation to survive, got %q", got)
+	}
+	if got := adopted.Spec.PodTemplate.ObjectMeta.Annotations["test-annotation"]; got != "template" {
+		t.Fatalf("expected template annotation to be restored on adoption, got %q", got)
+	}
 }
 
 func TestRecordCreationLatencyMetric(t *testing.T) {