asepharyana
2efd40ca88
- Added module-level documentation for memory tools (`remember`, `recall`, `forget`) to clarify their purpose. - Improved documentation in `recall.rs` and `remember.rs` to describe the functionality and flow of memory entry operations. - Updated `mod.rs` to include descriptions for the tool trait and execution context. - Enhanced `plan.rs` with detailed comments on plan-mode signaling tools. - Documented text search tools in `search.rs` to explain their functionality. - Improved sequential-thinking tool documentation in `seqthink.rs`. - Added safety filter documentation in `shell_filter` for credential and git operations. - Enhanced utility tools documentation, including `cd`, `dir_cache_update`, and `todowrite`. - Improved rendering documentation in view modules (`chat`, `markdown`, `status`, `workflow`) to clarify rendering flows and purposes.
408 lines
15 KiB
Rust
408 lines
15 KiB
Rust
//! MCP server connection management: spawning/talking to stdio child
|
|
//! processes and HTTP endpoints, and adapting their advertised tools to
|
|
//! the crate's `Tool` trait.
|
|
|
|
use serde_json::{json, Value};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::io::{BufRead, BufReader, Write};
|
|
use std::sync::{Arc, Mutex, OnceLock};
|
|
|
|
|
|
const MCP_CONNECT_TIMEOUT_MS: u64 = 20_000;
|
|
const MCP_CALL_TIMEOUT_MS: u64 = 60_000;
|
|
|
|
/// Global cache for `&'static str` names/descriptions of MCP tools, so we
|
|
/// never need `Box::leak`. Entries are never removed (small, bounded by the
|
|
/// number of MCP tools ever registered in a session).
|
|
fn mcp_static_str(s: &str) -> &'static str {
|
|
static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
|
|
let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
|
|
if let Some(&existing) = cache.iter().find(|e| **e == s) {
|
|
return existing;
|
|
}
|
|
let leaked: &'static str = Box::leak(s.to_string().into_boxed_str());
|
|
cache.push(leaked);
|
|
leaked
|
|
}
|
|
|
|
/// How an MCP server is reached: a spawned child process talking
|
|
/// newline-delimited JSON-RPC over stdio, or a remote HTTP endpoint.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum McpTransport {
|
|
Stdio {
|
|
command: String,
|
|
args: Vec<String>,
|
|
},
|
|
StreamableHttp {
|
|
url: String,
|
|
},
|
|
}
|
|
|
|
/// A single tool advertised by an MCP server, as returned by `tools/list`.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct McpToolInfo {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub input_schema: Value,
|
|
}
|
|
|
|
/// A connected MCP server: its transport, advertised tools, and (for stdio)
|
|
/// a live handle to the child process.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct McpServer {
|
|
pub name: String,
|
|
pub transport: McpTransport,
|
|
pub tools: Vec<McpToolInfo>,
|
|
/// Held child-process handle so subsequent tool calls reuse the same
|
|
/// connection instead of spawning a new child each time. Not serialized
|
|
/// because the child only lives in this process.
|
|
#[serde(skip)]
|
|
pub child_handle: Option<Arc<Mutex<StdioChild>>>,
|
|
}
|
|
|
|
/// Live handle to an MCP server child process communicating over stdio
|
|
/// via newline-delimited JSON-RPC 2.0.
|
|
#[derive(Debug)]
|
|
pub struct StdioChild {
|
|
stdin: std::process::ChildStdin,
|
|
stdout: BufReader<std::process::ChildStdout>,
|
|
next_id: u64,
|
|
}
|
|
|
|
impl StdioChild {
|
|
/// Send a JSON-RPC request to the child and block for its matching response.
|
|
///
|
|
/// Flow: assign the next request id → write request + newline to stdin →
|
|
/// loop reading lines from stdout until one has a matching `id` or the
|
|
/// timeout elapses → return its `result` (or error out on an `error` field).
|
|
///
|
|
/// Why: the child may interleave unrelated/malformed lines, so blank
|
|
/// lines are skipped and non-matching ids are ignored rather than
|
|
/// treated as a protocol violation.
|
|
///
|
|
/// Return: the `result` value of the matching response, or `Err` on
|
|
/// timeout, EOF, JSON-RPC error, or I/O failure.
|
|
pub fn call(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
|
|
self.next_id += 1;
|
|
let id = self.next_id;
|
|
let req = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": id,
|
|
"method": method,
|
|
"params": params
|
|
});
|
|
let mut line = serde_json::to_string(&req)?;
|
|
line.push('\n');
|
|
self.stdin.write_all(line.as_bytes())?;
|
|
self.stdin.flush()?;
|
|
|
|
let mut response_line = String::new();
|
|
let deadline = std::time::Instant::now()
|
|
+ std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS);
|
|
loop {
|
|
if std::time::Instant::now() > deadline {
|
|
anyhow::bail!("MCP call timed out after {}ms", MCP_CALL_TIMEOUT_MS);
|
|
}
|
|
response_line.clear();
|
|
match self.stdout.read_line(&mut response_line) {
|
|
Ok(0) => anyhow::bail!("MCP stdio child process closed unexpectedly"),
|
|
Ok(_) => {
|
|
let trimmed = response_line.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
let resp: Value = serde_json::from_str(trimmed)
|
|
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
|
|
if resp.get("id") == Some(&json!(id)) {
|
|
if let Some(err) = resp.get("error") {
|
|
anyhow::bail!("MCP error: {}", err);
|
|
}
|
|
return Ok(resp.get("result").cloned().unwrap_or_else(|| {
|
|
tracing::warn!("[mcp] stdio response missing 'result' field: {}", trimmed);
|
|
Value::Null
|
|
}));
|
|
}
|
|
}
|
|
Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow::Result<StdioChild> {
|
|
let parts: Vec<&str> = command.split_whitespace().collect();
|
|
let (prog, prog_args) = parts.split_first()
|
|
.ok_or_else(|| anyhow::anyhow!("MCP stdio command is empty"))?;
|
|
|
|
let mut cmd = std::process::Command::new(prog);
|
|
cmd.args(prog_args);
|
|
cmd.args(extra_args);
|
|
cmd.stdin(std::process::Stdio::piped());
|
|
cmd.stdout(std::process::Stdio::piped());
|
|
cmd.stderr(std::process::Stdio::null());
|
|
|
|
let mut child = cmd.spawn()
|
|
.map_err(|e| anyhow::anyhow!("failed to spawn MCP stdio server '{}': {}", command, e))?;
|
|
|
|
let stdin = child.stdin.take()
|
|
.ok_or_else(|| anyhow::anyhow!("failed to get stdin for MCP server"))?;
|
|
let stdout = child.stdout.take()
|
|
.ok_or_else(|| anyhow::anyhow!("failed to get stdout for MCP server"))?;
|
|
|
|
let mut mcp = StdioChild {
|
|
stdin,
|
|
stdout: BufReader::new(stdout),
|
|
next_id: 0,
|
|
};
|
|
|
|
let deadline = std::time::Instant::now()
|
|
+ std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS);
|
|
|
|
let init_result = mcp.call("initialize", json!({
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {},
|
|
"clientInfo": {
|
|
"name": "zesdex",
|
|
"version": "0.1.0"
|
|
}
|
|
}));
|
|
|
|
if std::time::Instant::now() > deadline {
|
|
anyhow::bail!("MCP initialize timed out");
|
|
}
|
|
|
|
init_result.map_err(|e| anyhow::anyhow!("MCP initialize failed: {}", e))?;
|
|
|
|
let _ = mcp.call("notifications/initialized", json!({}));
|
|
|
|
Ok(mcp)
|
|
}
|
|
|
|
fn call_via_stdio(
|
|
existing_handle: Option<&Mutex<StdioChild>>,
|
|
command: &str,
|
|
extra_args: &[String],
|
|
tool_name: &str,
|
|
tool_args: &Value,
|
|
) -> anyhow::Result<String> {
|
|
// Reuse the persistent child handle if available; otherwise spawn a new one.
|
|
let mut guard;
|
|
let child: &mut StdioChild = if let Some(mtx) = existing_handle {
|
|
guard = mtx.lock().map_err(|e| anyhow::anyhow!("MCP handle lock: {}", e))?;
|
|
&mut guard
|
|
} else {
|
|
let mut fresh = spawn_stdio_child(command, extra_args)?;
|
|
let result = fresh.call("tools/call", json!({
|
|
"name": tool_name,
|
|
"arguments": tool_args
|
|
}))?;
|
|
return extract_text_content(&result);
|
|
};
|
|
|
|
let result = child.call("tools/call", json!({
|
|
"name": tool_name,
|
|
"arguments": tool_args
|
|
}))?;
|
|
|
|
extract_text_content(&result)
|
|
}
|
|
|
|
fn call_via_http(url: &str, tool_name: &str, tool_args: &Value) -> anyhow::Result<String> {
|
|
let client = reqwest::blocking::Client::builder()
|
|
.timeout(std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS))
|
|
.connect_timeout(std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS))
|
|
.build()
|
|
.unwrap_or_else(|e| {
|
|
tracing::warn!("[mcp] HTTP client builder failed: {}, using default client without timeouts", e);
|
|
reqwest::blocking::Client::new()
|
|
});
|
|
|
|
let request_id: u64 = 1;
|
|
let body = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": tool_name,
|
|
"arguments": tool_args
|
|
}
|
|
});
|
|
|
|
let resp = client.post(url)
|
|
.header("Content-Type", "application/json")
|
|
.json(&body)
|
|
.send()
|
|
.map_err(|e| anyhow::anyhow!("MCP HTTP request failed: {}", e))?;
|
|
|
|
if !resp.status().is_success() {
|
|
let status = resp.status();
|
|
let text = resp.text().unwrap_or_else(|e| {
|
|
tracing::warn!("[mcp] failed to read HTTP response body: {}", e);
|
|
String::new()
|
|
});
|
|
anyhow::bail!("MCP HTTP server returned {}: {}", status, text);
|
|
}
|
|
|
|
let response: Value = resp.json()
|
|
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP HTTP server: {}", e))?;
|
|
|
|
if let Some(err) = response.get("error") {
|
|
anyhow::bail!("MCP HTTP error: {}", err);
|
|
}
|
|
|
|
let result = response.get("result").cloned().unwrap_or_else(|| {
|
|
tracing::warn!("[mcp] HTTP response missing 'result' field");
|
|
Value::Null
|
|
});
|
|
extract_text_content(&result)
|
|
}
|
|
|
|
fn extract_text_content(result: &Value) -> anyhow::Result<String> {
|
|
if let Some(content) = result.get("content") {
|
|
if let Some(arr) = content.as_array() {
|
|
let text: Vec<String> = arr.iter().filter_map(|item| {
|
|
if item.get("type").and_then(|t| t.as_str()) == Some("text") {
|
|
item.get("text").and_then(|t| t.as_str()).map(|s| s.to_string())
|
|
} else {
|
|
None
|
|
}
|
|
}).collect();
|
|
if !text.is_empty() {
|
|
return Ok(text.join("\n"));
|
|
}
|
|
}
|
|
}
|
|
Ok(serde_json::to_string_pretty(result).unwrap_or_else(|e| {
|
|
tracing::warn!("[mcp] failed to pretty-print result: {}", e);
|
|
result.to_string()
|
|
}))
|
|
}
|
|
|
|
/// Registry of connected MCP servers and their tools for the current session.
|
|
#[derive(Debug, Clone)]
|
|
pub struct McpManager {
|
|
pub servers: Vec<McpServer>,
|
|
}
|
|
|
|
/// Adapts a single MCP-advertised tool to the crate's `Tool` trait so it can
|
|
/// be dispatched through the same execution path as built-in tools.
|
|
pub struct McpToolAdapter {
|
|
pub tool_name: String,
|
|
pub server_name: String,
|
|
pub transport: McpTransport,
|
|
pub description: String,
|
|
pub parameters: Value,
|
|
/// Shared handle to a persistent child process (stdio transport only).
|
|
pub child_handle: Option<Arc<Mutex<StdioChild>>>,
|
|
}
|
|
|
|
impl crate::tool::Tool for McpToolAdapter {
|
|
fn name(&self) -> &'static str {
|
|
mcp_static_str(&format!("mcp__{}__{}", self.server_name, self.tool_name))
|
|
}
|
|
|
|
fn description(&self) -> &'static str {
|
|
mcp_static_str(&self.description)
|
|
}
|
|
|
|
fn parameters(&self) -> Value {
|
|
self.parameters.clone()
|
|
}
|
|
|
|
fn run(&self, _ctx: &crate::tool::ToolCtx, args: &Value) -> anyhow::Result<String> {
|
|
match &self.transport {
|
|
McpTransport::Stdio { command, args: extra_args } => {
|
|
call_via_stdio(self.child_handle.as_ref().map(|h| h.as_ref()), command, extra_args, &self.tool_name, args)
|
|
}
|
|
McpTransport::StreamableHttp { url } => {
|
|
call_via_http(url, &self.tool_name, args)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl McpManager {
|
|
/// Create an empty manager with no connected servers.
|
|
pub fn new() -> Self {
|
|
McpManager {
|
|
servers: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Flatten all connected servers' tools into a single list of `Tool` trait objects.
|
|
///
|
|
/// Flow: for each server, clone its child handle → wrap each of its
|
|
/// `McpToolInfo` entries in an `McpToolAdapter` sharing that handle.
|
|
///
|
|
/// Why: the handle is cloned (Arc) per tool so every adapter for a given
|
|
/// stdio server reuses the same persistent child process/connection.
|
|
///
|
|
/// Return: boxed `Tool` trait objects ready to merge into the harness's tool list.
|
|
pub fn as_tools(&self) -> Vec<Box<dyn crate::tool::Tool>> {
|
|
self.servers.iter().flat_map(|server| {
|
|
let handle = server.child_handle.clone();
|
|
server.tools.iter().map(move |info| {
|
|
Box::new(McpToolAdapter {
|
|
tool_name: info.name.clone(),
|
|
server_name: server.name.clone(),
|
|
transport: server.transport.clone(),
|
|
description: info.description.clone(),
|
|
parameters: info.input_schema.clone(),
|
|
child_handle: handle.clone(),
|
|
}) as Box<dyn crate::tool::Tool>
|
|
})
|
|
}).collect()
|
|
}
|
|
|
|
/// Connects to an MCP server via stdio by spawning the child process, running
|
|
/// the `initialize` handshake, calling `tools/list`, and registering the server
|
|
/// with its advertised tools in `self.servers`. The child process stays alive
|
|
/// for subsequent `tools/call` invocations via the stored `McpServer.tools`.
|
|
pub fn connect_stdio(&mut self, name: &str, command: &str, extra_args: &[String]) -> anyhow::Result<()> {
|
|
let transport = McpTransport::Stdio {
|
|
command: command.to_string(),
|
|
args: extra_args.to_vec(),
|
|
};
|
|
|
|
let mut child = spawn_stdio_child(command, extra_args)?;
|
|
let result = child.call("tools/list", json!({}))?;
|
|
|
|
let tools = if let Some(tool_list) = result.get("tools").and_then(|v| v.as_array()) {
|
|
tool_list.iter().filter_map(|t| {
|
|
Some(McpToolInfo {
|
|
name: t.get("name")?.as_str()?.to_string(),
|
|
description: t.get("description").and_then(|v| v.as_str()).unwrap_or_else(|| {
|
|
tracing::warn!("[mcp] tool {} missing description", t.get("name").and_then(|n| n.as_str()).unwrap_or("?"));
|
|
""
|
|
}).to_string(),
|
|
input_schema: t.get("inputSchema").cloned().unwrap_or_else(|| {
|
|
tracing::warn!("[mcp] tool {} missing inputSchema", t.get("name").and_then(|n| n.as_str()).unwrap_or("?"));
|
|
serde_json::Value::Null
|
|
}),
|
|
})
|
|
}).collect()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
let handle = Arc::new(Mutex::new(child));
|
|
|
|
self.servers.push(McpServer {
|
|
name: name.to_string(),
|
|
transport,
|
|
tools,
|
|
child_handle: Some(handle),
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Removes a server by name. Returns `true` if a server was found and removed.
|
|
#[allow(dead_code)]
|
|
pub fn disconnect(&mut self, name: &str) -> bool {
|
|
let len = self.servers.len();
|
|
self.servers.retain(|s| s.name != name);
|
|
self.servers.len() < len
|
|
}
|
|
}
|