use std::collections::HashMap;
use std::io::Write as _;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{Mutex, mpsc, oneshot};
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeRequest {
pub protocol_version: String,
#[serde(default)]
pub capabilities: ClientCapabilities,
pub client_info: Implementation,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeResult {
pub protocol_version: String,
pub capabilities: ServerCapabilities,
pub server_info: Implementation,
#[serde(default)]
pub instructions: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
pub struct ClientCapabilities {
#[serde(skip_serializing_if = "Option::is_none")]
pub roots: Option<RootsCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling: Option<SamplingCapability>,
#[serde(skip_serializing_if = "std::collections::BTreeMap::is_empty", default)]
pub experimental: std::collections::BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RootsCapability {
#[serde(default)]
pub list_changed: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
pub struct SamplingCapability {}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
pub struct ServerCapabilities {
#[serde(default)]
pub prompts: Option<PromptsCapability>,
#[serde(default)]
pub resources: Option<ResourcesCapability>,
#[serde(default)]
pub tools: Option<ToolsCapability>,
#[serde(default)]
pub logging: Option<LoggingCapability>,
#[serde(default)]
pub experimental: std::collections::BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PromptsCapability {
#[serde(default)]
pub list_changed: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ResourcesCapability {
#[serde(default)]
pub list_changed: bool,
#[serde(default)]
pub subscribe: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolsCapability {
#[serde(default)]
pub list_changed: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize)]
pub struct LoggingCapability {}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Implementation {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
pub version: String,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListToolsResult {
pub tools: Vec<Tool>,
#[serde(default)]
pub next_cursor: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Tool {
pub name: String,
#[serde(default)]
pub description: Option<String>,
pub input_schema: Value,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CallToolRequest {
pub name: String,
#[serde(default)]
pub arguments: Value,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CallToolResult {
pub content: Vec<ContentBlock>,
#[serde(default)]
pub is_error: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
Text {
text: String,
},
Image {
#[serde(rename = "mimeType")]
mime_type: String,
data: String,
},
Audio {
#[serde(rename = "mimeType")]
mime_type: String,
data: String,
},
Resource {
uri: String,
#[serde(rename = "mimeType", default)]
mime_type: Option<String>,
#[serde(default)]
text: Option<String>,
},
}
struct Pending {
tx: oneshot::Sender<Result<Value, McpPeerError>>,
}
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
pub struct McpStdioPeer {
pending: Arc<Mutex<HashMap<u64, Pending>>>,
next_id: AtomicU64,
write_tx: mpsc::Sender<Vec<u8>>,
timeout: Duration,
#[expect(dead_code, reason = "kill_on_drop machinery; never read directly")]
child: Child,
}
impl std::fmt::Debug for McpStdioPeer {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("McpStdioPeer")
.field("pending_len", &"*")
.field("next_id", &self.next_id)
.field("write_tx", &"*")
.field("timeout", &self.timeout)
.field("child", &"*")
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone)]
pub struct SpawnSpec {
pub command: String,
pub args: Vec<String>,
pub timeout: Option<Duration>,
}
impl SpawnSpec {
pub fn new(command: impl Into<String>, args: Vec<String>) -> Self {
Self {
command: command.into(),
args,
timeout: None,
}
}
#[must_use]
pub const fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
}
#[derive(Debug, thiserror::Error)]
pub enum McpPeerError {
#[error("spawning the child process: {0}")]
Spawn(#[source] std::io::Error),
#[error("encoding or decoding the JSON-RPC envelope: {0}")]
Decode(#[source] serde_json::Error),
#[error("the reply did not arrive within {0:?}")]
Timeout(Duration),
#[error("the writer task is gone")]
Write(#[source] std::io::Error),
#[error("the reader task is gone")]
ReadGone,
#[error("the server returned a JSON-RPC error: code={code} message={message}")]
JsonRpc {
code: i64,
message: String,
},
}
impl McpStdioPeer {
pub fn spawn(spec: &SpawnSpec) -> Result<Self, McpPeerError> {
let mut command = Command::new(&spec.command);
command
.args(&spec.args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.kill_on_drop(true);
let mut child = command.spawn().map_err(McpPeerError::Spawn)?;
let stdin = child.stdin.take().ok_or_else(|| {
McpPeerError::Spawn(std::io::Error::other("child stdin was not piped"))
})?;
let stdout = child.stdout.take().ok_or_else(|| {
McpPeerError::Spawn(std::io::Error::other("child stdout was not piped"))
})?;
let pending: Arc<Mutex<HashMap<u64, Pending>>> = Arc::new(Mutex::new(HashMap::new()));
let (write_tx, write_rx) = mpsc::channel::<Vec<u8>>(32);
spawn_reader(BufReader::new(stdout), Arc::clone(&pending));
spawn_writer(stdin, write_rx);
Ok(Self {
pending,
next_id: AtomicU64::new(1),
write_tx,
timeout: spec.timeout.unwrap_or(DEFAULT_TIMEOUT),
child,
})
}
pub async fn send_typed_request<Req, Resp>(
&self,
method: &str,
params: Req,
) -> Result<Resp, McpPeerError>
where
Req: Serialize,
Resp: DeserializeOwned,
{
let value = self.send_request(method, params).await?;
serde_json::from_value(value).map_err(McpPeerError::Decode)
}
pub async fn send_request<Req>(&self, method: &str, params: Req) -> Result<Value, McpPeerError>
where
Req: Serialize,
{
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let params_value: Value = serde_json::to_value(¶ms).map_err(McpPeerError::Decode)?;
let envelope = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params_value,
});
let bytes = serde_json::to_vec(&envelope).map_err(McpPeerError::Decode)?;
let mut framed = bytes;
framed.push(b'\n');
self.write_tx
.send(framed)
.await
.map_err(|_| McpPeerError::Write(io_other("the writer task is gone")))?;
let (tx, rx) = oneshot::channel();
self.pending.lock().await.insert(id, Pending { tx });
match tokio::time::timeout(self.timeout, rx).await {
Ok(Ok(result)) => result,
Ok(Err(_oneshot_dropped)) => {
self.pending.lock().await.remove(&id);
Err(McpPeerError::Timeout(self.timeout))
}
Err(_elapsed) => {
self.pending.lock().await.remove(&id);
Err(McpPeerError::Timeout(self.timeout))
}
}
}
}
fn io_other(message: &str) -> std::io::Error {
std::io::Error::other(message)
}
fn spawn_writer(mut stdin: ChildStdin, mut write_rx: mpsc::Receiver<Vec<u8>>) {
tokio::spawn(async move {
while let Some(mut bytes) = write_rx.recv().await {
if !bytes.ends_with(b"\n") {
bytes.push(b'\n');
}
if let Err(error) = stdin.write_all(&bytes).await {
let _ = std::io::stderr()
.lock()
.write_all(format!("mcp writer task: write failed: {error}\n").as_bytes());
return;
}
if let Err(error) = stdin.flush().await {
let _ = std::io::stderr()
.lock()
.write_all(format!("mcp writer task: flush failed: {error}\n").as_bytes());
return;
}
}
});
}
fn spawn_reader(mut reader: BufReader<ChildStdout>, pending: Arc<Mutex<HashMap<u64, Pending>>>) {
tokio::spawn(async move {
let mut line = String::new();
loop {
line.clear();
match reader.read_line(&mut line).await {
Ok(0) | Err(_) => return,
Ok(_) => {}
}
let trimmed = line.trim_end_matches('\n');
if trimmed.is_empty() {
continue;
}
let value: Value = match serde_json::from_str(trimmed) {
Ok(value) => value,
Err(_) => continue,
};
let id = value.get("id").and_then(Value::as_u64).unwrap_or(0);
if id == 0 {
continue;
}
let result = value.get("error").map_or_else(
|| {
value
.get("result")
.cloned()
.ok_or_else(|| McpPeerError::JsonRpc {
code: -32600,
message: "the reply had no `result` and no `error`".to_owned(),
})
},
|error_value| {
let code = error_value.get("code").and_then(Value::as_i64).unwrap_or(0);
let message = error_value
.get("message")
.and_then(Value::as_str)
.unwrap_or("")
.to_owned();
Err(McpPeerError::JsonRpc { code, message })
},
);
let entry = pending.lock().await.remove(&id);
if let Some(entry) = entry {
let _ = entry.tx.send(result);
}
}
});
}
#[must_use]
pub fn default_client_capabilities() -> ClientCapabilities {
ClientCapabilities::default()
}
#[must_use]
pub fn initialize_request_template() -> InitializeRequest {
InitializeRequest {
protocol_version: "2025-06-18".to_owned(),
capabilities: ClientCapabilities::default(),
client_info: Implementation {
name: "protocol-droid".to_owned(),
title: Some("protocol-droid MCP client".to_owned()),
version: "0.1.0".to_owned(),
},
}
}
#[cfg(test)]
mod typed_shape_round_trip {
use super::*;
#[test]
fn client_capabilities_default_serialises() -> Result<(), serde_json::Error> {
let caps = ClientCapabilities::default();
let value = serde_json::to_value(&caps)?;
if let Some(object) = value.as_object() {
assert!(
!object.contains_key("roots"),
"the default form omits `roots`"
);
} else {
assert!(
value.is_object(),
"the serialised value should be an object"
);
}
Ok(())
}
#[test]
fn client_capabilities_omits_roots_when_none() -> Result<(), serde_json::Error> {
let caps = ClientCapabilities {
roots: None,
sampling: None,
experimental: std::collections::BTreeMap::new(),
};
let value = serde_json::to_value(&caps)?;
if let Some(object) = value.as_object() {
assert!(!object.contains_key("roots"));
} else {
assert!(
value.is_object(),
"the serialised value should be an object"
);
}
Ok(())
}
#[test]
fn implementation_round_trip() -> Result<(), serde_json::Error> {
let impl_value = Implementation {
name: "test-client".to_owned(),
title: None,
version: "0.0.1".to_owned(),
};
let value = serde_json::to_value(&impl_value)?;
if let Some(object) = value.as_object() {
assert!(!object.contains_key("title"));
} else {
assert!(
value.is_object(),
"the serialised value should be an object"
);
}
let decoded: Implementation = serde_json::from_value(value)?;
assert_eq!(impl_value, decoded);
Ok(())
}
#[test]
fn call_tool_request_serialises_name() -> Result<(), serde_json::Error> {
let request = CallToolRequest {
name: "echo".to_owned(),
arguments: serde_json::json!({}),
};
let value = serde_json::to_value(&request)?;
assert_eq!(
value.get("name").and_then(serde_json::Value::as_str),
Some("echo"),
);
Ok(())
}
#[test]
fn initialize_request_uses_camel_case_on_the_wire() -> Result<(), serde_json::Error> {
let request = InitializeRequest {
protocol_version: "2025-06-18".to_owned(),
capabilities: ClientCapabilities::default(),
client_info: Implementation {
name: "test-client".to_owned(),
title: None,
version: "0.0.1".to_owned(),
},
};
let value = serde_json::to_value(&request)?;
if let Some(object) = value.as_object() {
assert!(object.contains_key("protocolVersion"));
assert!(!object.contains_key("protocol_version"));
} else {
assert!(
value.is_object(),
"the serialised value should be an object"
);
}
Ok(())
}
}