From 2de974982f2b9b9494ace3d9e9735bbb514a5ecf Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 09:30:56 +0300 Subject: [PATCH 1/6] feat: add MCP server config types and rmcp dependency --- Cargo.lock | 274 ++++++++++++++++++++++++++++++--- Cargo.toml | 1 + src/config/client.rs | 1 + src/config/config.toml.default | 13 ++ src/config/mod.rs | 2 +- src/config/resolve.rs | 3 + src/config/types.rs | 20 ++- 7 files changed, 288 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c8f2f51..8566cab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -459,6 +459,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.43" @@ -535,7 +541,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -761,7 +767,7 @@ checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" dependencies = [ "cfg-if", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -837,6 +843,21 @@ dependencies = [ "libc", ] +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.31" @@ -844,6 +865,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -852,6 +874,34 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.31" @@ -870,8 +920,13 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", @@ -1320,10 +1375,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.83" +version = "0.3.97" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" +checksum = "a1840c94c045fbcf8ba2812c95db44499f7c64910a912551aaaa541decebcacf" dependencies = [ + "cfg-if", + "futures-util", "once_cell", "wasm-bindgen", ] @@ -1362,9 +1419,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.180" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libredox" @@ -1510,7 +1567,19 @@ checksum = "ab2156c4fce2f8df6c499cc1c763e4394b7482525bf2a9701c9d79d215f519e4" dependencies = [ "bitflags 2.10.0", "cfg-if", - "cfg_aliases", + "cfg_aliases 0.1.1", + "libc", +] + +[[package]] +name = "nix" +version = "0.31.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3" +dependencies = [ + "bitflags 2.10.0", + "cfg-if", + "cfg_aliases 0.2.1", "libc", ] @@ -1590,7 +1659,8 @@ dependencies = [ "indicatif", "notify", "once_cell", - "reqwest", + "reqwest 0.12.28", + "rmcp", "rustyline", "serde", "serde_json", @@ -1740,6 +1810,20 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "process-wrap" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +dependencies = [ + "futures", + "indexmap", + "nix 0.31.2", + "tokio", + "tracing", + "windows", +] + [[package]] name = "quote" version = "1.0.43" @@ -1898,6 +1982,40 @@ dependencies = [ "web-sys", ] +[[package]] +name = "reqwest" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62e0021ea2c22aed41653bc7e1419abb2c97e038ff2c33d0e1309e49a97deec0" +dependencies = [ + "base64", + "bytes", + "futures-core", + "futures-util", + "http 1.4.0", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "serde", + "serde_json", + "sync_wrapper", + "tokio", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -1912,6 +2030,29 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmcp" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e12ca9067b5ebfbd5b3fcdc4acfceb81aa7d5ab2a879dff7cb75d22434276aad" +dependencies = [ + "async-trait", + "chrono", + "futures", + "http 1.4.0", + "pin-project-lite", + "process-wrap", + "reqwest 0.13.3", + "serde", + "serde_json", + "sse-stream", + "thiserror", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", +] + [[package]] name = "rustc_version" version = "0.4.1" @@ -1987,7 +2128,7 @@ dependencies = [ "libc", "log", "memchr", - "nix", + "nix 0.28.0", "radix_trie", "unicode-segmentation", "unicode-width 0.1.14", @@ -2192,6 +2333,19 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "sse-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3962b63f038885f15bce2c6e02c0e7925c072f1ac86bb60fd44c5c6b762fb72" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -2393,6 +2547,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -2692,9 +2857,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.106" +version = "0.2.120" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" +checksum = "df52b6d9b87e0c74c9edfa1eb2d9bf85e5d63515474513aa50fa181b3c4f5db1" dependencies = [ "cfg-if", "once_cell", @@ -2705,22 +2870,19 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.56" +version = "0.4.70" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836d9622d604feee9e5de25ac10e3ea5f2d65b41eac0d9ce72eb5deae707ce7c" +checksum = "af934872acec734c2d80e6617bbb5ff4f12b052dd8e6332b0817bce889516084" dependencies = [ - "cfg-if", "js-sys", - "once_cell", "wasm-bindgen", - "web-sys", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.106" +version = "0.2.120" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" +checksum = "78b1041f495fb322e64aca85f5756b2172e35cd459376e67f2a6c9dffcedb103" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2728,9 +2890,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.106" +version = "0.2.120" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" +checksum = "9dcd0ff20416988a18ac686d4d4d0f6aae9ebf08a389ff5d29012b05af2a1b41" dependencies = [ "bumpalo", "proc-macro2", @@ -2741,18 +2903,31 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.106" +version = "0.2.120" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" +checksum = "49757b3c82ebf16c57d69365a142940b384176c24df52a087fb748e2085359ea" dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" -version = "0.3.83" +version = "0.3.97" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" +checksum = "2eadbac71025cd7b0834f20d1fe8472e8495821b4e9801eb0a60bd1f19827602" dependencies = [ "js-sys", "wasm-bindgen", @@ -2777,6 +2952,27 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -2790,6 +2986,17 @@ dependencies = [ "windows-strings", ] +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core", + "windows-link", + "windows-threading", +] + [[package]] name = "windows-implement" version = "0.60.2" @@ -2818,6 +3025,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core", + "windows-link", +] + [[package]] name = "windows-registry" version = "0.6.1" @@ -2940,6 +3157,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" diff --git a/Cargo.toml b/Cargo.toml index 32dafcc..aaba877 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ toml = "0.8" dirs = "6.0" uuid = { version = "1", features = ["v4", "serde"] } chrono = { version = "0.4", features = ["serde"] } +rmcp = { version = "1.6.0", default-features = false, features = ["client", "transport-child-process", "transport-streamable-http-client-reqwest"] } [dev-dependencies] tempfile = "3" diff --git a/src/config/client.rs b/src/config/client.rs index b9c00c0..ba1e272 100644 --- a/src/config/client.rs +++ b/src/config/client.rs @@ -115,6 +115,7 @@ mod tests { default_provider: "openai".into(), max_iterations: 10, providers, + mcp_servers: BTreeMap::new(), } } diff --git a/src/config/config.toml.default b/src/config/config.toml.default index dbfdf8f..ab77dd0 100644 --- a/src/config/config.toml.default +++ b/src/config/config.toml.default @@ -51,3 +51,16 @@ env_var = "ANTHROPIC_API_KEY" # api_base = "http://localhost:11434/v1" # default_model = "llama2" # models = ["llama2", "mistral", "codellama", "mixtral"] + +# --- MCP Server Configuration (Model Context Protocol) --- +# Tools are exposed to the LLM as "{server_name}__{tool_name}". +# +# Stdio transport — spawn a local process: +# [mcp_servers.filesystem] +# command = "npx" +# args = ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"] +# env = { } +# +# Streamable HTTP transport — connect to a running server: +# [mcp_servers.remote-tools] +# url = "http://localhost:8080/mcp" diff --git a/src/config/mod.rs b/src/config/mod.rs index f63f118..e10017b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -3,7 +3,7 @@ mod resolve; mod types; pub use client::{build_http_client, create_client, resolve_client_and_config}; -pub use types::{AgentConfig, AppConfig, ProviderConfig}; +pub use types::{AgentConfig, AppConfig, McpServerConfig, ProviderConfig}; use std::path::PathBuf; diff --git a/src/config/resolve.rs b/src/config/resolve.rs index df0b0d7..2f0ebd2 100644 --- a/src/config/resolve.rs +++ b/src/config/resolve.rs @@ -123,6 +123,7 @@ mod tests { default_provider: "openai".into(), max_iterations: 5, providers, + mcp_servers: BTreeMap::new(), } } @@ -161,6 +162,7 @@ mod tests { default_provider: "nonexistent".into(), max_iterations: 10, providers: BTreeMap::new(), + mcp_servers: BTreeMap::new(), }; let err = config.resolve(None).unwrap_err(); assert!(err.to_string().contains("nonexistent")); @@ -183,6 +185,7 @@ mod tests { default_provider: "openai".into(), max_iterations: 10, providers: BTreeMap::new(), + mcp_servers: BTreeMap::new(), }; let err = config.list_models().unwrap_err(); assert!(err.to_string().contains("No providers configured")); diff --git a/src/config/types.rs b/src/config/types.rs index e24bff7..e7608f3 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -1,5 +1,5 @@ use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; /// Top-level configuration loaded from ~/.openheim/config.toml @@ -10,6 +10,24 @@ pub struct AppConfig { pub max_iterations: usize, #[serde(default)] pub providers: BTreeMap, + #[serde(default)] + pub mcp_servers: BTreeMap, +} + +/// Configuration for a single MCP server connection. +/// The map key in `[mcp_servers.]` is used as the server name and tool-name prefix. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpServerConfig { + /// Binary to spawn for stdio transport (e.g. `"npx"`, `"uvx"`). + pub command: Option, + /// Arguments passed to `command`. + #[serde(default)] + pub args: Vec, + /// Extra environment variables for the spawned process. + #[serde(default)] + pub env: HashMap, + /// Base URL for Streamable HTTP transport (e.g. `"http://localhost:8080/mcp"`). + pub url: Option, } fn default_max_iterations() -> usize { From 7202a284a50a244190e2a5e8e0e7b423bf18dcbf Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 09:31:01 +0300 Subject: [PATCH 2/6] feat: implement MCP client and tool handler --- src/mcp/client.rs | 95 +++++++++++++++++++++++++++++++++++++++++ src/mcp/mod.rs | 50 ++++++++++++++++++++++ src/mcp/tool_handler.rs | 52 ++++++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 src/mcp/client.rs create mode 100644 src/mcp/mod.rs create mode 100644 src/mcp/tool_handler.rs diff --git a/src/mcp/client.rs b/src/mcp/client.rs new file mode 100644 index 0000000..e16bd6f --- /dev/null +++ b/src/mcp/client.rs @@ -0,0 +1,95 @@ +use rmcp::{ + ServiceExt, + model::{CallToolRequestParams, Tool}, + service::{RoleClient, RunningService}, + transport::{TokioChildProcess, streamable_http_client::StreamableHttpClientTransport}, +}; + +use crate::{ + config::McpServerConfig, + error::{Error, Result}, +}; + +pub struct McpClient { + service: RunningService, + pub server_name: String, +} + +impl McpClient { + pub async fn connect(name: &str, config: &McpServerConfig) -> Result { + if let Some(ref url) = config.url { + let transport = StreamableHttpClientTransport::from_uri(url.as_str()); + let service = () + .serve(transport) + .await + .map_err(|e| Error::Other(format!("MCP HTTP connect to '{}' failed: {}", name, e)))?; + Ok(Self { service, server_name: name.to_string() }) + } else if let Some(ref command) = config.command { + let mut cmd = tokio::process::Command::new(command); + cmd.args(&config.args); + for (k, v) in &config.env { + cmd.env(k, v); + } + let transport = TokioChildProcess::new(cmd) + .map_err(|e| Error::Other(format!("MCP spawn '{}' failed: {}", name, e)))?; + let service = () + .serve(transport) + .await + .map_err(|e| Error::Other(format!("MCP stdio connect to '{}' failed: {}", name, e)))?; + Ok(Self { service, server_name: name.to_string() }) + } else { + Err(Error::ConfigError(format!( + "MCP server '{}' must have either 'command' (stdio) or 'url' (HTTP)", + name + ))) + } + } + + pub async fn list_tools(&self) -> Result> { + let result = self + .service + .peer() + .list_tools(Default::default()) + .await + .map_err(|e| Error::Other(format!("MCP list_tools failed for '{}': {}", self.server_name, e)))?; + Ok(result.tools.into_iter().collect()) + } + + pub async fn call_tool(&self, name: &str, args_json: &str) -> Result { + let params = build_call_params(name, args_json)?; + + let result = self + .service + .peer() + .call_tool(params) + .await + .map_err(|e| Error::ToolExecutionError(format!("MCP tool '{}' on '{}' failed: {}", name, self.server_name, e)))?; + + if result.is_error.unwrap_or(false) { + return Err(Error::ToolExecutionError(extract_text_content(&result.content))); + } + + Ok(extract_text_content(&result.content)) + } +} + +fn build_call_params(name: &str, args_json: &str) -> Result { + let trimmed = args_json.trim(); + if trimmed.is_empty() || trimmed == "{}" { + return Ok(CallToolRequestParams::new(name.to_string())); + } + let map: serde_json::Map = serde_json::from_str(trimmed)?; + Ok(CallToolRequestParams::new(name.to_string()).with_arguments(map)) +} + +fn extract_text_content(content: &[impl serde::Serialize]) -> String { + content + .iter() + .filter_map(|c| { + serde_json::to_value(c) + .ok() + .and_then(|v| v.get("text")?.as_str().map(String::from)) + }) + .collect::>() + .join("\n") +} diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs new file mode 100644 index 0000000..a488d6c --- /dev/null +++ b/src/mcp/mod.rs @@ -0,0 +1,50 @@ +mod client; +mod tool_handler; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use client::McpClient; +use tool_handler::McpToolHandler; + +use crate::{config::McpServerConfig, error::Result, tools::ToolHandler}; + +pub(crate) async fn load_mcp_tools(configs: &BTreeMap) -> Vec> { + let mut handlers: Vec> = Vec::new(); + + for (name, config) in configs { + match connect_server(name, config).await { + Ok(server_handlers) => { + tracing::info!( + server = %name, + count = server_handlers.len(), + "MCP server connected" + ); + handlers.extend(server_handlers); + } + Err(e) => { + tracing::warn!(server = %name, error = %e, "MCP server failed to connect"); + } + } + } + + handlers +} + +async fn connect_server(name: &str, config: &McpServerConfig) -> Result>> { + let client = Arc::new(McpClient::connect(name, config).await?); + let tools = client.list_tools().await?; + + // Sanitise the prefix: hyphens and spaces become underscores so the + // combined name is a valid identifier for tool-call APIs. + let prefix = name.replace(['-', ' '], "_"); + + let handlers = tools + .iter() + .map(|tool| -> Box { + Box::new(McpToolHandler::new(Arc::clone(&client), tool, &prefix)) + }) + .collect(); + + Ok(handlers) +} diff --git a/src/mcp/tool_handler.rs b/src/mcp/tool_handler.rs new file mode 100644 index 0000000..8636bcb --- /dev/null +++ b/src/mcp/tool_handler.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use rmcp::model::Tool as McpTool; + +use crate::{ + core::models::{FunctionDefinition, Tool}, + error::Result, + tools::ToolHandler, +}; + +use super::client::McpClient; + +pub struct McpToolHandler { + client: Arc, + /// Original tool name as reported by the MCP server. + tool_name: String, + /// Name exposed to the LLM: `{server_prefix}__{tool_name}`. + prefixed_name: String, + description: String, + schema: serde_json::Value, +} + +impl McpToolHandler { + pub fn new(client: Arc, tool: &McpTool, server_prefix: &str) -> Self { + let tool_name = tool.name.to_string(); + let prefixed_name = format!("{}__{}", server_prefix, tool_name); + let description = tool.description.as_deref().unwrap_or("").to_string(); + let schema = serde_json::to_value(&tool.input_schema) + .unwrap_or_else(|_| serde_json::json!({"type": "object", "properties": {}})); + + Self { client, tool_name, prefixed_name, description, schema } + } +} + +#[async_trait] +impl ToolHandler for McpToolHandler { + fn definition(&self) -> Tool { + Tool { + tool_type: "function".to_string(), + function: FunctionDefinition { + name: self.prefixed_name.clone(), + description: self.description.clone(), + parameters: self.schema.clone(), + }, + } + } + + async fn execute(&self, args: &str) -> Result { + self.client.call_tool(&self.tool_name, args).await + } +} From 40b284ddf27ef5297e0fd9ec793d1d77e233630e Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 09:31:05 +0300 Subject: [PATCH 3/6] feat: wire MCP tools into executor, agent, API, and CLI --- src/api/mod.rs | 3 +- src/cli/mod.rs | 9 +++--- src/core/agent.rs | 8 ++++-- src/lib.rs | 2 +- src/tools/mod.rs | 73 ++++++++++++++++++++++------------------------- 5 files changed, 48 insertions(+), 47 deletions(-) diff --git a/src/api/mod.rs b/src/api/mod.rs index 7dc4db6..3bd9f77 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -30,7 +30,8 @@ pub async fn start_api_server( tracing::info!(" WS /ws"); let llm_client: Arc = create_client(&config, &client); - let tool_executor: Arc = Arc::new(SystemToolExecutor::new()); + let tool_executor: Arc = + Arc::new(SystemToolExecutor::build(&app_config.mcp_servers).await); let rag_context = RagContext::new()?; let server = HttpServer::new(move || { diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 24456c9..22db096 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -72,7 +72,7 @@ struct SessionContext { } impl SessionContext { - fn new( + async fn new( client: &Client, agent_config: &AgentConfig, app_config: &AppConfig, @@ -88,7 +88,8 @@ impl SessionContext { create_client(agent_config, client), agent_config, )?; - let tool_executor: Arc = Arc::new(SystemToolExecutor::new()); + let tool_executor: Arc = + Arc::new(SystemToolExecutor::build(&app_config.mcp_servers).await); let rag = RagContext::new()?; let (conversation, prompt_builder) = rag.prepare( chat_id, @@ -177,7 +178,7 @@ pub async fn run_agent_mode( let chat_id = resolve_chat_id(chat_id, continue_last)?; let mut session = SessionContext::new( client, config, app_config, model_name, max_iterations, chat_id, skill_names, - )?; + ).await?; let chat_id_short = &session.conversation.meta.id.to_string()[..8]; let chat_status = if session.conversation.messages.is_empty() { @@ -290,7 +291,7 @@ pub async fn run_single_prompt( let mut session = SessionContext::new( client, config, app_config, model_name, max_iterations, chat_id, skill_names, - )?; + ).await?; session.conversation.messages.push(Message::user(prompt.clone())); diff --git a/src/core/agent.rs b/src/core/agent.rs index 204b509..b193545 100644 --- a/src/core/agent.rs +++ b/src/core/agent.rs @@ -5,7 +5,7 @@ use crate::core::llm::LlmClient; use crate::core::models::*; use crate::error::Result; use crate::rag::PromptBuilder; -use crate::tools::{get_available_tools, ToolExecutor}; +use crate::tools::ToolExecutor; async fn call_llm( llm: &Arc, @@ -31,7 +31,7 @@ async fn run_agent_loop( verbose: bool, mut callback: Option<&mut dyn FnMut(StreamEvent)>, ) -> Result { - let tools = get_available_tools(); + let tools = tool_executor.list_tools(); let mut steps = Vec::new(); let mut final_response = String::new(); @@ -284,6 +284,10 @@ mod tests { #[async_trait] impl ToolExecutor for MockToolExecutor { + fn list_tools(&self) -> Vec { + vec![] + } + async fn execute(&self, name: &str, args_json: &str) -> Result { self.calls.lock().unwrap().push((name.into(), args_json.into())); Ok(self.result.clone()) diff --git a/src/lib.rs b/src/lib.rs index 6bd50f2..16f9d4b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod cli; pub mod config; pub mod core; pub mod error; +pub mod mcp; pub mod rag; pub mod tools; @@ -11,6 +12,5 @@ pub use core::{agent, llm, models}; pub use error::{Error, Result}; pub use models::*; -pub use tools::{execute_tool, get_available_tools}; pub use llm::{LlmClient, OpenAiClient, OpenAiCompatibleClient, AnthropicClient, GeminiClient}; pub use rag::{RagContext, PromptBuilder}; diff --git a/src/tools/mod.rs b/src/tools/mod.rs index d662909..01c52f9 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -2,11 +2,11 @@ mod execute_command; mod read_file; mod write_file; -use std::collections::HashMap; -use std::sync::OnceLock; +use std::collections::{BTreeMap, HashMap}; use async_trait::async_trait; +use crate::config::McpServerConfig; use crate::core::models::Tool; use crate::error::{Error, Result}; @@ -21,6 +21,7 @@ pub trait ToolHandler: Send + Sync { #[async_trait] pub trait ToolExecutor: Send + Sync { + fn list_tools(&self) -> Vec; async fn execute(&self, name: &str, args_json: &str) -> Result; } @@ -30,17 +31,29 @@ pub struct SystemToolExecutor { impl SystemToolExecutor { pub fn new() -> Self { - let mut executor = Self { - handlers: HashMap::new(), - }; - executor.register(Box::new(execute_command::ExecuteCommandTool)); - executor.register(Box::new(read_file::ReadFileTool)); - executor.register(Box::new(write_file::WriteFileTool)); + Self { handlers: HashMap::new() } + } + + pub async fn build(mcp_configs: &BTreeMap) -> Self { + let mut executor = Self::new(); + executor.register_builtins(); + for handler in crate::mcp::load_mcp_tools(mcp_configs).await { + executor.register(handler); + } executor } - fn register(&mut self, handler: Box) { + pub fn register_builtins(&mut self) { + self.register(Box::new(execute_command::ExecuteCommandTool)); + self.register(Box::new(read_file::ReadFileTool)); + self.register(Box::new(write_file::WriteFileTool)); + } + + pub fn register(&mut self, handler: Box) { let name = handler.definition().function.name.clone(); + if self.handlers.contains_key(&name) { + tracing::warn!(name = %name, "Tool name collision: overwriting existing handler"); + } self.handlers.insert(name, handler); } } @@ -53,6 +66,10 @@ impl Default for SystemToolExecutor { #[async_trait] impl ToolExecutor for SystemToolExecutor { + fn list_tools(&self) -> Vec { + self.handlers.values().map(|h| h.definition()).collect() + } + async fn execute(&self, name: &str, args_json: &str) -> Result { let handler = self .handlers @@ -62,32 +79,20 @@ impl ToolExecutor for SystemToolExecutor { } } -/// Global executor used by the free-standing helper functions. -static GLOBAL_EXECUTOR: OnceLock = OnceLock::new(); - -fn global_executor() -> &'static SystemToolExecutor { - GLOBAL_EXECUTOR.get_or_init(SystemToolExecutor::new) -} - -pub fn get_available_tools() -> Vec { - global_executor() - .handlers - .values() - .map(|h| h.definition()) - .collect() -} - -pub async fn execute_tool(name: &str, arguments: &str) -> Result { - global_executor().execute(name, arguments).await -} - #[cfg(test)] mod tests { use super::*; #[test] - fn executor_registers_all_tools() { + fn new_executor_is_empty() { let executor = SystemToolExecutor::new(); + assert_eq!(executor.handlers.len(), 0); + } + + #[test] + fn register_builtins_adds_three_tools() { + let mut executor = SystemToolExecutor::new(); + executor.register_builtins(); assert!(executor.handlers.contains_key("execute_command")); assert!(executor.handlers.contains_key("read_file")); assert!(executor.handlers.contains_key("write_file")); @@ -101,14 +106,4 @@ mod tests { assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Unknown tool")); } - - #[test] - fn get_available_tools_returns_three_definitions() { - let tools = get_available_tools(); - assert_eq!(tools.len(), 3); - let names: Vec<&str> = tools.iter().map(|t| t.function.name.as_str()).collect(); - assert!(names.contains(&"execute_command")); - assert!(names.contains(&"read_file")); - assert!(names.contains(&"write_file")); - } } From 0324bafe4700a4114c37de95ad09b7586c3d21fc Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 16:55:53 +0300 Subject: [PATCH 4/6] fix(mcp): use list_all_tools instead of list_tools --- src/mcp/client.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/mcp/client.rs b/src/mcp/client.rs index e16bd6f..8ecdc23 100644 --- a/src/mcp/client.rs +++ b/src/mcp/client.rs @@ -46,13 +46,10 @@ impl McpClient { } pub async fn list_tools(&self) -> Result> { - let result = self - .service - .peer() - .list_tools(Default::default()) + self.service + .list_all_tools() .await - .map_err(|e| Error::Other(format!("MCP list_tools failed for '{}': {}", self.server_name, e)))?; - Ok(result.tools.into_iter().collect()) + .map_err(|e| Error::Other(format!("MCP list_tools failed for '{}': {}", self.server_name, e))) } pub async fn call_tool(&self, name: &str, args_json: &str) -> Result { From 8775ac2da535e88490ce201c8cfd7ebeff7c259d Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 17:03:34 +0300 Subject: [PATCH 5/6] fix(mcp): better extract_text_content parsing --- src/mcp/client.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/mcp/client.rs b/src/mcp/client.rs index 8ecdc23..d024a42 100644 --- a/src/mcp/client.rs +++ b/src/mcp/client.rs @@ -1,6 +1,6 @@ use rmcp::{ ServiceExt, - model::{CallToolRequestParams, Tool}, + model::{CallToolRequestParams, Content, RawContent, ResourceContents, Tool}, service::{RoleClient, RunningService}, transport::{TokioChildProcess, streamable_http_client::StreamableHttpClientTransport}, }; @@ -79,13 +79,20 @@ fn build_call_params(name: &str, args_json: &str) -> Result String { +fn extract_text_content(content: &[Content]) -> String { content .iter() - .filter_map(|c| { - serde_json::to_value(c) - .ok() - .and_then(|v| v.get("text")?.as_str().map(String::from)) + .map(|item| match &**item { + RawContent::Text(t) => t.text.clone(), + RawContent::Image(i) => format!("[image: {}]", i.mime_type), + RawContent::Audio(a) => format!("[audio: {}]", a.mime_type), + RawContent::Resource(r) => match &r.resource { + ResourceContents::TextResourceContents { text, .. } => text.clone(), + ResourceContents::BlobResourceContents { uri, mime_type, .. } => { + format!("[blob: {} ({})]", uri, mime_type.as_deref().unwrap_or("unknown")) + } + }, + RawContent::ResourceLink(l) => format!("[resource: {}]", l.uri), }) .collect::>() .join("\n") From 3b968e148ca9826aa75e21f1c73680178c9fa1f9 Mon Sep 17 00:00:00 2001 From: themartto Date: Sun, 3 May 2026 17:12:55 +0300 Subject: [PATCH 6/6] fix(mcp): fix sanitise prefix --- src/mcp/mod.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index a488d6c..09307f7 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -37,7 +37,10 @@ async fn connect_server(name: &str, config: &McpServerConfig) -> Result