fix: restore caller panic hook when terminal session ends

src/tui.rs

@@ -1,4 +1,8 @@
-use std::io::{self, Stdout};
+use std::{
+    io::{self, Stdout},
+    panic::{self, AssertUnwindSafe, PanicHookInfo},
+    sync::{Arc, Mutex},
+};
 
 use crossterm::{
     cursor::Show,
@@ -28,7 +32,16 @@ use crate::{
 pub const CAPABILITIES: &[&str] = &["text", "styled_text"];
 const TAB_WIDTH: usize = 4;
 
-struct TerminalSession;
+type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + Send + Sync + 'static>;
+
+struct HookState {
+    previous: Option<PanicHook>,
+    active: bool,
+}
+
+struct TerminalSession {
+    hook: Arc<Mutex<HookState>>,
+}
 
 fn restore_terminal() {
     let _ = disable_raw_mode();
@@ -41,45 +54,82 @@ fn restore_terminal() {
 }
 
 impl TerminalSession {
-    fn enter() -> io::Result<Self> {
-        enable_raw_mode()?;
-        let session = Self;
-        execute!(io::stdout(), EnterAlternateScreen, EnableBracketedPaste)?;
-        Ok(session)
+    fn run(body: impl FnOnce() -> io::Result<()>) -> io::Result<()> {
+        let hook = Arc::new(Mutex::new(HookState {
+            previous: Some(panic::take_hook()),
+            active: false,
+        }));
+        let during_session = Arc::clone(&hook);
+        panic::set_hook(Box::new(move |info| {
+            let mut state = during_session
+                .lock()
+                .unwrap_or_else(|error| error.into_inner());
+            if state.active {
+                restore_terminal();
+                state.active = false;
+            }
+            if let Some(previous) = &state.previous {
+                previous(info);
+            }
+        }));
+        let session = Self { hook };
+        // set_hook is forbidden on a panicking thread. Catch first, drop the
+        // owner while not panicking, then resume the original panic unchanged.
+        let result = panic::catch_unwind(AssertUnwindSafe(|| {
+            enable_raw_mode()?;
+            session
+                .hook
+                .lock()
+                .unwrap_or_else(|error| error.into_inner())
+                .active = true;
+            execute!(io::stdout(), EnterAlternateScreen, EnableBracketedPaste)?;
+            body()
+        }));
+        drop(session);
+        match result {
+            Ok(result) => result,
+            Err(payload) => panic::resume_unwind(payload),
+        }
     }
 }
 
 impl Drop for TerminalSession {
     fn drop(&mut self) {
-        restore_terminal();
+        let mut state = self.hook.lock().unwrap_or_else(|error| error.into_inner());
+        if state.active {
+            restore_terminal();
+            state.active = false;
+        }
+        let previous = state
+            .previous
+            .take()
+            .expect("terminal session owns its prior hook");
+        drop(state);
+        panic::set_hook(previous);
     }
 }
 
 /// Renders session state and translates terminal input; the session owns plugin
 /// execution and retirement, so the adapter needs no lifecycle callback.
 pub fn run(session: &mut EditorSession) -> io::Result<()> {
-    let previous_hook = std::panic::take_hook();
-    std::panic::set_hook(Box::new(move |info| {
-        restore_terminal();
-        previous_hook(info);
-    }));
-    let _session = TerminalSession::enter()?;
-    let mut terminal: Terminal<CrosstermBackend<Stdout>> =
-        Terminal::new(CrosstermBackend::new(io::stdout()))?;
-    let (mut top, mut left) = (0, 0);
-    while !session.editor().should_quit() {
-        terminal.draw(|frame| draw(frame, session.editor(), &mut top, &mut left))?;
-        match event::read()? {
-            Event::Key(key) if key.kind != KeyEventKind::Release => {
-                if let Some(input) = translate(key) {
-                    session.handle(input);
+    TerminalSession::run(|| {
+        let mut terminal: Terminal<CrosstermBackend<Stdout>> =
+            Terminal::new(CrosstermBackend::new(io::stdout()))?;
+        let (mut top, mut left) = (0, 0);
+        while !session.editor().should_quit() {
+            terminal.draw(|frame| draw(frame, session.editor(), &mut top, &mut left))?;
+            match event::read()? {
+                Event::Key(key) if key.kind != KeyEventKind::Release => {
+                    if let Some(input) = translate(key) {
+                        session.handle(input);
+                    }
                 }
+                Event::Paste(text) => session.handle(Input::Paste(text)),
+                _ => {}
             }
-            Event::Paste(text) => session.handle(Input::Paste(text)),
-            _ => {}
         }
-    }
-    Ok(())
+        Ok(())
+    })
 }
 
 fn translate(key: KeyEvent) -> Option<Input> {
@@ -231,4 +281,119 @@ mod tests {
         assert_eq!(visible_line(text, 2, 5), "  界e\u{301}");
         assert_eq!(visible_line(text, 5, 2), " e\u{301}");
     }
+
+    #[cfg(unix)]
+    #[test]
+    fn terminal_session_scopes_the_callers_panic_hook() {
+        const CHILD: &str = "VAMP_TERMINAL_HOOK_PROBE";
+        if let Ok(mode) = std::env::var(CHILD) {
+            std::panic::set_hook(Box::new(|_| {
+                assert!(!crossterm::terminal::is_raw_mode_enabled().unwrap());
+                println!("CALLER_HOOK");
+            }));
+            if mode == "entry" {
+                let mut editor = EditorSession::open(None, vec![], None).unwrap();
+                assert!(run(&mut editor).is_err());
+            } else {
+                let result = panic::catch_unwind(|| {
+                    TerminalSession::run(|| {
+                        assert!(crossterm::terminal::is_raw_mode_enabled().unwrap());
+                        match mode.as_str() {
+                            "normal" => Ok(()),
+                            "error" => Err(io::Error::from(io::ErrorKind::BrokenPipe)),
+                            "panic" => panic::panic_any(73_u64),
+                            _ => unreachable!(),
+                        }
+                    })
+                });
+                match mode.as_str() {
+                    "normal" => result.unwrap().unwrap(),
+                    "error" => assert_eq!(
+                        result.unwrap().unwrap_err().kind(),
+                        io::ErrorKind::BrokenPipe
+                    ),
+                    "panic" => assert_eq!(result.unwrap_err().downcast_ref::<u64>(), Some(&73)),
+                    _ => unreachable!(),
+                }
+            }
+            assert!(!crossterm::terminal::is_raw_mode_enabled().unwrap());
+            println!("AFTER_TERMINAL_RETURN");
+            let _ = panic::catch_unwind(|| panic!("outside the terminal session"));
+            return;
+        }
+
+        use std::os::unix::process::CommandExt;
+        use std::process::{Command, Stdio};
+        unsafe extern "C" {
+            fn setsid() -> i32;
+        }
+        let mut command = Command::new(std::env::current_exe().unwrap());
+        command
+            .args([
+                "--exact",
+                "tui::tests::terminal_session_scopes_the_callers_panic_hook",
+                "--nocapture",
+                "--color",
+                "never",
+            ])
+            .env(CHILD, "entry")
+            .stdin(Stdio::null());
+        // SAFETY: setsid is async-signal-safe and detaches the child from any
+        // controlling terminal without touching inherited Rust state.
+        unsafe {
+            command.pre_exec(|| {
+                if setsid() == -1 {
+                    Err(io::Error::last_os_error())
+                } else {
+                    Ok(())
+                }
+            });
+        }
+        let check = |output: std::process::Output| {
+            assert!(output.status.success(), "{output:?}");
+            let stdout = String::from_utf8(output.stdout).unwrap();
+            let after = stdout.split_once("AFTER_TERMINAL_RETURN").unwrap().1;
+            assert!(after.contains("CALLER_HOOK"), "{after:?}");
+            assert!(
+                !after.contains('\u{1b}'),
+                "terminal hook outlived its session: {after:?}"
+            );
+        };
+        check(command.output().unwrap());
+
+        // Ruby's standard-library PTY gives each positive-entry scenario a real
+        // terminal, keeping process-global hooks isolated from the test runner.
+        for mode in ["normal", "error", "panic"] {
+            let output =
+                Command::new(std::env::var_os("VAMP_RUBY").unwrap_or_else(|| "ruby".into()))
+                    .args([
+                        "-rpty",
+                        "-e",
+                        r#"
+PTY.spawn(*ARGV) do |output, input, pid|
+  input.close
+  begin
+    loop { STDOUT.write(output.readpartial(4096)) }
+  rescue EOFError, Errno::EIO
+  end
+  Process.wait(pid)
+  exit($?.exitstatus || 1)
+end
+"#,
+                        "--",
+                    ])
+                    .arg(std::env::current_exe().unwrap())
+                    .args([
+                        "--exact",
+                        "tui::tests::terminal_session_scopes_the_callers_panic_hook",
+                        "--nocapture",
+                        "--color",
+                        "never",
+                    ])
+                    .env(CHILD, mode)
+                    .output()
+                    .unwrap();
+            check(output);
+        }
+    }
 }