atman_runtime/
mcp.rs

1use std::collections::{HashMap, HashSet};
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    /// Send a JSON-RPC *notification* — no `id`, no response expected.
26    fn notify<'a>(
27        &'a self,
28        method: &'a str,
29        params: serde_json::Value,
30    ) -> BoxFut<'a, Result<(), McpError>>;
31
32    fn kind(&self) -> &'static str;
33}
34
35#[derive(Debug, Serialize, Deserialize, Clone)]
36pub struct JsonRpcRequest {
37    pub jsonrpc: &'static str,
38    pub id: u64,
39    pub method: String,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub params: Option<serde_json::Value>,
42}
43
44#[derive(Debug, Serialize, Deserialize, Clone)]
45pub struct JsonRpcResponse {
46    #[allow(dead_code)]
47    pub jsonrpc: String,
48    pub id: Option<u64>,
49    #[serde(default)]
50    pub result: Option<serde_json::Value>,
51    #[serde(default)]
52    pub error: Option<JsonRpcError>,
53}
54
55#[derive(Debug, Serialize, Deserialize, Clone)]
56pub struct JsonRpcError {
57    pub code: i64,
58    pub message: String,
59    #[serde(default)]
60    pub data: Option<serde_json::Value>,
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct McpToolSchema {
65    pub name: String,
66    #[serde(default)]
67    pub description: Option<String>,
68    #[serde(default, rename = "inputSchema")]
69    pub input_schema: Option<serde_json::Value>,
70}
71
72#[derive(Debug, Clone, thiserror::Error)]
73pub enum McpError {
74    #[error("mcp io: {0}")]
75    Io(String),
76    #[error("mcp protocol: {0}")]
77    Protocol(String),
78    #[error("mcp server error {code}: {message}")]
79    ServerError { code: i64, message: String },
80    #[error("mcp timeout ({timeout_ms}ms) on {method}")]
81    Timeout { timeout_ms: u64, method: String },
82    #[error("mcp disconnected")]
83    Disconnected,
84}
85
86impl From<McpError> for RuntimeError {
87    fn from(e: McpError) -> Self {
88        RuntimeError::ToolFailed(format!("{e}"))
89    }
90}
91
92pub struct McpStdioTransport {
93    stdin: Arc<Mutex<ChildStdin>>,
94    pending: PendingCalls,
95    next_id: Arc<Mutex<u64>>,
96    #[allow(dead_code)]
97    child: Arc<Mutex<Child>>,
98    #[allow(dead_code)]
99    reader_task: Arc<Mutex<Option<JoinHandle<()>>>>,
100    timeout_ms: u64,
101}
102
103impl McpTransport for McpStdioTransport {
104    fn call<'a>(
105        &'a self,
106        method: &'a str,
107        params: serde_json::Value,
108    ) -> BoxFut<'a, Result<serde_json::Value, McpError>> {
109        Box::pin(self.call_stdio(method, params))
110    }
111
112    fn notify<'a>(
113        &'a self,
114        method: &'a str,
115        params: serde_json::Value,
116    ) -> BoxFut<'a, Result<(), McpError>> {
117        Box::pin(self.notify_stdio(method, params))
118    }
119
120    fn kind(&self) -> &'static str {
121        "stdio"
122    }
123}
124
125impl McpStdioTransport {
126    pub async fn spawn(cmd: &str, args: &[String], timeout_ms: u64) -> Result<Self, McpError> {
127        let mut child = Command::new(cmd)
128            .args(args)
129            .stdin(Stdio::piped())
130            .stdout(Stdio::piped())
131            .stderr(Stdio::piped())
132            .spawn()
133            .map_err(|e| McpError::Io(format!("spawn {cmd}: {e}")))?;
134        let stdin = child
135            .stdin
136            .take()
137            .ok_or_else(|| McpError::Io("no stdin".into()))?;
138        let stdout = child
139            .stdout
140            .take()
141            .ok_or_else(|| McpError::Io("no stdout".into()))?;
142
143        let pending: PendingCalls = Arc::new(Mutex::new(HashMap::new()));
144        let pending_for_reader = pending.clone();
145        let reader_task = tokio::spawn(async move {
146            reader_loop(stdout, pending_for_reader).await;
147        });
148
149        Ok(Self {
150            stdin: Arc::new(Mutex::new(stdin)),
151            pending,
152            next_id: Arc::new(Mutex::new(1)),
153            child: Arc::new(Mutex::new(child)),
154            reader_task: Arc::new(Mutex::new(Some(reader_task))),
155            timeout_ms,
156        })
157    }
158
159    async fn call_stdio(
160        &self,
161        method: &str,
162        params: serde_json::Value,
163    ) -> Result<serde_json::Value, McpError> {
164        let id = {
165            let mut n = self.next_id.lock().await;
166            let v = *n;
167            *n += 1;
168            v
169        };
170        let (tx, rx) = oneshot::channel();
171        self.pending.lock().await.insert(id, tx);
172
173        let req = JsonRpcRequest {
174            jsonrpc: "2.0",
175            id,
176            method: method.into(),
177            params: Some(params),
178        };
179        let line = serde_json::to_string(&req)
180            .map_err(|e| McpError::Protocol(format!("serialize: {e}")))?;
181        {
182            let mut stdin = self.stdin.lock().await;
183            stdin
184                .write_all(line.as_bytes())
185                .await
186                .map_err(|e| McpError::Io(format!("write: {e}")))?;
187            stdin
188                .write_all(b"\n")
189                .await
190                .map_err(|e| McpError::Io(format!("write newline: {e}")))?;
191            stdin
192                .flush()
193                .await
194                .map_err(|e| McpError::Io(format!("flush: {e}")))?;
195        }
196
197        match tokio::time::timeout(Duration::from_millis(self.timeout_ms), rx).await {
198            Ok(Ok(inner)) => inner,
199            Ok(Err(_)) => {
200                self.pending.lock().await.remove(&id);
201                Err(McpError::Disconnected)
202            }
203            Err(_) => {
204                self.pending.lock().await.remove(&id);
205                Err(McpError::Timeout {
206                    timeout_ms: self.timeout_ms,
207                    method: method.into(),
208                })
209            }
210        }
211    }
212
213    async fn notify_stdio(&self, method: &str, params: serde_json::Value) -> Result<(), McpError> {
214        let req = serde_json::json!({
215            "jsonrpc": "2.0",
216            "method": method,
217            "params": params,
218        });
219        let line = serde_json::to_string(&req)
220            .map_err(|e| McpError::Protocol(format!("serialize: {e}")))?;
221        let mut stdin = self.stdin.lock().await;
222        stdin
223            .write_all(line.as_bytes())
224            .await
225            .map_err(|e| McpError::Io(format!("write: {e}")))?;
226        stdin
227            .write_all(b"\n")
228            .await
229            .map_err(|e| McpError::Io(format!("write newline: {e}")))?;
230        stdin
231            .flush()
232            .await
233            .map_err(|e| McpError::Io(format!("flush: {e}")))?;
234        Ok(())
235    }
236}
237
238async fn reader_loop(stdout: ChildStdout, pending: PendingCalls) {
239    let reader = BufReader::new(stdout);
240    let mut lines = reader.lines();
241    loop {
242        match lines.next_line().await {
243            Ok(Some(line)) => {
244                if line.trim().is_empty() {
245                    continue;
246                }
247                let parsed: JsonRpcResponse = match serde_json::from_str(&line) {
248                    Ok(r) => r,
249                    Err(_) => continue,
250                };
251                let Some(id) = parsed.id else {
252                    continue;
253                };
254                let sender = pending.lock().await.remove(&id);
255                if let Some(sender) = sender {
256                    let outcome = if let Some(err) = parsed.error {
257                        Err(McpError::ServerError {
258                            code: err.code,
259                            message: err.message,
260                        })
261                    } else {
262                        Ok(parsed.result.unwrap_or(serde_json::Value::Null))
263                    };
264                    let _ = sender.send(outcome);
265                }
266            }
267            Ok(None) => break,
268            Err(_) => break,
269        }
270    }
271    let mut pending = pending.lock().await;
272    for (_, tx) in pending.drain() {
273        let _ = tx.send(Err(McpError::Disconnected));
274    }
275}
276
277pub struct McpHttpTransport {
278    url: String,
279    auth_token: Option<String>,
280    client: reqwest::Client,
281    next_id: AtomicU64,
282    timeout_ms: u64,
283    retry_attempts: u32,
284}
285
286impl McpHttpTransport {
287    pub fn new(url: impl Into<String>, auth_token: Option<String>, timeout_ms: u64) -> Self {
288        Self {
289            url: url.into(),
290            auth_token,
291            client: reqwest::Client::new(),
292            next_id: AtomicU64::new(1),
293            timeout_ms,
294            retry_attempts: 3,
295        }
296    }
297
298    #[doc(hidden)]
299    pub fn with_client(mut self, client: reqwest::Client) -> Self {
300        self.client = client;
301        self
302    }
303
304    async fn call_http(
305        &self,
306        method: &str,
307        params: serde_json::Value,
308    ) -> Result<serde_json::Value, McpError> {
309        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
310        let body = serde_json::json!({
311            "jsonrpc": "2.0",
312            "id": id,
313            "method": method,
314            "params": params,
315        });
316        let call_timeout = Duration::from_millis(self.timeout_ms);
317        let attempt_result = tokio::time::timeout(call_timeout, async {
318            let mut delay_ms = 100u64;
319            let mut last_err: Option<McpError> = None;
320            for attempt in 0..=self.retry_attempts {
321                let mut req = self.client.post(&self.url).json(&body);
322                if let Some(t) = &self.auth_token {
323                    req = req.bearer_auth(t);
324                }
325                match req.send().await {
326                    Ok(resp) => {
327                        let status = resp.status();
328                        if status.is_success() {
329                            let text = resp
330                                .text()
331                                .await
332                                .map_err(|e| McpError::Io(format!("mcp http body: {e}")))?;
333                            let parsed: JsonRpcResponse = serde_json::from_str(&text)
334                                .map_err(|e| McpError::Protocol(format!("mcp http parse: {e}")))?;
335                            if let Some(err) = parsed.error {
336                                return Err(McpError::ServerError {
337                                    code: err.code,
338                                    message: err.message,
339                                });
340                            }
341                            return Ok(parsed.result.unwrap_or(serde_json::Value::Null));
342                        }
343                        if status.is_server_error() {
344                            last_err =
345                                Some(McpError::Io(format!("mcp http {status}: server error")));
346                            if attempt < self.retry_attempts {
347                                tokio::time::sleep(Duration::from_millis(delay_ms)).await;
348                                delay_ms = (delay_ms * 5).min(2000);
349                                continue;
350                            }
351                        }
352                        let body_text = resp.text().await.unwrap_or_default();
353                        return Err(McpError::Io(format!("mcp http {status}: {body_text}")));
354                    }
355                    Err(e) => {
356                        last_err = Some(McpError::Io(format!("mcp http send: {e}")));
357                        if attempt < self.retry_attempts {
358                            tokio::time::sleep(Duration::from_millis(delay_ms)).await;
359                            delay_ms = (delay_ms * 5).min(2000);
360                            continue;
361                        }
362                    }
363                }
364            }
365            Err(last_err.unwrap_or(McpError::Disconnected))
366        })
367        .await;
368        match attempt_result {
369            Ok(inner) => inner,
370            Err(_) => Err(McpError::Timeout {
371                timeout_ms: self.timeout_ms,
372                method: method.into(),
373            }),
374        }
375    }
376
377    async fn notify_http(&self, method: &str, params: serde_json::Value) -> Result<(), McpError> {
378        let body = serde_json::json!({
379            "jsonrpc": "2.0",
380            "method": method,
381            "params": params,
382        });
383        let mut req = self.client.post(&self.url).json(&body);
384        if let Some(t) = &self.auth_token {
385            req = req.bearer_auth(t);
386        }
387        // Fire and forget — don't wait for response body.
388        let _ = req.send().await;
389        Ok(())
390    }
391}
392
393impl McpTransport for McpHttpTransport {
394    fn call<'a>(
395        &'a self,
396        method: &'a str,
397        params: serde_json::Value,
398    ) -> BoxFut<'a, Result<serde_json::Value, McpError>> {
399        Box::pin(self.call_http(method, params))
400    }
401
402    fn notify<'a>(
403        &'a self,
404        method: &'a str,
405        params: serde_json::Value,
406    ) -> BoxFut<'a, Result<(), McpError>> {
407        Box::pin(self.notify_http(method, params))
408    }
409
410    fn kind(&self) -> &'static str {
411        "http"
412    }
413}
414
415pub struct McpClient {
416    pub name: String,
417    transport: std::sync::Mutex<Arc<dyn McpTransport>>,
418    pub tools: Vec<McpToolSchema>,
419    reconnect: Option<ReconnectConfig>,
420}
421
422/// Config needed to re-create the transport on disconnect.
423enum ReconnectConfig {
424    Stdio {
425        cmd: String,
426        args: Vec<String>,
427        timeout_ms: u64,
428    },
429    Http {
430        url: String,
431        auth_token: Option<String>,
432        timeout_ms: u64,
433    },
434}
435
436impl McpClient {
437    pub async fn connect_stdio(
438        name: impl Into<String>,
439        cmd: &str,
440        args: &[String],
441        timeout_ms: u64,
442    ) -> Result<Self, McpError> {
443        let transport: Arc<dyn McpTransport> =
444            Arc::new(McpStdioTransport::spawn(cmd, args, timeout_ms).await?);
445        let name = name.into();
446        let mut client = Self::finish_connect(name.clone(), transport).await?;
447        client.reconnect = Some(ReconnectConfig::Stdio {
448            cmd: cmd.to_string(),
449            args: args.to_vec(),
450            timeout_ms,
451        });
452        Ok(client)
453    }
454
455    pub async fn connect_http(
456        name: impl Into<String>,
457        url: impl Into<String>,
458        auth_token: Option<String>,
459        timeout_ms: u64,
460    ) -> Result<Self, McpError> {
461        let url_str: String = url.into();
462        let transport: Arc<dyn McpTransport> = Arc::new(McpHttpTransport::new(
463            url_str.clone(),
464            auth_token.clone(),
465            timeout_ms,
466        ));
467        let name = name.into();
468        let mut client = Self::finish_connect(name.clone(), transport).await?;
469        client.reconnect = Some(ReconnectConfig::Http {
470            url: url_str,
471            auth_token,
472            timeout_ms,
473        });
474        Ok(client)
475    }
476
477    pub async fn connect_with_transport(
478        name: impl Into<String>,
479        transport: Arc<dyn McpTransport>,
480    ) -> Result<Self, McpError> {
481        Self::finish_connect(name.into(), transport).await
482    }
483
484    async fn finish_connect(
485        name: String,
486        transport: Arc<dyn McpTransport>,
487    ) -> Result<Self, McpError> {
488        let init_params = serde_json::json!({
489            "protocolVersion": "2024-11-05",
490            "capabilities": {},
491            "clientInfo": {"name": "atman", "version": env!("CARGO_PKG_VERSION")}
492        });
493        transport.call("initialize", init_params).await?;
494        transport
495            .notify("notifications/initialized", serde_json::Value::Null)
496            .await?;
497        let list = transport.call("tools/list", serde_json::json!({})).await?;
498        let tools = parse_tools_list(&list)?;
499        Ok(Self {
500            name,
501            transport: std::sync::Mutex::new(transport),
502            tools,
503            reconnect: None,
504        })
505    }
506
507    pub fn transport_kind(&self) -> &'static str {
508        let t = self.transport.lock().unwrap();
509        t.kind()
510    }
511
512    async fn reconnect(&self) -> Result<(), McpError> {
513        let cfg = self.reconnect.as_ref().ok_or(McpError::Disconnected)?;
514        let new_transport: Arc<dyn McpTransport> = match cfg {
515            ReconnectConfig::Stdio {
516                cmd,
517                args,
518                timeout_ms,
519            } => Arc::new(McpStdioTransport::spawn(cmd, args, *timeout_ms).await?),
520            ReconnectConfig::Http {
521                url,
522                auth_token,
523                timeout_ms,
524            } => Arc::new(McpHttpTransport::new(url, auth_token.clone(), *timeout_ms)),
525        };
526        // Re-initialize the new transport.
527        let init_params = serde_json::json!({
528            "protocolVersion": "2024-11-05",
529            "capabilities": {},
530            "clientInfo": {"name": "atman", "version": env!("CARGO_PKG_VERSION")}
531        });
532        new_transport.call("initialize", init_params).await?;
533        new_transport
534            .notify("notifications/initialized", serde_json::Value::Null)
535            .await?;
536        *self.transport.lock().unwrap() = new_transport;
537        Ok(())
538    }
539
540    pub async fn call_tool(
541        &self,
542        tool_name: &str,
543        arguments: serde_json::Value,
544    ) -> Result<crate::value::Value, McpError> {
545        let params = serde_json::json!({
546            "name": tool_name,
547            "arguments": arguments,
548        });
549        let transport = { self.transport.lock().unwrap().clone() };
550        let result = transport.call("tools/call", params.clone()).await;
551        match result {
552            Err(McpError::Disconnected) => {
553                self.reconnect().await?;
554                let transport = { self.transport.lock().unwrap().clone() };
555                let retry = transport.call("tools/call", params).await?;
556                Ok(mcp_result_to_value(retry))
557            }
558            other => Ok(mcp_result_to_value(other?)),
559        }
560    }
561}
562
563fn parse_tools_list(v: &serde_json::Value) -> Result<Vec<McpToolSchema>, McpError> {
564    let arr = v.get("tools").and_then(|t| t.as_array()).ok_or_else(|| {
565        McpError::Protocol(format!("tools/list response missing `tools` array: {v}"))
566    })?;
567    let mut out = Vec::with_capacity(arr.len());
568    for item in arr {
569        let schema: McpToolSchema = serde_json::from_value(item.clone())
570            .map_err(|e| McpError::Protocol(format!("tool schema: {e}")))?;
571        out.push(schema);
572    }
573    Ok(out)
574}
575
576pub fn mcp_result_to_value(result: serde_json::Value) -> crate::value::Value {
577    if let Some(content) = result.get("content").and_then(|c| c.as_array()) {
578        let text_parts: Vec<String> = content
579            .iter()
580            .filter_map(|item| {
581                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
582                    item.get("text").and_then(|t| t.as_str()).map(String::from)
583                } else {
584                    None
585                }
586            })
587            .collect();
588        let is_error = result
589            .get("isError")
590            .and_then(|b| b.as_bool())
591            .unwrap_or(false);
592        return crate::value::Value::Struct(vec![
593            ("text".into(), crate::value::Value::Str(text_parts.join(""))),
594            ("is_error".into(), crate::value::Value::Bool(is_error)),
595            ("raw".into(), crate::value::Value::from_json(result)),
596        ]);
597    }
598    crate::value::Value::from_json(result)
599}
600
601pub fn value_to_mcp_args(args: &crate::tool::ToolArgs) -> serde_json::Value {
602    let mut map = serde_json::Map::new();
603    for (name, value) in &args.named {
604        map.insert(name.clone(), value.to_json());
605    }
606    serde_json::Value::Object(map)
607}
608
609// Some MCP servers use a schema like {action: enum, params: {type: "object"}}
610// where params is a catch-all for every argument not named at the top level.
611// atman sends flat arguments and Zod drops everything except action+params,
612// so we detect the pattern and repack unmatched keys into the container.
613
614#[derive(Debug, Clone)]
615struct ReconcilePlan {
616    explicit_keys: HashSet<String>,
617    container_key: String,
618    container_required: bool,
619}
620
621fn build_reconcile_plan(schema: &serde_json::Value) -> Option<ReconcilePlan> {
622    let props = schema.get("properties")?.as_object()?;
623    let required: HashSet<&str> = schema
624        .get("required")
625        .and_then(|r| r.as_array())
626        .map(|a| a.iter().filter_map(|v| v.as_str()).collect())
627        .unwrap_or_default();
628
629    let mut explicit_keys = HashSet::new();
630    let mut catch_all: Option<(&str, bool)> = None;
631
632    for (name, prop) in props {
633        let is_obj = prop.get("type").and_then(|t| t.as_str()) == Some("object");
634        let has_sub_props = prop.get("properties").is_some_and(|p| p.is_object());
635        let add_props_false =
636            prop.get("additionalProperties").and_then(|a| a.as_bool()) == Some(false);
637
638        if is_obj && !has_sub_props && !add_props_false {
639            if catch_all.is_some() {
640                return None;
641            }
642            catch_all = Some((name.as_str(), required.contains(name.as_str())));
643        } else {
644            explicit_keys.insert(name.clone());
645        }
646    }
647
648    let (container_key, container_required) = catch_all?;
649    Some(ReconcilePlan {
650        explicit_keys,
651        container_key: container_key.to_owned(),
652        container_required,
653    })
654}
655
656fn reconcile(
657    plan: &ReconcilePlan,
658    flat: serde_json::Map<String, serde_json::Value>,
659) -> serde_json::Value {
660    let mut container = serde_json::Map::new();
661    let mut out = serde_json::Map::new();
662
663    for (k, v) in flat {
664        if plan.explicit_keys.contains(&k) {
665            out.insert(k, v);
666        } else {
667            container.insert(k, v);
668        }
669    }
670
671    if !container.is_empty() || plan.container_required {
672        out.insert(
673            plan.container_key.clone(),
674            serde_json::Value::Object(container),
675        );
676    }
677
678    serde_json::Value::Object(out)
679}
680
681pub struct McpToolAdapter {
682    qualified_name: String,
683    tool_name: String,
684    tier: crate::tool::Tier,
685    client: Arc<McpClient>,
686    reconcile: Option<ReconcilePlan>,
687}
688
689impl McpToolAdapter {
690    pub fn new(
691        client: Arc<McpClient>,
692        tool_name: impl Into<String>,
693        tier: crate::tool::Tier,
694        schema: Option<&serde_json::Value>,
695    ) -> Self {
696        let tool_name = tool_name.into();
697        let qualified_name = format!("mcp.{}.{}", client.name, tool_name);
698        let reconcile = schema.and_then(build_reconcile_plan);
699        Self {
700            qualified_name,
701            tool_name,
702            tier,
703            client,
704            reconcile,
705        }
706    }
707}
708
709impl crate::tool::Tool for McpToolAdapter {
710    fn name(&self) -> &str {
711        &self.qualified_name
712    }
713
714    fn tier(&self) -> crate::tool::Tier {
715        self.tier
716    }
717
718    fn call<'a>(
719        &'a self,
720        args: crate::tool::ToolArgs,
721        _ctx: &'a crate::tool::ToolCtx,
722    ) -> crate::tool::BoxFut<'a, crate::tool::ToolResult> {
723        Box::pin(async move {
724            let params = value_to_mcp_args(&args);
725            let params = if let Some(plan) = &self.reconcile {
726                let map = match params {
727                    serde_json::Value::Object(m) => m,
728                    _ => return Err(RuntimeError::ToolFailed("mcp: expected object args".into())),
729                };
730                reconcile(plan, map)
731            } else {
732                params
733            };
734            let v = self.client.call_tool(&self.tool_name, params).await?;
735            Ok(v)
736        })
737    }
738}
739
740#[derive(Debug, Clone, PartialEq, Eq)]
741pub enum McpServerState {
742    Disabled,
743    Pending,
744    Connecting,
745    Connected { tool_count: usize },
746    Error { message: String },
747    Disconnected { message: String },
748    Timeout { message: String },
749}
750
751#[derive(Debug, Clone, PartialEq, Eq)]
752pub struct McpServerStatus {
753    pub name: String,
754    pub transport: TransportKind,
755    pub state: McpServerState,
756}
757
758impl McpServerStatus {
759    pub fn is_ok(&self) -> bool {
760        matches!(self.state, McpServerState::Connected { .. })
761    }
762}
763
764/// Derive ok/total counts from a list of server statuses.
765pub fn mcp_counts(servers: &[McpServerStatus]) -> (u16, u16) {
766    let total = servers.len() as u16;
767    let ok = servers.iter().filter(|s| s.is_ok()).count() as u16;
768    (ok, total)
769}
770
771#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
772pub enum TransportKind {
773    #[default]
774    Stdio,
775    Http,
776}
777
778pub struct McpServerConfig {
779    pub name: String,
780    #[allow(clippy::field_reassign_with_default)]
781    pub transport: TransportKind,
782    pub command: String,
783    pub args: Vec<String>,
784    pub url: Option<String>,
785    pub auth_token: Option<String>,
786    pub tier: crate::tool::Tier,
787    pub timeout_ms: u64,
788    pub disabled: bool,
789}
790
791impl McpServerConfig {
792    pub fn stdio(
793        name: impl Into<String>,
794        command: impl Into<String>,
795        args: Vec<String>,
796        tier: crate::tool::Tier,
797        timeout_ms: u64,
798    ) -> Self {
799        Self {
800            name: name.into(),
801            transport: TransportKind::Stdio,
802            command: command.into(),
803            args,
804            url: None,
805            auth_token: None,
806            tier,
807            timeout_ms,
808            disabled: false,
809        }
810    }
811
812    pub fn http(
813        name: impl Into<String>,
814        url: impl Into<String>,
815        auth_token: Option<String>,
816        tier: crate::tool::Tier,
817        timeout_ms: u64,
818    ) -> Self {
819        Self {
820            name: name.into(),
821            transport: TransportKind::Http,
822            command: String::new(),
823            args: Vec::new(),
824            url: Some(url.into()),
825            auth_token,
826            tier,
827            timeout_ms,
828            disabled: false,
829        }
830    }
831}
832
833pub async fn register_from_configs(
834    reg: &mut crate::tool::ToolRegistry,
835    configs: &[McpServerConfig],
836) -> Vec<Result<McpClientStatus, McpBootError>> {
837    let mut out = Vec::with_capacity(configs.len());
838    for cfg in configs {
839        let outcome = match cfg.transport {
840            TransportKind::Stdio => {
841                McpClient::connect_stdio(&cfg.name, &cfg.command, &cfg.args, cfg.timeout_ms).await
842            }
843            TransportKind::Http => match cfg.url.as_deref() {
844                Some(url) => {
845                    McpClient::connect_http(&cfg.name, url, cfg.auth_token.clone(), cfg.timeout_ms)
846                        .await
847                }
848                None => Err(McpError::Protocol("http transport requires `url`".into())),
849            },
850        };
851        match outcome {
852            Ok(client) => {
853                let tool_count = client.tools.len();
854                let transport_kind = client.transport_kind();
855                let arc_client = Arc::new(client);
856                for tool in &arc_client.tools {
857                    let adapter = McpToolAdapter::new(
858                        arc_client.clone(),
859                        &tool.name,
860                        cfg.tier,
861                        tool.input_schema.as_ref(),
862                    );
863                    reg.register(Arc::new(adapter));
864                }
865                out.push(Ok(McpClientStatus {
866                    name: cfg.name.clone(),
867                    tool_count,
868                    transport: transport_kind,
869                }));
870            }
871            Err(e) => out.push(Err(McpBootError {
872                name: cfg.name.clone(),
873                error: e,
874            })),
875        }
876    }
877    out
878}
879
880#[derive(Debug)]
881pub struct McpClientStatus {
882    pub name: String,
883    pub tool_count: usize,
884    pub transport: &'static str,
885}
886
887#[derive(Debug, thiserror::Error)]
888#[error("mcp `{name}` failed: {error}")]
889pub struct McpBootError {
890    pub name: String,
891    pub error: McpError,
892}
893
894#[cfg(test)]
895mod tests {
896    use super::*;
897
898    #[tokio::test]
899    async fn stdio_transport_call_returns_result() {
900        let script = r#"
901import sys, json
902for line in sys.stdin:
903    req = json.loads(line)
904    resp = {"jsonrpc": "2.0", "id": req["id"], "result": {"echo": req.get("params")}}
905    print(json.dumps(resp), flush=True)
906"#;
907        let dir = tempfile::tempdir().unwrap();
908        let script_path = dir.path().join("mcp_echo.py");
909        std::fs::write(&script_path, script).unwrap();
910
911        let transport =
912            McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 5000)
913                .await
914                .unwrap();
915
916        let result = transport
917            .call("hello", serde_json::json!({"x": 1}))
918            .await
919            .unwrap();
920        assert_eq!(result, serde_json::json!({"echo": {"x": 1}}));
921    }
922
923    #[tokio::test]
924    async fn stdio_transport_propagates_server_error() {
925        let script = r#"
926import sys, json
927for line in sys.stdin:
928    req = json.loads(line)
929    resp = {"jsonrpc": "2.0", "id": req["id"], "error": {"code": -32601, "message": "not found"}}
930    print(json.dumps(resp), flush=True)
931"#;
932        let dir = tempfile::tempdir().unwrap();
933        let script_path = dir.path().join("mcp_err.py");
934        std::fs::write(&script_path, script).unwrap();
935
936        let transport =
937            McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 5000)
938                .await
939                .unwrap();
940        let err = transport
941            .call("boom", serde_json::json!({}))
942            .await
943            .unwrap_err();
944        assert!(matches!(err, McpError::ServerError { code: -32601, .. }));
945    }
946
947    #[tokio::test]
948    async fn stdio_transport_call_times_out() {
949        let script = r#"
950import sys
951for line in sys.stdin:
952    pass
953"#;
954        let dir = tempfile::tempdir().unwrap();
955        let script_path = dir.path().join("mcp_silent.py");
956        std::fs::write(&script_path, script).unwrap();
957
958        let transport =
959            McpStdioTransport::spawn("python3", &[script_path.display().to_string()], 200)
960                .await
961                .unwrap();
962        let err = transport
963            .call("hangs", serde_json::json!({}))
964            .await
965            .unwrap_err();
966        assert!(matches!(err, McpError::Timeout { .. }));
967    }
968
969    #[test]
970    fn tool_schema_deserialize_from_mcp_list_tools_response() {
971        let json = serde_json::json!({
972            "name": "read_file",
973            "description": "reads a file",
974            "inputSchema": {"type": "object", "properties": {"path": {"type": "string"}}}
975        });
976        let s: McpToolSchema = serde_json::from_value(json).unwrap();
977        assert_eq!(s.name, "read_file");
978        assert!(s.description.as_deref().unwrap().contains("reads"));
979    }
980
981    #[test]
982    fn no_catch_all_returns_none() {
983        let schema = serde_json::json!({
984            "type": "object",
985            "properties": {
986                "document_id": {"type": "string"}
987            },
988            "required": ["document_id"]
989        });
990        assert!(build_reconcile_plan(&schema).is_none());
991    }
992
993    #[test]
994    fn single_catch_all_builds_plan() {
995        let schema = serde_json::json!({
996            "type": "object",
997            "properties": {
998                "action": {"type": "string", "enum": ["get", "list", "update"]},
999                "params": {"type": "object"}
1000            },
1001            "required": ["action"]
1002        });
1003        let plan = build_reconcile_plan(&schema).expect("should build a plan");
1004        assert!(plan.explicit_keys.contains("action"));
1005        assert!(!plan.explicit_keys.contains("params"));
1006        assert_eq!(plan.container_key, "params");
1007        assert!(!plan.container_required);
1008    }
1009
1010    #[test]
1011    fn container_required_is_detected() {
1012        let schema = serde_json::json!({
1013            "type": "object",
1014            "properties": {
1015                "action": {"type": "string"},
1016                "params": {"type": "object"}
1017            },
1018            "required": ["action", "params"]
1019        });
1020        let plan = build_reconcile_plan(&schema).expect("should build");
1021        assert!(plan.container_required);
1022    }
1023
1024    #[test]
1025    fn reconcile_moves_unmatched_into_container() {
1026        let plan = ReconcilePlan {
1027            explicit_keys: ["action".into()].into(),
1028            container_key: "params".into(),
1029            container_required: false,
1030        };
1031        let flat: serde_json::Map<_, _> = serde_json::json!({
1032            "action": "update",
1033            "guid": "abc",
1034            "due": "2026-07-17"
1035        })
1036        .as_object()
1037        .unwrap()
1038        .clone();
1039
1040        let out = reconcile(&plan, flat);
1041        assert_eq!(out["action"], "update");
1042        assert_eq!(out["params"]["guid"], "abc");
1043        assert_eq!(out["params"]["due"], "2026-07-17");
1044    }
1045
1046    #[test]
1047    fn reconcile_leaves_standard_schema_untouched() {
1048        let schema = serde_json::json!({
1049            "type": "object",
1050            "properties": {
1051                "action": {"type": "string"},
1052                "guid": {"type": "string"}
1053            },
1054            "required": ["action", "guid"]
1055        });
1056        assert!(build_reconcile_plan(&schema).is_none());
1057    }
1058
1059    #[test]
1060    fn two_catch_alls_bails() {
1061        let schema = serde_json::json!({
1062            "type": "object",
1063            "properties": {
1064                "action": {"type": "string"},
1065                "params": {"type": "object"},
1066                "extras": {"type": "object"}
1067            }
1068        });
1069        assert!(build_reconcile_plan(&schema).is_none());
1070    }
1071
1072    #[test]
1073    fn sealed_object_not_catch_all() {
1074        let schema = serde_json::json!({
1075            "type": "object",
1076            "properties": {
1077                "action": {"type": "string"},
1078                "lock": {"type": "object", "additionalProperties": false}
1079            }
1080        });
1081        assert!(build_reconcile_plan(&schema).is_none());
1082    }
1083
1084    #[test]
1085    fn nested_object_with_properties_not_catch_all() {
1086        let schema = serde_json::json!({
1087            "type": "object",
1088            "properties": {
1089                "action": {"type": "string"},
1090                "address": {
1091                    "type": "object",
1092                    "properties": {
1093                        "street": {"type": "string"},
1094                        "city": {"type": "string"}
1095                    }
1096                }
1097            }
1098        });
1099        assert!(build_reconcile_plan(&schema).is_none());
1100    }
1101
1102    #[test]
1103    fn required_empty_container_still_emitted() {
1104        let plan = ReconcilePlan {
1105            explicit_keys: ["action".into()].into(),
1106            container_key: "body".into(),
1107            container_required: true,
1108        };
1109        let flat: serde_json::Map<_, _> = serde_json::json!({"action": "ping"})
1110            .as_object()
1111            .unwrap()
1112            .clone();
1113
1114        let out = reconcile(&plan, flat);
1115        assert_eq!(out["action"], "ping");
1116        assert!(out.get("body").and_then(|v| v.as_object()).is_some());
1117    }
1118
1119    #[test]
1120    fn empty_non_required_container_omitted() {
1121        let plan = ReconcilePlan {
1122            explicit_keys: ["action".into()].into(),
1123            container_key: "params".into(),
1124            container_required: false,
1125        };
1126        let flat: serde_json::Map<_, _> = serde_json::json!({"action": "list"})
1127            .as_object()
1128            .unwrap()
1129            .clone();
1130
1131        let out = reconcile(&plan, flat);
1132        assert_eq!(out["action"], "list");
1133        assert!(out.get("params").is_none());
1134    }
1135}
atman_runtime/mcp.rs

atman_runtime/
mcp.rs
