Fix coordinator worker integration

coordinator/internal/domain/task.go

@@ -117,8 +117,10 @@ const DefaultMaxAttempts = 3
 func (t *Task) CanRetry() bool { return t.Attempt < t.MaxAttempts }
 
 // IsLeaseHeldBy reports whether worker currently holds this task at attempt.
-func (t *Task) IsLeaseHeldBy(worker string, attempt int) bool {
-	return t.LeaseOwner != nil && *t.LeaseOwner == worker && t.Attempt == attempt
+func (t *Task) IsLeaseHeldBy(worker string, attempt int, now time.Time) bool {
+	return t.LeaseOwner != nil && t.LeaseExpiresAt != nil && now.Before(*t.LeaseExpiresAt) &&
+		*t.LeaseOwner == worker && t.Attempt == attempt &&
+		(t.Status == TaskLeased || t.Status == TaskRunning)
 }
 
 // AsClaimed projects the task into the trimmed view handed to a worker:
@@ -146,7 +148,7 @@ func (t *Task) AsClaimed() ClaimedTask {
 
 // verifyLease is the guard every worker-driven transition shares: the caller
 // must own the lease and reference the attempt it was granted.
-func (t *Task) verifyLease(worker string, attempt int) error {
+func (t *Task) verifyLease(worker string, attempt int, now time.Time) error {
 	// A task is worker-owned while leased or running: the first heartbeat moves
 	// it from leased to running, but ownership rules are identical for both.
 	if t.Status != TaskLeased && t.Status != TaskRunning {
@@ -158,13 +160,16 @@ func (t *Task) verifyLease(worker string, attempt int) error {
 	if t.Attempt != attempt {
 		return ErrStaleAttempt
 	}
+	if t.LeaseExpiresAt == nil || !now.Before(*t.LeaseExpiresAt) {
+		return ErrLeaseConflict
+	}
 	return nil
 }
 
 // RenewLease extends the lease of the worker that holds it. The first heartbeat
 // also acknowledges start, moving the task from leased to running.
-func (t *Task) RenewLease(worker string, attempt int, until time.Time) error {
-	if err := t.verifyLease(worker, attempt); err != nil {
+func (t *Task) RenewLease(worker string, attempt int, now, until time.Time) error {
+	if err := t.verifyLease(worker, attempt, now); err != nil {
 		return err
 	}
 	t.LeaseExpiresAt = &until
@@ -195,7 +200,7 @@ func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any,
 		return ErrResultConflict
 	}
 
-	if err := t.verifyLease(worker, attempt); err != nil {
+	if err := t.verifyLease(worker, attempt, now); err != nil {
 		return err
 	}
 
@@ -214,7 +219,7 @@ func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any,
 // Fail records a worker-reported failure. A retryable failure with attempts
 // left returns the task to the queue; otherwise it terminates as failed.
 func (t *Task) Fail(worker string, attempt int, code, message string, retryable bool, now time.Time) error {
-	if err := t.verifyLease(worker, attempt); err != nil {
+	if err := t.verifyLease(worker, attempt, now); err != nil {
 		return err
 	}
 	t.ErrorCode = &code

coordinator/internal/domain/task_test.go

@@ -172,14 +172,14 @@ func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
 	task := leasedTask(1, 3)
 	until := testLater.Add(time.Hour)
 
-	if err := task.RenewLease(testWorker, 1, until); err != nil {
+	if err := task.RenewLease(testWorker, 1, testNow, until); err != nil {
 		t.Fatal(err)
 	}
 	if task.Status != TaskRunning {
 		t.Errorf("status = %q, want running after first heartbeat", task.Status)
 	}
 	// A second heartbeat keeps it running.
-	if err := task.RenewLease(testWorker, 1, until); err != nil {
+	if err := task.RenewLease(testWorker, 1, testNow, until); err != nil {
 		t.Fatal(err)
 	}
 	if task.Status != TaskRunning {
@@ -190,15 +190,15 @@ func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
 func TestRunningTaskCanBeCompletedAndExpired(t *testing.T) {
 	// Complete works from running.
 	task := leasedTask(1, 3)
-	_ = task.RenewLease(testWorker, 1, testLater) // -> running
+	_ = task.RenewLease(testWorker, 1, testNow, testLater) // -> running
 	if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil {
 		t.Errorf("complete from running: %v", err)
 	}
 
 	// Expire reclaims a running task too.
 	task2 := leasedTask(1, 3)
-	_ = task2.RenewLease(testWorker, 1, testLater) // -> running
-	task2.ExpireLease(testNow)
+	_ = task2.RenewLease(testWorker, 1, testNow, testLater) // -> running
+	task2.ExpireLease(testLater)
 	if task2.Status != TaskPending {
 		t.Errorf("status = %q, want pending after a running lease expires", task2.Status)
 	}
@@ -208,14 +208,29 @@ func TestRenewLeaseExtendsOnlyForHolder(t *testing.T) {
 	task := leasedTask(1, 3)
 	until := testLater.Add(time.Hour)
 
-	if err := task.RenewLease(testWorker, 1, until); err != nil {
+	if err := task.RenewLease(testWorker, 1, testNow, until); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 	if !task.LeaseExpiresAt.Equal(until) {
 		t.Error("lease must be extended")
 	}
 
-	if err := task.RenewLease("worker-2", 1, until); !errors.Is(err, ErrLeaseConflict) {
+	if err := task.RenewLease("worker-2", 1, testNow, until); !errors.Is(err, ErrLeaseConflict) {
 		t.Errorf("err = %v, want ErrLeaseConflict", err)
 	}
 }
+
+func TestExpiredLeaseRejectsRenewalCompletionAndFailure(t *testing.T) {
+	task := leasedTask(1, 3)
+	expired := testLater.Add(time.Nanosecond)
+
+	if err := task.RenewLease(testWorker, 1, expired, expired.Add(time.Minute)); !errors.Is(err, ErrLeaseConflict) {
+		t.Errorf("renew expired lease: err = %v, want ErrLeaseConflict", err)
+	}
+	if err := task.CompleteWith(testResult, nil, testWorker, 1, expired); !errors.Is(err, ErrLeaseConflict) {
+		t.Errorf("complete expired lease: err = %v, want ErrLeaseConflict", err)
+	}
+	if err := task.Fail(testWorker, 1, "timeout", "expired", true, expired); !errors.Is(err, ErrLeaseConflict) {
+		t.Errorf("fail expired lease: err = %v, want ErrLeaseConflict", err)
+	}
+}