feat(mcp): improve MCP server handling and tool information retrieval

- Refactor CustomMcpServerManager to store server information in a map
- Update McpFunctionProvider to handle multiple servers and tools
- Enhance toolInfos function to properly assign mcpGroup based on server name
- Optimize isApplicable and execute functions for better tool lookup

Author: phodal

core/src/main/kotlin/cc/unitmesh/devti/mcp/client/CustomMcpServerManager.kt

@@ -30,18 +30,19 @@ import com.intellij.execution.configurations.GeneralCommandLine
 
 @Service(Service.Level.PROJECT)
 class CustomMcpServerManager(val project: Project) {
-    val cached = mutableMapOf<String, List<Tool>>()
+    val cached = mutableMapOf<String, Map<String, List<Tool>>>()
     val toolClientMap = mutableMapOf<Tool, Client>()
 
-    suspend fun collectServerInfos(): List<Tool> {
+    suspend fun collectServerInfos(): Map<String, List<Tool>> {
         val mcpServerConfig = project.customizeSetting.mcpServerConfig
-        if (mcpServerConfig.isEmpty()) return emptyList()
+        if (mcpServerConfig.isEmpty()) return emptyMap()
         if (cached.containsKey(mcpServerConfig)) return cached[mcpServerConfig]!!
         val mcpConfig = McpServer.load(mcpServerConfig)
-        if (mcpConfig == null) return emptyList()
+        if (mcpConfig == null) return emptyMap()
 
-        val tools: List<Tool> = mcpConfig.mcpServers.map { entry ->
-            if (entry.value.disabled == true) return@map emptyList<Tool>()
+        val toolsMap = mutableMapOf<String, List<Tool>>()
+        mcpConfig.mcpServers.forEach { entry ->
+            if (entry.value.disabled == true) return@forEach
             val resolvedCommand = resolveCommand(entry.value.command)
             logger<CustomMcpServerManager>().info("Found MCP command: $resolvedCommand")
             val client = Client(clientInfo = Implementation(name = entry.key, version = "1.0.0"))
@@ -62,16 +63,15 @@ class CustomMcpServerManager(val project: Project) {
                 }
                 listTools?.tools ?: emptyList()
             } catch (e: Exception) {
-                throw e
                 logger<CustomMcpServerManager>().warn("Failed to list tools from ${entry.key}: $e")
                 emptyList<Tool>()
             }
+            
+            toolsMap[entry.key] = tools
+        }
 
-            return@map tools
-        }.flatMap { it }
-
-        cached[mcpServerConfig] = tools
-        return tools
+        cached[mcpServerConfig] = toolsMap
+        return toolsMap
     }
 
     fun execute(project: Project, tool: Tool, map: String): String {
@@ -165,3 +165,4 @@ fun resolveCommand(command: String): String {
     }
     return command
 }
+

core/src/main/kotlin/cc/unitmesh/devti/mcp/client/McpFunctionProvider.kt

@@ -4,34 +4,48 @@ import cc.unitmesh.devti.agent.tool.AgentTool
 import cc.unitmesh.devti.provider.toolchain.ToolchainFunctionProvider
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectManager
+import io.modelcontextprotocol.kotlin.sdk.Tool
 import io.modelcontextprotocol.kotlin.sdk.Tool.Input
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.encodeToString
 
 class McpFunctionProvider : ToolchainFunctionProvider {
     override suspend fun funcNames(): List<String> {
         val project = ProjectManager.getInstance().openProjects.firstOrNull() ?: return emptyList()
-        return CustomMcpServerManager.instance(project).collectServerInfos().map { it.name }
+        return CustomMcpServerManager.instance(project).collectServerInfos().values
+            .flatMap { it }
+            .map { it.name }
+            .distinct()
     }
 
     override suspend fun toolInfos(project: Project): List<AgentTool> {
         val manager = CustomMcpServerManager.instance(project)
-        return manager.collectServerInfos().map {
-            val schemaJson = Json.encodeToString<Input>(it.inputSchema)
-            val mockData = Json.encodeToString(MockDataGenerator.generateMockData(it.inputSchema))
-            AgentTool(
-                it.name,
-                it.description ?: "",
-                "Here is command and JSON schema\n/${it.name}\n```json\n$schemaJson\n```",
-                isMcp = true,
-                mcpGroup = it.name,
-                completion = mockData
-            )
+        val toolsMap = manager.collectServerInfos()
+        
+        val agentTools = mutableListOf<AgentTool>()
+        for ((serverName, tools) in toolsMap) {
+            for (tool in tools) {
+                val schemaJson = Json.encodeToString<Input>(tool.inputSchema)
+                val mockData = Json.encodeToString(MockDataGenerator.generateMockData(tool.inputSchema))
+                agentTools.add(
+                    AgentTool(
+                        tool.name,
+                        tool.description ?: "",
+                        "Here is command and JSON schema\n/${tool.name}\n```json\n$schemaJson\n```",
+                        isMcp = true,
+                        mcpGroup = serverName,
+                        completion = mockData
+                    )
+                )
+            }
         }
+        
+        return agentTools
     }
 
     override suspend fun isApplicable(project: Project, funcName: String): Boolean {
-        return CustomMcpServerManager.instance(project).collectServerInfos().any { it.name == funcName }
+        val toolsMap = CustomMcpServerManager.instance(project).collectServerInfos()
+        return toolsMap.any { (_, tools) -> tools.any { it.name == funcName } }
     }
 
     override suspend fun execute(
@@ -41,7 +55,10 @@ class McpFunctionProvider : ToolchainFunctionProvider {
         allVariables: Map<String, Any?>,
         commandName: String
     ): Any {
-        val tool = CustomMcpServerManager.instance(project).collectServerInfos().firstOrNull { it.name == commandName }
+        val toolsMap = CustomMcpServerManager.instance(project).collectServerInfos()
+        val tool = toolsMap.values.flatMap { it }
+            .firstOrNull { it.name == commandName }
+            
         if (tool == null) {
             return "No MCP such tool: $prop"
         }