diff --git a/src/main.rs b/src/main.rs index 39d214c..29b8229 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,9 +17,13 @@ mod task; mod transport; mod ui; -use std::io; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; +use std::{ + io, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; use crossterm::{ cursor::{Hide, Show}, @@ -37,6 +41,13 @@ use crossterm::{ use app::App; +/// Whether `PushKeyboardEnhancementFlags` actually executed, so restore pops +/// only what setup pushed. A static atomic because the panic hook is installed +/// before the terminal is touched and can fire on any thread; the hook and +/// `TerminalGuard` read the same source of truth. Only atomicity matters here +/// (the flag orders nothing else), so `Relaxed` suffices. +static KITTY_PUSHED: AtomicBool = AtomicBool::new(false); + const USAGE: &str = "\ fleetcom - a fleet-view supervisor for arbitrary shell commands @@ -155,6 +166,11 @@ fn main() -> io::Result<()> { let mut out = io::stdout(); enable_raw_mode()?; + // Armed the moment raw mode is on: every exit past this point (the `?`s + // below, a panic unwind, the normal return) must restore the terminal, + // or the shell is left in raw mode with a hidden cursor. The panic hook + // covers panics only, not `Err` returns. + let guard = TerminalGuard; // Probe for the kitty keyboard protocol before entering the alternate // screen: the query round-trips through the tty, and raw mode (just // enabled) is what keeps the reply out of the line discipline. `false` on @@ -171,6 +187,9 @@ fn main() -> io::Result<()> { out, PushKeyboardEnhancementFlags(KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES) )?; + // Recorded only after the execute succeeds: restore mirrors what + // setup actually did, not what it attempted. + KITTY_PUSHED.store(true, Ordering::Relaxed); } execute!(out, EnableBracketedPaste, Print("\x1b[?1007s\x1b[?1007h"))?; // `fleetcom [--foreground] ` loads that session at startup; the @@ -181,23 +200,52 @@ fn main() -> io::Result<()> { install_signal_handlers(app.signal_flag())?; let result = app.run(&mut out); - restore_terminal(&mut out); + // Consuming the guard restores here, at the same point the happy path + // always restored; the drop-on-unwind path exists for the `?`s above. + drop(guard); result } +/// Restores the terminal when dropped. Constructed only after +/// `enable_raw_mode` succeeds: before that there is nothing to undo. +struct TerminalGuard; + +impl Drop for TerminalGuard { + fn drop(&mut self) { + restore_terminal(&mut io::stdout()); + } +} + /// Restore terminal modes changed by the application. Cleanup is best-effort /// so a failed terminal write does not prevent raw mode from being disabled. +/// +/// May run twice: on a panic the hook restores first (so the message prints +/// on the normal screen), then `TerminalGuard`'s drop restores again during +/// unwind. Every step tolerates the repeat (leaving the alternate screen +/// twice, disabling raw mode twice, and popping past an empty kitty stack are +/// all no-ops or ignored), so don't "fix" the double restore by dropping one. fn restore_terminal(out: &mut io::Stdout) { - let _ = execute!( + let _ = emit_restore_sequences(out, KITTY_PUSHED.load(Ordering::Relaxed)); + let _ = disable_raw_mode(); +} + +/// Emit the escape sequences that undo terminal setup, mirroring what setup +/// actually did: the keyboard-enhancement pop only if the flags were pushed. +/// `disable_raw_mode` lives in `restore_terminal`, not here: it mutates +/// process-global tty state, and this function stays a pure emission so tests +/// can drive it against a buffer. +fn emit_restore_sequences(out: &mut impl io::Write, kitty_pushed: bool) -> io::Result<()> { + if kitty_pushed { + execute!(out, PopKeyboardEnhancementFlags)?; + } + execute!( out, - PopKeyboardEnhancementFlags, DisableMouseCapture, DisableBracketedPaste, Print("\x1b[?1007r"), Show, LeaveAlternateScreen - ); - let _ = disable_raw_mode(); + ) } /// Route external termination signals into the app's quit flag so the loop @@ -213,7 +261,9 @@ fn install_signal_handlers(flag: Arc) -> io::Result<()> { } /// Restore the terminal on panic. Otherwise a crash leaves the user in raw -/// mode on the alternate screen with no cursor. +/// mode on the alternate screen with no cursor. Restoring *before* the +/// default hook is what puts the panic message on the normal screen instead +/// of the vanishing alternate screen. fn install_panic_hook() { let default = std::panic::take_hook(); std::panic::set_hook(Box::new(move |info| { @@ -295,4 +345,42 @@ mod tests { fn second_session_name_is_rejected() { assert!(parse(&["one", "two"]).is_err()); } + + fn contains(haystack: &[u8], needle: &[u8]) -> bool { + haystack.windows(needle.len()).any(|w| w == needle) + } + + /// Render a single crossterm command to bytes, so the assertions below + /// track crossterm's actual encoding instead of hardcoding it. + fn encode(cmd: impl crossterm::Command) -> Vec { + let mut buf = Vec::new(); + execute!(buf, cmd).unwrap(); + buf + } + + /// Restore must mirror setup: the keyboard-enhancement pop is emitted + /// only when the flags were pushed, and its absence must not take the + /// rest of the restore down with it. + #[test] + fn restore_pops_kitty_flags_only_when_pushed() { + let pop = encode(PopKeyboardEnhancementFlags); + + let mut pushed = Vec::new(); + emit_restore_sequences(&mut pushed, true).unwrap(); + let mut unpushed = Vec::new(); + emit_restore_sequences(&mut unpushed, false).unwrap(); + + assert!(contains(&pushed, &pop)); + assert!(!contains(&unpushed, &pop)); + + // Both variants still emit the full remaining restore. + let leave = encode(LeaveAlternateScreen); + let show = encode(Show); + for out in [&pushed, &unpushed] { + assert!(contains(out, &leave)); + assert!(contains(out, &show)); + // Alternate-scroll restore is a raw Print, not a crossterm command. + assert!(contains(out, b"\x1b[?1007r")); + } + } }