From 2a3e0017cf3c03d9d021c800b14857453c7d8358 Mon Sep 17 00:00:00 2001
From: Ayman Bagabas <ayman.bagabas@gmail.com>
Date: Wed, 11 Sep 2024 11:57:21 -0400
Subject: [PATCH] fix: restore the terminal on kill

Supersedes: https://github.com/charmbracelet/bubbletea/pull/1133
Fixes: https://github.com/charmbracelet/bubbletea/issues/1127
---
 tea.go | 51 +++++++++++++++++++++++++++++----------------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/tea.go b/tea.go
index e7865a8ec5..87211ba299 100644
--- a/tea.go
+++ b/tea.go
@@ -128,6 +128,10 @@ func (h channelHandlers) shutdown() {
 type Program struct {
 	initialModel Model
 
+	// handlers is a list of channels that need to be waited on before the
+	// program can exit.
+	handlers channelHandlers
+
 	// Configuration options that will set as the program is initializing,
 	// treated as bits. These options can be set via various ProgramOptions.
 	startupOptions startupOptions
@@ -465,7 +469,7 @@ func (p *Program) eventLoop(model Model, cmds chan Cmd) (Model, error) {
 // terminated by either [Program.Quit], [Program.Kill], or its signal handler.
 // Returns the final model.
 func (p *Program) Run() (Model, error) {
-	handlers := channelHandlers{}
+	p.handlers = channelHandlers{}
 	cmds := make(chan Cmd)
 	p.errs = make(chan error)
 	p.finished = make(chan struct{}, 1)
@@ -512,7 +516,7 @@ func (p *Program) Run() (Model, error) {
 
 	// Handle signals.
 	if !p.startupOptions.has(withoutSignalHandler) {
-		handlers.add(p.handleSignals())
+		p.handlers.add(p.handleSignals())
 	}
 
 	// Recover from panics.
@@ -559,7 +563,7 @@ func (p *Program) Run() (Model, error) {
 	model := p.initialModel
 	if initCmd := model.Init(); initCmd != nil {
 		ch := make(chan struct{})
-		handlers.add(ch)
+		p.handlers.add(ch)
 
 		go func() {
 			defer close(ch)
@@ -582,10 +586,10 @@ func (p *Program) Run() (Model, error) {
 	}
 
 	// Handle resize events.
-	handlers.add(p.handleResize())
+	p.handlers.add(p.handleResize())
 
 	// Process commands.
-	handlers.add(p.handleCommands(cmds))
+	p.handlers.add(p.handleCommands(cmds))
 
 	// Run event loop, handle updates and draw.
 	model, err := p.eventLoop(model, cmds)
@@ -597,21 +601,6 @@ func (p *Program) Run() (Model, error) {
 		p.renderer.write(model.View())
 	}
 
-	// Tear down.
-	p.cancel()
-
-	// Check if the cancel reader has been setup before waiting and closing.
-	if p.cancelReader != nil {
-		// Wait for input loop to finish.
-		if p.cancelReader.Cancel() {
-			p.waitForReadLoop()
-		}
-		_ = p.cancelReader.Close()
-	}
-
-	// Wait for all handlers to finish.
-	handlers.shutdown()
-
 	// Restore terminal state.
 	p.shutdown(killed)
 
@@ -666,7 +655,7 @@ func (p *Program) Quit() {
 // The final render that you would normally see when quitting will be skipped.
 // [program.Run] returns a [ErrProgramKilled] error.
 func (p *Program) Kill() {
-	p.cancel()
+	p.shutdown(true)
 }
 
 // Wait waits/blocks until the underlying Program finished shutting down.
@@ -677,6 +666,22 @@ func (p *Program) Wait() {
 // shutdown performs operations to free up resources and restore the terminal
 // to its original state.
 func (p *Program) shutdown(kill bool) {
+	p.cancel()
+
+	// Wait for all handlers to finish.
+	p.handlers.shutdown()
+
+	// Check if the cancel reader has been setup before waiting and closing.
+	if p.cancelReader != nil {
+		// Wait for input loop to finish.
+		if p.cancelReader.Cancel() {
+			if !kill {
+				p.waitForReadLoop()
+			}
+		}
+		_ = p.cancelReader.Close()
+	}
+
 	if p.renderer != nil {
 		if kill {
 			p.renderer.kill()
@@ -686,7 +691,9 @@ func (p *Program) shutdown(kill bool) {
 	}
 
 	_ = p.restoreTerminalState()
-	p.finished <- struct{}{}
+	if !kill {
+		p.finished <- struct{}{}
+	}
 }
 
 // recoverFromPanic recovers from a panic, prints the stack trace, and restores