feat(api): add in-progress and runtime-sandbox roots to the template-deletion guard

packages/api/internal/handlers/sandbox_kill.go

@@ -33,6 +33,10 @@ func (a *APIStore) deleteSnapshot(ctx context.Context, sandboxID string, teamID
 	a.templateCache.InvalidateAliasesByTemplateID(context.WithoutCancel(ctx), snapshot.TemplateID, aliasKeys)
 	a.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID)
 
+	// Delete GCS/S3 artifacts in the background; DB deletion is already committed
+	// so artifact cleanup failures must not block or fail the response.
+	go a.deleteTemplateArtifacts(context.WithoutCancel(ctx), teamID, snapshot.TemplateID)
+
 	return nil
 }
 

packages/api/internal/handlers/sandbox_kill_test.go

@@ -0,0 +1,102 @@
+package handlers
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	templatecache "github.com/e2b-dev/infra/packages/api/internal/cache/templates"
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	dbtypes "github.com/e2b-dev/infra/packages/db/pkg/types"
+	"github.com/e2b-dev/infra/packages/db/queries"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
+)
+
+// createTestSnapshotForHandler inserts a snapshot and returns (sandboxID, templateID).
+// Uses UpsertSnapshot which also creates the env and env_build_assignment rows that
+// GetExclusiveBuildsForTemplateDeletion relies on.
+func createTestSnapshotForHandler(t *testing.T, db *testutils.Database, teamID uuid.UUID, baseEnvID string) (string, string) {
+	t.Helper()
+
+	sandboxID := "sbx-" + uuid.New().String()[:8]
+	templateID := "env-" + uuid.New().String()
+	envdVersion := "v1.0.0"
+	totalDisk := int64(1024)
+	allowInternet := true
+
+	sourceBuildID := testutils.CreateTestBuild(t, t.Context(), db, baseEnvID, "success")
+
+	result, err := db.SqlcClient.UpsertSnapshot(t.Context(), queries.UpsertSnapshotParams{
+		TemplateID:          templateID,
+		TeamID:              teamID,
+		SandboxID:           sandboxID,
+		BaseTemplateID:      baseEnvID,
+		StartedAt:           pgtype.Timestamptz{Time: time.Now(), Valid: true},
+		Vcpu:                2,
+		RamMb:               2048,
+		TotalDiskSizeMb:     &totalDisk,
+		Metadata:            dbtypes.JSONBStringMap{},
+		KernelVersion:       "6.1.0",
+		FirecrackerVersion:  "1.4.0",
+		EnvdVersion:         &envdVersion,
+		Secure:              true,
+		AllowInternetAccess: &allowInternet,
+		AutoPause:           false,
+		OriginNodeID:        "test-node",
+		SourceBuildID:       sourceBuildID,
+		Status:              "success",
+	})
+	require.NoError(t, err)
+
+	return sandboxID, result.TemplateID
+}
+
+// TestDeleteTemplateArtifacts_SnapshotBuildIsDeleted verifies that the snapshot
+// path (paused sandbox deletion) also triggers artifact cleanup.
+// UpsertSnapshot creates an env + env_build_assignment row, so the SQL query
+// picks up the build and the deleter is called.
+func TestDeleteTemplateArtifacts_SnapshotBuildIsDeleted(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	redis := redis_utils.SetupInstance(t)
+	ctx := t.Context()
+
+	teamID := testutils.CreateTestTeam(t, db)
+	baseEnvID := testutils.CreateTestTemplate(t, db, teamID)
+	_, snapshotTemplateID := createTestSnapshotForHandler(t, db, teamID, baseEnvID)
+
+	// Simulate what deleteSnapshot does: soft-delete then call the artifact cleanup.
+	err := db.SqlcClient.TestsRawSQL(ctx,
+		`UPDATE public.envs SET deleted_at = NOW() WHERE id = $1`,
+		snapshotTemplateID,
+	)
+	require.NoError(t, err)
+
+	store := &APIStore{
+		sqlcDB:           db.SqlcClient,
+		templateCache:    templatecache.NewTemplateCache(db.SqlcClient, redis),
+		buildAncestors:   stubAncestorReader{},
+		runtimeSandboxes: stubSandboxLister{},
+	}
+
+	var mu sync.Mutex
+	deleted := []uuid.UUID{}
+	deleteFn := func(_ context.Context, buildID uuid.UUID, _ string, _ uuid.UUID, _ string) error {
+		mu.Lock()
+		deleted = append(deleted, buildID)
+		mu.Unlock()
+
+		return nil
+	}
+
+	store.deleteTemplateArtifactsWithDeleter(ctx, teamID, snapshotTemplateID, deleteFn)
+
+	assert.Len(t, deleted, 1, "snapshot build should be passed to the deleter")
+}

packages/api/internal/handlers/store.go

@@ -88,6 +88,14 @@ type APIStore struct {
 	snapshotUpsertSem     *sharedutils.AdjustableSemaphore
 	sandboxListSem        *sharedutils.AdjustableSemaphore
 	snapshotBuildQuerySem *sharedutils.AdjustableSemaphore
+	// buildAncestors resolves a build's header-mapping ancestor closure for the
+	// template-deletion reachability guard. Nil when template storage is not
+	// configured for the API, in which case artifact cleanup fails closed.
+	buildAncestors buildAncestorReader
+	// runtimeSandboxes lists the team's live sandboxes so their pinned builds
+	// are protected roots during template-deletion cleanup. Nil makes the guard
+	// fail closed.
+	runtimeSandboxes runtimeSandboxLister
 }
 
 func NewAPIStore(ctx context.Context, tel *telemetry.Client, redisClient redis.UniversalClient, featureFlags *featureflags.Client, config cfg.Config) *APIStore {
@@ -286,6 +294,8 @@ func NewAPIStore(ctx context.Context, tel *telemetry.Client, redisClient redis.U
 		snapshotUpsertSem:     snapshotUpsertSem,
 		sandboxListSem:        sandboxListSem,
 		snapshotBuildQuerySem: snapshotBuildQuerySem,
+		buildAncestors:        newTemplateBuildAncestorReader(ctx),
+		runtimeSandboxes:      orch,
 	}
 
 	go a.updateDBThrottleLimits(ctx)

packages/api/internal/handlers/template_delete.go

@@ -8,11 +8,13 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel/attribute"
+	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/db/pkg/dberrors"
 	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/clusters"
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -54,6 +56,180 @@ func (a *APIStore) softDeleteTemplate(ctx context.Context, teamID uuid.UUID, tem
 	return aliasKeys, nil
 }
 
+type deleteBuildFn func(ctx context.Context, buildID uuid.UUID, templateID string, clusterID uuid.UUID, nodeID string) error
+
+// deleteTemplateArtifacts removes GCS/S3 build artifacts for builds that are
+// exclusively owned by the given template. It is called after the template has
+// been soft-deleted from the database; failures are logged but do not affect the
+// HTTP response because the DB deletion is already committed.
+func (a *APIStore) deleteTemplateArtifacts(ctx context.Context, teamID uuid.UUID, templateID string) {
+	a.deleteTemplateArtifactsWithDeleter(ctx, teamID, templateID, a.templateManager.DeleteBuild)
+}
+
+func (a *APIStore) deleteTemplateArtifactsWithDeleter(ctx context.Context, teamID uuid.UUID, templateID string, deleteFn deleteBuildFn) {
+	builds, err := a.sqlcDB.GetExclusiveBuildsForTemplateDeletion(ctx, templateID)
+	if err != nil {
+		logger.L().Error(ctx, "Failed to query exclusive builds for artifact cleanup",
+			logger.WithTemplateID(templateID),
+			zap.Error(err),
+		)
+
+		return
+	}
+
+	// A build being "exclusive" in the database (no env_build_assignments row
+	// for any other template) does NOT make it safe to delete: intermediate
+	// cache layers are never written to the DB, so a build that is DB-exclusive
+	// to this template can still be a header-mapping ancestor of another live
+	// template's build. Deleting its storage prefix would orphan that ancestor
+	// and surface later as a silent, unrecoverable "Input/output error" when
+	// the other template's image is reconstructed. Guard against that by only
+	// deleting candidates that are NOT reachable from a live root's header
+	// closure; fail closed (delete nothing) if reachability cannot be proven.
+	candidateIDs := make([]uuid.UUID, 0, len(builds))
+	for _, b := range builds {
+		candidateIDs = append(candidateIDs, b.BuildID)
+	}
+
+	deletable, err := a.deletableAfterReachabilityGuard(ctx, teamID, templateID, candidateIDs)
+	if err != nil {
+		logger.L().Warn(ctx, "Skipping template artifact cleanup: could not prove build reachability",
+			logger.WithTemplateID(templateID),
+			logger.WithTeamID(teamID.String()),
+			zap.Error(err),
+		)
+
+		return
+	}
+
+	for _, b := range builds {
+		if b.ClusterNodeID == nil {
+			continue
+		}
+
+		if _, ok := deletable[b.BuildID]; !ok {
+			logger.L().Warn(ctx, "Keeping build: still referenced by a live template's header closure",
+				logger.WithTemplateID(templateID),
+				logger.WithBuildID(b.BuildID.String()),
+			)
+
+			continue
+		}
+
+		clusterID := clusters.WithClusterFallback(b.ClusterID)
+		if err := deleteFn(ctx, b.BuildID, templateID, clusterID, *b.ClusterNodeID); err != nil {
+			logger.L().Warn(ctx, "Failed to delete build artifacts",
+				logger.WithTemplateID(templateID),
+				zap.String("build_id", b.BuildID.String()),
+				zap.Error(err),
+			)
+		} else {
+			logger.L().Info(ctx, "Deleted build artifacts",
+				logger.WithTemplateID(templateID),
+				zap.String("build_id", b.BuildID.String()),
+			)
+		}
+	}
+}
+
+// deletableAfterReachabilityGuard returns the subset of candidate builds that
+// are safe to physically delete: those NOT reachable from the header closure of
+// any other live template of the team. It fails closed (returns an error) when
+// the ancestor reader is unavailable or any live root's headers cannot be read,
+// so the caller deletes nothing rather than risk corrupting a shared ancestor.
+func (a *APIStore) deletableAfterReachabilityGuard(ctx context.Context, teamID uuid.UUID, templateID string, candidates []uuid.UUID) (map[uuid.UUID]struct{}, error) {
+	if len(candidates) == 0 {
+		return map[uuid.UUID]struct{}{}, nil
+	}
+
+	if a.buildAncestors == nil {
+		return nil, fmt.Errorf("build ancestor reader not configured; cannot prove reachability")
+	}
+
+	liveRoots, err := a.collectLiveRoots(ctx, teamID, templateID)
+	if err != nil {
+		return nil, err
+	}
+
+	keep, err := reachableKeepSet(ctx, a.buildAncestors, liveRoots)
+	if err != nil {
+		return nil, err
+	}
+
+	deletable, _ := partitionDeletableBuilds(candidates, keep)
+
+	out := make(map[uuid.UUID]struct{}, len(deletable))
+	for _, b := range deletable {
+		out[b] = struct{}{}
+	}
+
+	return out, nil
+}
+
+// liveSandboxStates are the sandbox states whose build must be protected: any
+// state in which the sandbox may still read the build's data. A sandbox in one
+// of these states can pin a build that is no longer any template's current
+// assignment.
+var liveSandboxStates = []sandbox.State{
+	sandbox.StateRunning,
+	sandbox.StatePausing,
+	sandbox.StateSnapshotting,
+	sandbox.StateKilling,
+}
+
+// collectLiveRoots unions the live-root build IDs from all three sources whose
+// data a deletion must not orphan: (1) builds assigned to the team's other
+// active templates (DB), (2) the team's in-progress builds still writing to
+// storage (DB), and (3) builds pinned by the team's live sandboxes in the
+// orchestrator (Redis) — which may reference a build no longer tracked as any
+// template's current assignment. Any source failing returns an error so the
+// caller fails closed.
+func (a *APIStore) collectLiveRoots(ctx context.Context, teamID uuid.UUID, templateID string) ([]uuid.UUID, error) {
+	seen := make(map[uuid.UUID]struct{})
+	var roots []uuid.UUID
+	add := func(id uuid.UUID) {
+		if _, ok := seen[id]; ok {
+			return
+		}
+
+		seen[id] = struct{}{}
+		roots = append(roots, id)
+	}
+
+	dbRoots, err := a.sqlcDB.GetLiveBuildRootsForTeamExcludingTemplate(ctx, queries.GetLiveBuildRootsForTeamExcludingTemplateParams{
+		TeamID:     teamID,
+		TemplateID: templateID,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("query live build roots: %w", err)
+	}
+	for _, id := range dbRoots {
+		add(id)
+	}
+
+	inProgress, err := a.sqlcDB.GetInProgressBuildRootsForTeam(ctx, teamID)
+	if err != nil {
+		return nil, fmt.Errorf("query in-progress build roots: %w", err)
+	}
+	for _, id := range inProgress {
+		add(id)
+	}
+
+	if a.runtimeSandboxes == nil {
+		return nil, fmt.Errorf("runtime sandbox lister not configured; cannot prove reachability")
+	}
+
+	sandboxes, err := a.runtimeSandboxes.GetSandboxes(ctx, teamID, liveSandboxStates)
+	if err != nil {
+		return nil, fmt.Errorf("list live sandboxes: %w", err)
+	}
+	for _, sbx := range sandboxes {
+		add(sbx.BuildID)
+	}
+
+	return roots, nil
+}
+
 // DeleteTemplatesTemplateID serves to delete a template (e.g. in CLI)
 func (a *APIStore) DeleteTemplatesTemplateID(c *gin.Context, aliasOrTemplateID api.TemplateID) {
 	ctx := c.Request.Context()
@@ -129,6 +305,10 @@ func (a *APIStore) DeleteTemplatesTemplateID(c *gin.Context, aliasOrTemplateID a
 	a.templateCache.InvalidateAllTags(context.WithoutCancel(ctx), templateID)
 	a.templateCache.InvalidateAliasesByTemplateID(context.WithoutCancel(ctx), templateID, aliasKeys)
 
+	// Delete GCS/S3 artifacts in the background; DB deletion is already committed
+	// so artifact cleanup failures must not block or fail the HTTP response.
+	go a.deleteTemplateArtifacts(context.WithoutCancel(ctx), team.ID, templateID)
+
 	telemetry.ReportEvent(ctx, "deleted template from db")
 
 	properties := a.posthog.GetPackageToPosthogProperties(&c.Request.Header)