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

feat: Secret injection into the MCP configuration

Browse Source
This commit is contained in:
Alex Clarke
2025-10-15 16:06:59 -06:00
parent df8b326d89
commit 39fc863e22
5 changed files with 4410 additions and 4380 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/agent.rs
+3 -4
View File
@@ -9,8 +9,7 @@ use anyhow::{Context, Result};
use inquire::{validator::Validation, Text};
use rust_embed::Embed;
use serde::{Deserialize, Serialize};
use std::ffi::OsStr;
use std::{fs, fs::read_to_string, path::Path};
use std::{ffi::OsStr, path::Path};
const DEFAULT_AGENT_NAME: &str = "rag";
@@ -68,7 +67,7 @@ impl Agent {
#[cfg(unix)]
if is_script {
use std::os::unix::fs::PermissionsExt;
use std::{fs, os::unix::fs::PermissionsExt};
fs::set_permissions(&file_path, fs::Permissions::from_mode(0o755))?;
}
}
@@ -104,7 +103,7 @@ impl Agent {
McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;
if !new_mcp_registry.is_empty() {
functions.append_mcp_meta_functions(new_mcp_registry.list_servers());
functions.append_mcp_meta_functions(new_mcp_registry.list_started_servers());
}
config.write().mcp_registry = Some(new_mcp_registry);
src/config/mod.rs
+23 -25
View File
@@ -723,7 +723,7 @@ impl Config {
Ok(output)
}
pub fn update(config: &GlobalConfig, data: &str) -> 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.");
@@ -745,7 +745,21 @@ impl Config {
}
"use_mcp_servers" => {
let value = parse_value(value)?;
config.write().set_use_mcp_servers(value);
config.write().set_use_mcp_servers(value.clone());
config.write().functions.clear_mcp_meta_functions();
let registry = config.write()
.mcp_registry
.take()
.expect("MCP registry should be initialized");
let new_mcp_registry =
McpRegistry::reinit(registry, value, abort_signal.clone()).await?;
if !new_mcp_registry.is_empty() {
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)?;
@@ -1019,7 +1033,7 @@ impl Config {
if !new_mcp_registry.is_empty() {
self.functions
.append_mcp_meta_functions(new_mcp_registry.list_servers());
.append_mcp_meta_functions(new_mcp_registry.list_started_servers());
}
self.mcp_registry = Some(new_mcp_registry);
@@ -2105,29 +2119,12 @@ impl Config {
if prefix.is_empty() {
values.push("all".to_string());
}
if let Some(registry) = &self.mcp_registry {
values.extend(
self.functions
.declarations()
.iter()
.filter(|v| {
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|| v.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
})
.map(|v| {
v.name
.strip_prefix(
format!("{MCP_LIST_META_FUNCTION_NAME_PREFIX}_").as_str(),
)
.or_else(|| {
v.name.strip_prefix(
format!("{MCP_INVOKE_META_FUNCTION_NAME_PREFIX}_")
.as_str(),
)
})
.unwrap()
.to_string()
}),
registry
.list_configured_servers()
);
}
values.extend(self.mapping_mcp_servers.keys().map(|v| v.to_string()));
values
.into_iter()
@@ -2698,12 +2695,13 @@ impl Config {
start_mcp_servers,
self.use_mcp_servers.clone(),
abort_signal.clone(),
self,
)
.await?;
match mcp_registry.is_empty() {
false => {
self.functions
.append_mcp_meta_functions(mcp_registry.list_servers());
.append_mcp_meta_functions(mcp_registry.list_started_servers());
}
_ => debug!(
"Skipping global MCP functions registration since start_mcp_servers was 'false'"
src/mcp/mod.rs
+32 -6
View File
@@ -14,8 +14,9 @@ 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";
@@ -57,6 +58,7 @@ impl McpRegistry {
start_mcp_servers: bool,
use_mcp_servers: Option<String>,
abort_signal: AbortSignal,
config: &Config,
) -> Result<Self> {
let mut registry = Self {
log_path,
@@ -83,7 +85,24 @@ impl McpRegistry {
let content = tokio::fs::read_to_string(Config::mcp_config_file())
.await
.with_context(err)?;
let config: McpServersConfig = serde_json::from_str(&content).with_context(err)?;
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);
match secret {
Ok(s) => s,
Err(_) => {
missing_secrets.push(caps[1].to_string());
"".to_string()
}
}
});
if !missing_secrets.is_empty() {
return Err(anyhow!("MCP config file references secrets that are missing from the vault: {:?}", missing_secrets));
}
let config: McpServersConfig = serde_json::from_str(&parsed_content).with_context(err)?;
registry.config = Some(config);
if start_mcp_servers {
@@ -124,13 +143,12 @@ impl McpRegistry {
async fn start_select_mcp_servers(&mut self, use_mcp_servers: Option<String>) -> Result<()> {
if self.config.is_none() {
debug!("MCP config is not present; assuming MCP servers are disabled globally. skipping MCP initialization");
debug!("MCP config is not present; assuming MCP servers are disabled globally. Skipping MCP initialization");
return Ok(());
}
debug!("Starting selected MCP servers: {:?}", use_mcp_servers);
if let Some(servers) = use_mcp_servers {
debug!("Starting selected MCP servers: {:?}", servers);
let config = self
.config
.as_ref()
@@ -231,10 +249,18 @@ impl McpRegistry {
Ok(self)
}
pub fn list_servers(&self) -> Vec<String> {
pub fn list_started_servers(&self) -> Vec<String> {
self.servers.keys().cloned().collect()
}
pub fn list_configured_servers(&self) -> Vec<String> {
if let Some(config) = &self.config {
config.mcp_servers.keys().cloned().collect()
} else {
vec![]
}
}
pub fn catalog(&self) -> BoxFuture<'static, Result<Value>> {
let servers: Vec<(String, Arc<ConnectedServer>)> = self
.servers
src/repl/mod.rs
+6 -6
View File
@@ -184,12 +184,12 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 37]> = LazyLock::new(|| {
"Delete roles, sessions, RAGs, or agents",
AssertState::pass(),
),
ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
ReplCommand::new(
".vault",
"View or modify the Loki vault",
AssertState::pass(),
),
ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
]
});
static COMMAND_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^\s*(\.\S*)\s*").unwrap());
@@ -663,7 +663,7 @@ pub async fn run_repl_command(
}
".set" => match args {
Some(args) => {
Config::update(config, args)?;
Config::update(config, args, abort_signal).await?;
}
_ => {
println!("Usage: .set <key> <value>...")
@@ -708,7 +708,7 @@ pub async fn run_repl_command(
config
.write()
.functions
.append_mcp_meta_functions(registry.list_servers());
.append_mcp_meta_functions(registry.list_started_servers());
}
config.write().mcp_registry = Some(registry);
}
@@ -730,7 +730,7 @@ pub async fn run_repl_command(
config
.write()
.functions
.append_mcp_meta_functions(registry.list_servers());
.append_mcp_meta_functions(registry.list_started_servers());
}
config.write().mcp_registry = Some(registry);
} else {
@@ -757,7 +757,7 @@ pub async fn run_repl_command(
config
.write()
.functions
.append_mcp_meta_functions(registry.list_servers());
.append_mcp_meta_functions(registry.list_started_servers());
}
config.write().mcp_registry = Some(registry);
}
@@ -782,7 +782,7 @@ pub async fn run_repl_command(
}
Some(("get", name)) => {
if let Some(name) = name {
config.read().vault.get_secret(name)?;
config.read().vault.get_secret(name, true)?;
} else {
println!("Usage: .vault get <name>");
}
src/vault/mod.rs
+11 -4
View File
@@ -1,14 +1,18 @@
mod utils;
use std::sync::LazyLock;
use crate::cli::Cli;
use crate::config::Config;
use crate::vault::utils::ensure_password_file_initialized;
use anyhow::{Context, Result};
use fancy_regex::Regex;
use gman::providers::local::LocalProvider;
use gman::providers::SecretProvider;
use inquire::{required, Password, PasswordDisplayMode};
use tokio::runtime::Handle;
pub static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{\{(.+)}}").unwrap());
#[derive(Debug, Default, Clone)]
pub struct Vault {
local_provider: LocalProvider,
@@ -45,14 +49,17 @@ impl Vault {
Ok(())
}
pub fn get_secret(&self, secret_name: &str) -> Result<()> {
pub fn get_secret(&self, secret_name: &str, display_output: bool) -> Result<String> {
let h = Handle::current();
let secret = tokio::task::block_in_place(|| {
h.block_on(self.local_provider.get_secret(secret_name))
})?;
println!("{}", secret);
Ok(())
if display_output {
println!("{}", secret);
}
Ok(secret)
}
pub fn update_secret(&self, secret_name: &str) -> Result<()> {
@@ -105,7 +112,7 @@ impl Vault {
}
if let Some(secret_name) = cli.get_secret {
config.vault.get_secret(&secret_name)?;
config.vault.get_secret(&secret_name, true)?;
}
if let Some(secret_name) = cli.update_secret {