fix(orchestrator): enforce sandbox TTL on the node via WaitForExit

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -1720,20 +1720,33 @@ func (s *Sandbox) WaitForExit(ctx context.Context) error {
 	ctx, span := tracer.Start(ctx, "sandbox-wait-for-exit")
 	defer span.End()
 
-	timeout := time.Until(s.GetEndAt())
+	for {
+		endAt := s.GetEndAt()
+		timeout := time.Until(endAt)
+		if timeout <= 0 {
+			return errors.New("sandbox TTL expired")
+		}
 
-	select {
-	case <-time.After(timeout):
-		return errors.New("waiting for exit took too long")
-	case <-ctx.Done():
-		return nil
-	case <-s.exit.Done():
-		err := s.exit.Error()
-		if err == nil {
+		timer := time.NewTimer(timeout)
+		select {
+		case <-timer.C:
+			// Re-check in case SetEndAt extended the deadline (e.g. via KeepAlive).
+			if s.GetEndAt().After(endAt) {
+				continue
+			}
+			return errors.New("waiting for exit took too long")
+		case <-ctx.Done():
+			timer.Stop()
 			return nil
-		}
+		case <-s.exit.Done():
+			timer.Stop()
+			err := s.exit.Error()
+			if err == nil {
+				return nil
+			}
 
-		return fmt.Errorf("fc process exited prematurely: %w", err)
+			return fmt.Errorf("fc process exited prematurely: %w", err)
+		}
 	}
 }
 

packages/orchestrator/pkg/sandbox/wait_for_exit_test.go

@@ -0,0 +1,137 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+func newTestSandboxForExit(endAt time.Time) *Sandbox {
+	sbx := &Sandbox{
+		Metadata: &Metadata{},
+		exit:     utils.NewErrorOnce(),
+	}
+	sbx.SetEndAt(endAt)
+
+	return sbx
+}
+
+// TestWaitForExit_AlreadyExpired verifies that WaitForExit returns an error
+// immediately when endAt is already in the past.
+func TestWaitForExit_AlreadyExpired(t *testing.T) {
+	t.Parallel()
+
+	sbx := newTestSandboxForExit(time.Now().Add(-1 * time.Second))
+
+	err := sbx.WaitForExit(context.Background())
+
+	require.Error(t, err)
+}
+
+// TestWaitForExit_ExitBeforeTTL verifies that WaitForExit returns nil when
+// the sandbox exits cleanly before its TTL expires.
+func TestWaitForExit_ExitBeforeTTL(t *testing.T) {
+	t.Parallel()
+
+	sbx := newTestSandboxForExit(time.Now().Add(2 * time.Second))
+
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		sbx.exit.SetSuccess()
+	}()
+
+	err := sbx.WaitForExit(context.Background())
+
+	require.NoError(t, err)
+}
+
+// TestWaitForExit_ExitWithError verifies that WaitForExit returns a wrapped
+// error when the Firecracker process exits with an error before the TTL.
+func TestWaitForExit_ExitWithError(t *testing.T) {
+	t.Parallel()
+
+	sbx := newTestSandboxForExit(time.Now().Add(2 * time.Second))
+
+	fcErr := errors.New("firecracker process killed")
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		sbx.exit.SetError(fcErr)
+	}()
+
+	err := sbx.WaitForExit(context.Background())
+
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "fc process exited prematurely")
+	assert.ErrorContains(t, err, fcErr.Error())
+}
+
+// TestWaitForExit_KeepAliveExtendsTTL verifies that a KeepAlive call that
+// extends endAt via SetEndAt resets the node-side timer. Without this, the
+// sandbox would be killed at the original TTL even though it was extended.
+func TestWaitForExit_KeepAliveExtendsTTL(t *testing.T) {
+	t.Parallel()
+
+	// Short initial TTL — would fire after 80 ms without the fix.
+	sbx := newTestSandboxForExit(time.Now().Add(80 * time.Millisecond))
+
+	done := make(chan error, 1)
+	go func() {
+		done <- sbx.WaitForExit(context.Background())
+	}()
+
+	// Simulate KeepAlive: extend endAt before the original TTL fires.
+	time.Sleep(40 * time.Millisecond)
+	sbx.SetEndAt(time.Now().Add(2 * time.Second))
+
+	// Wait past the original TTL; WaitForExit must NOT have returned yet.
+	time.Sleep(100 * time.Millisecond)
+	select {
+	case err := <-done:
+		t.Fatalf("WaitForExit returned early after KeepAlive extension: %v", err)
+	default:
+		// Good — still waiting.
+	}
+
+	// Now signal a clean exit; WaitForExit should return nil.
+	sbx.exit.SetSuccess()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("WaitForExit did not return after exit signal")
+	}
+}
+
+// TestWaitForExit_ContextCancelled verifies that WaitForExit returns nil when
+// the context is cancelled before either TTL expiry or process exit.
+func TestWaitForExit_ContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	sbx := newTestSandboxForExit(time.Now().Add(10 * time.Second))
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	done := make(chan error, 1)
+	go func() {
+		done <- sbx.WaitForExit(ctx)
+	}()
+
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("WaitForExit did not return after context cancellation")
+	}
+}

packages/orchestrator/pkg/server/sandboxes.go

@@ -1041,9 +1041,12 @@ func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox
 		ctx, childSpan := tracer.Start(context.WithoutCancel(ctx), "stop sandbox-lifecycle", trace.WithNewRoot())
 		defer childSpan.End()
 
-		waitErr := sbx.Wait(ctx)
+		waitErr := sbx.WaitForExit(ctx)
 		if waitErr != nil {
-			sbxlogger.I(sbx).Error(ctx, "failed to wait for sandbox, cleaning up", zap.Error(waitErr))
+			sbxlogger.I(sbx).Error(ctx, "sandbox exit wait failed, stopping", zap.Error(waitErr))
+			if stopErr := sbx.Stop(ctx); stopErr != nil {
+				sbxlogger.I(sbx).Error(ctx, "failed to stop timed-out sandbox", zap.Error(stopErr))
+			}
 		}
 
 		cleanupErr := sbx.Close(ctx)