fix(orchestrator): refuse a memory pause when envd is unresponsive

packages/orchestrator/pkg/sandbox/pause_envd_health.go

@@ -0,0 +1,89 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"context"
+	"time"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+	"go.uber.org/zap"
+)
+
+// awaitEnvdHealthy is the envd half of the snapshot-admission pre-flight.
+//
+// A memory snapshot does not restart envd, it restores the running process
+// mid-execution. So an envd that is already unresponsive when the snapshot is
+// taken is captured in that state and faithfully replayed on every later
+// resume: resume-fc succeeds, the VM comes back, and then sandbox-wait-for-start
+// retries envd's /init until the request budget is gone. Nothing records the
+// snapshot as bad, so every retry repeats it on any node, indefinitely.
+//
+// The pause path already had two signals that envd was gone — the pre-pause
+// freeze and heap collapse both burning their full timeouts — but both are
+// best-effort by design and cannot fail a pause. This probe turns the same
+// question into an admission decision, asked BEFORE any destructive step.
+//
+// Deliberately retryable rather than terminal. An unanswered health probe is a
+// strong predictor that the snapshot would be unresumable, not a certainty:
+// sandboxes whose freeze and collapse both timed out have still gone on to
+// resume cleanly. Deferring the pause costs a retry; condemning the sandbox
+// cannot be undone.
+//
+// Returns a nil error when the pause may proceed, including when the probe is
+// disabled or cannot be run.
+// AwaitEnvdAdmission runs as its own pre-flight, before (and independently of)
+// the durable-header admission wait, so the two can be rolled out separately.
+// A negative timeout disables the probe. A nil error means the pause may
+// proceed.
+//
+// The timeout is passed in rather than read here, matching how the caller gates
+// the durable-header wait on PauseAdmissionGraceMs: the flag lookup belongs to
+// the server's feature-flag client, and keeping this method free of one leaves
+// it pure and callable from a Sandbox built without flags.
+func (s *Sandbox) AwaitEnvdAdmission(ctx context.Context, timeout time.Duration) (SnapshotAdmissionOutcome, time.Duration, error) {
+	if timeout < 0 {
+		return SnapshotAdmissionReady, 0, nil
+	}
+
+	ctx, span := tracer.Start(ctx, "envd-admission")
+	defer span.End()
+
+	return s.awaitEnvdHealthy(ctx, timeout)
+}
+
+func (s *Sandbox) awaitEnvdHealthy(ctx context.Context, timeout time.Duration) (SnapshotAdmissionOutcome, time.Duration, error) {
+	// Checks owns the probe and its HTTP client. It is stopped inside Pause,
+	// well after admission, but a sandbox torn down concurrently can leave it
+	// nil — in which case say nothing rather than refuse on missing evidence.
+	if s.Checks == nil {
+		return SnapshotAdmissionReady, 0, nil
+	}
+
+	start := time.Now()
+	healthy, err := s.Checks.getHealth(ctx, timeout)
+	waited := time.Since(start)
+
+	span := trace.SpanFromContext(ctx)
+	span.SetAttributes(
+		attribute.Bool("admission.envd_healthy", healthy),
+		attribute.Int64("admission.envd_probe_ms", waited.Milliseconds()),
+		attribute.Int64("admission.envd_probe_timeout_ms", timeout.Milliseconds()),
+	)
+
+	if healthy {
+		return SnapshotAdmissionReady, waited, nil
+	}
+
+	// The caller's context ending is not evidence about envd. Let the existing
+	// mid-wait handling own it: nothing was decided, the sandbox is untouched.
+	if ctx.Err() != nil {
+		return "", waited, ctx.Err()
+	}
+
+	s.log().Warn(ctx, "refusing pause: envd is not answering health checks",
+		zap.Error(err), zap.Duration("probe", waited))
+
+	return SnapshotAdmissionEnvdUnhealthy, waited, ErrSnapshotAdmissionEnvdUnhealthy
+}

packages/orchestrator/pkg/sandbox/pause_envd_health_test.go

@@ -0,0 +1,171 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/network"
+	"github.com/e2b-dev/infra/packages/shared/pkg/sandboxtypes"
+)
+
+// probeTimeout is what a rollout would set: generous next to the 100ms
+// monitoring probe, because here a false refusal costs a customer a pause
+// rather than a log line.
+const probeTimeout = 500 * time.Millisecond
+
+// stubRoundTripper answers every request with a canned result, so the health
+// probe can be exercised without a real guest. getHealth builds its URL from
+// the slot IP and the fixed envd port, so intercepting at the transport is the
+// only way in.
+type stubRoundTripper struct {
+	status int
+	err    error
+	calls  int
+}
+
+func (rt *stubRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	rt.calls++
+	if rt.err != nil {
+		return nil, rt.err
+	}
+
+	return &http.Response{StatusCode: rt.status, Body: http.NoBody, Request: r}, nil
+}
+
+func newHealthProbeSandbox(t *testing.T) *Sandbox {
+	t.Helper()
+
+	sbx := &Sandbox{
+		Metadata: &Metadata{
+			Config:  NewConfig(Config{}),
+			Runtime: sandboxtypes.RuntimeMetadata{SandboxID: "test-sandbox"},
+		},
+		Resources: &Resources{Slot: &network.Slot{HostIP: net.IPv4(127, 0, 0, 1)}},
+	}
+	sbx.Checks = NewChecks(sbx)
+
+	return sbx
+}
+
+func withStubTransport(t *testing.T, rt http.RoundTripper) {
+	t.Helper()
+
+	orig := sandboxHttpClient
+	sandboxHttpClient = http.Client{Transport: rt}
+	t.Cleanup(func() { sandboxHttpClient = orig })
+}
+
+// The flag defaults to -1, which must leave the pause path exactly as it was:
+// no probe, no refusal, and crucially no dial of the guest.
+//
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_DisabledByDefault(t *testing.T) {
+	rt := &stubRoundTripper{err: errors.New("envd is gone")}
+	withStubTransport(t, rt)
+
+	outcome, _, err := newHealthProbeSandbox(t).AwaitEnvdAdmission(t.Context(), -1)
+	require.NoError(t, err, "a disabled probe must never refuse a pause")
+	assert.Equal(t, SnapshotAdmissionReady, outcome)
+	assert.Zero(t, rt.calls, "a disabled probe must not dial the guest")
+}
+
+// The case this change exists for: envd does not answer, so the pause is
+// refused BEFORE any destructive step rather than producing a snapshot that
+// records a wedged agent and can never be resumed.
+//
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_UnresponsiveEnvdRefusesRetryably(t *testing.T) {
+	rt := &stubRoundTripper{err: errors.New("connection refused")}
+	withStubTransport(t, rt)
+
+	outcome, _, err := newHealthProbeSandbox(t).AwaitEnvdAdmission(t.Context(), probeTimeout)
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrSnapshotAdmissionEnvdUnhealthy)
+	assert.Equal(t, SnapshotAdmissionEnvdUnhealthy, outcome)
+	assert.Equal(t, 1, rt.calls)
+}
+
+// A healthy envd must be admitted. envd answers /health with 204; anything else
+// is a failure, so this also pins that a 200 is NOT mistaken for healthy.
+//
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_HealthyEnvdAdmits(t *testing.T) {
+	withStubTransport(t, &stubRoundTripper{status: http.StatusNoContent})
+
+	outcome, _, err := newHealthProbeSandbox(t).AwaitEnvdAdmission(t.Context(), probeTimeout)
+	require.NoError(t, err)
+	assert.Equal(t, SnapshotAdmissionReady, outcome)
+}
+
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_UnexpectedStatusRefuses(t *testing.T) {
+	withStubTransport(t, &stubRoundTripper{status: http.StatusInternalServerError})
+
+	outcome, _, err := newHealthProbeSandbox(t).AwaitEnvdAdmission(t.Context(), probeTimeout)
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrSnapshotAdmissionEnvdUnhealthy)
+	assert.Equal(t, SnapshotAdmissionEnvdUnhealthy, outcome)
+}
+
+// A sandbox torn down concurrently can leave Checks nil. Missing evidence is
+// not evidence of a wedged envd, so the pause proceeds.
+//
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_NilChecksAdmits(t *testing.T) {
+	withStubTransport(t, &stubRoundTripper{err: errors.New("envd is gone")})
+
+	sbx := newHealthProbeSandbox(t)
+	sbx.Checks = nil
+
+	outcome, _, err := sbx.AwaitEnvdAdmission(t.Context(), probeTimeout)
+	require.NoError(t, err)
+	assert.Equal(t, SnapshotAdmissionReady, outcome)
+}
+
+// A context that ends mid-probe says nothing about envd. It must surface as the
+// caller's context error with an EMPTY outcome, which the handlers map to
+// "nothing was decided, the sandbox is untouched" — never as a refusal.
+//
+//nolint:paralleltest // overrides the package-level sandboxHttpClient
+func TestAwaitEnvdHealthy_ContextCancelledIsNotARefusal(t *testing.T) {
+	withStubTransport(t, &stubRoundTripper{err: errors.New("cancelled")})
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	outcome, _, err := newHealthProbeSandbox(t).AwaitEnvdAdmission(ctx, probeTimeout)
+	require.Error(t, err)
+	require.NotErrorIs(t, err, ErrSnapshotAdmissionEnvdUnhealthy,
+		"a cancelled context must not be reported as an unhealthy envd")
+	assert.Equal(t, SnapshotAdmissionOutcome(""), outcome,
+		"an empty outcome is what tells the handler nothing was decided")
+}
+
+// The refusal must be its own sentinel and must NOT satisfy the pending one.
+// The two take different branches in the Pause handler: pending and
+// envd-unhealthy both return a retryable ResourceExhausted, while anything
+// falling through to the default case is treated as a latched error and KILLS
+// the sandbox. Conflating them would turn an unresponsive envd into a kill,
+// which is precisely what the retryable refusal exists to avoid — envd often
+// recovers, and sandboxes whose pre-pause freeze and collapse both timed out
+// have still gone on to resume cleanly.
+func TestEnvdUnhealthySentinelIsDistinctAndNotLatched(t *testing.T) {
+	t.Parallel()
+
+	require.NotErrorIs(t, ErrSnapshotAdmissionEnvdUnhealthy, ErrSnapshotAdmissionPending)
+	require.NotErrorIs(t, ErrSnapshotAdmissionPending, ErrSnapshotAdmissionEnvdUnhealthy)
+
+	// Outcome labels are metric dimensions; keep them distinct and stable.
+	assert.NotEqual(t, SnapshotAdmissionRefused, SnapshotAdmissionEnvdUnhealthy)
+	assert.NotEqual(t, SnapshotAdmissionLatchedError, SnapshotAdmissionEnvdUnhealthy)
+	assert.Equal(t, SnapshotAdmissionEnvdUnhealthy, SnapshotAdmissionOutcome("envd_unhealthy"))
+}

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -3602,12 +3602,22 @@ const (
 	// SnapshotAdmissionLatchedError: a latched seal failure means no valid
 	// snapshot can ever be produced; not retryable.
 	SnapshotAdmissionLatchedError SnapshotAdmissionOutcome = "latched_error"
+	// SnapshotAdmissionEnvdUnhealthy: envd did not answer /health, so a memory
+	// snapshot taken now would capture an unresponsive agent and never resume.
+	SnapshotAdmissionEnvdUnhealthy SnapshotAdmissionOutcome = "envd_unhealthy"
 )
 
 // ErrSnapshotAdmissionPending marks a retryable admission refusal: the parent
 // memfile header was still deduplicating when the grace elapsed.
 var ErrSnapshotAdmissionPending = errors.New("parent memfile header is still deduplicating")
 
+// ErrSnapshotAdmissionEnvdUnhealthy marks a retryable admission refusal: envd
+// did not answer its health probe, so a memory snapshot taken now would record
+// an unresponsive agent. Retryable because envd frequently recovers on its own
+// — the probe is a strong signal that the snapshot would be unresumable, not a
+// certainty, so the pause is deferred rather than the sandbox condemned.
+var ErrSnapshotAdmissionEnvdUnhealthy = errors.New("envd is not responding to health checks")
+
 // AwaitSnapshotAdmission is the pre-destructive snapshot-admission pre-flight:
 // "can this sandbox produce a valid snapshot right now?". It folds the
 // EnsurePausable latched-error checks together with the durable-parent

packages/orchestrator/pkg/server/sandboxes.go

@@ -877,6 +877,27 @@ func (s *Server) Pause(ctx context.Context, in *orchestrator.SandboxPauseRequest
 		telemetry.WithEnvdVersion(sbx.Config.Envd.Version),
 	)
 
+	// Flag-gated envd pre-flight, independent of the durable-header wait below
+	// and only for memory snapshots: a memory snapshot restores envd
+	// mid-execution, so one taken while envd is unresponsive is replayed wedged
+	// on every later resume and the sandbox never comes back.
+	if healthMs := s.featureFlags.IntFlag(ctx, featureflags.PauseEnvdHealthTimeoutMs); healthMs >= 0 && !in.GetFilesystemOnly() {
+		outcome, waited, admitErr := sbx.AwaitEnvdAdmission(ctx, time.Duration(healthMs)*time.Millisecond)
+		switch {
+		case errors.Is(admitErr, sandbox.ErrSnapshotAdmissionEnvdUnhealthy):
+			s.recordPauseAdmission(ctx, "pause", outcome, waited)
+			// Retryable, and deliberately NOT the latched/kill path below: an
+			// unanswered probe predicts an unresumable snapshot, it does not
+			// prove one, and envd often recovers on its own.
+			sbxlogger.E(sbx).Warn(ctx, "Refusing pause: envd is not answering health checks", zap.Duration("probe", waited))
+
+			return nil, status.Errorf(codes.ResourceExhausted, "sandbox '%s' guest agent is not responding, please retry", in.GetSandboxId())
+		case admitErr != nil:
+			// Context ended mid-probe: nothing decided, sandbox untouched.
+			return nil, status.FromContextError(admitErr).Err()
+		}
+	}
+
 	// Flag-gated admission pre-flight: refuse retryably BEFORE any destructive
 	// step while the parent memfile header is still deduplicating.
 	var latchedErr error
@@ -1047,6 +1068,21 @@ func (s *Server) Checkpoint(ctx context.Context, in *orchestrator.SandboxCheckpo
 		return nil, status.Errorf(codes.FailedPrecondition, "%s", err.Error())
 	}
 
+	// The same envd pre-flight as Pause. A checkpoint always takes a full
+	// memory snapshot, so it can inherit a wedged envd the same way.
+	if healthMs := s.featureFlags.IntFlag(ctx, featureflags.PauseEnvdHealthTimeoutMs); healthMs >= 0 {
+		outcome, waited, admitErr := sbx.AwaitEnvdAdmission(ctx, time.Duration(healthMs)*time.Millisecond)
+		switch {
+		case errors.Is(admitErr, sandbox.ErrSnapshotAdmissionEnvdUnhealthy):
+			s.recordPauseAdmission(ctx, "checkpoint", outcome, waited)
+			sbxlogger.E(sbx).Warn(ctx, "Refusing checkpoint: envd is not answering health checks", zap.Duration("probe", waited))
+
+			return nil, status.Errorf(codes.ResourceExhausted, "sandbox '%s' guest agent is not responding, please retry", in.GetSandboxId())
+		case admitErr != nil:
+			return nil, status.FromContextError(admitErr).Err()
+		}
+	}
+
 	// The same flag-gated admission pre-flight as Pause (a checkpoint always
 	// takes a full memory snapshot, on both the in-place and resume-fresh
 	// paths); before waitForAcquire so a grace wait never holds a start slot.

packages/shared/pkg/featureflags/flags.go

@@ -462,6 +462,24 @@ var (
 	// 0 probes the parent header's readiness without waiting; a positive value
 	// waits up to that long before refusing retryably.
 	PauseAdmissionGraceMs = NewIntFlag("pause-admission-grace-milliseconds", -1)
+	// PauseEnvdHealthTimeoutMs gates the snapshot-admission envd health probe,
+	// in milliseconds. Same sign convention as PauseAdmissionGraceMs: negative
+	// (default) disables the probe; 0 or more probes envd's /health with that
+	// timeout and refuses the pause retryably when it does not answer.
+	//
+	// A memory snapshot restores envd mid-execution rather than restarting it,
+	// so an envd that is already unresponsive when the snapshot is taken is
+	// recorded in that state and replayed on every later resume: the resume
+	// reaches envd-init, never gets an answer, and burns the whole request
+	// budget. Nothing marks the snapshot bad, so the sandbox retries forever.
+	// Probing here catches that BEFORE the destructive steps, while the pause
+	// can still be refused retryably.
+	//
+	// Pair with PauseRefusalRestoreFlag: without it a refusal still ends as a
+	// removed record and an orphan-reaped VM, so the probe only converts one
+	// bad outcome into another. With it, a refused pause keeps the sandbox
+	// running and the customer retries.
+	PauseEnvdHealthTimeoutMs = NewIntFlag("pause-envd-health-timeout-milliseconds", -1)
 	// PauseRefusalRestoreFlag gates the API-side restore of a retryably
 	// refused pause: record kept, routing re-registered, state back to
 	// Running. Off (default), a refused pause still ends today's way — the