fix(agents): apply custom tool limits to mcp tools

Author: OnlyYu1996

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 {

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