Update sandbox metadata

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

@@ -122,6 +122,14 @@ func (i *InstanceInfo) LoggerMetadata() sbxlogger.SandboxMetadata {
 	}
 }
 
+func (i *InstanceInfo) Lock() {
+	i.mu.Lock()
+}
+
+func (i *InstanceInfo) Unlock() {
+	i.mu.Unlock()
+}
+
 func (i *InstanceInfo) IsExpired() bool {
 	i.mu.RLock()
 	defer i.mu.RUnlock()

packages/api/internal/handlers/sandbox_update.go

@@ -0,0 +1,92 @@
+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.StatusUnauthorized, fmt.Sprintf("Error updating sandbox - sandbox '%s' does not belong to your team '%s'", sandboxID, teamID.String()))
+
+			return
+		}
+
+		sbx.Lock()
+		defer sbx.Unlock()
+
+		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),
+			zap.String("team.id", 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/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.Metadata = 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;

packages/orchestrator/internal/sandbox/sandbox.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -89,6 +90,8 @@ type Sandbox struct {
 	*Resources
 	*Metadata
 
+	Mu sync.Mutex
+
 	files   *storage.SandboxFiles
 	cleanup *Cleanup
 

packages/orchestrator/internal/server/sandboxes.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"time"
@@ -181,6 +182,11 @@ func (s *server) Create(ctxConn context.Context, req *orchestrator.SandboxCreate
 	}, nil
 }
 
+type eventData struct {
+	SetTimeout string            `json:"set_timeout,omitempty"`
+	Metadata   map[string]string `json:"metadata,omitempty"`
+}
+
 func (s *server) Update(ctx context.Context, req *orchestrator.SandboxUpdateRequest) (*emptypb.Empty, error) {
 	ctx, childSpan := s.tracer.Start(ctx, "sandbox-update")
 	defer childSpan.End()
@@ -197,18 +203,37 @@ func (s *server) Update(ctx context.Context, req *orchestrator.SandboxUpdateRequ
 		return nil, status.Error(codes.NotFound, "sandbox not found")
 	}
 
-	item.EndAt = req.EndTime.AsTime()
+	item.Mu.Lock()
+	defer item.Mu.Unlock()
+
+	updated := false
+
+	event := &eventData{}
+	// Update end time if provided
+	if req.EndTime != nil {
+		item.EndAt = req.EndTime.AsTime()
+		updated = true
+		event.SetTimeout = req.EndTime.AsTime().Format(time.RFC3339)
+
+		telemetry.ReportEvent(ctx, "Updated sandbox timeout")
+	}
+
+	// Update metadata if provided
+	if req.Metadata != nil {
+		item.APIStoredConfig.Metadata = req.Metadata
+		updated = true
+		event.Metadata = req.Metadata
 
-	// TODO: adapt to new types of update events
-	eventData := fmt.Sprintf(`{"set_timeout": "%s"}`, req.EndTime.AsTime().Format(time.RFC3339))
+		telemetry.ReportEvent(ctx, "Updated sandbox metadata")
+	}
 
 	sandboxLifeCycleEventsWriteFlag, flagErr := s.featureFlags.BoolFlag(
 		featureflags.SandboxLifeCycleEventsWriteFlagName, item.Runtime.SandboxID)
 	if flagErr != nil {
 		zap.L().Error("soft failing during sandbox lifecycle events write feature flag receive", zap.Error(flagErr))
 	}
-	if sandboxLifeCycleEventsWriteFlag {
-		go func(eventData string) {
+	if updated && sandboxLifeCycleEventsWriteFlag {
+		go func() {
 			buildId := ""
 			if item.APIStoredConfig != nil {
 				buildId = item.APIStoredConfig.BuildId
@@ -220,6 +245,12 @@ func (s *server) Update(ctx context.Context, req *orchestrator.SandboxUpdateRequ
 				return
 			}
 
+			eventDataMarshalled, err := json.Marshal(event)
+			if err != nil {
+				sbxlogger.I(item).Error("error marshaling sandbox lifecycle event", zap.Error(err))
+				return
+			}
+
 			err = s.sandboxEventBatcher.Push(clickhouse.SandboxEvent{
 				Timestamp:          time.Now().UTC(),
 				SandboxID:          item.Runtime.SandboxID,
@@ -229,13 +260,13 @@ func (s *server) Update(ctx context.Context, req *orchestrator.SandboxUpdateRequ
 				SandboxExecutionID: item.Runtime.ExecutionID,
 				EventCategory:      string(clickhouse.SandboxEventCategoryLifecycle),
 				EventLabel:         string(clickhouse.SandboxEventLabelUpdate),
-				EventData:          sql.NullString{String: eventData, Valid: true},
+				EventData:          sql.NullString{String: string(eventDataMarshalled), Valid: true},
 			})
 			if err != nil {
 				sbxlogger.I(item).Error(
 					"error inserting sandbox lifecycle event", zap.String("event_label", string(clickhouse.SandboxEventLabelUpdate)), zap.Error(err))
 			}
-		}(eventData)
+		}()
 	}
 
 	return &emptypb.Empty{}, nil