Attenpts to mitigate an NPC's entity becoming unavailable

src/main/java/net/knarcraft/blacksmith/trait/Session.java

@@ -26,6 +26,7 @@ public abstract class Session implements Runnable {
     protected static final Random random = new Random();
     protected final Player player;
     protected final NPC npc;
+    protected final Entity entity;
     protected long finishTime;
     protected int taskId;
 
@@ -38,6 +39,10 @@ public abstract class Session implements Runnable {
     public Session(@NotNull Player player, @NotNull NPC npc) {
         this.player = player;
         this.npc = npc;
+        this.entity = npc.getEntity();
+        if (entity == null) {
+            throw new IllegalArgumentException("Tried to start session for an NPC without an assigned entity.");
+        }
     }
 
     /**
@@ -103,9 +108,9 @@ public abstract class Session implements Runnable {
 
         BlacksmithPlugin instance = BlacksmithPlugin.getInstance();
         if (this instanceof ReforgeSession) {
-            instance.callEvent(new BlacksmithReforgeStartEvent(npc, player, actionDelayTicks, getCraftingStation()));
+            instance.callEvent(new BlacksmithReforgeStartEvent(this.npc, this.entity, this.player, actionDelayTicks, getCraftingStation()));
         } else if (this instanceof SalvageSession) {
-            instance.callEvent(new ScrapperSalvageStartEvent(npc, player, actionDelayTicks, getCraftingStation()));
+            instance.callEvent(new ScrapperSalvageStartEvent(this.npc, this.entity, this.player, actionDelayTicks, getCraftingStation()));
         }
 
         taskId = scheduler.scheduleSyncDelayedTask(BlacksmithPlugin.getInstance(), this, actionDelayTicks);
@@ -116,14 +121,13 @@ public abstract class Session implements Runnable {
      *
      * @param hasDelay     <p>Whether the session was delayed, or if the item was processed instantly</p>
      * @param dropItem     <p>Whether the item should be dropped on the ground</p>
-     * @param npc          <p>The NPC holding the item</p>
      * @param itemToReturn <p>The item to return to the player</p>
      */
-    protected void giveResultingItem(boolean hasDelay, boolean dropItem, @NotNull NPC npc, @NotNull ItemStack itemToReturn) {
+    protected void giveResultingItem(boolean hasDelay, boolean dropItem, @NotNull ItemStack itemToReturn) {
         if (hasDelay) {
             //If the player isn't online, or the player cannot fit the item, drop the item to prevent it from disappearing
             if (dropItem || !this.player.isOnline() || !ItemHelper.canFitItem(this.player.getInventory(), itemToReturn)) {
-                this.player.getWorld().dropItemNaturally(npc.getEntity().getLocation(), itemToReturn);
+                this.player.getWorld().dropItemNaturally(this.entity.getLocation(), itemToReturn);
             } else {
                 this.player.getInventory().addItem(itemToReturn);
             }
@@ -172,7 +176,7 @@ public abstract class Session implements Runnable {
      * @param sound <p>The sound to play</p>
      */
     protected void playSound(Sound sound) {
-        playSound(this.npc.getEntity(), sound);
+        playSound(this.entity, sound);
     }
 
     /**
@@ -187,14 +191,13 @@ public abstract class Session implements Runnable {
             return;
         }
 
-        NPCSoundEvent event = new NPCSoundEvent(this.npc, this.player, SoundCategory.AMBIENT, sound, 0.5f, 1.0f);
+        NPCSoundEvent event = new NPCSoundEvent(this.npc, this.entity, this.player, SoundCategory.AMBIENT, sound, 0.5f, 1.0f);
         BlacksmithPlugin.getInstance().callEvent(event);
         if (event.isCancelled()) {
             return;
         }
 
-        world.playSound(event.getNpc().getEntity(), event.getSound(), event.getSoundCategory(), event.getVolume(),
-                event.getPitch());
+        world.playSound(event.getEntity(), event.getSound(), event.getSoundCategory(), event.getVolume(), event.getPitch());
     }
 
 }