Hack extensions language servers to record their IDs

Lukas Wirth created 3 months ago

Change summary

crates/extension/src/extension_host_proxy.rs           |  6 +-
crates/language_extension/src/extension_lsp_adapter.rs | 29 ++++++++++-
crates/language_extension/src/language_extension.rs    | 10 +++
3 files changed, 36 insertions(+), 9 deletions(-)

Detailed changes

crates/extension/src/extension_host_proxy.rs

@@ -281,7 +281,7 @@ pub trait ExtensionLanguageServerProxy: Send + Sync + 'static {
 
     fn update_language_server_status(
         &self,
-        language_server_id: LanguageServerName,
+        language_server_name: LanguageServerName,
         status: BinaryStatus,
     );
 }
@@ -315,14 +315,14 @@ impl ExtensionLanguageServerProxy for ExtensionHostProxy {
 
     fn update_language_server_status(
         &self,
-        language_server_id: LanguageServerName,
+        language_server_name: LanguageServerName,
         status: BinaryStatus,
     ) {
         let Some(proxy) = self.language_server_proxy.read().clone() else {
             return;
         };
 
-        proxy.update_language_server_status(language_server_id, status)
+        proxy.update_language_server_status(language_server_name, status)
     }
 }
 

crates/language_extension/src/extension_lsp_adapter.rs

@@ -1,8 +1,8 @@
-use std::any::Any;
 use std::ops::Range;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::{any::Any, sync::OnceLock};
 
 use anyhow::{Context as _, Result};
 use async_trait::async_trait;
@@ -61,12 +61,18 @@ impl ExtensionLanguageServerProxy for LanguageServerRegistryProxy {
         language_server_name: LanguageServerName,
         language: LanguageName,
     ) {
+        let language_server_id = Arc::new(OnceLock::new());
+        self.language_server_ids
+            .write()
+            .unwrap()
+            .insert(language_server_name.clone(), language_server_id.clone());
         self.language_registry.register_lsp_adapter(
             language.clone(),
             Arc::new(ExtensionLspAdapter::new(
                 extension,
                 language_server_name,
                 language,
+                language_server_id,
             )),
         );
     }
@@ -125,8 +131,16 @@ impl ExtensionLanguageServerProxy for LanguageServerRegistryProxy {
         language_server_name: LanguageServerName,
         status: BinaryStatus,
     ) {
-        // self.language_registry
-        //     .update_lsp_binary_status(language_server_name, status);
+        if let Some(id) = self
+            .language_server_ids
+            .read()
+            .unwrap()
+            .get(&language_server_name)
+            && let Some(&id) = id.get()
+        {
+            self.language_registry
+                .update_lsp_binary_status(id, language_server_name, status);
+        }
     }
 }
 
@@ -134,6 +148,7 @@ struct ExtensionLspAdapter {
     extension: Arc<dyn Extension>,
     language_server_name: LanguageServerName,
     language_name: LanguageName,
+    language_server_id: Arc<OnceLock<LanguageServerId>>,
 }
 
 impl ExtensionLspAdapter {
@@ -141,11 +156,13 @@ impl ExtensionLspAdapter {
         extension: Arc<dyn Extension>,
         language_server_name: LanguageServerName,
         language_name: LanguageName,
+        language_server_id: Arc<OnceLock<LanguageServerId>>,
     ) -> Self {
         Self {
             extension,
             language_server_name,
             language_name,
+            language_server_id,
         }
     }
 }
@@ -162,7 +179,7 @@ impl LspAdapter for ExtensionLspAdapter {
         _: Option<Toolchain>,
         _: LanguageServerBinaryOptions,
         _: futures::lock::MutexGuard<'a, Option<LanguageServerBinary>>,
-        _: LanguageServerId,
+        language_server_id: LanguageServerId,
         _: &'a mut AsyncApp,
     ) -> Pin<Box<dyn 'a + Future<Output = Result<LanguageServerBinary>>>> {
         async move {
@@ -192,6 +209,10 @@ impl LspAdapter for ExtensionLspAdapter {
                     .await
                     .context("failed to set file permissions")?;
             }
+            self.language_server_id
+                .set(language_server_id)
+                .ok()
+                .context("failed to set language server id")?;
 
             Ok(LanguageServerBinary {
                 path,

crates/language_extension/src/language_extension.rs

@@ -1,12 +1,14 @@
 mod extension_lsp_adapter;
 
-use std::path::PathBuf;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
+use std::{path::PathBuf, sync::RwLock};
 
 use anyhow::Result;
+use collections::FxHashMap;
 use extension::{ExtensionGrammarProxy, ExtensionHostProxy, ExtensionLanguageProxy};
 use gpui::{App, Entity};
 use language::{LanguageMatcher, LanguageName, LanguageRegistry, LoadedLanguage};
+use lsp::{LanguageServerId, LanguageServerName};
 use project::LspStore;
 
 #[derive(Clone)]
@@ -23,6 +25,7 @@ pub fn init(
 ) {
     let language_server_registry_proxy = LanguageServerRegistryProxy {
         language_registry,
+        language_server_ids: Default::default(),
         lsp_access,
     };
     extension_host_proxy.register_grammar_proxy(language_server_registry_proxy.clone());
@@ -33,6 +36,9 @@ pub fn init(
 #[derive(Clone)]
 struct LanguageServerRegistryProxy {
     language_registry: Arc<LanguageRegistry>,
+    // extensions can only spawn one language server currently
+    language_server_ids:
+        Arc<RwLock<FxHashMap<LanguageServerName, Arc<OnceLock<LanguageServerId>>>>>,
     lsp_access: LspAccess,
 }