add version support to textDocument/publishDiagnostics

pyrefly/lib/lsp/non_wasm/server.rs

@@ -306,14 +306,16 @@ impl ServerConnection {
         &self,
         diags: SmallMap<PathBuf, Vec<Diagnostic>>,
         notebook_cell_urls: SmallMap<PathBuf, Url>,
+        version_info: &HashMap<PathBuf, i32>,
     ) {
         for (path, diags) in diags {
             if let Some(url) = notebook_cell_urls.get(&path) {
                 self.publish_diagnostics_for_uri(url.clone(), diags, None)
             } else {
                 let path = path.absolutize();
+                let version = version_info.get(&path).copied();
                 match Url::from_file_path(&path) {
-                    Ok(uri) => self.publish_diagnostics_for_uri(uri, diags, None),
+                    Ok(uri) => self.publish_diagnostics_for_uri(uri, diags, version),
                     Err(_) => eprint!("Unable to convert path to uri: {path:?}"),
                 }
             }
@@ -1467,8 +1469,11 @@ impl Server {
                 Self::append_unused_parameter_diagnostics(transaction, &handle, diagnostics);
                 Self::append_unused_import_diagnostics(transaction, &handle, diagnostics);
             }
-            self.connection
-                .publish_diagnostics(diags, notebook_cell_urls);
+            self.connection.publish_diagnostics(
+                diags,
+                notebook_cell_urls,
+                &*self.version_info.lock(),
+            );
             if self
                 .initialize_params
                 .capabilities
@@ -1979,17 +1984,21 @@ impl Server {
         let Some(path) = self.path_for_uri(&url) else {
             return;
         };
-        self.version_info.lock().remove(&path);
+        let version = self
+            .version_info
+            .lock()
+            .remove(&uri)
+            .map(|version| version + 1);
         let open_files = self.open_files.dupe();
         if let Some(LspFile::Notebook(notebook)) = open_files.write().remove(&path).as_deref() {
             for cell in notebook.cell_urls() {
                 self.connection
-                    .publish_diagnostics_for_uri(cell.clone(), Vec::new(), None);
+                    .publish_diagnostics_for_uri(cell.clone(), Vec::new(), version);
                 self.open_notebook_cells.write().remove(cell);
             }
         } else {
             self.connection
-                .publish_diagnostics_for_uri(url.clone(), Vec::new(), None);
+                .publish_diagnostics_for_uri(url, Vec::new(), version);
         }
         self.unsaved_file_tracker.forget_uri_path(&url);
         let state = self.state.dupe();

pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs

@@ -5,13 +5,18 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use lsp_server::Message;
+use lsp_server::Notification;
 use lsp_server::RequestId;
+use lsp_server::Response;
+use lsp_types::Url;
 use lsp_types::request::DocumentDiagnosticRequest;
 use pyrefly_config::environment::environment::PythonEnvironment;
 use serde_json::json;
 
 use crate::test::lsp::lsp_interaction::object_model::InitializeSettings;
 use crate::test::lsp::lsp_interaction::object_model::LspInteraction;
+use crate::test::lsp::lsp_interaction::object_model::ValidationResult;
 use crate::test::lsp::lsp_interaction::util::get_test_files_root;
 
 #[test]
@@ -760,3 +765,107 @@ fn test_shows_stdlib_errors_when_explicitly_included_in_project_includes() {
 
     interaction.shutdown();
 }
+
+#[test]
+fn test_version_support_publish_diagnostics() {
+    let test_files_root = get_test_files_root();
+    let root = test_files_root.path().to_path_buf();
+    let mut file = root.clone();
+    file.push("text_document.py");
+    let uri = Url::from_file_path(file).unwrap();
+    let mut interaction = LspInteraction::new();
+    interaction.set_root(root);
+    interaction.initialize(InitializeSettings {
+        configuration: Some(None),
+        capabilities: Some(serde_json::json!({
+            "textDocument": {
+                "publishDiagnostics": {
+                    "versionSupport": true,
+                },
+            },
+        })),
+        ..Default::default()
+    });
+
+    let gen_validator = |expected_version: i64| {
+        let actual_uri = uri.as_str();
+        move |msg: &Message| {
+            let Message::Notification(Notification { method, params }) = msg else {
+                return ValidationResult::Skip;
+            };
+            let Some(uri_val) = params.get("uri") else {
+                return ValidationResult::Skip;
+            };
+            let Some(expected_uri) = uri_val.as_str() else {
+                return ValidationResult::Skip;
+            };
+            if expected_uri == actual_uri && method == "textDocument/publishDiagnostics" {
+                if let Some(actual_version) = params.get("version") {
+                    if let Some(actual_version) = actual_version.as_i64() {
+                        assert!(
+                            actual_version <= expected_version,
+                            "expected version: {}, actual version: {}",
+                            expected_version,
+                            actual_version
+                        );
+                        return match actual_version.cmp(&expected_version) {
+                            std::cmp::Ordering::Less => ValidationResult::Skip,
+                            std::cmp::Ordering::Equal => ValidationResult::Pass,
+                            std::cmp::Ordering::Greater => ValidationResult::Fail,
+                        };
+                    }
+                }
+            }
+            ValidationResult::Skip
+        }
+    };
+
+    interaction.server.did_open("text_document.py");
+
+    let version = 1;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
+
+    interaction.server.did_change("text_document.py", "a = b");
+
+    let version = 2;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
+
+    interaction
+        .server
+        .send_message(Message::Notification(Notification {
+            method: "textDocument/didClose".to_owned(),
+            params: serde_json::json!({
+                "textDocument": {
+                    "uri": uri.as_str(),
+                    "languageId": "python",
+                    "version": 3
+                },
+            }),
+        }));
+
+    let version = 3;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
+
+    interaction.shutdown();
+}