diff --git a/src/main/java/frc/robot/managers/superstructure/SuperstructureManager.java b/src/main/java/frc/robot/managers/superstructure/SuperstructureManager.java
index b3f5a25..b1af01c 100644
--- a/src/main/java/frc/robot/managers/superstructure/SuperstructureManager.java
+++ b/src/main/java/frc/robot/managers/superstructure/SuperstructureManager.java
@@ -13,6 +13,7 @@
 import frc.robot.util.scheduling.LifecycleSubsystem;
 import frc.robot.util.scheduling.SubsystemPriority;
 import frc.robot.wrist.WristSubsystem;
+import java.util.function.Supplier;
 
 public class SuperstructureManager extends LifecycleSubsystem {
   private SuperstructureMotionManager motionManager;
@@ -64,8 +65,12 @@ public Command setIntakeOverrideCommand(IntakeState intakeState) {
   }
 
   public Command setStateCommand(SuperstructureState newGoalState) {
-    return Commands.runOnce(() -> setGoal(newGoalState))
-        .andThen(Commands.waitUntil(() -> atGoal(newGoalState)));
+    return setStateCommand(() -> newGoalState);
+  }
+
+  public Command setStateCommand(Supplier<SuperstructureState> newGoalState) {
+    return Commands.runOnce(() -> setGoal(newGoalState.get()))
+        .andThen(Commands.waitUntil(() -> atGoal(newGoalState.get())));
   }
 
   public void setMode(HeldGamePiece mode) {
@@ -84,25 +89,23 @@ public HeldGamePiece getMode() {
   }
 
   public Command getIntakeFloorCommand() {
-    // TODO: Refactor to use setStateCommand, since this command never finishes atm
-    return Commands.run(
+    return setStateCommand(
         () -> {
           if (mode == HeldGamePiece.CONE) {
-            setGoal(States.INTAKING_CONE_FLOOR);
+            return States.INTAKING_CONE_FLOOR;
           } else {
-            setGoal(States.INTAKING_CUBE_FLOOR);
+            return States.INTAKING_CUBE_FLOOR;
           }
         });
   }
 
   public Command getIntakeShelfCommand() {
-    // TODO: Refactor to use setStateCommand, since this command never finishes atm
-    return Commands.run(
+    return setStateCommand(
         () -> {
           if (mode == HeldGamePiece.CONE) {
-            setGoal(States.INTAKING_CONE_SHELF);
+            return States.INTAKING_CONE_SHELF;
           } else {
-            setGoal(States.INTAKING_CUBE_SHELF);
+            return States.INTAKING_CUBE_SHELF;
           }
         });
   }