Fix pet entities disappearing by adding setRemoveWhenFarAway(false) and chunk-load persistence

src/main/java/dansplugins/wildpets/WildPets.java

@@ -146,6 +146,7 @@ private void registerEventHandlers() {
         listeners.add(new JoinAndQuitHandler(petListRepository, ephemeralData));
         listeners.add(new MoveHandler(petListRepository));
         listeners.add(new BreedEventHandler(petListRepository, petRecordRepository, configService, ephemeralData));
+        listeners.add(new ChunkLoadHandler(petListRepository));
         EventHandlerRegistry eventHandlerRegistry = new EventHandlerRegistry();
         eventHandlerRegistry.registerEventHandlers(listeners, this);
     }

src/main/java/dansplugins/wildpets/listeners/ChunkLoadHandler.java

@@ -0,0 +1,32 @@
+package dansplugins.wildpets.listeners;
+
+import dansplugins.wildpets.pet.Pet;
+import dansplugins.wildpets.pet.list.PetListRepository;
+import org.bukkit.entity.Entity;
+import org.bukkit.event.EventHandler;
+import org.bukkit.event.Listener;
+import org.bukkit.event.world.ChunkLoadEvent;
+
+/**
+ * Handles chunk load events to ensure pet entities retain their
+ * persistence flags, preventing them from despawning.
+ *
+ * @author Daniel McCoy Stephenson
+ */
+public class ChunkLoadHandler implements Listener {
+    private final PetListRepository petListRepository;
+
+    public ChunkLoadHandler(PetListRepository petListRepository) {
+        this.petListRepository = petListRepository;
+    }
+
+    @EventHandler()
+    public void handle(ChunkLoadEvent event) {
+        for (Entity entity : event.getChunk().getEntities()) {
+            Pet pet = petListRepository.getPet(entity);
+            if (pet != null) {
+                pet.ensurePersistence();
+            }
+        }
+    }
+}

src/main/java/dansplugins/wildpets/pet/Pet.java

@@ -8,6 +8,7 @@
 import dansplugins.wildpets.location.WpLocation;
 import org.bukkit.Server;
 import org.bukkit.entity.Entity;
+import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Mob;
 import dansplugins.wildpets.data.Lockable;
 import dansplugins.wildpets.data.Savable;
@@ -123,6 +124,32 @@ public void applyAIState() {
         }
     }
 
+    /**
+     * Ensures that, for a currently loaded pet entity, persistence and "do not remove when far away"
+     * flags are applied.
+     * This helps prevent the entity from being removed by distance-based despawn mechanics while it
+     * is loaded. It does not prevent chunk unloading itself; these flags may need to be re-applied
+     * after the entity is reloaded.
+     * Returns silently if the entity is not currently loaded.
+     */
+    public void ensurePersistence() {
+        if (serverProvider == null) {
+            return;
+        }
+        Server server = serverProvider.get();
+        if (server == null) {
+            return;
+        }
+        Entity entity = server.getEntity(uniqueID);
+        if (entity == null) {
+            return;
+        }
+        entity.setPersistent(true);
+        if (entity instanceof LivingEntity) {
+            ((LivingEntity) entity).setRemoveWhenFarAway(false);
+        }
+    }
+
     public String getMovementState() {
         return movementState;
     }

src/main/java/dansplugins/wildpets/pet/list/PetList.java

@@ -4,6 +4,7 @@
 import org.bukkit.Bukkit;
 import org.bukkit.ChatColor;
 import org.bukkit.entity.Entity;
+import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
 
 import dansplugins.wildpets.config.ConfigService;
@@ -62,6 +63,9 @@ public boolean removePet(Pet petToRemove) {
         if (entity != null) {
             entity.setCustomName("");
             entity.setPersistent(false);
+            if (entity instanceof LivingEntity) {
+                ((LivingEntity) entity).setRemoveWhenFarAway(true);
+            }
             entity.setInvulnerable(false);
         }
         return getPets().remove(petToRemove);

src/main/java/dansplugins/wildpets/pet/list/PetListRepository.java

@@ -11,6 +11,7 @@
 import org.bukkit.entity.Player;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.UUID;
 
 /**
@@ -20,6 +21,7 @@ public class PetListRepository {
     private final ConfigService configService;
 
     private final ArrayList<PetList> petLists = new ArrayList<>();
+    private final HashMap<UUID, Pet> petsByEntityUUID = new HashMap<>();
 
     public PetListRepository(ConfigService configService) {
         this.configService = configService;
@@ -37,27 +39,23 @@ public boolean addNewPet(Player player, Entity entity) {
         WpLocation wpLocation = new WpLocation(bukkitLocation.getX(), bukkitLocation.getY(), bukkitLocation.getZ());
         newPet.setLastKnownLocation(wpLocation);
         entity.setCustomName(ChatColor.GREEN + newPet.getName());
-        entity.setPersistent(true);
+        newPet.ensurePersistence();
         entity.playEffect(EntityEffect.LOVE_HEARTS);
 
         // add pet to pet list
         PetList petList = getPetList(player.getUniqueId());
         petList.addPet(newPet);
+        petsByEntityUUID.put(newPet.getUniqueID(), newPet);
         return true;
     }
 
     public boolean removePet(Pet petToRemove) {
+        petsByEntityUUID.remove(petToRemove.getUniqueID());
         return getPetList(petToRemove.getOwnerUUID()).removePet(petToRemove);
     }
 
     public Pet getPet(Entity entity) {
-        for (PetList petList : getPetLists()) {
-            Pet pet = petList.getPet(entity.getUniqueId());
-            if (pet != null) {
-                return pet;
-            }
-        }
-        return null;
+        return petsByEntityUUID.get(entity.getUniqueId());
     }
 
     public PetList getPetList(UUID playerUUID) {
@@ -74,6 +72,26 @@ public void createPetListForPlayer(UUID playerUUID) {
         getPetLists().add(newPetList);
     }
 
+    /**
+     * Adds an already-constructed Pet to the appropriate PetList and the UUID index.
+     * Used when loading pets from storage.
+     */
+    public void addExistingPet(Pet pet) {
+        if (getPetList(pet.getOwnerUUID()) == null) {
+            createPetListForPlayer(pet.getOwnerUUID());
+        }
+        getPetList(pet.getOwnerUUID()).addPet(pet);
+        petsByEntityUUID.put(pet.getUniqueID(), pet);
+    }
+
+    /**
+     * Clears all pet lists and the UUID index.
+     */
+    public void clearAll() {
+        petLists.clear();
+        petsByEntityUUID.clear();
+    }
+
     public Pet getPlayersPet(Player player, Entity entity) {
         PetList petList = getPetList(player.getUniqueId());
         return petList.getPet(entity.getUniqueId());

src/main/java/dansplugins/wildpets/storage/StorageService.java

@@ -91,7 +91,7 @@ private void writeOutFiles(List<Map<String, String>> saveData, String fileName)
 
     private void loadPets() {
         // load each pet individually and reconstruct pet list objects
-        petListRepository.getPetLists().clear();
+        petListRepository.clearAll();
 
         ArrayList<HashMap<String, String>> data = loadDataFromFilename(FILE_PATH + PETS_FILE_NAME);
 
@@ -103,11 +103,7 @@ private void loadPets() {
         }
 
         for (Pet pet : allPets) {
-            if (petListRepository.getPetList(pet.getOwnerUUID()) == null) {
-                petListRepository.createPetListForPlayer(pet.getOwnerUUID());
-
-            }
-            petListRepository.getPetList(pet.getOwnerUUID()).addPet(pet);
+            petListRepository.addExistingPet(pet);
             petRecordRepository.addPetRecord(pet); // will not result in duplicates because petRecords is a hashset
         }
 
@@ -134,6 +130,7 @@ private void applyAIStateToPets() {
         wildPets.getServer().getScheduler().runTaskLater(wildPets, () -> {
             for (Pet pet : petListRepository.getAllPets()) {
                 pet.applyAIState();
+                pet.ensurePersistence();
             }
         }, delayTicks); // Default: wait 100 ticks (5 seconds at 20 TPS) for entities to be loaded
     }

src/test/java/dansplugins/wildpets/pet/PetTest.java

@@ -303,6 +303,62 @@ public void testApplyAIStateWithNullEntity() {
         assertEquals("Staying", pet.getMovementState());
     }
 
+    @Test
+    public void testEnsurePersistenceWithMob() {
+        // prepare - Mob extends LivingEntity
+        when(mockServer.getEntity(entityUniqueId)).thenReturn(mockMob);
+
+        // Create pet
+        Pet pet = new Pet(entityUniqueId, playerOwnerUniqueId, playerOwnerName, mockServerProvider);
+
+        // execute
+        pet.ensurePersistence();
+
+        // verify
+        verify(mockMob).setPersistent(true);
+        verify(mockMob).setRemoveWhenFarAway(false);
+    }
+
+    @Test
+    public void testEnsurePersistenceWithNonLivingEntity() {
+        // prepare - mockEntity is just Entity, not LivingEntity
+        when(mockServer.getEntity(entityUniqueId)).thenReturn(mockEntity);
+
+        // Create pet
+        Pet pet = new Pet(entityUniqueId, playerOwnerUniqueId, playerOwnerName, mockServerProvider);
+
+        // execute
+        pet.ensurePersistence();
+
+        // verify setPersistent is called, but setRemoveWhenFarAway is not (not a LivingEntity)
+        verify(mockEntity).setPersistent(true);
+    }
+
+    @Test
+    public void testEnsurePersistenceWithNullEntity() {
+        // prepare - entity not found (e.g., in unloaded chunk)
+        when(mockServer.getEntity(entityUniqueId)).thenReturn(null);
+
+        // Create pet
+        Pet pet = new Pet(entityUniqueId, playerOwnerUniqueId, playerOwnerName, mockServerProvider);
+
+        // execute - should not throw an exception
+        pet.ensurePersistence();
+
+        // verify no interactions attempted on null entity
+    }
+
+    @Test
+    public void testEnsurePersistenceWithNullServerProvider() {
+        // Create pet with null server provider
+        Pet pet = new Pet(entityUniqueId, playerOwnerUniqueId, playerOwnerName, null);
+
+        // execute - should not throw an exception
+        pet.ensurePersistence();
+
+        // verify no exception occurred
+    }
+
     // Helper method to create test data
     private Map<String, String> createBasicPetData(String movementState) {
         Map<String, String> petData = new HashMap<>();