Update sandbox metadata

packages/api/internal/cache/instance/instance.go

@@ -60,7 +60,7 @@ func NewInstanceInfo(
 		ExecutionID:         ExecutionID,
 		TeamID:              TeamID,
 		BuildID:             BuildID,
-		Metadata:            Metadata,
+		metadata:            Metadata,
 		MaxInstanceLength:   MaxInstanceLength,
 		StartTime:           StartTime,
 		endTime:             endTime,
@@ -77,7 +77,6 @@ func NewInstanceInfo(
 		AutoPause:           atomic.Bool{},
 		Pausing:             utils.NewSetOnce[string](),
 		BaseTemplateID:      BaseTemplateID,
-		mu:                  sync.RWMutex{},
 	}
 
 	instance.AutoPause.Store(AutoPause)
@@ -95,7 +94,7 @@ type InstanceInfo struct {
 	TeamID              uuid.UUID
 	BuildID             uuid.UUID
 	BaseTemplateID      string
-	Metadata            map[string]string
+	metadata            map[string]string
 	MaxInstanceLength   time.Duration
 	StartTime           time.Time
 	endTime             time.Time
@@ -111,7 +110,7 @@ type InstanceInfo struct {
 	ClusterID           uuid.UUID
 	AutoPause           atomic.Bool
 	Pausing             *utils.SetOnce[string]
-	mu                  sync.RWMutex
+	sync.RWMutex
 }
 
 func (i *InstanceInfo) LoggerMetadata() sbxlogger.SandboxMetadata {
@@ -123,22 +122,36 @@ func (i *InstanceInfo) LoggerMetadata() sbxlogger.SandboxMetadata {
 }
 
 func (i *InstanceInfo) IsExpired() bool {
-	i.mu.RLock()
-	defer i.mu.RUnlock()
+	i.RLock()
+	defer i.RUnlock()
 
 	return time.Now().After(i.endTime)
 }
 
+func (i *InstanceInfo) Metadata() map[string]string {
+	i.RLock()
+	defer i.RUnlock()
+
+	return i.metadata
+}
+
+func (i *InstanceInfo) UpdateMetadata(metadata map[string]string) {
+	i.Lock()
+	defer i.Unlock()
+
+	i.metadata = metadata
+}
+
 func (i *InstanceInfo) GetEndTime() time.Time {
-	i.mu.RLock()
-	defer i.mu.RUnlock()
+	i.RLock()
+	defer i.RUnlock()
 
 	return i.endTime
 }
 
 func (i *InstanceInfo) SetEndTime(endTime time.Time) {
-	i.mu.Lock()
-	defer i.mu.Unlock()
+	i.Lock()
+	defer i.Unlock()
 
 	i.endTime = endTime
 }

packages/api/internal/handlers/sandbox_get.go

@@ -41,35 +41,39 @@ func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) {
 	}
 
 	// Try to get the running sandbox first
-	info, err := a.orchestrator.GetInstance(ctx, sandboxId)
+	sbx, err := a.orchestrator.GetInstance(ctx, sandboxId)
 	if err == nil {
+		sbx.RLock()
+		defer sbx.RUnlock()
+
 		// Check if sandbox belongs to the team
-		if info.TeamID != team.ID {
+		if sbx.TeamID != team.ID {
 			telemetry.ReportCriticalError(ctx, fmt.Sprintf("sandbox '%s' doesn't belong to team '%s'", sandboxId, team.ID.String()), nil)
 			a.sendAPIStoreError(c, http.StatusNotFound, fmt.Sprintf("sandbox \"%s\" doesn't exist or you don't have access to it", id))
 
 			return
 		}
 
-		// Sandbox exists and belongs to the team - return running sandbox info
+		// Sandbox exists and belongs to the team - return running sandbox sbx
 		sandbox := api.SandboxDetail{
-			ClientID:        info.ClientID,
-			TemplateID:      info.TemplateID,
-			Alias:           info.Alias,
-			SandboxID:       info.SandboxID,
-			StartedAt:       info.StartTime,
-			CpuCount:        api.CPUCount(info.VCpu),
-			MemoryMB:        api.MemoryMB(info.RamMB),
-			DiskSizeMB:      api.DiskSizeMB(info.TotalDiskSizeMB),
-			EndAt:           info.GetEndTime(),
+			ClientID:        sbx.ClientID,
+			TemplateID:      sbx.TemplateID,
+			Alias:           sbx.Alias,
+			SandboxID:       sbx.SandboxID,
+			StartedAt:       sbx.StartTime,
+			CpuCount:        api.CPUCount(sbx.VCpu),
+			MemoryMB:        api.MemoryMB(sbx.RamMB),
+			DiskSizeMB:      api.DiskSizeMB(sbx.TotalDiskSizeMB),
+			EndAt:           sbx.GetEndTime(),
 			State:           api.Running,
-			EnvdVersion:     info.EnvdVersion,
-			EnvdAccessToken: info.EnvdAccessToken,
+			EnvdVersion:     sbx.EnvdVersion,
+			EnvdAccessToken: sbx.EnvdAccessToken,
 			Domain:          sbxDomain,
 		}
 
-		if info.Metadata != nil {
-			meta := api.SandboxMetadata(info.Metadata)
+		metadata := sbx.Metadata()
+		if metadata != nil {
+			meta := api.SandboxMetadata(metadata)
 			sandbox.Metadata = &meta
 		}
 

packages/api/internal/handlers/sandbox_update.go

@@ -0,0 +1,89 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/utils"
+	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/db/types"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+func (a *APIStore) PutSandboxesSandboxIDMetadata(
+	c *gin.Context,
+	sandboxID api.SandboxID,
+) {
+	ctx := c.Request.Context()
+	sandboxID = utils.ShortID(sandboxID)
+
+	// Get team from context
+	teamID := a.GetTeamInfo(c).Team.ID
+
+	metadata, err := utils.ParseBody[api.PutSandboxesSandboxIDMetadataJSONRequestBody](ctx, c)
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing request: %s", err))
+
+		telemetry.ReportCriticalError(ctx, "error when parsing request", err)
+
+		return
+	}
+
+	telemetry.SetAttributes(ctx,
+		telemetry.WithSandboxID(sandboxID),
+		telemetry.WithTeamID(teamID.String()),
+	)
+
+	sbx, err := a.orchestrator.GetSandbox(sandboxID)
+	if err == nil {
+		// Verify the sandbox belongs to the team
+		if sbx.TeamID != teamID {
+			telemetry.ReportCriticalError(ctx, fmt.Sprintf("sandbox '%s' doesn't belong to team '%s'", sandboxID, teamID.String()), nil)
+			a.sendAPIStoreError(c, http.StatusNotFound, fmt.Sprintf("Sandbox '%s' not found", sandboxID))
+
+			return
+		}
+
+		apiErr := a.orchestrator.UpdateSandboxMetadata(ctx, sbx, metadata)
+		if apiErr != nil {
+			telemetry.ReportError(ctx, "error when updating sandbox metadata", apiErr.Err)
+			a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
+
+			return
+		}
+	} else {
+		// Sandbox not found in cache, might be paused
+		// Try to update the snapshot metadata in the database
+		zap.L().Debug("Sandbox not found in cache, checking if it's paused",
+			logger.WithSandboxID(sandboxID),
+			logger.WithTeamID(teamID.String()))
+
+		updated, err := a.sqlcDB.UpdateSnapshotMetadata(ctx, queries.UpdateSnapshotMetadataParams{
+			SandboxID: sandboxID,
+			TeamID:    teamID,
+			Metadata:  types.JSONBStringMap(metadata),
+		})
+		if err != nil {
+			telemetry.ReportCriticalError(ctx, "error when updating paused sandbox metadata", err)
+			a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error when updating paused sandbox metadata: %s", err))
+
+			return
+		}
+
+		if len(updated) == 0 {
+			telemetry.ReportCriticalError(ctx, "sandbox not found", nil)
+			a.sendAPIStoreError(c, http.StatusNotFound, fmt.Sprintf("Sandbox '%s' not found", sandboxID))
+
+			return
+		}
+
+		telemetry.ReportEvent(ctx, "Updated paused sandbox metadata")
+	}
+
+	c.Status(http.StatusOK)
+}

packages/api/internal/handlers/sandboxes_list.go

@@ -286,8 +286,9 @@ func instanceInfoToPaginatedSandboxes(runningSandboxes []*instance.InstanceInfo)
 			PaginationTimestamp: info.StartTime,
 		}
 
-		if info.Metadata != nil {
-			meta := api.SandboxMetadata(info.Metadata)
+		metadata := info.Metadata()
+		if metadata != nil {
+			meta := api.SandboxMetadata(metadata)
 			sandbox.Metadata = &meta
 		}
 

packages/api/internal/orchestrator/admin.go

@@ -82,8 +82,10 @@ func (o *Orchestrator) AdminNodeDetail(nomadNodeShortID string) (*api.NodeDetail
 	for _, sbx := range o.instanceCache.Items() {
 		if sbx.NodeID == n.ID && sbx.ClusterID == n.ClusterID {
 			var metadata *api.SandboxMetadata
-			if sbx.Metadata != nil {
-				meta := api.SandboxMetadata(sbx.Metadata)
+
+			sbxMetadata := sbx.Metadata()
+			if sbxMetadata != nil {
+				meta := api.SandboxMetadata(sbxMetadata)
 				metadata = &meta
 			}
 

packages/api/internal/orchestrator/pause_instance.go

@@ -37,7 +37,7 @@ func (o *Orchestrator) PauseInstance(
 		VCPU:                sbx.VCpu,
 		RAMMB:               sbx.RamMB,
 		TotalDiskSizeMB:     sbx.TotalDiskSizeMB,
-		Metadata:            sbx.Metadata,
+		Metadata:            sbx.Metadata(),
 		KernelVersion:       sbx.KernelVersion,
 		FirecrackerVersion:  sbx.FirecrackerVersion,
 		EnvdVersion:         sbx.EnvdVersion,

packages/api/internal/orchestrator/update_sandbox_metadata.go

@@ -0,0 +1,55 @@
+package orchestrator
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
+	"github.com/e2b-dev/infra/packages/api/internal/utils"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+func (o *Orchestrator) UpdateSandboxMetadata(
+	ctx context.Context,
+	sbx *instance.InstanceInfo,
+	metadata map[string]string,
+) *api.APIError {
+	childCtx, childSpan := o.tracer.Start(ctx, "update-sandbox-metadata")
+	defer childSpan.End()
+
+	client, childCtx, err := o.GetClient(childCtx, sbx.ClusterID, sbx.NodeID)
+	if err != nil {
+		return &api.APIError{
+			Err:       err,
+			ClientMsg: "Failed to connect to sandbox node",
+			Code:      http.StatusInternalServerError,
+		}
+	}
+
+	// Replace metadata completely instead of merging
+	// Send metadata to the orchestrator first - don't update local cache until success
+	_, err = client.Sandbox.Update(
+		childCtx, &orchestrator.SandboxUpdateRequest{
+			SandboxId: sbx.SandboxID,
+			Metadata:  metadata,
+		},
+	)
+
+	err = utils.UnwrapGRPCError(err)
+	if err != nil {
+		return &api.APIError{
+			Err:       err,
+			ClientMsg: "Failed to update sandbox metadata",
+			Code:      http.StatusInternalServerError,
+		}
+	}
+
+	// Only update local cache after orchestrator call succeeds
+	sbx.UpdateMetadata(metadata)
+
+	telemetry.ReportEvent(childCtx, "Updated running sandbox metadata")
+
+	return nil
+}

packages/db/queries/update_snapshot_metadata.sql

@@ -0,0 +1,7 @@
+-- name: UpdateSnapshotMetadata :many
+UPDATE "public"."snapshots" s
+SET metadata = $3
+FROM "public"."envs" e
+WHERE s.env_id = e.id
+  AND s.sandbox_id = $1
+  AND e.team_id = $2 RETURNING s.id;