fix: bug with WaitScreens

cfoust · Nov 9, 2023 · c100577 · c100577
1 parent dd05c97
commit c100577
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pkg/taro/screens.go b/pkg/taro/screens.go
@@ -10,7 +10,6 @@ type ScreenUpdate struct{}
 
 func WaitScreens(ctx context.Context, screens ...mux.Screen) Cmd {
 	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
 	done := make(chan struct{})
 	for _, screen := range screens {
@@ -32,6 +31,7 @@ func WaitScreens(ctx context.Context, screens ...mux.Screen) Cmd {
 	}
 
 	return func() Msg {
+		defer cancel()
 		select {
 		case <-done:
 			return ScreenUpdate{}