1use std::collections::BTreeMap;
4use std::time::Duration;
5
6use rmcp::{
7 transport::{ConfigureCommandExt, TokioChildProcess},
8 ServiceExt,
9};
10use serde_json::{Map, Value};
11use tokio::process::Command;
12
13use crate::cache::load_cached;
14use crate::error::{Error, Result};
15use crate::mcp::common::{call_tool_on, list_tools_on, McpClient};
16use crate::tools_index::save_tools_and_index;
17
18pub async fn fetch_mcp_tools_stdio(
19 command_str: &str,
20 env_vars: &BTreeMap<String, String>,
21 cache_key: &str,
22 ttl: u64,
23 refresh: bool,
24) -> Result<Vec<Value>> {
25 let tools_key = format!("{cache_key}_tools");
26 if !refresh {
27 if let Some(cached) = load_cached(&tools_key, ttl)? {
28 if let Some(arr) = cached.as_array() {
29 let _ = crate::tools_index::save_index(&tools_key, arr);
31 return Ok(arr.clone());
32 }
33 }
34 }
35
36 let tools = list_tools_stdio(command_str, env_vars).await?;
37 save_tools_and_index(&tools_key, &tools)?;
38 Ok(tools)
39}
40
41pub async fn list_tools_stdio(
42 command_str: &str,
43 env_vars: &BTreeMap<String, String>,
44) -> Result<Vec<Value>> {
45 let client = connect_stdio(command_str, env_vars).await?;
46 let tools = list_tools_on(&client).await?;
47 let _ = client.cancel().await;
48 Ok(tools)
49}
50
51pub async fn call_tool_stdio(
52 command_str: &str,
53 env_vars: &BTreeMap<String, String>,
54 tool_name: &str,
55 arguments: Map<String, Value>,
56 full_envelope: bool,
57) -> Result<Value> {
58 let client = connect_stdio(command_str, env_vars).await?;
59 let result = call_tool_on(&client, tool_name, arguments, full_envelope).await?;
60 let _ = client.cancel().await;
61 Ok(result)
62}
63
64pub async fn connect_stdio(
65 command_str: &str,
66 env_vars: &BTreeMap<String, String>,
67) -> Result<McpClient> {
68 connect_stdio_with(command_str, env_vars, false).await
69}
70
71pub async fn connect_stdio_with(
74 command_str: &str,
75 env_vars: &BTreeMap<String, String>,
76 clean_env: bool,
77) -> Result<McpClient> {
78 let parts = shell_words::split(command_str)
79 .map_err(|e| Error::runtime(format!("invalid --mcp-stdio command: {e}")))?;
80 if parts.is_empty() {
81 return Err(Error::usage("--mcp-stdio command is empty"));
82 }
83
84 let transport = TokioChildProcess::new(Command::new(&parts[0]).configure(|c| {
85 c.args(&parts[1..]).kill_on_drop(true);
86 if clean_env {
87 c.env_clear();
88 for key in ["PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "TMP", "TEMP"] {
89 if let Ok(v) = std::env::var(key) {
90 c.env(key, v);
91 }
92 }
93 }
94 for (k, v) in env_vars {
95 c.env(k, v);
96 }
97 }))
98 .map_err(|e| Error::runtime(format!("failed to start MCP stdio process: {e}")))?;
99
100 tokio::time::timeout(Duration::from_secs(30), ().serve(transport))
101 .await
102 .map_err(|_| Error::runtime("MCP initialize timed out after 30s"))?
103 .map_err(|e| Error::runtime(format!("MCP initialize failed: {e}")))
104}