diff --git a/src/main/java/com/lambda/mixin/gui/MixinGuiContainer.java b/src/main/java/com/lambda/mixin/gui/MixinGuiContainer.java
index de1f01712..d3b7b938f 100644
--- a/src/main/java/com/lambda/mixin/gui/MixinGuiContainer.java
+++ b/src/main/java/com/lambda/mixin/gui/MixinGuiContainer.java
@@ -12,43 +12,45 @@
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
-import java.io.IOException;
-
 @Mixin(GuiContainer.class)
 public class MixinGuiContainer extends GuiScreen {
 
     @Shadow protected int guiLeft;
     @Shadow protected int guiTop;
     @Shadow protected int xSize;
+    private final GuiButton stealButton = new LambdaGuiStealButton(guiLeft + xSize + 2, guiTop + 2);
+    private final GuiButton storeButton = new LambdaGuiStoreButton(guiLeft + xSize + 2, guiTop + 4 + stealButton.height);
 
-    private final GuiButton stealButton = new LambdaGuiStealButton(this.guiLeft + this.xSize + 2, this.guiTop + 2);
-    private final GuiButton storeButton = new LambdaGuiStoreButton(this.guiLeft + this.xSize + 2, this.guiTop + 4 + stealButton.height);
-
-    @Inject(method = "initGui", at = @At("HEAD"))
-    public void initGui(CallbackInfo ci) {
-        if (ChestStealer.INSTANCE.isValidGui()) {
-            this.buttonList.add(stealButton);
-            this.buttonList.add(storeButton);
-            ChestStealer.updateButton(stealButton, this.guiLeft, this.xSize, this.guiTop);
-            ChestStealer.updateButton(storeButton, this.guiLeft, this.xSize, this.guiTop);
-        }
-    }
+    @Inject(method = "mouseClicked", at = @At("TAIL"))
+    public void mouseClicked(int x, int y, int button, CallbackInfo ci) {
+        if (button != 0) return;
 
-    @Override
-    protected void actionPerformed(GuiButton button) throws IOException {
-        if (button.id == 696969) {
-            ChestStealer.INSTANCE.setStealing(!ChestStealer.INSTANCE.getStealing());
-        } else if (button.id == 420420) {
+        if (storeButton.mousePressed(mc, x, y)) {
             ChestStealer.INSTANCE.setStoring(!ChestStealer.INSTANCE.getStoring());
-        } else {
-            super.actionPerformed(button);
+        } else if (stealButton.mousePressed(mc, x, y)) {
+            ChestStealer.INSTANCE.setStealing(!ChestStealer.INSTANCE.getStealing());
         }
     }
 
-    @Inject(method = "updateScreen", at = @At("HEAD"))
+    @Inject(method = "updateScreen", at = @At("TAIL"))
     public void updateScreen(CallbackInfo ci) {
-        ChestStealer.updateButton(stealButton, this.guiLeft, this.xSize, this.guiTop);
-        ChestStealer.updateButton(storeButton, this.guiLeft, this.xSize, this.guiTop);
+        if (!ChestStealer.INSTANCE.isValidGui()) return;
+
+        if (ChestStealer.INSTANCE.isDisabled()) {
+            buttonList.remove(storeButton);
+            buttonList.remove(storeButton);
+            return;
+        }
+
+        if (!buttonList.contains(stealButton)) {
+            buttonList.add(stealButton);
+        }
+        if (!buttonList.contains(storeButton)) {
+            buttonList.add(storeButton);
+        }
+
+        ChestStealer.updateButton(stealButton, guiLeft, xSize, guiTop);
+        ChestStealer.updateButton(storeButton, guiLeft, xSize, guiTop);
     }
 
 }
diff --git a/src/main/kotlin/com/lambda/client/module/modules/player/ChestStealer.kt b/src/main/kotlin/com/lambda/client/module/modules/player/ChestStealer.kt
index f7834216e..404e62dba 100644
--- a/src/main/kotlin/com/lambda/client/module/modules/player/ChestStealer.kt
+++ b/src/main/kotlin/com/lambda/client/module/modules/player/ChestStealer.kt
@@ -89,29 +89,21 @@ object ChestStealer : Module(
         runSafe {
             if (isEnabled && isContainerOpen()) {
                 if (button.id == 696969) {
-                    val str = if (stealing) {
-                        "Stop"
-                    } else {
-                        "Steal"
-                    }
+                    val name = if (stealing) "Stop" else "Steal"
 
                     button.x = left + size + 2
                     button.y = top + 2
                     button.enabled = canSteal() and !storing
                     button.visible = true
-                    button.displayString = str
+                    button.displayString = name
                 } else if (button.id == 420420) {
-                    val str = if (storing) {
-                        "Stop"
-                    } else {
-                        "Store"
-                    }
+                    val name = if (storing) "Stop" else "Store"
 
                     button.x = left + size + 2
                     button.y = top + 24
                     button.enabled = canStore() and !stealing
                     button.visible = true
-                    button.displayString = str
+                    button.displayString = name
                 }
             } else {
                 button.visible = false