Cancel ongoing pull diagnostic events

Author: SofusA

helix-lsp/src/client.rs

@@ -62,6 +62,7 @@ pub struct Client {
     initialize_notify: Arc<Notify>,
     /// workspace folders added while the server is still initializing
     req_timeout: u64,
+    ongoing_work: Mutex<Vec<(lsp::TextDocumentIdentifier, lsp::ProgressToken)>>,
 }
 
 impl Client {
@@ -234,6 +235,7 @@ impl Client {
             root_uri,
             workspace_folders: Mutex::new(workspace_folders),
             initialize_notify: initialize_notify.clone(),
+            ongoing_work: Mutex::new(Default::default()),
         };
 
         Ok((client, server_rx, initialize_notify))
@@ -1227,13 +1229,39 @@ impl Client {
         Some(self.call::<lsp::request::RangeFormatting>(params))
     }
 
+    pub fn mark_work_as_done(&self, token: lsp::ProgressToken) {
+        self.ongoing_work.lock().retain(|x| x.1 != token);
+    }
+
     pub fn text_document_diagnostic(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         previous_result_id: Option<String>,
-    ) -> Option<impl Future<Output = Result<lsp::DocumentDiagnosticReportResult>>> {
+    ) -> Option<(
+        impl Future<Output = Result<lsp::DocumentDiagnosticReportResult>>,
+        lsp::ProgressToken,
+    )> {
         let capabilities = self.capabilities();
 
+        let ongoing_work = {
+            let ongoing_work_lock = self.ongoing_work.lock();
+
+            ongoing_work_lock
+                .iter()
+                .filter(|x| x.0 == text_document)
+                .cloned()
+                .collect::<Vec<_>>()
+        };
+
+        if !ongoing_work.is_empty() {
+            for id in ongoing_work.iter().map(|x| x.1.clone()) {
+                self.notify::<lsp::notification::Cancel>(lsp::CancelParams { id: id.clone() });
+                self.mark_work_as_done(id);
+            }
+
+            return None;
+        }
+
         // Return early if the server does not support pull diagnostic.
         let identifier = match capabilities.diagnostic_provider.as_ref()? {
             lsp::DiagnosticServerCapabilities::Options(cap) => cap.identifier.clone(),
@@ -1242,15 +1270,30 @@ impl Client {
             }
         };
 
+        let work_done_token = match self.next_request_id() {
+            jsonrpc::Id::Null => lsp::ProgressToken::Number(1),
+            jsonrpc::Id::Num(num) => lsp::ProgressToken::Number(num as i32),
+            jsonrpc::Id::Str(str) => lsp::ProgressToken::String(str),
+        };
+
         let params = lsp::DocumentDiagnosticParams {
-            text_document,
+            text_document: text_document.clone(),
             identifier,
             previous_result_id,
-            work_done_progress_params: lsp::WorkDoneProgressParams::default(),
+            work_done_progress_params: lsp::WorkDoneProgressParams {
+                work_done_token: Some(work_done_token.clone()),
+            },
             partial_result_params: lsp::PartialResultParams::default(),
         };
 
-        Some(self.call::<lsp::request::DocumentDiagnosticRequest>(params))
+        self.ongoing_work
+            .lock()
+            .push((text_document, work_done_token.clone()));
+
+        Some((
+            self.call::<lsp::request::DocumentDiagnosticRequest>(params),
+            work_done_token,
+        ))
     }
 
     pub fn text_document_document_highlight(

helix-term/src/handlers/diagnostics.rs

@@ -1,4 +1,5 @@
 use std::time::Duration;
+use tokio::time::Instant;
 
 use helix_core::diagnostic::DiagnosticProvider;
 use helix_core::syntax::LanguageServerFeature;
@@ -13,7 +14,6 @@ use helix_view::handlers::diagnostics::DiagnosticEvent;
 use helix_view::handlers::lsp::PullDiagnosticsEvent;
 use helix_view::handlers::Handlers;
 use helix_view::{DocumentId, Editor};
-use tokio::time::Instant;
 
 use crate::events::OnModeSwitch;
 use crate::job;
@@ -74,87 +74,30 @@ pub(super) fn register_hooks(handlers: &Handlers) {
 }
 
 #[derive(Debug)]
-pub(super) struct PullDiagnosticsHandler {
-    no_inter_file_dependency_timeout: Option<tokio::time::Instant>,
-}
+pub(super) struct PullDiagnosticsHandler {}
 
 impl PullDiagnosticsHandler {
-    pub fn new() -> PullDiagnosticsHandler {
-        PullDiagnosticsHandler {
-            no_inter_file_dependency_timeout: None,
-        }
+    pub fn new() -> Self {
+        PullDiagnosticsHandler {}
     }
 }
 
-const TIMEOUT: Duration = Duration::from_millis(500);
-const TIMEOUT_NO_INTER_FILE_DEPENDENCY: Duration = Duration::from_millis(125);
-
 impl helix_event::AsyncHook for PullDiagnosticsHandler {
     type Event = PullDiagnosticsEvent;
 
     fn handle_event(
         &mut self,
-        event: Self::Event,
-        timeout: Option<tokio::time::Instant>,
+        _event: Self::Event,
+        _timeout: Option<tokio::time::Instant>,
     ) -> Option<tokio::time::Instant> {
-        if timeout.is_none() {
-            dispatch_pull_diagnostic_for_document(event.document_id, false);
-            self.no_inter_file_dependency_timeout = Some(Instant::now());
-        }
-
-        if self
-            .no_inter_file_dependency_timeout
-            .is_some_and(|nifd_timeout| {
-                nifd_timeout.duration_since(Instant::now()) > TIMEOUT_NO_INTER_FILE_DEPENDENCY
-            })
-        {
-            dispatch_pull_diagnostic_for_document(event.document_id, true);
-            self.no_inter_file_dependency_timeout = Some(Instant::now());
-        };
-
-        Some(Instant::now() + TIMEOUT)
+        Some(Instant::now() + Duration::from_millis(150))
     }
 
     fn finish_debounce(&mut self) {
         dispatch_pull_diagnostic_for_open_documents();
     }
 }
 
-fn dispatch_pull_diagnostic_for_document(
-    document_id: DocumentId,
-    exclude_language_servers_without_inter_file_dependency: bool,
-) {
-    job::dispatch_blocking(move |editor, _| {
-        let Some(doc) = editor.document(document_id) else {
-            return;
-        };
-
-        let language_servers = doc
-            .language_servers_with_feature(LanguageServerFeature::PullDiagnostics)
-            .filter(|ls| ls.is_initialized())
-            .filter(|ls| {
-                if !exclude_language_servers_without_inter_file_dependency {
-                    return true;
-                };
-                ls.capabilities()
-                    .diagnostic_provider
-                    .as_ref()
-                    .is_some_and(|dp| match dp {
-                        lsp::DiagnosticServerCapabilities::Options(options) => {
-                            options.inter_file_dependencies
-                        }
-                        lsp::DiagnosticServerCapabilities::RegistrationOptions(options) => {
-                            options.diagnostic_options.inter_file_dependencies
-                        }
-                    })
-            });
-
-        for language_server in language_servers {
-            pull_diagnostics_for_document(doc, language_server);
-        }
-    })
-}
-
 fn dispatch_pull_diagnostic_for_open_documents() {
     job::dispatch_blocking(move |editor, _| {
         let documents = editor.documents.values();
@@ -204,9 +147,14 @@ pub fn pull_diagnostics_for_document(
     let document_id = doc.id();
 
     tokio::spawn(async move {
-        match future.await {
+        match future.0.await {
             Ok(result) => {
                 job::dispatch(move |editor, _| {
+                    if let Some(language_server) = editor.language_server_by_id(language_server_id)
+                    {
+                        language_server.mark_work_as_done(future.1);
+                    };
+
                     handle_pull_diagnostics_response(editor, result, provider, uri, document_id)
                 })
                 .await
@@ -230,6 +178,7 @@ pub fn pull_diagnostics_for_document(
                                 editor.document(document_id),
                                 editor.language_server_by_id(language_server_id),
                             ) {
+                                language_server.mark_work_as_done(future.1);
                                 pull_diagnostics_for_document(doc, language_server);
                             }
                         })