feat(mcp): add mcp_tools allowlist config surfaces across roles, sessions, agents, graphs, and skills
This commit is contained in:
@@ -682,6 +682,10 @@ impl RoleLike for Agent {
|
||||
Some(self.config.mcp_servers.clone())
|
||||
}
|
||||
|
||||
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>> {
|
||||
self.config.mcp_tools.clone()
|
||||
}
|
||||
|
||||
fn set_model(&mut self, model: Model) {
|
||||
self.config.model_id = Some(model.id());
|
||||
self.model = model;
|
||||
@@ -723,6 +727,10 @@ impl RoleLike for Agent {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>) {
|
||||
self.config.mcp_tools = value;
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
@@ -774,6 +782,8 @@ pub struct AgentConfig {
|
||||
pub version: String,
|
||||
#[serde(default)]
|
||||
pub mcp_servers: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
#[serde(default)]
|
||||
pub global_tools: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -846,6 +856,7 @@ impl AgentConfig {
|
||||
description: graph.description.clone(),
|
||||
global_tools: graph.global_tools.clone(),
|
||||
mcp_servers: graph.mcp_servers.clone(),
|
||||
mcp_tools: graph.mcp_tools.clone(),
|
||||
skills_enabled: graph.skills_enabled,
|
||||
enabled_skills: graph.enabled_skills.clone(),
|
||||
inject_skill_instructions: graph.inject_skill_instructions.unwrap_or(true),
|
||||
@@ -1285,6 +1296,33 @@ variables:
|
||||
assert_eq!(config.enabled_macros, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_config_parses_mcp_tools() {
|
||||
let yaml =
|
||||
"name: minimal\ninstructions: hi\nmcp_tools:\n github:\n - get_*\n - list_*\n";
|
||||
let config: AgentConfig = serde_yaml::from_str(yaml).unwrap();
|
||||
|
||||
let mcp_tools = config.mcp_tools.unwrap();
|
||||
assert_eq!(
|
||||
mcp_tools.get("github"),
|
||||
Some(&vec!["get_*".to_string(), "list_*".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_mcp_tools_role_like_round_trip() {
|
||||
let config: AgentConfig =
|
||||
serde_yaml::from_str("name: minimal\ninstructions: hi\n").unwrap();
|
||||
let mut agent = Agent::test_new(config);
|
||||
assert_eq!(agent.mcp_tools(), None);
|
||||
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]);
|
||||
agent.set_mcp_tools(Some(mcp_tools.clone()));
|
||||
|
||||
assert_eq!(agent.mcp_tools(), Some(mcp_tools));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_config_enabled_macros_empty_list_is_some_empty() {
|
||||
let yaml = "name: minimal\ninstructions: hi\nenabled_macros: []\n";
|
||||
|
||||
@@ -50,6 +50,7 @@ pub struct AppConfig {
|
||||
pub mapping_mcp_servers: IndexMap<String, String>,
|
||||
#[serde(default, deserialize_with = "super::deserialize_csv_or_vec")]
|
||||
pub enabled_mcp_servers: Option<Vec<String>>,
|
||||
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
|
||||
pub auto_continue: bool,
|
||||
pub max_auto_continues: usize,
|
||||
@@ -136,6 +137,7 @@ impl Default for AppConfig {
|
||||
mcp_server_support: true,
|
||||
mapping_mcp_servers: Default::default(),
|
||||
enabled_mcp_servers: None,
|
||||
mcp_tools: None,
|
||||
|
||||
auto_continue: false,
|
||||
max_auto_continues: 10,
|
||||
@@ -223,6 +225,7 @@ impl AppConfig {
|
||||
mcp_server_support: config.mcp_server_support,
|
||||
mapping_mcp_servers: config.mapping_mcp_servers,
|
||||
enabled_mcp_servers: config.enabled_mcp_servers,
|
||||
mcp_tools: config.mcp_tools,
|
||||
|
||||
auto_continue: config.auto_continue,
|
||||
max_auto_continues: config.max_auto_continues,
|
||||
@@ -786,6 +789,22 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_config_copies_mcp_tools() {
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]);
|
||||
let cfg = Config {
|
||||
model_id: "test-model".to_string(),
|
||||
clients: vec![ClientConfig::default()],
|
||||
mcp_tools: Some(mcp_tools.clone()),
|
||||
..Config::default()
|
||||
};
|
||||
|
||||
let app = AppConfig::from_config(cfg).unwrap();
|
||||
|
||||
assert_eq!(app.mcp_tools, Some(mcp_tools));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn from_config_copies_enabled_macros() {
|
||||
|
||||
@@ -251,6 +251,7 @@ pub struct Config {
|
||||
pub mapping_mcp_servers: IndexMap<String, String>,
|
||||
#[serde(default, deserialize_with = "deserialize_csv_or_vec")]
|
||||
pub enabled_mcp_servers: Option<Vec<String>>,
|
||||
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
|
||||
pub auto_continue: bool,
|
||||
pub max_auto_continues: usize,
|
||||
@@ -334,6 +335,7 @@ impl Default for Config {
|
||||
mcp_server_support: true,
|
||||
mapping_mcp_servers: Default::default(),
|
||||
enabled_mcp_servers: None,
|
||||
mcp_tools: None,
|
||||
|
||||
auto_continue: false,
|
||||
max_auto_continues: 10,
|
||||
@@ -1129,6 +1131,17 @@ clients:
|
||||
assert!(validate_no_template_in_secrets_provider(yaml).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_yaml_parses_mcp_tools() {
|
||||
let cfg: Config = serde_yaml::from_str("mcp_tools:\n github:\n - get_*\n").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cfg.mcp_tools.as_ref().unwrap().get("github"),
|
||||
Some(&vec!["get_*".to_string()])
|
||||
);
|
||||
assert_eq!(Config::default().mcp_tools, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_defaults_match_expected() {
|
||||
let cfg = Config::default();
|
||||
|
||||
@@ -30,6 +30,7 @@ pub trait RoleLike {
|
||||
fn top_p(&self) -> Option<f64>;
|
||||
fn enabled_tools(&self) -> Option<Vec<String>>;
|
||||
fn enabled_mcp_servers(&self) -> Option<Vec<String>>;
|
||||
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>>;
|
||||
fn set_model(&mut self, model: Model);
|
||||
fn set_temperature(&mut self, value: Option<f64>);
|
||||
fn reasoning_effort(&self) -> Option<String>;
|
||||
@@ -37,6 +38,7 @@ pub trait RoleLike {
|
||||
fn set_reasoning_effort(&mut self, value: Option<String>);
|
||||
fn set_enabled_tools(&mut self, value: Option<Vec<String>>);
|
||||
fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>);
|
||||
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
@@ -67,6 +69,8 @@ pub struct Role {
|
||||
deserialize_with = "super::deserialize_csv_or_vec"
|
||||
)]
|
||||
enabled_mcp_servers: Option<Vec<String>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
skills_enabled: Option<bool>,
|
||||
#[serde(
|
||||
@@ -133,6 +137,7 @@ impl Role {
|
||||
"enabled_mcp_servers" => {
|
||||
role.enabled_mcp_servers = parse_string_or_array(value)
|
||||
}
|
||||
"mcp_tools" => role.mcp_tools = parse_mcp_tools_map(value),
|
||||
"skills_enabled" => role.skills_enabled = value.as_bool(),
|
||||
"enabled_skills" => role.enabled_skills = parse_string_or_array(value),
|
||||
"enabled_macros" => role.enabled_macros = parse_string_or_array(value),
|
||||
@@ -196,6 +201,10 @@ impl Role {
|
||||
serde_json::to_string(enabled_mcp_servers).unwrap_or_else(|_| "[]".to_string());
|
||||
metadata.push(format!("enabled_mcp_servers: {inline}"));
|
||||
}
|
||||
if let Some(mcp_tools) = &self.mcp_tools {
|
||||
let inline = serde_json::to_string(mcp_tools).unwrap_or_else(|_| "{}".to_string());
|
||||
metadata.push(format!("mcp_tools: {inline}"));
|
||||
}
|
||||
if let Some(skills_enabled) = self.skills_enabled {
|
||||
metadata.push(format!("skills_enabled: {skills_enabled}"));
|
||||
}
|
||||
@@ -279,6 +288,10 @@ impl Role {
|
||||
enabled_tools,
|
||||
enabled_mcp_servers,
|
||||
);
|
||||
let mcp_tools = role_like.mcp_tools();
|
||||
if mcp_tools.is_some() {
|
||||
self.set_mcp_tools(mcp_tools);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn batch_set(
|
||||
@@ -453,6 +466,10 @@ impl RoleLike for Role {
|
||||
self.enabled_mcp_servers.clone()
|
||||
}
|
||||
|
||||
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>> {
|
||||
self.mcp_tools.clone()
|
||||
}
|
||||
|
||||
fn set_model(&mut self, model: Model) {
|
||||
if !self.model().id().is_empty() {
|
||||
self.model_id = Some(model.id().to_string());
|
||||
@@ -479,6 +496,10 @@ impl RoleLike for Role {
|
||||
fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>) {
|
||||
self.enabled_mcp_servers = value;
|
||||
}
|
||||
|
||||
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>) {
|
||||
self.mcp_tools = value;
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_or_array(value: &Value) -> Option<Vec<String>> {
|
||||
@@ -503,6 +524,19 @@ fn parse_string_or_array(value: &Value) -> Option<Vec<String>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_mcp_tools_map(value: &Value) -> Option<IndexMap<String, Vec<String>>> {
|
||||
let map = value.as_object()?;
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
for (server, tools) in map {
|
||||
if tools.is_null() {
|
||||
mcp_tools.insert(server.clone(), Vec::new());
|
||||
} else if let Some(tools) = parse_string_or_array(tools) {
|
||||
mcp_tools.insert(server.clone(), tools);
|
||||
}
|
||||
}
|
||||
Some(mcp_tools)
|
||||
}
|
||||
|
||||
fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) {
|
||||
let mut text = prompt;
|
||||
let mut search_input = true;
|
||||
@@ -652,6 +686,69 @@ mod tests {
|
||||
assert_eq!(role.enabled_macros, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_new_parses_mcp_tools_list_and_csv_values() {
|
||||
let content = "---\nmcp_tools:\n github: [get_*, list_*, search_code]\n slack: conversations_history,conversations_replies\n---\nPrompt";
|
||||
|
||||
let role = Role::new("test", content);
|
||||
|
||||
let mcp_tools = role.mcp_tools().unwrap();
|
||||
assert_eq!(
|
||||
mcp_tools.get("github"),
|
||||
Some(&vec![
|
||||
"get_*".to_string(),
|
||||
"list_*".to_string(),
|
||||
"search_code".to_string()
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
mcp_tools.get("slack"),
|
||||
Some(&vec![
|
||||
"conversations_history".to_string(),
|
||||
"conversations_replies".to_string()
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_new_mcp_tools_empty_list_server_is_some_empty() {
|
||||
let role = Role::new("test", "---\nmcp_tools:\n github: []\n---\nPrompt");
|
||||
|
||||
assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_new_mcp_tools_per_server_null_is_some_empty() {
|
||||
let role = Role::new("test", "---\nmcp_tools:\n github:\n---\nPrompt");
|
||||
|
||||
assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_new_mcp_tools_absent_is_none() {
|
||||
let role = Role::new("test", "---\ntemperature: 0.5\n---\nPrompt");
|
||||
|
||||
assert_eq!(role.mcp_tools(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_new_mcp_tools_null_is_none() {
|
||||
let role = Role::new("test", "---\nmcp_tools: null\n---\nPrompt");
|
||||
|
||||
assert_eq!(role.mcp_tools(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_export_mcp_tools_round_trips() {
|
||||
let content = "---\nmcp_tools:\n github: [get_issue]\n slack: a,b\n---\nPrompt";
|
||||
let role = Role::new("test", content);
|
||||
|
||||
let reparsed = Role::new("test", &role.export());
|
||||
|
||||
assert_eq!(reparsed.mcp_tools(), role.mcp_tools());
|
||||
assert!(role.mcp_tools().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn role_export_includes_enabled_macros() {
|
||||
let role = Role::new("test", "---\nenabled_macros: [a]\n---\nPrompt");
|
||||
|
||||
@@ -40,6 +40,8 @@ pub struct Session {
|
||||
deserialize_with = "super::deserialize_csv_or_vec"
|
||||
)]
|
||||
enabled_mcp_servers: Option<Vec<String>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
skills_enabled: Option<bool>,
|
||||
#[serde(
|
||||
@@ -249,6 +251,9 @@ impl Session {
|
||||
if let Some(enabled_mcp_servers) = self.enabled_mcp_servers() {
|
||||
data["enabled_mcp_servers"] = json!(enabled_mcp_servers);
|
||||
}
|
||||
if let Some(mcp_tools) = self.mcp_tools() {
|
||||
data["mcp_tools"] = json!(mcp_tools);
|
||||
}
|
||||
if let Some(skills_enabled) = self.skills_enabled() {
|
||||
data["skills_enabled"] = skills_enabled.into();
|
||||
}
|
||||
@@ -329,6 +334,13 @@ impl Session {
|
||||
items.push(("enabled_mcp_servers", enabled_mcp_servers.join(",")));
|
||||
}
|
||||
|
||||
if let Some(mcp_tools) = self.mcp_tools() {
|
||||
items.push((
|
||||
"mcp_tools",
|
||||
serde_json::to_string(&mcp_tools).unwrap_or_default(),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some(skills_enabled) = self.skills_enabled() {
|
||||
items.push(("skills_enabled", skills_enabled.to_string()));
|
||||
}
|
||||
@@ -870,6 +882,10 @@ impl RoleLike for Session {
|
||||
self.enabled_mcp_servers.clone()
|
||||
}
|
||||
|
||||
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>> {
|
||||
self.mcp_tools.clone()
|
||||
}
|
||||
|
||||
fn set_model(&mut self, model: Model) {
|
||||
if self.model().id() != model.id() {
|
||||
self.model_id = model.id();
|
||||
@@ -913,6 +929,13 @@ impl RoleLike for Session {
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
|
||||
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>) {
|
||||
if self.mcp_tools != value {
|
||||
self.mcp_tools = value;
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
@@ -1044,6 +1067,45 @@ mod tests {
|
||||
assert_eq!(session.enabled_macros, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_mcp_tools_survives_yaml_round_trip() {
|
||||
let mut session = Session::default();
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]);
|
||||
mcp_tools.insert("slack".to_string(), vec![]);
|
||||
session.set_mcp_tools(Some(mcp_tools.clone()));
|
||||
|
||||
let yaml = serde_yaml::to_string(&session).unwrap();
|
||||
let reloaded: Session = serde_yaml::from_str(&yaml).unwrap();
|
||||
|
||||
assert_eq!(reloaded.mcp_tools(), Some(mcp_tools));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_set_role_does_not_copy_mcp_tools() {
|
||||
let role = Role::new(
|
||||
"test",
|
||||
"---\nmcp_tools:\n github: [get_issue]\n---\nPrompt",
|
||||
);
|
||||
assert!(role.mcp_tools().is_some());
|
||||
let mut session = Session::default();
|
||||
|
||||
session.set_role(role);
|
||||
|
||||
assert_eq!(session.mcp_tools(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_set_mcp_tools_marks_dirty() {
|
||||
let mut session = Session::default();
|
||||
assert!(!session.dirty());
|
||||
|
||||
session.set_mcp_tools(Some(IndexMap::new()));
|
||||
|
||||
assert!(session.dirty());
|
||||
assert_eq!(session.mcp_tools(), Some(IndexMap::new()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_enabled_macros_empty_list_is_some_empty() {
|
||||
let session: Session =
|
||||
|
||||
@@ -37,6 +37,8 @@ pub struct Skill {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
enabled_mcp_servers: Option<Vec<String>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
auto_unload: Option<bool>,
|
||||
}
|
||||
|
||||
@@ -74,6 +76,9 @@ impl Skill {
|
||||
"enabled_mcp_servers" => {
|
||||
skill.enabled_mcp_servers = parse_skill_string_or_array(value);
|
||||
}
|
||||
"mcp_tools" => {
|
||||
skill.mcp_tools = parse_skill_mcp_tools_map(value);
|
||||
}
|
||||
"auto_unload" => {
|
||||
skill.auto_unload = value.as_bool();
|
||||
}
|
||||
@@ -147,6 +152,11 @@ impl Skill {
|
||||
self.enabled_mcp_servers.as_deref()
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn mcp_tools(&self) -> Option<&IndexMap<String, Vec<String>>> {
|
||||
self.mcp_tools.as_ref()
|
||||
}
|
||||
|
||||
pub fn auto_unload(&self) -> bool {
|
||||
self.auto_unload.unwrap_or(false)
|
||||
}
|
||||
@@ -185,6 +195,19 @@ fn parse_skill_string_or_array(value: &Value) -> Option<Vec<String>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_skill_mcp_tools_map(value: &Value) -> Option<IndexMap<String, Vec<String>>> {
|
||||
let map = value.as_object()?;
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
for (server, tools) in map {
|
||||
if tools.is_null() {
|
||||
mcp_tools.insert(server.clone(), Vec::new());
|
||||
} else if let Some(tools) = parse_skill_string_or_array(tools) {
|
||||
mcp_tools.insert(server.clone(), tools);
|
||||
}
|
||||
}
|
||||
Some(mcp_tools)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -198,6 +221,44 @@ mod tests {
|
||||
assert_eq!(skill.description(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_new_parses_mcp_tools_list_and_csv_values() {
|
||||
let content = "---\nmcp_tools:\n github: [get_*, list_*]\n slack: a,b\n---\nBody";
|
||||
|
||||
let skill = Skill::new("test", content);
|
||||
|
||||
let mcp_tools = skill.mcp_tools().unwrap();
|
||||
assert_eq!(
|
||||
mcp_tools.get("github"),
|
||||
Some(&vec!["get_*".to_string(), "list_*".to_string()])
|
||||
);
|
||||
assert_eq!(
|
||||
mcp_tools.get("slack"),
|
||||
Some(&vec!["a".to_string(), "b".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_new_mcp_tools_absent_is_none() {
|
||||
let skill = Skill::new("test", "---\ndescription: d\n---\nBody");
|
||||
|
||||
assert_eq!(skill.mcp_tools(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_new_mcp_tools_empty_list_server_is_some_empty() {
|
||||
let skill = Skill::new("test", "---\nmcp_tools:\n github: []\n---\nBody");
|
||||
|
||||
assert_eq!(skill.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_new_mcp_tools_per_server_null_is_some_empty() {
|
||||
let skill = Skill::new("test", "---\nmcp_tools:\n github:\n---\nBody");
|
||||
|
||||
assert_eq!(skill.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_new_parses_full_metadata() {
|
||||
let content = "---\n\
|
||||
|
||||
@@ -506,6 +506,7 @@ mod tests {
|
||||
instructions: Some("sys".into()),
|
||||
prompt: "user".into(),
|
||||
tools: None,
|
||||
mcp_tools: None,
|
||||
model: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
|
||||
@@ -37,6 +37,9 @@ pub struct Graph {
|
||||
#[serde(default)]
|
||||
pub mcp_servers: Vec<String>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub skills_enabled: Option<bool>,
|
||||
|
||||
@@ -285,6 +288,9 @@ pub struct LlmNode {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<String>>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
|
||||
@@ -530,6 +536,47 @@ nodes:
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserializes_mcp_tools_at_graph_and_node_level() {
|
||||
let yaml = r#"
|
||||
name: g
|
||||
mcp_servers: [github]
|
||||
mcp_tools:
|
||||
github:
|
||||
- get_*
|
||||
- list_*
|
||||
start: l
|
||||
nodes:
|
||||
l:
|
||||
id: l
|
||||
type: llm
|
||||
prompt: hi
|
||||
mcp_tools:
|
||||
github:
|
||||
- search_code
|
||||
next: e
|
||||
e:
|
||||
id: e
|
||||
type: end
|
||||
output: done
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
graph.mcp_tools.as_ref().unwrap().get("github"),
|
||||
Some(&vec!["get_*".to_string(), "list_*".to_string()])
|
||||
);
|
||||
match &graph.get_node("l").unwrap().node_type {
|
||||
NodeType::Llm(llm) => {
|
||||
assert_eq!(
|
||||
llm.mcp_tools.as_ref().unwrap().get("github"),
|
||||
Some(&vec!["search_code".to_string()])
|
||||
);
|
||||
}
|
||||
_ => panic!("expected Llm variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserializes_every_node_type() {
|
||||
let yaml = r#"
|
||||
|
||||
@@ -214,6 +214,19 @@ impl GraphValidator {
|
||||
return;
|
||||
};
|
||||
|
||||
let expand_alias = |name: &str| {
|
||||
ctx.app_config
|
||||
.mapping_mcp_servers
|
||||
.get(name)
|
||||
.into_iter()
|
||||
.flat_map(|mapped| mapped.split(',').map(|s| s.trim().to_string()))
|
||||
};
|
||||
let mut enabled_servers: HashSet<String> = ctx.mcp_servers.clone();
|
||||
for server in &ctx.mcp_servers {
|
||||
enabled_servers.extend(expand_alias(server));
|
||||
}
|
||||
let all_servers_enabled = ctx.mcp_servers.iter().any(|s| s.trim() == "all");
|
||||
|
||||
for (node_id, node) in &graph.nodes {
|
||||
let NodeType::Llm(llm) = &node.node_type else {
|
||||
continue;
|
||||
@@ -237,6 +250,23 @@ impl GraphValidator {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(mcp_tools) = &llm.mcp_tools
|
||||
&& !all_servers_enabled
|
||||
{
|
||||
for key in mcp_tools.keys() {
|
||||
let enabled = enabled_servers.contains(key)
|
||||
|| expand_alias(key).any(|id| enabled_servers.contains(&id));
|
||||
if !enabled {
|
||||
result.error(ValidationError::with_node(
|
||||
node_id,
|
||||
format!(
|
||||
"llm node 'mcp_tools' references MCP server '{key}' not enabled by this graph"
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(model_id) = &llm.model
|
||||
&& Model::retrieve_model(ctx.app_config.as_ref(), model_id, ModelType::Chat)
|
||||
.is_err()
|
||||
@@ -1001,6 +1031,7 @@ mod tests {
|
||||
max_concurrent_jobs: None,
|
||||
global_tools: Vec::new(),
|
||||
mcp_servers: Vec::new(),
|
||||
mcp_tools: None,
|
||||
skills_enabled: None,
|
||||
enabled_skills: None,
|
||||
inject_skill_instructions: None,
|
||||
@@ -1099,6 +1130,7 @@ mod tests {
|
||||
instructions: None,
|
||||
prompt: "p".into(),
|
||||
tools: None,
|
||||
mcp_tools: None,
|
||||
model: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
@@ -1257,6 +1289,19 @@ mod tests {
|
||||
node
|
||||
}
|
||||
|
||||
fn llm_node_with_mcp_tools(id: &str, servers: &[&str]) -> Node {
|
||||
let mut node = llm_node(id, None, Some("end"));
|
||||
if let NodeType::Llm(ref mut n) = node.node_type {
|
||||
let mut mcp_tools = IndexMap::new();
|
||||
for server in servers {
|
||||
mcp_tools.insert(server.to_string(), vec!["get_*".to_string()]);
|
||||
}
|
||||
n.mcp_tools = Some(mcp_tools);
|
||||
}
|
||||
|
||||
node
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_unknown_tool_is_an_error() {
|
||||
let graph = graph_with(
|
||||
@@ -1340,6 +1385,103 @@ mod tests {
|
||||
assert!(result.is_valid());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_mcp_tools_enabled_server_passes() {
|
||||
let graph = graph_with(
|
||||
vec![
|
||||
("l", llm_node_with_mcp_tools("l", &["github"])),
|
||||
("end", end_node("end")),
|
||||
],
|
||||
"l",
|
||||
);
|
||||
|
||||
let result = validator()
|
||||
.with_agent_context(agent_ctx(&[], &["github"]))
|
||||
.validate(&graph);
|
||||
|
||||
assert!(result.is_valid());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_mcp_tools_unknown_server_is_an_error() {
|
||||
let graph = graph_with(
|
||||
vec![
|
||||
("l", llm_node_with_mcp_tools("l", &["slack"])),
|
||||
("end", end_node("end")),
|
||||
],
|
||||
"l",
|
||||
);
|
||||
|
||||
let result = validator()
|
||||
.with_agent_context(agent_ctx(&[], &["github"]))
|
||||
.validate(&graph);
|
||||
|
||||
assert!(!result.is_valid());
|
||||
assert!(
|
||||
result
|
||||
.errors
|
||||
.iter()
|
||||
.any(|e| e.message.contains("'slack' not enabled"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_mcp_tools_alias_key_passes() {
|
||||
let graph = graph_with(
|
||||
vec![
|
||||
("l", llm_node_with_mcp_tools("l", &["gh"])),
|
||||
("end", end_node("end")),
|
||||
],
|
||||
"l",
|
||||
);
|
||||
let mut ctx = agent_ctx(&[], &["github-mcp"]);
|
||||
let mut app = AppConfig::default();
|
||||
app.mapping_mcp_servers
|
||||
.insert("gh".to_string(), "github-mcp".to_string());
|
||||
ctx.app_config = Arc::new(app);
|
||||
|
||||
let result = validator().with_agent_context(ctx).validate(&graph);
|
||||
|
||||
assert!(result.is_valid(), "errors: {:?}", result.errors);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_mcp_tools_key_matching_alias_expansion_passes() {
|
||||
let graph = graph_with(
|
||||
vec![
|
||||
("l", llm_node_with_mcp_tools("l", &["github-mcp"])),
|
||||
("end", end_node("end")),
|
||||
],
|
||||
"l",
|
||||
);
|
||||
let mut ctx = agent_ctx(&[], &["gh"]);
|
||||
let mut app = AppConfig::default();
|
||||
app.mapping_mcp_servers
|
||||
.insert("gh".to_string(), "github-mcp".to_string());
|
||||
ctx.app_config = Arc::new(app);
|
||||
|
||||
let result = validator().with_agent_context(ctx).validate(&graph);
|
||||
|
||||
assert!(result.is_valid(), "errors: {:?}", result.errors);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_mcp_tools_with_all_sentinel_passes() {
|
||||
let graph = graph_with(
|
||||
vec![
|
||||
("l", llm_node_with_mcp_tools("l", &["github"])),
|
||||
("end", end_node("end")),
|
||||
],
|
||||
"l",
|
||||
);
|
||||
|
||||
let result = validator()
|
||||
.with_agent_context(agent_ctx(&[], &["all"]))
|
||||
.validate(&graph);
|
||||
|
||||
assert!(result.is_valid(), "errors: {:?}", result.errors);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llm_node_unknown_model_is_an_error() {
|
||||
let graph = graph_with(
|
||||
|
||||
Reference in New Issue
Block a user