Add rootfs reboot fallback for snapshots

packages/api/internal/handlers/proxy_grpc.go

@@ -257,6 +257,7 @@ func (s *SandboxService) ResumeSandbox(ctx context.Context, req *proxygrpc.Sandb
 		s.api.buildResumeSandboxData(sandboxID, nil),
 		&headers,
 		true,
+		false,
 		nil, // mcp
 	)
 	if apiErr != nil {

packages/api/internal/handlers/sandbox.go

@@ -28,6 +28,7 @@ func (a *APIStore) startSandbox(
 	getSandboxData orchestrator.SandboxDataFetcher,
 	requestHeader *http.Header,
 	isResume bool,
+	rebootFromRootfs bool,
 	mcp api.Mcp,
 ) (*api.Sandbox, *api.APIError) {
 	sbx, apiErr := a.startSandboxInternal(
@@ -38,6 +39,7 @@ func (a *APIStore) startSandbox(
 		getSandboxData,
 		requestHeader,
 		isResume,
+		rebootFromRootfs,
 		mcp,
 	)
 	if apiErr != nil {
@@ -56,6 +58,7 @@ func (a *APIStore) startSandboxInternal(
 	getSandboxData orchestrator.SandboxDataFetcher,
 	requestHeader *http.Header,
 	isResume bool,
+	rebootFromRootfs bool,
 	mcp api.Mcp,
 ) (sandbox.Sandbox, *api.APIError) {
 	startTime := time.Now()
@@ -77,6 +80,7 @@ func (a *APIStore) startSandboxInternal(
 		timeout,
 		isResume,
 		creationMeta,
+		rebootFromRootfs,
 	)
 	if instanceErr != nil {
 		telemetry.ReportError(ctx, "error when creating instance", instanceErr.Err)

packages/api/internal/handlers/sandbox_connect.go

@@ -48,6 +48,7 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
 
 		return
 	}
+	rebootFromRootfs := body.Reboot != nil && *body.Reboot
 
 	teamID := teamInfo.Team.ID
 
@@ -152,6 +153,7 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
 		a.buildResumeSandboxData(sandboxID, nil),
 		&c.Request.Header,
 		true,
+		rebootFromRootfs,
 		nil, // mcp
 	)
 	if createErr != nil {

packages/api/internal/handlers/sandbox_create.go

@@ -265,6 +265,7 @@ func (a *APIStore) PostSandboxes(c *gin.Context) {
 		getSandboxData,
 		&c.Request.Header,
 		false,
+		false,
 		mcp,
 	)
 	if createErr != nil {

packages/api/internal/handlers/sandbox_pause.go

@@ -18,10 +18,15 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/auth/pkg/auth"
+	"github.com/e2b-dev/infra/packages/shared/pkg/ginutils"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+type pauseSandboxBody struct {
+	Memory *bool `json:"memory,omitempty"`
+}
+
 func (a *APIStore) PostSandboxesSandboxIDPause(c *gin.Context, sandboxID api.SandboxID) {
 	ctx := c.Request.Context()
 	// Get team from context, use TeamContextKey
@@ -42,9 +47,22 @@ func (a *APIStore) PostSandboxesSandboxIDPause(c *gin.Context, sandboxID api.San
 	traceID := span.SpanContext().TraceID().String()
 	c.Set("traceID", traceID)
 
+	memory := true
+	if c.Request.ContentLength != 0 {
+		body, err := ginutils.ParseBody[pauseSandboxBody](ctx, c)
+		if err != nil {
+			a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing request: %s", err))
+
+			return
+		}
+		if body.Memory != nil {
+			memory = *body.Memory
+		}
+	}
+
 	pause.LogInitiated(ctx, sandboxID, teamID.String(), pause.ReasonRequest)
 
-	err = a.orchestrator.RemoveSandbox(ctx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionPause})
+	err = a.orchestrator.RemoveSandbox(ctx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionPause, SkipMemory: !memory})
 	var transErr *sandbox.InvalidStateTransitionError
 
 	switch {

packages/api/internal/handlers/sandbox_resume.go

@@ -54,6 +54,7 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
 	}
 
 	telemetry.ReportEvent(ctx, "Parsed body")
+	rebootFromRootfs := body.Reboot != nil && *body.Reboot
 
 	timeout := sandbox.SandboxTimeoutDefault
 	if body.Timeout != nil {
@@ -160,6 +161,7 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
 		a.buildResumeSandboxData(sandboxID, body.AutoPause),
 		&c.Request.Header,
 		true,
+		rebootFromRootfs,
 		nil, // mcp
 	)
 	if createErr != nil {

packages/api/internal/handlers/snapshot_template_create.go

@@ -59,6 +59,9 @@ func (a *APIStore) PostSandboxesSandboxIDSnapshots(c *gin.Context, sandboxID api
 	opts := orchestrator.SnapshotTemplateOpts{
 		Tag: id.DefaultTag,
 	}
+	if body.Memory != nil {
+		opts.SkipMemory = !*body.Memory
+	}
 
 	if body.Name != nil {
 		identifier, tag, err := id.ParseName(*body.Name)

packages/api/internal/orchestrator/create_instance.go

@@ -10,6 +10,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 	"go.uber.org/zap"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
@@ -127,6 +128,7 @@ func (o *Orchestrator) CreateSandbox(
 	timeout time.Duration,
 	isResume bool,
 	creationMeta sandbox.CreationMetadata,
+	rebootFromRootfsOpt ...bool,
 ) (sbx sandbox.Sandbox, apiErr *api.APIError) {
 	ctx, childSpan := tracer.Start(ctx, "create-sandbox")
 	defer childSpan.End()
@@ -252,6 +254,10 @@ func (o *Orchestrator) CreateSandbox(
 			TimeoutSeconds: sbxData.AutoResume.Timeout,
 		}
 	}
+	rebootFromRootfs := len(rebootFromRootfsOpt) > 0 && rebootFromRootfsOpt[0]
+	if rebootFromRootfs {
+		ctx = metadata.AppendToOutgoingContext(ctx, orchestrator.SandboxRebootFromRootfsGRPCMetadataKey, "true")
+	}
 
 	sbxRequest := &orchestrator.SandboxCreateRequest{
 		Sandbox: &orchestrator.SandboxConfig{
@@ -365,12 +371,10 @@ func (o *Orchestrator) CreateSandbox(
 		// Copy to a new variable to avoid race conditions
 		sbxToRemove := sbx
 		go func() {
-			killErr := o.removeSandboxFromNode(
-				context.WithoutCancel(ctx),
-				sbxToRemove,
-				sandbox.StateActionKill,
-				sandbox.KillReasonUnknown,
-			)
+			killErr := o.removeSandboxFromNode(context.WithoutCancel(ctx), sbxToRemove, sandbox.RemoveOpts{
+				Action: sandbox.StateActionKill,
+				Reason: sandbox.KillReasonUnknown,
+			})
 			if killErr != nil {
 				logger.L().Error(ctx, "Error removing sandbox",
 					zap.Error(killErr),

packages/api/internal/orchestrator/delete_instance.go

@@ -104,7 +104,7 @@ func (o *Orchestrator) RemoveSandbox(ctx context.Context, teamID uuid.UUID, sand
 	defer func() { go o.analyticsRemove(context.WithoutCancel(ctx), sbx, opts.Action) }()
 	// Once we start the removal process, we want to make sure it gets removed from the store
 	defer o.sandboxStore.Remove(context.WithoutCancel(ctx), teamID, sandboxID)
-	err = o.removeSandboxFromNode(ctx, sbx, opts.Action, opts.Reason)
+	err = o.removeSandboxFromNode(ctx, sbx, opts)
 	if err != nil {
 		fields := []zap.Field{
 			zap.String("state_action", opts.Action.Name),
@@ -123,12 +123,7 @@ func (o *Orchestrator) RemoveSandbox(ctx context.Context, teamID uuid.UUID, sand
 	return nil
 }
 
-func (o *Orchestrator) removeSandboxFromNode(
-	ctx context.Context,
-	sbx sandbox.Sandbox,
-	stateAction sandbox.StateAction,
-	reason sandbox.KillReason,
-) error {
+func (o *Orchestrator) removeSandboxFromNode(ctx context.Context, sbx sandbox.Sandbox, opts sandbox.RemoveOpts) error {
 	ctx, span := tracer.Start(ctx, "remove-sandbox-from-node")
 	defer span.End()
 
@@ -137,8 +132,8 @@ func (o *Orchestrator) removeSandboxFromNode(
 		fields := []zap.Field{
 			logger.WithNodeID(sbx.NodeID),
 		}
-		if stateAction == sandbox.StateActionKill {
-			fields = append(fields, zap.String("kill_reason", reason.String()))
+		if opts.Action == sandbox.StateActionKill {
+			fields = append(fields, zap.String("kill_reason", opts.Reason.String()))
 		}
 
 		logger.L().Error(ctx, "failed to get node", fields...)
@@ -156,8 +151,8 @@ func (o *Orchestrator) removeSandboxFromNode(
 				zap.Error(err),
 				logger.WithSandboxID(sbx.SandboxID),
 			}
-			if stateAction == sandbox.StateActionKill {
-				fields = append(fields, zap.String("kill_reason", reason.String()))
+			if opts.Action == sandbox.StateActionKill {
+				fields = append(fields, zap.String("kill_reason", opts.Reason.String()))
 			}
 
 			logger.L().Error(ctx, "error removing routing record from catalog", fields...)
@@ -166,12 +161,12 @@ func (o *Orchestrator) removeSandboxFromNode(
 
 	sbxlogger.I(sbx).Debug(ctx, "Removing sandbox",
 		zap.Bool("auto_pause", sbx.AutoPause),
-		zap.String("state_action", stateAction.Name),
+		zap.String("state_action", opts.Action.Name),
 	)
 
-	switch stateAction {
+	switch opts.Action {
 	case sandbox.StateActionPause:
-		err := o.pauseSandbox(ctx, node, sbx)
+		err := o.pauseSandbox(ctx, node, sbx, opts.SkipMemory)
 		if err != nil {
 			if dberrors.IsForeignKeyViolation(err) {
 				killErr := o.killSandboxOnNode(ctx, node, sbx, sandbox.KillReasonBaseTemplateMissing)
@@ -191,7 +186,7 @@ func (o *Orchestrator) removeSandboxFromNode(
 
 		return nil
 	case sandbox.StateActionKill:
-		return o.killSandboxOnNode(ctx, node, sbx, reason)
+		return o.killSandboxOnNode(ctx, node, sbx, opts.Reason)
 	}
 
 	return nil

packages/api/internal/orchestrator/pause_instance.go

@@ -10,6 +10,7 @@
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5/pgtype"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
@@ -26,18 +27,18 @@
 	return "The pause queue is exhausted"
 }
 
-func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox) error {
+func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox, skipMemory bool) error {
 	ctx, span := tracer.Start(ctx, "pause-sandbox")
 	defer span.End()
 
-	result, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node))
+	result, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node, skipMemory))
 	if err != nil {
 		telemetry.ReportCriticalError(ctx, "error inserting snapshot for env", err)
 
 		return err
 	}
 
-	err = snapshotInstance(ctx, node, sbx, result.TemplateID, result.BuildID.String())
+	err = snapshotInstance(ctx, node, sbx, result.TemplateID, result.BuildID.String(), skipMemory)
 	if errors.Is(err, PauseQueueExhaustedError{}) {
 		telemetry.ReportCriticalError(ctx, "pause queue exhausted", err)
 
@@ -68,11 +69,14 @@
 	return nil
 }
 
-func snapshotInstance(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox, templateID, buildID string) error {
+func snapshotInstance(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox, templateID, buildID string, skipMemory bool) error {
 	childCtx, childSpan := tracer.Start(ctx, "snapshot-instance")
 	defer childSpan.End()
 
 	client, childCtx := node.GetSandboxDeleteCtx(childCtx, sbx.SandboxID, sbx.ExecutionID)
+	if skipMemory {
+		childCtx = metadata.AppendToOutgoingContext(childCtx, orchestrator.SandboxMemorySnapshotGRPCMetadataKey, "false")
+	}
 	_, err := client.Sandbox.Pause(
 		childCtx, &orchestrator.SandboxPauseRequest{
 			SandboxId:  sbx.SandboxID,
@@ -103,15 +107,14 @@
 	return o.sandboxStore.WaitForStateChange(ctx, teamID, sandboxID)
 }
 
-func buildUpsertSnapshotParams(sbx sandbox.Sandbox, node *nodemanager.Node) queries.UpsertSnapshotParams {
+func buildUpsertSnapshotParams(sbx sandbox.Sandbox, node *nodemanager.Node, skipMemory bool) queries.UpsertSnapshotParams {
 	machineInfo := node.MachineInfo()
 
 	metadata := types.JSONBStringMap(sbx.Metadata)
 	if metadata == nil {
 		metadata = types.JSONBStringMap{}
 	}
-
 	return queries.UpsertSnapshotParams{
 		// Used if there's no snapshot for this sandbox yet
 		TemplateID:     id.Generate(),
 		TeamID:         sbx.TeamID,

packages/api/internal/orchestrator/snapshot_template.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
@@ -30,6 +31,8 @@ type SnapshotTemplateOpts struct {
 	Namespace *string
 	// Tag is the build tag parsed from the name, defaults to "default".
 	Tag string
+	// SkipMemory controls whether the snapshot template should omit VM memory.
+	SkipMemory bool
 }
 
 // CreateSnapshotTemplate creates a persistent snapshot template from a running sandbox and immediately resumes it.
@@ -67,7 +70,7 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U
 		return SnapshotTemplateResult{}, fmt.Errorf("node '%s' not found", sbx.NodeID)
 	}
 
-	upsertResult, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node))
+	upsertResult, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node, opts.SkipMemory))
 	if err != nil {
 		return SnapshotTemplateResult{}, fmt.Errorf("error upserting snapshot: %w", err)
 	}
@@ -84,6 +87,9 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U
 	// kills the sandbox itself; RemoveSandbox is still needed to clean up
 	// API-side state (store, routing, analytics).
 	client, childCtx := node.GetClient(ctx)
+	if opts.SkipMemory {
+		childCtx = metadata.AppendToOutgoingContext(childCtx, orchestrator.SandboxMemorySnapshotGRPCMetadataKey, "false")
+	}
 	_, err = client.Sandbox.Checkpoint(childCtx, &orchestrator.SandboxCheckpointRequest{
 		SandboxId: sbx.SandboxID,
 		BuildId:   upsertResult.BuildID.String(),

packages/api/internal/sandbox/sandboxtypes/states.go

@@ -70,9 +70,10 @@ func (r KillReason) String() string {
 
 // RemoveOpts bundles the parameters that control sandbox removal.
 type RemoveOpts struct {
-	Action   StateAction
-	Eviction bool
-	Reason   KillReason
+	Action     StateAction
+	Eviction   bool
+	SkipMemory bool
+	Reason     KillReason
 }
 
 var AllowedTransitions = map[State]map[State]bool{