feat(tools): native MCP client — mount stdio MCP servers as Correx tools
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__<server>__<tool>. 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 <noreply@anthropic.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -265,9 +265,11 @@ object ConfigLoader {
|
||||
val providers = mutableListOf<MutableMap<String, Any>>()
|
||||
val models = mutableListOf<MutableMap<String, Any>>()
|
||||
val artifacts = mutableListOf<MutableMap<String, Any>>()
|
||||
val mcpServers = mutableListOf<MutableMap<String, Any>>()
|
||||
var currentProvider: MutableMap<String, Any>? = null
|
||||
var currentModel: MutableMap<String, Any>? = null
|
||||
var currentArtifact: MutableMap<String, Any>? = null
|
||||
var currentMcp: MutableMap<String, Any>? = 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<Map<String, Any>> = emptyList(),
|
||||
modelsList: List<Map<String, Any>> = emptyList(),
|
||||
artifactsList: List<Map<String, Any>> = emptyList(),
|
||||
mcpList: List<Map<String, Any>> = 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<String, Any>)?.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,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -19,6 +19,19 @@ data class CorrexConfig(
|
||||
val sampling: SamplingConfig = SamplingConfig(),
|
||||
val health: HealthConfig = HealthConfig(),
|
||||
val git: GitConfig = GitConfig(),
|
||||
val mcp: List<McpServerConfig> = emptyList(),
|
||||
)
|
||||
|
||||
/**
|
||||
* An MCP server to mount at startup. Its `tools/list` becomes Correx tools (`mcp__<id>__<name>`),
|
||||
* 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<String>,
|
||||
val env: Map<String, String> = emptyMap(),
|
||||
val tier: String = "T2",
|
||||
)
|
||||
|
||||
/**
|
||||
|
||||
@@ -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 = """
|
||||
|
||||
+40
@@ -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<Tool>,
|
||||
) : 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<String>,
|
||||
env: Map<String, String> = 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)
|
||||
}
|
||||
}
|
||||
+68
@@ -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<McpToolDescriptor> = 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<McpContentBlock> = 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")
|
||||
}
|
||||
+169
@@ -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<McpToolDescriptor> {
|
||||
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<String>, env: Map<String, String>): 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<String, Any?>): 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())
|
||||
}
|
||||
@@ -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__<server>__<remote>` — 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<ToolCapability> = 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)
|
||||
+112
@@ -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<String, JsonRpcResponse>) : McpTransport {
|
||||
val notified = mutableListOf<String>()
|
||||
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"),
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user