Move panic hook out of messenger

crates/ruff_server/src/message.rs

@@ -1,6 +1,6 @@
-use std::sync::OnceLock;
-
+use anyhow::Context;
 use lsp_types::notification::Notification;
+use std::sync::OnceLock;
 
 use crate::server::ClientSender;
 
@@ -10,55 +10,31 @@ pub(crate) fn init_messenger(client_sender: ClientSender) {
     MESSENGER
         .set(client_sender)
         .expect("messenger should only be initialized once");
-
-    // unregister any previously registered panic hook
-    let _ = std::panic::take_hook();
-
-    // When we panic, try to notify the client.
-    std::panic::set_hook(Box::new(move |panic_info| {
-        use std::io::Write;
-
-        if let Some(messenger) = MESSENGER.get() {
-            let _ = messenger.send(lsp_server::Message::Notification(
-                lsp_server::Notification {
-                    method: lsp_types::notification::ShowMessage::METHOD.into(),
-                    params: serde_json::to_value(lsp_types::ShowMessageParams {
-                        typ: lsp_types::MessageType::ERROR,
-                        message: String::from(
-                            "The Ruff language server exited with a panic. See the logs for more details."
-                        ),
-                    })
-                    .unwrap_or_default(),
-                },
-            ));
-        }
-
-        let backtrace = std::backtrace::Backtrace::force_capture();
-        tracing::error!("{panic_info}\n{backtrace}");
-
-        // we also need to print to stderr directly in case tracing hasn't
-        // been initialized.
-        // But use `writeln` instead of `eprintln` to avoid panicking when the stderr pipe is broken.
-        let mut stderr = std::io::stderr().lock();
-        writeln!(stderr, "{panic_info}\n{backtrace}").ok();
-    }));
 }
 
 pub(crate) fn show_message(message: String, message_type: lsp_types::MessageType) {
+    try_show_message(message, message_type).unwrap();
+}
+
+pub(super) fn try_show_message(
+    message: String,
+    message_type: lsp_types::MessageType,
+) -> crate::Result<()> {
     MESSENGER
         .get()
-        .expect("messenger should be initialized")
+        .ok_or_else(|| anyhow::anyhow!("messenger not initialized"))?
         .send(lsp_server::Message::Notification(
             lsp_server::Notification {
                 method: lsp_types::notification::ShowMessage::METHOD.into(),
                 params: serde_json::to_value(lsp_types::ShowMessageParams {
                     typ: message_type,
                     message,
-                })
-                .unwrap(),
+                })?,
             },
         ))
-        .expect("message should send");
+        .context("Failed to send message")?;
+
+    Ok(())
 }
 
 /// Sends an error to the client with a formatted message. The error is sent in a

crates/ruff_server/src/server.rs

@@ -1,10 +1,10 @@
 //! Scheduling, I/O, and API endpoints.
 
-use std::num::NonZeroUsize;
-use std::str::FromStr;
-
 use lsp_server as lsp;
 use lsp_types as types;
+use std::num::NonZeroUsize;
+use std::panic::PanicInfo;
+use std::str::FromStr;
 use types::ClientCapabilities;
 use types::CodeActionKind;
 use types::CodeActionOptions;
@@ -36,6 +36,7 @@ mod client;
 mod connection;
 mod schedule;
 
+use crate::message::try_show_message;
 pub(crate) use connection::ClientSender;
 
 pub(crate) type Result<T> = std::result::Result<T, api::Error>;
@@ -133,6 +134,46 @@ impl Server {
     }
 
     pub fn run(self) -> crate::Result<()> {
+        type PanicHook = Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>;
+        struct RestorePanicHook {
+            hook: Option<PanicHook>,
+        }
+
+        impl Drop for RestorePanicHook {
+            fn drop(&mut self) {
+                if let Some(hook) = self.hook.take() {
+                    std::panic::set_hook(hook);
+                }
+            }
+        }
+
+        // unregister any previously registered panic hook
+        // The hook will be restored when this function exits.
+        let _ = RestorePanicHook {
+            hook: Some(std::panic::take_hook()),
+        };
+
+        // When we panic, try to notify the client.
+        std::panic::set_hook(Box::new(move |panic_info| {
+            use std::io::Write;
+
+            let backtrace = std::backtrace::Backtrace::force_capture();
+            tracing::error!("{panic_info}\n{backtrace}");
+
+            // we also need to print to stderr directly for when using `$logTrace` because
+            // the message won't be sent to the client.
+            // But don't use `eprintln` because `eprintln` itself may panic if the pipe is broken.
+            let mut stderr = std::io::stderr().lock();
+            writeln!(stderr, "{panic_info}\n{backtrace}").ok();
+
+            try_show_message(
+                "The Ruff language server exited with a panic. See the logs for more details."
+                    .to_string(),
+                lsp_types::MessageType::ERROR,
+            )
+            .ok();
+        }));
+
         event_loop_thread(move || {
             Self::event_loop(
                 &self.connection,