From c100577e5da90171b09bcb4c8aa51a3c73c39156 Mon Sep 17 00:00:00 2001
From: Caleb Foust <cfoust@sqweebloid.com>
Date: Thu, 9 Nov 2023 06:07:15 -0500
Subject: [PATCH] fix: bug with WaitScreens

---
 pkg/taro/screens.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/taro/screens.go b/pkg/taro/screens.go
index 2db24557..078ab1bf 100644
--- 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{}