agent: Store if context server should be enabled/disabled in the settings (#32994)

Bennet Bo Fenner and Danilo Leal created

- [x] Show disabled MCP servers in the list so you can enable them again
- [x] If MCP is not present in the settings, add it

Closes #ISSUE

Release Notes:

- agent: Allow to enable/disable context servers permanently in the
agent configuration view

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>

Change summary

crates/agent/src/agent_configuration.rs                                |  69 
crates/agent/src/agent_configuration/configure_context_server_modal.rs |  25 
crates/extension_host/src/wasm_host/wit/since_v0_6_0.rs                |   2 
crates/project/src/context_server_store.rs                             | 182 
crates/project/src/project_settings.rs                                 |  23 
5 files changed, 260 insertions(+), 41 deletions(-)

Detailed changes

crates/agent/src/agent_configuration.rs 🔗

@@ -20,7 +20,7 @@ use language_model::{LanguageModelProvider, LanguageModelProviderId, LanguageMod
 use notifications::status_toast::{StatusToast, ToastIcon};
 use project::{
     context_server_store::{ContextServerConfiguration, ContextServerStatus, ContextServerStore},
-    project_settings::ProjectSettings,
+    project_settings::{ContextServerSettings, ProjectSettings},
 };
 use settings::{Settings, update_settings_file};
 use ui::{
@@ -741,24 +741,59 @@ impl AgentConfiguration {
                                         let context_server_manager =
                                             self.context_server_store.clone();
                                         let context_server_id = context_server_id.clone();
+                                        let fs = self.fs.clone();
+
+                                        move |state, _window, cx| {
+                                            let is_enabled = match state {
+                                                ToggleState::Unselected
+                                                | ToggleState::Indeterminate => {
+                                                    context_server_manager.update(
+                                                        cx,
+                                                        |this, cx| {
+                                                            this.stop_server(
+                                                                &context_server_id,
+                                                                cx,
+                                                            )
+                                                            .log_err();
+                                                        },
+                                                    );
+                                                    false
+                                                }
+                                                ToggleState::Selected => {
+                                                    context_server_manager.update(
+                                                        cx,
+                                                        |this, cx| {
+                                                            if let Some(server) =
+                                                                this.get_server(&context_server_id)
+                                                            {
+                                                                this.start_server(server, cx);
+                                                            }
+                                                        },
+                                                    );
+                                                    true
+                                                }
+                                            };
+                                            update_settings_file::<ProjectSettings>(
+                                                fs.clone(),
+                                                cx,
+                                                {
+                                                    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_id, cx)
-                                                        .log_err();
-                                                });
-                                            }
-                                            ToggleState::Selected => {
-                                                context_server_manager.update(cx, |this, cx| {
-                                                    if let Some(server) =
-                                                        this.get_server(&context_server_id)
-                                                    {
-                                                        this.start_server(server, cx);
+                                                    move |settings, _| {
+                                                        settings
+                                                            .context_servers
+                                                            .entry(context_server_id.0)
+                                                            .or_insert_with(|| {
+                                                                ContextServerSettings::Extension {
+                                                                    enabled: is_enabled,
+                                                                    settings: serde_json::json!({}),
+                                                                }
+                                                            })
+                                                            .set_enabled(is_enabled);
                                                     }
-                                                })
-                                            }
+                                                },
+                                            );
                                         }
                                     }),
                             ),

crates/agent/src/agent_configuration/configure_context_server_modal.rs 🔗

@@ -140,8 +140,15 @@ impl ConfigurationSource {
     fn output(&self, cx: &mut App) -> Result<(ContextServerId, ContextServerSettings)> {
         match self {
             ConfigurationSource::New { editor } | ConfigurationSource::Existing { editor } => {
-                parse_input(&editor.read(cx).text(cx))
-                    .map(|(id, command)| (id, ContextServerSettings::Custom { command }))
+                parse_input(&editor.read(cx).text(cx)).map(|(id, command)| {
+                    (
+                        id,
+                        ContextServerSettings::Custom {
+                            enabled: true,
+                            command,
+                        },
+                    )
+                })
             }
             ConfigurationSource::Extension {
                 id,
@@ -160,7 +167,13 @@ impl ConfigurationSource {
                         return Err(anyhow::anyhow!(error.to_string()));
                     }
                 }
-                Ok((id.clone(), ContextServerSettings::Extension { settings }))
+                Ok((
+                    id.clone(),
+                    ContextServerSettings::Extension {
+                        enabled: true,
+                        settings,
+                    },
+                ))
             }
         }
     }
@@ -283,6 +296,7 @@ impl ConfigureContextServerModal {
                     .read(cx)
                     .context_server_descriptor(&server_id.0)
                     .map(|_| ContextServerSettings::Extension {
+                        enabled: true,
                         settings: serde_json::json!({}),
                     })
             })
@@ -292,7 +306,10 @@ impl ConfigureContextServerModal {
 
         window.spawn(cx, async move |cx| {
             let target = match settings {
-                ContextServerSettings::Custom { command } => Some(ConfigurationTarget::Existing {
+                ContextServerSettings::Custom {
+                    enabled: _,
+                    command,
+                } => Some(ConfigurationTarget::Existing {
                     id: server_id,
                     command,
                 }),

crates/extension_host/src/wasm_host/wit/since_v0_6_0.rs 🔗

@@ -949,6 +949,7 @@ impl ExtensionImports for WasmState {
 
                         match settings {
                             project::project_settings::ContextServerSettings::Custom {
+                                enabled: _,
                                 command,
                             } => Ok(serde_json::to_string(&settings::ContextServerSettings {
                                 command: Some(settings::CommandSettings {
@@ -959,6 +960,7 @@ impl ExtensionImports for WasmState {
                                 settings: None,
                             })?),
                             project::project_settings::ContextServerSettings::Extension {
+                                enabled: _,
                                 settings,
                             } => Ok(serde_json::to_string(&settings::ContextServerSettings {
                                 command: None,

crates/project/src/context_server_store.rs 🔗

@@ -36,13 +36,8 @@ impl ContextServerStatus {
         match state {
             ContextServerState::Starting { .. } => ContextServerStatus::Starting,
             ContextServerState::Running { .. } => ContextServerStatus::Running,
-            ContextServerState::Stopped { error, .. } => {
-                if let Some(error) = error {
-                    ContextServerStatus::Error(error.clone())
-                } else {
-                    ContextServerStatus::Stopped
-                }
-            }
+            ContextServerState::Stopped { .. } => ContextServerStatus::Stopped,
+            ContextServerState::Error { error, .. } => ContextServerStatus::Error(error.clone()),
         }
     }
 }
@@ -60,7 +55,11 @@ enum ContextServerState {
     Stopped {
         server: Arc<ContextServer>,
         configuration: Arc<ContextServerConfiguration>,
-        error: Option<Arc<str>>,
+    },
+    Error {
+        server: Arc<ContextServer>,
+        configuration: Arc<ContextServerConfiguration>,
+        error: Arc<str>,
     },
 }
 
@@ -70,6 +69,7 @@ impl ContextServerState {
             ContextServerState::Starting { server, .. } => server.clone(),
             ContextServerState::Running { server, .. } => server.clone(),
             ContextServerState::Stopped { server, .. } => server.clone(),
+            ContextServerState::Error { server, .. } => server.clone(),
         }
     }
 
@@ -78,11 +78,12 @@ impl ContextServerState {
             ContextServerState::Starting { configuration, .. } => configuration.clone(),
             ContextServerState::Running { configuration, .. } => configuration.clone(),
             ContextServerState::Stopped { configuration, .. } => configuration.clone(),
+            ContextServerState::Error { configuration, .. } => configuration.clone(),
         }
     }
 }
 
-#[derive(PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum ContextServerConfiguration {
     Custom {
         command: ContextServerCommand,
@@ -109,10 +110,14 @@ impl ContextServerConfiguration {
         cx: &AsyncApp,
     ) -> Option<Self> {
         match settings {
-            ContextServerSettings::Custom { command } => {
-                Some(ContextServerConfiguration::Custom { command })
-            }
-            ContextServerSettings::Extension { settings } => {
+            ContextServerSettings::Custom {
+                enabled: _,
+                command,
+            } => Some(ContextServerConfiguration::Custom { command }),
+            ContextServerSettings::Extension {
+                enabled: _,
+                settings,
+            } => {
                 let descriptor = cx
                     .update(|cx| registry.read(cx).context_server_descriptor(&id.0))
                     .ok()
@@ -272,6 +277,10 @@ impl ContextServerStore {
                 .flatten()
                 .context("Failed to get context server settings")?;
 
+            if !settings.enabled() {
+                return Ok(());
+            }
+
             let (registry, worktree_store) = this.update(cx, |this, _| {
                 (this.registry.clone(), this.worktree_store.clone())
             })?;
@@ -293,6 +302,13 @@ impl ContextServerStore {
     }
 
     pub fn stop_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {
+        if matches!(
+            self.servers.get(id),
+            Some(ContextServerState::Stopped { .. })
+        ) {
+            return Ok(());
+        }
+
         let state = self
             .servers
             .remove(id)
@@ -311,7 +327,6 @@ impl ContextServerStore {
             ContextServerState::Stopped {
                 configuration,
                 server,
-                error: None,
             },
             cx,
         );
@@ -371,10 +386,10 @@ impl ContextServerStore {
                         this.update(cx, |this, cx| {
                             this.update_server_state(
                                 id.clone(),
-                                ContextServerState::Stopped {
+                                ContextServerState::Error {
                                     configuration,
                                     server,
-                                    error: Some(err.to_string().into()),
+                                    error: err.to_string().into(),
                                 },
                                 cx,
                             )
@@ -491,11 +506,17 @@ impl ContextServerStore {
             configured_servers
                 .entry(id)
                 .or_insert(ContextServerSettings::Extension {
+                    enabled: true,
                     settings: serde_json::json!({}),
                 });
         }
 
-        let configured_servers = join_all(configured_servers.into_iter().map(|(id, settings)| {
+        let (enabled_servers, disabled_servers): (HashMap<_, _>, HashMap<_, _>) =
+            configured_servers
+                .into_iter()
+                .partition(|(_, settings)| settings.enabled());
+
+        let configured_servers = join_all(enabled_servers.into_iter().map(|(id, settings)| {
             let id = ContextServerId(id);
             ContextServerConfiguration::from_settings(
                 settings,
@@ -520,13 +541,19 @@ impl ContextServerStore {
                 // All servers that are not in desired_servers should be removed from the store.
                 // This can happen if the user removed a server from the context server settings.
                 if !configured_servers.contains_key(&server_id) {
-                    servers_to_remove.insert(server_id.clone());
+                    if disabled_servers.contains_key(&server_id.0) {
+                        servers_to_stop.insert(server_id.clone());
+                    } else {
+                        servers_to_remove.insert(server_id.clone());
+                    }
                 }
             }
 
             for (id, config) in configured_servers {
-                let existing_config = this.servers.get(&id).map(|state| state.configuration());
-                if existing_config.as_deref() != Some(&config) {
+                let state = this.servers.get(&id);
+                let is_stopped = matches!(state, Some(ContextServerState::Stopped { .. }));
+                let existing_config = state.as_ref().map(|state| state.configuration());
+                if existing_config.as_deref() != Some(&config) || is_stopped {
                     let config = Arc::new(config);
                     if let Some(server) = this
                         .create_context_server(id.clone(), config.clone())
@@ -773,6 +800,7 @@ mod tests {
             vec![(
                 SERVER_1_ID.into(),
                 ContextServerSettings::Extension {
+                    enabled: true,
                     settings: json!({
                         "somevalue": true
                     }),
@@ -829,6 +857,7 @@ mod tests {
                 vec![(
                     server_1_id.0.clone(),
                     ContextServerSettings::Extension {
+                        enabled: true,
                         settings: json!({
                             "somevalue": false
                         }),
@@ -847,6 +876,7 @@ mod tests {
                 vec![(
                     server_1_id.0.clone(),
                     ContextServerSettings::Extension {
+                        enabled: true,
                         settings: json!({
                             "somevalue": false
                         }),
@@ -873,6 +903,7 @@ mod tests {
                     (
                         server_1_id.0.clone(),
                         ContextServerSettings::Extension {
+                            enabled: true,
                             settings: json!({
                                 "somevalue": false
                             }),
@@ -881,6 +912,7 @@ mod tests {
                     (
                         server_2_id.0.clone(),
                         ContextServerSettings::Custom {
+                            enabled: true,
                             command: ContextServerCommand {
                                 path: "somebinary".to_string(),
                                 args: vec!["arg".to_string()],
@@ -911,6 +943,7 @@ mod tests {
                     (
                         server_1_id.0.clone(),
                         ContextServerSettings::Extension {
+                            enabled: true,
                             settings: json!({
                                 "somevalue": false
                             }),
@@ -919,6 +952,7 @@ mod tests {
                     (
                         server_2_id.0.clone(),
                         ContextServerSettings::Custom {
+                            enabled: true,
                             command: ContextServerCommand {
                                 path: "somebinary".to_string(),
                                 args: vec!["anotherArg".to_string()],
@@ -944,6 +978,7 @@ mod tests {
                 vec![(
                     server_1_id.0.clone(),
                     ContextServerSettings::Extension {
+                        enabled: true,
                         settings: json!({
                             "somevalue": false
                         }),
@@ -960,6 +995,112 @@ mod tests {
         }
     }
 
+    #[gpui::test]
+    async fn test_context_server_enabled_disabled(cx: &mut TestAppContext) {
+        const SERVER_1_ID: &'static str = "mcp-1";
+
+        let server_1_id = ContextServerId(SERVER_1_ID.into());
+
+        let (_fs, project) = setup_context_server_test(
+            cx,
+            json!({"code.rs": ""}),
+            vec![(
+                SERVER_1_ID.into(),
+                ContextServerSettings::Custom {
+                    enabled: true,
+                    command: ContextServerCommand {
+                        path: "somebinary".to_string(),
+                        args: vec!["arg".to_string()],
+                        env: None,
+                    },
+                },
+            )],
+        )
+        .await;
+
+        let executor = cx.executor();
+        let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
+        let store = cx.new(|cx| {
+            ContextServerStore::test_maintain_server_loop(
+                Box::new(move |id, _| {
+                    Arc::new(ContextServer::new(
+                        id.clone(),
+                        Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
+                    ))
+                }),
+                registry.clone(),
+                project.read(cx).worktree_store(),
+                cx,
+            )
+        });
+
+        // Ensure that mcp-1 starts up
+        {
+            let _server_events = assert_server_events(
+                &store,
+                vec![
+                    (server_1_id.clone(), ContextServerStatus::Starting),
+                    (server_1_id.clone(), ContextServerStatus::Running),
+                ],
+                cx,
+            );
+            cx.run_until_parked();
+        }
+
+        // Ensure that mcp-1 is stopped once it is disabled.
+        {
+            let _server_events = assert_server_events(
+                &store,
+                vec![(server_1_id.clone(), ContextServerStatus::Stopped)],
+                cx,
+            );
+            set_context_server_configuration(
+                vec![(
+                    server_1_id.0.clone(),
+                    ContextServerSettings::Custom {
+                        enabled: false,
+                        command: ContextServerCommand {
+                            path: "somebinary".to_string(),
+                            args: vec!["arg".to_string()],
+                            env: None,
+                        },
+                    },
+                )],
+                cx,
+            );
+
+            cx.run_until_parked();
+        }
+
+        // Ensure that mcp-1 is started once it is enabled again.
+        {
+            let _server_events = assert_server_events(
+                &store,
+                vec![
+                    (server_1_id.clone(), ContextServerStatus::Starting),
+                    (server_1_id.clone(), ContextServerStatus::Running),
+                ],
+                cx,
+            );
+            set_context_server_configuration(
+                vec![(
+                    server_1_id.0.clone(),
+                    ContextServerSettings::Custom {
+                        enabled: true,
+                        command: ContextServerCommand {
+                            path: "somebinary".to_string(),
+                            args: vec!["arg".to_string()],
+                            env: None,
+                        },
+                    },
+                )],
+                cx,
+            );
+
+            cx.run_until_parked();
+        }
+    }
+
     fn set_context_server_configuration(
         context_servers: Vec<(Arc<str>, ContextServerSettings)>,
         cx: &mut TestAppContext,
@@ -997,6 +1138,7 @@ mod tests {
 
     fn dummy_server_settings() -> ContextServerSettings {
         ContextServerSettings::Custom {
+            enabled: true,
             command: ContextServerCommand {
                 path: "somebinary".to_string(),
                 args: vec!["arg".to_string()],

crates/project/src/project_settings.rs 🔗

@@ -88,12 +88,18 @@ pub struct DapSettings {
 #[serde(tag = "source", rename_all = "snake_case")]
 pub enum ContextServerSettings {
     Custom {
+        /// Whether the context server is enabled.
+        #[serde(default = "default_true")]
+        enabled: bool,
         /// The command to run this context server.
         ///
         /// This will override the command set by an extension.
         command: ContextServerCommand,
     },
     Extension {
+        /// Whether the context server is enabled.
+        #[serde(default = "default_true")]
+        enabled: bool,
         /// The settings for this context server specified by the extension.
         ///
         /// Consult the documentation for the context server to see what settings
@@ -102,6 +108,22 @@ pub enum ContextServerSettings {
     },
 }
 
+impl ContextServerSettings {
+    pub fn enabled(&self) -> bool {
+        match self {
+            ContextServerSettings::Custom { enabled, .. } => *enabled,
+            ContextServerSettings::Extension { enabled, .. } => *enabled,
+        }
+    }
+
+    pub fn set_enabled(&mut self, enabled: bool) {
+        match self {
+            ContextServerSettings::Custom { enabled: e, .. } => *e = enabled,
+            ContextServerSettings::Extension { enabled: e, .. } => *e = enabled,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
 pub struct NodeBinarySettings {
     /// The path to the Node binary.
@@ -478,6 +500,7 @@ impl Settings for ProjectSettings {
                     Some((
                         k.clone().into(),
                         ContextServerSettings::Custom {
+                            enabled: true,
                             command: serde_json::from_value::<VsCodeContextServerCommand>(
                                 v.clone(),
                             )