Dark-Alex-17/loki
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Improved MCP handling toggle handling

Browse Source
This commit is contained in:
Alex Clarke
2025-10-15 18:36:54 -06:00
parent 39fc863e22
commit a10948614d
8 changed files with 4569 additions and 4423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/agent.rs
+16 -3
View File
@@ -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 =
(!agent_config.mcp_servers.is_empty()).then(|| agent_config.mcp_servers.join(","));
let mcp_servers = if config.read().mcp_servers {
(!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
+134 -21
View File
@@ -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() {
bail!("MCP servers cannot be enabled because no MCP servers are installed.")
config.write().functions.clear_mcp_meta_functions();
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(
registry
.list_configured_servers()
);
values.extend(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
-7
View File
@@ -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
+6 -3
View File
@@ -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
.write()
.use_session(session.as_ref().map(|v| v.as_str()))?;
Config::use_session_safely(
&config,
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
+20 -8
View File
@@ -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
.get_secret(&caps[1], false);
let secret = config.vault.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)?;
registry.config = Some(config);
let mcp_servers_config: McpServersConfig =
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
+16 -4
View File
@@ -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 -1
View File
@@ -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());