fix(rome_lsp): lsp friendly catch unwind

crates/rome_diagnostics/src/lib.rs

@@ -9,6 +9,7 @@ pub mod diagnostic;
 pub mod display;
 pub mod error;
 pub mod location;
+pub mod panic;
 pub mod serde;
 
 mod suggestion;

crates/rome_diagnostics/src/panic.rs

@@ -0,0 +1,50 @@
+use std::panic::UnwindSafe;
+
+#[derive(Default, Debug)]
+pub struct PanicError {
+    pub info: String,
+    pub backtrace: Option<std::backtrace::Backtrace>,
+}
+
+thread_local! {
+    static LAST_PANIC: std::cell::Cell<Option<PanicError>> = std::cell::Cell::new(None);
+}
+
+impl std::fmt::Display for PanicError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let r = f.write_fmt(format_args!("{}\n", self.info));
+        if let Some(backtrace) = &self.backtrace {
+            f.write_fmt(format_args!("Backtrace: {}", backtrace))
+        } else {
+            r
+        }
+    }
+}
+
+/// Take and set a specific panic hook before calling `f` inside a `catch_unwind`, then
+/// return the old set_hook.
+///
+/// If `f` panicks am `Error` with the panic message plus backtrace will be returned.
+pub fn catch_unwind<F, R>(f: F) -> Result<R, PanicError>
+where
+    F: FnOnce() -> R + UnwindSafe,
+{
+    let prev = std::panic::take_hook();
+    std::panic::set_hook(Box::new(|info| {
+        let info = info.to_string();
+        let backtrace = std::backtrace::Backtrace::capture();
+        LAST_PANIC.with(|cell| {
+            cell.set(Some(PanicError {
+                info,
+                backtrace: Some(backtrace),
+            }))
+        })
+    }));
+
+    let result = std::panic::catch_unwind(f)
+        .map_err(|_| LAST_PANIC.with(|cell| cell.take()).unwrap_or_default());
+
+    std::panic::set_hook(prev);
+
+    result
+}

crates/rome_lsp/src/server.rs

@@ -10,6 +10,7 @@ use rome_fs::CONFIG_NAME;
 use rome_service::workspace::{RageEntry, RageParams, RageResult};
 use rome_service::{workspace, Workspace};
 use std::collections::HashMap;
+use std::panic::RefUnwindSafe;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -32,6 +33,8 @@ pub struct LSPServer {
     is_initialized: Arc<AtomicBool>,
 }
 
+impl RefUnwindSafe for LSPServer {}
+
 impl LSPServer {
     fn new(
         session: SessionHandle,
@@ -302,44 +305,59 @@ impl LanguageServer for LSPServer {
     }
 
     async fn code_action(&self, params: CodeActionParams) -> LspResult<Option<CodeActionResponse>> {
-        handlers::analysis::code_actions(&self.session, params).map_err(into_lsp_error)
+        rome_diagnostics::panic::catch_unwind(move || {
+            handlers::analysis::code_actions(&self.session, params).map_err(into_lsp_error)
+        })
+        .map_err(into_lsp_error)?
     }
 
     async fn formatting(
         &self,
         params: DocumentFormattingParams,
     ) -> LspResult<Option<Vec<TextEdit>>> {
-        handlers::formatting::format(&self.session, params).map_err(into_lsp_error)
+        rome_diagnostics::panic::catch_unwind(move || {
+            handlers::formatting::format(&self.session, params).map_err(into_lsp_error)
+        })
+        .map_err(into_lsp_error)?
     }
 
     async fn range_formatting(
         &self,
         params: DocumentRangeFormattingParams,
     ) -> LspResult<Option<Vec<TextEdit>>> {
-        handlers::formatting::format_range(&self.session, params).map_err(into_lsp_error)
+        rome_diagnostics::panic::catch_unwind(move || {
+            handlers::formatting::format_range(&self.session, params).map_err(into_lsp_error)
+        })
+        .map_err(into_lsp_error)?
     }
 
     async fn on_type_formatting(
         &self,
         params: DocumentOnTypeFormattingParams,
     ) -> LspResult<Option<Vec<TextEdit>>> {
-        handlers::formatting::format_on_type(&self.session, params).map_err(into_lsp_error)
+        rome_diagnostics::panic::catch_unwind(move || {
+            handlers::formatting::format_on_type(&self.session, params).map_err(into_lsp_error)
+        })
+        .map_err(into_lsp_error)?
     }
 
     async fn rename(&self, params: RenameParams) -> LspResult<Option<WorkspaceEdit>> {
-        let rename_enabled = self
-            .session
-            .config
-            .read()
-            .ok()
-            .and_then(|config| config.settings.rename)
-            .unwrap_or(false);
-
-        if rename_enabled {
-            handlers::rename::rename(&self.session, params).map_err(into_lsp_error)
-        } else {
-            Ok(None)
-        }
+        rome_diagnostics::panic::catch_unwind(move || {
+            let rename_enabled = self
+                .session
+                .config
+                .read()
+                .ok()
+                .and_then(|config| config.settings.rename)
+                .unwrap_or(false);
+
+            if rename_enabled {
+                handlers::rename::rename(&self.session, params).map_err(into_lsp_error)
+            } else {
+                Ok(None)
+            }
+        })
+        .map_err(into_lsp_error)?
     }
 }