Allow to temporarily stop LSP servers (#28034)

Kirill Bulatov and Michael Sloan created

Same as `editor::RestartLanguageServer`, now there's an
`editor::StopLanguageServer` action that stops all language servers,
related to the currently opened editor.

Opening another singleton editor with the same language or changing
selections in a multi buffer will bring the servers back up.

Release Notes:

- Added a way to temporarily stop LSP servers

---------

Co-authored-by: Michael Sloan <mgsloan@gmail.com>

Change summary

crates/collab/src/rpc.rs        |  1 
crates/editor/src/actions.rs    |  1 
crates/editor/src/editor.rs     | 19 ++++++
crates/editor/src/element.rs    |  1 
crates/project/src/lsp_store.rs | 98 +++++++++++++++++++++++++---------
crates/project/src/project.rs   | 10 +++
crates/proto/proto/zed.proto    |  9 ++
crates/proto/src/proto.rs       |  3 +
8 files changed, 114 insertions(+), 28 deletions(-)

Detailed changes

crates/collab/src/rpc.rs 🔗

@@ -356,6 +356,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)
             .add_request_handler(forward_mutating_project_request::<proto::MultiLspQuery>)
             .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>)
+            .add_request_handler(forward_mutating_project_request::<proto::StopLanguageServers>)
             .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>)
             .add_message_handler(create_buffer_for_peer)
             .add_request_handler(update_buffer)

crates/editor/src/actions.rs 🔗

@@ -412,6 +412,7 @@ actions!(
         SortLinesCaseInsensitive,
         SortLinesCaseSensitive,
         SplitSelectionIntoLines,
+        StopLanguageServer,
         SwitchSourceHeader,
         Tab,
         Backtab,

crates/editor/src/editor.rs 🔗

@@ -14149,6 +14149,25 @@ impl Editor {
         }
     }
 
+    fn stop_language_server(
+        &mut self,
+        _: &StopLanguageServer,
+        _: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        if let Some(project) = self.project.clone() {
+            self.buffer.update(cx, |multi_buffer, cx| {
+                project.update(cx, |project, cx| {
+                    project.stop_language_servers_for_buffers(
+                        multi_buffer.all_buffers().into_iter().collect(),
+                        cx,
+                    );
+                    cx.emit(project::Event::RefreshInlayHints);
+                });
+            });
+        }
+    }
+
     fn cancel_language_server_work(
         workspace: &mut Workspace,
         _: &actions::CancelLanguageServerWork,

crates/editor/src/element.rs 🔗

@@ -452,6 +452,7 @@ impl EditorElement {
             }
         });
         register_action(editor, window, Editor::restart_language_server);
+        register_action(editor, window, Editor::stop_language_server);
         register_action(editor, window, Editor::show_character_palette);
         register_action(editor, window, |editor, action, window, cx| {
             if let Some(task) = editor.confirm_completion(action, window, cx) {

crates/project/src/lsp_store.rs 🔗

@@ -3391,6 +3391,7 @@ impl LspStore {
     pub fn init(client: &AnyProtoClient) {
         client.add_entity_request_handler(Self::handle_multi_lsp_query);
         client.add_entity_request_handler(Self::handle_restart_language_servers);
+        client.add_entity_request_handler(Self::handle_stop_language_servers);
         client.add_entity_request_handler(Self::handle_cancel_language_server_work);
         client.add_entity_message_handler(Self::handle_start_language_server);
         client.add_entity_message_handler(Self::handle_update_language_server);
@@ -7970,6 +7971,19 @@ impl LspStore {
         Ok(proto::Ack {})
     }
 
+    pub async fn handle_stop_language_servers(
+        this: Entity<Self>,
+        envelope: TypedEnvelope<proto::StopLanguageServers>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::Ack> {
+        this.update(&mut cx, |this, cx| {
+            let buffers = this.buffer_ids_to_buffers(envelope.payload.buffer_ids.into_iter(), cx);
+            this.stop_language_servers_for_buffers(buffers, cx);
+        })?;
+
+        Ok(proto::Ack {})
+    }
+
     pub async fn handle_cancel_language_server_work(
         this: Entity<Self>,
         envelope: TypedEnvelope<proto::CancelLanguageServerWork>,
@@ -8352,6 +8366,7 @@ impl LspStore {
         self.buffer_store.update(cx, |buffer_store, cx| {
             for buffer in buffer_store.buffers() {
                 buffer.update(cx, |buffer, cx| {
+                    // TODO kb clean inlays
                     buffer.update_diagnostics(server_id, DiagnosticSet::new([], buffer), cx);
                     buffer.set_completion_triggers(server_id, Default::default(), cx);
                 });
@@ -8418,34 +8433,9 @@ impl LspStore {
             });
             cx.background_spawn(request).detach_and_log_err(cx);
         } else {
-            let Some(local) = self.as_local_mut() else {
-                return;
-            };
-            let language_servers_to_stop = buffers
-                .iter()
-                .flat_map(|buffer| {
-                    buffer.update(cx, |buffer, cx| {
-                        local.language_server_ids_for_buffer(buffer, cx)
-                    })
-                })
-                .collect::<BTreeSet<_>>();
-            local.lsp_tree.update(cx, |this, _| {
-                this.remove_nodes(&language_servers_to_stop);
-            });
-            let tasks = language_servers_to_stop
-                .into_iter()
-                .map(|server| {
-                    let name = self
-                        .language_server_statuses
-                        .get(&server)
-                        .map(|state| state.name.as_str().into())
-                        .unwrap_or_else(|| LanguageServerName::from("Unknown"));
-                    self.stop_local_language_server(server, name, cx)
-                })
-                .collect::<Vec<_>>();
-
+            let stop_task = self.stop_local_language_servers_for_buffers(&buffers, cx);
             cx.spawn(async move |this, cx| {
-                cx.background_spawn(futures::future::join_all(tasks)).await;
+                stop_task.await;
                 this.update(cx, |this, cx| {
                     for buffer in buffers {
                         this.register_buffer_with_language_servers(&buffer, true, cx);
@@ -8457,6 +8447,60 @@ impl LspStore {
         }
     }
 
+    pub fn stop_language_servers_for_buffers(
+        &mut self,
+        buffers: Vec<Entity<Buffer>>,
+        cx: &mut Context<Self>,
+    ) {
+        if let Some((client, project_id)) = self.upstream_client() {
+            let request = client.request(proto::StopLanguageServers {
+                project_id,
+                buffer_ids: buffers
+                    .into_iter()
+                    .map(|b| b.read(cx).remote_id().to_proto())
+                    .collect(),
+            });
+            cx.background_spawn(request).detach_and_log_err(cx);
+        } else {
+            self.stop_local_language_servers_for_buffers(&buffers, cx)
+                .detach();
+        }
+    }
+
+    fn stop_local_language_servers_for_buffers(
+        &mut self,
+        buffers: &[Entity<Buffer>],
+        cx: &mut Context<Self>,
+    ) -> Task<()> {
+        let Some(local) = self.as_local_mut() else {
+            return Task::ready(());
+        };
+        let language_servers_to_stop = buffers
+            .iter()
+            .flat_map(|buffer| {
+                buffer.update(cx, |buffer, cx| {
+                    local.language_server_ids_for_buffer(buffer, cx)
+                })
+            })
+            .collect::<BTreeSet<_>>();
+        local.lsp_tree.update(cx, |this, _| {
+            this.remove_nodes(&language_servers_to_stop);
+        });
+        let tasks = language_servers_to_stop
+            .into_iter()
+            .map(|server| {
+                let name = self
+                    .language_server_statuses
+                    .get(&server)
+                    .map(|state| state.name.as_str().into())
+                    .unwrap_or_else(|| LanguageServerName::from("Unknown"));
+                self.stop_local_language_server(server, name, cx)
+            })
+            .collect::<Vec<_>>();
+
+        cx.background_spawn(futures::future::join_all(tasks).map(|_| ()))
+    }
+
     fn get_buffer<'a>(&self, abs_path: &Path, cx: &'a App) -> Option<&'a Buffer> {
         let (worktree, relative_path) =
             self.worktree_store.read(cx).find_worktree(&abs_path, cx)?;

crates/project/src/project.rs 🔗

@@ -3015,6 +3015,16 @@ impl Project {
         })
     }
 
+    pub fn stop_language_servers_for_buffers(
+        &mut self,
+        buffers: Vec<Entity<Buffer>>,
+        cx: &mut Context<Self>,
+    ) {
+        self.lsp_store.update(cx, |lsp_store, cx| {
+            lsp_store.stop_language_servers_for_buffers(buffers, cx)
+        })
+    }
+
     pub fn cancel_language_server_work_for_buffers(
         &mut self,
         buffers: impl IntoIterator<Item = Entity<Buffer>>,

crates/proto/proto/zed.proto 🔗

@@ -367,7 +367,9 @@ message Envelope {
         LanguageServerIdForNameResponse language_server_id_for_name_response = 333; // current max
 
         LoadCommitDiff load_commit_diff = 334;
-        LoadCommitDiffResponse load_commit_diff_response = 335; // current max
+        LoadCommitDiffResponse load_commit_diff_response = 335;
+
+        StopLanguageServers stop_language_servers = 336; // current max
     }
 
     reserved 87 to 88;
@@ -2445,6 +2447,11 @@ message RestartLanguageServers {
     repeated uint64 buffer_ids = 2;
 }
 
+message StopLanguageServers {
+    uint64 project_id = 1;
+    repeated uint64 buffer_ids = 2;
+}
+
 message MultiLspQueryResponse {
     repeated LspResponse responses = 1;
 }

crates/proto/src/proto.rs 🔗

@@ -400,6 +400,7 @@ messages!(
     (RespondToChannelInvite, Foreground),
     (RespondToContactRequest, Foreground),
     (RestartLanguageServers, Foreground),
+    (StopLanguageServers, Background),
     (RoomUpdated, Foreground),
     (SaveBuffer, Foreground),
     (SendChannelMessage, Background),
@@ -593,6 +594,7 @@ request_messages!(
     (RejoinRemoteProjects, RejoinRemoteProjectsResponse),
     (MultiLspQuery, MultiLspQueryResponse),
     (RestartLanguageServers, Ack),
+    (StopLanguageServers, Ack),
     (OpenContext, OpenContextResponse),
     (CreateContext, CreateContextResponse),
     (SynchronizeContexts, SynchronizeContextsResponse),
@@ -674,6 +676,7 @@ entity_messages!(
     LoadCommitDiff,
     MultiLspQuery,
     RestartLanguageServers,
+    StopLanguageServers,
     OnTypeFormatting,
     OpenNewBuffer,
     OpenBufferById,