Add CheckPrerequisite to WorkflowStep to validate state machine edges

cmd/ateapi/internal/controlapi/functional_test.go

@@ -1733,7 +1733,8 @@ func TestUpdateActor_NotFound(t *testing.T) {
 // TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible verifies that
 // a worker claimed by a failed resume attempt is released back to the free
 // pool if, by the next resume attempt, the actor's worker_selector has
-// changed such that the worker's pool is no longer eligible.
+// changed such that the worker's pool is no longer eligible. The actor
+// itself is crashed rather than transparently migrated to another pool.
 // Workflow:
 //  1. Creates pool-a (tier=a) and pool-b (tier=b), and an actor narrowed to
 //     tier=a.
@@ -1742,8 +1743,9 @@ func TestUpdateActor_NotFound(t *testing.T) {
 //     worker is claimed, leaving worker-a's actor assignment set and the actor
 //     stuck in RESUMING.
 //  3. Updates the actor's selector to tier=b, making pool-a ineligible.
-//  4. Resumes again; asserts it succeeds onto worker-b, and that worker-a
-//     has been released (actor assignment cleared) rather than left dangling.
+//  4. Resumes again; asserts it fails and the actor is CRASHED, that worker-a
+//     has been released (actor assignment cleared) rather than left dangling,
+//     and that worker-b remains free (the crashed actor must not claim it).
 func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T) {
 	ns := namespaceForTest("ns-resume-release-stale")
 	tc := setupTest(t, ns)
@@ -1782,19 +1784,16 @@ func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T)
 		t.Fatalf("UpdateActor failed: %v", err)
 	}
 
-	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}}); err != nil {
-		t.Fatalf("second ResumeActor failed: %v", err)
+	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}}); err == nil {
+		t.Fatalf("expected second ResumeActor to fail: the assigned worker's pool is no longer eligible")
 	}
 
 	getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}})
 	if err != nil {
 		t.Fatalf("GetActor failed: %v", err)
 	}
-	if got := getResp.GetWorkerPoolName(); got != "pool-b" {
-		t.Errorf("expected actor to land on pool-b, got worker_pool_name=%q", got)
-	}
-	if got := getResp.GetStatus(); got != ateapipb.Actor_STATUS_RUNNING {
-		t.Errorf("expected actor status RUNNING, got %v", got)
+	if got := getResp.GetStatus(); got != ateapipb.Actor_STATUS_CRASHED {
+		t.Errorf("expected actor status CRASHED, got %v", got)
 	}
 
 	listResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
@@ -1815,16 +1814,12 @@ func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T)
 				t.Errorf("expected worker-a (now-ineligible pool-a) to be released, got actor name=%q", got)
 			}
 		case "pool-b":
-			if wass := w.Assignment; wass == nil {
-				t.Errorf("expected worker-b to be claimed by %q, got nil assignment", name)
-			} else {
-				if wact := wass.Actor; wact == nil {
-					t.Errorf("expected worker-b to be claimed by %q, got nil assignment.actor", name)
-				} else {
-					if got := wact.Name; got != name {
-						t.Errorf("expected worker-b to be claimed by %q, got actor name=%q", name, got)
-					}
+			if wass := w.Assignment; wass != nil {
+				got := "<nil-actor>"
+				if wass.Actor != nil {
+					got = wass.Actor.Name
 				}
+				t.Errorf("expected worker-b to stay free (actor crashed, not migrated), got actor name=%q", got)
 			}
 		}
 	}

cmd/ateapi/internal/controlapi/workflow.go

@@ -44,6 +44,14 @@ type WorkflowStep[Params any, Context any] interface {
 	// If it returns true, the engine skips Execute() and fast-forwards to the next step.
 	IsComplete(ctx context.Context, params Params, wCtx Context) (bool, error)
 
+	// CheckPrerequisite validates that the current state permits executing this
+	// step (e.g. the actor's status allows this state-machine edge). The engine
+	// calls it only when IsComplete returned false, immediately before Execute,
+	// so completed steps of a retried workflow fast-forward without
+	// re-validation. Return a gRPC status error with
+	// codes.FailedPrecondition to abort the workflow if prereqs are not met.
+	CheckPrerequisite(ctx context.Context, params Params, wCtx Context) error
+
 	// Execute performs the step's business logic and persists any state changes.
 	// If an error is returned, the workflow stops and relies on the client to retry.
 	Execute(ctx context.Context, params Params, wCtx Context) error
@@ -79,6 +87,13 @@ func RunWorkflow[Params any, Context any](ctx context.Context, params Params, wC
 			continue
 		}
 
+		if err := step.CheckPrerequisite(ctx, params, wCtx); err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+			span.End()
+			return fmt.Errorf("prerequisite not met at step %s: %w", step.Name(), err)
+		}
+
 		err = runStep(ctx, params, wCtx, step)
 		if err != nil {
 			span.RecordError(err)

cmd/ateapi/internal/controlapi/workflow_pause.go

@@ -27,6 +27,8 @@ 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"
 	"k8s.io/apimachinery/pkg/util/wait"
 )
 
@@ -52,6 +54,9 @@ func (s *LoadActorForPauseStep) IsComplete(ctx context.Context, input *PauseInpu
 	// Always run to get the freshest state
 	return false, nil
 }
+func (s *LoadActorForPauseStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error {
+	return nil
+}
 func (s *LoadActorForPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName)
 	if err != nil {
@@ -79,11 +84,14 @@ func (s *MarkPausingStep) IsComplete(ctx context.Context, input *PauseInput, sta
 	// Fast forward if we've already marked our intent or if we are further along.
 	return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, nil
 }
-func (s *MarkPausingStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
+func (s *MarkPausingStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error {
+	// The pause edge only exists from RUNNING; PAUSING/PAUSED are fast-forwarded by IsComplete.
 	if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING {
-		return nil
+		return status.Errorf(codes.FailedPrecondition, "MarkPausingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING)
 	}
-
+	return nil
+}
+func (s *MarkPausingStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	state.Actor.Status = ateapipb.Actor_STATUS_PAUSING
 	state.Actor.InProgressSnapshot = fmt.Sprintf("%s-%s-%s", state.Actor.GetMetadata().GetName(), time.Now().Format(time.RFC3339), rand.Text())
 	updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion())
@@ -106,14 +114,20 @@ func (s *CallAteletPauseStep) IsComplete(ctx context.Context, input *PauseInput,
 	// If we are already PAUSED, we've already called Atelet
 	return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, nil
 }
-func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
+func (s *CallAteletPauseStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error {
+	if state.Actor.GetStatus() != ateapipb.Actor_STATUS_PAUSING {
+		return status.Errorf(codes.FailedPrecondition, "CallAteletPauseStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING)
+	}
 	if state.Actor.GetAteomPodNamespace() == "" || state.Actor.GetAteomPodName() == "" {
 		if err := crashActor(ctx, s.store, state.Actor.GetMetadata().GetAtespace(), state.Actor.GetMetadata().GetName()); err != nil {
 			slog.Error("Failed to crash actor", slog.String("err", err.Error()))
 		}
-		return fmt.Errorf("actor is CRASHED because it was in PAUSING state but has no active worker")
+		return status.Errorf(codes.FailedPrecondition, "CallAteletPauseStep prerequisite not met for Actor: %s. AteomPodNamespace: %s, GetAteomPodName %s", input.ActorName, state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName())
 	}
+	return nil
+}
 
+func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName())
 	if err != nil {
 		if errors.Is(err, ErrWorkerPodNotFound) {
@@ -164,6 +178,13 @@ func (s *FinalizePausedStep) IsComplete(ctx context.Context, input *PauseInput,
 	terminal := status == ateapipb.Actor_STATUS_PAUSED || status == ateapipb.Actor_STATUS_CRASHED
 	return terminal && state.Actor.GetAteomPodNamespace() == "", nil
 }
+
+func (s *FinalizePausedStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error {
+	if state.Actor.GetStatus() != ateapipb.Actor_STATUS_PAUSING {
+		return status.Errorf(codes.FailedPrecondition, "FinalizePausedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING)
+	}
+	return nil
+}
 func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName)
 	if err != nil {

cmd/ateapi/internal/controlapi/workflow_pause_test.go

@@ -20,6 +20,8 @@ import (
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // TestFinalizePausedStep_WorkerGone reproduces the scenario where the worker pod
@@ -114,3 +116,150 @@ func TestFindFreeWorker_EmptyNodeRestriction(t *testing.T) {
 		t.Error("expected a worker with nil restrictions, got nil")
 	}
 }
+
+// TestPauseActorWorkflow_RejectedAndIdempotentPaths covers the two
+// short-circuit paths of the pause workflow: rejection by MarkPausingStep's
+// CheckPrerequisite and the IsComplete idempotent fast-forward.
+func TestPauseActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) {
+	tests := []struct {
+		name       string
+		seedStatus ateapipb.Actor_Status
+		// wantErr true means PauseActor must fail with FailedPrecondition.
+		wantErr bool
+		// wantStatus is the stored status after the call.
+		wantStatus ateapipb.Actor_Status
+	}{
+		{
+			// Pausing a SUSPENDED actor is rejected by MarkPausingStep's
+			// CheckPrerequisite and the actor's status is left untouched.
+			name:       "not running rejected",
+			seedStatus: ateapipb.Actor_STATUS_SUSPENDED,
+			wantErr:    true,
+			wantStatus: ateapipb.Actor_STATUS_SUSPENDED,
+		},
+		{
+			// Pausing a PAUSED actor succeeds idempotently via IsComplete
+			// fast-forward without calling atelet.
+			name:       "already paused succeeds",
+			seedStatus: ateapipb.Actor_STATUS_PAUSED,
+			wantStatus: ateapipb.Actor_STATUS_PAUSED,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			st, cleanup := storetest.SetupTestStore(t)
+			defer cleanup()
+			w := newTestActorWorkflow(t, st, "ns", "tmpl1")
+
+			seedWorkflowActor(t, ctx, st, "team-a", "id1", "ns", "tmpl1", tc.seedStatus)
+
+			actor, err := w.PauseActor(ctx, "team-a", "id1")
+			if tc.wantErr {
+				if got := status.Code(err); got != codes.FailedPrecondition {
+					t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("PauseActor failed: %v", err)
+				}
+				if actor.GetStatus() != tc.wantStatus {
+					t.Errorf("returned status = %v, want %v", actor.GetStatus(), tc.wantStatus)
+				}
+			}
+
+			got, err := st.GetActor(ctx, "team-a", "id1")
+			if err != nil {
+				t.Fatalf("GetActor failed: %v", err)
+			}
+			if got.GetStatus() != tc.wantStatus {
+				t.Errorf("stored status = %v, want %v", got.GetStatus(), tc.wantStatus)
+			}
+		})
+	}
+}
+
+// TestPauseSteps_CheckPrerequisite verifies each pause step's CheckPrerequisite
+// against every actor status: nil for the step's allowed statuses,
+// FailedPrecondition for all others.
+func TestPauseSteps_CheckPrerequisite(t *testing.T) {
+	tests := []struct {
+		name string
+		step WorkflowStep[*PauseInput, *PauseState]
+		// allowed lists the statuses CheckPrerequisite accepts; nil means
+		// every status is accepted.
+		allowed map[ateapipb.Actor_Status]bool
+	}{
+		{
+			// Loading has no prerequisite: it is allowed from every status.
+			name:    "LoadActorForPauseStep",
+			step:    &LoadActorForPauseStep{},
+			allowed: nil,
+		},
+		{
+			// Pausing is allowed only from RUNNING.
+			name: "MarkPausingStep",
+			step: &MarkPausingStep{},
+			allowed: map[ateapipb.Actor_Status]bool{
+				ateapipb.Actor_STATUS_RUNNING: true,
+			},
+		},
+		{
+			// The checkpoint call is allowed only from PAUSING (PAUSED is
+			// fast-forwarded by IsComplete).
+			name: "CallAteletPauseStep",
+			step: &CallAteletPauseStep{},
+			allowed: map[ateapipb.Actor_Status]bool{
+				ateapipb.Actor_STATUS_PAUSING: true,
+			},
+		},
+		{
+			// Finalizing is allowed only from PAUSING: a persisted PAUSED
+			// actor always has its worker pod fields cleared and is
+			// fast-forwarded by IsComplete.
+			name: "FinalizePausedStep",
+			step: &FinalizePausedStep{},
+			allowed: map[ateapipb.Actor_Status]bool{
+				ateapipb.Actor_STATUS_PAUSING: true,
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			for _, st := range allActorStatuses {
+				// Worker pod fields are populated so CallAteletPauseStep's
+				// missing-worker crash branch is not taken; this test only
+				// verifies status gating.
+				err := tc.step.CheckPrerequisite(ctx, &PauseInput{ActorName: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st, AteomPodNamespace: "ns", AteomPodName: "worker-1"}})
+				assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st])
+			}
+		})
+	}
+}
+
+// TestPauseActor_CrashesWhenPausingActorMissingWorkerPod verifies that a
+// PAUSING actor with no worker pod recorded is moved to CRASHED by
+// CallAteletPauseStep's prerequisite check and the pause fails with
+// FailedPrecondition.
+func TestPauseActor_CrashesWhenPausingActorMissingWorkerPod(t *testing.T) {
+	ctx := context.Background()
+	st, cleanup := storetest.SetupTestStore(t)
+	defer cleanup()
+	w := newTestActorWorkflow(t, st, "ns", "tmpl1")
+
+	seedWorkflowActor(t, ctx, st, "team-a", "id1", "ns", "tmpl1", ateapipb.Actor_STATUS_PAUSING)
+
+	_, err := w.PauseActor(ctx, "team-a", "id1")
+	if got := status.Code(err); got != codes.FailedPrecondition {
+		t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err)
+	}
+
+	got, err := st.GetActor(ctx, "team-a", "id1")
+	if err != nil {
+		t.Fatalf("GetActor failed: %v", err)
+	}
+	if got.GetStatus() != ateapipb.Actor_STATUS_CRASHED {
+		t.Errorf("stored status = %v, want %v", got.GetStatus(), ateapipb.Actor_STATUS_CRASHED)
+	}
+}