Add DialForNode and always finalize in FinalizeSuspendedStep

cmd/ateapi/internal/controlapi/dialer.go

@@ -36,6 +36,12 @@ import (
 
 var ErrWorkerPodNotFound = errors.New("worker pod not found")
 
+// ErrNoAteletOnNode reports that the informer cache holds no atelet pod for
+// the requested node — e.g. the atelet is restarting, or the node is gone.
+// Distinct from ErrWorkerPodNotFound, which callers treat as crash-worthy;
+// this one is retryable.
+var ErrNoAteletOnNode = errors.New("no atelet pod found on node")
+
 // The SPIFFE identity that atelet serving certs carry, as minted by the
 // podidentity signer (cmd/podcertcontroller/internal/podidentitysigner).
 // The namespace part is ateletNamespace, declared in informer.go.
@@ -91,13 +97,30 @@ func (d *AteletDialer) DialForWorker(workerPodNamespace, workerPodName string) (
 
 	selectedWorker := matchingPods[0].(*corev1.Pod)
 
-	matchingAtelets, err := d.ateletIndexer.ByIndex(byNode, selectedWorker.Spec.NodeName)
+	conn, err := d.DialForAteletOnNode(selectedWorker.Spec.NodeName)
+	if err != nil {
+		return nil, fmt.Errorf("for worker pod %q: %w", workerPodKey, err)
+	}
+	return conn, nil
+}
+
+// DialForAteletOnNode resolves the single atelet pod on nodeName and dials it
+// with per-atelet pod-UID-pinned credentials, caching the connection by the
+// atelet's pod UID. Used directly when an actor has no worker assignment but
+// its state is pinned to a node — e.g. a PAUSED actor whose local snapshot
+// lives there. Returns ErrNoAteletOnNode if the informer cache holds no
+// atelet pod for the node.
+func (d *AteletDialer) DialForAteletOnNode(nodeName string) (*grpc.ClientConn, error) {
+	matchingAtelets, err := d.ateletIndexer.ByIndex(byNode, nodeName)
 	if err != nil {
-		return nil, fmt.Errorf("while finding atelet for worker pod %q on node %q: %w", workerPodKey, selectedWorker.Spec.NodeName, err)
+		return nil, fmt.Errorf("while finding atelet on node %q: %w", nodeName, err)
 	}
 
-	if len(matchingAtelets) != 1 {
-		return nil, fmt.Errorf("found %d atelet pods on node %q, expected 1", len(matchingAtelets), selectedWorker.Spec.NodeName)
+	if len(matchingAtelets) == 0 {
+		return nil, fmt.Errorf("%w: %q", ErrNoAteletOnNode, nodeName)
+	}
+	if len(matchingAtelets) > 1 {
+		return nil, fmt.Errorf("found %d atelet pods on node %q, expected 1", len(matchingAtelets), nodeName)
 	}
 
 	selectedAtelet := matchingAtelets[0].(*corev1.Pod)

cmd/ateapi/internal/controlapi/dialer_test.go

@@ -21,6 +21,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
+	"errors"
 	"math/big"
 	"net/url"
 	"testing"
@@ -29,6 +30,12 @@ import (
 	"github.com/agent-substrate/substrate/internal/substratex509"
 	"github.com/spiffe/go-spiffe/v2/bundle/x509bundle"
 	"github.com/spiffe/go-spiffe/v2/spiffeid"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/tools/cache"
 )
 
 const testAteletSPIFFEID = "spiffe://cluster.local/ns/ate-system/sa/atelet"
@@ -214,3 +221,74 @@ func TestVerifyAteletServerCert(t *testing.T) {
 		}
 	})
 }
+
+// newTestAteletIndexer builds an indexer with the production byNode index
+// shape, holding the given atelet pods.
+func newTestAteletIndexer(t *testing.T, pods ...*corev1.Pod) cache.Indexer {
+	t.Helper()
+	idx := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{
+		byNode: func(obj any) ([]string, error) {
+			return []string{obj.(*corev1.Pod).Spec.NodeName}, nil
+		},
+	})
+	for _, p := range pods {
+		if err := idx.Add(p); err != nil {
+			t.Fatalf("adding pod to indexer: %v", err)
+		}
+	}
+	return idx
+}
+
+func TestDialForAteletOnNode(t *testing.T) {
+	ateletPod := func(name, uid, node, ip string) *corev1.Pod {
+		return &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{Namespace: "ate-system", Name: name, UID: types.UID(uid)},
+			Spec:       corev1.PodSpec{NodeName: node},
+			Status:     corev1.PodStatus{PodIPs: []corev1.PodIP{{IP: ip}}},
+		}
+	}
+
+	t.Run("no atelet on node", func(t *testing.T) {
+		d := NewAteletDialer(nil, newTestAteletIndexer(t), "", "")
+		if _, err := d.DialForAteletOnNode("node1"); !errors.Is(err, ErrNoAteletOnNode) {
+			t.Fatalf("DialForAteletOnNode = %v, want ErrNoAteletOnNode", err)
+		}
+	})
+
+	t.Run("more than one atelet on node", func(t *testing.T) {
+		d := NewAteletDialer(nil, newTestAteletIndexer(t,
+			ateletPod("atelet-1", "uid-1", "node1", "10.0.0.1"),
+			ateletPod("atelet-2", "uid-2", "node1", "10.0.0.2"),
+		), "", "")
+		_, err := d.DialForAteletOnNode("node1")
+		if err == nil || errors.Is(err, ErrNoAteletOnNode) {
+			t.Fatalf("DialForAteletOnNode = %v, want a non-ErrNoAteletOnNode error", err)
+		}
+	})
+
+	t.Run("dials and caches the node's atelet", func(t *testing.T) {
+		d := NewAteletDialer(nil, newTestAteletIndexer(t,
+			ateletPod("atelet-1", "uid-1", "node1", "10.0.0.1"),
+		), "", "")
+		var credsUID string
+		d.dialCredentials = func(expectedPodUID string) (credentials.TransportCredentials, error) {
+			credsUID = expectedPodUID
+			return insecure.NewCredentials(), nil
+		}
+
+		conn, err := d.DialForAteletOnNode("node1")
+		if err != nil {
+			t.Fatalf("DialForAteletOnNode: %v", err)
+		}
+		if credsUID != "uid-1" {
+			t.Errorf("credentials pinned to pod UID %q, want %q", credsUID, "uid-1")
+		}
+		again, err := d.DialForAteletOnNode("node1")
+		if err != nil {
+			t.Fatalf("DialForAteletOnNode (cached): %v", err)
+		}
+		if again != conn {
+			t.Error("second DialForAteletOnNode returned a different connection, want the cached one")
+		}
+	})
+}

cmd/ateapi/internal/controlapi/workflow_suspend.go

@@ -253,7 +253,7 @@ func (s *FinalizeSuspendedStep) Execute(ctx context.Context, input *SuspendInput
 		return err
 	}
 
-	// 1. Free the worker (if it hasn't been freed yet)
+	// 1. Free the worker (if the actor has one and it hasn't been freed yet)
 	if assignment := latestActor.GetWorkerAssignment(); assignment != nil {
 		workerPod := assignment.GetWorkerPod()
 
@@ -276,46 +276,50 @@ func (s *FinalizeSuspendedStep) Execute(ctx context.Context, input *SuspendInput
 			}
 		}
 
-		// 2. Clear the actor's assignment, now that the worker is freed
+		// Re-fetch the actor now that the worker is freed.
 		latestActor, err = s.store.GetActor(ctx, input.ActorRef)
 		if err != nil {
 			return err
 		}
-		latestActor.Status = ateapipb.Actor_STATUS_SUSPENDED
-		if latestActor.InProgressSnapshotName != "" {
-			snapshotName := latestActor.InProgressSnapshotName
-			// The same inputs CallAteletSuspend used, so the recorded URI is
-			// where the bytes were actually written.
-			snapshotURI, err := inProgressSnapshotURI(state, input.ActorRef.Atespace, snapshotName)
-			if err != nil {
-				return err
-			}
-			snapshot := &ateapipb.ActorSnapshot{
-				Metadata:               &ateapipb.ResourceMetadata{Atespace: input.ActorRef.Atespace, Name: snapshotName},
-				SourceActor:            input.ActorRef.ToObjectRef(),
-				SourceActorUid:         latestActor.GetMetadata().GetUid(),
-				SourceActorVersion:     state.SourceVersion,
-				ActorTemplateNamespace: latestActor.GetActorTemplateNamespace(),
-				ActorTemplateName:      latestActor.GetActorTemplateName(),
-				ActorTemplateUid:       string(state.ActorTemplate.GetUID()),
-				ContentScope:           toActorSnapshotContentScope(commitSnapshotScope(input.ActorRef.Atespace, state.ActorTemplate)),
-				SnapshotUri:            snapshotURI.String(),
-			}
-			if _, err := s.store.CreateActorSnapshot(ctx, snapshot); err != nil && !errors.Is(err, store.ErrAlreadyExists) {
-				return err
-			}
-			latestActor.LatestSnapshot = &ateapipb.ObjectRef{Atespace: input.ActorRef.Atespace, Name: snapshotName}
-			latestActor.InProgressSnapshotName = ""
-			latestActor.InProgressSnapshotSourceActorVersion = 0
-		}
-		latestActor.WorkerAssignment = nil
-		latestActor.LocalSnapshotInfo = nil
-		updatedActor, err := s.store.UpdateActor(ctx, latestActor, latestActor.GetMetadata().GetVersion())
+	}
+
+	// 2. Finalize the actor: record the snapshot and mark it SUSPENDED. This
+	// must run even with no worker assignment (nothing to free), or the actor
+	// would be left SUSPENDING forever with the workflow reporting success.
+	latestActor.Status = ateapipb.Actor_STATUS_SUSPENDED
+	if latestActor.InProgressSnapshotName != "" {
+		snapshotName := latestActor.InProgressSnapshotName
+		// The same inputs CallAteletSuspend used, so the recorded URI is
+		// where the bytes were actually written.
+		snapshotURI, err := inProgressSnapshotURI(state, input.ActorRef.Atespace, snapshotName)
 		if err != nil {
 			return err
 		}
-		latestActor = updatedActor
+		snapshot := &ateapipb.ActorSnapshot{
+			Metadata:               &ateapipb.ResourceMetadata{Atespace: input.ActorRef.Atespace, Name: snapshotName},
+			SourceActor:            input.ActorRef.ToObjectRef(),
+			SourceActorUid:         latestActor.GetMetadata().GetUid(),
+			SourceActorVersion:     state.SourceVersion,
+			ActorTemplateNamespace: latestActor.GetActorTemplateNamespace(),
+			ActorTemplateName:      latestActor.GetActorTemplateName(),
+			ActorTemplateUid:       string(state.ActorTemplate.GetUID()),
+			ContentScope:           toActorSnapshotContentScope(commitSnapshotScope(input.ActorRef.Atespace, state.ActorTemplate)),
+			SnapshotUri:            snapshotURI.String(),
+		}
+		if _, err := s.store.CreateActorSnapshot(ctx, snapshot); err != nil && !errors.Is(err, store.ErrAlreadyExists) {
+			return err
+		}
+		latestActor.LatestSnapshot = &ateapipb.ObjectRef{Atespace: input.ActorRef.Atespace, Name: snapshotName}
+		latestActor.InProgressSnapshotName = ""
+		latestActor.InProgressSnapshotSourceActorVersion = 0
+	}
+	latestActor.WorkerAssignment = nil
+	latestActor.LocalSnapshotInfo = nil
+	updatedActor, err := s.store.UpdateActor(ctx, latestActor, latestActor.GetMetadata().GetVersion())
+	if err != nil {
+		return err
 	}
+	latestActor = updatedActor
 
 	state.Actor = latestActor
 	return nil

cmd/ateapi/internal/controlapi/workflow_suspend_test.go

@@ -301,6 +301,74 @@ func TestCallAteletSuspendStep_DanglingWorkerDoesNotRecordPhantomSnapshot(t *tes
 	}
 }
 
+// TestFinalizeSuspendedStep_NoAssignment verifies finalization runs even when
+// the actor has no worker assignment: the ActorSnapshot must be recorded and
+// the actor moved to SUSPENDED rather than silently left SUSPENDING. This is
+// the shape a paused-origin suspend (#791) produces — a PAUSED actor has no
+// worker — and the regression test for finalization previously living inside
+// the worker-freeing branch.
+func TestFinalizeSuspendedStep_NoAssignment(t *testing.T) {
+	ctx := context.Background()
+	persistence := newTestPersistence(t)
+
+	const snapshotName = "2026-01-01t00-00-00z-abc"
+	actor := &ateapipb.Actor{
+		Metadata:               &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor-1"},
+		Status:                 ateapipb.Actor_STATUS_SUSPENDING,
+		InProgressSnapshotName: snapshotName,
+		LocalSnapshotInfo: &ateapipb.LocalSnapshotInfo{
+			SnapshotName:              "actor-1-pause-snapshot",
+			NodeVmsWithLocalSnapshots: []string{"node1"},
+		},
+	}
+	created, err := persistence.CreateActor(ctx, actor)
+	if err != nil {
+		t.Fatalf("CreateActor: %v", err)
+	}
+
+	step := &FinalizeSuspendedStep{store: persistence}
+	input := &SuspendInput{ActorRef: resources.ActorRef{Atespace: "team-a", Name: "actor-1"}}
+	state := &SuspendState{
+		Actor:         created,
+		ActorTemplate: &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SnapshotsConfig: atev1alpha1.SnapshotsConfig{Location: "gs://snapshots"}}},
+		SourceVersion: created.GetMetadata().GetVersion(),
+	}
+	if err := step.Execute(ctx, input, state); err != nil {
+		t.Fatalf("Execute: %v", err)
+	}
+
+	stored, err := persistence.GetActor(ctx, resources.ActorRef{Atespace: "team-a", Name: "actor-1"})
+	if err != nil {
+		t.Fatalf("GetActor: %v", err)
+	}
+	if stored.GetStatus() != ateapipb.Actor_STATUS_SUSPENDED {
+		t.Errorf("status = %v, want SUSPENDED", stored.GetStatus())
+	}
+	if got := stored.GetLatestSnapshot().GetName(); got != snapshotName {
+		t.Errorf("LatestSnapshot = %q, want %q", got, snapshotName)
+	}
+	if got := stored.GetInProgressSnapshotName(); got != "" {
+		t.Errorf("InProgressSnapshotName = %q, want cleared", got)
+	}
+	if stored.GetLocalSnapshotInfo() != nil {
+		t.Errorf("LocalSnapshotInfo = %v, want cleared", stored.GetLocalSnapshotInfo())
+	}
+	snapshot, err := persistence.GetActorSnapshot(ctx, "team-a", snapshotName)
+	if err != nil {
+		t.Fatalf("GetActorSnapshot: %v", err)
+	}
+	wantURI, err := resources.NewSnapshotURI("gs://snapshots", "team-a", snapshotName)
+	if err != nil {
+		t.Fatalf("NewSnapshotURI: %v", err)
+	}
+	if got := snapshot.GetSnapshotUri(); got != wantURI.String() {
+		t.Errorf("snapshot URI = %q, want %q", got, wantURI.String())
+	}
+	if got := snapshot.GetSourceActorUid(); got != created.GetMetadata().GetUid() {
+		t.Errorf("snapshot SourceActorUid = %q, want %q", got, created.GetMetadata().GetUid())
+	}
+}
+
 func TestFinalizeSuspendedStep_ReleasesOnlyOwnWorker(t *testing.T) {
 	tests := []struct {
 		name               string