diff --git a/app/src/main/scala/tripaint/MainStage.scala b/app/src/main/scala/tripaint/MainStage.scala
index b168c34..a191e75 100644
--- a/app/src/main/scala/tripaint/MainStage.scala
+++ b/app/src/main/scala/tripaint/MainStage.scala
@@ -32,9 +32,8 @@ import java.io.File
 import scala.language.implicitConversions
 import scala.util.Try
 
-class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
-    extends Stage
-    with TriPaintView {
+class MainStage(controls: TriPaintViewListener, model: TriPaintModel, stage: Stage)
+    extends TriPaintView {
   private val (currentEditMode, setCurrentEditMode) = createResource(EditMode.Draw)
 
   private val imageDisplay: ImageGridPane = new ImageGridPane(model.imageGrid, currentEditMode)
@@ -49,8 +48,8 @@ class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
   private var currentFolder: Option[File] = None
 
   {
-    this.setTitle("TriPaint")
-    this.setOnCloseRequest(e => {
+    stage.setTitle("TriPaint")
+    stage.setOnCloseRequest(e => {
       if (!controls.requestExit()) e.consume()
     })
 
@@ -84,7 +83,7 @@ class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
       }
     }
 
-    this.setScene(scene)
+    stage.setScene(scene)
   }
 
   private def makeColorBox() = {
@@ -119,7 +118,7 @@ class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
     currentFolder.foreach(chooser.setInitialDirectory)
     chooser.setTitle("Save file")
     chooser.getExtensionFilters.add(new ExtensionFilter("PNG", "*.png"))
-    val result = Option(chooser.showSaveDialog(this))
+    val result = Option(chooser.showSaveDialog(stage))
     result.foreach(r => currentFolder = Some(r.getParentFile))
     result
   }
@@ -139,7 +138,7 @@ class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
   override def askForFileToOpen(): Option[File] = {
     val chooser = new FileChooser
     chooser.setTitle("Open file")
-    val result = Option(chooser.showOpenDialog(this))
+    val result = Option(chooser.showOpenDialog(stage))
     result.foreach(r => currentFolder = Some(r.getParentFile))
     result
   }
@@ -300,4 +299,6 @@ class MainStage(controls: TriPaintViewListener, model: TriPaintModel)
     val sizeStr = dialog.showAndWait().toScala
     sizeStr.map(str => Try(str.toInt).toOption.filter(_ > 0).getOrElse(32))
   }
+
+  override def close(): Unit = stage.close()
 }
diff --git a/app/src/main/scala/tripaint/TriPaint.scala b/app/src/main/scala/tripaint/TriPaint.scala
index b74240e..c4b7bbe 100644
--- a/app/src/main/scala/tripaint/TriPaint.scala
+++ b/app/src/main/scala/tripaint/TriPaint.scala
@@ -12,13 +12,7 @@ object TriPaint {
   class App extends Application {
     override def start(stage: Stage): Unit = {
       val model: TriPaintModel = TriPaintModel.create()
-      val controller = new TriPaintController(model, new MainStage(_, _))
-      val s = controller.view.asInstanceOf[Stage]
-
-      // TODO: send stage to the UI instead of copying over the data
-      stage.setTitle(s.getTitle)
-      stage.setScene(s.getScene)
-      stage.setOnCloseRequest(s.getOnCloseRequest)
+      val controller = new TriPaintController(model, new MainStage(_, _, stage))
 
       Platform.runLater(() =>
         model.imageGrid.setImageSizeIfEmpty(controller.view.askForImageSize().getOrElse(32))