feat(api): per-region volume type defaults from node-derived region

packages/api/internal/cfg/model.go

@@ -112,6 +112,13 @@ type Config struct {
 
 	DefaultPersistentVolumeType string `env:"DEFAULT_PERSISTENT_VOLUME_TYPE"`
 
+	// DefaultPersistentVolumeTypeByRegion is the per-region default volume type,
+	// e.g. "us-west3:zonalfilestore-us-west3". A team's region is resolved from
+	// the region= node labels of the nodes its scheduling labels select; Terraform
+	// derives this map from the volume types themselves and fails the plan when a
+	// region with several types lacks an explicit default.
+	DefaultPersistentVolumeTypeByRegion map[string]string `env:"DEFAULT_PERSISTENT_VOLUME_TYPE_BY_REGION"`
+
 	DomainName string `env:"DOMAIN_NAME" envDefault:""`
 }
 

packages/api/internal/cfg/model_test.go

@@ -43,6 +43,17 @@ func TestParse(t *testing.T) {
 		assert.Equal(t, content, result.VolumesToken.SigningKey)
 	})
 
+	t.Run("default persistent volume type by region is parsed as a map", func(t *testing.T) {
+		t.Setenv("DEFAULT_PERSISTENT_VOLUME_TYPE_BY_REGION", "us-west3:zonalfilestore-us-west3,other:other-type")
+
+		result, err := Parse()
+		require.NoError(t, err)
+		assert.Equal(t, map[string]string{
+			"us-west3": "zonalfilestore-us-west3",
+			"other":    "other-type",
+		}, result.DefaultPersistentVolumeTypeByRegion)
+	})
+
 	t.Run("invalid service discovery provider exposes failure condition", func(t *testing.T) {
 		t.Setenv("SERVICE_DISCOVERY_PROVIDER", "invalid")
 

packages/api/internal/handlers/volume_create.go

@@ -6,12 +6,14 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/clusters"
+	"github.com/e2b-dev/infra/packages/auth/pkg/types"
 	"github.com/e2b-dev/infra/packages/db/pkg/dberrors"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	clustershared "github.com/e2b-dev/infra/packages/shared/pkg/clusters"
@@ -72,9 +74,12 @@ func (a *APIStore) PostVolumes(c *gin.Context) {
 		return
 	}
 
-	ctx = featureflags.AddToContext(ctx, featureflags.VolumeContext(body.Name))
+	ctx = featureflags.AddToContext(ctx,
+		featureflags.VolumeContext(body.Name),
+		featureflags.TeamContext(team.ID.String()),
+	)
 
-	volumeType := a.getVolumeType(ctx)
+	volumeType := a.getVolumeType(ctx, team)
 	if volumeType == "" {
 		a.sendAPIStoreError(c, http.StatusInternalServerError, "No persistent volume type is configured")
 		telemetry.ReportCriticalError(ctx, "default persistent volume type is not configured", nil)
@@ -174,13 +179,86 @@ func (a *APIStore) PostVolumes(c *gin.Context) {
 	c.JSON(http.StatusCreated, result)
 }
 
-func (a *APIStore) getVolumeType(ctx context.Context) string {
-	volumeType := a.featureFlags.StringFlag(ctx, featureflags.DefaultPersistentVolumeType)
-	if volumeType == "" {
-		volumeType = a.config.DefaultPersistentVolumeType
+const (
+	// regionNodeLabelPrefix marks the node label naming the region a node runs
+	// in, e.g. "region=us-west3". Regions live on nodes, never on teams:
+	// Terraform appends the label to every client cluster automatically.
+	regionNodeLabelPrefix = "region="
+
+	// defaultSchedulingLabel is the pool a team without scheduling labels of
+	// its own lands on.
+	defaultSchedulingLabel = "default"
+)
+
+// getVolumeType resolves the volume type a new volume of the given team gets,
+// in order of precedence: the LaunchDarkly override, the per-region default of
+// the region the team schedules into, and finally the deployment-wide default.
+//
+// Node labels only answer *where* the team runs; *what* a new volume there
+// should be is policy and comes exclusively from the region map. A region
+// mounting several volume types therefore never needs runtime guessing - the
+// map names its default explicitly.
+func (a *APIStore) getVolumeType(ctx context.Context, team *types.Team) string {
+	if volumeType := a.featureFlags.StringFlag(ctx, featureflags.DefaultPersistentVolumeType); volumeType != "" {
+		return volumeType
+	}
+
+	// Regional defaulting is opt-in: without a map there is nothing to look
+	// up, so don't walk the cluster.
+	if len(a.config.DefaultPersistentVolumeTypeByRegion) == 0 || a.orchestrator == nil {
+		return a.config.DefaultPersistentVolumeType
+	}
+
+	// Mirrors generateRequiredNodeLabels: a team without labels of its own
+	// runs on the "default" pool, so we resolve the region over the same set
+	// of nodes that placement would choose from.
+	requiredLabels := team.SandboxSchedulingLabels
+	if len(requiredLabels) == 0 {
+		requiredLabels = []string{defaultSchedulingLabel}
+	}
+
+	// Collect the distinct regions advertised by the ready nodes carrying all
+	// of requiredLabels; the same subset semantics as sandbox placement.
+	regions := make(map[string]struct{})
+	for _, node := range a.orchestrator.GetClusterNodes(clustershared.WithClusterFallback(team.ClusterID)) {
+		if node.Status() != api.NodeStatusReady {
+			continue
+		}
+
+		labels := node.Labels()
+		if !hasAllLabels(labels, requiredLabels) {
+			continue
+		}
+
+		for label := range labels {
+			if region, ok := strings.CutPrefix(label, regionNodeLabelPrefix); ok && region != "" {
+				regions[region] = struct{}{}
+			}
+		}
+	}
+
+	// Exactly one region with a mapped default is the only affirmative
+	// answer. Zero regions (labels match nothing yet) or several (labels do
+	// not pin a region) fall open to the deployment-wide default.
+	if len(regions) == 1 {
+		for region := range regions {
+			if volumeType, ok := a.config.DefaultPersistentVolumeTypeByRegion[region]; ok {
+				return volumeType
+			}
+		}
+	}
+
+	return a.config.DefaultPersistentVolumeType
+}
+
+func hasAllLabels(labels map[string]struct{}, requiredLabels []string) bool {
+	for _, required := range requiredLabels {
+		if _, ok := labels[required]; !ok {
+			return false
+		}
 	}
 
-	return volumeType
+	return true
 }
 
 var validVolumeNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)

packages/api/internal/handlers/volume_test.go

@@ -3,7 +3,16 @@ package handlers
 import (
 	"testing"
 
+	"github.com/google/uuid"
+	"github.com/launchdarkly/go-sdk-common/v3/ldvalue"
+	"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/api/internal/cfg"
+	"github.com/e2b-dev/infra/packages/auth/pkg/types"
+	authqueries "github.com/e2b-dev/infra/packages/db/pkg/auth/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 )
 
 func TestIsValidVolumeName(t *testing.T) {
@@ -135,3 +144,88 @@ func TestIsValidVolumeName(t *testing.T) {
 		})
 	}
 }
+
+// TestGetVolumeType covers the precedence paths that don't need a cluster;
+// APIStore.orchestrator is a concrete *orchestrator.Orchestrator with no test
+// constructor, so the node-derived resolution is not covered here.
+func TestGetVolumeType(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		// flagValue is the LaunchDarkly default-persistent-volume-type
+		// override; empty leaves the flag unset.
+		flagValue        string
+		regionVolumeType map[string]string
+		defaultType      string
+		expected         string
+	}{
+		{
+			name:        "no region map falls back to global default",
+			defaultType: "global-type",
+			expected:    "global-type",
+		},
+		{
+			// The orchestrator is only nil in tests, but a missing cluster
+			// view must not be what decides whether a volume can be created.
+			name:             "no orchestrator falls back to global default",
+			regionVolumeType: map[string]string{"us-west3": "zonalfilestore-us-west3"},
+			defaultType:      "global-type",
+			expected:         "global-type",
+		},
+		{
+			name:             "feature flag wins over region map and global default",
+			flagValue:        "flag-type",
+			regionVolumeType: map[string]string{"us-west3": "zonalfilestore-us-west3"},
+			defaultType:      "global-type",
+			expected:         "flag-type",
+		},
+		{
+			name:     "no default configured at all",
+			expected: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			// Each subtest gets its own datasource/client so parallel runs
+			// don't race on the shared flag value.
+			td := ldtestdata.DataSource()
+			ff, err := featureflags.NewClientWithDatasource(td)
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				assert.NoError(t, ff.Close(t.Context()))
+			})
+
+			if tt.flagValue != "" {
+				td.Update(td.Flag(featureflags.DefaultPersistentVolumeType.Key()).
+					ValueForAll(ldvalue.String(tt.flagValue)))
+			}
+
+			store := &APIStore{
+				featureFlags: ff,
+				config: cfg.Config{
+					DefaultPersistentVolumeType:         tt.defaultType,
+					DefaultPersistentVolumeTypeByRegion: tt.regionVolumeType,
+				},
+			}
+			team := &types.Team{Team: &authqueries.Team{ID: uuid.New()}}
+
+			assert.Equal(t, tt.expected, store.getVolumeType(t.Context(), team))
+		})
+	}
+}
+
+func TestHasAllLabels(t *testing.T) {
+	t.Parallel()
+
+	labels := map[string]struct{}{"gpu": {}, "highmem": {}, "region=us-west3": {}}
+
+	assert.True(t, hasAllLabels(labels, nil))
+	assert.True(t, hasAllLabels(labels, []string{"gpu"}))
+	assert.True(t, hasAllLabels(labels, []string{"gpu", "highmem"}))
+	assert.False(t, hasAllLabels(labels, []string{"gpu", "default"}))
+	assert.False(t, hasAllLabels(map[string]struct{}{}, []string{"default"}))
+}