From 26f6751c9b99926f304f5f8451e83fb9abd0c98f Mon Sep 17 00:00:00 2001
From: Vitaly Slobodin <vitaliy.slobodin@gmail.com>
Date: Wed, 25 Dec 2024 20:30:11 +0100
Subject: [PATCH] lsp: handle pull-based diagnostics model for related
 documents

---
 crates/lsp/src/lsp.rs             |  2 +-
 crates/project/src/lsp_command.rs | 50 +++++++++++++++++++++++--------
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/crates/lsp/src/lsp.rs b/crates/lsp/src/lsp.rs
index a929ef723fc960..54f7caee9698ff 100644
--- a/crates/lsp/src/lsp.rs
+++ b/crates/lsp/src/lsp.rs
@@ -759,7 +759,7 @@ impl LanguageServer {
                     }),
                     diagnostic: Some(DiagnosticClientCapabilities {
                         dynamic_registration: Some(false),
-                        related_document_support: Some(false),
+                        related_document_support: Some(true),
                     }),
                     ..TextDocumentClientCapabilities::default()
                 }),
diff --git a/crates/project/src/lsp_command.rs b/crates/project/src/lsp_command.rs
index 68f4860fbee354..b45be0675d8111 100644
--- a/crates/project/src/lsp_command.rs
+++ b/crates/project/src/lsp_command.rs
@@ -3260,36 +3260,60 @@ impl LspCommand for GetDocumentDiagnostics {
     async fn response_from_lsp(
         self,
         message: lsp::DocumentDiagnosticReportResult,
-        _: Model<LspStore>,
+        lsp_store: Model<LspStore>,
         buffer: Model<Buffer>,
         server_id: LanguageServerId,
-        cx: AsyncAppContext,
+        mut cx: AsyncAppContext,
     ) -> Result<Self::Response> {
         let uri = buffer.read_with(&cx, |buffer, cx| {
-            let file = buffer.file().and_then(|file| file.as_local())?;
-            Some(lsp::Url::from_file_path(file.abs_path(cx).clone()).unwrap())
+            buffer
+                .file()
+                .and_then(|file| file.as_local())
+                .map(|file| lsp::Url::from_file_path(file.abs_path(cx).clone()).unwrap())
         })?;
 
-        if uri.is_none() {
+        let Some(uri) = uri else {
             return Ok(None);
-        }
+        };
 
         match message {
             lsp::DocumentDiagnosticReportResult::Report(report) => match report {
-                lsp::DocumentDiagnosticReport::Full(report) => Ok(Some(LspDiagnostics {
-                    server_id,
-                    uri,
-                    diagnostics: Some(report.full_document_diagnostic_report.items.clone()),
-                })),
+                lsp::DocumentDiagnosticReport::Full(report) => {
+                    lsp_store.update(&mut cx, |store, cx| {
+                        for (uri, report) in report.related_documents.into_iter().flatten() {
+                            match report {
+                                lsp::DocumentDiagnosticReportKind::Full(report) => {
+                                    store.update_diagnostics(
+                                        server_id,
+                                        lsp::PublishDiagnosticsParams {
+                                            diagnostics: report.items.clone(),
+                                            uri,
+                                            version: None,
+                                        },
+                                        &[],
+                                        cx,
+                                    ).expect("Failed to update diagnostics for the related document");
+                                }
+                                lsp::DocumentDiagnosticReportKind::Unchanged(_) => (),
+                            }
+                        }
+                    }).expect("Failed to update diagnostics for related documents");
+
+                    Ok(Some(LspDiagnostics {
+                        server_id,
+                        uri: Some(uri),
+                        diagnostics: Some(report.full_document_diagnostic_report.items.clone()),
+                    }))
+                }
                 lsp::DocumentDiagnosticReport::Unchanged(_) => Ok(Some(LspDiagnostics {
                     server_id,
-                    uri,
+                    uri: Some(uri),
                     diagnostics: None,
                 })),
             },
             lsp::DocumentDiagnosticReportResult::Partial(_) => Ok(Some(LspDiagnostics {
                 server_id,
-                uri,
+                uri: Some(uri),
                 diagnostics: None,
             })),
         }