Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions pkg/tui/tui.go
Original file line number Diff line number Diff line change
Expand Up @@ -2089,13 +2089,31 @@ func (m *appModel) windowTitle() string {
return title
}

// exitFunc is the function called by the shutdown safety net when the
// graceful exit times out. It defaults to os.Exit but can be replaced
// in tests.
var exitFunc = os.Exit

var shutdownTimeout = 5 * time.Second

// cleanupAll cleans up all sessions, editors, and resources.
func (m *appModel) cleanupAll() {
m.transcriber.Stop()
m.closeTranscriptCh()
for _, ed := range m.editors {
ed.Cleanup()
}

// Safety net: force-exit if bubbletea's shutdown gets stuck.
// This can happen when the renderer's flush goroutine blocks on a
// stdout write (terminal buffer full) while holding the renderer
// mutex, preventing the event loop from completing the render call
// that follows tea.Quit.
go func() {
time.Sleep(shutdownTimeout)
slog.Warn("Graceful shutdown timed out, forcing exit")
exitFunc(0)
}()
}

// persistedSessionID returns the session-store ID that should be used for
Expand Down
281 changes: 281 additions & 0 deletions pkg/tui/tui_exit_test.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
package tui

import (
"bytes"
"reflect"
"sync"
"sync/atomic"
"testing"
"time"

"charm.land/bubbles/v2/help"
"charm.land/bubbles/v2/key"
Expand Down Expand Up @@ -143,8 +147,19 @@ func newTestModel() (*appModel, *mockEditor) {
return m, ed
}

// neutralizeExitFunc replaces the package-level exitFunc with a no-op for the
// duration of the test so that the safety-net goroutine spawned by cleanupAll
// doesn't call os.Exit.
func neutralizeExitFunc(t *testing.T) {
t.Helper()
orig := exitFunc
exitFunc = func(int) {}
t.Cleanup(func() { exitFunc = orig })
}

func TestExitSessionMsg_ExitsImmediately(t *testing.T) {
t.Parallel()
neutralizeExitFunc(t)

m, ed := newTestModel()

Expand All @@ -158,6 +173,7 @@ func TestExitSessionMsg_ExitsImmediately(t *testing.T) {

func TestExitConfirmedMsg_ExitsImmediately(t *testing.T) {
t.Parallel()
neutralizeExitFunc(t)

m, ed := newTestModel()

Expand All @@ -168,3 +184,268 @@ func TestExitConfirmedMsg_ExitsImmediately(t *testing.T) {
msgs := collectMsgs(cmd)
assert.True(t, hasMsg[tea.QuitMsg](msgs), "should produce tea.QuitMsg")
}

// blockingWriter is an io.Writer whose Write blocks until unblocked.
type blockingWriter struct {
mu sync.Mutex
blocked chan struct{} // closed once the first Write starts blocking
gate chan struct{} // Write blocks until this is closed
}

func newBlockingWriter() *blockingWriter {
return &blockingWriter{
blocked: make(chan struct{}),
gate: make(chan struct{}),
}
}

func (w *blockingWriter) Write(p []byte) (int, error) {
w.mu.Lock()
select {
case <-w.blocked:
default:
close(w.blocked)
}
gate := w.gate
w.mu.Unlock()

<-gate
return len(p), nil
}

// reblock installs a new gate so that subsequent writes block again.
func (w *blockingWriter) reblock() {
w.mu.Lock()
w.gate = make(chan struct{})
w.mu.Unlock()
}

// unblock releases all pending and future writes.
func (w *blockingWriter) unblock() {
w.mu.Lock()
select {
case <-w.gate:
default:
close(w.gate)
}
w.mu.Unlock()
}

// quitModel is a minimal bubbletea model that requests alt-screen output
// and quits in response to a trigger message. An optional onQuit callback
// runs inside Update before tea.Quit is returned.
type quitModel struct {
onQuit func()
}

type triggerQuitMsg struct{}

func (m *quitModel) Init() tea.Cmd { return nil }

func (m *quitModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
if _, ok := msg.(triggerQuitMsg); ok {
if m.onQuit != nil {
m.onQuit()
}
return m, tea.Quit
}
return m, nil
}

func (m *quitModel) View() tea.View {
v := tea.NewView("hello world")
v.AltScreen = true
return v
}

// initBlockingBubbletea creates a bubbletea program whose output writer
// blocks. It lets the initial render complete (so the event loop is ready)
// then re-blocks the writer. Returns the program and the writer.
func initBlockingBubbletea(t *testing.T, model tea.Model) (*tea.Program, *blockingWriter, <-chan struct{}) {
t.Helper()

w := newBlockingWriter()
var in bytes.Buffer

p := tea.NewProgram(model,
tea.WithContext(t.Context()),
tea.WithInput(&in),
tea.WithOutput(w),
)

runDone := make(chan struct{})
go func() {
defer close(runDone)
_, _ = p.Run()
}()

// Wait for the initial render to hit the blocking writer.
select {
case <-w.blocked:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for initial write to block")
}

// Let the initial writes through so the event loop starts.
w.unblock()
time.Sleep(200 * time.Millisecond)

// Re-block so the next renderer flush will stall.
w.reblock()

return p, w, runDone
}

// TestCleanupAll_SpawnsSafetyNet verifies that cleanupAll spawns a goroutine
// that calls exitFunc after shutdownTimeout. Without the safety net, the
// process would hang when bubbletea's renderer deadlocks on exit.
func TestCleanupAll_SpawnsSafetyNet(t *testing.T) {
origTimeout := shutdownTimeout
origExitFunc := exitFunc
t.Cleanup(func() {
shutdownTimeout = origTimeout
exitFunc = origExitFunc
})
shutdownTimeout = 200 * time.Millisecond

exitDone := make(chan int, 1)
exitFunc = func(code int) {
exitDone <- code
}

m, _ := newTestModel()
m.cleanupAll()

select {
case code := <-exitDone:
assert.Equal(t, 0, code)
case <-time.After(shutdownTimeout + time.Second):
t.Fatal("exitFunc was not called — safety net is missing from cleanupAll")
}
}

// TestExitDeadlock_BlockedStdout proves that bubbletea's p.Run() hangs when
// stdout blocks during the final render after tea.Quit. This is the underlying
// bug that the safety net in cleanupAll works around.
func TestExitDeadlock_BlockedStdout(t *testing.T) {
t.Parallel()

model := &quitModel{}
p, w, runDone := initBlockingBubbletea(t, model)

// Trigger quit — the event loop will deadlock trying to render.
p.Send(triggerQuitMsg{})

// Verify that p.Run() does NOT return within a reasonable window.
select {
case <-runDone:
t.Skip("bubbletea returned without deadlocking; upstream fix may have landed")
case <-time.After(2 * time.Second):
// Expected: p.Run() is stuck.
}

// Unblock everything to let goroutines drain.
w.unblock()
}

// TestExitSafetyNet_BlockedStdout verifies that when bubbletea's renderer
// is stuck writing to stdout (terminal buffer full), the shutdown safety net
// forces the process to exit.
//
// Background: bubbletea's cursed renderer holds a mutex during io.Copy to
// stdout. If stdout blocks (e.g. full PTY buffer), the event loop's final
// render call after tea.Quit deadlocks on the same mutex. Without the safety
// net the process hangs forever.
func TestExitSafetyNet_BlockedStdout(t *testing.T) {
t.Parallel()

const safetyNetTimeout = 500 * time.Millisecond
var exitCalled atomic.Bool
exitDone := make(chan int, 1)
testExitFunc := func(code int) {
exitCalled.Store(true)
exitDone <- code
}

model := &quitModel{
onQuit: func() {
go func() {
time.Sleep(safetyNetTimeout)
testExitFunc(0)
}()
},
}
p, w, runDone := initBlockingBubbletea(t, model)
defer w.unblock()

// Trigger quit — the model's onQuit starts the safety net.
p.Send(triggerQuitMsg{})

select {
case code := <-exitDone:
assert.True(t, exitCalled.Load())
assert.Equal(t, 0, code)
case <-runDone:
// p.Run() returned on its own — also acceptable.
case <-time.After(safetyNetTimeout + 2*time.Second):
t.Fatal("neither p.Run() returned nor safety-net exitFunc fired within the deadline")
}
}

// TestExitSafetyNet_GracefulShutdown verifies that when bubbletea shuts down
// normally (no blocked stdout), p.Run() returns before the safety net fires.
func TestExitSafetyNet_GracefulShutdown(t *testing.T) {
t.Parallel()

const safetyNetTimeout = 2 * time.Second
var exitCalled atomic.Bool
testExitFunc := func(int) {
exitCalled.Store(true)
}

var mu sync.Mutex
cleanupCalled := false

model := &quitModel{
onQuit: func() {
mu.Lock()
cleanupCalled = true
mu.Unlock()
go func() {
time.Sleep(safetyNetTimeout)
testExitFunc(0)
}()
},
}
var buf bytes.Buffer
var in bytes.Buffer

p := tea.NewProgram(model,
tea.WithContext(t.Context()),
tea.WithInput(&in),
tea.WithOutput(&buf),
)

runDone := make(chan error, 1)
go func() {
_, err := p.Run()
runDone <- err
}()

// Give bubbletea time to initialise.
time.Sleep(200 * time.Millisecond)

p.Send(triggerQuitMsg{})

select {
case err := <-runDone:
require.NoError(t, err)
case <-time.After(3 * time.Second):
t.Fatal("p.Run() did not return within deadline for graceful shutdown")
}

mu.Lock()
assert.True(t, cleanupCalled, "cleanup should have been called")
mu.Unlock()
assert.False(t, exitCalled.Load(), "exitFunc should NOT fire during graceful shutdown")
}
Loading