diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 4499ac6a..3764a468 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -157,8 +157,11 @@ 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 + // The workflow is done once the worker is freed and the actor reached PAUSED, + // or CRASHED (node name was lost, so it can never be safely resumed). + status := state.Actor.GetStatus() + terminal := status == ateapipb.Actor_STATUS_PAUSED || status == ateapipb.Actor_STATUS_CRASHED + return terminal && state.Actor.GetAteomPodNamespace() == "", nil } func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) @@ -202,19 +205,24 @@ func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, sta return err } latestActor.Status = ateapipb.Actor_STATUS_PAUSED - // TODO(dberkov) - what if we still don't know the node name? Maybe move to CRASHED status? if nodeName == "" { - slog.Warn("Node name not found during finalize pause", "actor", input.ActorName) + // Without a node name we cannot record where the local snapshot lives, + // so the actor can never be resumed (findFreeWorker would search for a + // worker on an unknown node forever). Crash it instead of leaving it + // stuck in PAUSED. + slog.Warn("Node name not found during finalize pause, crashing actor", "actor", input.ActorName) + latestActor.Status = ateapipb.Actor_STATUS_CRASHED } // TODO(dberkov) - what if InProgressSnapshot is empty? That shouldn't be possible. if latestActor.InProgressSnapshot != "" { + localInfo := &ateapipb.LocalSnapshotInfo{ + SnapshotPrefix: latestActor.InProgressSnapshot, + } + if nodeName != "" { + localInfo.NodeVmsWithLocalSnapshots = []string{nodeName} + } latestActor.LatestSnapshotInfo = &ateapipb.SnapshotInfo{ - Data: &ateapipb.SnapshotInfo_Local{ - Local: &ateapipb.LocalSnapshotInfo{ - SnapshotPrefix: latestActor.InProgressSnapshot, - NodeVmsWithLocalSnapshots: []string{nodeName}, - }, - }, + Data: &ateapipb.SnapshotInfo_Local{Local: localInfo}, } latestActor.InProgressSnapshot = "" } 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 00000000..daa1b31a --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_pause_test.go @@ -0,0 +1,116 @@ +// 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" +) + +// TestFinalizePausedStep_WorkerGone reproduces the scenario where the worker pod +// disappears from the DB during pause finalization, so the node it ran on is +// unknown. +// +// Old behavior: NodeVmsWithLocalSnapshots = []string{""}, which made +// findFreeWorker search for a worker with node name "", never found, a +// permanent "no free workers available" on resume. +// +// Current behavior: NodeVmsWithLocalSnapshots is left nil, and the actor is +// crashed instead of left PAUSED, since a local snapshot with an unknown node +// can never be safely resumed. +func TestFinalizePausedStep_WorkerGone(t *testing.T) { + st, cleanup := storetest.SetupTestStore(t) + defer cleanup() + + ctx := context.Background() + const atespace, actorName = "team-a", "actor-1" + + actor := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, Name: actorName}, + Status: ateapipb.Actor_STATUS_PAUSING, + AteomPodNamespace: "default", + AteomPodName: "worker-pod-1", + WorkerPoolName: "pool1", + InProgressSnapshot: "snap-prefix", + } + if _, err := st.CreateActor(ctx, actor); err != nil { + t.Fatalf("CreateActor: %v", err) + } + // Intentionally NOT creating the worker in store, simulates worker already gone. + + step := &FinalizePausedStep{store: st} + input := &PauseInput{Atespace: atespace, ActorName: actorName} + state := &PauseState{} + if err := step.Execute(ctx, input, state); err != nil { + t.Fatalf("Execute: %v", err) + } + + got, err := st.GetActor(ctx, atespace, actorName) + if err != nil { + t.Fatalf("GetActor: %v", err) + } + + if got.GetStatus() != ateapipb.Actor_STATUS_CRASHED { + t.Errorf("status = %v, want CRASHED (node name unknown, cannot resume safely)", got.GetStatus()) + } + for _, n := range got.GetLatestSnapshotInfo().GetLocal().GetNodeVmsWithLocalSnapshots() { + if n == "" { + t.Errorf("BUG: empty string in NodeVmsWithLocalSnapshots, findFreeWorker would never match") + } + } + + state.Actor = got + done, err := step.IsComplete(ctx, input, state) + if err != nil { + t.Fatalf("IsComplete: %v", err) + } + if !done { + t.Error("IsComplete = false, want true once the actor is CRASHED and the worker is freed") + } +} + +// TestFindFreeWorker_EmptyNodeRestriction shows the root symptom the fix +// avoids: old code wrote []string{""} into NodeVmsWithLocalSnapshots when the +// node name was unknown, and findFreeWorker required worker.NodeName == "", +// which never matches a real worker. +func TestFindFreeWorker_EmptyNodeRestriction(t *testing.T) { + workers := []*ateapipb.Worker{ + {WorkerNamespace: "default", WorkerPool: "pool1", WorkerPod: "w1", NodeName: "node1"}, + {WorkerNamespace: "default", WorkerPool: "pool1", WorkerPod: "w2", NodeName: "node2"}, + } + + s := &AssignWorkerStep{} + + // Old behavior: []string{""}, no worker has NodeName == "", returns nil. + got, err := s.findFreeWorker(workers, "", nil, nil, []string{""}) + if err != nil { + t.Fatalf("findFreeWorker: %v", err) + } + if got != nil { + t.Errorf("expected nil with old buggy input, got %v", got) + } + + // Fixed behavior: nil restrictions, any free worker matches. + got, err = s.findFreeWorker(workers, "", nil, nil, nil) + if err != nil { + t.Fatalf("findFreeWorker: %v", err) + } + if got == nil { + t.Error("expected a worker with nil restrictions, got nil") + } +}