feat: Improved MCP handling toggle handling

src/config/agent.rs

@@ -92,13 +92,26 @@ impl Agent {
         let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;
 
         config.write().functions.clear_mcp_meta_functions();
-		let mcp_servers =
+        let mcp_servers = if config.read().mcp_servers {
-			(!agent_config.mcp_servers.is_empty()).then(|| agent_config.mcp_servers.join(","));
+            (!agent_config.mcp_servers.is_empty()).then(|| agent_config.mcp_servers.join(","))
+        } else {
+            eprintln!(
+                "{}",
+                formatdoc!(
+                    "
+										This agent uses MCP servers, but MCP support is disabled.
+										To enable it, exit the agent and set 'mcp_servers: true', then try again
+										"
+                )
+            );
+            None
+        };
+
         let registry = config
             .write()
             .mcp_registry
             .take()
-			.expect("MCP registry should be initialized");
+            .with_context(|| "MCP registry should be populated")?;
         let new_mcp_registry =
             McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;
 

src/config/mod.rs

@@ -27,6 +27,7 @@ use crate::mcp::{
 use crate::vault::Vault;
 use anyhow::{anyhow, bail, Context, Result};
 use indexmap::IndexMap;
+use indoc::formatdoc;
 use inquire::{list_option::ListOption, validator::Validation, Confirm, MultiSelect, Select, Text};
 use log::LevelFilter;
 use parking_lot::RwLock;
@@ -723,7 +724,11 @@ impl Config {
         Ok(output)
     }
 
-	pub async fn update(config: &GlobalConfig, data: &str, abort_signal: AbortSignal) -> Result<()> {
+    pub async fn update(
+        config: &GlobalConfig,
+        data: &str,
+        abort_signal: AbortSignal,
+    ) -> Result<()> {
         let parts: Vec<&str> = data.split_whitespace().collect();
         if parts.len() != 2 {
             bail!("Usage: .set <key> <value>. If value is null, unset key.");
@@ -744,10 +749,27 @@ impl Config {
                 config.write().set_use_tools(value);
             }
             "use_mcp_servers" => {
-				let value = parse_value(value)?;
+                let value: Option<String> = parse_value(value)?;
+                if let Some(servers) = value.as_ref() {
+                    if let Some(registry) = &config.read().mcp_registry {
+                        if registry.list_configured_servers().is_empty() {
+                            bail!("No MCP servers are configured. Please configure MCP servers first before setting 'use_mcp_servers'.");
+                        }
+
+                        if servers.split(',').all(|s| {
+                            !registry
+                                .list_configured_servers()
+                                .contains(&s.trim().to_string())
+                        }) {
+                            bail!("None of the specified MCP servers in 'use_mcp_servers' are configured. Please check your MCP server configuration.");
+                        }
+                    }
+                }
                 config.write().set_use_mcp_servers(value.clone());
+                if config.read().mcp_servers {
                     config.write().functions.clear_mcp_meta_functions();
-				let registry = config.write()
+                    let registry = config
+                        .write()
                         .mcp_registry
                         .take()
                         .expect("MCP registry should be initialized");
@@ -755,12 +777,15 @@ impl Config {
                         McpRegistry::reinit(registry, value, abort_signal.clone()).await?;
 
                     if !new_mcp_registry.is_empty() {
-					config.write().functions
+                        config
+                            .write()
+                            .functions
                             .append_mcp_meta_functions(new_mcp_registry.list_started_servers());
                     }
 
                     config.write().mcp_registry = Some(new_mcp_registry);
                 }
+            }
             "max_output_tokens" => {
                 let value = parse_value(value)?;
                 config.write().set_max_output_tokens(value);
@@ -794,9 +819,27 @@ impl Config {
             }
             "mcp_servers" => {
                 let value = value.parse().with_context(|| "Invalid value")?;
-				if value && !config.write().functions.has_mcp_functions() {
+                config.write().functions.clear_mcp_meta_functions();
-					bail!("MCP servers cannot be enabled because no MCP servers are installed.")
+
+                let registry = config
+                    .write()
+                    .mcp_registry
+                    .take()
+                    .expect("MCP registry should be initialized");
+                let use_mcp_servers = if value {
+                    config.read().use_mcp_servers.clone()
+                } else {
+                    None
+                };
+                let new_registry =
+                    McpRegistry::reinit(registry, use_mcp_servers, abort_signal.clone()).await?;
+                if !new_registry.is_empty() && value {
+                    config
+                        .write()
+                        .functions
+                        .append_mcp_meta_functions(new_registry.list_started_servers());
                 }
+                config.write().mcp_registry = Some(new_registry);
                 config.write().mcp_servers = value;
             }
             "stream" => {
@@ -1022,12 +1065,25 @@ impl Config {
 
     pub async fn use_role(&mut self, name: &str, abort_signal: AbortSignal) -> Result<()> {
         let role = self.retrieve_role(name)?;
+        let mcp_servers = if self.mcp_servers {
+            role.use_mcp_servers()
+        } else {
+            eprintln!(
+                "{}",
+                formatdoc!(
+                    "
+										This role uses MCP servers, but MCP support is disabled.
+										To enable it, exit the role and set 'mcp_servers: true', then try again
+										"
+                )
+            );
+            None
+        };
         self.functions.clear_mcp_meta_functions();
-		let mcp_servers = role.use_mcp_servers();
         let registry = self
             .mcp_registry
             .take()
-			.expect("MCP registry should be initialized");
+            .with_context(|| "MCP registry should be populated")?;
         let new_mcp_registry =
             McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;
 
@@ -1234,7 +1290,31 @@ impl Config {
         names.contains(&name.to_string())
     }
 
-	pub fn use_session(&mut self, session_name: Option<&str>) -> Result<()> {
+    pub async fn use_session_safely(
+        config: &GlobalConfig,
+        session_name: Option<&str>,
+        abort_signal: AbortSignal,
+    ) -> Result<()> {
+        let mut cfg = {
+            let mut guard = config.write();
+            take(&mut *guard)
+        };
+
+        cfg.use_session(session_name, abort_signal.clone()).await?;
+
+        {
+            let mut guard = config.write();
+            *guard = cfg;
+        }
+
+        Ok(())
+    }
+
+    pub async fn use_session(
+        &mut self,
+        session_name: Option<&str>,
+        abort_signal: AbortSignal,
+    ) -> Result<()> {
         if self.session.is_some() {
             bail!(
                 "Already in a session, please run '.exit session' first to exit the current session."
@@ -1262,6 +1342,34 @@ impl Config {
         }
         let mut new_session = false;
         if let Some(session) = session.as_mut() {
+            let mcp_servers = if self.mcp_servers {
+                session.use_mcp_servers()
+            } else {
+                eprintln!(
+                    "{}",
+                    formatdoc!(
+                        "
+												This session uses MCP servers, but MCP support is disabled.
+												To enable it, exit the session and set 'mcp_servers: true', then try again
+												"
+                    )
+                );
+                None
+            };
+            self.functions.clear_mcp_meta_functions();
+            let registry = self
+                .mcp_registry
+                .take()
+                .with_context(|| "MCP registry should be populated")?;
+            let new_mcp_registry =
+                McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;
+
+            if !new_mcp_registry.is_empty() {
+                self.functions
+                    .append_mcp_meta_functions(new_mcp_registry.list_started_servers());
+            }
+
+            self.mcp_registry = Some(new_mcp_registry);
             if session.is_empty() {
                 new_session = true;
                 if let Some(LastMessage {
@@ -1661,7 +1769,7 @@ impl Config {
         if config.read().agent.is_some() {
             bail!("Already in an agent, please run '.exit agent' first to exit the current agent.");
         }
-		let agent = Agent::init(config, agent_name, abort_signal).await?;
+        let agent = Agent::init(config, agent_name, abort_signal.clone()).await?;
         let session = session_name.map(|v| v.to_string()).or_else(|| {
             if config.read().macro_flag {
                 None
@@ -1672,7 +1780,7 @@ impl Config {
         config.write().rag = agent.rag();
         config.write().agent = Some(agent);
         if let Some(session) = session {
-			config.write().use_session(Some(&session))?;
+            Config::use_session_safely(config, Some(&session), abort_signal).await?;
         } else {
             config.write().init_agent_shared_variables()?;
         }
@@ -1801,10 +1909,14 @@ impl Config {
                     .with_context(err_msg)?;
             }
             Some(("session", name)) => {
-				self.use_session(Some(name)).with_context(err_msg)?;
+                self.use_session(Some(name), abort_signal)
+                    .await
+                    .with_context(err_msg)?;
             }
             Some((session_name, role_name)) => {
-				self.use_session(Some(session_name)).with_context(err_msg)?;
+                self.use_session(Some(session_name), abort_signal.clone())
+                    .await
+                    .with_context(err_msg)?;
                 if let Some(true) = self.session.as_ref().map(|v| v.is_empty()) {
                     self.use_role(role_name, abort_signal)
                         .await
@@ -2119,12 +2231,11 @@ impl Config {
                     if prefix.is_empty() {
                         values.push("all".to_string());
                     }
+
                     if let Some(registry) = &self.mcp_registry {
-						values.extend(
+                        values.extend(registry.list_configured_servers());
-							registry
-								.list_configured_servers()
-						);
                     }
+
                     values.extend(self.mapping_mcp_servers.keys().map(|v| v.to_string()));
                     values
                         .into_iter()
@@ -2686,10 +2797,6 @@ impl Config {
         start_mcp_servers: bool,
         abort_signal: AbortSignal,
     ) -> Result<()> {
-		if !self.mcp_servers {
-			return Ok(());
-		}
-
         let mcp_registry = McpRegistry::init(
             log_path,
             start_mcp_servers,
@@ -2700,8 +2807,14 @@ impl Config {
         .await?;
         match mcp_registry.is_empty() {
             false => {
+                if self.mcp_servers {
                     self.functions
                         .append_mcp_meta_functions(mcp_registry.list_started_servers());
+                } else {
+                    debug!(
+                        "Skipping global MCP functions registration since mcp_servers was 'false'"
+                    );
+                }
             }
             _ => debug!(
                 "Skipping global MCP functions registration since start_mcp_servers was 'false'"

src/function.rs

@@ -247,13 +247,6 @@ impl Functions {
         self.declarations.is_empty()
     }
 
-    pub fn has_mcp_functions(&self) -> bool {
-        self.declarations.iter().any(|d| {
-            d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
-                || d.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
-        })
-    }
-
     pub fn clear_mcp_meta_functions(&mut self) {
         self.declarations.retain(|d| {
             !d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)

src/main.rs

@@ -176,9 +176,12 @@ async fn run(
             Config::use_role_safely(&config, CODE_ROLE, abort_signal.clone()).await?;
         }
         if let Some(session) = &cli.session {
-            config
+            Config::use_session_safely(
-                .write()
+                &config,
-                .use_session(session.as_ref().map(|v| v.as_str()))?;
+                session.as_ref().map(|v| v.as_str()),
+                abort_signal.clone(),
+            )
+            .await?;
         }
         if let Some(rag) = &cli.rag {
             Config::use_rag(&config, Some(rag), abort_signal.clone()).await?;

src/mcp/mod.rs

@@ -1,8 +1,10 @@
 use crate::config::Config;
 use crate::utils::{abortable_run_with_spinner, AbortSignal};
+use crate::vault::SECRET_RE;
 use anyhow::{anyhow, Context, Result};
 use futures_util::future::BoxFuture;
 use futures_util::{stream, StreamExt, TryStreamExt};
+use indoc::formatdoc;
 use rmcp::model::{CallToolRequestParam, CallToolResult};
 use rmcp::service::RunningService;
 use rmcp::transport::TokioChildProcess;
@@ -14,9 +16,8 @@ use std::collections::{HashMap, HashSet};
 use std::fs::OpenOptions;
 use std::path::PathBuf;
 use std::process::Stdio;
-use std::sync::{Arc};
+use std::sync::Arc;
 use tokio::process::Command;
-use crate::vault::SECRET_RE;
 
 pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke";
 pub const MCP_LIST_META_FUNCTION_NAME_PREFIX: &str = "mcp_list";
@@ -85,10 +86,15 @@ impl McpRegistry {
         let content = tokio::fs::read_to_string(Config::mcp_config_file())
             .await
             .with_context(err)?;
+
+        if content.trim().is_empty() {
+            debug!("MCP config file is empty, skipping MCP initialization");
+            return Ok(registry);
+        }
+
         let mut missing_secrets = vec![];
         let parsed_content = SECRET_RE.replace_all(&content, |caps: &fancy_regex::Captures<'_>| {
-			let secret = config.vault
+            let secret = config.vault.get_secret(&caps[1], false);
-				.get_secret(&caps[1], false);
             match secret {
                 Ok(s) => s,
                 Err(_) => {
@@ -99,13 +105,19 @@ impl McpRegistry {
         });
 
         if !missing_secrets.is_empty() {
-			return Err(anyhow!("MCP config file references secrets that are missing from the vault: {:?}", missing_secrets));
+            return Err(anyhow!(formatdoc!(
+                "
+								MCP config file references secrets that are missing from the vault: {:?}
+								Please add these secrets to the vault and try again.",
+                missing_secrets
+            )));
         }
 
-		let config: McpServersConfig = serde_json::from_str(&parsed_content).with_context(err)?;
+        let mcp_servers_config: McpServersConfig =
-		registry.config = Some(config);
+            serde_json::from_str(&parsed_content).with_context(err)?;
+        registry.config = Some(mcp_servers_config);
 
-		if start_mcp_servers {
+        if start_mcp_servers && config.mcp_servers {
             abortable_run_with_spinner(
                 registry.start_select_mcp_servers(use_mcp_servers),
                 "Loading MCP servers",

src/repl/mod.rs

@@ -450,7 +450,7 @@ pub async fn run_repl_command(
                 ),
             },
             ".session" => {
-				config.write().use_session(args)?;
+                Config::use_session_safely(config, args, abort_signal.clone()).await?;
                 Config::maybe_autoname_session(config.clone());
             }
             ".rag" => {
@@ -700,7 +700,11 @@ pub async fn run_repl_command(
                         .mcp_registry
                         .take()
                         .expect("MCP registry should exist");
-					let use_mcp_servers = config.read().use_mcp_servers.clone();
+                    let use_mcp_servers = if config.read().mcp_servers {
+                        config.read().use_mcp_servers.clone()
+                    } else {
+                        None
+                    };
                     let registry =
                         McpRegistry::reinit(registry, use_mcp_servers, abort_signal.clone())
                             .await?;
@@ -722,7 +726,11 @@ pub async fn run_repl_command(
                             .mcp_registry
                             .take()
                             .expect("MCP registry should exist");
-						let use_mcp_servers = config.read().use_mcp_servers.clone();
+                        let use_mcp_servers = if config.read().mcp_servers {
+                            config.read().use_mcp_servers.clone()
+                        } else {
+                            None
+                        };
                         let registry =
                             McpRegistry::reinit(registry, use_mcp_servers, abort_signal.clone())
                                 .await?;
@@ -749,7 +757,11 @@ pub async fn run_repl_command(
                         .mcp_registry
                         .take()
                         .expect("MCP registry should exist");
-					let use_mcp_servers = config.read().use_mcp_servers.clone();
+                    let use_mcp_servers = if config.read().mcp_servers {
+                        config.read().use_mcp_servers.clone()
+                    } else {
+                        None
+                    };
                     let registry =
                         McpRegistry::reinit(registry, use_mcp_servers, abort_signal.clone())
                             .await?;

src/vault/mod.rs

@@ -1,6 +1,5 @@
 mod utils;
 
-use std::sync::LazyLock;
 use crate::cli::Cli;
 use crate::config::Config;
 use crate::vault::utils::ensure_password_file_initialized;
@@ -9,6 +8,7 @@ use fancy_regex::Regex;
 use gman::providers::local::LocalProvider;
 use gman::providers::SecretProvider;
 use inquire::{required, Password, PasswordDisplayMode};
+use std::sync::LazyLock;
 use tokio::runtime::Handle;
 
 pub static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{\{(.+)}}").unwrap());