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/context_servers/src/context_servers.rs

@@ -1,13 +1,16 @@
+pub mod client;
+pub mod manager;
+pub mod protocol;
+mod registry;
+pub mod types;
+
 use command_palette_hooks::CommandPaletteFilter;
 use gpui::{actions, AppContext};
 use settings::Settings;
 
+pub use crate::manager::ContextServer;
 use crate::manager::ContextServerSettings;
-
-pub mod client;
-pub mod manager;
-pub mod protocol;
-pub mod types;
+pub use crate::registry::ContextServerFactoryRegistry;
 
 actions!(context_servers, [Restart]);
 
@@ -16,6 +19,7 @@ pub const CONTEXT_SERVERS_NAMESPACE: &'static str = "context_servers";
 
 pub fn init(cx: &mut AppContext) {
     ContextServerSettings::register(cx);
+    ContextServerFactoryRegistry::default_global(cx);
 
     CommandPaletteFilter::update_global(cx, |filter, _cx| {
         filter.hide_namespace(CONTEXT_SERVERS_NAMESPACE);

crates/context_servers/src/manager.rs

@@ -79,7 +79,7 @@ pub struct NativeContextServer {
 }
 
 impl NativeContextServer {
-    fn new(config: Arc<ServerConfig>) -> Self {
+    pub fn new(config: Arc<ServerConfig>) -> Self {
         Self {
             id: config.id.clone().into(),
             config,
@@ -151,13 +151,13 @@ impl ContextServer for NativeContextServer {
 /// must go through the `GlobalContextServerManager` which holds
 /// a model to the ContextServerManager.
 pub struct ContextServerManager {
-    servers: HashMap<String, Arc<dyn ContextServer>>,
-    pending_servers: HashSet<String>,
+    servers: HashMap<Arc<str>, Arc<dyn ContextServer>>,
+    pending_servers: HashSet<Arc<str>>,
 }
 
 pub enum Event {
-    ServerStarted { server_id: String },
-    ServerStopped { server_id: String },
+    ServerStarted { server_id: Arc<str> },
+    ServerStopped { server_id: Arc<str> },
 }
 
 impl EventEmitter<Event> for ContextServerManager {}
@@ -178,10 +178,10 @@ impl ContextServerManager {
 
     pub fn add_server(
         &mut self,
-        config: Arc<ServerConfig>,
+        server: Arc<dyn ContextServer>,
         cx: &ModelContext<Self>,
     ) -> Task<anyhow::Result<()>> {
-        let server_id = config.id.clone();
+        let server_id = server.id();
 
         if self.servers.contains_key(&server_id) || self.pending_servers.contains(&server_id) {
             return Task::ready(Ok(()));
@@ -190,7 +190,6 @@ impl ContextServerManager {
         let task = {
             let server_id = server_id.clone();
             cx.spawn(|this, mut cx| async move {
-                let server = Arc::new(NativeContextServer::new(config));
                 server.clone().start(&cx).await?;
                 this.update(&mut cx, |this, cx| {
                     this.servers.insert(server_id.clone(), server);
@@ -211,14 +210,20 @@ impl ContextServerManager {
         self.servers.get(id).cloned()
     }
 
-    pub fn remove_server(&mut self, id: &str, cx: &ModelContext<Self>) -> Task<anyhow::Result<()>> {
-        let id = id.to_string();
+    pub fn remove_server(
+        &mut self,
+        id: &Arc<str>,
+        cx: &ModelContext<Self>,
+    ) -> Task<anyhow::Result<()>> {
+        let id = id.clone();
         cx.spawn(|this, mut cx| async move {
-            if let Some(server) = this.update(&mut cx, |this, _cx| this.servers.remove(&id))? {
+            if let Some(server) =
+                this.update(&mut cx, |this, _cx| this.servers.remove(id.as_ref()))?
+            {
                 server.stop()?;
             }
             this.update(&mut cx, |this, cx| {
-                this.pending_servers.remove(&id);
+                this.pending_servers.remove(id.as_ref());
                 cx.emit(Event::ServerStopped {
                     server_id: id.clone(),
                 })
@@ -232,7 +237,7 @@ impl ContextServerManager {
         id: &Arc<str>,
         cx: &mut ModelContext<Self>,
     ) -> Task<anyhow::Result<()>> {
-        let id = id.to_string();
+        let id = id.clone();
         cx.spawn(|this, mut cx| async move {
             if let Some(server) = this.update(&mut cx, |this, _cx| this.servers.remove(&id))? {
                 server.stop()?;
@@ -284,7 +289,8 @@ impl ContextServerManager {
 
         log::trace!("servers_to_add={:?}", servers_to_add);
         for config in servers_to_add {
-            self.add_server(Arc::new(config), cx).detach_and_log_err(cx);
+            let server = Arc::new(NativeContextServer::new(Arc::new(config)));
+            self.add_server(server, cx).detach_and_log_err(cx);
         }
 
         for id in servers_to_remove {

crates/context_servers/src/registry.rs

@@ -0,0 +1,72 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use collections::HashMap;
+use gpui::{AppContext, AsyncAppContext, ReadGlobal};
+use gpui::{Global, Task};
+use parking_lot::RwLock;
+
+use crate::ContextServer;
+
+pub type ContextServerFactory =
+    Arc<dyn Fn(&AsyncAppContext) -> Task<Result<Arc<dyn ContextServer>>> + Send + Sync + 'static>;
+
+#[derive(Default)]
+struct GlobalContextServerFactoryRegistry(Arc<ContextServerFactoryRegistry>);
+
+impl Global for GlobalContextServerFactoryRegistry {}
+
+#[derive(Default)]
+struct ContextServerFactoryRegistryState {
+    context_servers: HashMap<Arc<str>, ContextServerFactory>,
+}
+
+#[derive(Default)]
+pub struct ContextServerFactoryRegistry {
+    state: RwLock<ContextServerFactoryRegistryState>,
+}
+
+impl ContextServerFactoryRegistry {
+    /// Returns the global [`ContextServerFactoryRegistry`].
+    pub fn global(cx: &AppContext) -> Arc<Self> {
+        GlobalContextServerFactoryRegistry::global(cx).0.clone()
+    }
+
+    /// Returns the global [`ContextServerFactoryRegistry`].
+    ///
+    /// Inserts a default [`ContextServerFactoryRegistry`] if one does not yet exist.
+    pub fn default_global(cx: &mut AppContext) -> Arc<Self> {
+        cx.default_global::<GlobalContextServerFactoryRegistry>()
+            .0
+            .clone()
+    }
+
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            state: RwLock::new(ContextServerFactoryRegistryState {
+                context_servers: HashMap::default(),
+            }),
+        })
+    }
+
+    pub fn context_server_factories(&self) -> Vec<(Arc<str>, ContextServerFactory)> {
+        self.state
+            .read()
+            .context_servers
+            .iter()
+            .map(|(id, factory)| (id.clone(), factory.clone()))
+            .collect()
+    }
+
+    /// Registers the provided [`ContextServerFactory`].
+    pub fn register_server_factory(&self, id: Arc<str>, factory: ContextServerFactory) {
+        let mut state = self.state.write();
+        state.context_servers.insert(id, factory);
+    }
+
+    /// Unregisters the [`ContextServerFactory`] for the server with the given ID.
+    pub fn unregister_server_factory_by_id(&self, server_id: &str) {
+        let mut state = self.state.write();
+        state.context_servers.remove(server_id);
+    }
+}