lsp: Fix dynamic registration of document diagnostics (#41929)

Piotr Osiewicz created

- lsp: Fix dynamic registration of diagnostic capabilities not taking
effect when an initial capability is not specified
Gist of the issue lies within use of .get_mut instead of .entry. If we
had not created any dynamic capability beforehand, we'd miss a
registration, essentially

- **Determine whether to update remote caps in a smarter manner**

Release Notes:

- Fixed document diagnostics with Ty language server.

Change summary

crates/project/src/lsp_store.rs | 55 +++++++++++++++++++++++-----------
1 file changed, 37 insertions(+), 18 deletions(-)

Detailed changes

crates/project/src/lsp_store.rs 🔗

@@ -238,7 +238,7 @@ pub struct DocumentDiagnostics {
     version: Option<i32>,
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct DynamicRegistrations {
     did_change_watched_files: HashMap<String, Vec<FileSystemWatcher>>,
     diagnostics: HashMap<Option<String>, DiagnosticServerCapabilities>,
@@ -12062,14 +12062,11 @@ impl LspStore {
                             .context("Could not obtain Language Servers state")?;
                         local
                             .language_server_dynamic_registrations
-                            .get_mut(&server_id)
-                            .and_then(|registrations| {
-                                registrations
-                                    .diagnostics
-                                    .insert(Some(reg.id.clone()), caps.clone())
-                            });
+                            .entry(server_id)
+                            .or_default()
+                            .diagnostics
+                            .insert(Some(reg.id.clone()), caps.clone());
 
-                        let mut can_now_provide_diagnostics = false;
                         if let LanguageServerState::Running {
                             workspace_diagnostics_refresh_tasks,
                             ..
@@ -12082,20 +12079,42 @@ impl LspStore {
                             )
                         {
                             workspace_diagnostics_refresh_tasks.insert(Some(reg.id), task);
-                            can_now_provide_diagnostics = true;
                         }
 
-                        // We don't actually care about capabilities.diagnostic_provider, but it IS relevant for the remote peer
-                        // to know that there's at least one provider. Otherwise, it will never ask us to issue documentdiagnostic calls on their behalf,
-                        // as it'll think that they're not supported.
-                        if can_now_provide_diagnostics {
-                            server.update_capabilities(|capabilities| {
-                                debug_assert!(capabilities.diagnostic_provider.is_none());
+                        let mut did_update_caps = false;
+                        server.update_capabilities(|capabilities| {
+                            if capabilities.diagnostic_provider.as_ref().is_none_or(
+                                |current_caps| {
+                                    let supports_workspace_diagnostics =
+                                        |capabilities: &DiagnosticServerCapabilities| {
+                                            match capabilities {
+                                            DiagnosticServerCapabilities::Options(
+                                                diagnostic_options,
+                                            ) => diagnostic_options.workspace_diagnostics,
+                                            DiagnosticServerCapabilities::RegistrationOptions(
+                                                diagnostic_registration_options,
+                                            ) => {
+                                                diagnostic_registration_options
+                                                    .diagnostic_options
+                                                    .workspace_diagnostics
+                                            }
+                                        }
+                                        };
+                                    // We don't actually care about capabilities.diagnostic_provider, but it IS relevant for the remote peer
+                                    // to know that there's at least one provider. Otherwise, it will never ask us to issue documentdiagnostic calls on their behalf,
+                                    // as it'll think that they're not supported.
+                                    // If we did not support any workspace diagnostics up to this point but now do, let's update.
+                                    !supports_workspace_diagnostics(current_caps)
+                                        & supports_workspace_diagnostics(&caps)
+                                },
+                            ) {
+                                did_update_caps = true;
                                 capabilities.diagnostic_provider = Some(caps);
-                            });
+                            }
+                        });
+                        if did_update_caps {
+                            notify_server_capabilities_updated(&server, cx);
                         }
-
-                        notify_server_capabilities_updated(&server, cx);
                     }
                 }
                 "textDocument/documentColor" => {