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
2025-05-05 21:36:12 +02:00
parent 8199664a5a
commit 9cb5ffac25
43 changed files with 1568 additions and 1047 deletions
@@ -3487,7 +3487,6 @@ fn open_editor_at_position(
 #[cfg(test)]
 mod tests {
    use assistant_tool::{ToolRegistry, ToolWorkingSet};
-    use context_server::ContextServerSettings;
    use editor::EditorSettings;
    use fs::FakeFs;
    use gpui::{TestAppContext, VisualTestContext};
@@ -3559,7 +3558,6 @@ mod tests {
            workspace::init_settings(cx);
            language_model::init_settings(cx);
            ThemeSettings::register(cx);
-            ContextServerSettings::register(cx);
            EditorSettings::register(cx);
            ToolRegistry::default_global(cx);
        });
@@ -1748,7 +1748,6 @@ mod tests {
    use crate::{Keep, ThreadStore, thread_store};
    use assistant_settings::AssistantSettings;
    use assistant_tool::ToolWorkingSet;
-    use context_server::ContextServerSettings;
    use editor::EditorSettings;
    use gpui::{TestAppContext, UpdateGlobal, VisualTestContext};
    use project::{FakeFs, Project};
@@ -1771,7 +1770,6 @@ mod tests {
            thread_store::init(cx);
            workspace::init_settings(cx);
            ThemeSettings::register(cx);
-            ContextServerSettings::register(cx);
            EditorSettings::register(cx);
            language_model::init_settings(cx);
        });
@@ -1928,7 +1926,6 @@ mod tests {
            thread_store::init(cx);
            workspace::init_settings(cx);
            ThemeSettings::register(cx);
-            ContextServerSettings::register(cx);
            EditorSettings::register(cx);
            language_model::init_settings(cx);
            workspace::register_project_item::<Editor>(cx);
@@ -7,6 +7,7 @@ mod buffer_codegen;
 mod context;
 mod context_picker;
 mod context_server_configuration;
+mod context_server_tool;
 mod context_store;
 mod context_strip;
 mod debug;
@@ -8,13 +8,14 @@ use std::{sync::Arc, time::Duration};
 use assistant_settings::AssistantSettings;
 use assistant_tool::{ToolSource, ToolWorkingSet};
 use collections::HashMap;
-use context_server::manager::{ContextServer, ContextServerManager, ContextServerStatus};
+use context_server::ContextServerId;
 use fs::Fs;
 use gpui::{
    Action, Animation, AnimationExt as _, AnyView, App, Entity, EventEmitter, FocusHandle,
    Focusable, ScrollHandle, Subscription, pulsating_between,
 };
 use language_model::{LanguageModelProvider, LanguageModelProviderId, LanguageModelRegistry};
+use project::context_server_store::{ContextServerStatus, ContextServerStore};
 use settings::{Settings, update_settings_file};
 use ui::{
    Disclosure, Divider, DividerColor, ElevationIndex, Indicator, Scrollbar, ScrollbarState,
@@ -33,8 +34,8 @@ pub struct AssistantConfiguration {
    fs: Arc<dyn Fs>,
    focus_handle: FocusHandle,
    configuration_views_by_provider: HashMap<LanguageModelProviderId, AnyView>,
-    context_server_manager: Entity<ContextServerManager>,
-    expanded_context_server_tools: HashMap<Arc<str>, bool>,
+    context_server_store: Entity<ContextServerStore>,
+    expanded_context_server_tools: HashMap<ContextServerId, bool>,
    tools: Entity<ToolWorkingSet>,
    _registry_subscription: Subscription,
    scroll_handle: ScrollHandle,
@@ -44,7 +45,7 @@ pub struct AssistantConfiguration {
 impl AssistantConfiguration {
    pub fn new(
        fs: Arc<dyn Fs>,
-        context_server_manager: Entity<ContextServerManager>,
+        context_server_store: Entity<ContextServerStore>,
        tools: Entity<ToolWorkingSet>,
        window: &mut Window,
        cx: &mut Context<Self>,
@@ -75,7 +76,7 @@ impl AssistantConfiguration {
            fs,
            focus_handle,
            configuration_views_by_provider: HashMap::default(),
-            context_server_manager,
+            context_server_store,
            expanded_context_server_tools: HashMap::default(),
            tools,
            _registry_subscription: registry_subscription,
@@ -306,7 +307,7 @@ impl AssistantConfiguration {
        window: &mut Window,
        cx: &mut Context<Self>,
    ) -> impl IntoElement {
-        let context_servers = self.context_server_manager.read(cx).all_servers().clone();
+        let context_server_ids = self.context_server_store.read(cx).all_server_ids().clone();

        const SUBHEADING: &str = "Connect to context servers via the Model Context Protocol either via Zed extensions or directly.";

@@ -322,9 +323,9 @@ impl AssistantConfiguration {
                    .child(Label::new(SUBHEADING).color(Color::Muted)),
            )
            .children(
-                context_servers
-                    .into_iter()
-                    .map(|context_server| self.render_context_server(context_server, window, cx)),
+                context_server_ids.into_iter().map(|context_server_id| {
+                    self.render_context_server(context_server_id, window, cx)
+                }),
            )
            .child(
                h_flex()
@@ -374,19 +375,20 @@ impl AssistantConfiguration {

    fn render_context_server(
        &self,
-        context_server: Arc<ContextServer>,
+        context_server_id: ContextServerId,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) -> impl use<> + IntoElement {
        let tools_by_source = self.tools.read(cx).tools_by_source(cx);
        let server_status = self
-            .context_server_manager
+            .context_server_store
            .read(cx)
-            .status_for_server(&context_server.id());
+            .status_for_server(&context_server_id)
+            .unwrap_or(ContextServerStatus::Stopped);

-        let is_running = matches!(server_status, Some(ContextServerStatus::Running));
+        let is_running = matches!(server_status, ContextServerStatus::Running);

-        let error = if let Some(ContextServerStatus::Error(error)) = server_status.clone() {
+        let error = if let ContextServerStatus::Error(error) = server_status.clone() {
            Some(error)
        } else {
            None
@@ -394,13 +396,13 @@ impl AssistantConfiguration {

        let are_tools_expanded = self
            .expanded_context_server_tools
-            .get(&context_server.id())
+            .get(&context_server_id)
            .copied()
            .unwrap_or_default();

        let tools = tools_by_source
            .get(&ToolSource::ContextServer {
-                id: context_server.id().into(),
+                id: context_server_id.0.clone().into(),
            })
            .map_or([].as_slice(), |tools| tools.as_slice());
        let tool_count = tools.len();
@@ -408,7 +410,7 @@ impl AssistantConfiguration {
        let border_color = cx.theme().colors().border.opacity(0.6);

        v_flex()
-            .id(SharedString::from(context_server.id()))
+            .id(SharedString::from(context_server_id.0.clone()))
            .border_1()
            .rounded_md()
            .border_color(border_color)
@@ -432,7 +434,7 @@ impl AssistantConfiguration {
                                )
                                .disabled(tool_count == 0)
                                .on_click(cx.listener({
-                                    let context_server_id = context_server.id();
+                                    let context_server_id = context_server_id.clone();
                                    move |this, _event, _window, _cx| {
                                        let is_open = this
                                            .expanded_context_server_tools
@@ -444,14 +446,14 @@ impl AssistantConfiguration {
                                })),
                            )
                            .child(match server_status {
-                                Some(ContextServerStatus::Starting) => {
+                                ContextServerStatus::Starting => {
                                    let color = Color::Success.color(cx);
                                    Indicator::dot()
                                        .color(Color::Success)
                                        .with_animation(
                                            SharedString::from(format!(
                                                "{}-starting",
-                                                context_server.id(),
+                                                context_server_id.0.clone(),
                                            )),
                                            Animation::new(Duration::from_secs(2))
                                                .repeat()
@@ -462,15 +464,17 @@ impl AssistantConfiguration {
                                        )
                                        .into_any_element()
                                }
-                                Some(ContextServerStatus::Running) => {
+                                ContextServerStatus::Running => {
                                    Indicator::dot().color(Color::Success).into_any_element()
                                }
-                                Some(ContextServerStatus::Error(_)) => {
+                                ContextServerStatus::Error(_) => {
                                    Indicator::dot().color(Color::Error).into_any_element()
                                }
-                                None => Indicator::dot().color(Color::Muted).into_any_element(),
+                                ContextServerStatus::Stopped => {
+                                    Indicator::dot().color(Color::Muted).into_any_element()
+                                }
                            })
-                            .child(Label::new(context_server.id()).ml_0p5())
+                            .child(Label::new(context_server_id.0.clone()).ml_0p5())
                            .when(is_running, |this| {
                                this.child(
                                    Label::new(if tool_count == 1 {
@@ -487,32 +491,22 @@ impl AssistantConfiguration {
                        Switch::new("context-server-switch", is_running.into())
                            .color(SwitchColor::Accent)
                            .on_click({
-                                let context_server_manager = self.context_server_manager.clone();
-                                let context_server = context_server.clone();
+                                let context_server_manager = self.context_server_store.clone();
+                                let context_server_id = context_server_id.clone();
                                move |state, _window, cx| match state {
                                    ToggleState::Unselected | ToggleState::Indeterminate => {
                                        context_server_manager.update(cx, |this, cx| {
-                                            this.stop_server(context_server.clone(), cx).log_err();
+                                            this.stop_server(&context_server_id, cx).log_err();
                                        });
                                    }
                                    ToggleState::Selected => {
-                                        cx.spawn({
-                                            let context_server_manager =
-                                                context_server_manager.clone();
-                                            let context_server = context_server.clone();
-                                            async move |cx| {
-                                                if let Some(start_server_task) =
-                                                    context_server_manager
-                                                        .update(cx, |this, cx| {
-                                                            this.start_server(context_server, cx)
-                                                        })
-                                                        .log_err()
-                                                {
-                                                    start_server_task.await.log_err();
-                                                }
+                                        context_server_manager.update(cx, |this, cx| {
+                                            if let Some(server) =
+                                                this.get_server(&context_server_id)
+                                            {
+                                                this.start_server(server, cx).log_err();
                                            }
                                        })
-                                        .detach();
                                    }
                                }
                            }),
@@ -1,5 +1,6 @@
-use context_server::{ContextServerSettings, ServerCommand, ServerConfig};
+use context_server::ContextServerCommand;
 use gpui::{DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, WeakEntity, prelude::*};
+use project::project_settings::{ContextServerConfiguration, ProjectSettings};
 use serde_json::json;
 use settings::update_settings_file;
 use ui::{KeyBinding, Modal, ModalFooter, ModalHeader, Section, Tooltip, prelude::*};
@@ -77,11 +78,11 @@ impl AddContextServerModal {
        if let Some(workspace) = self.workspace.upgrade() {
            workspace.update(cx, |workspace, cx| {
                let fs = workspace.app_state().fs.clone();
-                update_settings_file::<ContextServerSettings>(fs.clone(), cx, |settings, _| {
+                update_settings_file::<ProjectSettings>(fs.clone(), cx, |settings, _| {
                    settings.context_servers.insert(
                        name.into(),
-                        ServerConfig {
-                            command: Some(ServerCommand {
+                        ContextServerConfiguration {
+                            command: Some(ContextServerCommand {
                                path,
                                args,
                                env: None,
@@ -4,9 +4,8 @@ use std::{
 };

 use anyhow::Context as _;
-use context_server::manager::{ContextServerManager, ContextServerStatus};
+use context_server::ContextServerId;
 use editor::{Editor, EditorElement, EditorStyle};
-use extension::ContextServerConfiguration;
 use gpui::{
    Animation, AnimationExt, App, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, Task,
    TextStyle, TextStyleRefinement, Transformation, UnderlineStyle, WeakEntity, percentage,
@@ -14,6 +13,10 @@ use gpui::{
 use language::{Language, LanguageRegistry};
 use markdown::{Markdown, MarkdownElement, MarkdownStyle};
 use notifications::status_toast::{StatusToast, ToastIcon};
+use project::{
+    context_server_store::{ContextServerStatus, ContextServerStore},
+    project_settings::{ContextServerConfiguration, ProjectSettings},
+};
 use settings::{Settings as _, update_settings_file};
 use theme::ThemeSettings;
 use ui::{KeyBinding, Modal, ModalFooter, ModalHeader, Section, prelude::*};
@@ -23,11 +26,11 @@ use workspace::{ModalView, Workspace};
 pub(crate) struct ConfigureContextServerModal {
    workspace: WeakEntity<Workspace>,
    context_servers_to_setup: Vec<ConfigureContextServer>,
-    context_server_manager: Entity<ContextServerManager>,
+    context_server_store: Entity<ContextServerStore>,
 }

 struct ConfigureContextServer {
-    id: Arc<str>,
+    id: ContextServerId,
    installation_instructions: Entity<markdown::Markdown>,
    settings_validator: Option<jsonschema::Validator>,
    settings_editor: Entity<Editor>,
@@ -37,9 +40,9 @@ struct ConfigureContextServer {

 impl ConfigureContextServerModal {
    pub fn new(
-        configurations: impl Iterator<Item = (Arc<str>, ContextServerConfiguration)>,
+        configurations: impl Iterator<Item = (ContextServerId, extension::ContextServerConfiguration)>,
+        context_server_store: Entity<ContextServerStore>,
        jsonc_language: Option<Arc<Language>>,
-        context_server_manager: Entity<ContextServerManager>,
        language_registry: Arc<LanguageRegistry>,
        workspace: WeakEntity<Workspace>,
        window: &mut Window,
@@ -85,7 +88,7 @@ impl ConfigureContextServerModal {
        Some(Self {
            workspace,
            context_servers_to_setup,
-            context_server_manager,
+            context_server_store,
        })
    }
 }
@@ -126,14 +129,14 @@ impl ConfigureContextServerModal {
        }
        let id = configuration.id.clone();

-        let settings_changed = context_server::ContextServerSettings::get_global(cx)
+        let settings_changed = ProjectSettings::get_global(cx)
            .context_servers
-            .get(&id)
+            .get(&id.0)
            .map_or(true, |config| {
                config.settings.as_ref() != Some(&settings_value)
            });

-        let is_running = self.context_server_manager.read(cx).status_for_server(&id)
+        let is_running = self.context_server_store.read(cx).status_for_server(&id)
            == Some(ContextServerStatus::Running);

        if !settings_changed && is_running {
@@ -143,7 +146,7 @@ impl ConfigureContextServerModal {

        configuration.waiting_for_context_server = true;

-        let task = wait_for_context_server(&self.context_server_manager, id.clone(), cx);
+        let task = wait_for_context_server(&self.context_server_store, id.clone(), cx);
        cx.spawn({
            let id = id.clone();
            async move |this, cx| {
@@ -167,29 +170,25 @@ impl ConfigureContextServerModal {
        .detach();

        // When we write the settings to the file, the context server will be restarted.
-        update_settings_file::<context_server::ContextServerSettings>(
-            workspace.read(cx).app_state().fs.clone(),
-            cx,
-            {
-                let id = id.clone();
-                |settings, _| {
-                    if let Some(server_config) = settings.context_servers.get_mut(&id) {
-                        server_config.settings = Some(settings_value);
-                    } else {
-                        settings.context_servers.insert(
-                            id,
-                            context_server::ServerConfig {
-                                settings: Some(settings_value),
-                                ..Default::default()
-                            },
-                        );
-                    }
+        update_settings_file::<ProjectSettings>(workspace.read(cx).app_state().fs.clone(), cx, {
+            let id = id.clone();
+            |settings, _| {
+                if let Some(server_config) = settings.context_servers.get_mut(&id.0) {
+                    server_config.settings = Some(settings_value);
+                } else {
+                    settings.context_servers.insert(
+                        id.0,
+                        ContextServerConfiguration {
+                            settings: Some(settings_value),
+                            ..Default::default()
+                        },
+                    );
                }
-            },
-        );
+            }
+        });
    }

-    fn complete_setup(&mut self, id: Arc<str>, cx: &mut Context<Self>) {
+    fn complete_setup(&mut self, id: ContextServerId, cx: &mut Context<Self>) {
        self.context_servers_to_setup.remove(0);
        cx.notify();

@@ -223,31 +222,40 @@ impl ConfigureContextServerModal {
 }

 fn wait_for_context_server(
-    context_server_manager: &Entity<ContextServerManager>,
-    context_server_id: Arc<str>,
+    context_server_store: &Entity<ContextServerStore>,
+    context_server_id: ContextServerId,
    cx: &mut App,
 ) -> Task<Result<(), Arc<str>>> {
    let (tx, rx) = futures::channel::oneshot::channel();
    let tx = Arc::new(Mutex::new(Some(tx)));

-    let subscription = cx.subscribe(context_server_manager, move |_, event, _cx| match event {
-        context_server::manager::Event::ServerStatusChanged { server_id, status } => match status {
-            Some(ContextServerStatus::Running) => {
-                if server_id == &context_server_id {
-                    if let Some(tx) = tx.lock().unwrap().take() {
-                        let _ = tx.send(Ok(()));
+    let subscription = cx.subscribe(context_server_store, move |_, event, _cx| match event {
+        project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
+            match status {
+                ContextServerStatus::Running => {
+                    if server_id == &context_server_id {
+                        if let Some(tx) = tx.lock().unwrap().take() {
+                            let _ = tx.send(Ok(()));
+                        }
                    }
                }
-            }
-            Some(ContextServerStatus::Error(error)) => {
-                if server_id == &context_server_id {
-                    if let Some(tx) = tx.lock().unwrap().take() {
-                        let _ = tx.send(Err(error.clone()));
+                ContextServerStatus::Stopped => {
+                    if server_id == &context_server_id {
+                        if let Some(tx) = tx.lock().unwrap().take() {
+                            let _ = tx.send(Err("Context server stopped running".into()));
+                        }
                    }
                }
+                ContextServerStatus::Error(error) => {
+                    if server_id == &context_server_id {
+                        if let Some(tx) = tx.lock().unwrap().take() {
+                            let _ = tx.send(Err(error.clone()));
+                        }
+                    }
+                }
+                _ => {}
            }
-            _ => {}
-        },
+        }
    });

    cx.spawn(async move |_cx| {
@@ -1026,14 +1026,14 @@ impl AssistantPanel {
    }

    pub(crate) fn open_configuration(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        let context_server_manager = self.thread_store.read(cx).context_server_manager();
+        let context_server_store = self.project.read(cx).context_server_store();
        let tools = self.thread_store.read(cx).tools();
        let fs = self.fs.clone();

        self.set_active_view(ActiveView::Configuration, window, cx);
        self.configuration =
            Some(cx.new(|cx| {
-                AssistantConfiguration::new(fs, context_server_manager, tools, window, cx)
+                AssistantConfiguration::new(fs, context_server_store, tools, window, cx)
            }));

        if let Some(configuration) = self.configuration.as_ref() {
@@ -1095,9 +1095,7 @@ mod tests {

    #[gpui::test(iterations = 10)]
    async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
-        cx.set_global(cx.update(SettingsStore::test));
-        cx.update(language_model::LanguageModelRegistry::test);
-        cx.update(language_settings::init);
+        init_test(cx);

        let text = indoc! {"
            fn main() {
@@ -1167,8 +1165,7 @@ mod tests {
        cx: &mut TestAppContext,
        mut rng: StdRng,
    ) {
-        cx.set_global(cx.update(SettingsStore::test));
-        cx.update(language_settings::init);
+        init_test(cx);

        let text = indoc! {"
            fn main() {
@@ -1237,9 +1234,7 @@ mod tests {
        cx: &mut TestAppContext,
        mut rng: StdRng,
    ) {
-        cx.update(LanguageModelRegistry::test);
-        cx.set_global(cx.update(SettingsStore::test));
-        cx.update(language_settings::init);
+        init_test(cx);

        let text = concat!(
            "fn main() {\n",
@@ -1305,9 +1300,7 @@ mod tests {

    #[gpui::test(iterations = 10)]
    async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
-        cx.update(LanguageModelRegistry::test);
-        cx.set_global(cx.update(SettingsStore::test));
-        cx.update(language_settings::init);
+        init_test(cx);

        let text = indoc! {"
            func main() {
@@ -1367,9 +1360,7 @@ mod tests {

    #[gpui::test]
    async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
-        cx.update(LanguageModelRegistry::test);
-        cx.set_global(cx.update(SettingsStore::test));
-        cx.update(language_settings::init);
+        init_test(cx);

        let text = indoc! {"
            fn main() {
@@ -1473,6 +1464,13 @@ mod tests {
        }
    }

+    fn init_test(cx: &mut TestAppContext) {
+        cx.update(LanguageModelRegistry::test);
+        cx.set_global(cx.update(SettingsStore::test));
+        cx.update(Project::init_settings);
+        cx.update(language_settings::init);
+    }
+
    fn simulate_response_stream(
        codegen: Entity<CodegenAlternative>,
        cx: &mut TestAppContext,
@@ -1,14 +1,15 @@
 use std::sync::Arc;

 use anyhow::Context as _;
-use context_server::ContextServerDescriptorRegistry;
+use context_server::ContextServerId;
 use extension::ExtensionManifest;
 use language::LanguageRegistry;
+use project::context_server_store::registry::ContextServerDescriptorRegistry;
 use ui::prelude::*;
 use util::ResultExt;
 use workspace::Workspace;

-use crate::{AssistantPanel, assistant_configuration::ConfigureContextServerModal};
+use crate::assistant_configuration::ConfigureContextServerModal;

 pub(crate) fn init(language_registry: Arc<LanguageRegistry>, cx: &mut App) {
    cx.observe_new(move |_: &mut Workspace, window, cx| {
@@ -60,18 +61,10 @@ fn show_configure_mcp_modal(
    window: &mut Window,
    cx: &mut Context<'_, Workspace>,
 ) {
-    let Some(context_server_manager) = workspace.panel::<AssistantPanel>(cx).map(|panel| {
-        panel
-            .read(cx)
-            .thread_store()
-            .read(cx)
-            .context_server_manager()
-    }) else {
-        return;
-    };
+    let context_server_store = workspace.project().read(cx).context_server_store();

-    let registry = ContextServerDescriptorRegistry::global(cx).read(cx);
-    let project = workspace.project().clone();
+    let registry = ContextServerDescriptorRegistry::default_global(cx).read(cx);
+    let worktree_store = workspace.project().read(cx).worktree_store();
    let configuration_tasks = manifest
        .context_servers
        .keys()
@@ -80,15 +73,15 @@ fn show_configure_mcp_modal(
            |key| {
                let descriptor = registry.context_server_descriptor(&key)?;
                Some(cx.spawn({
-                    let project = project.clone();
+                    let worktree_store = worktree_store.clone();
                    async move |_, cx| {
                        descriptor
-                            .configuration(project, &cx)
+                            .configuration(worktree_store.clone(), &cx)
                            .await
                            .context("Failed to resolve context server configuration")
                            .log_err()
                            .flatten()
-                            .map(|config| (key, config))
+                            .map(|config| (ContextServerId(key), config))
                    }
                }))
            }
@@ -104,8 +97,8 @@ fn show_configure_mcp_modal(
        this.update_in(cx, |this, window, cx| {
            let modal = ConfigureContextServerModal::new(
                descriptors.into_iter().flatten(),
+                context_server_store,
                jsonc_language,
-                context_server_manager,
                language_registry,
                cx.entity().downgrade(),
                window,
@@ -0,0 +1,125 @@
+use std::sync::Arc;
+
+use anyhow::{Result, anyhow, bail};
+use assistant_tool::{ActionLog, Tool, ToolResult, ToolSource};
+use context_server::{ContextServerId, types};
+use gpui::{AnyWindowHandle, App, Entity, Task};
+use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
+use project::{Project, context_server_store::ContextServerStore};
+use ui::IconName;
+
+pub struct ContextServerTool {
+    store: Entity<ContextServerStore>,
+    server_id: ContextServerId,
+    tool: types::Tool,
+}
+
+impl ContextServerTool {
+    pub fn new(
+        store: Entity<ContextServerStore>,
+        server_id: ContextServerId,
+        tool: types::Tool,
+    ) -> Self {
+        Self {
+            store,
+            server_id,
+            tool,
+        }
+    }
+}
+
+impl Tool for ContextServerTool {
+    fn name(&self) -> String {
+        self.tool.name.clone()
+    }
+
+    fn description(&self) -> String {
+        self.tool.description.clone().unwrap_or_default()
+    }
+
+    fn icon(&self) -> IconName {
+        IconName::Cog
+    }
+
+    fn source(&self) -> ToolSource {
+        ToolSource::ContextServer {
+            id: self.server_id.clone().0.into(),
+        }
+    }
+
+    fn needs_confirmation(&self, _: &serde_json::Value, _: &App) -> bool {
+        true
+    }
+
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
+        let mut schema = self.tool.input_schema.clone();
+        assistant_tool::adapt_schema_to_format(&mut schema, format)?;
+        Ok(match schema {
+            serde_json::Value::Null => {
+                serde_json::json!({ "type": "object", "properties": [] })
+            }
+            serde_json::Value::Object(map) if map.is_empty() => {
+                serde_json::json!({ "type": "object", "properties": [] })
+            }
+            _ => schema,
+        })
+    }
+
+    fn ui_text(&self, _input: &serde_json::Value) -> String {
+        format!("Run MCP tool `{}`", self.tool.name)
+    }
+
+    fn run(
+        self: Arc<Self>,
+        input: serde_json::Value,
+        _messages: &[LanguageModelRequestMessage],
+        _project: Entity<Project>,
+        _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
+        cx: &mut App,
+    ) -> ToolResult {
+        if let Some(server) = self.store.read(cx).get_running_server(&self.server_id) {
+            let tool_name = self.tool.name.clone();
+            let server_clone = server.clone();
+            let input_clone = input.clone();
+
+            cx.spawn(async move |_cx| {
+                let Some(protocol) = server_clone.client() else {
+                    bail!("Context server not initialized");
+                };
+
+                let arguments = if let serde_json::Value::Object(map) = input_clone {
+                    Some(map.into_iter().collect())
+                } else {
+                    None
+                };
+
+                log::trace!(
+                    "Running tool: {} with arguments: {:?}",
+                    tool_name,
+                    arguments
+                );
+                let response = protocol.run_tool(tool_name, arguments).await?;
+
+                let mut result = String::new();
+                for content in response.content {
+                    match content {
+                        types::ToolResponseContent::Text { text } => {
+                            result.push_str(&text);
+                        }
+                        types::ToolResponseContent::Image { .. } => {
+                            log::warn!("Ignoring image content from tool response");
+                        }
+                        types::ToolResponseContent::Resource { .. } => {
+                            log::warn!("Ignoring resource content from tool response");
+                        }
+                    }
+                }
+                Ok(result)
+            })
+            .into()
+        } else {
+            Task::ready(Err(anyhow!("Context server not found"))).into()
+        }
+    }
+}
@@ -2660,7 +2660,6 @@ mod tests {
    use crate::{ThreadStore, context::load_context, context_store::ContextStore, thread_store};
    use assistant_settings::AssistantSettings;
    use assistant_tool::ToolRegistry;
-    use context_server::ContextServerSettings;
    use editor::EditorSettings;
    use gpui::TestAppContext;
    use language_model::fake_provider::FakeLanguageModel;
@@ -3082,7 +3081,6 @@ fn main() {{
            workspace::init_settings(cx);
            language_model::init_settings(cx);
            ThemeSettings::register(cx);
-            ContextServerSettings::register(cx);
            EditorSettings::register(cx);
            ToolRegistry::default_global(cx);
        });
@@ -9,8 +9,7 @@ use assistant_settings::{AgentProfile, AgentProfileId, AssistantSettings};
 use assistant_tool::{ToolId, ToolSource, ToolWorkingSet};
 use chrono::{DateTime, Utc};
 use collections::HashMap;
-use context_server::manager::{ContextServerManager, ContextServerStatus};
-use context_server::{ContextServerDescriptorRegistry, ContextServerTool};
+use context_server::ContextServerId;
 use futures::channel::{mpsc, oneshot};
 use futures::future::{self, BoxFuture, Shared};
 use futures::{FutureExt as _, StreamExt as _};
@@ -21,6 +20,7 @@ use gpui::{
 use heed::Database;
 use heed::types::SerdeBincode;
 use language_model::{LanguageModelToolUseId, Role, TokenUsage};
+use project::context_server_store::{ContextServerStatus, ContextServerStore};
 use project::{Project, ProjectItem, ProjectPath, Worktree};
 use prompt_store::{
    ProjectContext, PromptBuilder, PromptId, PromptStore, PromptsUpdatedEvent, RulesFileContext,
@@ -30,6 +30,7 @@ use serde::{Deserialize, Serialize};
 use settings::{Settings as _, SettingsStore};
 use util::ResultExt as _;

+use crate::context_server_tool::ContextServerTool;
 use crate::thread::{
    DetailedSummaryState, ExceededWindowError, MessageId, ProjectSnapshot, Thread, ThreadId,
 };
@@ -62,8 +63,7 @@ pub struct ThreadStore {
    tools: Entity<ToolWorkingSet>,
    prompt_builder: Arc<PromptBuilder>,
    prompt_store: Option<Entity<PromptStore>>,
-    context_server_manager: Entity<ContextServerManager>,
-    context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
+    context_server_tool_ids: HashMap<ContextServerId, Vec<ToolId>>,
    threads: Vec<SerializedThreadMetadata>,
    project_context: SharedProjectContext,
    reload_system_prompt_tx: mpsc::Sender<()>,
@@ -108,11 +108,6 @@ impl ThreadStore {
        prompt_store: Option<Entity<PromptStore>>,
        cx: &mut Context<Self>,
    ) -> (Self, oneshot::Receiver<()>) {
-        let context_server_factory_registry = ContextServerDescriptorRegistry::default_global(cx);
-        let context_server_manager = cx.new(|cx| {
-            ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
-        });
-
        let mut subscriptions = vec![
            cx.observe_global::<SettingsStore>(move |this: &mut Self, cx| {
                this.load_default_profile(cx);
@@ -159,7 +154,6 @@ impl ThreadStore {
            tools,
            prompt_builder,
            prompt_store,
-            context_server_manager,
            context_server_tool_ids: HashMap::default(),
            threads: Vec::new(),
            project_context: SharedProjectContext::default(),
@@ -354,10 +348,6 @@ impl ThreadStore {
        })
    }

-    pub fn context_server_manager(&self) -> Entity<ContextServerManager> {
-        self.context_server_manager.clone()
-    }
-
    pub fn prompt_store(&self) -> &Option<Entity<PromptStore>> {
        &self.prompt_store
    }
@@ -494,11 +484,17 @@ impl ThreadStore {
        });

        if profile.enable_all_context_servers {
-            for context_server in self.context_server_manager.read(cx).all_servers() {
+            for context_server_id in self
+                .project
+                .read(cx)
+                .context_server_store()
+                .read(cx)
+                .all_server_ids()
+            {
                self.tools.update(cx, |tools, cx| {
                    tools.enable_source(
                        ToolSource::ContextServer {
-                            id: context_server.id().into(),
+                            id: context_server_id.0.into(),
                        },
                        cx,
                    );
@@ -541,7 +537,7 @@ impl ThreadStore {

    fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
        cx.subscribe(
-            &self.context_server_manager.clone(),
+            &self.project.read(cx).context_server_store(),
            Self::handle_context_server_event,
        )
        .detach();
@@ -549,18 +545,19 @@ impl ThreadStore {

    fn handle_context_server_event(
        &mut self,
-        context_server_manager: Entity<ContextServerManager>,
-        event: &context_server::manager::Event,
+        context_server_store: Entity<ContextServerStore>,
+        event: &project::context_server_store::Event,
        cx: &mut Context<Self>,
    ) {
        let tool_working_set = self.tools.clone();
        match event {
-            context_server::manager::Event::ServerStatusChanged { server_id, status } => {
+            project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
                match status {
-                    Some(ContextServerStatus::Running) => {
-                        if let Some(server) = context_server_manager.read(cx).get_server(server_id)
+                    ContextServerStatus::Running => {
+                        if let Some(server) =
+                            context_server_store.read(cx).get_running_server(server_id)
                        {
-                            let context_server_manager = context_server_manager.clone();
+                            let context_server_manager = context_server_store.clone();
                            cx.spawn({
                                let server = server.clone();
                                let server_id = server_id.clone();
@@ -608,7 +605,7 @@ impl ThreadStore {
                            .detach();
                        }
                    }
-                    None => {
+                    ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
                        if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
                            tool_working_set.update(cx, |tool_working_set, _| {
                                tool_working_set.remove(&tool_ids);