diff --git a/core/src/mindustry/entities/units/UnitController.java b/core/src/mindustry/entities/units/UnitController.java
index b77a857a2f..240728b2cc 100644
--- a/core/src/mindustry/entities/units/UnitController.java
+++ b/core/src/mindustry/entities/units/UnitController.java
@@ -1,10 +1,11 @@
 package mindustry.entities.units;
 
+import arc.util.*;
 import mindustry.gen.*;
 
 public interface UnitController{
     void unit(Unit unit);
-    Unit unit();
+    @Nullable Unit unit();
 
     default void hit(Bullet bullet){
 
diff --git a/core/src/mindustry/input/InputHandler.java b/core/src/mindustry/input/InputHandler.java
index b5bacbe51e..1000494da4 100644
--- a/core/src/mindustry/input/InputHandler.java
+++ b/core/src/mindustry/input/InputHandler.java
@@ -151,6 +151,8 @@ public abstract class InputHandler implements InputProcessor, GestureListener{
             logicCutscene = false;
             itemDepositCooldown = 0f;
             Arrays.fill(controlGroups, null);
+            lastUnit = null;
+            lastPlans.clear();
         });
     }
 
@@ -814,6 +816,7 @@ public abstract class InputHandler implements InputProcessor, GestureListener{
 
         if(player.isBuilder()){
             if(player.unit() != lastUnit && player.unit().plans.size <= 1){
+                player.unit().plans.ensureCapacity(lastPlans.size);
                 for(var plan : lastPlans){
                     player.unit().plans.addLast(plan);
                 }