Skip to main content

skiff_cli/mcp/
stdio.rs

1//! MCP stdio client via rmcp.
2
3use 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                // Ensure index exists for warm search path.
30                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
71/// Connect to an MCP stdio server. When `clean_env` is true, clear the child
72/// environment and keep only PATH/HOME/LANG plus `env_vars`.
73pub 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}