Support textDocument/diagnostic specification (Pull diagnostics) (#11315)

Co-authored-by: Michael Davis <mcarsondavis@gmail.com>

helix-core/src/syntax/config.rs

@@ -270,6 +270,7 @@ pub enum LanguageServerFeature {
     WorkspaceSymbols,
     // Symbols, use bitflags, see above?
     Diagnostics,
+    PullDiagnostics,
     RenameSymbol,
     InlayHints,
     DocumentColors,
@@ -294,6 +295,7 @@ impl Display for LanguageServerFeature {
             DocumentSymbols => "document-symbols",
             WorkspaceSymbols => "workspace-symbols",
             Diagnostics => "diagnostics",
+            PullDiagnostics => "pull-diagnostics",
             RenameSymbol => "rename-symbol",
             InlayHints => "inlay-hints",
             DocumentColors => "document-colors",

helix-lsp/src/client.rs

@@ -372,6 +372,7 @@ impl Client {
                 Some(OneOf::Left(true) | OneOf::Right(_))
             ),
             LanguageServerFeature::Diagnostics => true, // there's no extra server capability
+            LanguageServerFeature::PullDiagnostics => capabilities.diagnostic_provider.is_some(),
             LanguageServerFeature::RenameSymbol => matches!(
                 capabilities.rename_provider,
                 Some(OneOf::Left(true)) | Some(OneOf::Right(_))
@@ -602,6 +603,9 @@ impl Client {
                         did_rename: Some(true),
                         ..Default::default()
                     }),
+                    diagnostic: Some(lsp::DiagnosticWorkspaceClientCapabilities {
+                        refresh_support: Some(true),
+                    }),
                     ..Default::default()
                 }),
                 text_document: Some(lsp::TextDocumentClientCapabilities {
@@ -679,6 +683,10 @@ impl Client {
                         }),
                         ..Default::default()
                     }),
+                    diagnostic: Some(lsp::DiagnosticClientCapabilities {
+                        dynamic_registration: Some(false),
+                        related_document_support: Some(true),
+                    }),
                     publish_diagnostics: Some(lsp::PublishDiagnosticsClientCapabilities {
                         version_support: Some(true),
                         tag_support: Some(lsp::TagSupport {
@@ -1229,6 +1237,32 @@ impl Client {
         Some(self.call::<lsp::request::RangeFormatting>(params))
     }
 
+    pub fn text_document_diagnostic(
+        &self,
+        text_document: lsp::TextDocumentIdentifier,
+        previous_result_id: Option<String>,
+    ) -> Option<impl Future<Output = Result<lsp::DocumentDiagnosticReportResult>>> {
+        let capabilities = self.capabilities();
+
+        // 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(),
+            lsp::DiagnosticServerCapabilities::RegistrationOptions(cap) => {
+                cap.diagnostic_options.identifier.clone()
+            }
+        };
+
+        let params = lsp::DocumentDiagnosticParams {
+            text_document,
+            identifier,
+            previous_result_id,
+            work_done_progress_params: lsp::WorkDoneProgressParams::default(),
+            partial_result_params: lsp::PartialResultParams::default(),
+        };
+
+        Some(self.call::<lsp::request::DocumentDiagnosticRequest>(params))
+    }
+
     pub fn text_document_document_highlight(
         &self,
         text_document: lsp::TextDocumentIdentifier,

helix-lsp/src/lib.rs

@@ -463,6 +463,7 @@ pub enum MethodCall {
     RegisterCapability(lsp::RegistrationParams),
     UnregisterCapability(lsp::UnregistrationParams),
     ShowDocument(lsp::ShowDocumentParams),
+    WorkspaceDiagnosticRefresh,
 }
 
 impl MethodCall {
@@ -494,6 +495,7 @@ impl MethodCall {
                 let params: lsp::ShowDocumentParams = params.parse()?;
                 Self::ShowDocument(params)
             }
+            lsp::request::WorkspaceDiagnosticRefresh::METHOD => Self::WorkspaceDiagnosticRefresh,
             _ => {
                 return Err(Error::Unhandled);
             }

helix-term/src/application.rs

@@ -1103,6 +1103,26 @@ impl Application {
                         let result = self.handle_show_document(params, offset_encoding);
                         Ok(json!(result))
                     }
+                    Ok(MethodCall::WorkspaceDiagnosticRefresh) => {
+                        let language_server = language_server!().id();
+
+                        let documents: Vec<_> = self
+                            .editor
+                            .documents
+                            .values()
+                            .filter(|x| x.supports_language_server(language_server))
+                            .map(|x| x.id())
+                            .collect();
+
+                        for document in documents {
+                            handlers::diagnostics::request_document_diagnostics(
+                                &mut self.editor,
+                                document,
+                            );
+                        }
+
+                        Ok(serde_json::Value::Null)
+                    }
                 };
 
                 let language_server = language_server!();

helix-term/src/handlers.rs

@@ -1,11 +1,13 @@
 use std::sync::Arc;
 
 use arc_swap::ArcSwap;
+use diagnostics::PullAllDocumentsDiagnosticHandler;
 use helix_event::AsyncHook;
 
 use crate::config::Config;
 use crate::events;
 use crate::handlers::auto_save::AutoSaveHandler;
+use crate::handlers::diagnostics::PullDiagnosticsHandler;
 use crate::handlers::signature_help::SignatureHelpHandler;
 
 pub use helix_view::handlers::{word_index, Handlers};
@@ -14,7 +16,7 @@ use self::document_colors::DocumentColorsHandler;
 
 mod auto_save;
 pub mod completion;
-mod diagnostics;
+pub mod diagnostics;
 mod document_colors;
 mod prompt;
 mod signature_help;
@@ -28,6 +30,8 @@ pub fn setup(config: Arc<ArcSwap<Config>>) -> Handlers {
     let auto_save = AutoSaveHandler::new().spawn();
     let document_colors = DocumentColorsHandler::default().spawn();
     let word_index = word_index::Handler::spawn();
+    let pull_diagnostics = PullDiagnosticsHandler::default().spawn();
+    let pull_all_documents_diagnostics = PullAllDocumentsDiagnosticHandler::default().spawn();
 
     let handlers = Handlers {
         completions: helix_view::handlers::completion::CompletionHandler::new(event_tx),
@@ -35,6 +39,8 @@ pub fn setup(config: Arc<ArcSwap<Config>>) -> Handlers {
         auto_save,
         document_colors,
         word_index,
+        pull_diagnostics,
+        pull_all_documents_diagnostics,
     };
 
     helix_view::handlers::register_hooks(&handlers);

helix-term/src/handlers/diagnostics.rs

@@ -1,12 +1,28 @@
-use helix_event::{register_hook, send_blocking};
+use futures_util::stream::FuturesUnordered;
+use std::collections::HashSet;
+use std::mem;
+use std::time::Duration;
+use tokio::time::Instant;
+use tokio_stream::StreamExt;
+
+use helix_core::diagnostic::DiagnosticProvider;
+use helix_core::syntax::config::LanguageServerFeature;
+use helix_core::Uri;
+use helix_event::{cancelable_future, register_hook, send_blocking};
+use helix_lsp::{lsp, LanguageServerId};
 use helix_view::document::Mode;
-use helix_view::events::DiagnosticsDidChange;
+use helix_view::events::{
+    DiagnosticsDidChange, DocumentDidChange, DocumentDidOpen, LanguageServerInitialized,
+};
 use helix_view::handlers::diagnostics::DiagnosticEvent;
+use helix_view::handlers::lsp::{PullAllDocumentsDiagnosticsEvent, PullDiagnosticsEvent};
 use helix_view::handlers::Handlers;
+use helix_view::{DocumentId, Editor};
 
 use crate::events::OnModeSwitch;
+use crate::job;
 
-pub(super) fn register_hooks(_handlers: &Handlers) {
+pub(super) fn register_hooks(handlers: &Handlers) {
     register_hook!(move |event: &mut DiagnosticsDidChange<'_>| {
         if event.editor.mode != Mode::Insert {
             for (view, _) in event.editor.tree.views_mut() {
@@ -21,4 +37,265 @@ pub(super) fn register_hooks(_handlers: &Handlers) {
         }
         Ok(())
     });
+
+    let tx = handlers.pull_diagnostics.clone();
+    let tx_all_documents = handlers.pull_all_documents_diagnostics.clone();
+    register_hook!(move |event: &mut DocumentDidChange<'_>| {
+        if event
+            .doc
+            .has_language_server_with_feature(LanguageServerFeature::PullDiagnostics)
+            && !event.ghost_transaction
+        {
+            // Cancel the ongoing request, if present.
+            event.doc.pull_diagnostic_controller.cancel();
+            let document_id = event.doc.id();
+            send_blocking(&tx, PullDiagnosticsEvent { document_id });
+
+            let inter_file_dependencies_language_servers = event
+                .doc
+                .language_servers_with_feature(LanguageServerFeature::PullDiagnostics)
+                .filter(|language_server| {
+                    language_server
+                        .capabilities()
+                        .diagnostic_provider
+                        .as_ref()
+                        .is_some_and(|diagnostic_provider| match diagnostic_provider {
+                            lsp::DiagnosticServerCapabilities::Options(options) => {
+                                options.inter_file_dependencies
+                            }
+
+                            lsp::DiagnosticServerCapabilities::RegistrationOptions(options) => {
+                                options.diagnostic_options.inter_file_dependencies
+                            }
+                        })
+                })
+                .map(|language_server| language_server.id())
+                .collect();
+
+            send_blocking(
+                &tx_all_documents,
+                PullAllDocumentsDiagnosticsEvent {
+                    language_servers: inter_file_dependencies_language_servers,
+                },
+            );
+        }
+        Ok(())
+    });
+
+    register_hook!(move |event: &mut DocumentDidOpen<'_>| {
+        request_document_diagnostics(event.editor, event.doc);
+
+        Ok(())
+    });
+
+    register_hook!(move |event: &mut LanguageServerInitialized<'_>| {
+        let doc_ids: Vec<_> = event.editor.documents.keys().copied().collect();
+
+        for doc_id in doc_ids {
+            request_document_diagnostics(event.editor, doc_id);
+        }
+
+        Ok(())
+    });
+}
+
+#[derive(Debug, Default)]
+pub(super) struct PullDiagnosticsHandler {
+    document_ids: HashSet<DocumentId>,
+}
+
+impl helix_event::AsyncHook for PullDiagnosticsHandler {
+    type Event = PullDiagnosticsEvent;
+
+    fn handle_event(
+        &mut self,
+        event: Self::Event,
+        _timeout: Option<tokio::time::Instant>,
+    ) -> Option<tokio::time::Instant> {
+        self.document_ids.insert(event.document_id);
+        Some(Instant::now() + Duration::from_millis(250))
+    }
+
+    fn finish_debounce(&mut self) {
+        let document_ids = mem::take(&mut self.document_ids);
+        job::dispatch_blocking(move |editor, _| {
+            for document_id in document_ids {
+                request_document_diagnostics(editor, document_id);
+            }
+        })
+    }
+}
+
+#[derive(Debug, Default)]
+pub(super) struct PullAllDocumentsDiagnosticHandler {
+    language_servers: HashSet<LanguageServerId>,
+}
+
+impl helix_event::AsyncHook for PullAllDocumentsDiagnosticHandler {
+    type Event = PullAllDocumentsDiagnosticsEvent;
+
+    fn handle_event(
+        &mut self,
+        event: Self::Event,
+        _timeout: Option<tokio::time::Instant>,
+    ) -> Option<tokio::time::Instant> {
+        self.language_servers.extend(&event.language_servers);
+        Some(Instant::now() + Duration::from_secs(1))
+    }
+
+    fn finish_debounce(&mut self) {
+        let language_servers = mem::take(&mut self.language_servers);
+        job::dispatch_blocking(move |editor, _| {
+            let documents: Vec<_> = editor.documents.keys().copied().collect();
+
+            for document in documents {
+                request_document_diagnostics_for_language_severs(
+                    editor,
+                    document,
+                    language_servers.clone(),
+                );
+            }
+        })
+    }
+}
+
+fn request_document_diagnostics_for_language_severs(
+    editor: &mut Editor,
+    doc_id: DocumentId,
+    language_servers: HashSet<LanguageServerId>,
+) {
+    let Some(doc) = editor.document_mut(doc_id) else {
+        return;
+    };
+
+    let cancel = doc.pull_diagnostic_controller.restart();
+
+    let mut futures: FuturesUnordered<_> = language_servers
+        .iter()
+        .filter_map(|x| doc.language_servers().find(|y| &y.id() == x))
+        .filter_map(|language_server| {
+            let future = language_server
+                .text_document_diagnostic(doc.identifier(), doc.previous_diagnostic_id.clone())?;
+
+            let identifier = language_server
+                .capabilities()
+                .diagnostic_provider
+                .as_ref()
+                .and_then(|diagnostic_provider| match diagnostic_provider {
+                    lsp::DiagnosticServerCapabilities::Options(options) => {
+                        options.identifier.clone()
+                    }
+                    lsp::DiagnosticServerCapabilities::RegistrationOptions(options) => {
+                        options.diagnostic_options.identifier.clone()
+                    }
+                });
+
+            let language_server_id = language_server.id();
+            let provider = DiagnosticProvider::Lsp {
+                server_id: language_server_id,
+                identifier,
+            };
+            let uri = doc.uri()?;
+
+            Some(async move {
+                let result = future.await;
+
+                (result, provider, uri)
+            })
+        })
+        .collect();
+
+    if futures.is_empty() {
+        return;
+    }
+
+    tokio::spawn(async move {
+        let mut retry_language_servers = HashSet::new();
+        loop {
+            match cancelable_future(futures.next(), &cancel).await {
+                Some(Some((Ok(result), provider, uri))) => {
+                    job::dispatch(move |editor, _| {
+                        handle_pull_diagnostics_response(editor, result, provider, uri, doc_id);
+                    })
+                    .await;
+                }
+                Some(Some((Err(err), DiagnosticProvider::Lsp { server_id, .. }, _))) => {
+                    let parsed_cancellation_data = if let helix_lsp::Error::Rpc(error) = err {
+                        error.data.and_then(|data| {
+                            serde_json::from_value::<lsp::DiagnosticServerCancellationData>(data)
+                                .ok()
+                        })
+                    } else {
+                        log::error!("Pull diagnostic request failed: {err}");
+                        continue;
+                    };
+                    if parsed_cancellation_data.is_some_and(|data| data.retrigger_request) {
+                        retry_language_servers.insert(server_id);
+                    }
+                }
+                Some(None) => break,
+                // The request was cancelled.
+                None => return,
+            }
+        }
+
+        if !retry_language_servers.is_empty() {
+            tokio::time::sleep(Duration::from_millis(500)).await;
+
+            job::dispatch(move |editor, _| {
+                request_document_diagnostics_for_language_severs(
+                    editor,
+                    doc_id,
+                    retry_language_servers,
+                );
+            })
+            .await;
+        }
+    });
+}
+
+pub fn request_document_diagnostics(editor: &mut Editor, doc_id: DocumentId) {
+    let Some(doc) = editor.document(doc_id) else {
+        return;
+    };
+
+    let language_servers = doc
+        .language_servers_with_feature(LanguageServerFeature::PullDiagnostics)
+        .map(|language_servers| language_servers.id())
+        .collect();
+
+    request_document_diagnostics_for_language_severs(editor, doc_id, language_servers);
+}
+
+fn handle_pull_diagnostics_response(
+    editor: &mut Editor,
+    result: lsp::DocumentDiagnosticReportResult,
+    provider: DiagnosticProvider,
+    uri: Uri,
+    document_id: DocumentId,
+) {
+    match result {
+        lsp::DocumentDiagnosticReportResult::Report(report) => {
+            let result_id = match report {
+                lsp::DocumentDiagnosticReport::Full(report) => {
+                    editor.handle_lsp_diagnostics(
+                        &provider,
+                        uri,
+                        None,
+                        report.full_document_diagnostic_report.items,
+                    );
+
+                    report.full_document_diagnostic_report.result_id
+                }
+                lsp::DocumentDiagnosticReport::Unchanged(report) => {
+                    Some(report.unchanged_document_diagnostic_report.result_id)
+                }
+            };
+
+            if let Some(doc) = editor.document_mut(document_id) {
+                doc.previous_diagnostic_id = result_id;
+            };
+        }
+        lsp::DocumentDiagnosticReportResult::Partial(_) => {}
+    };
 }

helix-view/src/document.rs

@@ -204,11 +204,14 @@ pub struct Document {
 
     pub readonly: bool,
 
+    pub previous_diagnostic_id: Option<String>,
+
     /// Annotations for LSP document color swatches
     pub color_swatches: Option<DocumentColorSwatches>,
     // NOTE: ideally this would live on the handler for color swatches. This is blocked on a
     // large refactor that would make `&mut Editor` available on the `DocumentDidChange` event.
     pub color_swatch_controller: TaskController,
+    pub pull_diagnostic_controller: TaskController,
 
     // NOTE: this field should eventually go away - we should use the Editor's syn_loader instead
     // of storing a copy on every doc. Then we can remove the surrounding `Arc` and use the
@@ -728,6 +731,8 @@ impl Document {
             color_swatches: None,
             color_swatch_controller: TaskController::new(),
             syn_loader,
+            previous_diagnostic_id: None,
+            pull_diagnostic_controller: TaskController::new(),
         }
     }
 
@@ -2284,6 +2289,10 @@ impl Document {
     pub fn reset_all_inlay_hints(&mut self) {
         self.inlay_hints = Default::default();
     }
+
+    pub fn has_language_server_with_feature(&self, feature: LanguageServerFeature) -> bool {
+        self.language_servers_with_feature(feature).next().is_some()
+    }
 }
 
 #[derive(Debug, Default)]

helix-view/src/handlers.rs

@@ -24,6 +24,8 @@ pub struct Handlers {
     pub auto_save: Sender<AutoSaveEvent>,
     pub document_colors: Sender<lsp::DocumentColorsEvent>,
     pub word_index: word_index::Handler,
+    pub pull_diagnostics: Sender<lsp::PullDiagnosticsEvent>,
+    pub pull_all_documents_diagnostics: Sender<lsp::PullAllDocumentsDiagnosticsEvent>,
 }
 
 impl Handlers {

helix-view/src/handlers/lsp.rs

@@ -1,4 +1,5 @@
 use std::collections::btree_map::Entry;
+use std::collections::HashSet;
 use std::fmt::Display;
 
 use crate::editor::Action;
@@ -30,6 +31,14 @@ pub enum SignatureHelpEvent {
     RequestComplete { open: bool },
 }
 
+pub struct PullDiagnosticsEvent {
+    pub document_id: DocumentId,
+}
+
+pub struct PullAllDocumentsDiagnosticsEvent {
+    pub language_servers: HashSet<LanguageServerId>,
+}
+
 #[derive(Debug)]
 pub struct ApplyEditError {
     pub kind: ApplyEditErrorKind,