From c2e79bf233e0ccb750183e05e73329155578a748 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 17 Jan 2026 17:02:32 +0000 Subject: [PATCH 1/3] Initial plan From 148d9794c8614ada97e85643a0fb54ba0ee2a32f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 17 Jan 2026 17:30:36 +0000 Subject: [PATCH 2/3] Integrate MCP tools with GGUF Co-authored-by: Godzilla675 <131464726+Godzilla675@users.noreply.github.com> --- app/build.gradle.kts | 6 +- .../com/dark/tool_neuron/di/AppContainer.kt | 15 +- .../tool_neuron/service/McpClientService.kt | 37 ++++- .../dark/tool_neuron/service/McpToolMapper.kt | 68 ++++++++ .../tool_neuron/viewmodel/ChatViewModel.kt | 148 +++++++++++++++++- .../viewmodel/factory/ChatViewModelFactory.kt | 16 +- .../dark/tool_neuron/worker/LlmModelWorker.kt | 21 ++- .../tool_neuron/service/McpToolMapperTest.kt | 38 +++++ gradle/libs.versions.toml | 4 +- 9 files changed, 333 insertions(+), 20 deletions(-) create mode 100644 app/src/main/java/com/dark/tool_neuron/service/McpToolMapper.kt create mode 100644 app/src/test/java/com/dark/tool_neuron/service/McpToolMapperTest.kt diff --git a/app/build.gradle.kts b/app/build.gradle.kts index a0366e91..f3b0a1b8 100644 --- a/app/build.gradle.kts +++ b/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\"" } -} \ No newline at end of file +} diff --git a/app/src/main/java/com/dark/tool_neuron/di/AppContainer.kt b/app/src/main/java/com/dark/tool_neuron/di/AppContainer.kt index 687be60e..14e24fbf 100644 --- a/app/src/main/java/com/dark/tool_neuron/di/AppContainer.kt +++ b/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 -} \ No newline at end of file +} diff --git a/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt b/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt index 63b87140..82632d53 100644 --- a/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt +++ b/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 = withContext(Dispatchers.IO) { + suspend fun listTools(server: McpServer): List = withContext(Dispatchers.IO) { try { val listToolsRequest = JSONObject().apply { put("jsonrpc", "2.0") @@ -304,6 +304,27 @@ class McpClientService @Inject constructor() { toolName: String, arguments: Map ): Result = withContext(Dispatchers.IO) { + return@withContext callToolInternal(server, toolName, JSONObject(arguments)) + } + + suspend fun callTool( + server: McpServer, + toolName: String, + argumentsJson: String + ): Result = withContext(Dispatchers.IO) { + val parsedArguments = try { + if (argumentsJson.isBlank()) JSONObject() else JSONObject(argumentsJson) + } catch (e: Exception) { + return@withContext Result.failure(Exception("Invalid tool arguments JSON: ${e.message}")) + } + return@withContext callToolInternal(server, toolName, parsedArguments) + } + + private fun callToolInternal( + server: McpServer, + toolName: String, + arguments: JSONObject + ): Result { try { val callToolRequest = JSONObject().apply { put("jsonrpc", "2.0") @@ -311,7 +332,7 @@ class McpClientService @Inject constructor() { put("method", "tools/call") put("params", JSONObject().apply { put("name", toolName) - put("arguments", JSONObject(arguments)) + put("arguments", arguments) }) } @@ -328,11 +349,11 @@ class McpClientService @Inject constructor() { val response = httpClient.newCall(requestBuilder.build()).execute() if (!response.isSuccessful) { - return@withContext Result.failure(Exception("Server returned: ${response.code}")) + return Result.failure(Exception("Server returned: ${response.code}")) } val rawResponseBody = response.body?.string() - ?: return@withContext Result.failure(Exception("Empty response")) + ?: return Result.failure(Exception("Empty response")) // Parse response based on transport type val responseBody = parseResponse(rawResponseBody, server.transportType) @@ -340,15 +361,15 @@ class McpClientService @Inject constructor() { if (jsonResponse.has("error")) { val error = jsonResponse.getJSONObject("error") - return@withContext Result.failure(Exception(error.optString("message", "Unknown error"))) + return Result.failure(Exception(error.optString("message", "Unknown error"))) } val result = jsonResponse.optJSONObject("result") - Result.success(result?.toString() ?: responseBody) - + return Result.success(result?.toString() ?: responseBody) + } catch (e: Exception) { Log.e(TAG, "Failed to call tool: ${e.message}", e) - Result.failure(e) + return Result.failure(e) } } } diff --git a/app/src/main/java/com/dark/tool_neuron/service/McpToolMapper.kt b/app/src/main/java/com/dark/tool_neuron/service/McpToolMapper.kt new file mode 100644 index 00000000..7495c1d8 --- /dev/null +++ b/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 +) + +object McpToolMapper { + fun sanitizeIdentifier(value: String): String { + return value.lowercase() + .replace(Regex("[^a-z0-9]+"), "_") + .trim('_') + } + + fun buildMapping(serverTools: Map>): McpToolMapping { + val toolsArray = JSONArray() + val registry = mutableMapOf() + + 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") + } + } + } +} diff --git a/app/src/main/java/com/dark/tool_neuron/viewmodel/ChatViewModel.kt b/app/src/main/java/com/dark/tool_neuron/viewmodel/ChatViewModel.kt index 0029b221..c0c3c929 100644 --- a/app/src/main/java/com/dark/tool_neuron/viewmodel/ChatViewModel.kt +++ b/app/src/main/java/com/dark/tool_neuron/viewmodel/ChatViewModel.kt @@ -13,7 +13,12 @@ import com.dark.tool_neuron.models.messages.MessageContent import com.dark.tool_neuron.models.messages.Messages import com.dark.tool_neuron.models.messages.RagResultItem import com.dark.tool_neuron.models.messages.Role +import com.dark.tool_neuron.models.table_schema.McpServer import com.dark.tool_neuron.models.table_schema.ModelConfig +import com.dark.tool_neuron.repo.McpServerRepository +import com.dark.tool_neuron.service.McpClientService +import com.dark.tool_neuron.service.McpToolMapper +import com.dark.tool_neuron.service.McpToolReference import com.dark.tool_neuron.state.AppStateManager import com.dark.tool_neuron.worker.ChatManager import com.dark.tool_neuron.worker.DiffusionConfig @@ -26,12 +31,15 @@ import jakarta.inject.Inject import kotlinx.coroutines.Job import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.StateFlow +import kotlinx.coroutines.flow.first import kotlinx.coroutines.launch @HiltViewModel class ChatViewModel @Inject constructor( private val chatManager: ChatManager, - private val generationManager: GenerationManager + private val generationManager: GenerationManager, + private val mcpServerRepository: McpServerRepository, + private val mcpClientService: McpClientService ) : ViewModel() { private val _messages = mutableStateListOf() @@ -103,6 +111,13 @@ class ChatViewModel @Inject constructor( private val _currentRagResults = MutableStateFlow>(emptyList()) val currentRagResults: StateFlow> = _currentRagResults + private data class ToolCallInfo( + val name: String, + val argsJson: String + ) + + private var mcpToolRegistry: Map = emptyMap() + // ==================== RAG Controls ==================== fun setRagEnabled(enabled: Boolean) { @@ -240,11 +255,13 @@ class ChatViewModel @Inject constructor( val tokenBatchSize = 3 try { + var pendingToolCall: ToolCallInfo? = null // Prepend RAG context if available val finalPrompt = _currentRagContext.value?.let { ragContext -> "$ragContext\n\n### User Query:\n$prompt" } ?: prompt + syncMcpTools() generationManager.generateTextStreaming(finalPrompt, maxTokens).collect { event -> when (event) { is GenerationEvent.Token -> { @@ -265,6 +282,11 @@ class ChatViewModel @Inject constructor( } is GenerationEvent.Done -> { + val toolCall = pendingToolCall + if (toolCall != null) { + handleToolCallForNewChat(prompt, toolCall) + return@collect + } _streamingAssistantMessage.value = currentGeneratedContent // Don't set _isGenerating.value = false here // It will be set in resetStreamingState() after messages are added @@ -279,7 +301,13 @@ class ChatViewModel @Inject constructor( currentMetrics = event.metrics } - is GenerationEvent.ToolCall -> {} + is GenerationEvent.ToolCall -> { + pendingToolCall = ToolCallInfo(event.name, event.args) + currentGeneratedContent = "" + tokenBuffer.clear() + tokenCount = 0 + _streamingAssistantMessage.value = "" + } } } } catch (e: Exception) { @@ -305,6 +333,7 @@ class ChatViewModel @Inject constructor( val tokenBatchSize = 3 try { + var pendingToolCall: ToolCallInfo? = null var conversationPrompt = generationManager.buildConversationPrompt( _messages, userMessage.content.content ) @@ -314,6 +343,7 @@ class ChatViewModel @Inject constructor( conversationPrompt = "$ragContext\n\n$conversationPrompt" } + syncMcpTools() generationManager.generateTextStreaming(conversationPrompt, maxTokens) .collect { event -> when (event) { @@ -335,6 +365,11 @@ class ChatViewModel @Inject constructor( } is GenerationEvent.Done -> { + val toolCall = pendingToolCall + if (toolCall != null) { + handleToolCallExistingChat(chatId, userMessage, toolCall) + return@collect + } _streamingAssistantMessage.value = currentGeneratedContent // Add user message first if not already added @@ -383,7 +418,13 @@ class ChatViewModel @Inject constructor( currentMetrics = event.metrics } - is GenerationEvent.ToolCall -> {} + is GenerationEvent.ToolCall -> { + pendingToolCall = ToolCallInfo(event.name, event.args) + currentGeneratedContent = "" + tokenBuffer.clear() + tokenCount = 0 + _streamingAssistantMessage.value = "" + } } } } catch (e: Exception) { @@ -524,6 +565,105 @@ class ChatViewModel @Inject constructor( } } + // ==================== MCP Tool Integration ==================== + + private suspend fun syncMcpTools() { + try { + val enabledServers = mcpServerRepository.getEnabledServers().first() + if (enabledServers.isEmpty()) { + mcpToolRegistry = emptyMap() + LlmModelWorker.clearGgufTools() + return + } + + val serverTools = mutableMapOf>() + 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) + } + } + + private suspend fun handleToolCallForNewChat(prompt: String, toolCall: ToolCallInfo) { + val response = resolveToolCallResponse(toolCall) + _streamingAssistantMessage.value = response + createChatWithMessages(prompt, response, null) + } + + private suspend fun handleToolCallExistingChat( + chatId: String, + userMessage: Messages, + toolCall: ToolCallInfo + ) { + val response = resolveToolCallResponse(toolCall) + _streamingAssistantMessage.value = response + + if (!userMessageAdded) { + _messages.add(userMessage) + userMessageAdded = true + } + + val assistantMessage = Messages( + role = Role.Assistant, + content = MessageContent( + contentType = ContentType.Text, + content = response + ) + ) + _messages.add(assistantMessage) + + chatManager.addAssistantMessage(chatId, response, null) + AppStateManager.setGenerationComplete() + resetStreamingState() + } + + private suspend fun resolveToolCallResponse(toolCall: ToolCallInfo): String { + return executeToolCall(toolCall).fold( + onSuccess = { result -> formatToolResult(toolCall.name, result) }, + onFailure = { error -> + val message = "Tool ${toolCall.name} failed: ${error.message ?: "Unknown error"}" + _error.value = message + AppStateManager.setError(message) + message + } + ) + } + + private suspend fun executeToolCall(toolCall: ToolCallInfo): Result { + val reference = mcpToolRegistry[toolCall.name] + ?: return Result.failure(Exception("Tool not found: ${toolCall.name}")) + return mcpClientService.callTool(reference.server, reference.toolName, toolCall.argsJson) + } + + private fun formatToolResult(toolName: String, result: String): String { + return "Tool $toolName result:\n$result" + } + private fun generateImageForNewChat( prompt: String, negativePrompt: String, @@ -1018,4 +1158,4 @@ class ChatViewModel @Inject constructor( fun hideModelList() { _showModelList.value = false } -} \ No newline at end of file +} diff --git a/app/src/main/java/com/dark/tool_neuron/viewmodel/factory/ChatViewModelFactory.kt b/app/src/main/java/com/dark/tool_neuron/viewmodel/factory/ChatViewModelFactory.kt index a134d583..e9b39c66 100644 --- a/app/src/main/java/com/dark/tool_neuron/viewmodel/factory/ChatViewModelFactory.kt +++ b/app/src/main/java/com/dark/tool_neuron/viewmodel/factory/ChatViewModelFactory.kt @@ -2,18 +2,28 @@ package com.dark.tool_neuron.viewmodel.factory import androidx.lifecycle.ViewModel import androidx.lifecycle.ViewModelProvider +import com.dark.tool_neuron.repo.McpServerRepository +import com.dark.tool_neuron.service.McpClientService import com.dark.tool_neuron.viewmodel.ChatViewModel import com.dark.tool_neuron.worker.ChatManager import com.dark.tool_neuron.worker.GenerationManager 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 create(modelClass: Class): 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 +} diff --git a/app/src/main/java/com/dark/tool_neuron/worker/LlmModelWorker.kt b/app/src/main/java/com/dark/tool_neuron/worker/LlmModelWorker.kt index ffbb9c93..7edeb1c2 100644 --- a/app/src/main/java/com/dark/tool_neuron/worker/LlmModelWorker.kt +++ b/app/src/main/java/com/dark/tool_neuron/worker/LlmModelWorker.kt @@ -309,6 +309,25 @@ object LlmModelWorker { return service?.modelInfoGguf } + 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 + } + } + + fun clearGgufTools() { + try { + service?.clearToolsGguf() + Log.i(TAG, "Cleared GGUF tools") + } catch (e: Exception) { + Log.e(TAG, "Failed to clear GGUF tools", e) + } + } + // ==================== Diffusion Methods ==================== /** @@ -618,4 +637,4 @@ object LlmModelWorker { Log.i(TAG, "Embedding model download started in background") } -} \ No newline at end of file +} diff --git a/app/src/test/java/com/dark/tool_neuron/service/McpToolMapperTest.kt b/app/src/test/java/com/dark/tool_neuron/service/McpToolMapperTest.kt new file mode 100644 index 00000000..2f975569 --- /dev/null +++ b/app/src/test/java/com/dark/tool_neuron/service/McpToolMapperTest.kt @@ -0,0 +1,38 @@ +package com.dark.tool_neuron.service + +import com.dark.tool_neuron.models.table_schema.McpServer +import com.dark.tool_neuron.models.table_schema.McpTransportType +import org.json.JSONArray +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNotNull +import org.junit.Test + +class McpToolMapperTest { + @Test + fun buildMappingCreatesToolRegistry() { + val server = McpServer( + id = "server-1", + name = "Zapier MCP", + url = "https://example.com/mcp", + transportType = McpTransportType.SSE + ) + val tool = McpToolInfo( + name = "send-email", + description = "Send an email", + inputSchema = """{"type":"object","properties":{"to":{"type":"string"}}}""" + ) + + val mapping = McpToolMapper.buildMapping(mapOf(server to listOf(tool))) + val toolsArray = JSONArray(mapping.toolsJson) + + assertEquals(1, toolsArray.length()) + val function = toolsArray.getJSONObject(0).getJSONObject("function") + assertEquals("zapier_mcp_send_email", function.getString("name")) + assertEquals("object", function.getJSONObject("parameters").getString("type")) + + val reference = mapping.toolRegistry["zapier_mcp_send_email"] + assertNotNull(reference) + assertEquals(server, reference?.server) + assertEquals("send-email", reference?.toolName) + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index f709a8d7..d3811ee3 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -33,6 +33,7 @@ xz = "1.11" androidx-espresso-core = "3.7.0" androidx-junit = "1.3.0" junit = "4.13.2" +org-json = "20240303" [libraries] androidx-activity-compose = { group = "androidx.activity", name = "activity-compose", version.ref = "androidx-activity-compose" } @@ -77,6 +78,7 @@ xz = { group = "org.tukaani", name = "xz", version.ref = "xz" } androidx-espresso-core = { group = "androidx.test.espresso", name = "espresso-core", version.ref = "androidx-espresso-core" } androidx-junit = { group = "androidx.test.ext", name = "junit", version.ref = "androidx-junit" } junit = { group = "junit", name = "junit", version.ref = "junit" } +org-json = { group = "org.json", name = "json", version.ref = "org-json" } [plugins] android-application = { id = "com.android.application", version.ref = "agp" } @@ -85,4 +87,4 @@ google-dagger-hilt = { id = "com.google.dagger.hilt.android", version.ref = "dag kotlin-android = { id = "org.jetbrains.kotlin.android", version.ref = "kotlin" } kotlin-compose = { id = "org.jetbrains.kotlin.plugin.compose", version.ref = "kotlin" } kotlin-ksp = { id = "com.google.devtools.ksp", version.ref = "ksp" } -kotlin-serialization = { id = "org.jetbrains.kotlin.plugin.serialization", version.ref = "kotlin" } \ No newline at end of file +kotlin-serialization = { id = "org.jetbrains.kotlin.plugin.serialization", version.ref = "kotlin" } From 89dde9bf01bfea44b3c4a0a5a188fba19e092535 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 17 Jan 2026 17:32:53 +0000 Subject: [PATCH 3/3] Refine MCP tool call execution Co-authored-by: Godzilla675 <131464726+Godzilla675@users.noreply.github.com> --- .../tool_neuron/service/McpClientService.kt | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt b/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt index 82632d53..52564f98 100644 --- a/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt +++ b/app/src/main/java/com/dark/tool_neuron/service/McpClientService.kt @@ -303,28 +303,26 @@ class McpClientService @Inject constructor() { server: McpServer, toolName: String, arguments: Map - ): Result = withContext(Dispatchers.IO) { - return@withContext callToolInternal(server, toolName, JSONObject(arguments)) - } + ): Result = callToolInternal(server, toolName, JSONObject(arguments)) suspend fun callTool( server: McpServer, toolName: String, argumentsJson: String - ): Result = withContext(Dispatchers.IO) { + ): Result { val parsedArguments = try { if (argumentsJson.isBlank()) JSONObject() else JSONObject(argumentsJson) } catch (e: Exception) { - return@withContext Result.failure(Exception("Invalid tool arguments JSON: ${e.message}")) + return Result.failure(Exception("Invalid tool arguments JSON: ${e.message}")) } - return@withContext callToolInternal(server, toolName, parsedArguments) + return callToolInternal(server, toolName, parsedArguments) } - private fun callToolInternal( + private suspend fun callToolInternal( server: McpServer, toolName: String, arguments: JSONObject - ): Result { + ): Result = withContext(Dispatchers.IO) { try { val callToolRequest = JSONObject().apply { put("jsonrpc", "2.0") @@ -347,13 +345,13 @@ class McpClientService @Inject constructor() { } val response = httpClient.newCall(requestBuilder.build()).execute() - + if (!response.isSuccessful) { - return Result.failure(Exception("Server returned: ${response.code}")) + return@withContext Result.failure(Exception("Server returned: ${response.code}")) } val rawResponseBody = response.body?.string() - ?: return Result.failure(Exception("Empty response")) + ?: return@withContext Result.failure(Exception("Empty response")) // Parse response based on transport type val responseBody = parseResponse(rawResponseBody, server.transportType) @@ -361,15 +359,15 @@ class McpClientService @Inject constructor() { if (jsonResponse.has("error")) { val error = jsonResponse.getJSONObject("error") - return Result.failure(Exception(error.optString("message", "Unknown error"))) + return@withContext Result.failure(Exception(error.optString("message", "Unknown error"))) } val result = jsonResponse.optJSONObject("result") - return 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) - return Result.failure(e) + return@withContext Result.failure(e) } } }