feat(api): add sandbox fork endpoint

packages/api/internal/handlers/sandbox_fork.go

@@ -0,0 +1,204 @@
+package handlers
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"go.opentelemetry.io/otel/trace"
+	"go.uber.org/zap"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/orchestrator"
+	"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/id"
+	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	sharedUtils "github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// maxForkCount caps how many sandboxes a single fork request can create,
+// bounding the parallel boots (and the result allocation) per request.
+const maxForkCount = 100
+
+// PostSandboxesSandboxIDFork forks a running sandbox: it checkpoints the
+// sandbox in place (snapshot it and resume it on its node, so the original
+// keeps running with its ID and expiration untouched) and creates count new
+// sandboxes from that snapshot under fresh IDs. Each fork succeeds or fails
+// independently: the response carries one result per requested fork, holding
+// either the created sandbox or the error that prevented it from starting.
+func (a *APIStore) PostSandboxesSandboxIDFork(c *gin.Context, sandboxID api.SandboxID) {
+	ctx := c.Request.Context()
+
+	teamInfo := auth.MustGetTeamInfo(c)
+	teamID := teamInfo.Team.ID
+
+	sandboxID, err := utils.ShortID(sandboxID)
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, "Invalid sandbox ID")
+
+		return
+	}
+
+	span := trace.SpanFromContext(ctx)
+	span.SetAttributes(telemetry.WithSandboxID(sandboxID))
+
+	traceID := span.SpanContext().TraceID().String()
+	c.Set("traceID", traceID)
+
+	body, err := ginutils.ParseOptionalBody[api.PostSandboxesSandboxIDForkJSONRequestBody](ctx, c)
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing request: %s", err))
+
+		return
+	}
+
+	forkTimeout := sandbox.SandboxTimeoutDefault
+	if body.Timeout != nil {
+		forkTimeout = time.Duration(*body.Timeout) * time.Second
+
+		if forkTimeout > time.Duration(teamInfo.Limits.MaxLengthHours)*time.Hour {
+			a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Timeout cannot be greater than %d hours", teamInfo.Limits.MaxLengthHours))
+
+			return
+		}
+	}
+
+	forkCount := 1
+	if body.Count != nil {
+		forkCount = int(*body.Count)
+	}
+
+	if forkCount < 1 {
+		a.sendAPIStoreError(c, http.StatusBadRequest, "Count must be at least 1")
+
+		return
+	}
+
+	if forkCount > maxForkCount {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Count cannot be greater than %d", maxForkCount))
+
+		return
+	}
+
+	// The original sandbox keeps running and holds one slot, so more forks
+	// than the concurrency limit can never succeed.
+	if int64(forkCount) >= teamInfo.Limits.SandboxConcurrency {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Count must be lower than the maximum number of concurrent sandboxes (%d)", teamInfo.Limits.SandboxConcurrency))
+
+		return
+	}
+
+	original, err := a.orchestrator.GetSandbox(ctx, teamID, sandboxID)
+	if err != nil {
+		if errors.Is(err, sandbox.ErrNotFound) {
+			apiErr := forkHandleNotRunningSandbox(ctx, a, sandboxID, teamID)
+			a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
+
+			return
+		}
+
+		telemetry.ReportError(ctx, "error getting sandbox for fork", err, telemetry.WithSandboxID(sandboxID))
+		a.sendAPIStoreError(c, http.StatusInternalServerError, "Error forking sandbox")
+
+		return
+	}
+
+	if err := sharedUtils.CheckEnvdVersionForSnapshot(original.EnvdVersion); err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, err.Error())
+
+		return
+	}
+
+	// Checkpoint the sandbox in place: it is briefly paused on its node,
+	// snapshotted, and resumed under the same execution ID, so the original
+	// keeps its ID, expiration, and concurrency slot.
+	err = a.orchestrator.CheckpointSandbox(ctx, teamID, sandboxID)
+	var transErr *sandbox.InvalidStateTransitionError
+
+	switch {
+	case err == nil:
+	case errors.Is(err, sandbox.ErrNotFound):
+		apiErr := forkHandleNotRunningSandbox(ctx, a, sandboxID, teamID)
+		a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
+
+		return
+	case errors.As(err, &transErr):
+		a.sendAPIStoreError(c, http.StatusConflict, fmt.Sprintf("Sandbox '%s' cannot be forked while in '%s' state", sandboxID, transErr.CurrentState))
+
+		return
+	case errors.Is(err, orchestrator.PauseQueueExhaustedError{}):
+		a.sendAPIStoreError(c, http.StatusServiceUnavailable, fmt.Sprintf("Sandbox '%s' cannot be forked right now because its node is busy, please retry", sandboxID))
+
+		return
+	default:
+		telemetry.ReportError(ctx, "error checkpointing sandbox for fork", err, telemetry.WithSandboxID(sandboxID))
+		a.sendAPIStoreError(c, http.StatusInternalServerError, "Error forking sandbox")
+
+		return
+	}
+
+	sbxlogger.E(&sbxlogger.SandboxMetadata{
+		SandboxID:  sandboxID,
+		TemplateID: original.TemplateID,
+		TeamID:     teamID.String(),
+	}).Debug(ctx, "Creating forked sandboxes from snapshot", zap.Int("count", forkCount))
+
+	// All forks boot in parallel from the same immutable snapshot, each
+	// succeeding or failing independently.
+	results := make([]api.SandboxForkResult, forkCount)
+
+	wg := errgroup.Group{}
+	for i := range forkCount {
+		wg.Go(func() error {
+			forkedSandboxID := InstanceIDPrefix + id.Generate()
+
+			forkedSbx, createErr := a.startSandbox(
+				ctx,
+				forkedSandboxID,
+				forkTimeout,
+				teamInfo,
+				a.buildResumeSandboxDataFromSnapshot(sandboxID, forkedSandboxID, nil),
+				&c.Request.Header,
+				true,
+				nil, // mcp
+			)
+			if createErr != nil {
+				telemetry.ReportError(ctx, "error creating forked sandbox", createErr.Err, telemetry.WithSandboxID(forkedSandboxID))
+				results[i] = api.SandboxForkResult{Error: &api.Error{Code: int32(createErr.Code), Message: createErr.ClientMsg}}
+
+				//nolint:nilerr // per-fork errors are reported in the result entry, not propagated
+				return nil
+			}
+
+			results[i] = api.SandboxForkResult{Sandbox: forkedSbx}
+
+			return nil
+		})
+	}
+	_ = wg.Wait()
+
+	c.JSON(http.StatusCreated, results)
+}
+
+// forkHandleNotRunningSandbox classifies a fork request for a sandbox that is
+// not running: 409 if it is paused (a snapshot exists), 404 otherwise.
+func forkHandleNotRunningSandbox(ctx context.Context, a *APIStore, sandboxID string, teamID uuid.UUID) api.APIError {
+	apiErr := pauseHandleNotRunningSandbox(ctx, a.snapshotCache, sandboxID, teamID)
+	switch apiErr.Code {
+	case http.StatusConflict:
+		apiErr.ClientMsg = fmt.Sprintf("Sandbox '%s' is paused and cannot be forked; resume it first", sandboxID)
+	case http.StatusInternalServerError:
+		apiErr.ClientMsg = "Error forking sandbox"
+	}
+
+	return apiErr
+}

packages/api/internal/handlers/sandbox_resume.go

@@ -186,17 +186,25 @@ func convertDatabaseMountsToOrchestratorMounts(volumes []*types.SandboxVolumeMou
 	return results
 }
 
-// buildResumeSandboxData returns a SandboxDataFetcher that fetches snapshot data
-// from the cache and builds SandboxMetadata for resume operations.
-// The returned callback is called inside the sandbox lock to prevent race conditions.
+// buildResumeSandboxData returns a SandboxDataFetcher for resuming a sandbox
+// from its own snapshot.
 func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *bool) orchestrator.SandboxDataFetcher {
+	return a.buildResumeSandboxDataFromSnapshot(sandboxID, sandboxID, autoPauseOverride)
+}
+
+// buildResumeSandboxDataFromSnapshot returns a SandboxDataFetcher that fetches
+// snapshot data for snapshotSandboxID from the cache and builds SandboxMetadata
+// for resume operations. sandboxID is the ID the sandbox will run under — it
+// differs from snapshotSandboxID when forking — and scopes the envd access token.
+// The returned callback is called inside the sandbox lock to prevent race conditions.
+func (a *APIStore) buildResumeSandboxDataFromSnapshot(snapshotSandboxID, sandboxID string, autoPauseOverride *bool) orchestrator.SandboxDataFetcher {
 	return func(ctx context.Context) (orchestrator.SandboxMetadata, *api.APIError) {
-		lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID)
+		lastSnapshot, err := a.snapshotCache.Get(ctx, snapshotSandboxID)
 		if err != nil {
 			return orchestrator.SandboxMetadata{}, &api.APIError{
 				Code:      http.StatusInternalServerError,
 				ClientMsg: "Error when getting snapshot",
-				Err:       fmt.Errorf("error getting last snapshot for sandbox '%s': %w", sandboxID, err),
+				Err:       fmt.Errorf("error getting last snapshot for sandbox '%s': %w", snapshotSandboxID, err),
 			}
 		}
 
@@ -253,6 +261,7 @@ func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *b
 			VolumeMounts:            convertDatabaseMountsToOrchestratorMounts(volumes),
 			EnvdAccessToken:         envdAccessToken,
 			NodeID:                  &nodeID,
+			SnapshotSandboxID:       snapshotSandboxID,
 		}, nil
 	}
 }

packages/api/internal/orchestrator/checkpoint_instance.go

@@ -0,0 +1,130 @@
+package orchestrator
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gogo/status"
+	"github.com/google/uuid"
+	"google.golang.org/grpc/codes"
+
+	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
+	"github.com/e2b-dev/infra/packages/db/pkg/types"
+	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/storageopts"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+// CheckpointSandbox snapshots a running sandbox in place: the sandbox is
+// briefly paused on its node, snapshotted, and resumed under the same
+// execution ID, so it keeps running with its ID, expiration, and reservation
+// untouched. The snapshot is written to the sandbox's own snapshots row, so
+// it can immediately be resumed or forked from.
+func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, sandboxID string) error {
+	ctx, span := tracer.Start(ctx, "checkpoint-sandbox")
+	defer span.End()
+
+	sbx, alreadyDone, finishSnapshotting, err := o.sandboxStore.StartRemoving(ctx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionSnapshot})
+	if err != nil {
+		return fmt.Errorf("failed to start snapshotting: %w", err)
+	}
+
+	// alreadyDone conflates joining a concurrent checkpoint that just
+	// succeeded with finding the sandbox stuck in Snapshotting after a failed
+	// one, where no fresh snapshot exists. Treating it as success could fork
+	// stale state, so report a conflict and let the caller retry (same
+	// behavior as CreateSnapshotTemplate).
+	if alreadyDone {
+		return &sandbox.InvalidStateTransitionError{
+			CurrentState: sandbox.StateSnapshotting,
+			TargetState:  sandbox.StateSnapshotting,
+		}
+	}
+
+	// finish completes the snapshotting transition exactly once.
+	// On success (nil) it restores the sandbox to Running.
+	// On error it leaves the state as Snapshotting so that
+	// RemoveSandbox can transition directly to Killing.
+	var once sync.Once
+	finish := func(err error) {
+		once.Do(func() {
+			finishSnapshotting(context.WithoutCancel(ctx), err)
+		})
+	}
+	defer finish(nil)
+
+	node := o.getOrConnectNode(ctx, sbx.ClusterID, sbx.NodeID)
+	if node == nil {
+		return fmt.Errorf("node '%s' not found", sbx.NodeID)
+	}
+
+	upsertResult, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node, false))
+	if err != nil {
+		return fmt.Errorf("error upserting snapshot: %w", err)
+	}
+
+	// Checkpoint pauses the sandbox, snapshots it, and resumes it on the
+	// orchestrator with the same ExecutionID. Once the pause has started, the
+	// orchestrator stops the old sandbox itself on error; RemoveSandbox is
+	// still needed to clean up API-side state (store, routing, analytics).
+	client, childCtx := node.GetClient(ctx)
+	_, err = client.Sandbox.Checkpoint(childCtx, &orchestrator.SandboxCheckpointRequest{
+		SandboxId: sbx.SandboxID,
+		BuildId:   upsertResult.BuildID.String(),
+		Metadata:  map[string]string{storageopts.ObjectMetadataTemplateID: upsertResult.TemplateID},
+	})
+	if err != nil {
+		// Cleanup must run even when the checkpoint failed because this
+		// request's context was cancelled (e.g. client disconnect mid-fork).
+		cleanupCtx := context.WithoutCancel(ctx)
+
+		o.failSnapshotBuild(cleanupCtx, upsertResult.BuildID, err)
+
+		// The orchestrator rejects these before pausing the VM (envd too old,
+		// starting-sandboxes queue full), so the sandbox is still running
+		// healthy on its node: restore it to Running instead of killing it.
+		if st, ok := status.FromError(err); ok {
+			switch st.Code() {
+			case codes.FailedPrecondition:
+				finish(nil)
+
+				return fmt.Errorf("checkpoint rejected: %w", err)
+			case codes.ResourceExhausted:
+				finish(nil)
+
+				return PauseQueueExhaustedError{}
+			}
+		}
+
+		// Complete the snapshotting transition with error — leaves state as
+		// Snapshotting (no restore to Running) and clears the transition key
+		// so RemoveSandbox can proceed without deadlock.
+		finish(err)
+
+		if killErr := o.RemoveSandbox(cleanupCtx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionKill}); killErr != nil {
+			telemetry.ReportError(cleanupCtx, "error killing sandbox after failed checkpoint", killErr)
+		}
+
+		return fmt.Errorf("checkpoint failed: %w", err)
+	}
+
+	now := time.Now()
+	err = o.sqlcDB.UpdateEnvBuildStatus(ctx, queries.UpdateEnvBuildStatusParams{
+		Status:     types.BuildStatusSuccess,
+		FinishedAt: &now,
+		Reason:     types.BuildReason{},
+		BuildID:    upsertResult.BuildID,
+	})
+	if err != nil {
+		return fmt.Errorf("error updating build status: %w", err)
+	}
+
+	o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID)
+
+	telemetry.ReportEvent(ctx, "Checkpointed sandbox")
+
+	return nil
+}

packages/api/internal/orchestrator/create_instance.go

@@ -54,6 +54,9 @@ type SandboxMetadata struct {
 	VolumeMounts            []*orchestrator.SandboxVolumeMount
 	EnvdAccessToken         *string
 	NodeID                  *string
+	// SnapshotSandboxID is the sandbox ID the resume snapshot is stored under.
+	// It differs from the ID of the sandbox being started when forking.
+	SnapshotSandboxID string
 }
 
 // buildEgressConfig constructs the orchestrator egress configuration from
@@ -322,7 +325,13 @@ func (o *Orchestrator) CreateSandbox(
 	placed, err := placement.PlaceSandbox(ctx, o.placementAlgorithm, clusterNodes, node, sbxRequest, builds.ToMachineInfo(sbxData.Build), labelFilteringEnabled, allLabels)
 	if err != nil {
 		if isResume && placed.TimedOut {
-			o.maybeRemapResumeOriginNode(ctx, sandboxID, team, sbxData.NodeID, placed.WarmedNode)
+			// Remap by the snapshot's own sandbox ID: when forking, the started
+			// sandbox ID has no snapshot row and the remap would silently no-op.
+			snapshotSandboxID := sandboxID
+			if sbxData.SnapshotSandboxID != "" {
+				snapshotSandboxID = sbxData.SnapshotSandboxID
+			}
+			o.maybeRemapResumeOriginNode(ctx, snapshotSandboxID, team, sbxData.NodeID, placed.WarmedNode)
 		}
 
 		return sandbox.Sandbox{}, &api.APIError{