diff --git a/src/cortex-agents/src/agent.rs b/src/cortex-agents/src/agent.rs index fbaefb6c0..aab1fdf86 100644 --- a/src/cortex-agents/src/agent.rs +++ b/src/cortex-agents/src/agent.rs @@ -127,7 +127,19 @@ impl AgentInfo { } pub fn is_tool_enabled(&self, tool: &str) -> bool { - self.tools.get(tool).copied().unwrap_or(true) + if let Some(enabled) = self.tools.get(tool) { + return *enabled; + } + + if tool.eq_ignore_ascii_case("mcp") { + return self.tools.get("mcp").copied().unwrap_or(true); + } + + if tool.to_ascii_lowercase().starts_with("mcp__") { + return self.tools.get("mcp").copied().unwrap_or(true); + } + + true } pub fn with_max_steps(mut self, max_steps: usize) -> Self { diff --git a/src/cortex-agents/src/custom/registry.rs b/src/cortex-agents/src/custom/registry.rs index 222ca49e2..d9f8c8f5c 100644 --- a/src/cortex-agents/src/custom/registry.rs +++ b/src/cortex-agents/src/custom/registry.rs @@ -184,6 +184,20 @@ fn custom_agent_to_agent_info(agent: &CustomAgentConfig) -> AgentInfo { } } + for tool in &allowed { + if tool.to_ascii_lowercase().starts_with("mcp__") { + info = info.enable_tool(tool); + } + } + + if !matches!( + agent.tools, + ToolsConfig::Category(ToolCategory::Mcp | ToolCategory::All) + ) && !allowed.iter().any(|t| t.eq_ignore_ascii_case("mcp")) + { + info = info.disable_tool("mcp"); + } + info } @@ -378,6 +392,63 @@ mod tests { assert!(info.model.is_none()); // Should not have a model set } + #[test] + fn test_read_only_agent_disables_mcp_tools() { + let mut registry = CustomAgentRegistry::new(); + registry.register(CustomAgentConfig::new("readonly").with_tools(ToolsConfig::read_only())); + + let info = registry.to_agent_info("readonly").unwrap(); + + assert!(!info.is_tool_enabled("mcp")); + assert!(!info.is_tool_enabled("mcp__postgres__query")); + assert!(!info.is_tool_enabled("MCP__postgres__query")); + assert!(info.is_tool_enabled("read")); + } + + #[test] + fn test_mcp_and_all_tool_configs_allow_mcp_tools() { + let mut registry = CustomAgentRegistry::new(); + registry.register( + CustomAgentConfig::new("mcp-agent") + .with_tools(ToolsConfig::Category(ToolCategory::Mcp)), + ); + registry.register(CustomAgentConfig::new("all-agent").with_tools(ToolsConfig::all())); + + let mcp_info = registry.to_agent_info("mcp-agent").unwrap(); + assert!(mcp_info.is_tool_enabled("mcp__postgres__query")); + + let all_info = registry.to_agent_info("all-agent").unwrap(); + assert!(all_info.is_tool_enabled("mcp__postgres__query")); + } + + #[test] + fn test_custom_tools_list_can_allow_mcp_tools() { + let mut registry = CustomAgentRegistry::new(); + registry.register( + CustomAgentConfig::new("custom-mcp").with_tools(ToolsConfig::custom(["Read", "mcp"])), + ); + + let info = registry.to_agent_info("custom-mcp").unwrap(); + + assert!(info.is_tool_enabled("mcp__postgres__query")); + assert!(info.is_tool_enabled("Read")); + assert!(!info.is_tool_enabled("Execute")); + } + + #[test] + fn test_custom_tools_list_can_allow_single_mcp_tool() { + let mut registry = CustomAgentRegistry::new(); + registry.register( + CustomAgentConfig::new("single-mcp") + .with_tools(ToolsConfig::custom(["Read", "mcp__postgres__query"])), + ); + + let info = registry.to_agent_info("single-mcp").unwrap(); + + assert!(info.is_tool_enabled("mcp__postgres__query")); + assert!(!info.is_tool_enabled("mcp__postgres__write")); + } + #[test] fn test_tools_to_permission() { // Read-only should have read_only permission