Record the SandboxConfigRef in snapshot storage instead of sandbox assets

cmd/ateapi/internal/controlapi/sandbox_assets.go

@@ -20,6 +20,10 @@ import (
 	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
 	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
 	listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
+	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/labels"
 )
 
@@ -63,6 +67,54 @@ func resolveSandboxAssets(
 	return sandboxAssetsProto(class, sc), nil
 }
 
+// resolveSandboxAssetsByRef resolves an actor's recorded SandboxConfig
+// reference. The object must still be exactly the recorded revision: a UID
+// mismatch means it was re-created under the same name, a resourceVersion
+// mismatch that it was updated in place — either way the binaries the
+// snapshot recorded are gone, and it must not silently resolve different
+// ones.
+func resolveSandboxAssetsByRef(
+	sandboxConfigLister listersv1alpha1.SandboxConfigLister,
+	ref *ateapipb.SandboxConfigRef,
+) (*ateletpb.SandboxAssets, error) {
+	sc, err := sandboxConfigLister.Get(ref.GetName())
+	if err != nil {
+		if k8serrors.IsNotFound(err) {
+			return nil, status.Errorf(codes.FailedPrecondition, "actor's SandboxConfig %s not found", ref.GetName())
+		}
+		return nil, fmt.Errorf("while getting SandboxConfig %q: %w", ref.GetName(), err)
+	}
+	if string(sc.UID) != ref.GetUid() {
+		return nil, status.Errorf(codes.FailedPrecondition,
+			"actor records SandboxConfig %s with uid %s, but the object now has uid %s",
+			ref.GetName(), ref.GetUid(), sc.UID)
+	}
+	if sc.ResourceVersion != ref.GetResourceVersion() {
+		return nil, status.Errorf(codes.FailedPrecondition,
+			"actor records SandboxConfig %s at resourceVersion %s, but the object is now at %s",
+			ref.GetName(), ref.GetResourceVersion(), sc.ResourceVersion)
+	}
+	class := sc.Spec.SandboxClass
+	if class == "" {
+		class = atev1alpha1.SandboxClassGvisor
+	}
+	return sandboxAssetsProto(class, sc), nil
+}
+
+// sandboxConfigRefFromAtelet converts the SandboxConfig reference atelet
+// reports (from its on-node record) into the control plane's proto. Nil when
+// atelet reports none — a record written before the reference existed.
+func sandboxConfigRefFromAtelet(ref *ateletpb.SandboxConfigRef) *ateapipb.SandboxConfigRef {
+	if ref == nil {
+		return nil
+	}
+	return &ateapipb.SandboxConfigRef{
+		Name:            ref.GetName(),
+		Uid:             ref.GetUid(),
+		ResourceVersion: ref.GetResourceVersion(),
+	}
+}
+
 // defaultSandboxConfig returns the single SandboxConfig marked Default for the
 // given class, erroring if there are zero or more than one.
 func defaultSandboxConfig(lister listersv1alpha1.SandboxConfigLister, class atev1alpha1.SandboxClass) (*atev1alpha1.SandboxConfig, error) {
@@ -92,6 +144,11 @@ func sandboxAssetsProto(class atev1alpha1.SandboxClass, sc *atev1alpha1.SandboxC
 		SandboxClass: string(class),
 		PauseImage:   sc.Spec.PauseImage,
 		Assets:       make(map[string]*ateletpb.ArchAssets, len(sc.Spec.Assets)),
+		SandboxConfigRef: &ateletpb.SandboxConfigRef{
+			Name:            sc.Name,
+			Uid:             string(sc.UID),
+			ResourceVersion: sc.ResourceVersion,
+		},
 	}
 	for arch, files := range sc.Spec.Assets {
 		archAssets := &ateletpb.ArchAssets{Files: make(map[string]*ateletpb.AssetFile, len(files))}

cmd/ateapi/internal/controlapi/sandbox_assets_test.go

@@ -19,6 +19,9 @@ import (
 
 	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
 	listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
+	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/tools/cache"
 )
@@ -101,3 +104,72 @@ func TestResolveSandboxAssetsCarriesPauseImage(t *testing.T) {
 		})
 	}
 }
+
+// TestResolveSandboxAssetsCarriesConfigRef pins that the resolved assets name
+// the SandboxConfig object they came from (name + UID + resourceVersion).
+func TestResolveSandboxAssetsCarriesConfigRef(t *testing.T) {
+	config := &atev1alpha1.SandboxConfig{
+		ObjectMeta: metav1.ObjectMeta{Name: "gvisor-prod", UID: "sandbox-uid-1", ResourceVersion: "42"},
+		Spec: atev1alpha1.SandboxConfigSpec{
+			SandboxClass: atev1alpha1.SandboxClassGvisor,
+			Default:      true,
+			PauseImage:   "registry.k8s.io/pause@sha256:abc",
+			Assets:       testAssets(),
+		},
+	}
+	pool := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{Name: "pool1", Namespace: "worker-ns"},
+	}
+	poolLister, configLister := listersFor(t, []*atev1alpha1.WorkerPool{pool}, []*atev1alpha1.SandboxConfig{config})
+
+	got, err := resolveSandboxAssets(poolLister, configLister, "worker-ns", "pool1")
+	if err != nil {
+		t.Fatalf("resolveSandboxAssets() error: %v", err)
+	}
+	ref := got.GetSandboxConfigRef()
+	if ref.GetName() != "gvisor-prod" || ref.GetUid() != "sandbox-uid-1" || ref.GetResourceVersion() != "42" {
+		t.Errorf("SandboxConfig ref = %s/%s@%s, want gvisor-prod/sandbox-uid-1@42", ref.GetName(), ref.GetUid(), ref.GetResourceVersion())
+	}
+}
+
+// TestResolveSandboxAssetsByRef pins that the named SandboxConfig is used
+// only while it is exactly the recorded object revision; a missing,
+// re-created, or in-place-updated object is a FailedPrecondition rather
+// than silently resolving different binaries.
+func TestResolveSandboxAssetsByRef(t *testing.T) {
+	config := &atev1alpha1.SandboxConfig{
+		ObjectMeta: metav1.ObjectMeta{Name: "gvisor-prod", UID: "sandbox-uid-1", ResourceVersion: "42"},
+		Spec: atev1alpha1.SandboxConfigSpec{
+			SandboxClass: atev1alpha1.SandboxClassGvisor,
+			PauseImage:   "registry.k8s.io/pause@sha256:abc",
+			Assets:       testAssets(),
+		},
+	}
+	_, configLister := listersFor(t, nil, []*atev1alpha1.SandboxConfig{config})
+
+	tests := []struct {
+		name     string
+		ref      *ateapipb.SandboxConfigRef
+		wantCode codes.Code
+	}{
+		{name: "match", ref: &ateapipb.SandboxConfigRef{Name: "gvisor-prod", Uid: "sandbox-uid-1", ResourceVersion: "42"}, wantCode: codes.OK},
+		{name: "object gone", ref: &ateapipb.SandboxConfigRef{Name: "gvisor-gone", Uid: "sandbox-uid-1", ResourceVersion: "42"}, wantCode: codes.FailedPrecondition},
+		{name: "uid mismatch (object recreated under the same name)", ref: &ateapipb.SandboxConfigRef{Name: "gvisor-prod", Uid: "sandbox-uid-0", ResourceVersion: "42"}, wantCode: codes.FailedPrecondition},
+		{name: "revision mismatch (object updated in place)", ref: &ateapipb.SandboxConfigRef{Name: "gvisor-prod", Uid: "sandbox-uid-1", ResourceVersion: "41"}, wantCode: codes.FailedPrecondition},
+		{name: "ref recorded without a revision", ref: &ateapipb.SandboxConfigRef{Name: "gvisor-prod", Uid: "sandbox-uid-1"}, wantCode: codes.FailedPrecondition},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := resolveSandboxAssetsByRef(configLister, tt.ref)
+			if code := status.Code(err); code != tt.wantCode {
+				t.Fatalf("status.Code = %v (err %v), want %v", code, err, tt.wantCode)
+			}
+			if tt.wantCode != codes.OK {
+				return
+			}
+			if got.GetSandboxConfigRef().GetName() != "gvisor-prod" || got.GetSandboxConfigRef().GetUid() != "sandbox-uid-1" {
+				t.Errorf("SandboxConfig ref = %s/%s, want gvisor-prod/sandbox-uid-1", got.GetSandboxConfigRef().GetName(), got.GetSandboxConfigRef().GetUid())
+			}
+		})
+	}
+}

cmd/ateapi/internal/controlapi/template_reconciler.go

@@ -206,8 +206,9 @@ func (r *ActorTemplateReconciler) reconcileOne(ctx context.Context, ref resource
 			// The golden snapshot exists already.
 			return 0, nil
 		}
-		// TODO: Freeze sandbox assets before creating the golden actor.
-
+		// The golden actor boots with the pool's SandboxConfig; the suspend
+		// below freezes the reference its checkpoint reports into the golden
+		// ActorSnapshot's status.
 		actor, err := r.ensureActorExists(ctx, tmpl, goldenActorRef)
 		if err != nil {
 			if status.Code(err) == codes.InvalidArgument {

cmd/ateapi/internal/controlapi/workflow_pause.go

@@ -73,7 +73,8 @@ func (w *ActorWorkflow) PauseActor(ctx context.Context, actorRef resources.Actor
 		return nil, err
 	}
 	actor = marked
-	if wireSnapshotScope, err = w.ensureAteletPaused(leaseCtx, actorRef, actor, actorTemplate); err != nil {
+	var sandboxConfigRef *ateapipb.SandboxConfigRef
+	if wireSnapshotScope, sandboxConfigRef, err = w.ensureAteletPaused(leaseCtx, actorRef, actor, actorTemplate); err != nil {
 		return nil, err
 	}
 	// TODO: There is no difference between suspend and pause for now, but we
@@ -86,7 +87,7 @@ func (w *ActorWorkflow) PauseActor(ctx context.Context, actorRef resources.Actor
 	// them here, as crash.go does for the crash counter.
 	finalAttrs = lifecycleOpAttrs(actor, actorTemplate, "", wireSnapshotScope)
 	var finalized *ateapipb.Actor
-	if finalized, err = w.ensurePausedFinalized(leaseCtx, actorRef, actorTemplate); err != nil {
+	if finalized, err = w.ensurePausedFinalized(leaseCtx, actorRef, actorTemplate, sandboxConfigRef); err != nil {
 		return nil, err
 	}
 	actor = finalized
@@ -147,12 +148,13 @@ func (w *ActorWorkflow) ensureMarkedPausing(ctx context.Context, actorRef resour
 }
 
 // ensureAteletPaused checkpoints the workload locally on the worker node
-// under the actor's persisted snapshot name. This is the atelet reentrancy
-// seam (#372): the request is keyed by the actor UID, the worker pod UID, and
-// the once-minted snapshot name, so a re-entered workflow re-sends the same
-// semantic request; once atelet's Checkpoint is idempotent on those keys this
-// step becomes fully reentrant with no changes here.
-func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resources.ActorRef, actor *ateapipb.Actor, actorTemplate *ateapipb.ActorTemplate) (wireSnapshotScope string, err error) {
+// under the actor's persisted snapshot name, returning the SandboxConfig
+// reference atelet reports the checkpoint was taken with. This is the atelet
+// reentrancy seam (#372): the request is keyed by the actor UID, the worker
+// pod UID, and the once-minted snapshot name, so a re-entered workflow
+// re-sends the same semantic request; once atelet's Checkpoint is idempotent
+// on those keys this step becomes fully reentrant with no changes here.
+func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resources.ActorRef, actor *ateapipb.Actor, actorTemplate *ateapipb.ActorTemplate) (wireSnapshotScope string, sandboxConfigRef *ateapipb.SandboxConfigRef, err error) {
 	ctx, done := stepSpan(ctx, "CallAteletPause")
 	defer func() { err = done(err) }()
 
@@ -162,7 +164,7 @@ func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resourc
 		if err := crashActor(ctx, w.store, actorRef, ateattr.OperationPause, ateattr.ReasonCorruptedAssignment); err != nil {
 			slog.ErrorContext(ctx, "Failed to crash actor", slog.String("err", err.Error()))
 		}
-		return "", status.Errorf(codes.FailedPrecondition, "CallAteletPause prerequisite not met for Actor: %s. No worker assignment", actorRef)
+		return "", nil, status.Errorf(codes.FailedPrecondition, "CallAteletPause prerequisite not met for Actor: %s. No worker assignment", actorRef)
 	}
 
 	ateletConn, err := w.dialer.DialForWorker(assignment.GetWorkerNamespace(), assignment.GetWorkerPod())
@@ -172,20 +174,20 @@ func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resourc
 			if err := crashActor(ctx, w.store, actorRef, ateattr.OperationPause, ateattr.ReasonWorkerPodGone); err != nil {
 				slog.ErrorContext(ctx, "Failed to crash actor", slog.String("err", err.Error()))
 			}
-			return "", fmt.Errorf("actor is CRASHED because its worker pod is gone and no snapshot was written")
+			return "", nil, fmt.Errorf("actor is CRASHED because its worker pod is gone and no snapshot was written")
 		}
-		return "", fmt.Errorf("while getting atelet conn for worker pod: %w", err)
+		return "", nil, fmt.Errorf("while getting atelet conn for worker pod: %w", err)
 	}
 	client := ateletpb.NewAteomHerderClient(ateletConn)
 
 	workloadSpec, err := workloadSpecFromActorTemplate(actorTemplate, actor)
 	if err != nil {
-		return "", err
+		return "", nil, err
 	}
 
-	// Checkpoint does not carry the sandbox config: atelet uses the version the
-	// actor is currently running (recorded on-node at Run/Restore) and pins it
-	// into the snapshot manifest.
+	// Checkpoint does not carry the sandbox config: atelet uses the version
+	// the actor is currently running (recorded on-node at Run/Restore) and
+	// pins its SandboxConfig reference into the snapshot manifest.
 	req := &ateletpb.CheckpointRequest{
 		TargetAteomUid:        assignment.GetWorkerPodUid(),
 		Atespace:              actor.GetMetadata().GetAtespace(),
@@ -204,8 +206,8 @@ func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resourc
 	}
 	wireSnapshotScope = ateattr.SnapshotScopeValue(req.Scope)
 
-	_, err = client.Checkpoint(ctx, req)
-	return wireSnapshotScope, maybeCrashActor(ctx, w.store, actorRef, err, "while checkpointing workload", ateattr.OperationPause)
+	resp, err := client.Checkpoint(ctx, req)
+	return wireSnapshotScope, sandboxConfigRefFromAtelet(resp.GetSandboxConfigRef()), maybeCrashActor(ctx, w.store, actorRef, err, "while checkpointing workload", ateattr.OperationPause)
 }
 
 // ensurePausedFinalized releases the actor's worker (only when it is still
@@ -215,7 +217,9 @@ func (w *ActorWorkflow) ensureAteletPaused(ctx context.Context, actorRef resourc
 // never be resumed. It re-reads the actor first so an out-of-band transition
 // (e.g. the syncer crashing the actor after its worker died) is not
 // overwritten: with no assignment left there is nothing to finalize.
-func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef resources.ActorRef, actorTemplate *ateapipb.ActorTemplate) (_ *ateapipb.Actor, err error) {
+// sandboxConfigRef is the reference the pause checkpoint reported (nil when
+// the checkpoint never completed or atelet's record predates the reference).
+func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef resources.ActorRef, actorTemplate *ateapipb.ActorTemplate, sandboxConfigRef *ateapipb.SandboxConfigRef) (_ *ateapipb.Actor, err error) {
 	ctx, done := stepSpan(ctx, "FinalizePaused")
 	defer func() { err = done(err) }()
 
@@ -284,6 +288,10 @@ func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef reso
 				localInfo := &ateapipb.LocalSnapshotInfo{
 					SnapshotName: toUpdate.GetStatus().GetInProgressLocalSnapshotName(),
 					ContentScope: contentScope,
+					// The snapshot record carries the SandboxConfig the
+					// checkpoint reported it was taken with; a later resume
+					// resolves the restore's assets from it.
+					SandboxConfigRef: sandboxConfigRef,
 				}
 				if newState != ateapipb.ActorState_ACTOR_STATE_CRASHED {
 					localInfo.NodeVmsWithLocalSnapshots = []string{nodeName}

cmd/ateapi/internal/controlapi/workflow_pause_test.go

@@ -59,7 +59,7 @@ func TestEnsurePausedFinalized_WorkerGone(t *testing.T) {
 	// Intentionally NOT creating the worker in store, simulates worker already gone.
 
 	w := &ActorWorkflow{store: st}
-	finalized, err := w.ensurePausedFinalized(ctx, actorRef, &ateapipb.ActorTemplate{})
+	finalized, err := w.ensurePausedFinalized(ctx, actorRef, &ateapipb.ActorTemplate{}, nil)
 	if err != nil {
 		t.Fatalf("ensurePausedFinalized: %v", err)
 	}
@@ -139,7 +139,7 @@ func TestEnsurePausedFinalized_RecordsContentScope(t *testing.T) {
 			tmpl := &ateapipb.ActorTemplate{
 				SnapshotsConfig: &ateapipb.SnapshotsConfig{OnPause: tc.onPause},
 			}
-			got, err := w.ensurePausedFinalized(ctx, actorRef, tmpl)
+			got, err := w.ensurePausedFinalized(ctx, actorRef, tmpl, &ateapipb.SandboxConfigRef{Name: "gvisor-prod", Uid: "sandbox-uid-1"})
 			if err != nil {
 				t.Fatalf("ensurePausedFinalized: %v", err)
 			}
@@ -150,6 +150,11 @@ func TestEnsurePausedFinalized_RecordsContentScope(t *testing.T) {
 			if scope := got.GetStatus().GetLocalSnapshotInfo().GetContentScope(); scope != tc.want {
 				t.Errorf("LocalSnapshotInfo.ContentScope = %v, want %v", scope, tc.want)
 			}
+			// The pause snapshot record carries the activation's SandboxConfig
+			// reference; a later resume resolves the restore's assets from it.
+			if ref := got.GetStatus().GetLocalSnapshotInfo().GetSandboxConfigRef(); ref.GetName() != "gvisor-prod" || ref.GetUid() != "sandbox-uid-1" {
+				t.Errorf("LocalSnapshotInfo.SandboxConfigRef = %s/%s, want gvisor-prod/sandbox-uid-1", ref.GetName(), ref.GetUid())
+			}
 		})
 	}
 }
@@ -283,7 +288,7 @@ func TestEnsureAteletPaused_DanglingWorkerDoesNotRecordPhantomSnapshot(t *testin
 			created := storetest.MustCreateActor(t, ctx, persistence, actor)
 
 			w := &ActorWorkflow{store: persistence, dialer: newDanglingDialer()}
-			if _, err := w.ensureAteletPaused(ctx, resources.ActorRef{Atespace: "team-a", Name: "actor-1"}, created, &ateapipb.ActorTemplate{}); err == nil {
+			if _, _, err := w.ensureAteletPaused(ctx, resources.ActorRef{Atespace: "team-a", Name: "actor-1"}, created, &ateapipb.ActorTemplate{}); err == nil {
 				t.Fatal("ensureAteletPaused: want error for dangling worker, got nil")
 			}