Add support for context server extensions (#20250)

This PR adds support for context servers provided by extensions.

To provide a context server from an extension, you need to list the
context servers in your `extension.toml`:

```toml
[context_servers.my-context-server]
```

And then implement the `context_server_command` method to return the
command that will be used to start the context server:

```rs
use zed_extension_api::{self as zed, Command, ContextServerId, Result};

struct ExampleContextServerExtension;

impl zed::Extension for ExampleContextServerExtension {
    fn new() -> Self {
        ExampleContextServerExtension
    }

    fn context_server_command(&mut self, _context_server_id: &ContextServerId) -> Result<Command> {
        Ok(Command {
            command: "node".to_string(),
            args: vec!["/path/to/example-context-server/index.js".to_string()],
            env: Vec::new(),
        })
    }
}

zed::register_extension!(ExampleContextServerExtension);
```

Release Notes:

- N/A

crates/extensions_ui/Cargo.toml

@@ -20,6 +20,7 @@ assistant_slash_command.workspace = true
 async-trait.workspace = true
 client.workspace = true
 collections.workspace = true
+context_servers.workspace = true
 db.workspace = true
 editor.workspace = true
 extension_host.workspace = true

crates/extensions_ui/src/extension_context_server.rs

@@ -0,0 +1,80 @@
+use std::pin::Pin;
+use std::sync::Arc;
+
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use context_servers::manager::{NativeContextServer, ServerConfig};
+use context_servers::protocol::InitializedContextServerProtocol;
+use context_servers::ContextServer;
+use extension_host::wasm_host::{WasmExtension, WasmHost};
+use futures::{Future, FutureExt};
+use gpui::AsyncAppContext;
+
+pub struct ExtensionContextServer {
+    #[allow(unused)]
+    pub(crate) extension: WasmExtension,
+    #[allow(unused)]
+    pub(crate) host: Arc<WasmHost>,
+    id: Arc<str>,
+    context_server: Arc<NativeContextServer>,
+}
+
+impl ExtensionContextServer {
+    pub async fn new(extension: WasmExtension, host: Arc<WasmHost>, id: Arc<str>) -> Result<Self> {
+        let command = extension
+            .call({
+                let id = id.clone();
+                |extension, store| {
+                    async move {
+                        let command = extension
+                            .call_context_server_command(store, id.clone())
+                            .await?
+                            .map_err(|e| anyhow!("{}", e))?;
+                        anyhow::Ok(command)
+                    }
+                    .boxed()
+                }
+            })
+            .await?;
+
+        let config = Arc::new(ServerConfig {
+            id: id.to_string(),
+            executable: command.command,
+            args: command.args,
+            env: Some(command.env.into_iter().collect()),
+        });
+
+        anyhow::Ok(Self {
+            extension,
+            host,
+            id,
+            context_server: Arc::new(NativeContextServer::new(config)),
+        })
+    }
+}
+
+#[async_trait(?Send)]
+impl ContextServer for ExtensionContextServer {
+    fn id(&self) -> Arc<str> {
+        self.id.clone()
+    }
+
+    fn config(&self) -> Arc<ServerConfig> {
+        self.context_server.config()
+    }
+
+    fn client(&self) -> Option<Arc<InitializedContextServerProtocol>> {
+        self.context_server.client()
+    }
+
+    fn start<'a>(
+        self: Arc<Self>,
+        cx: &'a AsyncAppContext,
+    ) -> Pin<Box<dyn 'a + Future<Output = Result<()>>>> {
+        self.context_server.clone().start(cx)
+    }
+
+    fn stop(&self) -> Result<()> {
+        self.context_server.stop()
+    }
+}

crates/extensions_ui/src/extension_registration_hooks.rs

@@ -2,6 +2,7 @@ use std::{path::PathBuf, sync::Arc};
 
 use anyhow::Result;
 use assistant_slash_command::SlashCommandRegistry;
+use context_servers::ContextServerFactoryRegistry;
 use extension_host::{extension_lsp_adapter::ExtensionLspAdapter, wasm_host};
 use fs::Fs;
 use gpui::{AppContext, BackgroundExecutor, Task};
@@ -11,6 +12,7 @@ use snippet_provider::SnippetRegistry;
 use theme::{ThemeRegistry, ThemeSettings};
 use ui::SharedString;
 
+use crate::extension_context_server::ExtensionContextServer;
 use crate::{extension_indexed_docs_provider, extension_slash_command::ExtensionSlashCommand};
 
 pub struct ConcreteExtensionRegistrationHooks {
@@ -19,6 +21,7 @@ pub struct ConcreteExtensionRegistrationHooks {
     indexed_docs_registry: Arc<IndexedDocsRegistry>,
     snippet_registry: Arc<SnippetRegistry>,
     language_registry: Arc<LanguageRegistry>,
+    context_server_factory_registry: Arc<ContextServerFactoryRegistry>,
     executor: BackgroundExecutor,
 }
 
@@ -29,6 +32,7 @@ impl ConcreteExtensionRegistrationHooks {
         indexed_docs_registry: Arc<IndexedDocsRegistry>,
         snippet_registry: Arc<SnippetRegistry>,
         language_registry: Arc<LanguageRegistry>,
+        context_server_factory_registry: Arc<ContextServerFactoryRegistry>,
         cx: &AppContext,
     ) -> Arc<dyn extension_host::ExtensionRegistrationHooks> {
         Arc::new(Self {
@@ -37,6 +41,7 @@ impl ConcreteExtensionRegistrationHooks {
             indexed_docs_registry,
             snippet_registry,
             language_registry,
+            context_server_factory_registry,
             executor: cx.background_executor().clone(),
         })
     }
@@ -69,6 +74,31 @@ impl extension_host::ExtensionRegistrationHooks for ConcreteExtensionRegistratio
         )
     }
 
+    fn register_context_server(
+        &self,
+        id: Arc<str>,
+        extension: wasm_host::WasmExtension,
+        host: Arc<wasm_host::WasmHost>,
+    ) {
+        self.context_server_factory_registry
+            .register_server_factory(
+                id.clone(),
+                Arc::new({
+                    move |cx| {
+                        let id = id.clone();
+                        let extension = extension.clone();
+                        let host = host.clone();
+                        cx.spawn(|_cx| async move {
+                            let context_server =
+                                ExtensionContextServer::new(extension, host, id).await?;
+
+                            anyhow::Ok(Arc::new(context_server) as _)
+                        })
+                    }
+                }),
+            );
+    }
+
     fn register_docs_provider(
         &self,
         extension: wasm_host::WasmExtension,

crates/extensions_ui/src/extension_store_test.rs

@@ -1,6 +1,7 @@
 use assistant_slash_command::SlashCommandRegistry;
 use async_compression::futures::bufread::GzipEncoder;
 use collections::BTreeMap;
+use context_servers::ContextServerFactoryRegistry;
 use extension_host::ExtensionSettings;
 use extension_host::SchemaVersion;
 use extension_host::{
@@ -161,6 +162,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
                         .into_iter()
                         .collect(),
                         language_servers: BTreeMap::default(),
+                        context_servers: BTreeMap::default(),
                         slash_commands: BTreeMap::default(),
                         indexed_docs_providers: BTreeMap::default(),
                         snippets: None,
@@ -187,6 +189,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
                         languages: Default::default(),
                         grammars: BTreeMap::default(),
                         language_servers: BTreeMap::default(),
+                        context_servers: BTreeMap::default(),
                         slash_commands: BTreeMap::default(),
                         indexed_docs_providers: BTreeMap::default(),
                         snippets: None,
@@ -264,6 +267,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
     let slash_command_registry = SlashCommandRegistry::new();
     let indexed_docs_registry = Arc::new(IndexedDocsRegistry::new(cx.executor()));
     let snippet_registry = Arc::new(SnippetRegistry::new());
+    let context_server_factory_registry = ContextServerFactoryRegistry::new();
     let node_runtime = NodeRuntime::unavailable();
 
     let store = cx.new_model(|cx| {
@@ -273,6 +277,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
             indexed_docs_registry.clone(),
             snippet_registry.clone(),
             language_registry.clone(),
+            context_server_factory_registry.clone(),
             cx,
         );
 
@@ -356,6 +361,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
                 languages: Default::default(),
                 grammars: BTreeMap::default(),
                 language_servers: BTreeMap::default(),
+                context_servers: BTreeMap::default(),
                 slash_commands: BTreeMap::default(),
                 indexed_docs_providers: BTreeMap::default(),
                 snippets: None,
@@ -406,6 +412,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
             indexed_docs_registry,
             snippet_registry,
             language_registry.clone(),
+            context_server_factory_registry.clone(),
             cx,
         );
 
@@ -500,6 +507,7 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
     let slash_command_registry = SlashCommandRegistry::new();
     let indexed_docs_registry = Arc::new(IndexedDocsRegistry::new(cx.executor()));
     let snippet_registry = Arc::new(SnippetRegistry::new());
+    let context_server_factory_registry = ContextServerFactoryRegistry::new();
     let node_runtime = NodeRuntime::unavailable();
 
     let mut status_updates = language_registry.language_server_binary_statuses();
@@ -596,6 +604,7 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
             indexed_docs_registry,
             snippet_registry,
             language_registry.clone(),
+            context_server_factory_registry.clone(),
             cx,
         );
         ExtensionStore::new(

crates/extensions_ui/src/extensions_ui.rs

@@ -1,4 +1,5 @@
 mod components;
+mod extension_context_server;
 mod extension_indexed_docs_provider;
 mod extension_registration_hooks;
 mod extension_slash_command;