use anyhow::{Context, Result, anyhow};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::io::{AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tokio::sync::{Mutex, oneshot};
use tokio::time::{Duration, timeout};
use crate::constants::MAX_MCP_FRAME_BYTES;
use crate::utils::{CappedLine, read_line_capped};
const REQUEST_TIMEOUT_SECS: u64 = 30;
const TOOL_CALL_TIMEOUT_SECS: u64 = 300;
const WRITE_TIMEOUT_SECS: u64 = 10;
#[cfg(not(windows))]
const SAFE_ENV_PASSTHROUGH: &[&str] = &[
"PATH", "HOME", "USER", "LOGNAME", "SHELL", "LANG", "TERM", "TMPDIR",
];
#[cfg(windows)]
const SAFE_ENV_PASSTHROUGH: &[&str] = &[
"PATH",
"HOME",
"USER",
"LOGNAME",
"SHELL",
"LANG",
"TERM",
"TMPDIR",
"SYSTEMROOT",
"SystemRoot",
"APPDATA",
"USERPROFILE",
"PATHEXT",
"COMSPEC",
];
pub struct StdioTransport {
stdin: Arc<Mutex<tokio::process::ChildStdin>>,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
next_id: AtomicU64,
child: Arc<Mutex<Child>>,
_reader_task: tokio::task::JoinHandle<()>,
}
impl StdioTransport {
pub async fn spawn(
command: &str,
args: &[String],
env: &HashMap<String, String>,
) -> Result<Self> {
let mut cmd = Command::new(command);
cmd.args(args)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
cmd.env_clear();
for key in SAFE_ENV_PASSTHROUGH {
if let Some(value) = std::env::var_os(key) {
cmd.env(key, value);
}
}
for (key, value) in std::env::vars_os() {
if key.to_str().is_some_and(|k| k.starts_with("LC_")) {
cmd.env(&key, &value);
}
}
for (key, value) in env {
cmd.env(key, value);
}
let mut child = cmd.spawn().with_context(|| {
format!(
"Failed to spawn MCP server: {} {}",
command,
crate::utils::redact_secrets(&args.join(" "))
)
})?;
let stdin = child
.stdin
.take()
.ok_or_else(|| anyhow!("Failed to capture MCP server stdin"))?;
let stdin = Arc::new(Mutex::new(stdin));
let stdout = child
.stdout
.take()
.ok_or_else(|| anyhow!("Failed to capture MCP server stdout"))?;
let stderr = child
.stderr
.take()
.ok_or_else(|| anyhow!("Failed to capture MCP server stderr"))?;
let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>> =
Arc::new(Mutex::new(HashMap::new()));
let pending_clone = Arc::clone(&pending);
let stdin_for_reader = Arc::clone(&stdin);
let reader_task = tokio::spawn(async move {
let mut reader = BufReader::new(stdout);
loop {
let line = match read_line_capped(&mut reader, MAX_MCP_FRAME_BYTES).await {
Ok(CappedLine::Line(bytes)) => match String::from_utf8(bytes) {
Ok(s) => s,
Err(_) => {
tracing::warn!("MCP: dropping non-UTF-8 stdout frame");
continue;
},
},
Ok(CappedLine::TooLong) => {
tracing::warn!("MCP: dropping oversize stdout frame (exceeds cap)");
continue;
},
Ok(CappedLine::Eof) | Err(_) => break,
};
if line.trim().is_empty() {
continue;
}
let Ok(msg) = serde_json::from_str::<Value>(&line) else {
let end = line.floor_char_boundary(200);
tracing::warn!("MCP: unparseable stdout line: {}", &line[..end]);
continue;
};
if is_response(&msg) {
if let Some(id) = msg.get("id").and_then(parse_response_id) {
let mut pending = pending_clone.lock().await;
if let Some(sender) = pending.remove(&id) {
let _ = sender.send(msg);
}
}
} else if msg.get("method").is_some() {
match msg.get("id") {
Some(id) if !id.is_null() => {
tracing::debug!(
method = msg.get("method").and_then(|m| m.as_str()).unwrap_or(""),
"MCP: replying method-not-supported to unsupported server request"
);
reply_method_not_supported(&stdin_for_reader, id).await;
},
_ => {
tracing::trace!(
method = msg.get("method").and_then(|m| m.as_str()).unwrap_or(""),
"MCP: ignoring server notification (no reply expected)"
);
},
}
} else {
tracing::trace!("MCP: ignoring unrecognized stdout message");
}
}
pending_clone.lock().await.clear();
});
tokio::spawn(async move {
let mut reader = BufReader::new(stderr);
loop {
match read_line_capped(&mut reader, MAX_MCP_FRAME_BYTES).await {
Ok(CappedLine::Line(bytes)) => {
let line = String::from_utf8_lossy(&bytes);
tracing::debug!("MCP stderr: {}", crate::utils::redact_secrets(&line));
},
Ok(CappedLine::TooLong) => continue,
Ok(CappedLine::Eof) | Err(_) => break,
}
}
});
Ok(Self {
stdin,
pending,
next_id: AtomicU64::new(1),
child: Arc::new(Mutex::new(child)),
_reader_task: reader_task,
})
}
pub async fn send_request(&self, method: &str, params: Value) -> Result<Value> {
self.send_request_with_timeout(method, params, REQUEST_TIMEOUT_SECS)
.await
}
pub fn tool_call_timeout_secs() -> u64 {
TOOL_CALL_TIMEOUT_SECS
}
pub async fn send_request_with_timeout(
&self,
method: &str,
params: Value,
response_timeout_secs: u64,
) -> Result<Value> {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params,
});
let msg = format!("{}\n", serde_json::to_string(&request)?);
let (tx, rx) = oneshot::channel();
{
let mut pending = self.pending.lock().await;
pending.insert(id, tx);
}
{
let mut stdin = self.stdin.lock().await;
match timeout(
Duration::from_secs(WRITE_TIMEOUT_SECS),
stdin.write_all(msg.as_bytes()),
)
.await
{
Ok(Ok(())) => {},
Ok(Err(e)) => {
drop(stdin);
self.pending.lock().await.remove(&id);
return Err(e).with_context(|| {
format!("Failed to write to MCP server stdin (method: {})", method)
});
},
Err(_) => {
drop(stdin);
self.pending.lock().await.remove(&id);
return Err(anyhow!(
"Timed out after {}s writing to MCP server stdin (method: {})",
WRITE_TIMEOUT_SECS,
method
));
},
}
match timeout(Duration::from_secs(WRITE_TIMEOUT_SECS), stdin.flush()).await {
Ok(Ok(())) => {},
Ok(Err(e)) => {
drop(stdin);
self.pending.lock().await.remove(&id);
return Err(e).with_context(|| {
format!("Failed to flush MCP server stdin (method: {})", method)
});
},
Err(_) => {
drop(stdin);
self.pending.lock().await.remove(&id);
return Err(anyhow!(
"Timed out after {}s flushing MCP server stdin (method: {})",
WRITE_TIMEOUT_SECS,
method
));
},
}
}
let response = match timeout(Duration::from_secs(response_timeout_secs), rx).await {
Ok(Ok(value)) => value,
Ok(Err(_)) => {
self.pending.lock().await.remove(&id);
return Err(anyhow!("MCP response channel closed unexpectedly"));
},
Err(_) => {
self.pending.lock().await.remove(&id);
return Err(anyhow!(
"MCP request timed out after {}s: {}",
response_timeout_secs,
method
));
},
};
if let Some(error) = response.get("error") {
let code = error.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
let message = error
.get("message")
.and_then(|m| m.as_str())
.unwrap_or("Unknown error");
return Err(anyhow!("MCP error (code {}): {}", code, message));
}
response
.get("result")
.cloned()
.ok_or_else(|| anyhow!("MCP response missing 'result' field"))
}
pub async fn send_notification(&self, method: &str, params: Value) -> Result<()> {
let notification = serde_json::json!({
"jsonrpc": "2.0",
"method": method,
"params": params,
});
let msg = format!("{}\n", serde_json::to_string(¬ification)?);
let mut stdin = self.stdin.lock().await;
timeout(
Duration::from_secs(WRITE_TIMEOUT_SECS),
stdin.write_all(msg.as_bytes()),
)
.await
.map_err(|_| anyhow!("Timed out writing MCP notification (method: {})", method))??;
timeout(Duration::from_secs(WRITE_TIMEOUT_SECS), stdin.flush())
.await
.map_err(|_| anyhow!("Timed out flushing MCP notification (method: {})", method))??;
Ok(())
}
pub async fn shutdown(&self) {
{
let mut stdin = self.stdin.lock().await;
let _ = stdin.shutdown().await;
}
let mut child = self.child.lock().await;
if tokio::time::timeout(Duration::from_secs(2), child.wait())
.await
.is_ok()
{
return;
}
#[cfg(unix)]
if let Some(pid) = child.id() {
terminate(pid).await;
if tokio::time::timeout(Duration::from_secs(2), child.wait())
.await
.is_ok()
{
return;
}
}
if let Err(e) = child.start_kill() {
tracing::debug!("MCP: start_kill failed: {}", e);
}
if tokio::time::timeout(Duration::from_secs(1), child.wait())
.await
.is_ok()
{
return;
}
let _ = child.kill().await;
let _ = child.wait().await;
}
}
fn parse_response_id(v: &Value) -> Option<u64> {
v.as_u64()
.or_else(|| v.as_str().and_then(|s| s.parse::<u64>().ok()))
}
fn is_response(msg: &Value) -> bool {
msg.get("id").and_then(parse_response_id).is_some() && msg.get("method").is_none()
}
async fn reply_method_not_supported(stdin: &Mutex<tokio::process::ChildStdin>, id: &Value) {
let reply = serde_json::json!({
"jsonrpc": "2.0",
"id": id.clone(),
"error": { "code": -32601, "message": "method not supported" },
});
let Ok(serialized) = serde_json::to_string(&reply) else {
return;
};
let mut line = serialized;
line.push('\n');
let mut guard = stdin.lock().await;
let _ = timeout(
Duration::from_secs(WRITE_TIMEOUT_SECS),
guard.write_all(line.as_bytes()),
)
.await;
let _ = timeout(Duration::from_secs(WRITE_TIMEOUT_SECS), guard.flush()).await;
}
#[cfg(unix)]
async fn terminate(pid: u32) {
let _ = Command::new("kill")
.arg(pid.to_string())
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
.await;
}
#[cfg(test)]
mod tests {
use super::{is_response, parse_response_id};
use serde_json::json;
#[test]
fn truncation_respects_char_boundary() {
let line = format!("{}你好", "a".repeat(199));
let end = line.floor_char_boundary(200);
let truncated = &line[..end]; assert!(end <= 200);
assert!(truncated.is_char_boundary(end));
}
#[test]
fn parse_response_id_accepts_integer() {
assert_eq!(parse_response_id(&json!(5)), Some(5));
assert_eq!(parse_response_id(&json!(0)), Some(0));
}
#[test]
fn parse_response_id_accepts_string_integer() {
assert_eq!(parse_response_id(&json!("5")), Some(5));
assert_eq!(parse_response_id(&json!("0")), Some(0));
}
#[test]
fn parse_response_id_rejects_non_numeric() {
assert_eq!(parse_response_id(&json!("abc")), None);
assert_eq!(parse_response_id(&json!(null)), None);
assert_eq!(parse_response_id(&json!({})), None);
assert_eq!(parse_response_id(&json!(-1)), None);
}
#[test]
fn is_response_true_for_result_or_error() {
assert!(is_response(
&json!({"jsonrpc": "2.0", "id": 1, "result": {}})
));
assert!(is_response(
&json!({"jsonrpc": "2.0", "id": "7", "error": {"code": -1, "message": "x"}})
));
}
#[test]
fn is_response_false_for_server_request() {
assert!(!is_response(
&json!({"jsonrpc": "2.0", "id": 1, "method": "sampling/createMessage", "params": {}})
));
}
#[test]
fn is_response_false_for_notification() {
assert!(!is_response(
&json!({"jsonrpc": "2.0", "method": "notifications/message", "params": {}})
));
}
#[test]
fn is_response_false_without_id() {
assert!(!is_response(&json!({"jsonrpc": "2.0", "result": {}})));
}
#[test]
fn invalid_utf8_frame_is_not_valid_json() {
let bad = [0xff, 0xfe, b'{', b'}'];
assert!(String::from_utf8(bad.to_vec()).is_err());
assert!(serde_json::from_slice::<serde_json::Value>(&bad).is_err());
}
#[cfg(unix)]
#[tokio::test]
async fn pending_request_fails_fast_when_server_closes_stdout() {
use super::StdioTransport;
use std::collections::HashMap;
let t = StdioTransport::spawn(
"sh",
&["-c".to_string(), "read x; exec 1>&-; sleep 2".to_string()],
&HashMap::new(),
)
.await
.expect("spawn");
let start = std::time::Instant::now();
let res = t.send_request("ping", json!({})).await;
assert!(
res.is_err(),
"request must fail once the server closes stdout"
);
assert!(
start.elapsed() < std::time::Duration::from_secs(5),
"must fail fast via the EOF drain, not wait REQUEST_TIMEOUT_SECS"
);
}
#[cfg(unix)]
#[tokio::test]
async fn child_env_is_cleared_except_allowlist_and_declared() {
use super::StdioTransport;
use std::collections::HashMap;
let script = r#"read line
printf '{"jsonrpc":"2.0","id":1,"result":{"declared":"%s","path":"%s","cargo":"%s"}}\n' "$DECLARED_TOKEN" "$([ -n "$PATH" ] && echo yes || echo no)" "$CARGO""#;
let mut env = HashMap::new();
env.insert("DECLARED_TOKEN".to_string(), "passed-through".to_string());
let t = StdioTransport::spawn("sh", &["-c".to_string(), script.to_string()], &env)
.await
.expect("spawn");
let res = t.send_request("ping", json!({})).await.expect("response");
assert_eq!(
res["declared"], "passed-through",
"declared env vars must still pass through"
);
assert_eq!(res["path"], "yes", "PATH must be in the allowlist");
assert_eq!(
res["cargo"], "",
"a non-allowlisted inherited var must be cleared by env_clear()"
);
}
#[cfg(unix)]
#[tokio::test]
async fn server_initiated_request_gets_method_not_supported_reply() {
use super::StdioTransport;
use std::collections::HashMap;
let script = r#"read ping
printf '{"jsonrpc":"2.0","id":99,"method":"sampling/createMessage","params":{}}\n'
read reply
case "$reply" in
*-32601*) printf '{"jsonrpc":"2.0","id":1,"result":{"replied":true}}\n' ;;
*) printf '{"jsonrpc":"2.0","id":1,"result":{"replied":false}}\n' ;;
esac"#;
let t = StdioTransport::spawn(
"sh",
&["-c".to_string(), script.to_string()],
&HashMap::new(),
)
.await
.expect("spawn");
let res = t
.send_request("ping", json!({}))
.await
.expect("ping response");
assert_eq!(
res["replied"], true,
"transport must reply with a -32601 error to a server-initiated request, \
not silently drop it"
);
}
#[cfg(unix)]
#[tokio::test]
async fn terminate_sends_real_sigterm() {
use std::os::unix::process::ExitStatusExt;
use tokio::process::Command;
let mut child = Command::new("sh")
.arg("-c")
.arg("sleep 10")
.spawn()
.expect("spawn sh");
let pid = child.id().expect("pid");
super::terminate(pid).await;
let status = tokio::time::timeout(std::time::Duration::from_secs(5), child.wait())
.await
.expect("child should die from the signal")
.expect("wait");
assert_eq!(
status.signal(),
Some(15),
"terminate() must send SIGTERM, not SIGKILL; got {status:?}"
);
}
}