fix(e2b): real command exit codes, pod-IP wait after resume, warm-pool status

pkg/sandbox/e2b/gateway.go

@@ -246,11 +246,78 @@ func (s *Server) wakeForTraffic(r *http.Request, sandboxID string) (*pb.GetSessi
 			return nil, false, connectError("unavailable", "auto-resume failed: "+rerr.Error())
 		}
 		s.registry.markPaused(sandboxID, false)
-		return sess, true, nil
+		// Resume re-creates the pod asynchronously, so the session view we
+		// already hold still reports the pre-pause state: PodIp is empty
+		// because the old pod is gone and the new one has not been assigned an
+		// address yet. Callers address sandboxd by pod IP, so returning that
+		// stale view made every post-resume call fail with
+		// "session ... has no pod IP" (E2B watch_dir, files.stat, ...).
+		//
+		// Poll until the new pod reports an address. Bounded: a resume that
+		// never schedules a pod is an infrastructure failure, and blocking the
+		// request forever is worse than reporting it.
+		fresh, perr := s.waitForPodIP(r, sandboxID)
+		if perr != nil {
+			return nil, false, connectError("unavailable", "sandbox resumed but has no pod IP: "+perr.Error())
+		}
+		return fresh, true, nil
+	}
+	// Running: normally return the live view. But a sandbox that an explicit
+	// resume() just brought back is already marked running while its pod is
+	// still being created, so its PodIp is momentarily empty. Callers address
+	// sandboxd by pod IP, so wait briefly rather than failing the request —
+	// otherwise the first files/process call after resume() reports
+	// "has no pod IP".
+	if st == stateRunning && sess.GetPodIp() == "" {
+		if fresh, perr := s.waitForPodIP(r, sandboxID); perr == nil {
+			return fresh, false, nil
+		}
 	}
 	return sess, false, nil
 }
 
+// waitForPodIP polls the session until the scheduler has assigned the
+// (re)created pod an address, returning the fresh session view.
+//
+// Poll rather than a single read because resume is asynchronous: the pod is
+// claimed and scheduled after ResumeSession returns, and the pod IP cache in
+// the sandboxd client is deliberately short-lived (e2bPodIPTTL) so it cannot
+// paper over the gap.
+func (s *Server) waitForPodIP(r *http.Request, sandboxID string) (*pb.GetSessionResponse, error) {
+	const (
+		// podIPWaitTimeout bounds the whole wait. Warm-pool claims normally
+		// land in well under a second; a cold pod needs scheduling time.
+		podIPWaitTimeout = 30 * time.Second
+		podIPWaitPoll    = 200 * time.Millisecond
+	)
+	deadline := time.Now().Add(podIPWaitTimeout)
+	for {
+		sess, err := s.gw.GetSession(r.Context(), &pb.GetSessionRequest{SessionId: sandboxID})
+		if err == nil && sess.GetPodIp() != "" {
+			// The pod changed, so any cached address for this session is stale.
+			s.invalidatePodIP(sandboxID)
+			return sess, nil
+		}
+		if !time.Now().Before(deadline) {
+			if err != nil {
+				return nil, err
+			}
+			return nil, fmt.Errorf("pod IP not assigned after %s", podIPWaitTimeout)
+		}
+		time.Sleep(podIPWaitPoll)
+	}
+}
+
+// invalidatePodIP drops any cached sandboxd address for a session. Pause,
+// resume and pod replacement all change the address, so the short pod-IP cache
+// must be dropped at those boundaries or it keeps answering with an IP that no
+// longer exists.
+func (s *Server) invalidatePodIP(sandboxID string) {
+	if s.sandboxd != nil {
+		s.sandboxd.invalidatePodIP(sandboxID)
+	}
+}
+
 // unixNow returns the current unix time in seconds.
 func unixNow() int64 { return time.Now().Unix() }
 

pkg/sandbox/e2b/gateway_podip_test.go

@@ -0,0 +1,158 @@
+package e2b
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+
+	pb "github.com/xiaods/k8e/pkg/sandboxmatrix/grpc/pb/sandbox/v1"
+)
+
+// asyncPodIPGateway models the real resume path: ResumeSession marks the
+// session Active but the pod IP is only assigned a few polls later, exactly as
+// the scheduler behaves. A single GetSession immediately after resume therefore
+// reports an empty PodIp — which is what made the E2B layer answer
+// "session ... has no pod IP".
+type asyncPodIPGateway struct {
+	*fakeGateway
+	mu       sync.Mutex
+	gets     int
+	assignAt int // GetSession call number at/after which PodIp appears
+}
+
+func (g *asyncPodIPGateway) GetSession(ctx context.Context, req *pb.GetSessionRequest) (*pb.GetSessionResponse, error) {
+	g.mu.Lock()
+	g.gets++
+	n := g.gets
+	assign := n >= g.assignAt
+	g.mu.Unlock()
+
+	sess, err := g.fakeGateway.GetSession(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	if assign && sess.GetPodIp() == "" {
+		// Simulate the scheduler finally assigning an address.
+		g.fakeGateway.mu.Lock()
+		g.fakeGateway.sessions[req.SessionId].PodIp = "10.42.0.99"
+		g.fakeGateway.mu.Unlock()
+		return g.fakeGateway.GetSession(ctx, req)
+	}
+	return sess, nil
+}
+
+// PauseSession drops the pod address, mirroring the real controller: a paused
+// sandbox has no running pod, so status.podIP is cleared until it is resumed.
+// Without this the fake would leave a stale address behind and the pod-IP race
+// would never be exercised.
+func (g *asyncPodIPGateway) PauseSession(ctx context.Context, req *pb.PauseSessionRequest) (*pb.PauseSessionResponse, error) {
+	resp, err := g.fakeGateway.PauseSession(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	g.fakeGateway.mu.Lock()
+	if s, ok := g.fakeGateway.sessions[req.SessionId]; ok {
+		s.PodIp = ""
+	}
+	g.fakeGateway.mu.Unlock()
+	return resp, nil
+}
+
+// listDir issues an authenticated envd ListDir, which is the request that
+// triggers wakeForTraffic's auto-resume and then must reach sandboxd by pod IP.
+func listDir(t *testing.T, ts *httptest.Server, s *Server, id string) {
+	t.Helper()
+	req, err := http.NewRequest(http.MethodPost,
+		ts.URL+"/e2b/envd/filesystem.Filesystem/ListDir",
+		strings.NewReader(`{"path":"/workspace","depth":1}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for k, v := range envdHeaders(t, s, id) {
+		req.Header.Set(k, v)
+	}
+	resp, err := ts.Client().Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("ListDir after resume: %d %s", resp.StatusCode, readBody(t, resp))
+	}
+}
+
+// TestWakeForTrafficWaitsForPodIPAfterResume is the regression test for the E2B
+// bug where the first call after a pause failed with "session has no pod IP".
+func TestWakeForTrafficWaitsForPodIPAfterResume(t *testing.T) {
+	base := newFakeGateway()
+	gw := &asyncPodIPGateway{fakeGateway: base, assignAt: 6}
+	stub := newSandboxdStub()
+	s, ts := testServerWithSandboxd(t, gw, stub)
+	defer ts.Close()
+
+	id := createSandboxID(t, ts)
+
+	// Seed the workspace so ListDir has something real to stat.
+	gw.mu.Lock()
+	gw.files["/workspace/probe.txt"] = "x"
+	gw.mu.Unlock()
+	stub.mu.Lock()
+	stub.dirs["/workspace/"] = true
+	stub.files["/workspace/probe.txt"] = "x"
+	stub.mu.Unlock()
+
+	// Pause: the pod address goes away and the phase becomes Paused.
+	if resp := controlReq(t, ts, "POST", "/sandboxes/"+id+"/pause", nil); resp.StatusCode != 204 {
+		t.Fatalf("pause: %d %s", resp.StatusCode, readBody(t, resp))
+	}
+
+	// The call must auto-resume AND resolve to a usable pod IP. Before the fix
+	// it returned the pre-resume view whose PodIp was empty, and the sandboxd
+	// client answered "session ... has no pod IP".
+	listDir(t, ts, s, id)
+
+	if got := len(base.resumed); got != 1 {
+		t.Fatalf("expected exactly one auto-resume, got %d", got)
+	}
+	if base.sessions[id].GetPodIp() == "" {
+		t.Fatal("auto-resume must leave the session with an assigned pod IP")
+	}
+}
+
+// TestWakeForTrafficRunningWithoutPodIPWaits covers the explicit-resume case:
+// the phase is already running while the pod is still being created.
+func TestWakeForTrafficRunningWithoutPodIPWaits(t *testing.T) {
+	base := newFakeGateway()
+	gw := &asyncPodIPGateway{fakeGateway: base, assignAt: 3}
+	stub := newSandboxdStub()
+	s, ts := testServerWithSandboxd(t, gw, stub)
+	defer ts.Close()
+
+	id := createSandboxID(t, ts)
+
+	gw.mu.Lock()
+	gw.files["/workspace/probe.txt"] = "x"
+	gw.mu.Unlock()
+	stub.mu.Lock()
+	stub.dirs["/workspace/"] = true
+	stub.files["/workspace/probe.txt"] = "x"
+	stub.mu.Unlock()
+
+	// A running session whose pod has no address yet (what an explicit resume
+	// leaves behind) must not fail the request outright.
+	base.mu.Lock()
+	base.sessions[id].PodIp = ""
+	base.mu.Unlock()
+
+	listDir(t, ts, s, id)
+
+	if len(base.resumed) != 0 {
+		t.Fatalf("a running session must not be resumed again, got %d resumes", len(base.resumed))
+	}
+	if base.sessions[id].GetPodIp() == "" {
+		t.Fatal("the call must have waited for the pod IP to be assigned")
+	}
+}

pkg/sandboxmatrix/controller.go

@@ -215,7 +215,7 @@ func reconcileWarmPools(ctx context.Context, k8s kubernetes.Interface, dyn dynam
 		boost = demand.observe(time.Now(), coldStarts)
 	}
 	for _, pool := range pools.Items {
-		reconcileSinglePool(ctx, k8s, pool, maxPods, cfg, boost)
+		reconcileSinglePool(ctx, k8s, dyn, pool, maxPods, cfg, boost)
 	}
 	recycleUnhealthyWarmPods(ctx, k8s, cfg.Namespace)
 	updateSandboxMatrixStatus(ctx, k8s, dyn, cfg, orch, maxPods)
@@ -357,7 +357,7 @@ func deleteWarmPod(ctx context.Context, k8s kubernetes.Interface, namespace stri
 }
 
 // reconcileSinglePool ensures one WarmPool CRD's target is met within capacity limits.
-func reconcileSinglePool(ctx context.Context, k8s kubernetes.Interface, pool unstructured.Unstructured, maxPods int64, cfg config.SandboxConfig, boost int64) {
+func reconcileSinglePool(ctx context.Context, k8s kubernetes.Interface, dyn dynamic.Interface, pool unstructured.Unstructured, maxPods int64, cfg config.SandboxConfig, boost int64) {
 	specMap, _ := pool.Object["spec"].(map[string]interface{})
 	configuredSize, _ := specMap["size"].(int64)
 	runtimeClass, _ := specMap["runtimeClass"].(string)
@@ -404,8 +404,10 @@ func reconcileSinglePool(ctx context.Context, k8s kubernetes.Interface, pool uns
 		// returning from resetting, or a target that was just lowered. Scale
 		// down now instead of waiting for the idle reaper's sessionTTL×2 sweep.
 		trimWarmPods(ctx, k8s, cfg.Namespace, warmPods.Items, -gap)
+		updateWarmPoolStatus(ctx, dyn, cfg, pool, warmPods.Items, 0)
 		return
 	case gap == 0:
+		updateWarmPoolStatus(ctx, dyn, cfg, pool, warmPods.Items, 0)
 		return
 	}
 	if maxPods >= 0 && int64(len(allPods.Items)) >= maxPods {
@@ -419,6 +421,36 @@ func reconcileSinglePool(ctx context.Context, k8s kubernetes.Interface, pool uns
 			logrus.Debugf("sandbox-matrix: create warm pod: %v", err)
 		}
 	}
+	// The pods we just asked for are still Pending, so report them as pending
+	// rather than ready: an operator watching the pool sees "1 pending" until
+	// the next pass flips it to ready, instead of an empty status that looks
+	// like a broken pool.
+	updateWarmPoolStatus(ctx, dyn, cfg, pool, warmPods.Items, gap)
+}
+
+// updateWarmPoolStatus writes the pool's ready/pending counts to its .status.
+//
+// The CRD has carried readyCount/pendingCount since KIP-30 but nothing ever
+// populated them, so `kubectl get sandboxwarmpools` reported an empty READY
+// column and operators had no in-cluster signal for "is my pool actually
+// serving?". Pending is derived from the gap this pass is closing rather than
+// counted from pod phases, because a pod created a moment ago has no readable
+// phase yet.
+func updateWarmPoolStatus(ctx context.Context, dyn dynamic.Interface, cfg config.SandboxConfig, pool unstructured.Unstructured, warmPods []corev1.Pod, pending int64) {
+	var ready int64
+	for i := range warmPods {
+		if sandboxgrpc.PodReadyCondition(&warmPods[i]) {
+			ready++
+		}
+	}
+	status, ok := pool.Object["status"].(map[string]interface{})
+	if !ok {
+		status = map[string]interface{}{}
+		pool.Object["status"] = status
+	}
+	status["readyCount"] = ready
+	status["pendingCount"] = pending
+	dyn.Resource(warmPoolGVR).Namespace(cfg.Namespace).UpdateStatus(ctx, &pool, metav1.UpdateOptions{}) //nolint:errcheck
 }
 
 // createWarmPodWithWorkspace provisions a blank workspace PVC and boots a warm