Fix remote clients unable to query custom, lsp_ext, commands (#27775)

Closes https://github.com/zed-industries/zed/issues/20583
Closes https://github.com/zed-industries/zed/issues/27133

A preparation for rust-analyzer's LSP tasks fetching, ensures all remote
clients are able to query custom, lsp_ext, commands.

Release Notes:

- Fixed remote clients unable to query custom, lsp_ext, commands

crates/editor/src/clangd_ext.rs

@@ -1,6 +1,8 @@
 use anyhow::Context as _;
 use gpui::{App, Context, Entity, Window};
 use language::Language;
+use project::lsp_store::lsp_ext_command::SwitchSourceHeaderResult;
+use rpc::proto;
 use url::Url;
 use workspace::{OpenOptions, OpenVisible};
 
@@ -27,34 +29,42 @@ pub fn switch_source_header(
         return;
     };
 
-    let Some((_, _, server_to_query, buffer)) =
-        find_specific_language_server_in_selection(editor, cx, is_c_language, CLANGD_SERVER_NAME)
-    else {
-        return;
-    };
-
+    let server_lookup =
+        find_specific_language_server_in_selection(editor, cx, is_c_language, CLANGD_SERVER_NAME);
     let project = project.clone();
-    let buffer_snapshot = buffer.read(cx).snapshot();
-    let source_file = buffer_snapshot
-        .file()
-        .unwrap()
-        .file_name(cx)
-        .to_str()
-        .unwrap()
-        .to_owned();
-
-    let switch_source_header_task = project.update(cx, |project, cx| {
-        project.request_lsp(
-            buffer,
-            project::LanguageServerToQuery::Other(server_to_query),
-            project::lsp_store::lsp_ext_command::SwitchSourceHeader,
-            cx,
-        )
-    });
+    let upstream_client = project.read(cx).lsp_store().read(cx).upstream_client();
     cx.spawn_in(window, async move |_editor, cx| {
-        let switch_source_header = switch_source_header_task
-            .await
-            .with_context(|| format!("Switch source/header LSP request for path \"{source_file}\" failed"))?;
+        let Some((_, _, server_to_query, buffer)) =
+            server_lookup.await
+        else {
+            return Ok(());
+        };
+        let source_file = buffer.update(cx, |buffer, _| {
+            buffer.file().map(|file| file.path()).map(|path| path.to_string_lossy().to_string()).unwrap_or_else(|| "Unknown".to_string())
+        })?;
+
+        let switch_source_header = if let Some((client, project_id)) = upstream_client {
+            let buffer_id = buffer.update(cx, |buffer, _| buffer.remote_id())?;
+            let request = proto::LspExtSwitchSourceHeader {
+                project_id,
+                buffer_id: buffer_id.to_proto(),
+            };
+            let response = client
+                .request(request)
+                .await
+                .context("lsp ext switch source header proto request")?;
+            SwitchSourceHeaderResult(response.target_file)
+        } else {
+            project.update(cx, |project, cx| {
+                project.request_lsp(
+                    buffer,
+                    project::LanguageServerToQuery::Other(server_to_query),
+                    project::lsp_store::lsp_ext_command::SwitchSourceHeader,
+                    cx,
+                )
+            })?.await.with_context(|| format!("Switch source/header LSP request for path \"{source_file}\" failed"))?
+        };
+
         if switch_source_header.0.is_empty() {
             log::info!("Clangd returned an empty string when requesting to switch source/header from \"{source_file}\"" );
             return Ok(());
@@ -87,10 +97,15 @@ pub fn switch_source_header(
 }
 
 pub fn apply_related_actions(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) {
-    if editor.update(cx, |e, cx| {
-        find_specific_language_server_in_selection(e, cx, is_c_language, CLANGD_SERVER_NAME)
-            .is_some()
-    }) {
-        register_action(editor, window, switch_source_header);
+    if editor
+        .read(cx)
+        .buffer()
+        .read(cx)
+        .all_buffers()
+        .into_iter()
+        .filter_map(|buffer| buffer.read(cx).language())
+        .any(|language| is_c_language(language))
+    {
+        register_action(&editor, window, switch_source_header);
     }
 }

crates/editor/src/element.rs

@@ -187,6 +187,7 @@ impl EditorElement {
 
         crate::rust_analyzer_ext::apply_related_actions(editor, window, cx);
         crate::clangd_ext::apply_related_actions(editor, window, cx);
+
         register_action(editor, window, Editor::open_context_menu);
         register_action(editor, window, Editor::move_left);
         register_action(editor, window, Editor::move_right);

crates/editor/src/lsp_ext.rs

@@ -1,9 +1,8 @@
-use std::collections::hash_map::Entry;
 use std::sync::Arc;
 
 use crate::Editor;
-use collections::HashMap;
-use gpui::{App, Entity};
+use gpui::{App, AppContext as _, Entity, Task};
+use itertools::Itertools;
 use language::Buffer;
 use language::Language;
 use lsp::LanguageServerId;
@@ -14,44 +13,50 @@ pub(crate) fn find_specific_language_server_in_selection<F>(
     cx: &mut App,
     filter_language: F,
     language_server_name: &str,
-) -> Option<(Anchor, Arc<Language>, LanguageServerId, Entity<Buffer>)>
+) -> Task<Option<(Anchor, Arc<Language>, LanguageServerId, Entity<Buffer>)>>
 where
     F: Fn(&Language) -> bool,
 {
     let Some(project) = &editor.project else {
-        return None;
+        return Task::ready(None);
     };
-    let mut language_servers_for = HashMap::default();
-    editor
+
+    let applicable_buffers = editor
         .selections
         .disjoint_anchors()
         .iter()
         .filter(|selection| selection.start == selection.end)
-        .filter_map(|selection| Some((selection.start.buffer_id?, selection.start)))
-        .find_map(|(buffer_id, trigger_anchor)| {
+        .filter_map(|selection| Some((selection.start, selection.start.buffer_id?)))
+        .filter_map(|(trigger_anchor, buffer_id)| {
             let buffer = editor.buffer().read(cx).buffer(buffer_id)?;
-            let server_id = *match language_servers_for.entry(buffer_id) {
-                Entry::Occupied(occupied_entry) => occupied_entry.into_mut(),
-                Entry::Vacant(vacant_entry) => {
-                    let language_server_id = buffer.update(cx, |buffer, cx| {
-                        project.update(cx, |project, cx| {
-                            project.language_server_id_for_name(buffer, language_server_name, cx)
-                        })
-                    });
-                    vacant_entry.insert(language_server_id)
-                }
-            }
-            .as_ref()?;
-
             let language = buffer.read(cx).language_at(trigger_anchor.text_anchor)?;
-            if !filter_language(&language) {
-                return None;
+            if filter_language(&language) {
+                Some((trigger_anchor, buffer, language))
+            } else {
+                None
             }
-            Some((
-                trigger_anchor,
-                Arc::clone(&language),
-                server_id,
-                buffer.clone(),
-            ))
         })
+        .unique_by(|(_, buffer, _)| buffer.read(cx).remote_id())
+        .collect::<Vec<_>>();
+
+    let applicable_buffer_tasks = applicable_buffers
+        .into_iter()
+        .map(|(trigger_anchor, buffer, language)| {
+            let task = buffer.update(cx, |buffer, cx| {
+                project.update(cx, |project, cx| {
+                    project.language_server_id_for_name(buffer, language_server_name, cx)
+                })
+            });
+            (trigger_anchor, buffer, language, task)
+        })
+        .collect::<Vec<_>>();
+    cx.background_spawn(async move {
+        for (trigger_anchor, buffer, language, task) in applicable_buffer_tasks {
+            if let Some(server_id) = task.await {
+                return Some((trigger_anchor, language, server_id, buffer));
+            }
+        }
+
+        None
+    })
 }

crates/editor/src/rust_analyzer_ext.rs

@@ -2,9 +2,13 @@ use std::{fs, path::Path};
 
 use anyhow::Context as _;
 use gpui::{App, AppContext as _, Context, Entity, Window};
-use language::{Capability, Language};
+use language::{proto::serialize_anchor, Capability, Language};
 use multi_buffer::MultiBuffer;
-use project::lsp_store::{lsp_ext_command::ExpandMacro, rust_analyzer_ext::RUST_ANALYZER_NAME};
+use project::lsp_store::{
+    lsp_ext_command::{DocsUrls, ExpandMacro, ExpandedMacro},
+    rust_analyzer_ext::RUST_ANALYZER_NAME,
+};
+use rpc::proto;
 use text::ToPointUtf16;
 
 use crate::{
@@ -18,13 +22,16 @@ fn is_rust_language(language: &Language) -> bool {
 
 pub fn apply_related_actions(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) {
     if editor
-        .update(cx, |e, cx| {
-            find_specific_language_server_in_selection(e, cx, is_rust_language, RUST_ANALYZER_NAME)
-        })
-        .is_some()
+        .read(cx)
+        .buffer()
+        .read(cx)
+        .all_buffers()
+        .into_iter()
+        .filter_map(|buffer| buffer.read(cx).language())
+        .any(|language| is_rust_language(language))
     {
-        register_action(editor, window, expand_macro_recursively);
-        register_action(editor, window, open_docs);
+        register_action(&editor, window, expand_macro_recursively);
+        register_action(&editor, window, open_docs);
     }
 }
 
@@ -44,32 +51,57 @@ pub fn expand_macro_recursively(
         return;
     };
 
-    let Some((trigger_anchor, rust_language, server_to_query, buffer)) =
-        find_specific_language_server_in_selection(
-            editor,
-            cx,
-            is_rust_language,
-            RUST_ANALYZER_NAME,
-        )
-    else {
-        return;
-    };
+    let server_lookup = find_specific_language_server_in_selection(
+        editor,
+        cx,
+        is_rust_language,
+        RUST_ANALYZER_NAME,
+    );
 
     let project = project.clone();
-    let buffer_snapshot = buffer.read(cx).snapshot();
-    let position = trigger_anchor.text_anchor.to_point_utf16(&buffer_snapshot);
-    let expand_macro_task = project.update(cx, |project, cx| {
-        project.request_lsp(
-            buffer,
-            project::LanguageServerToQuery::Other(server_to_query),
-            ExpandMacro { position },
-            cx,
-        )
-    });
+    let upstream_client = project.read(cx).lsp_store().read(cx).upstream_client();
     cx.spawn_in(window, async move |_editor, cx| {
-        let macro_expansion = expand_macro_task.await.context("expand macro")?;
+        let Some((trigger_anchor, rust_language, server_to_query, buffer)) = server_lookup.await
+        else {
+            return Ok(());
+        };
+
+        let macro_expansion = if let Some((client, project_id)) = upstream_client {
+            let buffer_id = buffer.update(cx, |buffer, _| buffer.remote_id())?;
+            let request = proto::LspExtExpandMacro {
+                project_id,
+                buffer_id: buffer_id.to_proto(),
+                position: Some(serialize_anchor(&trigger_anchor.text_anchor)),
+            };
+            let response = client
+                .request(request)
+                .await
+                .context("lsp ext expand macro proto request")?;
+            ExpandedMacro {
+                name: response.name,
+                expansion: response.expansion,
+            }
+        } else {
+            let buffer_snapshot = buffer.update(cx, |buffer, _| buffer.snapshot())?;
+            let position = trigger_anchor.text_anchor.to_point_utf16(&buffer_snapshot);
+            project
+                .update(cx, |project, cx| {
+                    project.request_lsp(
+                        buffer,
+                        project::LanguageServerToQuery::Other(server_to_query),
+                        ExpandMacro { position },
+                        cx,
+                    )
+                })?
+                .await
+                .context("expand macro")?
+        };
+
         if macro_expansion.is_empty() {
-            log::info!("Empty macro expansion for position {position:?}");
+            log::info!(
+                "Empty macro expansion for position {:?}",
+                trigger_anchor.text_anchor
+            );
             return Ok(());
         }
 
@@ -111,36 +143,57 @@ pub fn open_docs(editor: &mut Editor, _: &OpenDocs, window: &mut Window, cx: &mu
         return;
     };
 
-    let Some((trigger_anchor, _rust_language, server_to_query, buffer)) =
-        find_specific_language_server_in_selection(
-            editor,
-            cx,
-            is_rust_language,
-            RUST_ANALYZER_NAME,
-        )
-    else {
-        return;
-    };
+    let server_lookup = find_specific_language_server_in_selection(
+        editor,
+        cx,
+        is_rust_language,
+        RUST_ANALYZER_NAME,
+    );
 
     let project = project.clone();
-    let buffer_snapshot = buffer.read(cx).snapshot();
-    let position = trigger_anchor.text_anchor.to_point_utf16(&buffer_snapshot);
-    let open_docs_task = project.update(cx, |project, cx| {
-        project.request_lsp(
-            buffer,
-            project::LanguageServerToQuery::Other(server_to_query),
-            project::lsp_store::lsp_ext_command::OpenDocs { position },
-            cx,
-        )
-    });
-
+    let upstream_client = project.read(cx).lsp_store().read(cx).upstream_client();
     cx.spawn_in(window, async move |_editor, cx| {
-        let docs_urls = open_docs_task.await.context("open docs")?;
-        if docs_urls.is_empty() {
-            log::debug!("Empty docs urls for position {position:?}");
+        let Some((trigger_anchor, _, server_to_query, buffer)) = server_lookup.await else {
             return Ok(());
+        };
+
+        let docs_urls = if let Some((client, project_id)) = upstream_client {
+            let buffer_id = buffer.update(cx, |buffer, _| buffer.remote_id())?;
+            let request = proto::LspExtOpenDocs {
+                project_id,
+                buffer_id: buffer_id.to_proto(),
+                position: Some(serialize_anchor(&trigger_anchor.text_anchor)),
+            };
+            let response = client
+                .request(request)
+                .await
+                .context("lsp ext open docs proto request")?;
+            DocsUrls {
+                web: response.web,
+                local: response.local,
+            }
         } else {
-            log::debug!("{:?}", docs_urls);
+            let buffer_snapshot = buffer.update(cx, |buffer, _| buffer.snapshot())?;
+            let position = trigger_anchor.text_anchor.to_point_utf16(&buffer_snapshot);
+            project
+                .update(cx, |project, cx| {
+                    project.request_lsp(
+                        buffer,
+                        project::LanguageServerToQuery::Other(server_to_query),
+                        project::lsp_store::lsp_ext_command::OpenDocs { position },
+                        cx,
+                    )
+                })?
+                .await
+                .context("open docs")?
+        };
+
+        if docs_urls.is_empty() {
+            log::debug!(
+                "Empty docs urls for position {:?}",
+                trigger_anchor.text_anchor
+            );
+            return Ok(());
         }
 
         workspace.update(cx, |_workspace, cx| {

crates/editor/src/test.rs

@@ -14,6 +14,8 @@ use gpui::{
 use project::Project;
 use util::test::{marked_text_offsets, marked_text_ranges};
 
+pub use crate::rust_analyzer_ext::expand_macro_recursively;
+
 #[cfg(test)]
 #[ctor::ctor]
 fn init_logger() {