From 63e8b4f5d63bea3fca53a5c7301f832c6d3a2014 Mon Sep 17 00:00:00 2001 From: kami Date: Thu, 16 Jul 2026 19:26:01 +0400 Subject: [PATCH] =?UTF-8?q?feat(tools):=20native=20MCP=20client=20?= =?UTF-8?q?=E2=80=94=20mount=20stdio=20MCP=20servers=20as=20Correx=20tools?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an MCP host layer so external MCP servers (AST/LSP code-intelligence, package resolvers, etc.) can be mounted at startup and surface their tools/list as first-class Correx tools named mcp____. Each MCP tool is a Tool+ToolExecutor, so it rides the normal ToolExecutor path and inherits tier gating, receipts, and event-recorded execution (#5) with zero special-casing; replay reads the recorded receipt rather than re-calling the server (#8/#9). - McpProtocol/McpStdioClient/McpTool/McpMounter in infrastructure:tools (stdio JSON-RPC 2.0, tools/* slice only). Capabilities empty; safety via default T2 tier since external side effects are opaque. - [[mcp]] config (id/command/env/tier) + array-of-tables parsing. - Main wires mounted servers into extraTools (main + per-workspace paths) with a shutdown hook; a server that fails to start is logged and skipped. - Tests: in-process fake transport (handshake/list/call/error + tool mapping), [[mcp]] parser. Co-Authored-By: Claude Opus 4.8 --- .../kotlin/com/correx/apps/server/Main.kt | 54 ++++-- .../com/correx/core/config/ConfigLoader.kt | 29 ++- .../com/correx/core/config/CorrexConfig.kt | 13 ++ .../correx/core/config/ConfigLoaderTest.kt | 25 +++ .../infrastructure/tools/mcp/McpMounter.kt | 40 +++++ .../infrastructure/tools/mcp/McpProtocol.kt | 68 +++++++ .../tools/mcp/McpStdioClient.kt | 169 ++++++++++++++++++ .../infrastructure/tools/mcp/McpTool.kt | 64 +++++++ .../infrastructure/tools/mcp/McpClientTest.kt | 112 ++++++++++++ 9 files changed, 563 insertions(+), 11 deletions(-) create mode 100644 infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpMounter.kt create mode 100644 infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpProtocol.kt create mode 100644 infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpStdioClient.kt create mode 100644 infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpTool.kt create mode 100644 infrastructure/tools/src/test/kotlin/com/correx/infrastructure/tools/mcp/McpClientTest.kt diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 1a3ba9e0..9d802d56 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -210,14 +210,21 @@ fun main() { val explicitWorkspaceRoot = System.getenv("CORREX_WORKSPACE_ROOT") ?.let { Path.of(it) } ?: toolsConfig.workspaceRoot.takeIf { it.isNotEmpty() }?.let { Path.of(it) } - // workingDir and workspaceRoot must resolve to the same tree by default — the tool-call - // assessor's containment rules (PathContainmentRule, ManifestContainmentRule) only ever see - // workspaceRoot, while FileWriteTool resolves relative paths against workingDir. If only one - // is configured, each falls back to the other before falling back to process CWD, so the - // assessor's containment check and the actual write always share one resolved root. val shellAllowedExecutables = toolsConfig.shellAllowedExecutables.toSet() - val workspaceRoot = explicitWorkspaceRoot ?: explicitWorkingDir ?: Path.of("").toAbsolutePath() - val workingDir = explicitWorkingDir ?: workspaceRoot + val bootWorkspace = resolveBootWorkspace( + explicitWorkspaceRoot = explicitWorkspaceRoot, + explicitWorkingDir = explicitWorkingDir, + processWorkingDir = Path.of(""), + ) + val workspaceRoot = bootWorkspace.workspaceRoot + val workingDir = bootWorkspace.workingDir + if (bootWorkspace.workingDirWasClamped) { + log.warn( + "configured working_dir {} is outside workspace_root {}; clamping working_dir to workspace_root", + explicitWorkingDir, + workspaceRoot, + ) + } // One shared HTTP client backs both the default and per-workspace registries' research tools // (web_search/web_fetch). Built only when research is enabled, so the static path stays offline. // Lives for the process lifetime (shared across requests), so it's closed via shutdown hook @@ -262,7 +269,34 @@ fun main() { ) // Retrieves full tool output the kernel spilled to CAS on truncation (ref shown in-context). val toolOutputTool = com.correx.infrastructure.tools.ToolOutputTool(artifactStore) - val extraTools = taskTools + toolOutputTool + // Mount configured MCP servers: each server's tools/list becomes Correx tools that ride the normal + // ToolExecutor path (tier-gated, event-recorded, replay-safe — no special-casing). A server that + // fails to start is logged and skipped rather than aborting the whole process. + val mountedMcpServers = correxConfig.mcp.mapNotNull { mcpConfig -> + runCatching { + runBlocking { + com.correx.infrastructure.tools.mcp.McpMounter.mount( + serverId = mcpConfig.id, + command = mcpConfig.command, + env = mcpConfig.env, + tier = runCatching { com.correx.core.approvals.Tier.valueOf(mcpConfig.tier) } + .getOrDefault(com.correx.core.approvals.Tier.T2), + ) + } + }.onSuccess { log.info("Mounted MCP server '{}' ({} tools)", mcpConfig.id, it.tools.size) } + .onFailure { log.warn("Failed to mount MCP server '{}': {}", mcpConfig.id, it.message) } + .getOrNull() + } + if (mountedMcpServers.isNotEmpty()) { + Runtime.getRuntime().addShutdownHook(Thread { + mountedMcpServers.forEach { server -> + runCatching { server.close() } + .onFailure { log.warn("Error closing MCP server '{}': {}", server.serverId, it.message) } + } + }) + } + val mcpTools = mountedMcpServers.flatMap { it.tools } + val extraTools = taskTools + toolOutputTool + mcpTools val toolRegistry = InfrastructureModule.createToolRegistry( buildToolConfig( workspaceRoot, @@ -545,7 +579,7 @@ fun main() { fun buildProjectMemory(cfg: CorrexConfig): com.correx.apps.server.memory.ProjectMemoryService? = if (cfg.project.enabled) { com.correx.apps.server.memory.ProjectMemoryService( - config = cfg.project, + config = cfg.project.boundToWorkspace(workspaceRoot), embedder = embedder, l3MemoryStore = l3MemoryStore, journalRepository = decisionJournalRepository, @@ -871,7 +905,7 @@ private fun buildToolConfig( toolsConfig: com.correx.core.config.ToolsConfig, research: com.correx.infrastructure.tools.ResearchToolConfig, ): ToolConfig { - val allowed = setOf(workspaceRoot, workingDir) + val allowed = setOf(workspaceRoot) return ToolConfig( shell = ShellConfig( enabled = toolsConfig.shellEnabled, diff --git a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt index 4864508a..5ecaf9a4 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt @@ -265,9 +265,11 @@ object ConfigLoader { val providers = mutableListOf>() val models = mutableListOf>() val artifacts = mutableListOf>() + val mcpServers = mutableListOf>() var currentProvider: MutableMap? = null var currentModel: MutableMap? = null var currentArtifact: MutableMap? = null + var currentMcp: MutableMap? = null // Flush any open array-of-table entry into its list. Called whenever a new // table header (array or section) starts, so the previous entry is committed. @@ -275,9 +277,11 @@ object ConfigLoader { currentProvider?.let { providers.add(it) } currentModel?.let { models.add(it) } currentArtifact?.let { artifacts.add(it) } + currentMcp?.let { mcpServers.add(it) } currentProvider = null currentModel = null currentArtifact = null + currentMcp = null } for ((lineNum, line) in lines.withIndex()) { @@ -302,6 +306,11 @@ object ConfigLoader { currentArtifact = mutableMapOf() currentSection = "" } + trimmed == "[[mcp]]" -> { + flushTables() + currentMcp = mutableMapOf() + currentSection = "" + } trimmed.startsWith("[") && !trimmed.startsWith("[[") && trimmed.endsWith("]") -> { // Parse section headers like [server] or [tools.shell] flushTables() @@ -320,10 +329,12 @@ object ConfigLoader { val provider = currentProvider val model = currentModel val artifact = currentArtifact + val mcp = currentMcp when { provider != null -> provider[key] = parsedValue model != null -> model[key] = parsedValue artifact != null -> artifact[key] = parsedValue + mcp != null -> mcp[key] = parsedValue currentSection.isNotEmpty() -> sections[currentSection]?.put(key, parsedValue) } } @@ -334,7 +345,7 @@ object ConfigLoader { // Don't forget the last array-of-table entry if file ends with one flushTables() - return buildConfig(sections, providers, models, artifacts) + return buildConfig(sections, providers, models, artifacts, mcpServers) } private fun parseValue(valueStr: String, lineNum: Int): Any { @@ -452,6 +463,7 @@ object ConfigLoader { providersList: List> = emptyList(), modelsList: List> = emptyList(), artifactsList: List> = emptyList(), + mcpList: List> = emptyList(), ): CorrexConfig { val serverSection = sections["server"] ?: emptyMap() val tuiSection = sections["tui"] ?: emptyMap() @@ -735,6 +747,20 @@ object ConfigLoader { repeatPenalty = samplingSection["repeat_penalty"]?.let { asDouble(it) }, ) + val mcp = mcpList.mapNotNull { mcpMap -> + val id = asStringOrNull(mcpMap["id"]) ?: return@mapNotNull null + val command = asStringList(mcpMap["command"]) + if (command.isEmpty()) return@mapNotNull null + @Suppress("UNCHECKED_CAST") + val env = (mcpMap["env"] as? Map)?.mapValues { it.value.toString() } ?: emptyMap() + McpServerConfig( + id = id, + command = command, + env = env, + tier = asString(mcpMap["tier"], "T2"), + ) + } + val gitSection = sections["git"] ?: emptyMap() val git = GitConfig( enabled = asBoolean(gitSection["enabled"], false), @@ -758,6 +784,7 @@ object ConfigLoader { orchestration = orchestration, sampling = sampling, git = git, + mcp = mcp, ) } diff --git a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt index 99a7819d..aa6845de 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt @@ -19,6 +19,19 @@ data class CorrexConfig( val sampling: SamplingConfig = SamplingConfig(), val health: HealthConfig = HealthConfig(), val git: GitConfig = GitConfig(), + val mcp: List = emptyList(), +) + +/** + * An MCP server to mount at startup. Its `tools/list` becomes Correx tools (`mcp____`), + * gated at [tier] (default T2 = approval) since an external tool's side effects are opaque. + */ +@Serializable +data class McpServerConfig( + val id: String, + val command: List, + val env: Map = emptyMap(), + val tier: String = "T2", ) /** diff --git a/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt b/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt index 78a0e794..729df866 100644 --- a/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt +++ b/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt @@ -390,6 +390,31 @@ class ConfigLoaderTest { assertEquals(10001, result.modelsSettings.port) } + @Test + fun `parseToml parses mcp array-of-tables with command list and env`() { + val toml = """ + [[mcp]] + id = "codebase-memory" + command = ["codebase-memory-mcp", "serve"] + env = { RUST_LOG = "info" } + + [[mcp]] + id = "other" + command = ["other-server"] + """.trimIndent() + + val parseTomlMethod = ConfigLoader::class.java.getDeclaredMethod("parseToml", String::class.java) + parseTomlMethod.isAccessible = true + val result = parseTomlMethod.invoke(ConfigLoader, toml) as CorrexConfig + + assertEquals(2, result.mcp.size) + assertEquals("codebase-memory", result.mcp[0].id) + assertEquals(listOf("codebase-memory-mcp", "serve"), result.mcp[0].command) + assertEquals("info", result.mcp[0].env["RUST_LOG"]) + assertEquals("T2", result.mcp[0].tier) + assertEquals(listOf("other-server"), result.mcp[1].command) + } + @Test fun `parseToml returns empty models list and default modelsSettings when sections absent`() { val toml = """ diff --git a/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpMounter.kt b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpMounter.kt new file mode 100644 index 00000000..c2b8ed42 --- /dev/null +++ b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpMounter.kt @@ -0,0 +1,40 @@ +package com.correx.infrastructure.tools.mcp + +import com.correx.core.approvals.Tier +import com.correx.core.tools.contract.Tool + +/** + * A running MCP server mounted into Correx: its advertised tools (ready to hand to the tool + * registry as `extraTools`) plus the handle that shuts the child process down. The caller registers + * [close] with the server's shutdown hook, exactly as the managed-model path unloads llama-server. + */ +class MountedMcpServer internal constructor( + val serverId: String, + private val client: McpClient, + val tools: List, +) : AutoCloseable { + override fun close() = client.close() +} + +/** Spawn an MCP server over stdio, initialize it, and surface its `tools/list` as Correx tools. */ +object McpMounter { + suspend fun mount( + serverId: String, + command: List, + env: Map = emptyMap(), + tier: Tier = Tier.T2, + ): MountedMcpServer { + val client = McpClient.spawn(serverId, command, env) + val tools = client.listTools().map { descriptor -> + McpTool( + serverId = serverId, + remoteName = descriptor.name, + description = descriptor.description, + parametersSchema = descriptor.inputSchema, + tier = tier, + client = client, + ) + } + return MountedMcpServer(serverId, client, tools) + } +} diff --git a/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpProtocol.kt b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpProtocol.kt new file mode 100644 index 00000000..b29ef028 --- /dev/null +++ b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpProtocol.kt @@ -0,0 +1,68 @@ +package com.correx.infrastructure.tools.mcp + +import kotlinx.serialization.SerialName +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonObject + +/** + * Minimal JSON-RPC 2.0 + MCP wire types — just the tools slice Correx needs to mount an MCP + * server's tools as Correx tools. Resources, prompts, and sampling are intentionally absent (nothing + * consumes them yet). `ignoreUnknownKeys` on the codec lets servers send richer objects than these. + */ +@Serializable +internal data class JsonRpcRequest( + val method: String, + val params: JsonElement? = null, + val id: Long, + val jsonrpc: String = "2.0", +) + +@Serializable +internal data class JsonRpcNotification( + val method: String, + val params: JsonElement? = null, + val jsonrpc: String = "2.0", +) + +@Serializable +internal data class JsonRpcResponse( + val id: Long? = null, + val result: JsonElement? = null, + val error: JsonRpcError? = null, + val jsonrpc: String = "2.0", +) + +@Serializable +internal data class JsonRpcError( + val code: Int, + val message: String, + val data: JsonElement? = null, +) + +/** One tool advertised by an MCP server (`tools/list` entry). */ +@Serializable +internal data class McpToolDescriptor( + val name: String, + val description: String = "", + val inputSchema: JsonObject = JsonObject(emptyMap()), +) + +@Serializable +internal data class McpListToolsResult(val tools: List = emptyList()) + +/** One content block of a `tools/call` result. Only text blocks are surfaced to the model. */ +@Serializable +internal data class McpContentBlock( + val type: String, + val text: String? = null, +) + +@Serializable +internal data class McpCallToolResult( + val content: List = emptyList(), + @SerialName("isError") val isError: Boolean = false, +) { + /** Flatten text blocks into the single string the Correx tool contract returns. */ + fun text(): String = content.mapNotNull { it.text }.joinToString("\n") +} diff --git a/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpStdioClient.kt b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpStdioClient.kt new file mode 100644 index 00000000..6c959db1 --- /dev/null +++ b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpStdioClient.kt @@ -0,0 +1,169 @@ +package com.correx.infrastructure.tools.mcp + +import com.correx.core.tools.process.ChildProcess +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock +import kotlinx.coroutines.withContext +import kotlinx.serialization.json.Json +import kotlinx.serialization.json.JsonArray +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.decodeFromJsonElement +import kotlinx.serialization.json.put +import kotlinx.serialization.json.putJsonObject +import java.io.BufferedReader +import java.io.Writer +import java.util.concurrent.atomic.AtomicLong + +/** + * The transport an [McpClient] talks over. Extracted so tests can drive the client with an + * in-process fake instead of spawning a real subprocess (keeps the client's protocol logic + * deterministically testable — no process, no I/O). + */ +internal interface McpTransport { + suspend fun request(request: JsonRpcRequest): JsonRpcResponse + suspend fun notify(notification: JsonRpcNotification) + fun close() +} + +private val MCP_JSON = Json { + ignoreUnknownKeys = true + encodeDefaults = true +} + +/** + * A live MCP session: performs the `initialize` handshake, lists tools, and calls them. Ecosystem- + * agnostic — it neither knows nor cares whether the server behind [transport] resolves npm packages, + * indexes a codebase, or anything else. The nondeterminism of a `tools/call` is captured by the + * kernel's normal tool-execution recording (the call rides the ToolExecutor path), so nothing here + * needs to emit events: replay never re-calls because replay never executes tools. + */ +internal class McpClient( + val serverId: String, + private val transport: McpTransport, +) : McpClientHandle { + suspend fun initialize() { + val params = buildJsonObject { + put("protocolVersion", MCP_PROTOCOL_VERSION) + putJsonObject("capabilities") {} + putJsonObject("clientInfo") { + put("name", "correx") + put("version", "0") + } + } + transport.request(JsonRpcRequest(method = "initialize", params = params, id = 0)).orThrow(serverId) + transport.notify(JsonRpcNotification(method = "notifications/initialized")) + } + + suspend fun listTools(): List { + val result = transport.request(JsonRpcRequest(method = "tools/list", id = nextId())).orThrow(serverId) + return MCP_JSON.decodeFromJsonElement(McpListToolsResult.serializer(), result).tools + } + + override suspend fun callTool(remoteName: String, arguments: JsonObject): McpToolCallOutcome { + val params = buildJsonObject { + put("name", remoteName) + put("arguments", arguments) + } + val result = transport.request(JsonRpcRequest(method = "tools/call", params = params, id = nextId())) + .orThrow(serverId) + val decoded = MCP_JSON.decodeFromJsonElement(McpCallToolResult.serializer(), result) + return McpToolCallOutcome(text = decoded.text(), isError = decoded.isError) + } + + fun close() = transport.close() + + private val ids = AtomicLong(1) + private fun nextId() = ids.getAndIncrement() + + private fun JsonRpcResponse.orThrow(serverId: String): JsonElement { + error?.let { error("MCP server '$serverId' error ${it.code}: ${it.message}") } + return result ?: error("MCP server '$serverId' returned no result") + } + + companion object { + // The protocol revision Correx advertises. Servers negotiate down if they only speak older. + const val MCP_PROTOCOL_VERSION = "2024-11-05" + + /** + * Spawn [command] as a child process and return an initialized client speaking MCP over its + * stdio pipes. The caller owns [McpClient.close] (kills the process). Uses [ChildProcess] so + * the server inherits the same unattended/non-interactive execution guarantees as every other + * spawned child. + */ + suspend fun spawn(serverId: String, command: List, env: Map): McpClient { + val process = withContext(Dispatchers.IO) { + ChildProcess.builder(command).apply { environment().putAll(env) }.start() + } + val client = McpClient(serverId, StdioTransport(process)) + client.initialize() + return client + } + } +} + +/** + * Newline-delimited JSON-RPC over a child process's stdio (the MCP stdio transport). Requests are + * serialized through a [Mutex]: one line written, then response lines drained until the matching id + * arrives — server-initiated notifications/logs on the same pipe are skipped, not mistaken for the + * reply. Blocking pipe I/O is confined to [Dispatchers.IO]. + */ +private class StdioTransport(private val process: Process) : McpTransport { + private val lock = Mutex() + private val writer: Writer = process.outputStream.bufferedWriter() + private val reader: BufferedReader = process.inputStream.bufferedReader() + + override suspend fun request(request: JsonRpcRequest): JsonRpcResponse = lock.withLock { + withContext(Dispatchers.IO) { + writeLine(MCP_JSON.encodeToString(JsonRpcRequest.serializer(), request)) + while (true) { + val line = reader.readLine() + ?: error("MCP server '${request.method}' closed the connection") + if (line.isBlank()) continue + val response = runCatching { + MCP_JSON.decodeFromString(JsonRpcResponse.serializer(), line) + }.getOrNull() ?: continue // not a response we can parse (a notification/log line) — skip + if (response.id == request.id) return@withContext response + } + @Suppress("UNREACHABLE_CODE") + error("unreachable") + } + } + + override suspend fun notify(notification: JsonRpcNotification) = lock.withLock { + withContext(Dispatchers.IO) { + writeLine(MCP_JSON.encodeToString(JsonRpcNotification.serializer(), notification)) + } + } + + override fun close() { + runCatching { writer.close() } + runCatching { reader.close() } + process.destroy() + } + + private fun writeLine(json: String) { + writer.write(json) + writer.write("\n") + writer.flush() + } +} + +/** Encode a [ToolRequest]'s parameter map into the JsonObject `arguments` payload for a tools/call. */ +internal fun argumentsOf(parameters: Map): JsonObject = + JsonObject(parameters.mapValues { (_, value) -> anyToJsonElement(value) }) + +private fun anyToJsonElement(value: Any?): JsonElement = when (value) { + null -> JsonNull + is JsonElement -> value + is Boolean -> JsonPrimitive(value) + is Number -> JsonPrimitive(value) + is String -> JsonPrimitive(value) + is Map<*, *> -> JsonObject(value.entries.associate { (k, v) -> k.toString() to anyToJsonElement(v) }) + is Iterable<*> -> JsonArray(value.map { anyToJsonElement(it) }) + else -> JsonPrimitive(value.toString()) +} diff --git a/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpTool.kt b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpTool.kt new file mode 100644 index 00000000..50feb697 --- /dev/null +++ b/infrastructure/tools/src/main/kotlin/com/correx/infrastructure/tools/mcp/McpTool.kt @@ -0,0 +1,64 @@ +package com.correx.infrastructure.tools.mcp + +import com.correx.core.approvals.Tier +import com.correx.core.events.events.ToolRequest +import com.correx.core.tools.contract.Tool +import com.correx.core.tools.contract.ToolCapability +import com.correx.core.tools.contract.ToolExecutor +import com.correx.core.tools.contract.ToolResult +import com.correx.core.tools.contract.ValidationResult +import kotlinx.serialization.json.JsonObject + +/** + * One tool advertised by an MCP server, surfaced as a first-class Correx [Tool]/[ToolExecutor]. It + * is exactly like any other tool from the kernel's point of view — so it inherits tier gating, + * receipts, and event-recorded execution (#5) with zero special-casing, and replay reads the + * recorded receipt instead of ever re-calling the server (#8/#9). + * + * Naming: `mcp____` — namespaced by server so two servers can advertise a tool of + * the same name, and recognizable as MCP-backed at a glance (the same convention MCP hosts use). + * + * Capabilities are declared empty: an external MCP tool's side effects are opaque and cannot be + * classified into Correx's FILE/NETWORK/SHELL taxonomy, so it does not claim any — safety comes from + * [tier] (default T2 = approval-gated) rather than from a capability guess that would misfire the + * plane-2 rules (e.g. NETWORK_ACCESS expects a NETWORK_TARGET param an MCP tool doesn't have). + */ +class McpTool( + private val serverId: String, + private val remoteName: String, + override val description: String, + override val parametersSchema: JsonObject, + override val tier: Tier, + private val client: McpClientHandle, +) : Tool, ToolExecutor { + + override val name: String = "mcp__${serverId}__$remoteName" + override val requiredCapabilities: Set = emptySet() + + override fun validateRequest(request: ToolRequest): ValidationResult = ValidationResult.Valid + + override suspend fun execute(request: ToolRequest): ToolResult = + runCatching { client.callTool(remoteName, argumentsOf(request.parameters)) } + .fold( + onSuccess = { result -> + if (result.isError) { + // A tool-reported failure (e.g. "package not found") — recoverable so it flows + // back into the ReAct loop as evidence rather than aborting the stage. + ToolResult.Failure(request.invocationId, result.text, recoverable = true) + } else { + ToolResult.Success(request.invocationId, result.text) + } + }, + // A transport/protocol failure (server crashed, bad frame) is not the model's fault + // and not fixable by retrying the same call: non-recoverable. + onFailure = { ToolResult.Failure(request.invocationId, "MCP call failed: ${it.message}", recoverable = false) }, + ) +} + +/** The slice of [McpClient] an [McpTool] needs — narrowed to an interface so tools are unit-testable. */ +interface McpClientHandle { + suspend fun callTool(remoteName: String, arguments: JsonObject): McpToolCallOutcome +} + +/** Transport-agnostic result an [McpClientHandle] hands back (decouples the tool from wire types). */ +data class McpToolCallOutcome(val text: String, val isError: Boolean) diff --git a/infrastructure/tools/src/test/kotlin/com/correx/infrastructure/tools/mcp/McpClientTest.kt b/infrastructure/tools/src/test/kotlin/com/correx/infrastructure/tools/mcp/McpClientTest.kt new file mode 100644 index 00000000..5ca78f61 --- /dev/null +++ b/infrastructure/tools/src/test/kotlin/com/correx/infrastructure/tools/mcp/McpClientTest.kt @@ -0,0 +1,112 @@ +package com.correx.infrastructure.tools.mcp + +import com.correx.core.approvals.Tier +import com.correx.core.events.events.ToolRequest +import com.correx.core.events.types.SessionId +import com.correx.core.events.types.StageId +import com.correx.core.events.types.ToolInvocationId +import com.correx.core.tools.contract.ToolResult +import kotlinx.coroutines.runBlocking +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.put +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test + +/** + * Drives [McpClient] through an in-process [McpTransport] — no subprocess, no I/O. Verifies the + * protocol logic (handshake, tools/list decode, tools/call outcome, error surfacing) and the + * [McpTool] adapter's mapping onto the Correx tool contract, all deterministically. + */ +class McpClientTest { + + /** A scripted transport: answers each method from a canned map; records notifications. */ + private class FakeTransport(private val results: Map) : McpTransport { + val notified = mutableListOf() + override suspend fun request(request: JsonRpcRequest): JsonRpcResponse = + (results[request.method] ?: error("no scripted result for ${request.method}")) + .copy(id = request.id) + override suspend fun notify(notification: JsonRpcNotification) { notified.add(notification.method) } + override fun close() {} + } + + private fun ok(json: kotlinx.serialization.json.JsonElement) = JsonRpcResponse(result = json) + + @Test + fun `initialize handshakes then lists and calls tools`(): Unit = runBlocking { + val listResult = buildJsonObject { + put("tools", kotlinx.serialization.json.buildJsonArray { + add(buildJsonObject { + put("name", "search_code") + put("description", "search the graph") + }) + }) + } + val callResult = buildJsonObject { + put("content", kotlinx.serialization.json.buildJsonArray { + add(buildJsonObject { put("type", "text"); put("text", "hit one") }) + add(buildJsonObject { put("type", "text"); put("text", "hit two") }) + }) + } + val transport = FakeTransport( + mapOf( + "initialize" to ok(buildJsonObject {}), + "tools/list" to ok(listResult), + "tools/call" to ok(callResult), + ), + ) + val client = McpClient("cm", transport) + client.initialize() + assertEquals(listOf("notifications/initialized"), transport.notified) + + val tools = client.listTools() + assertEquals(listOf("search_code"), tools.map { it.name }) + + val outcome = client.callTool("search_code", JsonObject(emptyMap())) + assertEquals("hit one\nhit two", outcome.text) + assertTrue(!outcome.isError) + } + + @Test + fun `server error is surfaced as a thrown failure`(): Unit = runBlocking { + val transport = FakeTransport( + mapOf("tools/call" to JsonRpcResponse(error = JsonRpcError(code = -32000, message = "boom"))), + ) + val client = McpClient("cm", transport) + val ex = runCatching { client.callTool("x", JsonObject(emptyMap())) }.exceptionOrNull() + assertTrue(ex?.message?.contains("boom") == true, "expected server error message, got: ${ex?.message}") + } + + @Test + fun `McpTool maps an isError result to a recoverable failure`(): Unit = runBlocking { + val handle = object : McpClientHandle { + override suspend fun callTool(remoteName: String, arguments: JsonObject) = + McpToolCallOutcome(text = "package not found", isError = true) + } + val tool = McpTool("cm", "resolve", "", JsonObject(emptyMap()), Tier.T2, handle) + assertEquals("mcp__cm__resolve", tool.name) + val result = tool.execute(req(tool.name)) as ToolResult.Failure + assertEquals("package not found", result.reason) + assertTrue(result.recoverable) + } + + @Test + fun `McpTool maps a normal result to success`(): Unit = runBlocking { + val handle = object : McpClientHandle { + override suspend fun callTool(remoteName: String, arguments: JsonObject) = + McpToolCallOutcome(text = "ok", isError = false) + } + val tool = McpTool("cm", "resolve", "", JsonObject(emptyMap()), Tier.T2, handle) + val result = tool.execute(req(tool.name)) as ToolResult.Success + assertEquals("ok", result.output) + } + + private fun req(name: String) = ToolRequest( + invocationId = ToolInvocationId("inv-1"), + sessionId = SessionId("s-1"), + stageId = StageId("st-1"), + toolName = name, + parameters = mapOf("q" to "vite"), + ) +}