[WIP] Check AI models access to MCP servers and fix any issues

app/build.gradle.kts

@@ -143,6 +143,10 @@ dependencies {
 
     // Debug
     debugImplementation(libs.androidx.compose.ui.tooling)
+
+    // Tests
+    testImplementation(libs.junit)
+    testImplementation(libs.org.json)
 }
 
 fun getProperty(value: String): String {
@@ -154,4 +158,4 @@ fun getProperty(value: String): String {
     } else {
         System.getenv(value) ?: "\"sample_val\""
     }
-}
+}

app/src/main/java/com/dark/tool_neuron/di/AppContainer.kt

@@ -4,7 +4,9 @@ import android.app.Application
 import android.content.Context
 import com.dark.tool_neuron.database.AppDatabase
 import com.dark.tool_neuron.repo.ChatRepository
+import com.dark.tool_neuron.repo.McpServerRepository
 import com.dark.tool_neuron.repo.ModelRepository
+import com.dark.tool_neuron.service.McpClientService
 import com.dark.tool_neuron.vault.VaultHelper
 import com.dark.tool_neuron.viewmodel.factory.ChatListViewModelFactory
 import com.dark.tool_neuron.viewmodel.factory.ChatViewModelFactory
@@ -21,6 +23,8 @@ object AppContainer {
     private lateinit var database: AppDatabase
     private lateinit var modelRepository: ModelRepository
     private lateinit var chatRepository: ChatRepository
+    private lateinit var mcpServerRepository: McpServerRepository
+    private lateinit var mcpClientService: McpClientService
     private lateinit var llmModelViewModelFactory: LLMModelViewModelFactory
     private lateinit var chatListViewModelFactory: ChatListViewModelFactory
     private lateinit var chatViewModelFactory: ChatViewModelFactory
@@ -38,10 +42,17 @@ object AppContainer {
         )
 
         chatRepository = ChatRepository()
+        mcpServerRepository = McpServerRepository(database.mcpServerDao())
+        mcpClientService = McpClientService()
 
         llmModelViewModelFactory = LLMModelViewModelFactory(application, modelRepository)
         chatListViewModelFactory = ChatListViewModelFactory(chatManager)
-        chatViewModelFactory = ChatViewModelFactory(chatManager, generationManager)
+        chatViewModelFactory = ChatViewModelFactory(
+            chatManager,
+            generationManager,
+            mcpServerRepository,
+            mcpClientService
+        )
 
         initVault(context)
     }
@@ -84,4 +95,4 @@ object AppContainer {
 
 
     fun isVaultReady(): Boolean = vaultInitialized
-}
+}

app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt

@@ -241,7 +241,7 @@ class McpClientService @Inject constructor() {
     /**
      * List available tools from an MCP server
      */
-    private suspend fun listTools(server: McpServer): List<McpToolInfo> = withContext(Dispatchers.IO) {
+    suspend fun listTools(server: McpServer): List<McpToolInfo> = withContext(Dispatchers.IO) {
         try {
             val listToolsRequest = JSONObject().apply {
                 put("jsonrpc", "2.0")
@@ -303,6 +303,25 @@ class McpClientService @Inject constructor() {
         server: McpServer,
         toolName: String,
         arguments: Map<String, Any>
+    ): Result<String> = callToolInternal(server, toolName, JSONObject(arguments))
+
+    suspend fun callTool(
+        server: McpServer,
+        toolName: String,
+        argumentsJson: String
+    ): Result<String> {
+        val parsedArguments = try {
+            if (argumentsJson.isBlank()) JSONObject() else JSONObject(argumentsJson)
+        } catch (e: Exception) {
+            return Result.failure(Exception("Invalid tool arguments JSON: ${e.message}"))
+        }
+        return callToolInternal(server, toolName, parsedArguments)
+    }
+
+    private suspend fun callToolInternal(
+        server: McpServer,
+        toolName: String,
+        arguments: JSONObject
     ): Result<String> = withContext(Dispatchers.IO) {
         try {
             val callToolRequest = JSONObject().apply {
@@ -311,7 +330,7 @@ class McpClientService @Inject constructor() {
                 put("method", "tools/call")
                 put("params", JSONObject().apply {
                     put("name", toolName)
-                    put("arguments", JSONObject(arguments))
+                    put("arguments", arguments)
                 })
             }
 
@@ -326,7 +345,7 @@ class McpClientService @Inject constructor() {
             }
 
             val response = httpClient.newCall(requestBuilder.build()).execute()
-            
+
             if (!response.isSuccessful) {
                 return@withContext Result.failure(Exception("Server returned: ${response.code}"))
             }
@@ -344,11 +363,11 @@ class McpClientService @Inject constructor() {
             }
 
             val result = jsonResponse.optJSONObject("result")
-            Result.success(result?.toString() ?: responseBody)
-            
+            return@withContext Result.success(result?.toString() ?: responseBody)
+
         } catch (e: Exception) {
             Log.e(TAG, "Failed to call tool: ${e.message}", e)
-            Result.failure(e)
+            return@withContext Result.failure(e)
         }
     }
 }

app/src/main/java/com/dark/tool_neuron/service/McpToolMapper.kt

@@ -0,0 +1,68 @@
+package com.dark.tool_neuron.service
+
+import com.dark.tool_neuron.models.table_schema.McpServer
+import org.json.JSONArray
+import org.json.JSONObject
+
+data class McpToolReference(
+    val server: McpServer,
+    val toolName: String
+)
+
+data class McpToolMapping(
+    val toolsJson: String,
+    val toolRegistry: Map<String, McpToolReference>
+)
+
+object McpToolMapper {
+    fun sanitizeIdentifier(value: String): String {
+        return value.lowercase()
+            .replace(Regex("[^a-z0-9]+"), "_")
+            .trim('_')
+    }
+
+    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
+        )
+    }
+
+    private fun buildToolDefinition(toolId: String, tool: McpToolInfo): JSONObject {
+        val function = JSONObject().apply {
+            put("name", toolId)
+            tool.description?.takeIf { it.isNotBlank() }?.let { put("description", it) }
+            put("parameters", buildParameters(tool.inputSchema))
+        }
+
+        return JSONObject().apply {
+            put("type", "function")
+            put("function", function)
+        }
+    }
+
+    private fun buildParameters(inputSchema: String?): JSONObject {
+        val parsedSchema = inputSchema?.takeIf { it.isNotBlank() }?.let {
+            runCatching { JSONObject(it) }.getOrNull()
+        }
+
+        return (parsedSchema ?: JSONObject()).apply {
+            if (!has("type")) {
+                put("type", "object")
+            }
+        }
+    }
+}