From f5bfbfd4f8a18a42add7f57e65ef059cd1284409 Mon Sep 17 00:00:00 2001
From: EpicKnarvik97 <kristian.knarvik@knett.no>
Date: Fri, 27 Sep 2024 01:37:07 +0200
Subject: [PATCH] Attenpts to mitigate an NPC's entity becoming unavailable

---
 .../event/AbstractBlacksmithPluginEvent.java  | 12 ++++++++-
 .../event/BlacksmithPluginEvent.java          |  9 +++++++
 .../event/BlacksmithReforgeFailEvent.java     |  9 ++++---
 .../event/BlacksmithReforgeStartEvent.java    |  6 +++--
 .../event/BlacksmithReforgeSucceedEvent.java  |  6 +++--
 .../blacksmith/event/NPCSoundEvent.java       |  6 +++--
 .../event/ScrapperSalvageFailEvent.java       |  6 +++--
 .../event/ScrapperSalvageStartEvent.java      |  6 +++--
 .../event/ScrapperSalvageSucceedEvent.java    |  6 +++--
 .../blacksmith/trait/CustomTrait.java         | 25 +++++++++++++-----
 .../blacksmith/trait/ReforgeSession.java      | 22 +++++++++-------
 .../blacksmith/trait/SalvageSession.java      | 26 +++++++++++--------
 .../knarcraft/blacksmith/trait/Session.java   | 21 ++++++++-------
 13 files changed, 109 insertions(+), 51 deletions(-)

diff --git a/src/main/java/net/knarcraft/blacksmith/event/AbstractBlacksmithPluginEvent.java b/src/main/java/net/knarcraft/blacksmith/event/AbstractBlacksmithPluginEvent.java
index 3087ae2..5372d84 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/AbstractBlacksmithPluginEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/AbstractBlacksmithPluginEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.Event;
 import org.jetbrains.annotations.NotNull;
@@ -13,15 +14,18 @@ public abstract class AbstractBlacksmithPluginEvent extends Event implements Bla
 
     protected final NPC npc;
     protected final Player player;
+    protected final Entity entity;
 
     /**
      * Instantiates a new blacksmith plugin event
      *
      * @param npc    <p>The NPC involved in the event</p>
+     * @param entity <p>The entity of the NPC</p>
      * @param player <p>The player involved in the event</p>
      */
-    public AbstractBlacksmithPluginEvent(@NotNull NPC npc, @NotNull Player player) {
+    public AbstractBlacksmithPluginEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player) {
         this.npc = npc;
+        this.entity = entity;
         this.player = player;
     }
 
@@ -31,6 +35,12 @@ public abstract class AbstractBlacksmithPluginEvent extends Event implements Bla
         return this.npc;
     }
 
+    @Override
+    @NotNull
+    public Entity getEntity() {
+        return this.entity;
+    }
+
     @Override
     @NotNull
     public Player getPlayer() {
diff --git a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithPluginEvent.java b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithPluginEvent.java
index a5ae014..abee100 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithPluginEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithPluginEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.jetbrains.annotations.NotNull;
 
@@ -18,6 +19,14 @@ public interface BlacksmithPluginEvent {
     @NotNull
     NPC getNpc();
 
+    /**
+     * Gets the entity of the NPC
+     *
+     * @return <p>The NPC entity</p>
+     */
+    @NotNull
+    Entity getEntity();
+
     /**
      * Gets the player involved in the event
      *
diff --git a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeFailEvent.java b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeFailEvent.java
index d69877b..46791f2 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeFailEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeFailEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -15,10 +16,12 @@ public class BlacksmithReforgeFailEvent extends AbstractBlacksmithPluginEvent im
     /**
      * Instantiates a new blacksmith reforge fail event
      *
-     * @param npc <p>The NPC involved in the event</p>
+     * @param npc    <p>The NPC involved in the event</p>
+     * @param entity <p>The entity of the NPC</p>
+     * @param player <p>The player that initiated the session</p>
      */
-    public BlacksmithReforgeFailEvent(@NotNull NPC npc, @NotNull Player player) {
-        super(npc, player);
+    public BlacksmithReforgeFailEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player) {
+        super(npc, entity, player);
     }
 
     /**
diff --git a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeStartEvent.java b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeStartEvent.java
index 7c9061e..db905cc 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeStartEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeStartEvent.java
@@ -2,6 +2,7 @@ package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
 import org.bukkit.Material;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -19,13 +20,14 @@ public class BlacksmithReforgeStartEvent extends AbstractBlacksmithPluginEvent i
      * Instantiates a new blacksmith reforge start event
      *
      * @param npc             <p>The NPC involved in the event</p>
+     * @param entity          <p>The entity of the NPC</p>
      * @param player          <p>The player involved in the event</p>
      * @param durationTicks   <p>The duration of the reforge</p>
      * @param craftingStation <p>The appropriate crafting station for this reforging</p>
      */
-    public BlacksmithReforgeStartEvent(@NotNull NPC npc, @NotNull Player player, long durationTicks,
+    public BlacksmithReforgeStartEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player, long durationTicks,
                                        @NotNull Material craftingStation) {
-        super(npc, player);
+        super(npc, entity, player);
         this.durationTicks = durationTicks;
         this.craftingStation = craftingStation;
     }
diff --git a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeSucceedEvent.java b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeSucceedEvent.java
index 269a40c..fff4861 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeSucceedEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/BlacksmithReforgeSucceedEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -16,10 +17,11 @@ public class BlacksmithReforgeSucceedEvent extends AbstractBlacksmithPluginEvent
      * Instantiates a new blacksmith reforge succeed event
      *
      * @param npc    <p>The NPC involved in the event</p>
+     * @param entity <p>The entity of the NPC</p>
      * @param player <p>The player involved in the event</p>
      */
-    public BlacksmithReforgeSucceedEvent(@NotNull NPC npc, @NotNull Player player) {
-        super(npc, player);
+    public BlacksmithReforgeSucceedEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player) {
+        super(npc, entity, player);
     }
 
     /**
diff --git a/src/main/java/net/knarcraft/blacksmith/event/NPCSoundEvent.java b/src/main/java/net/knarcraft/blacksmith/event/NPCSoundEvent.java
index 9e41d22..e341514 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/NPCSoundEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/NPCSoundEvent.java
@@ -3,6 +3,7 @@ package net.knarcraft.blacksmith.event;
 import net.citizensnpcs.api.npc.NPC;
 import org.bukkit.Sound;
 import org.bukkit.SoundCategory;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.Cancellable;
 import org.bukkit.event.HandlerList;
@@ -25,15 +26,16 @@ public class NPCSoundEvent extends AbstractBlacksmithPluginEvent implements Canc
      * Instantiates a new NPC sound event
      *
      * @param npc           <p>The NPC playing the sound</p>
+     * @param entity        <p>The entity playing the sound</p>
      * @param player        <p>The player whose interaction triggered the sound</p>
      * @param soundCategory <p>The category the sound is to play in</p>
      * @param sound         <p>The sound to play</p>
      * @param volume        <p>The volume of the played sound</p>
      * @param pitch         <p>The pitch of the played sound</p>
      */
-    public NPCSoundEvent(@NotNull NPC npc, @NotNull Player player, @NotNull SoundCategory soundCategory,
+    public NPCSoundEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player, @NotNull SoundCategory soundCategory,
                          @NotNull Sound sound, float volume, float pitch) {
-        super(npc, player);
+        super(npc, entity, player);
         this.soundCategory = soundCategory;
         this.sound = sound;
         this.volume = volume;
diff --git a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageFailEvent.java b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageFailEvent.java
index 1441efa..129446d 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageFailEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageFailEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -16,10 +17,11 @@ public class ScrapperSalvageFailEvent extends AbstractBlacksmithPluginEvent impl
      * Instantiates a new scrapper salvage fail event
      *
      * @param npc    <p>The NPC involved in the event</p>
+     * @param entity <p>The entity of the NPC</p>
      * @param player <p>The player involved in the event</p>
      */
-    public ScrapperSalvageFailEvent(@NotNull NPC npc, @NotNull Player player) {
-        super(npc, player);
+    public ScrapperSalvageFailEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player) {
+        super(npc, entity, player);
     }
 
     /**
diff --git a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageStartEvent.java b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageStartEvent.java
index 0d3c7b9..94d7605 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageStartEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageStartEvent.java
@@ -2,6 +2,7 @@ package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
 import org.bukkit.Material;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -20,13 +21,14 @@ public class ScrapperSalvageStartEvent extends AbstractBlacksmithPluginEvent imp
      * Instantiates a new scrapper salvage start event
      *
      * @param npc             <p>The NPC involved in the event</p>
+     * @param entity          <p>The entity of the NPC</p>
      * @param player          <p>The player involved in the event</p>
      * @param durationTicks   <p>The duration of the salvage</p>
      * @param craftingStation <p>The appropriate crafting station for this salvaging</p>
      */
-    public ScrapperSalvageStartEvent(@NotNull NPC npc, @NotNull Player player, long durationTicks,
+    public ScrapperSalvageStartEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player, long durationTicks,
                                      @NotNull Material craftingStation) {
-        super(npc, player);
+        super(npc, entity, player);
         this.durationTicks = durationTicks;
         this.craftingStation = craftingStation;
     }
diff --git a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageSucceedEvent.java b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageSucceedEvent.java
index 4eb67a4..6f9c438 100644
--- a/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageSucceedEvent.java
+++ b/src/main/java/net/knarcraft/blacksmith/event/ScrapperSalvageSucceedEvent.java
@@ -1,6 +1,7 @@
 package net.knarcraft.blacksmith.event;
 
 import net.citizensnpcs.api.npc.NPC;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
 import org.bukkit.event.HandlerList;
 import org.jetbrains.annotations.NotNull;
@@ -16,10 +17,11 @@ public class ScrapperSalvageSucceedEvent extends AbstractBlacksmithPluginEvent i
      * Instantiates a new scrapper salvage succeed event
      *
      * @param npc    <p>The NPC involved in the event</p>
+     * @param entity <p>The entity of the NPC</p>
      * @param player <p>The player involved in the event</p>
      */
-    public ScrapperSalvageSucceedEvent(@NotNull NPC npc, @NotNull Player player) {
-        super(npc, player);
+    public ScrapperSalvageSucceedEvent(@NotNull NPC npc, @NotNull Entity entity, @NotNull Player player) {
+        super(npc, entity, player);
     }
 
     /**
diff --git a/src/main/java/net/knarcraft/blacksmith/trait/CustomTrait.java b/src/main/java/net/knarcraft/blacksmith/trait/CustomTrait.java
index 1860486..1312ecd 100644
--- a/src/main/java/net/knarcraft/blacksmith/trait/CustomTrait.java
+++ b/src/main/java/net/knarcraft/blacksmith/trait/CustomTrait.java
@@ -2,14 +2,17 @@ package net.knarcraft.blacksmith.trait;
 
 import net.citizensnpcs.api.npc.NPC;
 import net.citizensnpcs.api.trait.Trait;
+import net.knarcraft.blacksmith.BlacksmithPlugin;
 import net.knarcraft.blacksmith.config.Setting;
 import net.knarcraft.blacksmith.config.Settings;
 import net.knarcraft.blacksmith.config.TraitSettings;
 import net.knarcraft.blacksmith.formatting.TimeFormatter;
 import net.knarcraft.blacksmith.manager.EconomyManager;
 import net.knarcraft.knarlib.formatting.StringFormatter;
+import org.bukkit.entity.Entity;
 import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
+import org.bukkit.inventory.EntityEquipment;
 import org.bukkit.inventory.ItemStack;
 import org.bukkit.inventory.PlayerInventory;
 import org.jetbrains.annotations.NotNull;
@@ -18,7 +21,6 @@ import org.jetbrains.annotations.Nullable;
 import java.util.Calendar;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Objects;
 import java.util.UUID;
 
 import static net.knarcraft.blacksmith.formatting.BlacksmithStringFormatter.sendNPCMessage;
@@ -153,8 +155,13 @@ public abstract class CustomTrait<K extends Setting> extends Trait {
         }
 
         //If already in a session, but the player has failed to interact, and left the blacksmith, allow a new session
+        Entity entity = npc.getEntity();
+        if (entity == null) {
+            BlacksmithPlugin.error("NPC session could not be started, as the NPC does not have a valid entity");
+            return false;
+        }
         if (session != null && !session.isRunning() && (System.currentTimeMillis() > currentSessionStartTime + 10 * 1000 ||
-                this.npc.getEntity().getLocation().distance(session.getPlayer().getLocation()) > 20)) {
+                entity.getLocation().distance(session.getPlayer().getLocation()) > 20)) {
             session = null;
         }
         return true;
@@ -189,10 +196,16 @@ public abstract class CustomTrait<K extends Setting> extends Trait {
         ItemStack heldItem = player.getInventory().getItemInMainHand();
 
         //Display the item in the NPC's hand
-        if (npc.getEntity() instanceof Player) {
-            ((Player) npc.getEntity()).getInventory().setItemInMainHand(heldItem);
-        } else {
-            Objects.requireNonNull(((LivingEntity) npc.getEntity()).getEquipment()).setItemInMainHand(heldItem);
+        Entity entity = npc.getEntity();
+        if (entity instanceof Player playerNPC) {
+            playerNPC.getInventory().setItemInMainHand(heldItem);
+        } else if (entity instanceof LivingEntity) {
+            EntityEquipment equipment = ((LivingEntity) entity).getEquipment();
+            if (equipment != null) {
+                ((LivingEntity) entity).getEquipment().setItemInMainHand(heldItem);
+            } else {
+                BlacksmithPlugin.warn("Failed to update NPC item, as its equipment was irretrievable.");
+            }
         }
         //Remove the item from the player's inventory
         if (this.session instanceof SalvageSession salvageSession) {
diff --git a/src/main/java/net/knarcraft/blacksmith/trait/ReforgeSession.java b/src/main/java/net/knarcraft/blacksmith/trait/ReforgeSession.java
index 8fea898..d9b2a5b 100644
--- a/src/main/java/net/knarcraft/blacksmith/trait/ReforgeSession.java
+++ b/src/main/java/net/knarcraft/blacksmith/trait/ReforgeSession.java
@@ -15,13 +15,13 @@ import org.bukkit.Sound;
 import org.bukkit.enchantments.Enchantment;
 import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
+import org.bukkit.inventory.EntityEquipment;
 import org.bukkit.inventory.ItemStack;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.List;
-import java.util.Objects;
 import java.util.Random;
 
 import static net.knarcraft.blacksmith.formatting.BlacksmithStringFormatter.sendNPCMessage;
@@ -107,10 +107,15 @@ public class ReforgeSession extends Session implements Runnable {
         playSound(success ? Sound.BLOCK_ANVIL_USE : Sound.ENTITY_VILLAGER_NO);
 
         //Stop the reforged item from displaying in the blacksmith's hand
-        if (this.npc.getEntity() instanceof Player) {
-            ((Player) this.npc.getEntity()).getInventory().setItemInMainHand(null);
-        } else {
-            Objects.requireNonNull(((LivingEntity) this.npc.getEntity()).getEquipment()).setItemInMainHand(null);
+        if (this.entity instanceof Player playerNPC) {
+            playerNPC.getInventory().setItemInMainHand(null);
+        } else if (this.entity instanceof LivingEntity) {
+            EntityEquipment equipment = ((LivingEntity) this.entity).getEquipment();
+            if (equipment != null) {
+                equipment.setItemInMainHand(null);
+            } else {
+                BlacksmithPlugin.warn("Failed to clear Blacksmith item, as its equipment was irretrievable.");
+            }
         }
 
         //Give the item back to the player
@@ -129,8 +134,7 @@ public class ReforgeSession extends Session implements Runnable {
      * Gives the reforged item back to the player
      */
     private void giveReforgedItem() {
-        giveResultingItem(this.config.getMaxReforgeDelay() > 0, this.config.getDropItem(), this.npc,
-                this.itemToReforge);
+        giveResultingItem(this.config.getMaxReforgeDelay() > 0, this.config.getDropItem(), this.itemToReforge);
     }
 
     /**
@@ -141,11 +145,11 @@ public class ReforgeSession extends Session implements Runnable {
     private boolean reforgeItem() {
         if (random.nextInt(100) < this.config.getFailChance()) {
             failReforge();
-            BlacksmithPlugin.getInstance().callEvent(new BlacksmithReforgeFailEvent(this.npc, this.player));
+            BlacksmithPlugin.getInstance().callEvent(new BlacksmithReforgeFailEvent(this.npc, this.entity, this.player));
             return false;
         } else {
             succeedReforge();
-            BlacksmithPlugin.getInstance().callEvent(new BlacksmithReforgeSucceedEvent(this.npc, this.player));
+            BlacksmithPlugin.getInstance().callEvent(new BlacksmithReforgeSucceedEvent(this.npc, this.entity, this.player));
             return true;
         }
     }
diff --git a/src/main/java/net/knarcraft/blacksmith/trait/SalvageSession.java b/src/main/java/net/knarcraft/blacksmith/trait/SalvageSession.java
index 4785d24..b262aea 100644
--- a/src/main/java/net/knarcraft/blacksmith/trait/SalvageSession.java
+++ b/src/main/java/net/knarcraft/blacksmith/trait/SalvageSession.java
@@ -13,13 +13,13 @@ import org.bukkit.Material;
 import org.bukkit.Sound;
 import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
+import org.bukkit.inventory.EntityEquipment;
 import org.bukkit.inventory.ItemStack;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.List;
-import java.util.Objects;
 import java.util.Random;
 
 import static net.knarcraft.blacksmith.formatting.BlacksmithStringFormatter.sendNPCMessage;
@@ -105,10 +105,15 @@ public class SalvageSession extends Session implements Runnable {
         salvageItem();
 
         //Stop the reforged item from displaying in the scrapper's hand
-        if (this.npc.getEntity() instanceof Player) {
-            ((Player) this.npc.getEntity()).getInventory().setItemInMainHand(null);
-        } else {
-            Objects.requireNonNull(((LivingEntity) this.npc.getEntity()).getEquipment()).setItemInMainHand(null);
+        if (this.entity instanceof Player playerNPC) {
+            playerNPC.getInventory().setItemInMainHand(null);
+        } else if (this.entity instanceof LivingEntity) {
+            EntityEquipment equipment = ((LivingEntity) this.entity).getEquipment();
+            if (equipment != null) {
+                equipment.setItemInMainHand(null);
+            } else {
+                BlacksmithPlugin.warn("Failed to clear Scrapper item, as its equipment was irretrievable.");
+            }
         }
 
         //Mark this scrapper as available
@@ -136,11 +141,11 @@ public class SalvageSession extends Session implements Runnable {
         if (random.nextInt(100) < this.config.getFailChance()) {
             playSound(Sound.ENTITY_VILLAGER_NO);
             failSalvage();
-            BlacksmithPlugin.getInstance().callEvent(new ScrapperSalvageFailEvent(this.npc, this.player));
+            BlacksmithPlugin.getInstance().callEvent(new ScrapperSalvageFailEvent(this.npc, this.entity, this.player));
         } else {
             playSound(Sound.BLOCK_ANVIL_USE);
             giveSalvage();
-            BlacksmithPlugin.getInstance().callEvent(new ScrapperSalvageSucceedEvent(this.npc, this.player));
+            BlacksmithPlugin.getInstance().callEvent(new ScrapperSalvageSucceedEvent(this.npc, this.entity, this.player));
             sendNPCMessage(this.npc, this.player, this.config.getSuccessMessage());
         }
     }
@@ -155,14 +160,13 @@ public class SalvageSession extends Session implements Runnable {
         if (ItemHelper.getMaxDurability(this.itemToSalvage) > 0) {
             //Damage the item if possible
             damageItemRandomly(this.itemToSalvage);
-            giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), this.npc,
-                    this.itemToSalvage);
+            giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), this.itemToSalvage);
             sendNPCMessage(this.npc, this.player, this.config.getFailSalvageMessage());
         } else {
             // Give half the salvage
             List<ItemStack> halfSalvage = SalvageHelper.pickRandomSalvage(this.salvage, new ArrayList<>(), 0.5);
             for (ItemStack item : halfSalvage) {
-                giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), this.npc, item);
+                giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), item);
             }
             sendNPCMessage(this.npc, this.player, this.config.getFailExtendedSalvageMessage());
         }
@@ -177,7 +181,7 @@ public class SalvageSession extends Session implements Runnable {
         this.player.giveExpLevels(this.enchantmentLevels);
         BlacksmithPlugin.debug("Giving salvage " + this.salvage);
         for (ItemStack item : this.salvage) {
-            giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), this.npc, item);
+            giveResultingItem(this.config.getMaxSalvageDelay() > 0, this.config.getDropItem(), item);
         }
     }
 
diff --git a/src/main/java/net/knarcraft/blacksmith/trait/Session.java b/src/main/java/net/knarcraft/blacksmith/trait/Session.java
index c909ae7..5b686de 100644
--- a/src/main/java/net/knarcraft/blacksmith/trait/Session.java
+++ b/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());
     }
 
 }