add version support to textDocument/publishDiagnostics

pyrefly/lib/lsp/non_wasm/server.rs

@@ -294,11 +294,17 @@ impl ServerConnection {
         ));
     }
 
-    fn publish_diagnostics(&self, diags: SmallMap<PathBuf, Vec<Diagnostic>>) {
+    fn publish_diagnostics(
+        &self,
+        diags: SmallMap<PathBuf, Vec<Diagnostic>>,
+        version_info: &HashMap<PathBuf, i32>,
+    ) {
         for (path, diags) in diags {
             let path = path.absolutize();
             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_info.get(&path).copied())
+                }
                 Err(_) => eprint!("Unable to convert path to uri: {path:?}"),
             }
         }
@@ -1230,7 +1236,8 @@ impl Server {
                 let handle = make_open_handle(&self.state, path);
                 Self::append_unreachable_diagnostics(transaction, &handle, diagnostics);
             }
-            self.connection.publish_diagnostics(diags);
+            self.connection
+                .publish_diagnostics(diags, &*self.version_info.lock());
             if self
                 .initialize_params
                 .capabilities
@@ -1590,11 +1597,15 @@ impl Server {
 
     fn did_close(&self, params: DidCloseTextDocumentParams) {
         let uri = params.text_document.uri.to_file_path().unwrap();
-        self.version_info.lock().remove(&uri);
+        let version = self
+            .version_info
+            .lock()
+            .remove(&uri)
+            .map(|version| version + 1);
         let open_files = self.open_files.dupe();
         open_files.write().remove(&uri);
         self.connection
-            .publish_diagnostics_for_uri(params.text_document.uri, Vec::new(), None);
+            .publish_diagnostics_for_uri(params.text_document.uri, Vec::new(), version);
         let state = self.state.dupe();
         let lsp_queue = self.lsp_queue.dupe();
         let open_files = self.open_files.dupe();

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

@@ -5,11 +5,15 @@
  * 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 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]
@@ -239,3 +243,107 @@ fn test_publish_diagnostics_preserves_symlink_uri() {
 
     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();
+}