[ty] Gracefully handle salsa cancellations and panics in background request handlers (#18254)

crates/ruff_benchmark/benches/ty.rs

@@ -131,7 +131,7 @@ fn benchmark_incremental(criterion: &mut Criterion) {
     fn setup() -> Case {
         let case = setup_tomllib_case();
 
-        let result: Vec<_> = case.db.check().unwrap();
+        let result: Vec<_> = case.db.check();
 
         assert_diagnostics(&case.db, &result, EXPECTED_TOMLLIB_DIAGNOSTICS);
 
@@ -159,7 +159,7 @@ fn benchmark_incremental(criterion: &mut Criterion) {
             None,
         );
 
-        let result = db.check().unwrap();
+        let result = db.check();
 
         assert_eq!(result.len(), EXPECTED_TOMLLIB_DIAGNOSTICS.len());
     }
@@ -179,7 +179,7 @@ fn benchmark_cold(criterion: &mut Criterion) {
             setup_tomllib_case,
             |case| {
                 let Case { db, .. } = case;
-                let result: Vec<_> = db.check().unwrap();
+                let result: Vec<_> = db.check();
 
                 assert_diagnostics(db, &result, EXPECTED_TOMLLIB_DIAGNOSTICS);
             },
@@ -293,7 +293,7 @@ fn benchmark_many_string_assignments(criterion: &mut Criterion) {
             },
             |case| {
                 let Case { db, .. } = case;
-                let result = db.check().unwrap();
+                let result = db.check();
                 assert_eq!(result.len(), 0);
             },
             BatchSize::SmallInput,
@@ -339,7 +339,7 @@ fn benchmark_many_tuple_assignments(criterion: &mut Criterion) {
             },
             |case| {
                 let Case { db, .. } = case;
-                let result = db.check().unwrap();
+                let result = db.check();
                 assert_eq!(result.len(), 0);
             },
             BatchSize::SmallInput,

crates/ruff_db/src/panic.rs

@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::backtrace::BacktraceStatus;
 use std::cell::Cell;
 use std::panic::Location;
@@ -24,17 +25,25 @@ impl Payload {
             None
         }
     }
+
+    pub fn downcast_ref<R: Any>(&self) -> Option<&R> {
+        self.0.downcast_ref::<R>()
+    }
 }
 
 impl std::fmt::Display for PanicError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        writeln!(f, "panicked at")?;
+        write!(f, "panicked at")?;
         if let Some(location) = &self.location {
             write!(f, " {location}")?;
         }
         if let Some(payload) = self.payload.as_str() {
             write!(f, ":\n{payload}")?;
         }
+        if let Some(query_trace) = self.salsa_backtrace.as_ref() {
+            let _ = writeln!(f, "{query_trace}");
+        }
+
         if let Some(backtrace) = &self.backtrace {
             match backtrace.status() {
                 BacktraceStatus::Disabled => {
@@ -49,6 +58,7 @@ impl std::fmt::Display for PanicError {
                 _ => {}
             }
         }
+
         Ok(())
     }
 }

crates/ty/src/lib.rs

@@ -243,12 +243,14 @@ impl MainLoop {
                 MainLoopMessage::CheckWorkspace => {
                     let db = db.clone();
                     let sender = self.sender.clone();
-                    let mut reporter = R::default();
 
                     // Spawn a new task that checks the project. This needs to be done in a separate thread
                     // to prevent blocking the main loop here.
                     rayon::spawn(move || {
-                        match db.check_with_reporter(&mut reporter) {
+                        match salsa::Cancelled::catch(|| {
+                            let mut reporter = R::default();
+                            db.check_with_reporter(&mut reporter)
+                        }) {
                             Ok(result) => {
                                 // Send the result back to the main loop for printing.
                                 sender

crates/ty/tests/file_watching.rs

@@ -1135,7 +1135,7 @@ print(sys.last_exc, os.getegid())
         Ok(())
     })?;
 
-    let diagnostics = case.db.check().context("Failed to check project.")?;
+    let diagnostics = case.db.check();
 
     assert_eq!(diagnostics.len(), 2);
     assert_eq!(
@@ -1160,7 +1160,7 @@ print(sys.last_exc, os.getegid())
     })
     .expect("Search path settings to be valid");
 
-    let diagnostics = case.db.check().context("Failed to check project.")?;
+    let diagnostics = case.db.check();
     assert!(diagnostics.is_empty());
 
     Ok(())
@@ -1763,10 +1763,7 @@ fn changes_to_user_configuration() -> anyhow::Result<()> {
     let foo = case
         .system_file(case.project_path("foo.py"))
         .expect("foo.py to exist");
-    let diagnostics = case
+    let diagnostics = case.db().check_file(foo);
-        .db()
-        .check_file(foo)
-        .context("Failed to check project.")?;
 
     assert!(
         diagnostics.is_empty(),
@@ -1786,10 +1783,7 @@ fn changes_to_user_configuration() -> anyhow::Result<()> {
 
     case.apply_changes(changes);
 
-    let diagnostics = case
+    let diagnostics = case.db().check_file(foo);
-        .db()
-        .check_file(foo)
-        .context("Failed to check project.")?;
 
     assert!(
         diagnostics.len() == 1,

crates/ty_project/src/db.rs

@@ -8,8 +8,8 @@ use ruff_db::files::{File, Files};
 use ruff_db::system::System;
 use ruff_db::vendored::VendoredFileSystem;
 use ruff_db::{Db as SourceDb, Upcast};
+use salsa::Event;
 use salsa::plumbing::ZalsaDatabase;
-use salsa::{Cancelled, Event};
 use ty_ide::Db as IdeDb;
 use ty_python_semantic::lint::{LintRegistry, RuleSelection};
 use ty_python_semantic::{Db as SemanticDb, Program};
@@ -76,24 +76,21 @@ impl ProjectDatabase {
     }
 
     /// Checks all open files in the project and its dependencies.
-    pub fn check(&self) -> Result<Vec<Diagnostic>, Cancelled> {
+    pub fn check(&self) -> Vec<Diagnostic> {
         let mut reporter = DummyReporter;
         let reporter = AssertUnwindSafe(&mut reporter as &mut dyn Reporter);
-        self.with_db(|db| db.project().check(db, reporter))
+        self.project().check(self, reporter)
     }
 
     /// Checks all open files in the project and its dependencies, using the given reporter.
-    pub fn check_with_reporter(
+    pub fn check_with_reporter(&self, reporter: &mut dyn Reporter) -> Vec<Diagnostic> {
-        &self,
-        reporter: &mut dyn Reporter,
-    ) -> Result<Vec<Diagnostic>, Cancelled> {
         let reporter = AssertUnwindSafe(reporter);
-        self.with_db(|db| db.project().check(db, reporter))
+        self.project().check(self, reporter)
     }
 
     #[tracing::instrument(level = "debug", skip(self))]
-    pub fn check_file(&self, file: File) -> Result<Vec<Diagnostic>, Cancelled> {
+    pub fn check_file(&self, file: File) -> Vec<Diagnostic> {
-        self.with_db(|db| self.project().check_file(db, file))
+        self.project().check_file(self, file)
     }
 
     /// Returns a mutable reference to the system.
@@ -107,13 +104,6 @@ impl ProjectDatabase {
         Arc::get_mut(&mut self.system)
             .expect("ref count should be 1 because `zalsa_mut` drops all other DB references.")
     }
-
-    pub(crate) fn with_db<F, T>(&self, f: F) -> Result<T, Cancelled>
-    where
-        F: FnOnce(&ProjectDatabase) -> T + std::panic::UnwindSafe,
-    {
-        Cancelled::catch(|| f(self))
-    }
 }
 
 impl Upcast<dyn SemanticDb> for ProjectDatabase {

crates/ty_server/src/server.rs

@@ -1,10 +1,5 @@
 //! Scheduling, I/O, and API endpoints.
 
-use std::num::NonZeroUsize;
-// The new PanicInfoHook name requires MSRV >= 1.82
-#[expect(deprecated)]
-use std::panic::PanicInfo;
-
 use lsp_server::Message;
 use lsp_types::{
     ClientCapabilities, DiagnosticOptions, DiagnosticServerCapabilities,
@@ -13,6 +8,8 @@ use lsp_types::{
     TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
     TypeDefinitionProviderCapability, Url,
 };
+use std::num::NonZeroUsize;
+use std::panic::PanicHookInfo;
 
 use self::connection::{Connection, ConnectionInitializer};
 use self::schedule::event_loop_thread;
@@ -125,9 +122,7 @@ impl Server {
     }
 
     pub(crate) fn run(self) -> crate::Result<()> {
-        // The new PanicInfoHook name requires MSRV >= 1.82
+        type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + 'static + Sync + Send>;
-        #[expect(deprecated)]
-        type PanicHook = Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send>;
         struct RestorePanicHook {
             hook: Option<PanicHook>,
         }

crates/ty_server/src/server/api.rs

@@ -3,42 +3,41 @@ use crate::session::Session;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use anyhow::anyhow;
 use lsp_server as server;
+use lsp_server::RequestId;
 use lsp_types::notification::Notification;
+use ruff_db::panic::PanicError;
+use std::panic::UnwindSafe;
 
 mod diagnostics;
 mod notifications;
 mod requests;
 mod traits;
 
-use notifications as notification;
-use requests as request;
-
 use self::traits::{NotificationHandler, RequestHandler};
-
 use super::{Result, client::Responder, schedule::BackgroundSchedule};
 
 pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
     let id = req.id.clone();
 
     match req.method.as_str() {
-        request::DocumentDiagnosticRequestHandler::METHOD => background_request_task::<
+        requests::DocumentDiagnosticRequestHandler::METHOD => background_request_task::<
-            request::DocumentDiagnosticRequestHandler,
+            requests::DocumentDiagnosticRequestHandler,
         >(
             req, BackgroundSchedule::Worker
         ),
-        request::GotoTypeDefinitionRequestHandler::METHOD => background_request_task::<
+        requests::GotoTypeDefinitionRequestHandler::METHOD => background_request_task::<
-            request::GotoTypeDefinitionRequestHandler,
+            requests::GotoTypeDefinitionRequestHandler,
         >(
             req, BackgroundSchedule::Worker
         ),
-        request::HoverRequestHandler::METHOD => {
+        requests::HoverRequestHandler::METHOD => background_request_task::<
-            background_request_task::<request::HoverRequestHandler>(req, BackgroundSchedule::Worker)
+            requests::HoverRequestHandler,
-        }
-        request::InlayHintRequestHandler::METHOD => background_request_task::<
-            request::InlayHintRequestHandler,
         >(req, BackgroundSchedule::Worker),
-        request::CompletionRequestHandler::METHOD => background_request_task::<
+        requests::InlayHintRequestHandler::METHOD => background_request_task::<
-            request::CompletionRequestHandler,
+            requests::InlayHintRequestHandler,
+        >(req, BackgroundSchedule::Worker),
+        requests::CompletionRequestHandler::METHOD => background_request_task::<
+            requests::CompletionRequestHandler,
         >(
             req, BackgroundSchedule::LatencySensitive
         ),
@@ -64,23 +63,23 @@ pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
 
 pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
     match notif.method.as_str() {
-        notification::DidCloseTextDocumentHandler::METHOD => {
+        notifications::DidCloseTextDocumentHandler::METHOD => {
-            local_notification_task::<notification::DidCloseTextDocumentHandler>(notif)
+            local_notification_task::<notifications::DidCloseTextDocumentHandler>(notif)
         }
-        notification::DidOpenTextDocumentHandler::METHOD => {
+        notifications::DidOpenTextDocumentHandler::METHOD => {
-            local_notification_task::<notification::DidOpenTextDocumentHandler>(notif)
+            local_notification_task::<notifications::DidOpenTextDocumentHandler>(notif)
         }
-        notification::DidChangeTextDocumentHandler::METHOD => {
+        notifications::DidChangeTextDocumentHandler::METHOD => {
-            local_notification_task::<notification::DidChangeTextDocumentHandler>(notif)
+            local_notification_task::<notifications::DidChangeTextDocumentHandler>(notif)
         }
-        notification::DidOpenNotebookHandler::METHOD => {
+        notifications::DidOpenNotebookHandler::METHOD => {
-            local_notification_task::<notification::DidOpenNotebookHandler>(notif)
+            local_notification_task::<notifications::DidOpenNotebookHandler>(notif)
         }
-        notification::DidCloseNotebookHandler::METHOD => {
+        notifications::DidCloseNotebookHandler::METHOD => {
-            local_notification_task::<notification::DidCloseNotebookHandler>(notif)
+            local_notification_task::<notifications::DidCloseNotebookHandler>(notif)
         }
-        notification::DidChangeWatchedFiles::METHOD => {
+        notifications::DidChangeWatchedFiles::METHOD => {
-            local_notification_task::<notification::DidChangeWatchedFiles>(notif)
+            local_notification_task::<notifications::DidChangeWatchedFiles>(notif)
         }
         lsp_types::notification::SetTrace::METHOD => {
             tracing::trace!("Ignoring `setTrace` notification");
@@ -103,23 +102,25 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
 
 fn _local_request_task<'a, R: traits::SyncRequestHandler>(
     req: server::Request,
-) -> super::Result<Task<'a>> {
+) -> super::Result<Task<'a>>
+where
+    <<R as RequestHandler>::RequestType as lsp_types::request::Request>::Params: UnwindSafe,
+{
     let (id, params) = cast_request::<R>(req)?;
     Ok(Task::local(|session, notifier, requester, responder| {
-        let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered();
+        let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
         let result = R::run(session, notifier, requester, params);
         respond::<R>(id, result, &responder);
     }))
 }
 
-// TODO(micha): Calls to `db` could panic if the db gets mutated while this task is running.
-// We should either wrap `R::run_with_snapshot` with a salsa catch cancellation handler or
-// use `SemanticModel` instead of passing `db` which uses a Result for all it's methods
-// that propagate cancellations.
 fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
     req: server::Request,
     schedule: BackgroundSchedule,
-) -> super::Result<Task<'a>> {
+) -> super::Result<Task<'a>>
+where
+    <<R as RequestHandler>::RequestType as lsp_types::request::Request>::Params: UnwindSafe,
+{
     let (id, params) = cast_request::<R>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
         let url = R::document_url(&params).into_owned();
@@ -128,6 +129,7 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
             tracing::warn!("Ignoring request for invalid `{url}`");
             return Box::new(|_, _| {});
         };
+
         let db = match &path {
             AnySystemPath::System(path) => match session.project_db_for_path(path.as_std_path()) {
                 Some(db) => db.clone(),
@@ -142,19 +144,55 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         };
 
         Box::new(move |notifier, responder| {
-            let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered();
+            let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
-            let result = R::run_with_snapshot(&db, snapshot, notifier, params);
+            let result = ruff_db::panic::catch_unwind(|| {
-            respond::<R>(id, result, &responder);
+                R::run_with_snapshot(&db, snapshot, notifier, params)
+            });
+
+            if let Some(response) = request_result_to_response(&id, &responder, result) {
+                respond::<R>(id, response, &responder);
+            }
         })
     }))
 }
 
+fn request_result_to_response<R>(
+    id: &RequestId,
+    responder: &Responder,
+    result: std::result::Result<Result<R>, PanicError>,
+) -> Option<Result<R>> {
+    match result {
+        Ok(response) => Some(response),
+        Err(error) => {
+            if error.payload.downcast_ref::<salsa::Cancelled>().is_some() {
+                // Request was cancelled by Salsa. TODO: Retry
+                respond_silent_error(
+                    id.clone(),
+                    responder,
+                    Error {
+                        code: lsp_server::ErrorCode::ContentModified,
+                        error: anyhow!("content modified"),
+                    },
+                );
+                None
+            } else {
+                let message = format!("request handler {error}");
+
+                Some(Err(Error {
+                    code: lsp_server::ErrorCode::InternalError,
+                    error: anyhow!(message),
+                }))
+            }
+        }
+    }
+}
+
 fn local_notification_task<'a, N: traits::SyncNotificationHandler>(
     notif: server::Notification,
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_notification::<N>(notif)?;
     Ok(Task::local(move |session, notifier, requester, _| {
-        let _span = tracing::trace_span!("notification", method = N::METHOD).entered();
+        let _span = tracing::debug_span!("notification", method = N::METHOD).entered();
         if let Err(err) = N::run(session, notifier, requester, params) {
             tracing::error!("An error occurred while running {id}: {err}");
             show_err_msg!("ty encountered a problem. Check the logs for more details.");
@@ -163,10 +201,15 @@ fn local_notification_task<'a, N: traits::SyncNotificationHandler>(
 }
 
 #[expect(dead_code)]
-fn background_notification_thread<'a, N: traits::BackgroundDocumentNotificationHandler>(
+fn background_notification_thread<'a, N>(
     req: server::Notification,
     schedule: BackgroundSchedule,
-) -> super::Result<Task<'a>> {
+) -> super::Result<Task<'a>>
+where
+    N: traits::BackgroundDocumentNotificationHandler,
+    <<N as NotificationHandler>::NotificationType as lsp_types::notification::Notification>::Params:
+        UnwindSafe,
+{
     let (id, params) = cast_notification::<N>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
         let url = N::document_url(&params);
@@ -177,8 +220,20 @@ fn background_notification_thread<'a, N: traits::BackgroundDocumentNotificationH
             return Box::new(|_, _| {});
         };
         Box::new(move |notifier, _| {
-            let _span = tracing::trace_span!("notification", method = N::METHOD).entered();
+            let _span = tracing::debug_span!("notification", method = N::METHOD).entered();
-            if let Err(err) = N::run_with_snapshot(snapshot, notifier, params) {
+
+            let result = match ruff_db::panic::catch_unwind(|| {
+                N::run_with_snapshot(snapshot, notifier, params)
+            }) {
+                Ok(result) => result,
+                Err(panic) => {
+                    tracing::error!("An error occurred while running {id}: {panic}");
+                    show_err_msg!("ty encountered a panic. Check the logs for more details.");
+                    return;
+                }
+            };
+
+            if let Err(err) = result {
                 tracing::error!("An error occurred while running {id}: {err}");
                 show_err_msg!("ty encountered a problem. Check the logs for more details.");
             }
@@ -198,6 +253,7 @@ fn cast_request<Req>(
 )>
 where
     Req: traits::RequestHandler,
+    <<Req as RequestHandler>::RequestType as lsp_types::request::Request>::Params: UnwindSafe,
 {
     request
         .extract(Req::METHOD)
@@ -232,6 +288,14 @@ fn respond<Req>(
     }
 }
 
+/// Sends back an error response to the server using a [`Responder`] without showing a warning
+/// to the user.
+fn respond_silent_error(id: server::RequestId, responder: &Responder, error: Error) {
+    if let Err(err) = responder.respond::<()>(id, Err(error)) {
+        tracing::error!("Failed to send response: {err}");
+    }
+}
+
 /// Tries to cast a serialized request from the server into
 /// a parameter type for a specific request handler.
 fn cast_notification<N>(
@@ -240,7 +304,9 @@ fn cast_notification<N>(
     (
         &'static str,
         <<N as traits::NotificationHandler>::NotificationType as lsp_types::notification::Notification>::Params,
-)> where N: traits::NotificationHandler{
+    )> where
+    N: traits::NotificationHandler,
+{
     Ok((
         N::METHOD,
         notification

crates/ty_server/src/server/api/diagnostics.rs

@@ -41,13 +41,7 @@ pub(super) fn compute_diagnostics(
         return vec![];
     };
 
-    let diagnostics = match db.check_file(file) {
+    let diagnostics = db.check_file(file);
-        Ok(diagnostics) => diagnostics,
-        Err(cancelled) => {
-            tracing::info!("Diagnostics computation {cancelled}");
-            return vec![];
-        }
-    };
 
     diagnostics
         .as_slice()

crates/ty_wasm/src/lib.rs

@@ -187,14 +187,14 @@ impl Workspace {
     /// Checks a single file.
     #[wasm_bindgen(js_name = "checkFile")]
     pub fn check_file(&self, file_id: &FileHandle) -> Result<Vec<Diagnostic>, Error> {
-        let result = self.db.check_file(file_id.file).map_err(into_error)?;
+        let result = self.db.check_file(file_id.file);
 
         Ok(result.into_iter().map(Diagnostic::wrap).collect())
     }
 
     /// Checks all open files
     pub fn check(&self) -> Result<Vec<Diagnostic>, Error> {
-        let result = self.db.check().map_err(into_error)?;
+        let result = self.db.check();
 
         Ok(result.into_iter().map(Diagnostic::wrap).collect())
     }