2 files changed, 9 insertions, 1 deletions
diff --git a/src/main/java/gregtech/common/misc/spaceprojects/base/SP_Upgrade.java b/src/main/java/gregtech/common/misc/spaceprojects/base/SP_Upgrade.java
index 6a1fe82791..18540c429a 100644
--- a/src/main/java/gregtech/common/misc/spaceprojects/base/SP_Upgrade.java
+++ b/src/main/java/gregtech/common/misc/spaceprojects/base/SP_Upgrade.java
@@ -7,6 +7,7 @@ import net.minecraft.util.StatCollector;
 import net.minecraftforge.fluids.FluidStack;
 
 import gregtech.common.misc.spaceprojects.SpaceProjectManager;
+import gregtech.common.misc.spaceprojects.enums.SpaceBodyType;
 import gregtech.common.misc.spaceprojects.enums.UpgradeStatus;
 import gregtech.common.misc.spaceprojects.interfaces.ISpaceProject;
 import gregtech.common.misc.spaceprojects.interfaces.ISpaceProject.ISP_Upgrade;
@@ -293,7 +294,7 @@ public class SP_Upgrade implements ISP_Upgrade {
             return true;
         }
 
-        if (requirements.getBodyType() != null) {
+        if (requirements.getBodyType() != null && !requirements.getBodyType().equals(SpaceBodyType.NONE)) {
             if (!requirements.getBodyType().equals(projectBelongingTo.getProjectLocation().getType())) {
                 return false;
             }
@@ -335,6 +336,9 @@ public class SP_Upgrade implements ISP_Upgrade {
     @Override
     public void goToNextStage() {
         currentStage++;
+        if (isFinished()) {
+            projectBelongingTo.setBuiltUpgrade(this);
+        }
     }
 
     @Override
diff --git a/src/main/java/gregtech/common/misc/spaceprojects/base/SpaceProject.java b/src/main/java/gregtech/common/misc/spaceprojects/base/SpaceProject.java
index 57625f083e..0095ece5b1 100644
--- a/src/main/java/gregtech/common/misc/spaceprojects/base/SpaceProject.java
+++ b/src/main/java/gregtech/common/misc/spaceprojects/base/SpaceProject.java
@@ -310,6 +310,7 @@ public class SpaceProject implements ISpaceProject {
 
     public SpaceProject setProjectUpgrades(ISP_Upgrade... spaceProjectUpgrades) {
         for (ISP_Upgrade upgrade : spaceProjectUpgrades) {
+            upgrade.setUpgradeProject(this);
             upgradesAvailable.put(upgrade.getUpgradeName(), upgrade);
         }
         return this;
@@ -354,6 +355,9 @@ public class SpaceProject implements ISpaceProject {
         }
 
         for (ISP_Upgrade upgrade : upgrades) {
+            if (upgrade.equals(currentUpgrade)) {
+                currentUpgrade = null;
+            }
             upgradesInstalled.put(upgrade.getUpgradeName(), upgrade);
         }
     }