use crate::{
agent::cancellation::AgentCancellation,
config::{McpServerConfig, McpStdioServerConfig},
mcp::{
McpError, McpResult,
http::HttpConnection,
jsonrpc::RequestId,
protocol::{
CallToolParams, CallToolResult, InitializeRequestParams, InitializeResult,
ListToolsParams, ListToolsResult, METHOD_INITIALIZE, METHOD_INITIALIZED,
METHOD_TOOLS_CALL, METHOD_TOOLS_LIST, PROTOCOL_VERSION, Tool,
},
stdio::StdioConnection,
},
};
use serde::de::DeserializeOwned;
use serde_json::Value;
use std::{
collections::HashSet,
path::Path,
sync::{Arc, Mutex, atomic::AtomicU64},
};
const MAX_TOOLS_LIST_PAGES: usize = 100;
pub(crate) struct McpClient {
connection: McpConnection,
next_id: AtomicU64,
server_info: Arc<Mutex<Option<crate::mcp::protocol::Implementation>>>,
}
impl std::fmt::Debug for McpClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("McpClient")
.field("connection", &self.connection)
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for McpConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Stdio(_) => f.debug_tuple("Stdio").field(&"<stdio connection>").finish(),
Self::Http(connection) => f.debug_tuple("Http").field(connection).finish(),
}
}
}
enum McpConnection {
Stdio(StdioConnection),
Http(HttpConnection),
}
impl McpConnection {
fn send_request(
&self,
id: RequestId,
method: &str,
params: Option<Value>,
cancellation: Option<&AgentCancellation>,
) -> McpResult<Value> {
match self {
Self::Stdio(connection) => connection.send_request(id, method, params, cancellation),
Self::Http(connection) => connection.send_request(id, method, params, cancellation),
}
}
fn send_notification(&self, method: &str, params: Option<Value>) -> McpResult<()> {
match self {
Self::Stdio(connection) => connection.send_notification(method, params),
Self::Http(connection) => connection.send_notification(method, params),
}
}
fn shutdown(&mut self) {
match self {
Self::Stdio(connection) => connection.shutdown(),
Self::Http(connection) => connection.shutdown(),
}
}
#[cfg(test)]
fn is_process_running(&self) -> bool {
match self {
Self::Stdio(connection) => connection.is_process_running(),
Self::Http(_) => false,
}
}
}
impl McpClient {
pub(crate) fn connect_stdio(config: &McpStdioServerConfig) -> McpResult<Self> {
Ok(Self {
connection: McpConnection::Stdio(StdioConnection::connect(config)?),
next_id: AtomicU64::new(1),
server_info: Arc::new(Mutex::new(None)),
})
}
#[cfg(test)]
pub(crate) fn connect(config: &McpServerConfig) -> McpResult<Self> {
Self::connect_named(None, config, None)
}
pub(crate) fn connect_named(
server_name: Option<&str>,
config: &McpServerConfig,
mc_home: Option<&Path>,
) -> McpResult<Self> {
match config {
McpServerConfig::Stdio(config) => Self::connect_stdio(config),
McpServerConfig::Http(config) => Ok(Self {
connection: McpConnection::Http(HttpConnection::connect_named(
server_name,
config,
mc_home,
)?),
next_id: AtomicU64::new(1),
server_info: Arc::new(Mutex::new(None)),
}),
}
}
fn next_id(&self) -> RequestId {
let id = self
.next_id
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
RequestId::Number(i64::try_from(id).unwrap_or(i64::MAX))
}
fn send_request_inner(
&self,
method: &str,
params: Option<Value>,
cancellation: Option<&AgentCancellation>,
) -> McpResult<Value> {
let id = self.next_id();
self.connection
.send_request(id, method, params, cancellation)
}
#[cfg(test)]
fn is_process_running(&self) -> bool {
self.connection.is_process_running()
}
pub(crate) fn send_request(&self, method: &str, params: Option<Value>) -> McpResult<Value> {
self.send_request_inner(method, params, None)
}
pub(crate) fn send_request_cancellable(
&self,
method: &str,
params: Option<Value>,
cancellation: &AgentCancellation,
) -> McpResult<Value> {
self.send_request_inner(method, params, Some(cancellation))
}
pub(crate) fn send_notification(&self, method: &str, params: Option<Value>) -> McpResult<()> {
self.connection.send_notification(method, params)
}
pub(crate) fn initialize(&self) -> McpResult<InitializeResult> {
let params = serde_json::to_value(InitializeRequestParams::default())
.map_err(McpError::transport)?;
let result: InitializeResult = decode(self.send_request(METHOD_INITIALIZE, Some(params))?)?;
if result.protocol_version != PROTOCOL_VERSION {
return Err(McpError::Protocol {
code: -32602,
message: format!(
"unsupported MCP protocol version {}; expected {PROTOCOL_VERSION}",
result.protocol_version
),
});
}
if let Ok(mut server_info) = self.server_info.lock() {
*server_info = Some(result.server_info.clone());
}
self.send_notification(METHOD_INITIALIZED, None)?;
Ok(result)
}
pub(crate) fn list_tools(&self) -> McpResult<Vec<Tool>> {
let mut tools = Vec::new();
let mut cursor = None;
let mut seen_cursors = HashSet::new();
for page in 0..MAX_TOOLS_LIST_PAGES {
let params = serde_json::to_value(ListToolsParams {
cursor: cursor.clone(),
})
.map_err(McpError::transport)?;
let result: ListToolsResult =
decode(self.send_request(METHOD_TOOLS_LIST, Some(params))?)?;
tools.extend(result.tools);
cursor = result.next_cursor;
if let Some(next_cursor) = cursor.as_ref() {
if !seen_cursors.insert(next_cursor.clone()) {
return Err(McpError::Protocol {
code: -32603,
message: "MCP tools/list returned repeated pagination cursor".to_string(),
});
}
} else {
return Ok(tools);
}
if page + 1 == MAX_TOOLS_LIST_PAGES {
return Err(McpError::Protocol {
code: -32603,
message: format!("MCP tools/list exceeded {MAX_TOOLS_LIST_PAGES} pages"),
});
}
}
unreachable!("tools/list pagination loop returns from inside bounded range")
}
#[cfg(test)]
pub(crate) fn call_tool(
&self,
name: &str,
arguments: Option<Value>,
) -> McpResult<CallToolResult> {
let params = serde_json::to_value(CallToolParams {
name: name.to_string(),
arguments,
})
.map_err(McpError::transport)?;
decode(self.send_request(METHOD_TOOLS_CALL, Some(params))?)
}
pub(crate) fn call_tool_cancellable(
&self,
name: &str,
arguments: Option<Value>,
cancellation: &AgentCancellation,
) -> McpResult<CallToolResult> {
let params = serde_json::to_value(CallToolParams {
name: name.to_string(),
arguments,
})
.map_err(McpError::transport)?;
decode(self.send_request_cancellable(METHOD_TOOLS_CALL, Some(params), cancellation)?)
}
pub(crate) fn shutdown(&mut self) {
self.connection.shutdown();
}
}
impl Drop for McpClient {
fn drop(&mut self) {
self.shutdown();
}
}
fn decode<T: DeserializeOwned>(value: Value) -> McpResult<T> {
serde_json::from_value(value).map_err(McpError::transport)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::BTreeMap;
use std::sync::{Arc, atomic::AtomicBool};
use std::thread;
use std::time::{Duration, Instant};
fn mock_config(mode: &str, timeout: u64) -> McpStdioServerConfig {
let mut env = BTreeMap::new();
env.insert("MCP_MOCK_MODE".to_string(), mode.to_string());
McpStdioServerConfig {
command: "sh".to_string(),
args: vec![format!(
"{}/tests/fixtures/mcp/mock_stdio_server.sh",
env!("CARGO_MANIFEST_DIR")
)],
env,
enabled: true,
timeout: Some(timeout),
}
}
#[test]
fn mock_server_initializes_lists_and_calls_tool() {
let mut client = McpClient::connect_stdio(&mock_config("normal", 2)).unwrap();
let init = client.initialize().unwrap();
assert_eq!(init.server_info.name, "mock-mcp");
let tools = client.list_tools().unwrap();
assert_eq!(tools.len(), 2);
assert!(tools.iter().any(|tool| tool.name == "echo"));
let result = client
.call_tool("echo", Some(serde_json::json!({"text":"hello"})))
.unwrap();
assert_eq!(result.is_error, Some(false));
assert!(matches!(
&result.content[0],
crate::mcp::ContentBlock::Text { text } if text == "hello"
));
client.shutdown();
}
#[test]
fn mock_server_reports_protocol_and_tool_errors() {
let mut protocol = McpClient::connect_stdio(&mock_config("protocol_error", 2)).unwrap();
protocol.initialize().unwrap();
let error = protocol.list_tools().unwrap_err().to_string();
assert!(error.contains("mock protocol error"), "{error}");
protocol.shutdown();
let mut tool = McpClient::connect_stdio(&mock_config("tool_error", 2)).unwrap();
tool.initialize().unwrap();
let result = tool.call_tool("echo", Some(serde_json::json!({}))).unwrap();
assert_eq!(result.is_error, Some(true));
tool.shutdown();
}
#[test]
fn mock_server_timeout_is_bounded() {
let client = McpClient::connect_stdio(&mock_config("hang", 1)).unwrap();
let error = client.initialize().unwrap_err().to_string();
assert!(error.contains("timed out"), "{error}");
assert!(!client.is_process_running());
}
#[test]
fn canceled_call_terminates_hanging_server() {
let client = McpClient::connect_stdio(&mock_config("hang", 5)).unwrap();
let cancel_flag = Arc::new(AtomicBool::new(false));
let cancellation =
crate::agent::cancellation::AgentCancellation::new(Arc::clone(&cancel_flag));
let canceler = {
let cancel_flag = Arc::clone(&cancel_flag);
thread::spawn(move || {
std::thread::sleep(Duration::from_millis(150));
cancel_flag.store(true, std::sync::atomic::Ordering::SeqCst);
})
};
let error = client
.send_request_cancellable(METHOD_INITIALIZE, None, &cancellation)
.unwrap_err()
.to_string();
canceler.join().unwrap();
assert!(error.contains("prompt canceled"), "{error}");
assert!(!client.is_process_running());
}
#[test]
fn malformed_json_fails_pending_request_without_timeout_delay() {
let mut client = McpClient::connect_stdio(&mock_config("malformed_json", 5)).unwrap();
let started = Instant::now();
let error = client.initialize().unwrap_err().to_string();
assert!(
started.elapsed() < Duration::from_secs(2),
"elapsed: {:?}",
started.elapsed()
);
assert!(
error.contains("stdout closed") || error.contains("malformed"),
"{error}"
);
client.shutdown();
}
#[test]
fn list_tools_accepts_missing_input_schema() {
let mut client = McpClient::connect_stdio(&mock_config("missing_input_schema", 2)).unwrap();
client.initialize().unwrap();
let tools = client.list_tools().unwrap();
assert_eq!(tools[0].name, "no_schema");
assert_eq!(
tools[0].input_schema,
serde_json::json!({"type":"object","properties":{}})
);
client.shutdown();
}
#[test]
fn list_changed_notification_is_ignored() {
let mut client = McpClient::connect_stdio(&mock_config("list_changed", 2)).unwrap();
client.initialize().unwrap();
let tools = client.list_tools().unwrap();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "echo");
client.shutdown();
}
#[test]
fn list_tools_rejects_repeated_cursor() {
let mut client = McpClient::connect_stdio(&mock_config("repeated_cursor", 2)).unwrap();
client.initialize().unwrap();
let error = client.list_tools().unwrap_err().to_string();
assert!(error.contains("repeated pagination cursor"), "{error}");
client.shutdown();
}
#[test]
fn list_tools_rejects_excessive_pagination() {
let mut client =
McpClient::connect_stdio(&mock_config("infinite_unique_cursor", 2)).unwrap();
client.initialize().unwrap();
let error = client.list_tools().unwrap_err().to_string();
assert!(error.contains("exceeded 100 pages"), "{error}");
client.shutdown();
}
#[test]
fn initialize_rejects_unsupported_protocol_version() {
let mut client = McpClient::connect_stdio(&mock_config("old_protocol", 2)).unwrap();
let error = client.initialize().unwrap_err().to_string();
assert!(
error.contains("unsupported MCP protocol version"),
"{error}"
);
client.shutdown();
}
#[test]
fn http_connection_variant_connects_without_process() {
let config = crate::config::McpHttpServerConfig {
url: "https://mcp.example.test/mcp".to_string(),
headers: BTreeMap::new(),
oauth: None,
enabled: true,
timeout: Some(1),
};
let client = McpClient::connect(&McpServerConfig::Http(config)).unwrap();
assert!(!client.is_process_running());
}
}