diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index 5c193f992..496e1bb89 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/cmd/ateapi/internal/controlapi/workflow.go @@ -158,7 +158,7 @@ func (w *ActorWorkflow) ResumeActor(ctx context.Context, atespace, name string, // Acquire lock and get the timeout context for the workflow // Lock TTL is 30 seconds, with 2 seconds padding for workflow timeout - ctx, releaseLock, err := w.acquireActorLock(ctx, name, 30*time.Second, 2*time.Second) + ctx, releaseLock, err := w.acquireActorLock(ctx, atespace, name, 30*time.Second, 2*time.Second) if err != nil { return nil, err } @@ -188,7 +188,7 @@ func (w *ActorWorkflow) SuspendActor(ctx context.Context, atespace, name string) // Acquire lock and get the timeout context for the workflow // Lock TTL is 30 seconds, with 2 seconds padding for workflow timeout - ctx, releaseLock, err := w.acquireActorLock(ctx, name, 30*time.Second, 2*time.Second) + ctx, releaseLock, err := w.acquireActorLock(ctx, atespace, name, 30*time.Second, 2*time.Second) if err != nil { return nil, err } @@ -218,7 +218,7 @@ func (w *ActorWorkflow) PauseActor(ctx context.Context, atespace, name string) ( // Acquire lock and get the timeout context for the workflow // Lock TTL is 30 seconds, with 2 seconds padding for workflow timeout - ctx, releaseLock, err := w.acquireActorLock(ctx, name, 30*time.Second, 2*time.Second) + ctx, releaseLock, err := w.acquireActorLock(ctx, atespace, name, 30*time.Second, 2*time.Second) if err != nil { return nil, err } @@ -238,8 +238,8 @@ func (w *ActorWorkflow) PauseActor(ctx context.Context, atespace, name string) ( return state.Actor, nil } -func (w *ActorWorkflow) acquireActorLock(ctx context.Context, name string, ttl time.Duration, padding time.Duration) (context.Context, func(), error) { - lockKey := "lock:actor:" + name +func (w *ActorWorkflow) acquireActorLock(ctx context.Context, atespace, name string, ttl time.Duration, padding time.Duration) (context.Context, func(), error) { + lockKey := "lock:actor:" + atespace + ":" + name lockValue := uuid.New().String() // Create a child context for the workflow that expires BEFORE the lock diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 14c28ca5f..5ac831ce2 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -186,7 +186,7 @@ func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, sta // Only free it if it still belongs to us if wass := worker.Assignment; wass != nil { - if wass.Actor.Name == input.ActorName { + if wass.Actor.Atespace == input.Atespace && wass.Actor.Name == input.ActorName { worker.Assignment = nil err = s.store.UpdateWorker(ctx, worker, worker.Version) if err != nil { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 7cd768ac4..da7d742ba 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -132,7 +132,7 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat if worker.Assignment == nil { continue } - if worker.Assignment.Actor.Name != input.ActorName { + if worker.Assignment.Actor.Atespace != input.Atespace || worker.Assignment.Actor.Name != input.ActorName { continue } eligible, err := isWorkerEligibleForActor(worker, state.ActorTemplate.Spec.SandboxClass, state.ActorTemplate.Spec.WorkerSelector, state.Actor.GetWorkerSelector()) diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index b87bd2305..4bb4b3980 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -15,10 +15,15 @@ package controlapi import ( + "context" "testing" + "time" + "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" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) @@ -162,3 +167,52 @@ func TestIsWorkerEligibleForActor(t *testing.T) { }) } } + +func TestAssignWorkerStep_SkipsWorkerAssignedInOtherAtespace(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + // The only worker is held by a same-named actor in another atespace. It is + // eligible for the template, so a name-only match would adopt it. + worker := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool", + WorkerPod: "pod-1", + SandboxClass: "gvisor", + Assignment: &ateapipb.Assignment{ + Actor: &ateapipb.ActorRef{Atespace: "team-b", Name: "shared"}, + }, + } + if err := persistence.CreateWorker(ctx, worker); err != nil { + t.Fatalf("CreateWorker: %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: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "shared"}, + }, + ActorTemplate: &atev1alpha1.ActorTemplate{ + Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}, + }, + } + err := step.Execute(ctx, &ResumeInput{ActorName: "shared", Atespace: "team-a"}, state) + if status.Code(err) != codes.FailedPrecondition { + t.Fatalf("Execute() error = %v, want FailedPrecondition (no free workers)", err) + } + + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if got := stored.GetAssignment().GetActor().GetAtespace(); got != "team-b" { + t.Errorf("worker assignment atespace = %q, want %q (assignment: %v)", got, "team-b", stored.GetAssignment()) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index cacb8c307..28297d1fd 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -184,7 +184,7 @@ func (s *FinalizeSuspendedStep) Execute(ctx context.Context, input *SuspendInput } else { // Only free it if it still belongs to us if wass := worker.Assignment; wass != nil { - if wass.Actor.Name == input.ActorName { + if wass.Actor.Atespace == input.Atespace && wass.Actor.Name == input.ActorName { worker.Assignment = nil err = s.store.UpdateWorker(ctx, worker, worker.Version) if err != nil { diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go new file mode 100644 index 000000000..84e8083c8 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -0,0 +1,103 @@ +// 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" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" +) + +// newTestPersistence returns a store backed by a throwaway miniredis. +func newTestPersistence(t *testing.T) store.Interface { + t.Helper() + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("failed to start miniredis: %v", err) + } + t.Cleanup(mr.Close) + rdb := redis.NewClusterClient(&redis.ClusterOptions{Addrs: []string{mr.Addr()}}) + t.Cleanup(func() { rdb.Close() }) //nolint:errcheck // test cleanup + return ateredis.NewPersistence(rdb) +} + +func TestFinalizeSuspendedStep_ReleasesOnlyOwnWorker(t *testing.T) { + tests := []struct { + name string + assignmentAtespace string + wantReleased bool + }{ + { + name: "frees worker assigned to this actor", + assignmentAtespace: "team-a", + wantReleased: true, + }, + { + name: "keeps worker assigned to same-named actor in another atespace", + assignmentAtespace: "team-b", + wantReleased: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + worker := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool", + WorkerPod: "pod-1", + Assignment: &ateapipb.Assignment{ + Actor: &ateapipb.ActorRef{Atespace: tt.assignmentAtespace, Name: "shared"}, + }, + } + if err := persistence.CreateWorker(ctx, worker); err != nil { + t.Fatalf("CreateWorker: %v", err) + } + + actor := &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "shared"}, + Status: ateapipb.Actor_STATUS_SUSPENDING, + AteomPodNamespace: "worker-ns", + AteomPodName: "pod-1", + WorkerPoolName: "pool", + InProgressSnapshot: "gs://snapshots/shared/1", + } + if _, err := persistence.CreateActor(ctx, actor); err != nil { + t.Fatalf("CreateActor: %v", err) + } + + step := &FinalizeSuspendedStep{store: persistence} + input := &SuspendInput{ActorName: "shared", Atespace: "team-a"} + if err := step.Execute(ctx, input, &SuspendState{}); err != nil { + t.Fatalf("Execute: %v", err) + } + + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if released := stored.GetAssignment() == nil; released != tt.wantReleased { + t.Errorf("worker released = %t, want %t (assignment: %v)", released, tt.wantReleased, stored.GetAssignment()) + } + }) + } +}