1 files changed, 3 insertions, 1 deletions
diff --git a/src/Waylean.lean b/src/Waylean.lean
index 0cbd6f8..3f020a1 100644
--- a/src/Waylean.lean
+++ b/src/Waylean.lean
@@ -63,6 +63,7 @@ def WLDisplay.sync {σ : Type} (display : WLDisplay) : DisplayM σ WLCallback :=
 
 structure State where
   count : Int
+  callback : Option WLCallback
 
 partial def nextTimer : DisplayM Int Unit := do
   _ ← DisplayM.registerTimer 16ms do
@@ -81,6 +82,7 @@ def launch : DisplayM State Unit := do
   let callback : WLCallback ← disp.sync
   callback.listen fun ev ↦ match ev with
     | .done _ => DisplayM.liftIO (IO.println "Roundtrip complete")
+  set { (← get) with callback := callback }
 
 def main : IO Unit :=
-  DisplayM.run launch { count := 0 }
+  DisplayM.run launch { count := 0, callback := .none }