Preserve leases needing recovery

cmd/orchestra-worker/main.go

@@ -814,6 +814,15 @@ func (w *worker) once(ctx context.Context) error {
 				w.leases[e.TaskID] = l
 			}
 		}
+		if e.Type == "TaskNeedsAttention" {
+			// The diagnostic event increments the aggregate version but leaves
+			// ownership intact. Keep our locally persisted expected version in
+			// sync so a late, otherwise valid completion is not self-staled.
+			if l, ok := w.leases[e.TaskID]; ok {
+				l.Version = e.Version
+				w.leases[e.TaskID] = l
+			}
+		}
 		if e.Type == "TaskCompleted" {
 			delete(w.leases, e.TaskID)
 			if session, active := w.sessions[e.TaskID]; active {
@@ -932,7 +941,7 @@ func (w *worker) reconcileLeases(ctx context.Context) error {
 	active := make(map[string]lease)
 	for _, task := range tasks {
 		w.tasks[task.ID] = task
-		if task.State == domain.StateLeased && task.Lease != nil && task.Lease.HarnessID == w.harnessID {
+		if (task.State == domain.StateLeased || task.State == domain.StateNeedsAttention) && task.Lease != nil && task.Lease.HarnessID == w.harnessID {
 			active[task.ID] = lease{Epoch: task.Lease.Epoch, HandoffRef: task.HandoffRef, TransactionID: task.ReleaseTransaction, AnchorSHA: task.ReleaseAnchor, Version: task.Version, Until: task.Lease.Until}
 		}
 	}