Skip to main content

tower_mcp/transport/
childproc.rs

1//! Child process transport for MCP
2//!
3//! Spawns and communicates with subprocess MCP servers via stdio.
4//! Useful for:
5//! - Running untrusted MCP servers in isolation
6//! - Spawning tool-specific servers on demand
7//! - Testing
8//!
9//! # Example
10//!
11//! ```rust,no_run
12//! use tower_mcp::BoxError;
13//! use tower_mcp::transport::childproc::ChildProcessTransport;
14//!
15//! #[tokio::main]
16//! async fn main() -> Result<(), BoxError> {
17//!     // Spawn an MCP server as a child process
18//!     let mut transport = ChildProcessTransport::new("my-mcp-server")
19//!         .arg("--some-flag")
20//!         .spawn()
21//!         .await?;
22//!
23//!     // Send a request
24//!     let response = transport.send_request(
25//!         "initialize",
26//!         serde_json::json!({
27//!             "protocolVersion": "2025-11-25",
28//!             "capabilities": {},
29//!             "clientInfo": { "name": "my-client", "version": "1.0" }
30//!         })
31//!     ).await?;
32//!
33//!     // Shutdown
34//!     transport.shutdown().await?;
35//!     Ok(())
36//! }
37//! ```
38
39use std::process::Stdio;
40use std::sync::atomic::{AtomicI64, Ordering};
41
42use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
43use tokio::process::{Child, Command};
44
45use crate::error::{Error, Result};
46use crate::protocol::{JsonRpcRequest, JsonRpcResponse};
47
48/// Builder for child process transport
49pub struct ChildProcessTransport {
50    program: String,
51    args: Vec<String>,
52    envs: Vec<(String, String)>,
53}
54
55impl ChildProcessTransport {
56    /// Create a new child process transport builder
57    pub fn new(program: impl Into<String>) -> Self {
58        Self {
59            program: program.into(),
60            args: Vec::new(),
61            envs: Vec::new(),
62        }
63    }
64
65    /// Add a command-line argument
66    pub fn arg(mut self, arg: impl Into<String>) -> Self {
67        self.args.push(arg.into());
68        self
69    }
70
71    /// Add multiple command-line arguments
72    pub fn args<I, S>(mut self, args: I) -> Self
73    where
74        I: IntoIterator<Item = S>,
75        S: Into<String>,
76    {
77        self.args.extend(args.into_iter().map(|s| s.into()));
78        self
79    }
80
81    /// Set an environment variable
82    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
83        self.envs.push((key.into(), value.into()));
84        self
85    }
86
87    /// Spawn the child process
88    pub async fn spawn(self) -> Result<ChildProcessConnection> {
89        let mut cmd = Command::new(&self.program);
90        cmd.args(&self.args)
91            .stdin(Stdio::piped())
92            .stdout(Stdio::piped())
93            .stderr(Stdio::inherit());
94
95        for (key, value) in &self.envs {
96            cmd.env(key, value);
97        }
98
99        let child = cmd
100            .spawn()
101            .map_err(|e| Error::Transport(format!("Failed to spawn {}: {}", self.program, e)))?;
102
103        tracing::info!(program = %self.program, "Spawned child process");
104
105        ChildProcessConnection::new(child)
106    }
107}
108
109/// Active connection to a child MCP server process
110pub struct ChildProcessConnection {
111    child: Child,
112    stdin: tokio::process::ChildStdin,
113    stdout: BufReader<tokio::process::ChildStdout>,
114    request_id: AtomicI64,
115}
116
117impl ChildProcessConnection {
118    fn new(mut child: Child) -> Result<Self> {
119        let stdin = child
120            .stdin
121            .take()
122            .ok_or_else(|| Error::Transport("Failed to get child stdin".to_string()))?;
123        let stdout = child
124            .stdout
125            .take()
126            .ok_or_else(|| Error::Transport("Failed to get child stdout".to_string()))?;
127
128        Ok(Self {
129            child,
130            stdin,
131            stdout: BufReader::new(stdout),
132            request_id: AtomicI64::new(1),
133        })
134    }
135
136    /// Send a JSON-RPC request and wait for response
137    pub async fn send_request(
138        &mut self,
139        method: &str,
140        params: serde_json::Value,
141    ) -> Result<serde_json::Value> {
142        let id = self.request_id.fetch_add(1, Ordering::Relaxed);
143        let request = JsonRpcRequest::new(id, method).with_params(params);
144
145        // Send request
146        let request_json = serde_json::to_string(&request)
147            .map_err(|e| Error::Transport(format!("Failed to serialize request: {}", e)))?;
148
149        tracing::debug!(method = %method, id = %id, "Sending request to child");
150
151        self.stdin
152            .write_all(request_json.as_bytes())
153            .await
154            .map_err(|e| Error::Transport(format!("Failed to write to child stdin: {}", e)))?;
155        self.stdin
156            .write_all(b"\n")
157            .await
158            .map_err(|e| Error::Transport(format!("Failed to write newline: {}", e)))?;
159        self.stdin
160            .flush()
161            .await
162            .map_err(|e| Error::Transport(format!("Failed to flush stdin: {}", e)))?;
163
164        // Read response
165        let mut line = String::new();
166        self.stdout
167            .read_line(&mut line)
168            .await
169            .map_err(|e| Error::Transport(format!("Failed to read from child stdout: {}", e)))?;
170
171        if line.is_empty() {
172            return Err(Error::Transport("Child process closed stdout".to_string()));
173        }
174
175        tracing::debug!(response = %line.trim(), "Received response from child");
176
177        let response: JsonRpcResponse = serde_json::from_str(line.trim())
178            .map_err(|e| Error::Transport(format!("Failed to parse response: {}", e)))?;
179
180        match response {
181            JsonRpcResponse::Result(r) => Ok(r.result),
182            JsonRpcResponse::Error(e) => Err(Error::JsonRpc(e.error)),
183            _ => Err(Error::Transport("unexpected response variant".to_string())),
184        }
185    }
186
187    /// Send a notification (no response expected)
188    pub async fn send_notification(
189        &mut self,
190        method: &str,
191        params: serde_json::Value,
192    ) -> Result<()> {
193        let notification = serde_json::json!({
194            "jsonrpc": "2.0",
195            "method": method,
196            "params": params
197        });
198
199        let json = serde_json::to_string(&notification)
200            .map_err(|e| Error::Transport(format!("Failed to serialize notification: {}", e)))?;
201
202        tracing::debug!(method = %method, "Sending notification to child");
203
204        self.stdin
205            .write_all(json.as_bytes())
206            .await
207            .map_err(|e| Error::Transport(format!("Failed to write notification: {}", e)))?;
208        self.stdin
209            .write_all(b"\n")
210            .await
211            .map_err(|e| Error::Transport(format!("Failed to write newline: {}", e)))?;
212        self.stdin
213            .flush()
214            .await
215            .map_err(|e| Error::Transport(format!("Failed to flush stdin: {}", e)))?;
216
217        Ok(())
218    }
219
220    /// Initialize the MCP connection
221    pub async fn initialize(
222        &mut self,
223        client_name: &str,
224        client_version: &str,
225    ) -> Result<serde_json::Value> {
226        self.send_request(
227            "initialize",
228            serde_json::json!({
229                "protocolVersion": "2025-11-25",
230                "capabilities": {},
231                "clientInfo": {
232                    "name": client_name,
233                    "version": client_version
234                }
235            }),
236        )
237        .await
238    }
239
240    /// Send initialized notification
241    pub async fn send_initialized(&mut self) -> Result<()> {
242        self.send_notification("notifications/initialized", serde_json::json!({}))
243            .await
244    }
245
246    /// List available tools
247    pub async fn list_tools(&mut self) -> Result<serde_json::Value> {
248        self.send_request("tools/list", serde_json::json!({})).await
249    }
250
251    /// Call a tool
252    pub async fn call_tool(
253        &mut self,
254        name: &str,
255        arguments: serde_json::Value,
256    ) -> Result<serde_json::Value> {
257        self.send_request(
258            "tools/call",
259            serde_json::json!({
260                "name": name,
261                "arguments": arguments
262            }),
263        )
264        .await
265    }
266
267    /// Check if the child process is still running
268    pub fn is_running(&mut self) -> bool {
269        matches!(self.child.try_wait(), Ok(None))
270    }
271
272    /// Gracefully shutdown the child process
273    pub async fn shutdown(mut self) -> Result<()> {
274        // Close stdin to signal EOF
275        drop(self.stdin);
276
277        // Wait for process to exit with timeout
278        let result =
279            tokio::time::timeout(std::time::Duration::from_secs(5), self.child.wait()).await;
280
281        match result {
282            Ok(Ok(status)) => {
283                tracing::info!(status = ?status, "Child process exited");
284                Ok(())
285            }
286            Ok(Err(e)) => {
287                tracing::error!(error = %e, "Error waiting for child process");
288                Err(Error::Transport(format!("Child process error: {}", e)))
289            }
290            Err(_) => {
291                // Timeout - kill the process
292                tracing::warn!("Child process did not exit gracefully, killing");
293                self.child
294                    .kill()
295                    .await
296                    .map_err(|e| Error::Transport(format!("Failed to kill child: {}", e)))?;
297                Ok(())
298            }
299        }
300    }
301
302    /// Kill the child process immediately
303    pub async fn kill(mut self) -> Result<()> {
304        self.child
305            .kill()
306            .await
307            .map_err(|e| Error::Transport(format!("Failed to kill child: {}", e)))?;
308        tracing::info!("Child process killed");
309        Ok(())
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[tokio::test]
318    async fn test_transport_builder() {
319        let transport = ChildProcessTransport::new("echo")
320            .arg("hello")
321            .env("FOO", "bar");
322
323        assert_eq!(transport.program, "echo");
324        assert_eq!(transport.args, vec!["hello"]);
325        assert_eq!(transport.envs, vec![("FOO".to_string(), "bar".to_string())]);
326    }
327
328    #[tokio::test]
329    async fn test_transport_args() {
330        let transport = ChildProcessTransport::new("cmd").args(["--flag1", "--flag2"]);
331
332        assert_eq!(transport.args, vec!["--flag1", "--flag2"]);
333    }
334
335    #[tokio::test]
336    async fn test_transport_env() {
337        let transport = ChildProcessTransport::new("prog")
338            .env("KEY1", "val1")
339            .env("KEY2", "val2");
340
341        assert_eq!(transport.envs.len(), 2);
342        assert_eq!(transport.envs[0], ("KEY1".to_string(), "val1".to_string()));
343    }
344
345    #[tokio::test]
346    async fn test_spawn_nonexistent_fails() {
347        let result = ChildProcessTransport::new("nonexistent-program-xyz-123")
348            .spawn()
349            .await;
350        assert!(result.is_err());
351    }
352
353    #[tokio::test]
354    async fn test_spawn_and_communicate() {
355        // Use `cat` as a simple echo server
356        let mut conn = ChildProcessTransport::new("cat").spawn().await.unwrap();
357
358        assert!(conn.is_running());
359
360        // Send a JSON-RPC request
361        let response = conn
362            .send_request("echo", serde_json::json!({"msg": "hello"}))
363            .await;
364
365        // cat will echo our request back, but it won't be a valid JSON-RPC response.
366        // That's OK - we're testing that I/O works, not protocol correctness.
367        // The response will be a parse error since cat echoes the request verbatim.
368        assert!(response.is_err());
369    }
370
371    #[tokio::test]
372    async fn test_shutdown_graceful() {
373        let conn = ChildProcessTransport::new("cat").spawn().await.unwrap();
374        // Shutdown should succeed - cat exits when stdin is closed
375        conn.shutdown().await.unwrap();
376    }
377
378    #[tokio::test]
379    async fn test_is_running_after_exit() {
380        // `true` exits immediately, but spawn plus exit plus reap is not
381        // instantaneous, especially on loaded Windows CI runners where a
382        // fixed 100 ms sleep flaked. Poll with a generous bound instead.
383        let mut conn = ChildProcessTransport::new("true").spawn().await.unwrap();
384        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
385        while conn.is_running() {
386            assert!(
387                std::time::Instant::now() < deadline,
388                "child process still reported running 10s after exit"
389            );
390            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
391        }
392    }
393
394    #[tokio::test]
395    async fn test_send_notification() {
396        let mut conn = ChildProcessTransport::new("cat").spawn().await.unwrap();
397        // Notification should succeed (no response expected)
398        conn.send_notification("test/notify", serde_json::json!({"data": 1}))
399            .await
400            .unwrap();
401        conn.shutdown().await.unwrap();
402    }
403}