use crate::http::{HttpError, HttpTransport, McpEndpoint};
use crate::inbound;
use crate::rpc::{self, RpcError};
use crate::wire::{
CallToolResult, CompleteParams, CompleteResult, Era, GetPromptParams, GetPromptResult,
Implementation, LATEST_MODERN_VERSION, ListResourceTemplatesResult, Prompt, ReadResourceResult,
Resource, ResourceTemplate, ServerCapabilities, Task, Tool, as_task_result, method,
};
use crate::modern;
use serde::Serialize;
use serde_json::{Value, json};
use std::collections::{HashMap, VecDeque};
use std::fmt;
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::Duration;
#[derive(Debug)]
pub enum McpError {
Transport(String),
Rpc(RpcError),
Timeout(String),
Capability(String),
}
impl fmt::Display for McpError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
McpError::Transport(m) => write!(f, "mcp: transport: {m}"),
McpError::Rpc(e) => write!(f, "mcp: rpc error {}: {}", e.code, e.message),
McpError::Timeout(m) => write!(f, "mcp: timeout: {m}"),
McpError::Capability(m) => write!(f, "mcp: capability: {m}"),
}
}
}
impl std::error::Error for McpError {}
type NotifQueue = Arc<Mutex<VecDeque<rpc::Notification>>>;
struct InboundRouter {
http: Arc<HttpTransport>,
queue: NotifQueue,
caps: inbound::Capabilities,
handler: Option<Arc<dyn inbound::Handler>>,
timeout: Duration,
}
impl InboundRouter {
fn route(&self, frame: Value) {
if let Some(req) = inbound::as_request(&frame) {
let resp = inbound::answer(&req, self.caps, self.handler.as_deref());
if let Ok(body) = serde_json::to_vec(&resp) {
let _ = self.http.send(None, &body, self.timeout, &[], |_| {});
}
return;
}
queue_notification(&self.queue, frame);
}
}
pub struct McpClient {
name: String,
http: Arc<HttpTransport>,
notifications: NotifQueue,
events: Mutex<Option<EventStreamHandle>>,
tool_schemas: Mutex<HashMap<String, Value>>,
next_id: AtomicI64,
caps: ServerCapabilities,
protocol_version: Option<String>,
era: Era,
timeout: Duration,
rmcp: Option<crate::rmcp_client::RmcpClient>,
endpoint: String,
extra_headers: Vec<(String, String)>,
inbound_caps: inbound::Capabilities,
inbound_handler: Option<Arc<dyn inbound::Handler>>,
tool_meta: Option<Value>,
client_info: Implementation,
client_capabilities: Value,
}
struct EventStreamHandle {
stop: Arc<AtomicBool>,
handle: JoinHandle<()>,
}
impl McpClient {
pub fn connect(
name: &str,
endpoint: &str,
headers: Vec<(String, String)>,
timeout: Duration,
) -> Result<McpClient, McpError> {
Self::connect_signed(name, endpoint, headers, timeout, None)
}
pub fn connect_signed(
name: &str,
endpoint: &str,
headers: Vec<(String, String)>,
timeout: Duration,
signer: Option<Arc<dyn crate::http::RequestSigner>>,
) -> Result<McpClient, McpError> {
let ep = McpEndpoint::parse(endpoint)
.map_err(|e| McpError::Transport(format!("mcp server '{name}': {e}")))?;
Ok(McpClient {
name: name.to_string(),
http: Arc::new(HttpTransport::new(ep, headers.clone()).with_signer(signer)),
notifications: Arc::new(Mutex::new(VecDeque::new())),
events: Mutex::new(None),
tool_schemas: Mutex::new(HashMap::new()),
next_id: AtomicI64::new(1),
caps: ServerCapabilities::default(),
protocol_version: None,
era: Era::Legacy,
timeout,
rmcp: None,
endpoint: endpoint.to_string(),
extra_headers: headers,
inbound_caps: inbound::Capabilities::default(),
inbound_handler: None,
tool_meta: None,
client_info: Implementation {
name: "agentd".into(),
version: env!("CARGO_PKG_VERSION").into(),
title: None,
},
client_capabilities: json!({}),
})
}
pub fn with_client_info(mut self, info: Implementation) -> Self {
self.client_info = info;
self
}
pub fn with_elicitation(mut self, handler: Arc<dyn inbound::Handler>) -> Self {
self.inbound_caps.elicitation = true;
self.inbound_handler = Some(handler);
self
}
pub fn with_roots(mut self, handler: Arc<dyn inbound::Handler>) -> Self {
self.inbound_caps.roots = true;
self.inbound_handler = Some(handler);
self
}
fn router(&self) -> Arc<InboundRouter> {
Arc::new(InboundRouter {
http: Arc::clone(&self.http),
queue: Arc::clone(&self.notifications),
caps: self.inbound_caps,
handler: self.inbound_handler.clone(),
timeout: self.timeout,
})
}
fn declared_capabilities(&self) -> Value {
let mut caps = self.client_capabilities.clone();
let inbound = self.inbound_caps.to_json();
match (caps.as_object_mut(), inbound.as_object()) {
(Some(dst), Some(src)) => {
for (k, v) in src {
dst.insert(k.clone(), v.clone());
}
Value::Object(dst.clone())
}
_ => caps,
}
}
pub fn with_tasks(mut self) -> Self {
self.client_capabilities = json!({
"extensions": { crate::wire::TASKS_EXTENSION: {} }
});
self
}
#[cfg(feature = "tls")]
pub fn with_identity(mut self, identity: net::tls::ClientIdentity) -> Self {
if let Some(h) = Arc::get_mut(&mut self.http) {
h.set_identity(Some(identity));
}
self
}
pub fn name(&self) -> &str {
&self.name
}
pub fn capabilities(&self) -> &ServerCapabilities {
&self.caps
}
pub fn set_tool_meta(&mut self, meta: Value) {
self.tool_meta = Some(meta);
}
pub fn initialize(&mut self) -> Result<(), McpError> {
self.initialize_within(self.timeout)
}
pub fn initialize_within(&mut self, timeout: Duration) -> Result<(), McpError> {
{
let mut b = crate::rmcp_client::RmcpBuilder::new(
&self.name,
&self.endpoint,
self.extra_headers.clone(),
timeout,
)
.with_http(Arc::clone(&self.http))
.with_client_info(self.client_info.clone());
if self.inbound_caps.elicitation
&& let Some(h) = &self.inbound_handler
{
b = b.with_elicitation(Arc::clone(h));
}
let c = b.connect()?;
self.caps = c.capabilities().clone();
self.protocol_version = c.protocol_version().map(str::to_string);
self.era = c
.protocol_version()
.map(crate::version::era_of)
.unwrap_or(Era::Legacy);
self.rmcp = Some(c);
Ok(())
}
}
pub fn era(&self) -> Era {
self.era
}
pub fn protocol_version(&self) -> Option<&str> {
self.protocol_version.as_deref()
}
pub fn list_tools(&self) -> Result<Vec<Tool>, McpError> {
self.list_tools_within(self.timeout)
}
pub fn list_tools_within(&self, _timeout: Duration) -> Result<Vec<Tool>, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.list_tools()
}
pub fn call_tool(
&self,
name: &str,
arguments: Option<Value>,
) -> Result<CallToolResult, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
let raw = c.call_tool_with_meta(name, arguments.clone(), None)?;
serde_json::from_value(raw).map_err(|e| {
McpError::Transport(format!("bad tools/call result on '{}': {e}", self.name))
})
}
pub fn call_tool_with_meta(
&self,
name: &str,
arguments: Option<Value>,
extra_meta: Value,
) -> Result<CallToolResult, McpError> {
self.call_tool_with_meta_within(name, arguments, extra_meta, self.timeout)
}
pub fn call_tool_with_meta_within(
&self,
name: &str,
arguments: Option<Value>,
extra_meta: Value,
timeout: Duration,
) -> Result<CallToolResult, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
let _ = timeout; let raw = c.call_tool_with_meta(name, arguments.clone(), Some(extra_meta.clone()))?;
serde_json::from_value(raw).map_err(|e| {
McpError::Transport(format!("bad tools/call result on '{}': {e}", self.name))
})
}
pub fn list_resources(&self) -> Result<Vec<Resource>, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.list_resources()
}
pub fn list_prompts(&self) -> Result<Vec<Prompt>, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.list_prompts()
}
pub fn get_prompt(
&self,
name: &str,
arguments: Option<Value>,
) -> Result<GetPromptResult, McpError> {
if !self.caps.supports_prompts() {
return Err(McpError::Capability(format!(
"server '{}' has no prompts",
self.name
)));
}
let params = GetPromptParams {
name: name.to_string(),
arguments,
};
self.request_as(method::PROMPTS_GET, Some(to_value(¶ms)))
}
pub fn complete(&self, reference: Value, argument: Value) -> Result<CompleteResult, McpError> {
if !self.caps.supports_completions() {
return Err(McpError::Capability(format!(
"server '{}' has no completions",
self.name
)));
}
let params = CompleteParams {
reference,
argument,
context: None,
};
self.request_as(method::COMPLETION_COMPLETE, Some(to_value(¶ms)))
}
pub fn list_resource_templates(&self) -> Result<Vec<ResourceTemplate>, McpError> {
if !self.caps.supports_resources() {
return Ok(Vec::new());
}
let mut templates = Vec::new();
let mut cursor: Option<String> = None;
loop {
let params = cursor.as_ref().map(|c| json!({ "cursor": c }));
let page: ListResourceTemplatesResult =
self.request_as(method::RESOURCES_TEMPLATES_LIST, params)?;
templates.extend(page.resource_templates);
match page.next_cursor {
Some(c) => cursor = Some(c),
None => break,
}
}
Ok(templates)
}
pub fn ping(&self) -> Result<(), McpError> {
self.request_with_timeout(method::PING, None, self.timeout)?;
Ok(())
}
pub fn as_task(&self, result: &Value) -> Option<Task> {
as_task_result(result)
}
pub fn get_task(&self, task_id: &str) -> Result<Task, McpError> {
self.request_as(method::TASKS_GET, Some(json!({ "taskId": task_id })))
}
pub fn update_task(&self, task_id: &str, input_responses: Value) -> Result<(), McpError> {
self.request_with_timeout(
method::TASKS_UPDATE,
Some(json!({ "taskId": task_id, "inputResponses": input_responses })),
self.timeout,
)?;
Ok(())
}
pub fn cancel_task(&self, task_id: &str) -> Result<(), McpError> {
self.request_with_timeout(
method::TASKS_CANCEL,
Some(json!({ "taskId": task_id })),
self.timeout,
)?;
Ok(())
}
pub fn await_task(
&self,
task_id: &str,
deadline: std::time::Instant,
) -> Result<Task, McpError> {
loop {
let task = self.get_task(task_id)?;
if task.is_terminal() || task.needs_input() {
return Ok(task);
}
if std::time::Instant::now() >= deadline {
return Err(McpError::Timeout(format!(
"task '{task_id}' on '{}' did not finish before the deadline",
self.name
)));
}
let poll = task.poll_interval_ms.unwrap_or(500).clamp(50, 5_000);
std::thread::sleep(Duration::from_millis(poll));
}
}
pub fn read_resource(&self, uri: &str) -> Result<ReadResourceResult, McpError> {
self.read_resource_within(uri, self.timeout)
}
pub fn read_resource_within(
&self,
uri: &str,
_timeout: Duration,
) -> Result<ReadResourceResult, McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.read_resource(uri)
}
pub fn subscribe(&self, uri: &str) -> Result<(), McpError> {
self.subscribe_within(uri, self.timeout)
}
pub fn subscribe_within(&self, uri: &str, _timeout: Duration) -> Result<(), McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.subscribe(uri)
}
pub fn unsubscribe(&self, uri: &str) -> Result<(), McpError> {
self.unsubscribe_within(uri, self.timeout)
}
pub fn unsubscribe_within(&self, uri: &str, _timeout: Duration) -> Result<(), McpError> {
let Some(c) = &self.rmcp else {
return Err(McpError::Transport(
"the MCP connection is not established".into(),
));
};
c.unsubscribe(uri)
}
pub fn drain_notifications(&self) -> Vec<rpc::Notification> {
let Some(c) = &self.rmcp else {
return Vec::new();
};
c.drain_notifications()
}
fn request_as<T: serde::de::DeserializeOwned>(
&self,
method: &str,
params: Option<Value>,
) -> Result<T, McpError> {
self.request_as_within(method, params, self.timeout)
}
fn request_as_within<T: serde::de::DeserializeOwned>(
&self,
method: &str,
params: Option<Value>,
timeout: Duration,
) -> Result<T, McpError> {
let v = self.request_with_timeout(method, params, timeout)?;
serde_json::from_value(v)
.map_err(|e| McpError::Transport(format!("bad {method} result: {e}")))
}
fn request_with_timeout(
&self,
method: &str,
params: Option<Value>,
timeout: Duration,
) -> Result<Value, McpError> {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (params, routing) = if self.era == Era::Modern {
let mut p = params.unwrap_or_else(|| json!({}));
let version = self
.protocol_version
.as_deref()
.unwrap_or(LATEST_MODERN_VERSION);
modern::inject_client_meta(
&mut p,
version,
&self.client_info,
&self.declared_capabilities(),
);
let mut routing: Vec<(String, String)> = modern::routing_headers(method, &p)
.into_iter()
.map(|(k, v)| (k.to_string(), v))
.collect();
if method == method::TOOLS_CALL
&& let Some(name) = p.get("name").and_then(Value::as_str)
{
let schema = self
.tool_schemas
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(name)
.cloned();
if let Some(schema) = schema {
let args = p.get("arguments").cloned().unwrap_or_else(|| json!({}));
routing.extend(modern::param_headers(&schema, &args));
}
}
(Some(p), routing)
} else {
(params, Vec::new())
};
let refs: Vec<(&str, &str)> = routing
.iter()
.map(|(k, v)| (k.as_str(), v.as_str()))
.collect();
let req = rpc::Request::new(id, method, params);
let body = serde_json::to_vec(&req)
.map_err(|e| McpError::Transport(format!("encode {method}: {e}")))?;
let router = self.router();
let msg = self
.http
.send(Some(id), &body, timeout, &refs, |n| router.route(n))
.map_err(|e| http_err(&self.name, method, e))?
.ok_or_else(|| {
McpError::Transport(format!("no response to {method} on '{}'", self.name))
})?;
let resp: rpc::Response = serde_json::from_value(msg).map_err(|e| {
McpError::Transport(format!("bad {method} response on '{}': {e}", self.name))
})?;
match resp.error {
Some(err) => Err(McpError::Rpc(err)),
None => Ok(resp.result.unwrap_or(Value::Null)),
}
}
}
impl Drop for McpClient {
fn drop(&mut self) {
if let Some(ev) = self
.events
.get_mut()
.unwrap_or_else(|e| e.into_inner())
.take()
{
ev.stop.store(true, Ordering::SeqCst);
let _ = ev.handle.join();
}
}
}
fn http_err(name: &str, method: &str, e: HttpError) -> McpError {
use std::io::ErrorKind;
match e {
HttpError::Connect(io) | HttpError::Http(io) => match io.kind() {
ErrorKind::TimedOut | ErrorKind::WouldBlock => {
McpError::Timeout(format!("{method} on '{name}'"))
}
_ => McpError::Transport(format!("{method} on '{name}': {io}")),
},
HttpError::Status(code, _) => {
McpError::Transport(format!("{method} on '{name}': server returned HTTP {code}"))
}
HttpError::Unsupported(m) => McpError::Transport(m),
HttpError::NoResponse => {
McpError::Transport(format!("{method} on '{name}': no JSON-RPC response"))
}
}
}
fn queue_notification(queue: &Mutex<VecDeque<rpc::Notification>>, n: Value) {
if let Ok(note) = serde_json::from_value::<rpc::Notification>(n) {
queue
.lock()
.unwrap_or_else(|e| e.into_inner())
.push_back(note);
}
}
fn to_value<T: Serialize>(v: &T) -> Value {
serde_json::to_value(v).unwrap_or(Value::Null)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Read, Write};
use std::os::unix::net::{UnixListener, UnixStream};
#[test]
fn error_display() {
let e = McpError::Timeout("tools/call on 'fs'".into());
assert!(e.to_string().contains("timeout"));
}
#[test]
fn http_err_folds_socket_timeout_into_timeout_variant() {
use std::io::{Error, ErrorKind};
let e = http_err(
"fs",
"tools/call",
HttpError::Http(Error::new(ErrorKind::WouldBlock, "read timed out")),
);
assert!(matches!(e, McpError::Timeout(_)), "got {e:?}");
let e = http_err("fs", "initialize", HttpError::Status(503, Vec::new()));
assert!(matches!(e, McpError::Transport(_)), "got {e:?}");
}
#[test]
fn queue_notification_enqueues_notifications_and_drops_others() {
let q = Mutex::new(VecDeque::new());
queue_notification(
&q,
json!({"jsonrpc":"2.0","method":"notifications/resources/updated","params":{"uri":"x"}}),
);
queue_notification(&q, json!({"jsonrpc":"2.0","id":1,"result":{}}));
let drained: Vec<_> = q.lock().unwrap().drain(..).collect();
assert_eq!(drained.len(), 1);
assert_eq!(drained[0].method, "notifications/resources/updated");
}
#[test]
fn connect_rejects_a_bad_endpoint() {
match McpClient::connect("bad", "ftp://nope/", Vec::new(), Duration::from_secs(1)) {
Err(McpError::Transport(_)) => {}
Err(other) => panic!("expected a Transport error, got {other:?}"),
Ok(_) => panic!("expected connect to reject an unsupported scheme"),
}
}
fn spawn_silent_server() -> (String, std::thread::JoinHandle<()>) {
let path = std::env::temp_dir().join(format!(
"agentd-mcp-silent-{}-{}.sock",
std::process::id(),
line!()
));
let _ = std::fs::remove_file(&path);
let listener = UnixListener::bind(&path).expect("bind silent server");
let handle = std::thread::spawn(move || {
for conn in listener.incoming() {
let Ok(mut stream) = conn else { continue };
std::thread::spawn(move || {
let mut buf = [0u8; 256];
while let Ok(n) = stream.read(&mut buf) {
if n == 0 {
break;
}
}
});
}
});
(format!("unix:{}", path.display()), handle)
}
#[test]
fn management_timeout_bounds_a_call_on_a_silent_server() {
let (endpoint, _srv) = spawn_silent_server();
let client = McpClient::connect("silent", &endpoint, Vec::new(), Duration::from_secs(60))
.expect("connect");
let short = Duration::from_millis(300);
let started = std::time::Instant::now();
let r = client.request_with_timeout("ping", None, short);
let elapsed = started.elapsed();
assert!(
matches!(r, Err(McpError::Timeout(_))),
"expected a Timeout within the short bound, got {r:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"the short per-call timeout must govern (took {elapsed:?})"
);
}
#[test]
fn write_read_smoke_for_unix_stream() {
let path = std::env::temp_dir().join(format!("agentd-smoke-{}.sock", std::process::id()));
let _ = std::fs::remove_file(&path);
let listener = UnixListener::bind(&path).unwrap();
let p2 = path.clone();
let h = std::thread::spawn(move || {
let (mut s, _) = listener.accept().unwrap();
let _ = s.write_all(b"hi");
});
let mut c = UnixStream::connect(&p2).unwrap();
let mut buf = [0u8; 2];
c.read_exact(&mut buf).unwrap();
assert_eq!(&buf, b"hi");
h.join().unwrap();
let _ = std::fs::remove_file(&path);
}
}