diff --git a/cmd/ateapi/internal/controlapi/functional_test.go b/cmd/ateapi/internal/controlapi/functional_test.go index 57c08a9c5..29e87acdf 100644 --- a/cmd/ateapi/internal/controlapi/functional_test.go +++ b/cmd/ateapi/internal/controlapi/functional_test.go @@ -1731,7 +1731,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. @@ -1740,8 +1741,9 @@ func TestUpdateActor_NotFound(t *testing.T) { // worker is claimed, leaving worker-a's actor_id 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_id cleared) rather than left dangling. +// 4. Resumes again; asserts it fails and the actor is CRASHED, that +// worker-a has been released (actor_id cleared) rather than left +// dangling, and that worker-b was never claimed. func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T) { ns := namespaceForTest("ns-resume-release-stale") tc := setupTest(t, ns) @@ -1780,19 +1782,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{}) @@ -1813,16 +1812,12 @@ func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T) t.Errorf("expected worker-a (now-ineligible pool-a) to be released, got actor_id=%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_id=%q", name, got) - } + if wass := w.Assignment; wass != nil { + got := "" + if wass.Actor != nil { + got = wass.Actor.Name } + t.Errorf("expected worker-b to remain unclaimed (actor is crashed, not migrated), got actor_id=%q", got) } } } diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index 496e1bb89..099e9be5b 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/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) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 4499ac6a5..8ecf797aa 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/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) { @@ -157,8 +171,13 @@ type FinalizePausedStep struct { func (s *FinalizePausedStep) Name() string { return "FinalizePaused" } func (s *FinalizePausedStep) IsComplete(ctx context.Context, input *PauseInput, state *PauseState) (bool, error) { - // The workflow is completely done ONLY if the status is PAUSED *and* we've successfully freed the worker. - return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED && state.Actor.GetAteomPodNamespace() == "", nil + return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, 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) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause_test.go b/cmd/ateapi/internal/controlapi/workflow_pause_test.go new file mode 100644 index 000000000..176b8a52d --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_pause_test.go @@ -0,0 +1,172 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "testing" + + "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" +) + +// 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) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 3b7c9ab5f..3f4bb30e9 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -48,6 +48,7 @@ type ResumeInput struct { // ResumeState holds the mutable state loaded and modified during execution. type ResumeState struct { Actor *ateapipb.Actor + Worker *ateapipb.Worker ActorTemplate *atev1alpha1.ActorTemplate } @@ -61,6 +62,9 @@ func (s *LoadActorForResumeStep) IsComplete(ctx context.Context, input *ResumeIn // Always run this step to get the latest state from the DB return false, nil } +func (s *LoadActorForResumeStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + return nil +} func (s *LoadActorForResumeStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { @@ -77,6 +81,36 @@ func (s *LoadActorForResumeStep) Execute(ctx context.Context, input *ResumeInput } state.ActorTemplate = actorTemplate + // If the Actor is in Resuming state, it means a previous attempt crashed after AssignWorkerStep. + // We don't need to repeat the AssignWorkerStep, load the Worker now. + if actor.Status == ateapipb.Actor_STATUS_RESUMING { + allPopulated := actor.AteomPodUid != "" && actor.WorkerPoolName != "" && actor.AteomPodName != "" + if !allPopulated { + slog.ErrorContext(ctx, "expected all of AteomPodUid, WorkerPoolName and AteomPodName to be populated, found", + slog.String("AteomPodUid", actor.AteomPodUid), + slog.String("WorkerPoolName", actor.WorkerPoolName), + slog.String("AteomPodName", actor.AteomPodName)) + + // Crash the actor if its worker assignment is corrupted. We should never be in this state. + if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil { + return cerr + } + return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName) + } + + wk, err := s.store.GetWorker(ctx, actor.AteomPodNamespace, actor.WorkerPoolName, actor.AteomPodName) + if err != nil { + // Crash the actor if it was assigned to a deleted pod. + if errors.Is(err, ErrWorkerPodNotFound) { + if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil { + return cerr + } + return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName) + } + return fmt.Errorf("failed to get already assigned worker for actor %w", err) + } + state.Worker = wk + } return nil } @@ -113,21 +147,28 @@ type AssignWorkerStep struct { func (s *AssignWorkerStep) Name() string { return "AssignWorker" } func (s *AssignWorkerStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) { - return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil + // RESUMING means a previous attempt already assigned a worker (loaded by + // LoadActorForResumeStep); RUNNING is past this step entirely. + return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RESUMING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil } +func (s *AssignWorkerStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + switch state.Actor.GetStatus() { + case ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED: + return nil + default: + return status.Errorf(codes.FailedPrecondition, "AssignWorkerStep prerequisite not met for Actor: %s (got: %v, want %s or %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED) + } +} + func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { workers, err := s.workerCache.Workers() if err != nil { return fmt.Errorf("while listing workers: %w", err) } - var assignedWorker *ateapipb.Worker - // Check if we already have a worker assigned from a previous failed attempt. - // If that worker's pool is no longer eligible (e.g. the actor's - // worker_selector was updated after the failed attempt), release it back - // to the free pool instead of leaving it claimed forever — nothing else - // reclaims a healthy worker whose actor moved on to a different pool. + // This can happen if ateapi crashed after updating worker with actor assignment, + // but has not yet updated the actor. for _, worker := range workers { if worker.Assignment == nil { continue @@ -140,20 +181,28 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat return fmt.Errorf("while checking worker eligibility: %w", err) } if eligible { - assignedWorker = worker + state.Worker = worker break } // Workers() returns pointers directly from the cache so we need to clone before // mutating so that the cache is not corrupted if UpdateWorker fails. - release := proto.Clone(worker).(*ateapipb.Worker) - release.Assignment = nil - if err := s.store.UpdateWorker(ctx, release, release.Version); err != nil { - return fmt.Errorf("while releasing stale worker assignment: %w", err) - } + releaseWorker := proto.Clone(worker).(*ateapipb.Worker) + releaseWorker.Assignment = nil + // The claimed worker is no longer eligible (e.g. the actor's + // worker_selector changed after the failed attempt); release it back + // to the free pool — nothing else reclaims a healthy worker whose + // actor moved on to a different pool. Best effort in the background. + go func(release *ateapipb.Worker) { + bgCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + if err := s.store.UpdateWorker(bgCtx, release, release.Version); err != nil { + slog.ErrorContext(bgCtx, "Failed to release stale worker assignment", + slog.String("worker", release.GetWorkerNamespace()+"/"+release.GetWorkerPod()), + slog.Any("err", err)) + } + }(releaseWorker) } - - // If not, find a free one using randomized shuffling - if assignedWorker == nil { + if state.Worker == nil { pickedWorker, err := s.findFreeWorker(workers, state.ActorTemplate.Spec.SandboxClass, state.ActorTemplate.Spec.WorkerSelector, state.Actor.GetWorkerSelector(), state.Actor.GetLatestSnapshotInfo().GetLocal().GetNodeVmsWithLocalSnapshots()) if err != nil { return err @@ -162,14 +211,14 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat return status.Errorf(codes.FailedPrecondition, "no free workers available") } - assignedWorker = pickedWorker + state.Worker = pickedWorker slog.InfoContext(ctx, "Picked worker", slog.Any("worker", pickedWorker.String())) } // Workers() returns pointers directly from the cache so we need to clone before // mutating so that the cache is not corrupted if UpdateWorker fails. - assignedWorker = proto.Clone(assignedWorker).(*ateapipb.Worker) - assignedWorker.Assignment = &ateapipb.Assignment{ + state.Worker = proto.Clone(state.Worker).(*ateapipb.Worker) + state.Worker.Assignment = &ateapipb.Assignment{ ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ Namespace: state.Actor.GetActorTemplateNamespace(), Name: state.Actor.GetActorTemplateName(), @@ -180,16 +229,16 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat }, } - if err := s.store.UpdateWorker(ctx, assignedWorker, assignedWorker.Version); err != nil { + if err := s.store.UpdateWorker(ctx, state.Worker, state.Worker.Version); err != nil { return err } state.Actor.Status = ateapipb.Actor_STATUS_RESUMING - state.Actor.AteomPodNamespace = assignedWorker.GetWorkerNamespace() - state.Actor.AteomPodName = assignedWorker.GetWorkerPod() - state.Actor.AteomPodIp = assignedWorker.GetIp() - state.Actor.AteomPodUid = assignedWorker.GetWorkerPodUid() - state.Actor.WorkerPoolName = assignedWorker.GetWorkerPool() + state.Actor.AteomPodNamespace = state.Worker.GetWorkerNamespace() + state.Actor.AteomPodName = state.Worker.GetWorkerPod() + state.Actor.AteomPodIp = state.Worker.GetIp() + state.Actor.AteomPodUid = state.Worker.GetWorkerPodUid() + state.Actor.WorkerPoolName = state.Worker.GetWorkerPool() updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion()) if err != nil { @@ -254,6 +303,46 @@ func (s *CallAteletRestoreStep) Name() string { return "CallAteletRestore" } func (s *CallAteletRestoreStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) { return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil } +func (s *CallAteletRestoreStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING { + return status.Errorf(codes.FailedPrecondition, "CallAteletRestoreStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) + } + if state.Worker == nil { + return status.Errorf(codes.FailedPrecondition, "Assigned worker is nil") + } + // Verify if the worker is still assigned to the same Actor. + assigned := state.Worker.GetAssignment().GetActor() + if assigned.GetAtespace() != input.Atespace || assigned.GetName() != input.ActorName { + slog.ErrorContext(ctx, "crashing actor because its assigned worker no longer belongs to it", + slog.String("worker", state.Worker.GetWorkerPod()), + slog.Any("assignment", state.Worker.GetAssignment())) + if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil { + return fmt.Errorf("while crashing actor: %w", cerr) + } + return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName) + } + eligible, err := isWorkerEligibleForActor(state.Worker, state.ActorTemplate.Spec.SandboxClass, state.ActorTemplate.Spec.WorkerSelector, state.Actor.GetWorkerSelector()) + if err != nil { + return fmt.Errorf("while calling isWorkerEligbleForActor :%w", err) + } + if !eligible { + slog.ErrorContext(ctx, "crashing actor because previously assigned worker is not eligible anymore") + release := proto.Clone(state.Worker).(*ateapipb.Worker) + release.Assignment = nil + // If that worker's pool is no longer eligible (e.g. the actor's + // worker_selector was updated after the failed attempt), release it back + // to the free pool instead of leaving it claimed forever — nothing else + // reclaims a healthy worker whose actor moved on to a different pool. + if err := s.store.UpdateWorker(ctx, release, release.Version); err != nil { + return fmt.Errorf("while releasing stale worker assignment: %w", err) + } + if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil { + return fmt.Errorf("while crashing actor: %w", cerr) + } + return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName) + } + return nil +} func (s *CallAteletRestoreStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName()) if err != nil { @@ -358,6 +447,12 @@ func (s *FinalizeRunningStep) Name() string { return "FinalizeRunning" } func (s *FinalizeRunningStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) { return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil } +func (s *FinalizeRunningStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING { + return status.Errorf(codes.FailedPrecondition, "FinalizeRunningStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) + } + return nil +} func (s *FinalizeRunningStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 99922fea1..32ba25171 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -19,11 +19,13 @@ import ( "testing" "time" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) @@ -216,3 +218,367 @@ func TestAssignWorkerStep_SkipsWorkerAssignedInOtherAtespace(t *testing.T) { t.Errorf("worker assignment atespace = %q, want %q (assignment: %v)", got, "team-b", stored.GetAssignment()) } } + +// TestAssignWorkerStep_ReleasesIneligibleStaleWorkerInBackground verifies +// that a worker claimed by a previous failed attempt whose pool is no longer +// eligible is released back to the free pool asynchronously, without failing +// the resume, while a fresh eligible worker is assigned. +func TestAssignWorkerStep_ReleasesIneligibleStaleWorkerInBackground(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + // stale-pod is claimed by this actor from a failed attempt but its sandbox + // class no longer matches the template; free-pod is eligible and free. + stale := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool-a", + WorkerPod: "stale-pod", + SandboxClass: "microvm", + Assignment: &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "id1"}, + }, + } + free := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool-b", + WorkerPod: "free-pod", + SandboxClass: "gvisor", + } + for _, w := range []*ateapipb.Worker{stale, free} { + if err := persistence.CreateWorker(ctx, w); err != nil { + t.Fatalf("CreateWorker(%s): %v", w.GetWorkerPod(), err) + } + } + + actor, err := persistence.CreateActor(ctx, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "id1"}, + Status: ateapipb.Actor_STATUS_SUSPENDED, + }) + if err != nil { + t.Fatalf("CreateActor: %v", err) + } + + cacheCtx, cancel := context.WithCancel(ctx) + defer cancel() + wc := workercache.New(persistence, time.Minute) + if err := wc.Start(cacheCtx); err != nil { + t.Fatalf("workercache.Start: %v", err) + } + + step := &AssignWorkerStep{store: persistence, workerCache: wc} + state := &ResumeState{ + Actor: actor, + ActorTemplate: &atev1alpha1.ActorTemplate{ + Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}, + }, + } + if err := step.Execute(ctx, &ResumeInput{ActorName: "id1", Atespace: "team-a"}, state); err != nil { + t.Fatalf("Execute() error = %v, want nil (release must not fail the resume)", err) + } + + if got := state.Worker.GetWorkerPod(); got != "free-pod" { + t.Errorf("assigned worker = %q, want %q", got, "free-pod") + } + + // The stale worker is released in the background; poll until its + // assignment is cleared. + deadline := time.Now().Add(5 * time.Second) + for { + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool-a", "stale-pod") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if stored.GetAssignment() == nil { + break + } + if time.Now().After(deadline) { + t.Fatalf("stale worker still assigned after %v: %v", 5*time.Second, stored.GetAssignment()) + } + time.Sleep(10 * time.Millisecond) + } +} + +// TestResumeActorWorkflow_RejectedAndIdempotentPaths covers the two +// short-circuit paths of the resume workflow: rejection by AssignWorkerStep's +// CheckPrerequisite and the IsComplete idempotent fast-forward. +func TestResumeActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) { + tests := []struct { + name string + seedStatus ateapipb.Actor_Status + // wantErr true means ResumeActor must fail with FailedPrecondition. + wantErr bool + // wantStatus is the stored status after the call. + wantStatus ateapipb.Actor_Status + }{ + { + // The resume edge only exists from SUSPENDED, PAUSED, and + // RESUMING; a CRASHED actor is rejected by AssignWorkerStep's + // CheckPrerequisite and its status is left untouched. + name: "crashed rejected", + seedStatus: ateapipb.Actor_STATUS_CRASHED, + wantErr: true, + wantStatus: ateapipb.Actor_STATUS_CRASHED, + }, + { + // Resuming a RUNNING actor succeeds idempotently: every step + // fast-forwards via IsComplete. + name: "already running succeeds", + seedStatus: ateapipb.Actor_STATUS_RUNNING, + wantStatus: ateapipb.Actor_STATUS_RUNNING, + }, + } + 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, func(a *ateapipb.Actor) { + a.AteomPodNamespace = "wns" + a.AteomPodName = "wpod" + a.AteomPodIp = "1.2.3.4" + a.AteomPodUid = "uid" + a.WorkerPoolName = "pool1" + }) + + actor, err := w.ResumeActor(ctx, "team-a", "id1", false) + 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("ResumeActor 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) + } + }) + } +} + +// TestResumeSteps_CheckPrerequisite verifies each resume step's +// CheckPrerequisite against every actor status: nil for the step's allowed +// statuses, FailedPrecondition for all others. +func TestResumeSteps_CheckPrerequisite(t *testing.T) { + tests := []struct { + name string + step WorkflowStep[*ResumeInput, *ResumeState] + // 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: "LoadActorForResumeStep", + step: &LoadActorForResumeStep{}, + allowed: nil, + }, + { + // Resuming is allowed from SUSPENDED and PAUSED (RESUMING and + // RUNNING are fast-forwarded by IsComplete). + name: "AssignWorkerStep", + step: &AssignWorkerStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDED: true, + ateapipb.Actor_STATUS_PAUSED: true, + }, + }, + { + // The restore call is allowed only from RESUMING (RUNNING is + // fast-forwarded by IsComplete). + name: "CallAteletRestoreStep", + step: &CallAteletRestoreStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RESUMING: true, + }, + }, + { + // Finalizing transitions RESUMING -> RUNNING; RUNNING itself is + // fast-forwarded by IsComplete before the prerequisite is checked. + name: "FinalizeRunningStep", + step: &FinalizeRunningStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RESUMING: true, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + for _, st := range allActorStatuses { + // An eligible Worker assigned to this actor is provided so + // CallAteletRestoreStep's worker checks pass; this test only + // verifies status gating. + state := &ResumeState{ + Actor: &ateapipb.Actor{Status: st}, + Worker: &ateapipb.Worker{ + SandboxClass: string(atev1alpha1.SandboxClassGvisor), + Assignment: &ateapipb.Assignment{Actor: &ateapipb.ObjectRef{Name: "id1"}}, + }, + ActorTemplate: &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}}, + } + err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorName: "id1"}, state) + assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) + } + }) + } +} + +// TestResumeActor_CrashesOnCorruptWorkerAssignment verifies that a RESUMING +// actor with only some of its worker assignment fields populated is moved to +// CRASHED by LoadActorForResumeStep and the resume fails with Aborted. +func TestResumeActor_CrashesOnCorruptWorkerAssignment(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_RESUMING, func(a *ateapipb.Actor) { + a.AteomPodName = "worker-1" // AteomPodUid and WorkerPoolName left empty + }) + + _, err := w.ResumeActor(ctx, "team-a", "id1", false) + if got := status.Code(err); got != codes.Aborted { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.Aborted, 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) + } +} + +// TestCallAteletRestoreStep_CheckPrerequisite_WorkerOwnership verifies that +// the restore prerequisite only proceeds on a worker whose assignment still +// names this actor: the recovery path loads the worker by pod name only, so +// the assignment may have been cleared and the worker re-claimed by another +// actor in the meantime. On a mismatch the actor is crashed and the worker — +// which is not ours — must not be written. +func TestCallAteletRestoreStep_CheckPrerequisite_WorkerOwnership(t *testing.T) { + ownAssignment := &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "shared"}, + } + otherAssignment := &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-b", Name: "shared"}, + } + + tests := []struct { + name string + sandboxClass string + assignment *ateapipb.Assignment + // wantCode is codes.OK when CheckPrerequisite must return nil. + wantCode codes.Code + wantActorStatus ateapipb.Actor_Status + // wantAssignment is the assignment expected on the stored worker + // afterwards; wantWorkerWrite false additionally asserts the worker + // version did not move (no write at all). + wantAssignment *ateapipb.Assignment + wantWorkerWrite bool + }{ + { + name: "crashes actor and leaves worker untouched when assigned to another actor", + sandboxClass: "gvisor", + assignment: otherAssignment, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: otherAssignment, + }, + { + name: "crashes actor and leaves worker untouched when assignment is cleared", + sandboxClass: "gvisor", + assignment: nil, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: nil, + }, + { + name: "passes for own eligible worker", + sandboxClass: "gvisor", + assignment: ownAssignment, + wantCode: codes.OK, + wantActorStatus: ateapipb.Actor_STATUS_RESUMING, + wantAssignment: ownAssignment, + }, + { + name: "releases own ineligible worker and crashes actor", + sandboxClass: "microvm", + assignment: ownAssignment, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: nil, + wantWorkerWrite: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool", + WorkerPod: "pod-1", + SandboxClass: tt.sandboxClass, + Assignment: tt.assignment, + }); err != nil { + t.Fatalf("CreateWorker: %v", err) + } + // Re-fetch so state.Worker carries the stored version (needed by + // the release path's optimistic update). + seeded, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + + seedWorkflowActor(t, ctx, persistence, "team-a", "shared", "ns", "tmpl1", ateapipb.Actor_STATUS_RESUMING) + + step := &CallAteletRestoreStep{store: persistence} + state := &ResumeState{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "shared"}, + Status: ateapipb.Actor_STATUS_RESUMING, + }, + Worker: seeded, + ActorTemplate: &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}}, + } + err = step.CheckPrerequisite(ctx, &ResumeInput{Atespace: "team-a", ActorName: "shared"}, state) + if got := status.Code(err); got != tt.wantCode { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, tt.wantCode, err) + } + + actor, err := persistence.GetActor(ctx, "team-a", "shared") + if err != nil { + t.Fatalf("GetActor: %v", err) + } + if actor.GetStatus() != tt.wantActorStatus { + t.Errorf("stored actor status = %v, want %v", actor.GetStatus(), tt.wantActorStatus) + } + + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if !proto.Equal(stored.GetAssignment(), tt.wantAssignment) { + t.Errorf("stored worker assignment = %v, want %v", stored.GetAssignment(), tt.wantAssignment) + } + if !tt.wantWorkerWrite && stored.GetVersion() != seeded.GetVersion() { + t.Errorf("worker version moved %d -> %d, want no write", seeded.GetVersion(), stored.GetVersion()) + } + }) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index 6a4cc6cc4..095676618 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -28,6 +28,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" ) @@ -53,6 +55,9 @@ func (s *LoadActorForSuspendStep) IsComplete(ctx context.Context, input *Suspend // Always run to get the freshest state return false, nil } +func (s *LoadActorForSuspendStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + return nil +} func (s *LoadActorForSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { @@ -80,11 +85,13 @@ func (s *MarkSuspendingStep) IsComplete(ctx context.Context, input *SuspendInput // Fast forward if we've already marked our intent or if we are further along. return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil } -func (s *MarkSuspendingStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { +func (s *MarkSuspendingStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - return nil + return status.Errorf(codes.FailedPrecondition, "MarkSuspendingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) } - + return nil +} +func (s *MarkSuspendingStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { state.Actor.Status = ateapipb.Actor_STATUS_SUSPENDING snapshotID := time.Now().Format(time.RFC3339) + "-" + rand.Text() state.Actor.InProgressSnapshot = strings.TrimSuffix(state.ActorTemplate.Spec.SnapshotsConfig.Location, "/") + "/" + input.ActorName + "/" + snapshotID @@ -108,14 +115,19 @@ func (s *CallAteletSuspendStep) IsComplete(ctx context.Context, input *SuspendIn // If we are already SUSPENDED, we've already called Atelet return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil } -func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { +func (s *CallAteletSuspendStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { + return status.Errorf(codes.FailedPrecondition, "CallAteletSuspendStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + } 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 SUSPENDING state but has no active worker") } - + return nil +} +func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName()) if err != nil { if errors.Is(err, ErrWorkerPodNotFound) { @@ -159,8 +171,13 @@ type FinalizeSuspendedStep struct { func (s *FinalizeSuspendedStep) Name() string { return "FinalizeSuspended" } func (s *FinalizeSuspendedStep) IsComplete(ctx context.Context, input *SuspendInput, state *SuspendState) (bool, error) { - // The workflow is completely done ONLY if the status is SUSPENDED *and* we've successfully freed the worker. - return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED && state.Actor.GetAteomPodNamespace() == "", nil + return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil +} +func (s *FinalizeSuspendedStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { + return status.Errorf(codes.FailedPrecondition, "FinalizeSuspendedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + } + return nil } func (s *FinalizeSuspendedStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index 87c52ba2e..a3995793b 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -20,11 +20,161 @@ import ( "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/alicebob/miniredis/v2" "github.com/redis/go-redis/v9" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) +// TestSuspendActorWorkflow_RejectedAndIdempotentPaths covers the two +// short-circuit paths of the suspend workflow: rejection by +// MarkSuspendingStep's CheckPrerequisite and the IsComplete idempotent +// fast-forward. +func TestSuspendActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) { + tests := []struct { + name string + seedStatus ateapipb.Actor_Status + // wantErr true means SuspendActor must fail with FailedPrecondition. + wantErr bool + // wantStatus is the stored status after the call. + wantStatus ateapipb.Actor_Status + }{ + { + // The state machine's PAUSED->SUSPENDED commit edge is rejected + // (suspending needs a live worker to checkpoint from) and the + // actor's status is left untouched. + name: "paused rejected", + seedStatus: ateapipb.Actor_STATUS_PAUSED, + wantErr: true, + wantStatus: ateapipb.Actor_STATUS_PAUSED, + }, + { + // Suspending a SUSPENDED actor succeeds idempotently via + // IsComplete fast-forward without calling atelet. + name: "already suspended succeeds", + seedStatus: ateapipb.Actor_STATUS_SUSPENDED, + wantStatus: ateapipb.Actor_STATUS_SUSPENDED, + }, + } + 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.SuspendActor(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("SuspendActor 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) + } + }) + } +} + +// TestSuspendSteps_CheckPrerequisite verifies each suspend step's +// CheckPrerequisite against every actor status: nil for the step's allowed +// statuses, FailedPrecondition for all others. +func TestSuspendSteps_CheckPrerequisite(t *testing.T) { + tests := []struct { + name string + step WorkflowStep[*SuspendInput, *SuspendState] + // 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: "LoadActorForSuspendStep", + step: &LoadActorForSuspendStep{}, + allowed: nil, + }, + { + // Suspending is allowed only from RUNNING. + name: "MarkSuspendingStep", + step: &MarkSuspendingStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RUNNING: true, + }, + }, + { + // The checkpoint call is allowed only from SUSPENDING (SUSPENDED + // is fast-forwarded by IsComplete). + name: "CallAteletSuspendStep", + step: &CallAteletSuspendStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDING: true, + }, + }, + { + // Finalizing is allowed only from SUSPENDING: a persisted + // SUSPENDED actor always has its worker pod fields cleared and is + // fast-forwarded by IsComplete. + name: "FinalizeSuspendedStep", + step: &FinalizeSuspendedStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDING: 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 CallAteletSuspendStep's + // missing-worker crash branch is not taken; this test only + // verifies status gating. + err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorName: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st, AteomPodNamespace: "ns", AteomPodName: "worker-1"}}) + assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) + } + }) + } +} + +// TestSuspendActor_CrashesWhenSuspendingActorMissingWorkerPod verifies that a +// SUSPENDING actor with no worker pod recorded is moved to CRASHED by +// CallAteletSuspendStep's prerequisite check and the suspend fails. +func TestSuspendActor_CrashesWhenSuspendingActorMissingWorkerPod(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_SUSPENDING) + + if _, err := w.SuspendActor(ctx, "team-a", "id1"); err == nil { + t.Fatal("SuspendActor succeeded, want error for SUSPENDING actor with no worker pod") + } + + 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) + } +} + // newTestPersistence returns a store backed by a throwaway miniredis. func newTestPersistence(t *testing.T) store.Interface { t.Helper() diff --git a/cmd/ateapi/internal/controlapi/workflow_test.go b/cmd/ateapi/internal/controlapi/workflow_test.go new file mode 100644 index 000000000..f4a81b728 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_test.go @@ -0,0 +1,106 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "k8s.io/apimachinery/pkg/util/wait" +) + +// stubStep records the order of engine callbacks and returns canned results. +type stubStep struct { + name string + complete bool + prereqErr error + executeErr error + calls *[]string +} + +func (s *stubStep) Name() string { return s.name } +func (s *stubStep) IsComplete(ctx context.Context, params struct{}, wCtx struct{}) (bool, error) { + *s.calls = append(*s.calls, s.name+".IsComplete") + return s.complete, nil +} +func (s *stubStep) CheckPrerequisite(ctx context.Context, params struct{}, wCtx struct{}) error { + *s.calls = append(*s.calls, s.name+".CheckPrerequisite") + return s.prereqErr +} +func (s *stubStep) Execute(ctx context.Context, params struct{}, wCtx struct{}) error { + *s.calls = append(*s.calls, s.name+".Execute") + return s.executeErr +} +func (s *stubStep) RetryBackoff() *wait.Backoff { return nil } + +func TestRunWorkflow_CheckPrerequisiteOrdering(t *testing.T) { + ctx := context.Background() + + t.Run("prerequisite checked after IsComplete and before Execute", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", calls: &calls}, + } + if err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + want := []string{"s1.IsComplete", "s1.CheckPrerequisite", "s1.Execute"} + if len(calls) != len(want) { + t.Fatalf("calls = %v, want %v", calls, want) + } + for i := range want { + if calls[i] != want[i] { + t.Fatalf("calls = %v, want %v", calls, want) + } + } + }) + + t.Run("prerequisite skipped when step is complete", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", complete: true, prereqErr: status.Error(codes.FailedPrecondition, "must not be called"), calls: &calls}, + } + if err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + for _, c := range calls { + if c == "s1.CheckPrerequisite" || c == "s1.Execute" { + t.Fatalf("unexpected call %q; calls = %v", c, calls) + } + } + }) + + t.Run("prerequisite error aborts before Execute and preserves status code", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", prereqErr: status.Error(codes.FailedPrecondition, "nope"), calls: &calls}, + &stubStep{name: "s2", calls: &calls}, + } + err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps) + if err == nil { + t.Fatal("RunWorkflow: expected error, got nil") + } + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err) + } + for _, c := range calls { + if c == "s1.Execute" || c == "s2.IsComplete" { + t.Fatalf("unexpected call %q after failed prerequisite; calls = %v", c, calls) + } + } + }) +} diff --git a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go new file mode 100644 index 000000000..c6962694a --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go @@ -0,0 +1,99 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "slices" + "testing" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + 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" +) + +// newTestActorWorkflow builds an ActorWorkflow backed by the given store and a +// lister serving one minimal ActorTemplate. Dependencies the unit tests never +// reach (worker cache, atelet dialer, k8s clients) are nil, so a step that +// unexpectedly executes against them fails the test loudly. +func newTestActorWorkflow(t *testing.T, st store.Interface, tmplNamespace, tmplName string) *ActorWorkflow { + t.Helper() + indexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{}) + if err := indexer.Add(&atev1alpha1.ActorTemplate{ + ObjectMeta: metav1.ObjectMeta{Namespace: tmplNamespace, Name: tmplName}, + }); err != nil { + t.Fatalf("add template to indexer: %v", err) + } + return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil) +} + +// seedWorkflowActor stores an actor with the given status, bound to the given +// template (pass the same tmplNamespace/tmplName as newTestActorWorkflow). +// opts mutate the actor before it is stored. +func seedWorkflowActor(t *testing.T, ctx context.Context, st store.Interface, atespace, id, tmplNamespace, tmplName string, actorStatus ateapipb.Actor_Status, opts ...func(*ateapipb.Actor)) { + t.Helper() + actor := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Name: id, Atespace: atespace}, + Status: actorStatus, + ActorTemplateNamespace: tmplNamespace, + ActorTemplateName: tmplName, + } + for _, opt := range opts { + opt(actor) + } + if _, err := st.CreateActor(ctx, actor); err != nil { + t.Fatalf("seed actor: %v", err) + } +} + +// allActorStatuses enumerates every Actor_Status value, for exhaustive +// CheckPrerequisite table tests. It is derived from the generated enum map so +// statuses added to the proto are covered automatically. +var allActorStatuses = func() []ateapipb.Actor_Status { + nums := make([]int32, 0, len(ateapipb.Actor_Status_name)) + for n := range ateapipb.Actor_Status_name { + nums = append(nums, n) + } + slices.Sort(nums) + statuses := make([]ateapipb.Actor_Status, 0, len(nums)) + for _, n := range nums { + statuses = append(statuses, ateapipb.Actor_Status(n)) + } + return statuses +}() + +// assertPrerequisiteResult verifies a CheckPrerequisite outcome for one +// status: nil when allowed, FailedPrecondition otherwise. +func assertPrerequisiteResult(t *testing.T, st ateapipb.Actor_Status, err error, wantAllowed bool) { + t.Helper() + if wantAllowed { + if err != nil { + t.Errorf("status %v: CheckPrerequisite = %v, want nil", st, err) + } + return + } + if err == nil { + t.Errorf("status %v: CheckPrerequisite = nil, want FailedPrecondition", st) + return + } + if got := status.Code(err); got != codes.FailedPrecondition { + t.Errorf("status %v: status.Code = %v, want %v", st, got, codes.FailedPrecondition) + } +}