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")
|
val explicitWorkspaceRoot = System.getenv("CORREX_WORKSPACE_ROOT")
|
||||||
?.let { Path.of(it) }
|
?.let { Path.of(it) }
|
||||||
?: toolsConfig.workspaceRoot.takeIf { it.isNotEmpty() }?.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 shellAllowedExecutables = toolsConfig.shellAllowedExecutables.toSet()
|
||||||
val workspaceRoot = explicitWorkspaceRoot ?: explicitWorkingDir ?: Path.of("").toAbsolutePath()
|
val bootWorkspace = resolveBootWorkspace(
|
||||||
val workingDir = explicitWorkingDir ?: workspaceRoot
|
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
|
// 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.
|
// (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
|
// 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).
|
// Retrieves full tool output the kernel spilled to CAS on truncation (ref shown in-context).
|
||||||
val toolOutputTool = com.correx.infrastructure.tools.ToolOutputTool(artifactStore)
|
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(
|
val toolRegistry = InfrastructureModule.createToolRegistry(
|
||||||
buildToolConfig(
|
buildToolConfig(
|
||||||
workspaceRoot,
|
workspaceRoot,
|
||||||
@@ -545,7 +579,7 @@ fun main() {
|
|||||||
fun buildProjectMemory(cfg: CorrexConfig): com.correx.apps.server.memory.ProjectMemoryService? =
|
fun buildProjectMemory(cfg: CorrexConfig): com.correx.apps.server.memory.ProjectMemoryService? =
|
||||||
if (cfg.project.enabled) {
|
if (cfg.project.enabled) {
|
||||||
com.correx.apps.server.memory.ProjectMemoryService(
|
com.correx.apps.server.memory.ProjectMemoryService(
|
||||||
config = cfg.project,
|
config = cfg.project.boundToWorkspace(workspaceRoot),
|
||||||
embedder = embedder,
|
embedder = embedder,
|
||||||
l3MemoryStore = l3MemoryStore,
|
l3MemoryStore = l3MemoryStore,
|
||||||
journalRepository = decisionJournalRepository,
|
journalRepository = decisionJournalRepository,
|
||||||
@@ -871,7 +905,7 @@ private fun buildToolConfig(
|
|||||||
toolsConfig: com.correx.core.config.ToolsConfig,
|
toolsConfig: com.correx.core.config.ToolsConfig,
|
||||||
research: com.correx.infrastructure.tools.ResearchToolConfig,
|
research: com.correx.infrastructure.tools.ResearchToolConfig,
|
||||||
): ToolConfig {
|
): ToolConfig {
|
||||||
val allowed = setOf(workspaceRoot, workingDir)
|
val allowed = setOf(workspaceRoot)
|
||||||
return ToolConfig(
|
return ToolConfig(
|
||||||
shell = ShellConfig(
|
shell = ShellConfig(
|
||||||
enabled = toolsConfig.shellEnabled,
|
enabled = toolsConfig.shellEnabled,
|
||||||
|
|||||||
@@ -265,9 +265,11 @@ object ConfigLoader {
|
|||||||
val providers = mutableListOf<MutableMap<String, Any>>()
|
val providers = mutableListOf<MutableMap<String, Any>>()
|
||||||
val models = mutableListOf<MutableMap<String, Any>>()
|
val models = mutableListOf<MutableMap<String, Any>>()
|
||||||
val artifacts = mutableListOf<MutableMap<String, Any>>()
|
val artifacts = mutableListOf<MutableMap<String, Any>>()
|
||||||
|
val mcpServers = mutableListOf<MutableMap<String, Any>>()
|
||||||
var currentProvider: MutableMap<String, Any>? = null
|
var currentProvider: MutableMap<String, Any>? = null
|
||||||
var currentModel: MutableMap<String, Any>? = null
|
var currentModel: MutableMap<String, Any>? = null
|
||||||
var currentArtifact: 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
|
// 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.
|
// table header (array or section) starts, so the previous entry is committed.
|
||||||
@@ -275,9 +277,11 @@ object ConfigLoader {
|
|||||||
currentProvider?.let { providers.add(it) }
|
currentProvider?.let { providers.add(it) }
|
||||||
currentModel?.let { models.add(it) }
|
currentModel?.let { models.add(it) }
|
||||||
currentArtifact?.let { artifacts.add(it) }
|
currentArtifact?.let { artifacts.add(it) }
|
||||||
|
currentMcp?.let { mcpServers.add(it) }
|
||||||
currentProvider = null
|
currentProvider = null
|
||||||
currentModel = null
|
currentModel = null
|
||||||
currentArtifact = null
|
currentArtifact = null
|
||||||
|
currentMcp = null
|
||||||
}
|
}
|
||||||
|
|
||||||
for ((lineNum, line) in lines.withIndex()) {
|
for ((lineNum, line) in lines.withIndex()) {
|
||||||
@@ -302,6 +306,11 @@ object ConfigLoader {
|
|||||||
currentArtifact = mutableMapOf()
|
currentArtifact = mutableMapOf()
|
||||||
currentSection = ""
|
currentSection = ""
|
||||||
}
|
}
|
||||||
|
trimmed == "[[mcp]]" -> {
|
||||||
|
flushTables()
|
||||||
|
currentMcp = mutableMapOf()
|
||||||
|
currentSection = ""
|
||||||
|
}
|
||||||
trimmed.startsWith("[") && !trimmed.startsWith("[[") && trimmed.endsWith("]") -> {
|
trimmed.startsWith("[") && !trimmed.startsWith("[[") && trimmed.endsWith("]") -> {
|
||||||
// Parse section headers like [server] or [tools.shell]
|
// Parse section headers like [server] or [tools.shell]
|
||||||
flushTables()
|
flushTables()
|
||||||
@@ -320,10 +329,12 @@ object ConfigLoader {
|
|||||||
val provider = currentProvider
|
val provider = currentProvider
|
||||||
val model = currentModel
|
val model = currentModel
|
||||||
val artifact = currentArtifact
|
val artifact = currentArtifact
|
||||||
|
val mcp = currentMcp
|
||||||
when {
|
when {
|
||||||
provider != null -> provider[key] = parsedValue
|
provider != null -> provider[key] = parsedValue
|
||||||
model != null -> model[key] = parsedValue
|
model != null -> model[key] = parsedValue
|
||||||
artifact != null -> artifact[key] = parsedValue
|
artifact != null -> artifact[key] = parsedValue
|
||||||
|
mcp != null -> mcp[key] = parsedValue
|
||||||
currentSection.isNotEmpty() -> sections[currentSection]?.put(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
|
// Don't forget the last array-of-table entry if file ends with one
|
||||||
flushTables()
|
flushTables()
|
||||||
|
|
||||||
return buildConfig(sections, providers, models, artifacts)
|
return buildConfig(sections, providers, models, artifacts, mcpServers)
|
||||||
}
|
}
|
||||||
|
|
||||||
private fun parseValue(valueStr: String, lineNum: Int): Any {
|
private fun parseValue(valueStr: String, lineNum: Int): Any {
|
||||||
@@ -452,6 +463,7 @@ object ConfigLoader {
|
|||||||
providersList: List<Map<String, Any>> = emptyList(),
|
providersList: List<Map<String, Any>> = emptyList(),
|
||||||
modelsList: List<Map<String, Any>> = emptyList(),
|
modelsList: List<Map<String, Any>> = emptyList(),
|
||||||
artifactsList: List<Map<String, Any>> = emptyList(),
|
artifactsList: List<Map<String, Any>> = emptyList(),
|
||||||
|
mcpList: List<Map<String, Any>> = emptyList(),
|
||||||
): CorrexConfig {
|
): CorrexConfig {
|
||||||
val serverSection = sections["server"] ?: emptyMap()
|
val serverSection = sections["server"] ?: emptyMap()
|
||||||
val tuiSection = sections["tui"] ?: emptyMap()
|
val tuiSection = sections["tui"] ?: emptyMap()
|
||||||
@@ -735,6 +747,20 @@ object ConfigLoader {
|
|||||||
repeatPenalty = samplingSection["repeat_penalty"]?.let { asDouble(it) },
|
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 gitSection = sections["git"] ?: emptyMap()
|
||||||
val git = GitConfig(
|
val git = GitConfig(
|
||||||
enabled = asBoolean(gitSection["enabled"], false),
|
enabled = asBoolean(gitSection["enabled"], false),
|
||||||
@@ -758,6 +784,7 @@ object ConfigLoader {
|
|||||||
orchestration = orchestration,
|
orchestration = orchestration,
|
||||||
sampling = sampling,
|
sampling = sampling,
|
||||||
git = git,
|
git = git,
|
||||||
|
mcp = mcp,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -19,6 +19,19 @@ data class CorrexConfig(
|
|||||||
val sampling: SamplingConfig = SamplingConfig(),
|
val sampling: SamplingConfig = SamplingConfig(),
|
||||||
val health: HealthConfig = HealthConfig(),
|
val health: HealthConfig = HealthConfig(),
|
||||||
val git: GitConfig = GitConfig(),
|
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)
|
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
|
@Test
|
||||||
fun `parseToml returns empty models list and default modelsSettings when sections absent`() {
|
fun `parseToml returns empty models list and default modelsSettings when sections absent`() {
|
||||||
val toml = """
|
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