1use 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
48pub struct ChildProcessTransport {
50 program: String,
51 args: Vec<String>,
52 envs: Vec<(String, String)>,
53}
54
55impl ChildProcessTransport {
56 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 pub fn arg(mut self, arg: impl Into<String>) -> Self {
67 self.args.push(arg.into());
68 self
69 }
70
71 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 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 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
109pub 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 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 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 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 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(¬ification)
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 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 pub async fn send_initialized(&mut self) -> Result<()> {
242 self.send_notification("notifications/initialized", serde_json::json!({}))
243 .await
244 }
245
246 pub async fn list_tools(&mut self) -> Result<serde_json::Value> {
248 self.send_request("tools/list", serde_json::json!({})).await
249 }
250
251 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 pub fn is_running(&mut self) -> bool {
269 matches!(self.child.try_wait(), Ok(None))
270 }
271
272 pub async fn shutdown(mut self) -> Result<()> {
274 drop(self.stdin);
276
277 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 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 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 let mut conn = ChildProcessTransport::new("cat").spawn().await.unwrap();
357
358 assert!(conn.is_running());
359
360 let response = conn
362 .send_request("echo", serde_json::json!({"msg": "hello"}))
363 .await;
364
365 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 conn.shutdown().await.unwrap();
376 }
377
378 #[tokio::test]
379 async fn test_is_running_after_exit() {
380 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 conn.send_notification("test/notify", serde_json::json!({"data": 1}))
399 .await
400 .unwrap();
401 conn.shutdown().await.unwrap();
402 }
403}