1use std::collections::HashMap;
2use std::process::Stdio;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::Duration;
6
7use serde::{Deserialize, Serialize};
8use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
9use tokio::process::{Child, ChildStdin, ChildStdout, Command};
10use tokio::sync::{Mutex, oneshot};
11use tokio::task::JoinHandle;
12
13use crate::error::RuntimeError;
14use crate::tool::BoxFut;
15
16type PendingCalls = Arc<Mutex<HashMap<u64, oneshot::Sender<Result<serde_json::Value, McpError>>>>>;
17
18pub trait McpTransport: Send + Sync {
19 fn call<'a>(
20 &'a self,
21 method: &'a str,
22 params: serde_json::Value,
23 ) -> BoxFut<'a, Result<serde_json::Value, McpError>>;
24
25 fn kind(&self) -> &'static str;
26}
27
28#[derive(Debug, Serialize, Deserialize, Clone)]
29pub struct JsonRpcRequest {
30 pub jsonrpc: &'static str,
31 pub id: u64,
32 pub method: String,
33 #[serde(skip_serializing_if = "Option::is_none")]
34 pub params: Option<serde_json::Value>,
35}
36
37#[derive(Debug, Serialize, Deserialize, Clone)]
38pub struct JsonRpcResponse {
39 #[allow(dead_code)]
40 pub jsonrpc: String,
41 pub id: Option<u64>,
42 #[serde(default)]
43 pub result: Option<serde_json::Value>,
44 #[serde(default)]
45 pub error: Option<JsonRpcError>,
46}
47
48#[derive(Debug, Serialize, Deserialize, Clone)]
49pub struct JsonRpcError {
50 pub code: i64,
51 pub message: String,
52 #[serde(default)]
53 pub data: Option<serde_json::Value>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct McpToolSchema {
58 pub name: String,
59 #[serde(default)]
60 pub description: Option<String>,
61 #[serde(default, rename = "inputSchema")]
62 pub input_schema: Option<serde_json::Value>,
63}
64
65#[derive(Debug, Clone, thiserror::Error)]
66pub enum McpError {
67 #[error("mcp io: {0}")]
68 Io(String),
69 #[error("mcp protocol: {0}")]
70 Protocol(String),
71 #[error("mcp server error {code}: {message}")]
72 ServerError { code: i64, message: String },
73 #[error("mcp timeout ({timeout_ms}ms) on {method}")]
74 Timeout { timeout_ms: u64, method: String },
75 #[error("mcp disconnected")]
76 Disconnected,
77}
78
79impl From<McpError> for RuntimeError {
80 fn from(e: McpError) -> Self {
81 RuntimeError::ToolFailed(format!("{e}"))
82 }
83}
84
85pub struct McpStdioTransport {
86 stdin: Arc<Mutex<ChildStdin>>,
87 pending: PendingCalls,
88 next_id: Arc<Mutex<u64>>,
89 #[allow(dead_code)]
90 child: Arc<Mutex<Child>>,
91 #[allow(dead_code)]
92 reader_task: Arc<Mutex<Option<JoinHandle<()>>>>,
93 timeout_ms: u64,
94}
95
96impl McpTransport for McpStdioTransport {
97 fn call<'a>(
98 &'a self,
99 method: &'a str,
100 params: serde_json::Value,
101 ) -> BoxFut<'a, Result<serde_json::Value, McpError>> {
102 Box::pin(self.call_stdio(method, params))
103 }
104
105 fn kind(&self) -> &'static str {
106 "stdio"
107 }
108}
109
110impl McpStdioTransport {
111 pub async fn spawn(cmd: &str, args: &[String], timeout_ms: u64) -> Result<Self, McpError> {
112 let mut child = Command::new(cmd)
113 .args(args)
114 .stdin(Stdio::piped())
115 .stdout(Stdio::piped())
116 .stderr(Stdio::piped())
117 .spawn()
118 .map_err(|e| McpError::Io(format!("spawn {cmd}: {e}")))?;
119 let stdin = child
120 .stdin
121 .take()
122 .ok_or_else(|| McpError::Io("no stdin".into()))?;
123 let stdout = child
124 .stdout
125 .take()
126 .ok_or_else(|| McpError::Io("no stdout".into()))?;
127
128 let pending: PendingCalls = Arc::new(Mutex::new(HashMap::new()));
129 let pending_for_reader = pending.clone();
130 let reader_task = tokio::spawn(async move {
131 reader_loop(stdout, pending_for_reader).await;
132 });
133
134 Ok(Self {
135 stdin: Arc::new(Mutex::new(stdin)),
136 pending,
137 next_id: Arc::new(Mutex::new(1)),
138 child: Arc::new(Mutex::new(child)),
139 reader_task: Arc::new(Mutex::new(Some(reader_task))),
140 timeout_ms,
141 })
142 }
143
144 async fn call_stdio(
145 &self,
146 method: &str,
147 params: serde_json::Value,
148 ) -> Result<serde_json::Value, McpError> {
149 let id = {
150 let mut n = self.next_id.lock().await;
151 let v = *n;
152 *n += 1;
153 v
154 };
155 let (tx, rx) = oneshot::channel();
156 self.pending.lock().await.insert(id, tx);
157
158 let req = JsonRpcRequest {
159 jsonrpc: "2.0",
160 id,
161 method: method.into(),
162 params: Some(params),
163 };
164 let line = serde_json::to_string(&req)
165 .map_err(|e| McpError::Protocol(format!("serialize: {e}")))?;
166 {
167 let mut stdin = self.stdin.lock().await;
168 stdin
169 .write_all(line.as_bytes())
170 .await
171 .map_err(|e| McpError::Io(format!("write: {e}")))?;
172 stdin
173 .write_all(b"\n")
174 .await
175 .map_err(|e| McpError::Io(format!("write newline: {e}")))?;
176 stdin
177 .flush()
178 .await
179 .map_err(|e| McpError::Io(format!("flush: {e}")))?;
180 }
181
182 match tokio::time::timeout(Duration::from_millis(self.timeout_ms), rx).await {
183 Ok(Ok(inner)) => inner,
184 Ok(Err(_)) => {
185 self.pending.lock().await.remove(&id);
186 Err(McpError::Disconnected)
187 }
188 Err(_) => {
189 self.pending.lock().await.remove(&id);
190 Err(McpError::Timeout {
191 timeout_ms: self.timeout_ms,
192 method: method.into(),
193 })
194 }
195 }
196 }
197}
198
199async fn reader_loop(stdout: ChildStdout, pending: PendingCalls) {
200 let reader = BufReader::new(stdout);
201 let mut lines = reader.lines();
202 loop {
203 match lines.next_line().await {
204 Ok(Some(line)) => {
205 if line.trim().is_empty() {
206 continue;
207 }
208 let parsed: JsonRpcResponse = match serde_json::from_str(&line) {
209 Ok(r) => r,
210 Err(_) => continue,
211 };
212 let Some(id) = parsed.id else {
213 continue;
214 };
215 let sender = pending.lock().await.remove(&id);
216 if let Some(sender) = sender {
217 let outcome = if let Some(err) = parsed.error {
218 Err(McpError::ServerError {
219 code: err.code,
220 message: err.message,
221 })
222 } else {
223 Ok(parsed.result.unwrap_or(serde_json::Value::Null))
224 };
225 let _ = sender.send(outcome);
226 }
227 }
228 Ok(None) => break,
229 Err(_) => break,
230 }
231 }
232 let mut pending = pending.lock().await;
233 for (_, tx) in pending.drain() {
234 let _ = tx.send(Err(McpError::Disconnected));
235 }
236}
237
238pub struct McpHttpTransport {
239 url: String,
240 auth_token: Option<String>,
241 client: reqwest::Client,
242 next_id: AtomicU64,
243 timeout_ms: u64,
244 retry_attempts: u32,
245}
246
247impl McpHttpTransport {
248 pub fn new(url: impl Into<String>, auth_token: Option<String>, timeout_ms: u64) -> Self {
249 Self {
250 url: url.into(),
251 auth_token,
252 client: reqwest::Client::new(),
253 next_id: AtomicU64::new(1),
254 timeout_ms,
255 retry_attempts: 3,
256 }
257 }
258
259 #[doc(hidden)]
260 pub fn with_client(mut self, client: reqwest::Client) -> Self {
261 self.client = client;
262 self
263 }
264
265 async fn call_http(
266 &self,
267 method: &str,
268 params: serde_json::Value,
269 ) -> Result<serde_json::Value, McpError> {
270 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
271 let body = serde_json::json!({
272 "jsonrpc": "2.0",
273 "id": id,
274 "method": method,
275 "params": params,
276 });
277 let call_timeout = Duration::from_millis(self.timeout_ms);
278 let attempt_result = tokio::time::timeout(call_timeout, async {
279 let mut delay_ms = 100u64;
280 let mut last_err: Option<McpError> = None;
281 for attempt in 0..=self.retry_attempts {
282 let mut req = self.client.post(&self.url).json(&body);
283 if let Some(t) = &self.auth_token {
284 req = req.bearer_auth(t);
285 }
286 match req.send().await {
287 Ok(resp) => {
288 let status = resp.status();
289 if status.is_success() {
290 let text = resp
291 .text()
292 .await
293 .map_err(|e| McpError::Io(format!("mcp http body: {e}")))?;
294 let parsed: JsonRpcResponse = serde_json::from_str(&text)
295 .map_err(|e| McpError::Protocol(format!("mcp http parse: {e}")))?;
296 if let Some(err) = parsed.error {
297 return Err(McpError::ServerError {
298 code: err.code,
299 message: err.message,
300 });
301 }
302 return Ok(parsed.result.unwrap_or(serde_json::Value::Null));
303 }
304 if status.is_server_error() {
305 last_err =
306 Some(McpError::Io(format!("mcp http {status}: server error")));
307 if attempt < self.retry_attempts {
308 tokio::time::sleep(Duration::from_millis(delay_ms)).await;
309 delay_ms = (delay_ms * 5).min(2000);
310 continue;
311 }
312 }
313 let body_text = resp.text().await.unwrap_or_default();
314 return Err(McpError::Io(format!("mcp http {status}: {body_text}")));
315 }
316 Err(e) => {
317 last_err = Some(McpError::Io(format!("mcp http send: {e}")));
318 if attempt < self.retry_attempts {
319 tokio::time::sleep(Duration::from_millis(delay_ms)).await;
320 delay_ms = (delay_ms * 5).min(2000);
321 continue;
322 }
323 }
324 }
325 }
326 Err(last_err.unwrap_or(McpError::Disconnected))
327 })
328 .await;
329 match attempt_result {
330 Ok(inner) => inner,
331 Err(_) => Err(McpError::Timeout {
332 timeout_ms: self.timeout_ms,
333 method: method.into(),
334 }),
335 }
336 }
337}
338
339impl McpTransport for McpHttpTransport {
340 fn call<'a>(
341 &'a self,
342 method: &'a str,
343 params: serde_json::Value,
344 ) -> BoxFut<'a, Result<serde_json::Value, McpError>> {
345 Box::pin(self.call_http(method, params))
346 }
347
348 fn kind(&self) -> &'static str {
349 "http"
350 }
351}
352
353pub struct McpClient {
354 pub name: String,
355 transport: Arc<dyn McpTransport>,
356 pub tools: Vec<McpToolSchema>,
357}
358
359impl McpClient {
360 pub async fn connect_stdio(
361 name: impl Into<String>,
362 cmd: &str,
363 args: &[String],
364 timeout_ms: u64,
365 ) -> Result<Self, McpError> {
366 let transport: Arc<dyn McpTransport> =
367 Arc::new(McpStdioTransport::spawn(cmd, args, timeout_ms).await?);
368 Self::finish_connect(name.into(), transport).await
369 }
370
371 pub async fn connect_http(
372 name: impl Into<String>,
373 url: impl Into<String>,
374 auth_token: Option<String>,
375 timeout_ms: u64,
376 ) -> Result<Self, McpError> {
377 let transport: Arc<dyn McpTransport> =
378 Arc::new(McpHttpTransport::new(url, auth_token, timeout_ms));
379 Self::finish_connect(name.into(), transport).await
380 }
381
382 pub async fn connect_with_transport(
383 name: impl Into<String>,
384 transport: Arc<dyn McpTransport>,
385 ) -> Result<Self, McpError> {
386 Self::finish_connect(name.into(), transport).await
387 }
388
389 async fn finish_connect(
390 name: String,
391 transport: Arc<dyn McpTransport>,
392 ) -> Result<Self, McpError> {
393 let init_params = serde_json::json!({
394 "protocolVersion": "2024-11-05",
395 "capabilities": {},
396 "clientInfo": {"name": "atman", "version": env!("CARGO_PKG_VERSION")}
397 });
398 transport.call("initialize", init_params).await?;
399 let _ = transport
400 .call("notifications/initialized", serde_json::Value::Null)
401 .await;
402 let list = transport.call("tools/list", serde_json::json!({})).await?;
403 let tools = parse_tools_list(&list)?;
404 Ok(Self {
405 name,
406 transport,
407 tools,
408 })
409 }
410
411 pub fn transport_kind(&self) -> &'static str {
412 self.transport.kind()
413 }
414
415 pub async fn call_tool(
416 &self,
417 tool_name: &str,
418 arguments: serde_json::Value,
419 ) -> Result<crate::value::Value, McpError> {
420 let params = serde_json::json!({
421 "name": tool_name,
422 "arguments": arguments,
423 });
424 let result = self.transport.call("tools/call", params).await?;
425 Ok(mcp_result_to_value(result))
426 }
427}
428
429fn parse_tools_list(v: &serde_json::Value) -> Result<Vec<McpToolSchema>, McpError> {
430 let arr = v.get("tools").and_then(|t| t.as_array()).ok_or_else(|| {
431 McpError::Protocol(format!("tools/list response missing `tools` array: {v}"))
432 })?;
433 let mut out = Vec::with_capacity(arr.len());
434 for item in arr {
435 let schema: McpToolSchema = serde_json::from_value(item.clone())
436 .map_err(|e| McpError::Protocol(format!("tool schema: {e}")))?;
437 out.push(schema);
438 }
439 Ok(out)
440}
441
442pub fn mcp_result_to_value(result: serde_json::Value) -> crate::value::Value {
443 if let Some(content) = result.get("content").and_then(|c| c.as_array()) {
444 let text_parts: Vec<String> = content
445 .iter()
446 .filter_map(|item| {
447 if item.get("type").and_then(|t| t.as_str()) == Some("text") {
448 item.get("text").and_then(|t| t.as_str()).map(String::from)
449 } else {
450 None
451 }
452 })
453 .collect();
454 let is_error = result
455 .get("isError")
456 .and_then(|b| b.as_bool())
457 .unwrap_or(false);
458 return crate::value::Value::Struct(vec![
459 ("text".into(), crate::value::Value::Str(text_parts.join(""))),
460 ("is_error".into(), crate::value::Value::Bool(is_error)),
461 ("raw".into(), crate::value::Value::from_json(result)),
462 ]);
463 }
464 crate::value::Value::from_json(result)
465}
466
467pub fn value_to_mcp_args(args: &crate::tool::ToolArgs) -> serde_json::Value {
468 let mut map = serde_json::Map::new();
469 for (name, value) in &args.named {
470 map.insert(name.clone(), value.to_json());
471 }
472 serde_json::Value::Object(map)
473}
474
475pub struct McpToolAdapter {
476 qualified_name: String,
477 tool_name: String,
478 tier: crate::tool::Tier,
479 client: Arc<McpClient>,
480}
481
482impl McpToolAdapter {
483 pub fn new(
484 client: Arc<McpClient>,
485 tool_name: impl Into<String>,
486 tier: crate::tool::Tier,
487 ) -> Self {
488 let tool_name = tool_name.into();
489 let qualified_name = format!("{}.{}", client.name, tool_name);
490 Self {
491 qualified_name,
492 tool_name,
493 tier,
494 client,
495 }
496 }
497}
498
499impl crate::tool::Tool for McpToolAdapter {
500 fn name(&self) -> &str {
501 &self.qualified_name
502 }
503
504 fn tier(&self) -> crate::tool::Tier {
505 self.tier
506 }
507
508 fn call<'a>(
509 &'a self,
510 args: crate::tool::ToolArgs,
511 _ctx: &'a crate::tool::ToolCtx,
512 ) -> crate::tool::BoxFut<'a, crate::tool::ToolResult> {
513 Box::pin(async move {
514 let params = value_to_mcp_args(&args);
515 let v = self.client.call_tool(&self.tool_name, params).await?;
516 Ok(v)
517 })
518 }
519}
520
521#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
522pub enum TransportKind {
523 #[default]
524 Stdio,
525 Http,
526}
527
528pub struct McpServerConfig {
529 pub name: String,
530 #[allow(clippy::field_reassign_with_default)]
531 pub transport: TransportKind,
532 pub command: String,
533 pub args: Vec<String>,
534 pub url: Option<String>,
535 pub auth_token: Option<String>,
536 pub tier: crate::tool::Tier,
537 pub timeout_ms: u64,
538}
539
540impl McpServerConfig {
541 pub fn stdio(
542 name: impl Into<String>,
543 command: impl Into<String>,
544 args: Vec<String>,
545 tier: crate::tool::Tier,
546 timeout_ms: u64,
547 ) -> Self {
548 Self {
549 name: name.into(),
550 transport: TransportKind::Stdio,
551 command: command.into(),
552 args,
553 url: None,
554 auth_token: None,
555 tier,
556 timeout_ms,
557 }
558 }
559
560 pub fn http(
561 name: impl Into<String>,
562 url: impl Into<String>,
563 auth_token: Option<String>,
564 tier: crate::tool::Tier,
565 timeout_ms: u64,
566 ) -> Self {
567 Self {
568 name: name.into(),
569 transport: TransportKind::Http,
570 command: String::new(),
571 args: Vec::new(),
572 url: Some(url.into()),
573 auth_token,
574 tier,
575 timeout_ms,
576 }
577 }
578}
579
580pub async fn register_from_configs(
581 reg: &mut crate::tool::ToolRegistry,
582 configs: &[McpServerConfig],
583) -> Vec<Result<McpClientStatus, McpBootError>> {
584 let mut out = Vec::with_capacity(configs.len());
585 for cfg in configs {
586 let outcome = match cfg.transport {
587 TransportKind::Stdio => {
588 McpClient::connect_stdio(&cfg.name, &cfg.command, &cfg.args, cfg.timeout_ms).await
589 }
590 TransportKind::Http => match cfg.url.as_deref() {
591 Some(url) => {
592 McpClient::connect_http(&cfg.name, url, cfg.auth_token.clone(), cfg.timeout_ms)
593 .await
594 }
595 None => Err(McpError::Protocol("http transport requires `url`".into())),
596 },
597 };
598 match outcome {
599 Ok(client) => {
600 let tool_count = client.tools.len();
601 let transport_kind = client.transport_kind();
602 let arc_client = Arc::new(client);
603 for tool in &arc_client.tools {
604 let adapter = McpToolAdapter::new(arc_client.clone(), &tool.name, cfg.tier);
605 reg.register(Arc::new(adapter));
606 }
607 out.push(Ok(McpClientStatus {
608 name: cfg.name.clone(),
609 tool_count,
610 transport: transport_kind,
611 }));
612 }
613 Err(e) => out.push(Err(McpBootError {
614 name: cfg.name.clone(),
615 error: e,
616 })),
617 }
618 }
619 out
620}
621
622#[derive(Debug)]
623pub struct McpClientStatus {
624 pub name: String,
625 pub tool_count: usize,
626 pub transport: &'static str,
627}
628
629#[derive(Debug, thiserror::Error)]
630#[error("mcp `{name}` failed: {error}")]
631pub struct McpBootError {
632 pub name: String,
633 pub error: McpError,
634}
635
636#[cfg(test)]
637mod tests {
638 use super::*;
639
640 #[tokio::test]
641 async fn stdio_transport_call_returns_result() {
642 let script = r#"
643import sys, json
644for line in sys.stdin:
645 req = json.loads(line)
646 resp = {"jsonrpc": "2.0", "id": req["id"], "result": {"echo": req.get("params")}}
647 print(json.dumps(resp), flush=True)
648"#;
649 let dir = tempfile::tempdir().unwrap();
650 let script_path = dir.path().join("mcp_echo.py");
651 std::fs::write(&script_path, script).unwrap();
652
653 let transport =
654 McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 5000)
655 .await
656 .unwrap();
657
658 let result = transport
659 .call("hello", serde_json::json!({"x": 1}))
660 .await
661 .unwrap();
662 assert_eq!(result, serde_json::json!({"echo": {"x": 1}}));
663 }
664
665 #[tokio::test]
666 async fn stdio_transport_propagates_server_error() {
667 let script = r#"
668import sys, json
669for line in sys.stdin:
670 req = json.loads(line)
671 resp = {"jsonrpc": "2.0", "id": req["id"], "error": {"code": -32601, "message": "not found"}}
672 print(json.dumps(resp), flush=True)
673"#;
674 let dir = tempfile::tempdir().unwrap();
675 let script_path = dir.path().join("mcp_err.py");
676 std::fs::write(&script_path, script).unwrap();
677
678 let transport =
679 McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 5000)
680 .await
681 .unwrap();
682 let err = transport
683 .call("boom", serde_json::json!({}))
684 .await
685 .unwrap_err();
686 assert!(matches!(err, McpError::ServerError { code: -32601, .. }));
687 }
688
689 #[tokio::test]
690 async fn stdio_transport_call_times_out() {
691 let script = r#"
692import sys
693for line in sys.stdin:
694 pass
695"#;
696 let dir = tempfile::tempdir().unwrap();
697 let script_path = dir.path().join("mcp_silent.py");
698 std::fs::write(&script_path, script).unwrap();
699
700 let transport =
701 McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 200)
702 .await
703 .unwrap();
704 let err = transport
705 .call("hangs", serde_json::json!({}))
706 .await
707 .unwrap_err();
708 assert!(matches!(err, McpError::Timeout { .. }));
709 }
710
711 #[test]
712 fn tool_schema_deserialize_from_mcp_list_tools_response() {
713 let json = serde_json::json!({
714 "name": "read_file",
715 "description": "reads a file",
716 "inputSchema": {"type": "object", "properties": {"path": {"type": "string"}}}
717 });
718 let s: McpToolSchema = serde_json::from_value(json).unwrap();
719 assert_eq!(s.name, "read_file");
720 assert!(s.description.as_deref().unwrap().contains("reads"));
721 }
722}