diff --git a/projects/common-api/src/main/java/dan200/computercraft/api/upgrades/UpgradeBase.java b/projects/common-api/src/main/java/dan200/computercraft/api/upgrades/UpgradeBase.java
index ee83da691..5e16c281d 100644
--- a/projects/common-api/src/main/java/dan200/computercraft/api/upgrades/UpgradeBase.java
+++ b/projects/common-api/src/main/java/dan200/computercraft/api/upgrades/UpgradeBase.java
@@ -8,9 +8,11 @@ import dan200.computercraft.api.pocket.IPocketUpgrade;
 import dan200.computercraft.api.turtle.ITurtleUpgrade;
 import dan200.computercraft.impl.PlatformHelper;
 import net.minecraft.Util;
+import net.minecraft.nbt.CompoundTag;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.item.ItemStack;
 
+import javax.annotation.Nonnull;
 import java.util.Objects;
 
 /**
@@ -50,6 +52,14 @@ public interface UpgradeBase {
      */
     ItemStack getCraftingItem();
 
+    default @Nonnull ItemStack produceCraftingItem(@Nonnull CompoundTag upgradeData) {
+        return getCraftingItem();
+    }
+
+    default @Nonnull CompoundTag produceUpgradeData(@Nonnull ItemStack stack) {
+        return new CompoundTag();
+    }
+
     /**
      * Determine if an item is suitable for being used for this upgrade.
      * <p>
diff --git a/projects/common/src/main/java/dan200/computercraft/shared/pocket/apis/PocketAPI.java b/projects/common/src/main/java/dan200/computercraft/shared/pocket/apis/PocketAPI.java
index e1b1cb0ad..091d97407 100644
--- a/projects/common/src/main/java/dan200/computercraft/shared/pocket/apis/PocketAPI.java
+++ b/projects/common/src/main/java/dan200/computercraft/shared/pocket/apis/PocketAPI.java
@@ -61,17 +61,21 @@ public class PocketAPI implements ILuaAPI {
 
         // Attempt to find the upgrade, starting in the main segment, and then looking in the opposite
         // one. We start from the position the item is currently in and loop round to the start.
-        var newUpgrade = findUpgrade(inventory.items, inventory.selected, previousUpgrade);
-        if (newUpgrade == null) {
-            newUpgrade = findUpgrade(inventory.offhand, 0, previousUpgrade);
+        var newUpgradeItem = findUpgrade(inventory.items, inventory.selected, previousUpgrade);
+        if (newUpgradeItem == null) {
+            newUpgradeItem = findUpgrade(inventory.offhand, 0, previousUpgrade);
         }
+        if (newUpgradeItem == null) return new Object[]{ false, "Cannot find a valid upgrade" };
+        // Unpack new upgrade here and ensure some mystical error is not happened
+        var newUpgrade = PocketUpgrades.instance().get(newUpgradeItem);
         if (newUpgrade == null) return new Object[]{ false, "Cannot find a valid upgrade" };
 
         // Remove the current upgrade
-        if (previousUpgrade != null) storeItem(player, previousUpgrade.getCraftingItem().copy());
+        if (previousUpgrade != null) storeItem(player, previousUpgrade.produceCraftingItem(computer.getUpgradeNBTData()).copy());
 
         // Set the new upgrade
         computer.setUpgrade(newUpgrade);
+        computer.getUpgradeNBTData().merge(newUpgrade.produceUpgradeData(newUpgradeItem));
 
         return new Object[]{ true };
     }
@@ -105,7 +109,7 @@ public class PocketAPI implements ILuaAPI {
         }
     }
 
-    private static @Nullable IPocketUpgrade findUpgrade(NonNullList<ItemStack> inv, int start, @Nullable IPocketUpgrade previous) {
+    private static @Nullable ItemStack findUpgrade(NonNullList<ItemStack> inv, int start, @Nullable IPocketUpgrade previous) {
         for (var i = 0; i < inv.size(); i++) {
             var invStack = inv.get((i + start) % inv.size());
             if (!invStack.isEmpty()) {
@@ -117,7 +121,7 @@ public class PocketAPI implements ILuaAPI {
                     invStack.shrink(1);
                     inv.set((i + start) % inv.size(), invStack.isEmpty() ? ItemStack.EMPTY : invStack);
 
-                    return newUpgrade;
+                    return invStack.copyWithCount(1);
                 }
             }
         }
diff --git a/projects/common/src/main/java/dan200/computercraft/shared/turtle/core/TurtleEquipCommand.java b/projects/common/src/main/java/dan200/computercraft/shared/turtle/core/TurtleEquipCommand.java
index 1a15bb840..44d5e098d 100644
--- a/projects/common/src/main/java/dan200/computercraft/shared/turtle/core/TurtleEquipCommand.java
+++ b/projects/common/src/main/java/dan200/computercraft/shared/turtle/core/TurtleEquipCommand.java
@@ -7,6 +7,7 @@ package dan200.computercraft.shared.turtle.core;
 import dan200.computercraft.api.turtle.*;
 import dan200.computercraft.impl.TurtleUpgrades;
 import dan200.computercraft.shared.turtle.TurtleUtil;
+import net.minecraft.nbt.CompoundTag;
 
 public class TurtleEquipCommand implements TurtleCommand {
     private final TurtleSide side;
@@ -30,10 +31,16 @@ public class TurtleEquipCommand implements TurtleCommand {
             newUpgrade = null;
         }
 
+        CompoundTag upgradeData = null;
+
         // Do the swapping:
-        if (newUpgrade != null) turtle.getInventory().removeItem(turtle.getSelectedSlot(), 1);
-        if (oldUpgrade != null) TurtleUtil.storeItemOrDrop(turtle, oldUpgrade.getCraftingItem().copy());
+        if (newUpgrade != null) {
+            var upgradeItem = turtle.getInventory().removeItem(turtle.getSelectedSlot(), 1);
+            upgradeData = newUpgrade.produceUpgradeData(upgradeItem);
+        }
+        if (oldUpgrade != null) TurtleUtil.storeItemOrDrop(turtle, oldUpgrade.produceCraftingItem(turtle.getUpgradeNBTData(side)).copy());
         turtle.setUpgrade(side, newUpgrade);
+        if (upgradeData != null) turtle.getUpgradeNBTData(side).merge(upgradeData);
 
         // Animate
         if (newUpgrade != null || oldUpgrade != null) {