| 1 | package cli |
| 2 | |
| 3 | import ( |
| 4 | "errors" |
| 5 | "io" |
| 6 | "sync/atomic" |
| 7 | "testing" |
| 8 | "time" |
| 9 | |
| 10 | tea "charm.land/bubbletea/v2" |
| 11 | |
| 12 | "reasonix/internal/control" |
| 13 | "reasonix/internal/session" |
| 14 | ) |
| 15 | |
| 16 | type shutdownSnapshotSpy struct { |
| 17 | control.SessionAPI |
| 18 | err error |
| 19 | started chan<- struct{} |
| 20 | release <-chan struct{} |
| 21 | snapshotCalls atomic.Int32 |
| 22 | shutdownCalls atomic.Int32 |
| 23 | } |
| 24 | |
| 25 | // The spy stands for a controller whose session identity is gone: the reclaim |
| 26 | // path reads both locators and must find neither. |
| 27 | func (s *shutdownSnapshotSpy) SessionRef() (session.SessionRef, bool) { |
| 28 | return session.SessionRef{}, false |
| 29 | } |
| 30 | |
| 31 | func (s *shutdownSnapshotSpy) SessionPath() string { return "" } |
| 32 | |
| 33 | func (s *shutdownSnapshotSpy) Snapshot() error { |
| 34 | s.snapshotCalls.Add(1) |
| 35 | return nil |
| 36 | } |
| 37 | |
| 38 | func (s *shutdownSnapshotSpy) SnapshotForShutdown() error { |
| 39 | s.shutdownCalls.Add(1) |
| 40 | if s.started != nil { |
| 41 | s.started <- struct{}{} |
| 42 | } |
| 43 | if s.release != nil { |
| 44 | <-s.release |
| 45 | } |
| 46 | return s.err |
| 47 | } |
| 48 | |
| 49 | func TestTUIShutdownUsesRecoveringSnapshotAndKeepsFailure(t *testing.T) { |
| 50 | wantErr := errors.New("final snapshot failed") |
| 51 | ctrl := &shutdownSnapshotSpy{err: wantErr} |
| 52 | m := newTestChatTUI() |
| 53 | m.ctrl = ctrl |
| 54 | completion := newTUIShutdownCompletion() |
| 55 | |
| 56 | next, cmd := m.update(tuiShutdownMsg{completion: completion}) |
| 57 | got := next.(chatTUI) |
| 58 | if cmd == nil || cmd() != (tea.QuitMsg{}) { |
| 59 | t.Fatal("shutdown message did not return tea.Quit") |
| 60 | } |
| 61 | if calls := ctrl.snapshotCalls.Load(); calls != 0 { |
| 62 | t.Fatalf("plain Snapshot calls = %d, want 0", calls) |
| 63 | } |
| 64 | if calls := ctrl.shutdownCalls.Load(); calls != 1 { |
| 65 | t.Fatalf("SnapshotForShutdown calls = %d, want 1", calls) |
| 66 | } |
| 67 | if !errors.Is(got.shutdownErr, wantErr) { |
| 68 | t.Fatalf("shutdownErr = %v, want %v", got.shutdownErr, wantErr) |
| 69 | } |
| 70 | select { |
| 71 | case <-completion.done: |
| 72 | default: |
| 73 | t.Fatal("shutdown completion was not acknowledged after the final snapshot") |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | // TestTUIShutdownSignalQuitsAfterReclaim pins the post-reclaim contract: once |
| 78 | // the remote side owns the session, SIGHUP/SIGTERM terminate the process like |
| 79 | // any other exit, without snapshotting a session this TUI no longer writes. |
| 80 | // Consuming the signal here left an orphan after an SSH drop and forced |
| 81 | // SIGKILL under systemctl stop. |
| 82 | func TestTUIShutdownSignalQuitsAfterReclaim(t *testing.T) { |
| 83 | cases := []struct { |
| 84 | name string |
| 85 | setup func(m *chatTUI) |
| 86 | }{ |
| 87 | {"after the reclaim callback", func(m *chatTUI) { m.sessionReclaimed = true }}, |
| 88 | {"after the return but before the callback", func(m *chatTUI) { |
| 89 | m.takeover = newCLITakeoverManager(nil, nil) |
| 90 | m.takeover.returned.Store(true) |
| 91 | }}, |
| 92 | } |
| 93 | for _, tc := range cases { |
| 94 | t.Run(tc.name, func(t *testing.T) { |
| 95 | ctrl := &shutdownSnapshotSpy{} |
| 96 | m := newTestChatTUI() |
| 97 | m.ctrl = ctrl |
| 98 | tc.setup(&m) |
| 99 | completion := newTUIShutdownCompletion() |
| 100 | |
| 101 | _, cmd := m.update(tuiShutdownMsg{completion: completion}) |
| 102 | if cmd == nil || cmd() != (tea.QuitMsg{}) { |
| 103 | t.Fatal("signal shutdown after reclaim did not return tea.Quit") |
| 104 | } |
| 105 | if calls := ctrl.shutdownCalls.Load(); calls != 0 { |
| 106 | t.Fatalf("SnapshotForShutdown calls = %d, want 0 for a session the remote side owns", calls) |
| 107 | } |
| 108 | select { |
| 109 | case <-completion.done: |
| 110 | default: |
| 111 | t.Fatal("shutdown completion was not acknowledged") |
| 112 | } |
| 113 | }) |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | // TestTUIShutdownSignalDeferredWhileReclaiming keeps the in-flight guard: a |
| 118 | // signal during the handoff transaction must not race the manager's final |
| 119 | // snapshot, but it is honored as soon as the reclaim callback lands. |
| 120 | func TestTUIShutdownSignalDeferredWhileReclaiming(t *testing.T) { |
| 121 | ctrl := &shutdownSnapshotSpy{} |
| 122 | m := newTestChatTUI() |
| 123 | m.ctrl = ctrl |
| 124 | m.takeover = newCLITakeoverManager(nil, nil) |
| 125 | m.takeover.reclaiming.Store(true) |
| 126 | completion := newTUIShutdownCompletion() |
| 127 | |
| 128 | next, cmd := m.update(tuiShutdownMsg{completion: completion}) |
| 129 | if cmd != nil { |
| 130 | t.Fatalf("shutdown during reclaim returned %T, want no quit command", cmd) |
| 131 | } |
| 132 | got := next.(chatTUI) |
| 133 | if !got.takeover.Reclaiming() { |
| 134 | t.Fatal("reclaim marker was cleared by shutdown race") |
| 135 | } |
| 136 | if !got.shutdownAfterReclaim { |
| 137 | t.Fatal("signal during reclaim was dropped instead of deferred") |
| 138 | } |
| 139 | select { |
| 140 | case <-completion.done: |
| 141 | default: |
| 142 | t.Fatal("shutdown during reclaim was not acknowledged") |
| 143 | } |
| 144 | if calls := ctrl.shutdownCalls.Load(); calls != 0 { |
| 145 | t.Fatalf("SnapshotForShutdown calls = %d during reclaim, want 0", calls) |
| 146 | } |
| 147 | |
| 148 | // The handoff completes and its callback arrives: the deferred exit fires |
| 149 | // without snapshotting the session the remote side now owns. |
| 150 | got.takeover.reclaiming.Store(false) |
| 151 | got.takeover.returned.Store(true) |
| 152 | next, cmd = got.update(tuiSessionReclaimedMsg{}) |
| 153 | if cmd == nil || cmd() != (tea.QuitMsg{}) { |
| 154 | t.Fatal("deferred signal shutdown did not quit after the reclaim callback") |
| 155 | } |
| 156 | if !next.(chatTUI).sessionReclaimed { |
| 157 | t.Fatal("reclaim callback did not mark the session reclaimed before quitting") |
| 158 | } |
| 159 | if calls := ctrl.shutdownCalls.Load(); calls != 0 { |
| 160 | t.Fatalf("SnapshotForShutdown calls = %d after reclaim, want 0", calls) |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | func TestBubbleTeaKeepsRunningWhenShutdownRacesReclaim(t *testing.T) { |
| 165 | m := newTestChatTUI() |
| 166 | m.takeover = newCLITakeoverManager(nil, nil) |
| 167 | m.takeover.reclaiming.Store(true) |
| 168 | p := tea.NewProgram( |
| 169 | shutdownOnlyProgramModel{chatTUI: m}, |
| 170 | tea.WithInput(nil), |
| 171 | tea.WithOutput(io.Discard), |
| 172 | tea.WithoutRenderer(), |
| 173 | tea.WithoutSignals(), |
| 174 | ) |
| 175 | |
| 176 | type result struct { |
| 177 | model tea.Model |
| 178 | err error |
| 179 | } |
| 180 | done := make(chan result, 1) |
| 181 | go func() { |
| 182 | model, err := p.Run() |
| 183 | done <- result{model: model, err: err} |
| 184 | }() |
| 185 | assertRunning := func(stage string) { |
| 186 | t.Helper() |
| 187 | select { |
| 188 | case result := <-done: |
| 189 | t.Fatalf("shutdown ended Bubble Tea %s: model=%T err=%v", stage, result.model, result.err) |
| 190 | case <-time.After(100 * time.Millisecond): |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | // A signal that races the handoff transaction must not tear the program |
| 195 | // down underneath the manager's final snapshot... |
| 196 | p.Send(tuiShutdownMsg{}) |
| 197 | assertRunning("while reclaiming") |
| 198 | // ...but it is a real exit request: once the reclaim callback lands the |
| 199 | // program leaves gracefully instead of lingering as an orphan. |
| 200 | m.takeover.returned.Store(true) |
| 201 | m.takeover.reclaiming.Store(false) |
| 202 | p.Send(tuiSessionReclaimedMsg{}) |
| 203 | select { |
| 204 | case result := <-done: |
| 205 | if result.err != nil { |
| 206 | t.Fatalf("deferred shutdown error = %v", result.err) |
| 207 | } |
| 208 | case <-time.After(time.Second): |
| 209 | p.Kill() |
| 210 | t.Fatal("Bubble Tea did not honor the signal deferred across the reclaim") |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | // shutdownOnlyProgramModel suppresses chatTUI's unrelated rendering and Init |
| 215 | // work while delegating messages to the production shutdown handler. |
| 216 | type shutdownOnlyProgramModel struct{ chatTUI } |
| 217 | |
| 218 | func (shutdownOnlyProgramModel) Init() tea.Cmd { return nil } |
| 219 | |
| 220 | func (m shutdownOnlyProgramModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { |
| 221 | next, cmd := m.update(msg) |
| 222 | return shutdownOnlyProgramModel{chatTUI: next.(chatTUI)}, cmd |
| 223 | } |
| 224 | |
| 225 | func (shutdownOnlyProgramModel) View() tea.View { return tea.NewView("") } |
| 226 | |
| 227 | func TestWatchdogDoesNotReclassifyCompletedBubbleTeaShutdownAsKilled(t *testing.T) { |
| 228 | started := make(chan struct{}) |
| 229 | release := make(chan struct{}) |
| 230 | ctrl := &shutdownSnapshotSpy{started: started, release: release} |
| 231 | m := newTestChatTUI() |
| 232 | m.ctrl = ctrl |
| 233 | p := tea.NewProgram( |
| 234 | shutdownOnlyProgramModel{chatTUI: m}, |
| 235 | tea.WithInput(nil), |
| 236 | tea.WithOutput(io.Discard), |
| 237 | tea.WithoutRenderer(), |
| 238 | tea.WithoutSignals(), |
| 239 | ) |
| 240 | |
| 241 | type runResult struct { |
| 242 | model tea.Model |
| 243 | err error |
| 244 | } |
| 245 | runDone := make(chan runResult, 1) |
| 246 | go func() { |
| 247 | model, err := p.Run() |
| 248 | runDone <- runResult{model: model, err: err} |
| 249 | }() |
| 250 | |
| 251 | scheduled := make(chan func(), 1) |
| 252 | completionSeen := make(chan *tuiShutdownCompletion, 1) |
| 253 | var kills atomic.Int32 |
| 254 | d := &tuiDiagnostics{ |
| 255 | afterFunc: func(delay time.Duration, fn func()) { |
| 256 | if delay != watchdogKillFallbackDelay { |
| 257 | t.Errorf("fallback delay = %s, want %s", delay, watchdogKillFallbackDelay) |
| 258 | } |
| 259 | scheduled <- fn |
| 260 | }, |
| 261 | shutdownFn: func(completion *tuiShutdownCompletion) { |
| 262 | completionSeen <- completion |
| 263 | p.Send(tuiShutdownMsg{completion: completion}) |
| 264 | }, |
| 265 | killFn: func() { |
| 266 | kills.Add(1) |
| 267 | p.Kill() |
| 268 | }, |
| 269 | } |
| 270 | killRequestDone := make(chan struct{}) |
| 271 | go func() { |
| 272 | d.doKill() |
| 273 | close(killRequestDone) |
| 274 | }() |
| 275 | |
| 276 | var fallback func() |
| 277 | select { |
| 278 | case fallback = <-scheduled: |
| 279 | case <-time.After(time.Second): |
| 280 | t.Fatal("watchdog fallback was not armed before shutdown") |
| 281 | } |
| 282 | var completion *tuiShutdownCompletion |
| 283 | select { |
| 284 | case completion = <-completionSeen: |
| 285 | case <-time.After(time.Second): |
| 286 | t.Fatal("watchdog did not send a completion-bearing shutdown message") |
| 287 | } |
| 288 | select { |
| 289 | case <-started: |
| 290 | case <-time.After(time.Second): |
| 291 | t.Fatal("Bubble Tea did not enter the final snapshot") |
| 292 | } |
| 293 | close(release) |
| 294 | select { |
| 295 | case <-completion.done: |
| 296 | case <-time.After(time.Second): |
| 297 | t.Fatal("final snapshot completed without acknowledging shutdown") |
| 298 | } |
| 299 | |
| 300 | // Exercise the old failure window: Update has finished the snapshot, but |
| 301 | // Bubble Tea may not have consumed tea.Quit yet when the timer callback runs. |
| 302 | fallback() |
| 303 | select { |
| 304 | case <-killRequestDone: |
| 305 | case <-time.After(time.Second): |
| 306 | t.Fatal("watchdog shutdown request remained blocked") |
| 307 | } |
| 308 | select { |
| 309 | case result := <-runDone: |
| 310 | if result.err != nil { |
| 311 | t.Fatalf("Bubble Tea shutdown error = %v, want graceful nil", result.err) |
| 312 | } |
| 313 | if _, ok := result.model.(shutdownOnlyProgramModel); !ok { |
| 314 | t.Fatalf("final model = %T, want shutdownOnlyProgramModel", result.model) |
| 315 | } |
| 316 | case <-time.After(time.Second): |
| 317 | p.Kill() |
| 318 | t.Fatal("Bubble Tea program did not exit") |
| 319 | } |
| 320 | if got := kills.Load(); got != 0 { |
| 321 | t.Fatalf("hard-kill calls = %d, want 0 after completed snapshot", got) |
| 322 | } |
| 323 | } |
| 324 |