context_store: Refactor state management (#29910)

Because we instantiated `ContextServerManager` both in `agent` and
`assistant-context-editor`, and these two entities track the running MCP
servers separately, we were effectively running every MCP server twice.

This PR moves the `ContextServerManager` into the project crate (now
called `ContextServerStore`). The store can be accessed via a project
instance. This ensures that we only instantiate one `ContextServerStore`
per project.

Also, this PR adds a bunch of tests to ensure that the
`ContextServerStore` behaves correctly (Previously there were none).

Closes #28714
Closes #29530

Release Notes:

- N/A

crates/project/src/context_server_store/extension.rs

@@ -0,0 +1,121 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use context_server::ContextServerCommand;
+use extension::{
+    ContextServerConfiguration, Extension, ExtensionContextServerProxy, ExtensionHostProxy,
+    ProjectDelegate,
+};
+use gpui::{App, AsyncApp, Entity, Task};
+
+use crate::worktree_store::WorktreeStore;
+
+use super::registry::{self, ContextServerDescriptorRegistry};
+
+pub fn init(cx: &mut App) {
+    let proxy = ExtensionHostProxy::default_global(cx);
+    proxy.register_context_server_proxy(ContextServerDescriptorRegistryProxy {
+        context_server_factory_registry: ContextServerDescriptorRegistry::default_global(cx),
+    });
+}
+
+struct ExtensionProject {
+    worktree_ids: Vec<u64>,
+}
+
+impl ProjectDelegate for ExtensionProject {
+    fn worktree_ids(&self) -> Vec<u64> {
+        self.worktree_ids.clone()
+    }
+}
+
+struct ContextServerDescriptor {
+    id: Arc<str>,
+    extension: Arc<dyn Extension>,
+}
+
+fn extension_project(
+    worktree_store: Entity<WorktreeStore>,
+    cx: &mut AsyncApp,
+) -> Result<Arc<ExtensionProject>> {
+    worktree_store.update(cx, |worktree_store, cx| {
+        Arc::new(ExtensionProject {
+            worktree_ids: worktree_store
+                .visible_worktrees(cx)
+                .map(|worktree| worktree.read(cx).id().to_proto())
+                .collect(),
+        })
+    })
+}
+
+impl registry::ContextServerDescriptor for ContextServerDescriptor {
+    fn command(
+        &self,
+        worktree_store: Entity<WorktreeStore>,
+        cx: &AsyncApp,
+    ) -> Task<Result<ContextServerCommand>> {
+        let id = self.id.clone();
+        let extension = self.extension.clone();
+        cx.spawn(async move |cx| {
+            let extension_project = extension_project(worktree_store, cx)?;
+            let mut command = extension
+                .context_server_command(id.clone(), extension_project.clone())
+                .await?;
+            command.command = extension
+                .path_from_extension(command.command.as_ref())
+                .to_string_lossy()
+                .to_string();
+
+            log::info!("loaded command for context server {id}: {command:?}");
+
+            Ok(ContextServerCommand {
+                path: command.command,
+                args: command.args,
+                env: Some(command.env.into_iter().collect()),
+            })
+        })
+    }
+
+    fn configuration(
+        &self,
+        worktree_store: Entity<WorktreeStore>,
+        cx: &AsyncApp,
+    ) -> Task<Result<Option<ContextServerConfiguration>>> {
+        let id = self.id.clone();
+        let extension = self.extension.clone();
+        cx.spawn(async move |cx| {
+            let extension_project = extension_project(worktree_store, cx)?;
+            let configuration = extension
+                .context_server_configuration(id.clone(), extension_project)
+                .await?;
+
+            log::debug!("loaded configuration for context server {id}: {configuration:?}");
+
+            Ok(configuration)
+        })
+    }
+}
+
+struct ContextServerDescriptorRegistryProxy {
+    context_server_factory_registry: Entity<ContextServerDescriptorRegistry>,
+}
+
+impl ExtensionContextServerProxy for ContextServerDescriptorRegistryProxy {
+    fn register_context_server(&self, extension: Arc<dyn Extension>, id: Arc<str>, cx: &mut App) {
+        self.context_server_factory_registry
+            .update(cx, |registry, _| {
+                registry.register_context_server_descriptor(
+                    id.clone(),
+                    Arc::new(ContextServerDescriptor { id, extension })
+                        as Arc<dyn registry::ContextServerDescriptor>,
+                )
+            });
+    }
+
+    fn unregister_context_server(&self, server_id: Arc<str>, cx: &mut App) {
+        self.context_server_factory_registry
+            .update(cx, |registry, _| {
+                registry.unregister_context_server_descriptor_by_id(&server_id)
+            });
+    }
+}

crates/project/src/context_server_store/registry.rs

@@ -0,0 +1,77 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use collections::HashMap;
+use context_server::ContextServerCommand;
+use extension::ContextServerConfiguration;
+use gpui::{App, AppContext as _, AsyncApp, Entity, Global, Task};
+
+use crate::worktree_store::WorktreeStore;
+
+pub trait ContextServerDescriptor {
+    fn command(
+        &self,
+        worktree_store: Entity<WorktreeStore>,
+        cx: &AsyncApp,
+    ) -> Task<Result<ContextServerCommand>>;
+    fn configuration(
+        &self,
+        worktree_store: Entity<WorktreeStore>,
+        cx: &AsyncApp,
+    ) -> Task<Result<Option<ContextServerConfiguration>>>;
+}
+
+struct GlobalContextServerDescriptorRegistry(Entity<ContextServerDescriptorRegistry>);
+
+impl Global for GlobalContextServerDescriptorRegistry {}
+
+#[derive(Default)]
+pub struct ContextServerDescriptorRegistry {
+    context_servers: HashMap<Arc<str>, Arc<dyn ContextServerDescriptor>>,
+}
+
+impl ContextServerDescriptorRegistry {
+    /// Returns the global [`ContextServerDescriptorRegistry`].
+    ///
+    /// Inserts a default [`ContextServerDescriptorRegistry`] if one does not yet exist.
+    pub fn default_global(cx: &mut App) -> Entity<Self> {
+        if !cx.has_global::<GlobalContextServerDescriptorRegistry>() {
+            let registry = cx.new(|_| Self::new());
+            cx.set_global(GlobalContextServerDescriptorRegistry(registry));
+        }
+        cx.global::<GlobalContextServerDescriptorRegistry>()
+            .0
+            .clone()
+    }
+
+    pub fn new() -> Self {
+        Self {
+            context_servers: HashMap::default(),
+        }
+    }
+
+    pub fn context_server_descriptors(&self) -> Vec<(Arc<str>, Arc<dyn ContextServerDescriptor>)> {
+        self.context_servers
+            .iter()
+            .map(|(id, factory)| (id.clone(), factory.clone()))
+            .collect()
+    }
+
+    pub fn context_server_descriptor(&self, id: &str) -> Option<Arc<dyn ContextServerDescriptor>> {
+        self.context_servers.get(id).cloned()
+    }
+
+    /// Registers the provided [`ContextServerDescriptor`].
+    pub fn register_context_server_descriptor(
+        &mut self,
+        id: Arc<str>,
+        descriptor: Arc<dyn ContextServerDescriptor>,
+    ) {
+        self.context_servers.insert(id, descriptor);
+    }
+
+    /// Unregisters the [`ContextServerDescriptor`] for the server with the given ID.
+    pub fn unregister_context_server_descriptor_by_id(&mut self, server_id: &str) {
+        self.context_servers.remove(server_id);
+    }
+}