[ty] Initial support for workspace diagnostics

Author: dhruvmanila

crates/ty_project/src/db.rs

@@ -81,15 +81,20 @@ impl ProjectDatabase {
 
     /// Checks all open files in the project and its dependencies.
     pub fn check(&self) -> Vec<Diagnostic> {
-        let mut reporter = DummyReporter;
-        let reporter = AssertUnwindSafe(&mut reporter as &mut dyn Reporter);
-        self.project().check(self, reporter)
+        self.check_with_mode(CheckMode::OpenFiles)
     }
 
     /// Checks all open files in the project and its dependencies, using the given reporter.
     pub fn check_with_reporter(&self, reporter: &mut dyn Reporter) -> Vec<Diagnostic> {
         let reporter = AssertUnwindSafe(reporter);
-        self.project().check(self, reporter)
+        self.project().check(self, CheckMode::OpenFiles, reporter)
+    }
+
+    /// Check the project with the given mode.
+    pub fn check_with_mode(&self, mode: CheckMode) -> Vec<Diagnostic> {
+        let mut reporter = DummyReporter;
+        let reporter = AssertUnwindSafe(&mut reporter as &mut dyn Reporter);
+        self.project().check(self, mode, reporter)
     }
 
     #[tracing::instrument(level = "debug", skip(self))]
@@ -110,6 +115,17 @@ impl ProjectDatabase {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum CheckMode {
+    /// Checks only the open files in the project.
+    OpenFiles,
+
+    /// Checks all files in the project, ignoring the open file set.
+    ///
+    /// This includes virtual files, such as those created by the language server.
+    AllFiles,
+}
+
 impl Upcast<dyn SemanticDb> for ProjectDatabase {
     fn upcast(&self) -> &(dyn SemanticDb + 'static) {
         self

crates/ty_project/src/lib.rs

@@ -1,7 +1,7 @@
 use crate::glob::{GlobFilterCheckMode, IncludeResult};
 use crate::metadata::options::{OptionDiagnostic, ToSettingsError};
 use crate::walk::{ProjectFilesFilter, ProjectFilesWalker};
-pub use db::{Db, ProjectDatabase};
+pub use db::{CheckMode, Db, ProjectDatabase};
 use files::{Index, Indexed, IndexedFiles};
 use metadata::settings::Settings;
 pub use metadata::{ProjectMetadata, ProjectMetadataError};
@@ -213,6 +213,7 @@ impl Project {
     pub(crate) fn check(
         self,
         db: &ProjectDatabase,
+        mode: CheckMode,
         mut reporter: AssertUnwindSafe<&mut dyn Reporter>,
     ) -> Vec<Diagnostic> {
         let project_span = tracing::debug_span!("Project::check");
@@ -227,7 +228,11 @@ impl Project {
                 .map(OptionDiagnostic::to_diagnostic),
         );
 
-        let files = ProjectFiles::new(db, self);
+        let files = match mode {
+            CheckMode::OpenFiles => ProjectFiles::new(db, self),
+            // TODO: Consider open virtual files as well
+            CheckMode::AllFiles => ProjectFiles::Indexed(self.files(db)),
+        };
         reporter.set_files(files.len());
 
         diagnostics.extend(

crates/ty_server/src/server.rs

@@ -173,6 +173,7 @@ impl Server {
             diagnostic_provider: Some(DiagnosticServerCapabilities::Options(DiagnosticOptions {
                 identifier: Some(crate::DIAGNOSTIC_NAME.into()),
                 inter_file_dependencies: true,
+                workspace_diagnostics: true,
                 ..Default::default()
             })),
             text_document_sync: Some(TextDocumentSyncCapability::Options(

crates/ty_server/src/server/api.rs

@@ -28,23 +28,28 @@ pub(super) fn request(req: server::Request) -> Task {
     let id = req.id.clone();
 
     match req.method.as_str() {
-        requests::DocumentDiagnosticRequestHandler::METHOD => background_request_task::<
+        requests::DocumentDiagnosticRequestHandler::METHOD => background_document_request_task::<
             requests::DocumentDiagnosticRequestHandler,
         >(
             req, BackgroundSchedule::Worker
         ),
-        requests::GotoTypeDefinitionRequestHandler::METHOD => background_request_task::<
+        requests::WorkspaceDiagnosticRequestHandler::METHOD => background_request_task::<
+            requests::WorkspaceDiagnosticRequestHandler,
+        >(
+            req, BackgroundSchedule::Worker
+        ),
+        requests::GotoTypeDefinitionRequestHandler::METHOD => background_document_request_task::<
             requests::GotoTypeDefinitionRequestHandler,
         >(
             req, BackgroundSchedule::Worker
         ),
-        requests::HoverRequestHandler::METHOD => background_request_task::<
+        requests::HoverRequestHandler::METHOD => background_document_request_task::<
             requests::HoverRequestHandler,
         >(req, BackgroundSchedule::Worker),
-        requests::InlayHintRequestHandler::METHOD => background_request_task::<
+        requests::InlayHintRequestHandler::METHOD => background_document_request_task::<
             requests::InlayHintRequestHandler,
         >(req, BackgroundSchedule::Worker),
-        requests::CompletionRequestHandler::METHOD => background_request_task::<
+        requests::CompletionRequestHandler::METHOD => background_document_request_task::<
             requests::CompletionRequestHandler,
         >(
             req, BackgroundSchedule::LatencySensitive
@@ -135,7 +140,51 @@ where
     }))
 }
 
-fn background_request_task<R: traits::BackgroundDocumentRequestHandler>(
+fn background_request_task<R: traits::BackgroundRequestHandler>(
+    req: server::Request,
+    schedule: BackgroundSchedule,
+) -> Result<Task>
+where
+    <<R as RequestHandler>::RequestType as Request>::Params: UnwindSafe,
+{
+    let retry = R::RETRY_ON_CANCELLATION.then(|| req.clone());
+    let (id, params) = cast_request::<R>(req)?;
+
+    Ok(Task::background(schedule, move |session: &Session| {
+        let cancellation_token = session
+            .request_queue()
+            .incoming()
+            .cancellation_token(&id)
+            .expect("request should have been tested for cancellation before scheduling");
+
+        let snapshot = session.take_workspace_snapshot();
+
+        Box::new(move |client| {
+            let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
+
+            // Test again if the request was cancelled since it was scheduled on the background task
+            // and, if so, return early
+            if cancellation_token.is_cancelled() {
+                tracing::trace!(
+                    "Ignoring request id={id} method={} because it was cancelled",
+                    R::METHOD
+                );
+
+                // We don't need to send a response here because the `cancel` notification
+                // handler already responded with a message.
+                return;
+            }
+
+            let result = ruff_db::panic::catch_unwind(|| R::run(snapshot, client, params));
+
+            if let Some(response) = request_result_to_response::<R>(&id, client, result, retry) {
+                respond::<R>(&id, response, client);
+            }
+        })
+    }))
+}
+
+fn background_document_request_task<R: traits::BackgroundDocumentRequestHandler>(
     req: server::Request,
     schedule: BackgroundSchedule,
 ) -> Result<Task>
@@ -168,7 +217,7 @@ where
         };
 
         let Some(snapshot) = session.take_snapshot(url) else {
-            tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist.");
+            tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist");
             return Box::new(|_| {});
         };
 
@@ -209,7 +258,7 @@ fn request_result_to_response<R>(
     request: Option<lsp_server::Request>,
 ) -> Option<Result<<<R as RequestHandler>::RequestType as Request>::Result>>
 where
-    R: traits::BackgroundDocumentRequestHandler,
+    R: traits::RetriableRequestHandler,
 {
     match result {
         Ok(response) => Some(response),

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

@@ -172,7 +172,7 @@ pub(super) fn compute_diagnostics(
 
 /// Converts the tool specific [`Diagnostic`][ruff_db::diagnostic::Diagnostic] to an LSP
 /// [`Diagnostic`].
-fn to_lsp_diagnostic(
+pub(super) fn to_lsp_diagnostic(
     db: &dyn Db,
     diagnostic: &ruff_db::diagnostic::Diagnostic,
     encoding: PositionEncoding,

crates/ty_server/src/server/api/notifications/did_close.rs

@@ -41,6 +41,13 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             );
         }
 
-        clear_diagnostics(&key, client)
+        if !session.global_settings().diagnostic_mode().is_workspace() {
+            // The server needs to clear the diagnostics regardless of whether the client supports
+            // pull diagnostics or not. This is because the client only has the capability to fetch
+            // the diagnostics but does not automatically clear them when a document is closed.
+            clear_diagnostics(&key, client)?;
+        }
+
+        Ok(())
     }
 }

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

@@ -4,10 +4,12 @@ mod goto_type_definition;
 mod hover;
 mod inlay_hints;
 mod shutdown;
+mod workspace_diagnostic;
 
 pub(super) use completion::CompletionRequestHandler;
 pub(super) use diagnostic::DocumentDiagnosticRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
 pub(super) use inlay_hints::InlayHintRequestHandler;
 pub(super) use shutdown::ShutdownHandler;
+pub(super) use workspace_diagnostic::WorkspaceDiagnosticRequestHandler;

crates/ty_server/src/server/api/requests/completion.rs

@@ -8,7 +8,9 @@ use ty_project::ProjectDatabase;
 
 use crate::DocumentSnapshot;
 use crate::document::PositionExt;
-use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
 use crate::session::client::Client;
 
 pub(crate) struct CompletionRequestHandler;
@@ -18,8 +20,6 @@ impl RequestHandler for CompletionRequestHandler {
 }
 
 impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
-    const RETRY_ON_CANCELLATION: bool = true;
-
     fn document_url(params: &CompletionParams) -> Cow<Url> {
         Cow::Borrowed(&params.text_document_position.text_document.uri)
     }
@@ -65,3 +65,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
         Ok(Some(response))
     }
 }
+
+impl RetriableRequestHandler for CompletionRequestHandler {
+    const RETRY_ON_CANCELLATION: bool = true;
+}

crates/ty_server/src/server/api/requests/diagnostic.rs

@@ -8,7 +8,9 @@ use lsp_types::{
 
 use crate::server::Result;
 use crate::server::api::diagnostics::{Diagnostics, compute_diagnostics};
-use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 use ty_project::ProjectDatabase;
@@ -43,7 +45,9 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
             }),
         ))
     }
+}
 
+impl RetriableRequestHandler for DocumentDiagnosticRequestHandler {
     fn salsa_cancellation_error() -> lsp_server::ResponseError {
         lsp_server::ResponseError {
             code: lsp_server::ErrorCode::ServerCancelled as i32,

crates/ty_server/src/server/api/requests/goto_type_definition.rs

@@ -8,7 +8,9 @@ use ty_project::ProjectDatabase;
 
 use crate::DocumentSnapshot;
 use crate::document::{PositionExt, ToLink};
-use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
 use crate::session::client::Client;
 
 pub(crate) struct GotoTypeDefinitionRequestHandler;
@@ -70,3 +72,5 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
         }
     }
 }
+
+impl RetriableRequestHandler for GotoTypeDefinitionRequestHandler {}