fix(lsp): independent diagnostic publishing should include all diagnostic sources on each publish (#13483)

2 files changed, 122 insertions, 74 deletions

diff --git a/cli/lsp/diagnostics.rs b/cli/lsp/diagnostics.rs
index d1eef1d9f..5d474459a 100644
--- a/cli/lsp/diagnostics.rs
+++ b/cli/lsp/diagnostics.rs
@@ -40,6 +40,54 @@ pub type DiagnosticVec = Vec<DiagnosticRecord>;
 type DiagnosticMap =
   HashMap<ModuleSpecifier, (Option<i32>, Vec<lsp::Diagnostic>)>;
 type TsDiagnosticsMap = HashMap<String, Vec<diagnostics::Diagnostic>>;
+type DiagnosticsByVersionMap = HashMap<Option<i32>, Vec<lsp::Diagnostic>>;
+
+#[derive(Clone)]
+struct DiagnosticsPublisher {
+  client: Client,
+  all_diagnostics:
+    Arc<Mutex<HashMap<ModuleSpecifier, DiagnosticsByVersionMap>>>,
+}
+
+impl DiagnosticsPublisher {
+  pub fn new(client: Client) -> Self {
+    Self {
+      client,
+      all_diagnostics: Default::default(),
+    }
+  }
+
+  pub async fn publish(
+    &self,
+    diagnostics: DiagnosticVec,
+    token: &CancellationToken,
+  ) {
+    let mut all_diagnostics = self.all_diagnostics.lock().await;
+    for (specifier, version, diagnostics) in diagnostics {
+      if token.is_cancelled() {
+        return;
+      }
+
+      // the versions of all the published diagnostics should be the same, but just
+      // in case they're not keep track of that
+      let diagnostics_by_version =
+        all_diagnostics.entry(specifier.clone()).or_default();
+      let mut version_diagnostics =
+        diagnostics_by_version.entry(version).or_default();
+      version_diagnostics.extend(diagnostics);
+
+      self
+        .client
+        .publish_diagnostics(specifier, version_diagnostics.clone(), version)
+        .await;
+    }
+  }
+
+  pub async fn clear(&self) {
+    let mut all_diagnostics = self.all_diagnostics.lock().await;
+    all_diagnostics.clear();
+  }
+}
 
 #[derive(Debug)]
 pub(crate) struct DiagnosticsServer {
@@ -113,6 +161,7 @@ impl DiagnosticsServer {
         let mut ts_handle: Option<tokio::task::JoinHandle<()>> = None;
         let mut lint_handle: Option<tokio::task::JoinHandle<()>> = None;
         let mut deps_handle: Option<tokio::task::JoinHandle<()>> = None;
+        let diagnostics_publisher = DiagnosticsPublisher::new(client.clone());
 
         loop {
           match rx.recv().await {
@@ -122,6 +171,7 @@ impl DiagnosticsServer {
               // cancel the previous run
               token.cancel();
               token = CancellationToken::new();
+              diagnostics_publisher.clear().await;
 
               let (snapshot, config, maybe_lint_config) = {
                 let language_server = language_server.lock().await;
@@ -135,8 +185,8 @@ impl DiagnosticsServer {
               let previous_ts_handle = ts_handle.take();
               ts_handle = Some(tokio::spawn({
                 let performance = performance.clone();
+                let diagnostics_publisher = diagnostics_publisher.clone();
                 let ts_server = ts_server.clone();
-                let client = client.clone();
                 let token = token.clone();
                 let stored_ts_diagnostics = stored_ts_diagnostics.clone();
                 let snapshot = snapshot.clone();
@@ -184,12 +234,11 @@ impl DiagnosticsServer {
                         .collect();
                     }
 
-                    for (specifier, version, diagnostics) in diagnostics {
-                      client
-                        .publish_diagnostics(specifier, diagnostics, version)
-                        .await;
+                    diagnostics_publisher.publish(diagnostics, &token).await;
+
+                    if !token.is_cancelled() {
+                      performance.measure(mark);
                     }
-                    performance.measure(mark);
                   }
                 }
               }));
@@ -197,7 +246,7 @@ impl DiagnosticsServer {
               let previous_deps_handle = deps_handle.take();
               deps_handle = Some(tokio::spawn({
                 let performance = performance.clone();
-                let client = client.clone();
+                let diagnostics_publisher = diagnostics_publisher.clone();
                 let token = token.clone();
                 let snapshot = snapshot.clone();
                 let config = config.clone();
@@ -214,12 +263,9 @@ impl DiagnosticsServer {
                   )
                   .await;
 
+                  diagnostics_publisher.publish(diagnostics, &token).await;
+
                   if !token.is_cancelled() {
-                    for (specifier, version, diagnostics) in diagnostics {
-                      client
-                        .publish_diagnostics(specifier, diagnostics, version)
-                        .await;
-                    }
                     performance.measure(mark);
                   }
                 }
@@ -228,7 +274,7 @@ impl DiagnosticsServer {
               let previous_lint_handle = lint_handle.take();
               lint_handle = Some(tokio::spawn({
                 let performance = performance.clone();
-                let client = client.clone();
+                let diagnostics_publisher = diagnostics_publisher.clone();
                 let token = token.clone();
                 let snapshot = snapshot.clone();
                 let config = config.clone();
@@ -246,12 +292,9 @@ impl DiagnosticsServer {
                   )
                   .await;
 
+                  diagnostics_publisher.publish(diagnostics, &token).await;
+
                   if !token.is_cancelled() {
-                    for (specifier, version, diagnostics) in diagnostics {
-                      client
-                        .publish_diagnostics(specifier, diagnostics, version)
-                        .await;
-                    }
                     performance.measure(mark);
                   }
                 }
diff --git a/cli/tests/integration/lsp_tests.rs b/cli/tests/integration/lsp_tests.rs
index 35a2e29ac..12c50e757 100644
--- a/cli/tests/integration/lsp_tests.rs
+++ b/cli/tests/integration/lsp_tests.rs
@@ -10,6 +10,7 @@ use deno_core::serde_json::Value;
 use deno_core::url::Url;
 use lspower::lsp;
 use pretty_assertions::assert_eq;
+use std::collections::HashSet;
 use std::fs;
 use tempfile::TempDir;
 use test_util::deno_exe_path;
@@ -155,17 +156,32 @@ impl TestSession {
 struct CollectedDiagnostics(Vec<lsp::PublishDiagnosticsParams>);
 
 impl CollectedDiagnostics {
-  pub fn all(&self) -> Vec<lsp::Diagnostic> {
+  /// Gets the diagnostics that the editor will see after all the publishes.
+  pub fn viewed(&self) -> Vec<lsp::Diagnostic> {
     self
-      .0
-      .iter()
-      .flat_map(|p| p.diagnostics.clone().into_iter())
+      .viewed_messages()
+      .into_iter()
+      .flat_map(|m| m.diagnostics)
       .collect()
   }
 
+  /// Gets the messages that the editor will see after all the publishes.
+  pub fn viewed_messages(&self) -> Vec<lsp::PublishDiagnosticsParams> {
+    // go over the publishes in reverse order in order to get
+    // the final messages that will be shown in the editor
+    let mut messages = Vec::new();
+    let mut had_specifier = HashSet::new();
+    for message in self.0.iter().rev() {
+      if had_specifier.insert(message.uri.clone()) {
+        messages.insert(0, message.clone());
+      }
+    }
+    messages
+  }
+
   pub fn with_source(&self, source: &str) -> lsp::PublishDiagnosticsParams {
     self
-      .0
+      .viewed_messages()
       .iter()
       .find(|p| {
         p.diagnostics
@@ -183,7 +199,7 @@ impl CollectedDiagnostics {
   ) -> lsp::PublishDiagnosticsParams {
     let specifier = ModuleSpecifier::parse(specifier).unwrap();
     self
-      .0
+      .viewed_messages()
       .iter()
       .find(|p| {
         p.uri == specifier
@@ -2622,29 +2638,22 @@ fn lsp_code_actions() {
 
 #[test]
 fn lsp_code_actions_deno_cache() {
-  let mut client = init("initialize_params.json");
-  client
-    .write_notification("textDocument/didOpen", json!({
-      "textDocument": {
-        "uri": "file:///a/file.ts",
-        "languageId": "typescript",
-        "version": 1,
-        "text": "import * as a from \"https://deno.land/x/a/mod.ts\";\n\nconsole.log(a);\n"
-      }
-    }))
-    .unwrap();
-  let (id, method, _) = client.read_request::<Value>().unwrap();
-  assert_eq!(method, "workspace/configuration");
-  client
-    .write_response(id, json!([{ "enable": true }]))
-    .unwrap();
-  let diagnostics = read_diagnostics(&mut client);
+  let mut session = TestSession::from_file("initialize_params.json");
+  let diagnostics = session.did_open(json!({
+    "textDocument": {
+      "uri": "file:///a/file.ts",
+      "languageId": "typescript",
+      "version": 1,
+      "text": "import * as a from \"https://deno.land/x/a/mod.ts\";\n\nconsole.log(a);\n"
+    }
+  }));
   assert_eq!(
     diagnostics.with_source("deno"),
     load_fixture_as("diagnostics_deno_deps.json")
   );
 
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request(
       "textDocument/codeAction",
       load_fixture("code_action_params_cache.json"),
@@ -2655,7 +2664,7 @@ fn lsp_code_actions_deno_cache() {
     maybe_res,
     Some(load_fixture("code_action_response_cache.json"))
   );
-  shutdown(&mut client);
+  session.shutdown_and_exit();
 }
 
 #[test]
@@ -3215,26 +3224,22 @@ fn lsp_cache_location() {
   client
     .write_request::<_, _, Value>("initialize", params)
     .unwrap();
-
   client.write_notification("initialized", json!({})).unwrap();
-  did_open(
-    &mut client,
-    json!({
-      "textDocument": {
-        "uri": "file:///a/file_01.ts",
-        "languageId": "typescript",
-        "version": 1,
-        "text": "export const a = \"a\";\n",
-      }
-    }),
-  );
-  let diagnostics = did_open(
-    &mut client,
-    load_fixture("did_open_params_import_hover.json"),
-  );
-  let diagnostics = diagnostics.into_iter().flat_map(|x| x.diagnostics);
-  assert_eq!(diagnostics.count(), 7);
-  let (maybe_res, maybe_err) = client
+  let mut session = TestSession::from_client(client);
+
+  session.did_open(json!({
+    "textDocument": {
+      "uri": "file:///a/file_01.ts",
+      "languageId": "typescript",
+      "version": 1,
+      "text": "export const a = \"a\";\n",
+    }
+  }));
+  let diagnostics =
+    session.did_open(load_fixture("did_open_params_import_hover.json"));
+  assert_eq!(diagnostics.viewed().len(), 7);
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request::<_, _, Value>(
       "deno/cache",
       json!({
@@ -3247,7 +3252,8 @@ fn lsp_cache_location() {
     .unwrap();
   assert!(maybe_err.is_none());
   assert!(maybe_res.is_some());
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request(
       "textDocument/hover",
       json!({
@@ -3281,7 +3287,8 @@ fn lsp_cache_location() {
       }
     }))
   );
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request::<_, _, Value>(
       "textDocument/hover",
       json!({
@@ -3318,7 +3325,7 @@ fn lsp_cache_location() {
   let cache_path = temp_dir.path().join(".cache");
   assert!(cache_path.is_dir());
   assert!(cache_path.join("gen").is_dir());
-  shutdown(&mut client);
+  session.shutdown_and_exit();
 }
 
 /// Sets the TLS root certificate on startup, which allows the LSP to connect to
@@ -3351,7 +3358,7 @@ fn lsp_tls_cert() {
   }));
   let diagnostics =
     session.did_open(load_fixture("did_open_params_tls_cert.json"));
-  let diagnostics = diagnostics.all();
+  let diagnostics = diagnostics.viewed();
   assert_eq!(diagnostics.len(), 7);
   let (maybe_res, maybe_err) = session
     .client
@@ -3585,7 +3592,7 @@ fn lsp_diagnostics_deno_types() {
   assert!(maybe_res.is_some());
   assert!(maybe_err.is_none());
   let diagnostics = read_diagnostics(&mut client);
-  assert_eq!(diagnostics.all().len(), 5);
+  assert_eq!(diagnostics.viewed().len(), 5);
   shutdown(&mut client);
 }
 
@@ -3671,7 +3678,7 @@ fn lsp_diagnostics_refresh_dependents() {
     )
     .unwrap();
   let diagnostics = session.read_diagnostics();
-  assert_eq!(diagnostics.all().len(), 0); // no diagnostics now
+  assert_eq!(diagnostics.viewed().len(), 0); // no diagnostics now
 
   session.shutdown_and_exit();
   assert_eq!(session.client.queue_len(), 0);
@@ -4405,18 +4412,16 @@ fn lsp_lint_with_config() {
   client
     .write_request::<_, _, Value>("initialize", params)
     .unwrap();
+  let mut session = TestSession::from_client(client);
 
-  let diagnostics = did_open(&mut client, load_fixture("did_open_lint.json"));
-  let diagnostics = diagnostics
-    .into_iter()
-    .flat_map(|x| x.diagnostics)
-    .collect::<Vec<_>>();
+  let diagnostics = session.did_open(load_fixture("did_open_lint.json"));
+  let diagnostics = diagnostics.viewed();
   assert_eq!(diagnostics.len(), 1);
   assert_eq!(
     diagnostics[0].code,
     Some(lsp::NumberOrString::String("ban-untagged-todo".to_string()))
   );
-  shutdown(&mut client);
+  session.shutdown_and_exit();
 }
 
 #[test]