[WIP] Check AI models access to MCP servers and fix any issues#3
[WIP] Check AI models access to MCP servers and fix any issues#3Godzilla675 merged 3 commits intore-writefrom
Conversation
Co-authored-by: Godzilla675 <131464726+Godzilla675@users.noreply.github.com>
Co-authored-by: Godzilla675 <131464726+Godzilla675@users.noreply.github.com>
|
@copilot continue |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for MCP (Model Context Protocol) servers to enable tool usage by the AI models. The changes are comprehensive, touching upon dependency injection, view models, services, and the data layer to integrate this new functionality. A new McpToolMapper is introduced to translate tool definitions for the model, and a unit test for it has been added, which is great.
The overall implementation is solid and follows good practices for integrating such a feature in an Android app. I've identified one high-severity issue related to state management in the ChatViewModel that could lead to inconsistencies. My review includes a suggested refactoring to address this, which will also improve the code's maintainability and robustness.
| try { | ||
| val enabledServers = mcpServerRepository.getEnabledServers().first() | ||
| if (enabledServers.isEmpty()) { | ||
| mcpToolRegistry = emptyMap() | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val serverTools = mutableMapOf<McpServer, List<com.dark.tool_neuron.service.McpToolInfo>>() | ||
| enabledServers.forEach { server -> | ||
| val tools = mcpClientService.listTools(server) | ||
| if (tools.isNotEmpty()) { | ||
| serverTools[server] = tools | ||
| } | ||
| } | ||
|
|
||
| if (serverTools.isEmpty()) { | ||
| mcpToolRegistry = emptyMap() | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val mapping = McpToolMapper.buildMapping(serverTools) | ||
| mcpToolRegistry = mapping.toolRegistry | ||
|
|
||
| if (mapping.toolRegistry.isEmpty()) { | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val success = LlmModelWorker.setGgufToolsJson(mapping.toolsJson) | ||
| if (!success) { | ||
| LlmModelWorker.clearGgufTools() | ||
| } | ||
| } catch (e: Exception) { | ||
| val message = "Failed to refresh MCP tools: ${e.message}" | ||
| _error.value = message | ||
| AppStateManager.setError(message) | ||
| } |
There was a problem hiding this comment.
The syncMcpTools function contains repetitive logic for clearing tool configurations and has a bug that can lead to an inconsistent state. If LlmModelWorker.setGgufToolsJson() fails, the GGUF tools are cleared, but mcpToolRegistry in the ViewModel is not. This mismatch could cause runtime errors when trying to execute a tool.
I recommend refactoring this function to improve clarity, fix the bug, and enhance robustness. You can achieve this by extracting the tool-clearing logic into a private helper function and restructuring syncMcpTools to handle success, failure, and exceptions more cleanly.
You'll need to add the following helper function to ChatViewModel:
private fun clearMcpToolConfiguration() {
mcpToolRegistry = emptyMap()
LlmModelWorker.clearGgufTools()
}Then, you can use the suggested code below to replace the body of the syncMcpTools function.
try {
val enabledServers = mcpServerRepository.getEnabledServers().first()
if (enabledServers.isEmpty()) {
clearMcpToolConfiguration()
return
}
val serverTools = enabledServers.mapNotNull { server ->
val tools = mcpClientService.listTools(server)
if (tools.isNotEmpty()) server to tools else null
}.toMap()
if (serverTools.isEmpty()) {
clearMcpToolConfiguration()
return
}
val mapping = McpToolMapper.buildMapping(serverTools)
if (mapping.toolRegistry.isEmpty()) {
clearMcpToolConfiguration()
return
}
if (LlmModelWorker.setGgufToolsJson(mapping.toolsJson)) {
mcpToolRegistry = mapping.toolRegistry
} else {
clearMcpToolConfiguration()
}
} catch (e: Exception) {
val message = "Failed to refresh MCP tools: ${e.message}"
_error.value = message
AppStateManager.setError(message)
clearMcpToolConfiguration()
}There was a problem hiding this comment.
Pull request overview
This PR implements MCP (Model Context Protocol) tool integration to enable AI models to access and execute tools from connected MCP servers. The implementation adds tool mapping, synchronization with GGUF models, and handling of tool calls during chat generation.
Changes:
- Added
McpToolMapperutility to convert MCP server tools into GGUF-compatible JSON format with prefixed tool names - Integrated MCP tool synchronization into
ChatViewModelto configure tools before each generation - Extended
LlmModelWorkerwith methods to set and clear GGUF tool configurations
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 15 comments.
Show a summary per file
| File | Description |
|---|---|
| gradle/libs.versions.toml | Added org.json dependency for JSON processing in tests |
| app/build.gradle.kts | Added test dependencies for JUnit and org.json |
| app/src/test/java/com/dark/tool_neuron/service/McpToolMapperTest.kt | New unit test for McpToolMapper with basic happy path coverage |
| app/src/main/java/com/dark/tool_neuron/service/McpToolMapper.kt | New utility to map MCP tools to GGUF tool definitions with sanitized identifiers |
| app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt | Added overloaded callTool method accepting JSON string arguments and refactored internal implementation |
| app/src/main/java/com/dark/tool_neuron/worker/LlmModelWorker.kt | Added setGgufToolsJson and clearGgufTools methods for tool configuration |
| app/src/main/java/com/dark/tool_neuron/viewmodel/ChatViewModel.kt | Integrated MCP tool synchronization, tool call detection, and execution into chat generation flow |
| app/src/main/java/com/dark/tool_neuron/viewmodel/factory/ChatViewModelFactory.kt | Updated factory with MCP dependencies (unused since ChatViewModel uses Hilt) |
| app/src/main/java/com/dark/tool_neuron/di/AppContainer.kt | Added MCP repository and service initialization (unused since ChatViewModel uses Hilt) |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| private suspend fun syncMcpTools() { | ||
| try { | ||
| val enabledServers = mcpServerRepository.getEnabledServers().first() | ||
| if (enabledServers.isEmpty()) { | ||
| mcpToolRegistry = emptyMap() | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } |
There was a problem hiding this comment.
Before syncing MCP tools, check if a GGUF model is actually loaded using LlmModelWorker.isGgufModelLoaded. If no model is loaded, either skip tool configuration or show an appropriate message to the user. This would prevent the silent failure when setGgufToolsJson is called before a model is loaded.
| fun buildMapping(serverTools: Map<McpServer, List<McpToolInfo>>): McpToolMapping { | ||
| val toolsArray = JSONArray() | ||
| val registry = mutableMapOf<String, McpToolReference>() | ||
|
|
||
| serverTools.forEach { (server, tools) -> | ||
| val serverPrefix = sanitizeIdentifier(server.name).ifBlank { "mcp" } | ||
| tools.forEach { tool -> | ||
| val toolSlug = sanitizeIdentifier(tool.name).ifBlank { "tool" } | ||
| val toolId = "${serverPrefix}_${toolSlug}" | ||
| toolsArray.put(buildToolDefinition(toolId, tool)) | ||
| registry[toolId] = McpToolReference(server, tool.name) | ||
| } | ||
| } | ||
|
|
||
| return McpToolMapping( | ||
| toolsJson = toolsArray.toString(), | ||
| toolRegistry = registry | ||
| ) | ||
| } |
There was a problem hiding this comment.
The sanitizeIdentifier function could produce duplicate tool IDs if two tools with different names normalize to the same identifier (e.g., "send-email" and "send_email" both become "send_email"). This would cause the later tool to silently overwrite the earlier one in the registry. Consider detecting and handling collisions by appending a numeric suffix or throwing an error.
| mcpServerRepository = McpServerRepository(database.mcpServerDao()) | ||
| mcpClientService = McpClientService() | ||
|
|
||
| llmModelViewModelFactory = LLMModelViewModelFactory(application, modelRepository) | ||
| chatListViewModelFactory = ChatListViewModelFactory(chatManager) | ||
| chatViewModelFactory = ChatViewModelFactory(chatManager, generationManager) | ||
| chatViewModelFactory = ChatViewModelFactory( | ||
| chatManager, | ||
| generationManager, | ||
| mcpServerRepository, | ||
| mcpClientService | ||
| ) |
There was a problem hiding this comment.
The ChatViewModelFactory changes in AppContainer appear to be unused since ChatViewModel is obtained via Hilt's hiltViewModel() in MainActivity (line 99). The factory pattern here is redundant when using @hiltviewmodel. Consider removing these factory-related changes to avoid confusion about which dependency injection mechanism is being used.
| currentGeneratedContent = "" | ||
| tokenBuffer.clear() | ||
| tokenCount = 0 | ||
| _streamingAssistantMessage.value = "" |
There was a problem hiding this comment.
The tool call handling logic clears the streaming message but doesn't inform the user that a tool is being executed. Consider showing a loading state or status message to the user (e.g., "Executing tool...") before calling the tool, so users understand why the response is being generated.
| _streamingAssistantMessage.value = "" | |
| _streamingAssistantMessage.value = "Executing tool..." |
| } catch (e: Exception) { | ||
| val message = "Failed to refresh MCP tools: ${e.message}" | ||
| _error.value = message | ||
| AppStateManager.setError(message) | ||
| } |
There was a problem hiding this comment.
The error handling in syncMcpTools catches all exceptions and sets an error state, but doesn't log which specific operation failed (fetching enabled servers, listing tools, building mapping, or setting GGUF tools). Consider adding more granular error logging with specific context about where the failure occurred to aid debugging.
| } | ||
|
|
||
| val success = LlmModelWorker.setGgufToolsJson(mapping.toolsJson) | ||
| if (!success) { |
There was a problem hiding this comment.
The setToolsJson function in GGUFEngine returns false if the model isn't loaded, but this failure case isn't logged or communicated to the user. When syncMcpTools is called before the GGUF model is loaded, the tools will silently fail to configure. Consider checking if the model is loaded before calling setGgufToolsJson, or deferring tool configuration until after model load, or at minimum, logging this specific failure case.
| if (!success) { | |
| if (!success) { | |
| val message = "Failed to configure MCP tools: GGUF model may not be loaded." | |
| _error.value = message | |
| AppStateManager.setError(message) |
| onFailure = { error -> | ||
| val message = "Tool ${toolCall.name} failed: ${error.message ?: "Unknown error"}" | ||
| _error.value = message | ||
| AppStateManager.setError(message) |
There was a problem hiding this comment.
When a tool call fails, the error message is set but the generation state is not properly reset. The _isGenerating flag and other streaming state variables should be reset via resetStreamingState() to allow the user to retry or continue with a new message.
| AppStateManager.setError(message) | |
| AppStateManager.setError(message) | |
| AppStateManager.setGenerationComplete() | |
| resetStreamingState() |
| class ChatViewModelFactory( | ||
| private val chatManager: ChatManager, private val generationManager: GenerationManager | ||
| private val chatManager: ChatManager, | ||
| private val generationManager: GenerationManager, | ||
| private val mcpServerRepository: McpServerRepository, | ||
| private val mcpClientService: McpClientService | ||
| ) : ViewModelProvider.Factory { | ||
| @Suppress("UNCHECKED_CAST") | ||
| override fun <T : ViewModel> create(modelClass: Class<T>): T { | ||
| if (modelClass.isAssignableFrom(ChatViewModel::class.java)) { | ||
| return ChatViewModel(chatManager, generationManager) as T | ||
| return ChatViewModel( | ||
| chatManager, | ||
| generationManager, | ||
| mcpServerRepository, | ||
| mcpClientService | ||
| ) as T | ||
| } | ||
| throw IllegalArgumentException("Unknown ViewModel class") | ||
| } | ||
| } No newline at end of file | ||
| } |
There was a problem hiding this comment.
The ChatViewModelFactory is not being used since ChatViewModel is obtained via Hilt's hiltViewModel() in MainActivity. These changes to the factory are unnecessary when using @hiltviewmodel annotation. Consider removing or deprecating this factory class to avoid maintaining two different dependency injection approaches.
| suspend fun setGgufToolsJson(toolsJson: String): Boolean { | ||
| val svc = ensureServiceBound() | ||
| return try { | ||
| svc.setToolsJsonGguf(toolsJson) | ||
| } catch (e: Exception) { | ||
| Log.e(TAG, "Failed to set GGUF tools JSON", e) | ||
| false | ||
| } | ||
| } |
There was a problem hiding this comment.
The setGgufToolsJson function returns a Boolean indicating success/failure, but the error is only logged. When this returns false in ChatViewModel.syncMcpTools, the tools are cleared but the user isn't informed why the tools weren't configured. Consider propagating the error message or providing more specific feedback about why tool configuration failed.
| private suspend fun syncMcpTools() { | ||
| try { | ||
| val enabledServers = mcpServerRepository.getEnabledServers().first() | ||
| if (enabledServers.isEmpty()) { | ||
| mcpToolRegistry = emptyMap() | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val serverTools = mutableMapOf<McpServer, List<com.dark.tool_neuron.service.McpToolInfo>>() | ||
| enabledServers.forEach { server -> | ||
| val tools = mcpClientService.listTools(server) | ||
| if (tools.isNotEmpty()) { | ||
| serverTools[server] = tools | ||
| } | ||
| } | ||
|
|
||
| if (serverTools.isEmpty()) { | ||
| mcpToolRegistry = emptyMap() | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val mapping = McpToolMapper.buildMapping(serverTools) | ||
| mcpToolRegistry = mapping.toolRegistry | ||
|
|
||
| if (mapping.toolRegistry.isEmpty()) { | ||
| LlmModelWorker.clearGgufTools() | ||
| return | ||
| } | ||
|
|
||
| val success = LlmModelWorker.setGgufToolsJson(mapping.toolsJson) | ||
| if (!success) { | ||
| LlmModelWorker.clearGgufTools() | ||
| } | ||
| } catch (e: Exception) { | ||
| val message = "Failed to refresh MCP tools: ${e.message}" | ||
| _error.value = message | ||
| AppStateManager.setError(message) | ||
| } | ||
| } |
There was a problem hiding this comment.
The syncMcpTools function lacks logging for successful tool synchronization. Add log statements to track: how many servers are enabled, how many tools are loaded, whether the tools JSON was successfully set, and what the actual tool names are. This would help diagnose why the AI models might not be able to access the tools.
Original prompt
💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.