2 files changed, 10 insertions, 8 deletions
diff --git a/src/main/java/dev/isxander/yacl/api/ButtonOption.java b/src/main/java/dev/isxander/yacl/api/ButtonOption.java
index 346bfa4..f3d1e6b 100644
--- a/src/main/java/dev/isxander/yacl/api/ButtonOption.java
+++ b/src/main/java/dev/isxander/yacl/api/ButtonOption.java
@@ -1,5 +1,6 @@
 package dev.isxander.yacl.api;
 
+import dev.isxander.yacl.gui.YACLScreen;
 import dev.isxander.yacl.impl.ButtonOptionImpl;
 import net.minecraft.text.MutableText;
 import net.minecraft.text.Text;
@@ -8,10 +9,11 @@ import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
-public interface ButtonOption extends Option<Runnable> {
-    Runnable action();
+public interface ButtonOption extends Option<Consumer<YACLScreen>> {
+    Consumer<YACLScreen> action();
 
     static Builder createBuilder() {
         return new Builder();
@@ -20,8 +22,8 @@ public interface ButtonOption extends Option<Runnable> {
     class Builder {
         private Text name;
         private final List<Text> tooltipLines = new ArrayList<>();
-        private Function<ButtonOption, Controller<Runnable>> controlGetter;
-        private Runnable action;
+        private Function<ButtonOption, Controller<Consumer<YACLScreen>>> controlGetter;
+        private Consumer<YACLScreen> action;
 
         private Builder() {
 
@@ -41,14 +43,14 @@ public interface ButtonOption extends Option<Runnable> {
             return this;
         }
 
-        public Builder action(@NotNull Runnable action) {
+        public Builder action(@NotNull Consumer<YACLScreen> action) {
             Validate.notNull(action, "`action` cannot be null");
 
             this.action = action;
             return this;
         }
 
-        public Builder controller(@NotNull Function<ButtonOption, Controller<Runnable>> control) {
+        public Builder controller(@NotNull Function<ButtonOption, Controller<Consumer<YACLScreen>>> control) {
             Validate.notNull(control, "`control` cannot be null");
 
             this.controlGetter = control;
diff --git a/src/main/java/dev/isxander/yacl/api/Controller.java b/src/main/java/dev/isxander/yacl/api/Controller.java
index 2c0a5e1..198e5df 100644
--- a/src/main/java/dev/isxander/yacl/api/Controller.java
+++ b/src/main/java/dev/isxander/yacl/api/Controller.java
@@ -1,8 +1,8 @@
 package dev.isxander.yacl.api;
 
 import dev.isxander.yacl.api.utils.Dimension;
+import dev.isxander.yacl.gui.YACLScreen;
 import dev.isxander.yacl.gui.controllers.ControllerWidget;
-import net.minecraft.client.gui.screen.Screen;
 import net.minecraft.text.Text;
 import org.jetbrains.annotations.ApiStatus;
 
@@ -26,5 +26,5 @@ public interface Controller<T> {
      * @param screen parent screen
      */
     @ApiStatus.Internal
-    ControllerWidget<?> provideWidget(Screen screen, Dimension<Integer> widgetDimension);
+    ControllerWidget<?> provideWidget(YACLScreen screen, Dimension<Integer> widgetDimension);
 }