From f3320998a80d118e265ae4b98bc52f320c8a254a Mon Sep 17 00:00:00 2001
From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Date: Sun, 10 Nov 2024 10:29:10 +0100
Subject: [PATCH] lsp: Track completion triggers for each language separately
 (#20471)

This PR improves how we handle completions in buffers with multiple
LSPs.

Context: while working on
https://github.com/zed-industries/zed/issues/19777 with @mgsloan we
noticed that completion triggers coming from language servers are not
tracked properly. Namely, each buffer has `completion_triggers` field
which is read from the configuration of a language server. The problem
is, there can be multiple language servers for a single buffer, in which
case we'd just stick to the one that was registered last.

This PR makes the tracking a bit more fine-grained. We now track not
only what the completion triggers are, but also their origin server id.
Whenever completion triggers are updated, we recreate the completion
triggers set.
Release Notes:

- Fixed completions not triggering when multiple language servers are
used for a single file.
---
 crates/editor/src/editor.rs         |  5 +--
 crates/language/src/buffer.rs       | 64 +++++++++++++++++++++++------
 crates/language/src/proto.rs        |  5 ++-
 crates/languages/src/json.rs        |  1 -
 crates/project/src/lsp_store.rs     | 18 +++++++-
 crates/project/src/project_tests.rs | 21 ++++++++--
 crates/proto/proto/zed.proto        |  1 +
 7 files changed, 91 insertions(+), 24 deletions(-)

diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs
index 68a0e4044ec34a..3313fb100bda5a 100644
--- a/crates/editor/src/editor.rs
+++ b/crates/editor/src/editor.rs
@@ -13700,10 +13700,7 @@ impl CompletionProvider for Model<Project> {
             return true;
         }
 
-        buffer
-            .completion_triggers()
-            .iter()
-            .any(|string| string == text)
+        buffer.completion_triggers().contains(text)
     }
 }
 
diff --git a/crates/language/src/buffer.rs b/crates/language/src/buffer.rs
index 9568baeb5d303c..2f010a9c3fab1a 100644
--- a/crates/language/src/buffer.rs
+++ b/crates/language/src/buffer.rs
@@ -40,7 +40,7 @@ use std::{
     borrow::Cow,
     cell::Cell,
     cmp::{self, Ordering, Reverse},
-    collections::BTreeMap,
+    collections::{BTreeMap, BTreeSet},
     ffi::OsStr,
     fmt,
     future::Future,
@@ -126,7 +126,8 @@ pub struct Buffer {
     diagnostics: SmallVec<[(LanguageServerId, DiagnosticSet); 2]>,
     remote_selections: TreeMap<ReplicaId, SelectionSet>,
     diagnostics_timestamp: clock::Lamport,
-    completion_triggers: Vec<String>,
+    completion_triggers: BTreeSet<String>,
+    completion_triggers_per_language_server: HashMap<LanguageServerId, BTreeSet<String>>,
     completion_triggers_timestamp: clock::Lamport,
     deferred_ops: OperationQueue<Operation>,
     capability: Capability,
@@ -315,6 +316,8 @@ pub enum Operation {
         triggers: Vec<String>,
         /// The buffer's lamport timestamp.
         lamport_timestamp: clock::Lamport,
+        /// The language server ID.
+        server_id: LanguageServerId,
     },
 }
 
@@ -697,12 +700,15 @@ impl Buffer {
             }));
         }
 
-        operations.push(proto::serialize_operation(
-            &Operation::UpdateCompletionTriggers {
-                triggers: self.completion_triggers.clone(),
-                lamport_timestamp: self.completion_triggers_timestamp,
-            },
-        ));
+        for (server_id, completions) in &self.completion_triggers_per_language_server {
+            operations.push(proto::serialize_operation(
+                &Operation::UpdateCompletionTriggers {
+                    triggers: completions.iter().cloned().collect(),
+                    lamport_timestamp: self.completion_triggers_timestamp,
+                    server_id: *server_id,
+                },
+            ));
+        }
 
         let text_operations = self.text.operations().clone();
         cx.background_executor().spawn(async move {
@@ -774,6 +780,7 @@ impl Buffer {
             diagnostics: Default::default(),
             diagnostics_timestamp: Default::default(),
             completion_triggers: Default::default(),
+            completion_triggers_per_language_server: Default::default(),
             completion_triggers_timestamp: Default::default(),
             deferred_ops: OperationQueue::new(),
             has_conflict: false,
@@ -2229,8 +2236,21 @@ impl Buffer {
             Operation::UpdateCompletionTriggers {
                 triggers,
                 lamport_timestamp,
+                server_id,
             } => {
-                self.completion_triggers = triggers;
+                if triggers.is_empty() {
+                    self.completion_triggers_per_language_server
+                        .remove(&server_id);
+                    self.completion_triggers = self
+                        .completion_triggers_per_language_server
+                        .values()
+                        .flat_map(|triggers| triggers.into_iter().cloned())
+                        .collect();
+                } else {
+                    self.completion_triggers_per_language_server
+                        .insert(server_id, triggers.iter().cloned().collect());
+                    self.completion_triggers.extend(triggers);
+                }
                 self.text.lamport_clock.observe(lamport_timestamp);
             }
         }
@@ -2374,13 +2394,31 @@ impl Buffer {
     }
 
     /// Override current completion triggers with the user-provided completion triggers.
-    pub fn set_completion_triggers(&mut self, triggers: Vec<String>, cx: &mut ModelContext<Self>) {
-        self.completion_triggers.clone_from(&triggers);
+    pub fn set_completion_triggers(
+        &mut self,
+        server_id: LanguageServerId,
+        triggers: BTreeSet<String>,
+        cx: &mut ModelContext<Self>,
+    ) {
         self.completion_triggers_timestamp = self.text.lamport_clock.tick();
+        if triggers.is_empty() {
+            self.completion_triggers_per_language_server
+                .remove(&server_id);
+            self.completion_triggers = self
+                .completion_triggers_per_language_server
+                .values()
+                .flat_map(|triggers| triggers.into_iter().cloned())
+                .collect();
+        } else {
+            self.completion_triggers_per_language_server
+                .insert(server_id, triggers.clone());
+            self.completion_triggers.extend(triggers.iter().cloned());
+        }
         self.send_operation(
             Operation::UpdateCompletionTriggers {
-                triggers,
+                triggers: triggers.iter().cloned().collect(),
                 lamport_timestamp: self.completion_triggers_timestamp,
+                server_id,
             },
             true,
             cx,
@@ -2390,7 +2428,7 @@ impl Buffer {
 
     /// Returns a list of strings which trigger a completion menu for this language.
     /// Usually this is driven by LSP server which returns a list of trigger characters for completions.
-    pub fn completion_triggers(&self) -> &[String] {
+    pub fn completion_triggers(&self) -> &BTreeSet<String> {
         &self.completion_triggers
     }
 
diff --git a/crates/language/src/proto.rs b/crates/language/src/proto.rs
index a8aa8378b11c6d..ec864a9519eae1 100644
--- a/crates/language/src/proto.rs
+++ b/crates/language/src/proto.rs
@@ -79,11 +79,13 @@ pub fn serialize_operation(operation: &crate::Operation) -> proto::Operation {
             crate::Operation::UpdateCompletionTriggers {
                 triggers,
                 lamport_timestamp,
+                server_id,
             } => proto::operation::Variant::UpdateCompletionTriggers(
                 proto::operation::UpdateCompletionTriggers {
                     replica_id: lamport_timestamp.replica_id as u32,
                     lamport_timestamp: lamport_timestamp.value,
-                    triggers: triggers.clone(),
+                    triggers: triggers.iter().cloned().collect(),
+                    language_server_id: server_id.to_proto(),
                 },
             ),
         }),
@@ -326,6 +328,7 @@ pub fn deserialize_operation(message: proto::Operation) -> Result<crate::Operati
                         replica_id: message.replica_id as ReplicaId,
                         value: message.lamport_timestamp,
                     },
+                    server_id: LanguageServerId::from_proto(message.language_server_id),
                 }
             }
         },
diff --git a/crates/languages/src/json.rs b/crates/languages/src/json.rs
index d65b30ac79c499..924d3977043cfd 100644
--- a/crates/languages/src/json.rs
+++ b/crates/languages/src/json.rs
@@ -352,7 +352,6 @@ impl LspAdapter for NodeVersionAdapter {
             }
             remove_matching(&container_dir, |entry| entry != destination_path).await;
         }
-
         Ok(LanguageServerBinary {
             path: destination_path,
             env: None,
diff --git a/crates/project/src/lsp_store.rs b/crates/project/src/lsp_store.rs
index bc0bfdf4165064..57c046e282adc6 100644
--- a/crates/project/src/lsp_store.rs
+++ b/crates/project/src/lsp_store.rs
@@ -3351,11 +3351,17 @@ impl LspStore {
 
                     buffer_handle.update(cx, |buffer, cx| {
                         buffer.set_completion_triggers(
+                            server.server_id(),
                             server
                                 .capabilities()
                                 .completion_provider
                                 .as_ref()
-                                .and_then(|provider| provider.trigger_characters.clone())
+                                .and_then(|provider| {
+                                    provider
+                                        .trigger_characters
+                                        .as_ref()
+                                        .map(|characters| characters.iter().cloned().collect())
+                                })
                                 .unwrap_or_default(),
                             cx,
                         );
@@ -3394,6 +3400,7 @@ impl LspStore {
                 for adapter in self.languages.lsp_adapters(&language.name()) {
                     if let Some(server_id) = ids.get(&(worktree_id, adapter.name.clone())) {
                         buffer.update_diagnostics(*server_id, DiagnosticSet::new([], buffer), cx);
+                        buffer.set_completion_triggers(*server_id, Default::default(), cx);
                     }
                 }
             }
@@ -5797,6 +5804,7 @@ impl LspStore {
                                 DiagnosticSet::new([], buffer),
                                 cx,
                             );
+                            buffer.set_completion_triggers(server_id, Default::default(), cx);
                         });
                     }
                 });
@@ -6685,11 +6693,17 @@ impl LspStore {
 
                 buffer_handle.update(cx, |buffer, cx| {
                     buffer.set_completion_triggers(
+                        server_id,
                         language_server
                             .capabilities()
                             .completion_provider
                             .as_ref()
-                            .and_then(|provider| provider.trigger_characters.clone())
+                            .and_then(|provider| {
+                                provider
+                                    .trigger_characters
+                                    .as_ref()
+                                    .map(|characters| characters.iter().cloned().collect())
+                            })
                             .unwrap_or_default(),
                         cx,
                     )
diff --git a/crates/project/src/project_tests.rs b/crates/project/src/project_tests.rs
index a1ea2e4f3b48a5..49406b971585ef 100644
--- a/crates/project/src/project_tests.rs
+++ b/crates/project/src/project_tests.rs
@@ -481,7 +481,11 @@ async fn test_managing_language_servers(cx: &mut gpui::TestAppContext) {
     // The buffer is configured based on the language server's capabilities.
     rust_buffer.update(cx, |buffer, _| {
         assert_eq!(
-            buffer.completion_triggers(),
+            buffer
+                .completion_triggers()
+                .into_iter()
+                .cloned()
+                .collect::<Vec<_>>(),
             &[".".to_string(), "::".to_string()]
         );
     });
@@ -528,7 +532,14 @@ async fn test_managing_language_servers(cx: &mut gpui::TestAppContext) {
     // This buffer is configured based on the second language server's
     // capabilities.
     json_buffer.update(cx, |buffer, _| {
-        assert_eq!(buffer.completion_triggers(), &[":".to_string()]);
+        assert_eq!(
+            buffer
+                .completion_triggers()
+                .into_iter()
+                .cloned()
+                .collect::<Vec<_>>(),
+            &[":".to_string()]
+        );
     });
 
     // When opening another buffer whose language server is already running,
@@ -541,7 +552,11 @@ async fn test_managing_language_servers(cx: &mut gpui::TestAppContext) {
         .unwrap();
     rust_buffer2.update(cx, |buffer, _| {
         assert_eq!(
-            buffer.completion_triggers(),
+            buffer
+                .completion_triggers()
+                .into_iter()
+                .cloned()
+                .collect::<Vec<_>>(),
             &[".".to_string(), "::".to_string()]
         );
     });
diff --git a/crates/proto/proto/zed.proto b/crates/proto/proto/zed.proto
index c727296d6406f7..b86894bae1e361 100644
--- a/crates/proto/proto/zed.proto
+++ b/crates/proto/proto/zed.proto
@@ -1906,6 +1906,7 @@ message Operation {
         uint32 replica_id = 1;
         uint32 lamport_timestamp = 2;
         repeated string triggers = 3;
+        uint64 language_server_id = 4;
     }
 }