fix: simplify sandbox pod lookup

Signed-off-by: Krrish <krrishrastogi00@gmail.com>

Author: krrishrastogi05

pkg/workloadmanager/handlers.go

@@ -235,13 +235,9 @@ func (s *Server) createSandbox(ctx context.Context, dynamicClient dynamic.Interf
 		return nil, errSandboxCreationTimeout
 	}
 
-	// agent-sandbox create pod with same name as sandbox if no warmpool is used
-	// so here we try to get pod IP by sandbox name first
-	// if warmpool is used, the pod name is stored in sandbox's annotation `agents.x-k8s.io/sandbox-pod-name`
-	// https://github.com/kubernetes-sigs/agent-sandbox/blob/3ab7fbcd85ad0d75c6e632ecd14bcaeda5e76e1e/controllers/sandbox_controller.go#L465
-	sandboxPodName := sandbox.Name
-	if podName, exists := createdSandbox.Annotations[controllers.SandboxPodNameAnnotation]; exists {
-		sandboxPodName = podName
+	sandboxPodName := createdSandbox.Annotations[controllers.SandboxPodNameAnnotation]
+	if sandboxPodName == "" {
+		return nil, api.NewInternalError(fmt.Errorf("sandbox %s/%s missing pod name annotation %s", sandbox.Namespace, sandbox.Name, controllers.SandboxPodNameAnnotation))
 	}
 
 	podIP, err := s.k8sClient.GetSandboxPodIP(ctx, sandbox.Namespace, sandbox.Name, sandboxPodName)

pkg/workloadmanager/handlers_test.go

@@ -111,6 +111,7 @@ func TestServerCreateSandbox(t *testing.T) {
 		readyErr          error
 		updateErr         error
 		sendResult        bool
+		missingPodName    bool
 		expectErr         bool
 		expectCreateCalls int
 		expectClaimCalls  int
@@ -159,6 +160,14 @@ func TestServerCreateSandbox(t *testing.T) {
 			expectCreateCalls: 1,
 			expectDeleteCalls: 1,
 		},
+		{
+			name:              "missing pod name annotation triggers rollback",
+			sendResult:        true,
+			missingPodName:    true,
+			expectErr:         true,
+			expectCreateCalls: 1,
+			expectDeleteCalls: 1,
+		},
 		{
 			name:              "entrypoint readiness failure triggers rollback",
 			readyErr:          errors.New("connection refused"),
@@ -233,6 +242,9 @@ func TestServerCreateSandbox(t *testing.T) {
 
 			resultChan := make(chan SandboxStatusUpdate, 1)
 			sb := readySandbox()
+			if tt.missingPodName {
+				delete(sb.Annotations, controllers.SandboxPodNameAnnotation)
+			}
 			if tt.sendResult {
 				resultChan <- SandboxStatusUpdate{Sandbox: sb.DeepCopy()}
 			}

pkg/workloadmanager/k8s_client.go

@@ -21,14 +21,11 @@ import (
 	"fmt"
 	"time"
 
-	"k8s.io/klog/v2"
-
 	runtimev1alpha1 "github.com/volcano-sh/agentcube/pkg/apis/runtime/v1alpha1"
 	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/apis/meta/v1/unstructured"
-	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/dynamic"
 	"k8s.io/client-go/dynamic/dynamicinformer"
@@ -296,31 +293,19 @@ func (u *UserK8sClient) DeleteSandboxClaim(ctx context.Context, namespace, sandb
 
 // GetSandboxPodIP gets the IP address of the pod corresponding to the Sandbox
 func (c *K8sClient) GetSandboxPodIP(_ context.Context, namespace, sandboxName, podName string) (string, error) {
-	// If podName is provided, try to get it directly from cache first
-	if podName != "" {
-		pod, err := c.podLister.Pods(namespace).Get(podName)
-		if err == nil && pod != nil {
-			return validateAndGetPodIP(pod)
-		}
-		klog.Infof("failed to get sandbox pod %s/%s: %v, try get pod by sandbox-name label", namespace, podName, err)
+	if podName == "" {
+		return "", fmt.Errorf("sandbox %s has no pod name annotation", sandboxName)
 	}
-	// Find pod through label selector (sandbox-name label we set)
-	pods, err := c.podLister.Pods(namespace).List(labels.SelectorFromSet(map[string]string{SandboxNameLabelKey: sandboxName}))
+
+	pod, err := c.podLister.Pods(namespace).Get(podName)
 	if err != nil {
-		return "", fmt.Errorf("failed to list pods from cache: %w", err)
+		return "", fmt.Errorf("failed to get sandbox pod %s/%s: %w", namespace, podName, err)
 	}
-	// Find the pod that belongs to this sandbox by checking ownerReferences
-	for _, pod := range pods {
-		for _, ownerRef := range pod.OwnerReferences {
-			if ownerRef.Kind == "Sandbox" && ownerRef.Name == sandboxName {
-				if ownerRef.Controller == nil || *ownerRef.Controller {
-					return validateAndGetPodIP(pod)
-				}
-			}
-		}
+	if pod == nil {
+		return "", fmt.Errorf("sandbox pod %s/%s not found", namespace, podName)
 	}
 
-	return "", fmt.Errorf("no pod found for sandbox %s", sandboxName)
+	return validateAndGetPodIP(pod)
 }
 
 // validateAndGetPodIP validates pod status and returns IP

pkg/workloadmanager/k8s_client_test.go

@@ -22,28 +22,19 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	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/labels"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	listersv1 "k8s.io/client-go/listers/core/v1"
 )
 
-// Helper function to create a pod with owner reference
-func createPodWithOwner(name, namespace, sandboxName string, phase corev1.PodPhase, podIP string) *corev1.Pod {
-	controller := true
+// Helper function to create a pod.
+func createPod(name, namespace string, phase corev1.PodPhase, podIP string) *corev1.Pod {
 	return &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      name,
 			Namespace: namespace,
-			Labels: map[string]string{
-				SandboxNameLabelKey: sandboxName,
-			},
-			OwnerReferences: []metav1.OwnerReference{
-				{
-					Kind:       "Sandbox",
-					Name:       sandboxName,
-					Controller: &controller,
-				},
-			},
 		},
 		Status: corev1.PodStatus{
 			Phase: phase,
@@ -52,6 +43,14 @@ func createPodWithOwner(name, namespace, sandboxName string, phase corev1.PodPha
 	}
 }
 
+func createPodWithLabel(name, namespace, sandboxName string, phase corev1.PodPhase, podIP string) *corev1.Pod {
+	pod := createPod(name, namespace, phase, podIP)
+	pod.Labels = map[string]string{
+		SandboxNameLabelKey: sandboxName,
+	}
+	return pod
+}
+
 // mockPodNamespaceLister is a mock implementation of listersv1.PodNamespaceLister
 type mockPodNamespaceLister struct {
 	pods []*corev1.Pod
@@ -76,7 +75,7 @@ func (m *mockPodNamespaceLister) Get(name string) (*corev1.Pod, error) {
 			return pod, nil
 		}
 	}
-	return nil, nil
+	return nil, apierrors.NewNotFound(schema.GroupResource{Resource: "pods"}, name)
 }
 
 // mockPodLister is a mock implementation of listersv1.PodLister
@@ -118,7 +117,7 @@ func (m *mockPodLister) addPod(pod *corev1.Pod) {
 // TestGetSandboxPodIP_Success verifies GetSandboxPodIP returns IP when pod is present and valid
 func TestGetSandboxPodIP_Success(t *testing.T) {
 	// Setup: Create a mock pod lister with a valid running pod
-	pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", corev1.PodRunning, "10.0.0.1")
+	pod := createPod("test-pod", "test-namespace", corev1.PodRunning, "10.0.0.1")
 	mockPodLister := newMockPodLister()
 	mockPodLister.addPod(pod)
 
@@ -127,43 +126,55 @@ func TestGetSandboxPodIP_Success(t *testing.T) {
 	}
 
 	// Execute
-	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "")
+	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "test-pod")
 
 	// Verify
 	assert.NoError(t, err, "Expected no error for valid pod")
 	assert.Equal(t, "10.0.0.1", ip, "Expected IP to match pod IP")
 }
 
-// TestGetSandboxPodIP_PodNotFound verifies GetSandboxPodIP returns error when pod is not found
+func TestGetSandboxPodIP_MissingPodName(t *testing.T) {
+	client := &K8sClient{
+		podLister: newMockPodLister(),
+	}
+
+	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "")
+
+	assert.Error(t, err, "Expected error when pod name is empty")
+	assert.Empty(t, ip, "Expected empty IP when error occurs")
+	assert.Contains(t, err.Error(), "sandbox test-sandbox has no pod name annotation")
+}
+
+// TestGetSandboxPodIP_PodNotFound verifies GetSandboxPodIP returns error when annotated pod is not found.
 func TestGetSandboxPodIP_PodNotFound(t *testing.T) {
-	// Setup: Create a mock pod lister with pod that has wrong label
 	mockPodLister := newMockPodLister()
-	pod := &corev1.Pod{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      "test-pod",
-			Namespace: "test-namespace",
-			Labels: map[string]string{
-				SandboxNameLabelKey: "other-sandbox",
-			},
-		},
-		Status: corev1.PodStatus{
-			Phase: corev1.PodRunning,
-			PodIP: "10.0.0.1",
-		},
-	}
-	mockPodLister.addPod(pod)
 
 	client := &K8sClient{
 		podLister: mockPodLister,
 	}
 
 	// Execute
-	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "")
+	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "missing-pod")
 
 	// Verify
 	assert.Error(t, err, "Expected error when pod not found")
 	assert.Empty(t, ip, "Expected empty IP when error occurs")
-	assert.Contains(t, err.Error(), "no pod found for sandbox test-sandbox", "Error message should indicate pod not found")
+	assert.Contains(t, err.Error(), "failed to get sandbox pod test-namespace/missing-pod", "Error message should indicate pod not found")
+}
+
+func TestGetSandboxPodIP_DoesNotFallbackToLabelSelector(t *testing.T) {
+	mockPodLister := newMockPodLister()
+	mockPodLister.addPod(createPodWithLabel("label-matched-pod", "test-namespace", "test-sandbox", corev1.PodRunning, "10.0.0.1"))
+
+	client := &K8sClient{
+		podLister: mockPodLister,
+	}
+
+	ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "missing-pod")
+
+	assert.Error(t, err, "Expected error when annotated pod is missing")
+	assert.Empty(t, ip, "Expected empty IP when error occurs")
+	assert.Contains(t, err.Error(), "missing-pod")
 }
 
 // TestGetSandboxPodIP_InvalidPodStatus verifies GetSandboxPodIP returns error when pod status is invalid
@@ -192,7 +203,7 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			// Setup: Create a mock pod lister with invalid pod status
-			pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", tc.phase, tc.podIP)
+			pod := createPod("test-pod", "test-namespace", tc.phase, tc.podIP)
 			mockPodLister := newMockPodLister()
 			mockPodLister.addPod(pod)
 
@@ -201,7 +212,7 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) {
 			}
 
 			// Execute
-			ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "")
+			ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "test-pod")
 
 			// Verify
 			assert.Error(t, err, "Expected error for invalid pod status")

test/e2e/run_e2e.sh

@@ -6,7 +6,7 @@ IFS=$'\n\t'
 E2E_CLUSTER_NAME=${E2E_CLUSTER_NAME:-agentcube-e2e}
 E2E_CLEAN_CLUSTER=${E2E_CLEAN_CLUSTER:-true}
 E2E_SKIP_SETUP=${E2E_SKIP_SETUP:-false}
-AGENT_SANDBOX_VERSION=${AGENT_SANDBOX_VERSION:-v0.1.1}
+AGENT_SANDBOX_VERSION=${AGENT_SANDBOX_VERSION:-v0.3.10}
 WORKLOAD_MANAGER_IMAGE=${WORKLOAD_MANAGER_IMAGE:-workloadmanager:latest}
 ROUTER_IMAGE=${ROUTER_IMAGE:-agentcube-router:latest}
 PICOD_IMAGE=${PICOD_IMAGE:-picod:latest}