fix: make stdio shutdown more graceful

According to the protocol specifications

Signed-off-by: jokemanfire <[email protected]>

Author: jokemanfire

crates/rmcp/src/transport/child_process.rs

@@ -9,6 +9,7 @@ use tokio::{
 use super::{IntoTransport, Transport};
 use crate::service::ServiceRole;
 
+const MAX_WAIT_ON_DROP_SECS: u64 = 3;
 /// The parts of a child process.
 type ChildProcessParts = (
     Box<dyn TokioChildWrapper>,
@@ -41,13 +42,35 @@ pub struct TokioChildProcess {
 }
 
 pub struct ChildWithCleanup {
-    inner: Box<dyn TokioChildWrapper>,
+    inner: Option<Box<dyn TokioChildWrapper>>,
 }
 
 impl Drop for ChildWithCleanup {
     fn drop(&mut self) {
-        if let Err(e) = self.inner.start_kill() {
-            tracing::warn!("Failed to kill child process: {e}");
+        // Close child more graceful
+        if let Some(mut child) = self.inner.take() {
+            // Spawn a background task to clean up
+            tokio::spawn(async move {
+                // Wait will drop the stdin
+                let wait_fut = Box::into_pin(child.wait());
+                tokio::select! {
+                    _ = tokio::time::sleep(std::time::Duration::from_secs(MAX_WAIT_ON_DROP_SECS)) => {
+                        if let Err(e) = Box::into_pin(child.kill()).await{
+                            tracing::warn!("Error killing child: {e}");
+                        }
+                    },
+                    res = wait_fut => {
+                        match res {
+                            Ok(status) => {
+                                tracing::info!("Child exited gracefully {}", status);
+                            }
+                            Err(e) => {
+                                tracing::warn!("Error waiting for child: {e}");
+                            }
+                        }
+                    }
+                }
+            });
         }
     }
 }
@@ -64,7 +87,7 @@ pin_project_lite::pin_project! {
 impl TokioChildProcessOut {
     /// Get the process ID of the child process.
     pub fn id(&self) -> Option<u32> {
-        self.child.inner.id()
+        self.child.inner.as_ref()?.id()
     }
 }
 
@@ -92,7 +115,7 @@ impl TokioChildProcess {
 
     /// Get the process ID of the child process.
     pub fn id(&self) -> Option<u32> {
-        self.child.inner.id()
+        self.child.inner.as_ref()?.id()
     }
 
     /// Split this helper into a reader (stdout) and writer (stdin).
@@ -157,7 +180,7 @@ impl TokioChildProcessBuilder {
         let (child, stdout, stdin, stderr_opt) = child_process(self.cmd.spawn()?)?;
 
         let proc = TokioChildProcess {
-            child: ChildWithCleanup { inner: child },
+            child: ChildWithCleanup { inner: Some(child) },
             child_stdin: stdin,
             child_stdout: stdout,
         };
@@ -183,3 +206,45 @@ impl ConfigureCommandExt for tokio::process::Command {
         self
     }
 }
+
+#[cfg(unix)]
+#[cfg(test)]
+mod tests {
+    use tokio::process::Command;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_tokio_child_process_drop() {
+        let r = TokioChildProcess::new(Command::new("sleep").configure(|cmd| {
+            cmd.arg("30");
+        }));
+        assert!(r.is_ok());
+        let child_process = r.unwrap();
+        let id = child_process.id();
+        assert!(id.is_some());
+        let id = id.unwrap();
+        // Drop the child process
+        drop(child_process);
+        // Wait a moment to allow the cleanup task to run
+        tokio::time::sleep(std::time::Duration::from_secs(MAX_WAIT_ON_DROP_SECS + 1)).await;
+        // Check if the process is still running
+        let status = Command::new("ps")
+            .arg("-p")
+            .arg(id.to_string())
+            .status()
+            .await;
+        match status {
+            Ok(status) => {
+                assert!(
+                    !status.success(),
+                    "Process with PID {} is still running",
+                    id
+                );
+            }
+            Err(e) => {
+                panic!("Failed to check process status: {}", e);
+            }
+        }
+    }
+}