Make worker assignment checks atespace and actor name, also updated locks. (#411)

Current we only check & lock on actor name.

- [x] Tests pass
- [x] Appropriate changes to documentation are included in the PR

Author: HavenXia

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

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 {

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())

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())
+	}
+}

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 {

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())
+			}
+		})
+	}
+}