From 9daa491f7c27a5513fec554473a7122d88fca367 Mon Sep 17 00:00:00 2001 From: Cam moten Date: Fri, 7 Aug 2026 03:25:45 +0000 Subject: [PATCH] Harden local MCP server process tree cleanup (#37366) ## Why A local stdio MCP server can exit while its descendants remain alive. Cleaning up only the server process can therefore leak processes after client shutdown. ## What changed - Contain locally launched MCP servers in non-breakaway Windows job objects and terminate the job during shutdown so descendants exit with the server. - Preserve compatibility when job creation or assignment is unavailable by retrying without containment and terminating an owned process handle instead. - Use the process-group member fallback when terminating piped processes on macOS. ## Testing - Cover descendant cleanup and breakaway prevention for both local MCP protocol modes on Windows. - Cover the Windows nested-job assignment fallback and verify the suspended process resumes before handle-based termination. GitOrigin-RevId: 33d3cc122de8f7ce0ae69a649b12ea0ff64b7c6f --- .../rmcp-client/src/bin/test_stdio_server.rs | 42 +++++ .../rmcp-client/src/stdio_server_launcher.rs | 167 ++++++++++++------ .../rmcp-client/tests/stdio_message_limits.rs | 138 +++++++++++++++ codex-rs/utils/pty/src/pipe.rs | 7 +- codex-rs/utils/pty/src/win/job.rs | 82 ++++++++- codex-rs/utils/pty/src/windows_tests.rs | 59 +++++++ 6 files changed, 440 insertions(+), 55 deletions(-) diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index d19f48156e10..5c6a105d8e20 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -900,10 +900,52 @@ fn parse_data_url(url: &str) -> Option<(String, String)> { #[tokio::main] async fn main() -> Result<(), Box> { + #[cfg(windows)] + if std::env::var_os("MCP_TEST_DESCENDANT_ROLE").is_some() { + tokio::time::sleep(Duration::from_secs(30)).await; + return Ok(()); + } + eprintln!("starting rmcp test server"); if let Ok(pid_file) = std::env::var("MCP_TEST_PID_FILE") { std::fs::write(pid_file, std::process::id().to_string())?; } + #[cfg(windows)] + if let Ok(marker_file) = std::env::var("MCP_TEST_BREAKAWAY_DENIED_FILE") { + use std::os::windows::process::CommandExt; + + const CREATE_BREAKAWAY_FROM_JOB: u32 = 0x0100_0000; + const ERROR_ACCESS_DENIED: i32 = 5; + + let escaped = std::process::Command::new(std::env::current_exe()?) + .creation_flags(CREATE_BREAKAWAY_FROM_JOB) + .env("MCP_TEST_DESCENDANT_ROLE", "1") + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn(); + match escaped { + Err(error) if error.raw_os_error() == Some(ERROR_ACCESS_DENIED) => { + std::fs::write(marker_file, "denied")?; + } + Err(error) => return Err(error.into()), + Ok(mut child) => { + let _ = child.kill(); + let _ = child.wait(); + return Err("MCP descendant unexpectedly escaped its Windows job".into()); + } + } + } + #[cfg(windows)] + if let Ok(pid_file) = std::env::var("MCP_TEST_DESCENDANT_PID_FILE") { + let child = std::process::Command::new(std::env::current_exe()?) + .env("MCP_TEST_DESCENDANT_ROLE", "1") + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn()?; + std::fs::write(pid_file, child.id().to_string())?; + } // Run the server with STDIO transport. If the client disconnects we simply // bubble up the error so the process exits. let service = TestToolServer::new(); diff --git a/codex-rs/rmcp-client/src/stdio_server_launcher.rs b/codex-rs/rmcp-client/src/stdio_server_launcher.rs index dcf15735f55f..9eecfadce7af 100644 --- a/codex-rs/rmcp-client/src/stdio_server_launcher.rs +++ b/codex-rs/rmcp-client/src/stdio_server_launcher.rs @@ -15,6 +15,8 @@ use std::collections::HashMap; use std::ffi::OsString; use std::future::Future; use std::io; +#[cfg(windows)] +use std::os::windows::io::OwnedHandle; use std::path::Path; use std::path::PathBuf; use std::process::Stdio; @@ -226,8 +228,9 @@ struct LocalProcessTerminator { } #[cfg(windows)] -struct LocalProcessTerminator { - pid: u32, +enum LocalProcessTerminator { + Job(codex_utils_pty::JobObject), + Process(OwnedHandle), } #[cfg(not(any(unix, windows)))] @@ -274,45 +277,103 @@ impl LocalStdioServerLauncher { let resolved_program = program_resolver::resolve(program, &envs, &cwd).map_err(io::Error::other)?; - let mut command = Command::new(resolved_program); - command - .kill_on_drop(true) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .current_dir(cwd) - .env_clear() - .envs(envs) - .args(args); - #[cfg(unix)] - command.process_group(0); - - let (transport, stderr, process_id) = match protocol_mode { - McpProtocolMode::Legacy => { - let (transport, stderr) = TokioChildProcess::builder(command) - .stderr(Stdio::piped()) - .spawn()?; - let process_id = transport.id(); - ( - StdioServerTransportInner::LocalLegacy(transport), - stderr, - process_id, - ) + let build_command = || { + let mut command = Command::new(&resolved_program); + command + .kill_on_drop(true) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .current_dir(&cwd) + .env_clear() + .envs(&envs) + .args(&args); + #[cfg(unix)] + command.process_group(0); + command + }; + #[cfg(windows)] + let mut command = build_command(); + #[cfg(not(windows))] + let command = build_command(); + #[cfg(windows)] + let job = match codex_utils_pty::JobObject::create_without_breakaway() { + Ok(job) => { + job.prepare_suspended_spawn(&mut command); + Some(job) } - McpProtocolMode::V20260728 => { - let (transport, stderr) = - LocalStdioTransport::spawn(command, program_name.clone())?; - let process_id = transport.id(); - ( - StdioServerTransportInner::LocalModern(transport), - stderr, - process_id, - ) + Err(error) => { + warn!("Windows MCP process job containment unavailable: {error}"); + None } }; - let process = StdioServerProcessHandle::local( - program_name.clone(), - process_id.map(LocalProcessTerminator::new), - ); + + let spawn_transport = |command: Command| -> io::Result<( + StdioServerTransportInner, + Option, + Option, + )> { + match protocol_mode { + McpProtocolMode::Legacy => { + let (transport, stderr) = TokioChildProcess::builder(command) + .stderr(Stdio::piped()) + .spawn()?; + let process_id = transport.id(); + Ok(( + StdioServerTransportInner::LocalLegacy(transport), + stderr, + process_id, + )) + } + McpProtocolMode::V20260728 => { + let (transport, stderr) = + LocalStdioTransport::spawn(command, program_name.clone())?; + let process_id = transport.id(); + Ok(( + StdioServerTransportInner::LocalModern(transport), + stderr, + process_id, + )) + } + } + }; + let (transport, stderr, process_id) = spawn_transport(command)?; + #[cfg(windows)] + let (transport, stderr, process_id, job) = match job { + Some(job) => match process_id + .ok_or_else(|| io::Error::other("missing suspended MCP server process id")) + .and_then(|process_id| job.assign_and_resume_process(process_id)) + { + Ok(true) => (transport, stderr, process_id, Some(job)), + Ok(false) => (transport, stderr, process_id, None), + Err(error) => { + warn!( + "Windows MCP process job containment failed; retrying without it: {error}" + ); + drop(stderr); + drop(transport); + drop(job); + let (transport, stderr, process_id) = spawn_transport(build_command())?; + (transport, stderr, process_id, None) + } + }, + None => (transport, stderr, process_id, None), + }; + #[cfg(windows)] + let terminator = match job { + Some(job) => Some(LocalProcessTerminator::Job(job)), + None => process_id.and_then(|process_id| { + match codex_utils_pty::JobObject::open_process_handle(process_id) { + Ok(handle) => Some(LocalProcessTerminator::Process(handle)), + Err(error) => { + warn!("Windows MCP process handle unavailable: {error}"); + None + } + } + }), + }; + #[cfg(not(windows))] + let terminator = process_id.map(LocalProcessTerminator::new); + let process = StdioServerProcessHandle::local(program_name.clone(), terminator); if let Some(stderr) = stderr { tokio::spawn(async move { @@ -340,17 +401,12 @@ impl LocalStdioServerLauncher { } impl LocalProcessTerminator { + #[cfg(not(windows))] fn new(process_group_id: u32) -> Self { #[cfg(unix)] { Self { process_group_id } } - #[cfg(windows)] - { - Self { - pid: process_group_id, - } - } #[cfg(not(any(unix, windows)))] { let _ = process_group_id; @@ -380,15 +436,15 @@ impl LocalProcessTerminator { #[cfg(windows)] fn terminate(&self) { - let _ = std::process::Command::new("taskkill") - .arg("/PID") - .arg(self.pid.to_string()) - .arg("/T") - .arg("/F") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status(); + let result = match self { + Self::Job(job) => job.terminate(), + Self::Process(process_handle) => { + codex_utils_pty::JobObject::terminate_process_handle(process_handle) + } + }; + if let Err(error) = result { + warn!("Failed to terminate Windows MCP process: {error}"); + } } #[cfg(not(any(unix, windows)))] @@ -479,6 +535,11 @@ impl Drop for StdioServerProcessHandleInner { /// MCP framing still runs in the orchestrator. The executor only owns the /// child process and transports raw stdin/stdout/stderr bytes, so it does not /// need to know about MCP methods such as `initialize` or `tools/list`. +/// +/// Windows executor-backed servers retain the executor's normal descendant +/// lifetime. MCP-specific containment requires negotiated process ownership: +/// caller-controlled process IDs cannot safely select a destructive policy, +/// and a wrapper may exit while its descendants continue serving requests. #[derive(Clone)] pub struct ExecutorStdioServerLauncher { exec_backend: Arc, diff --git a/codex-rs/rmcp-client/tests/stdio_message_limits.rs b/codex-rs/rmcp-client/tests/stdio_message_limits.rs index 0ddc4c3ee67e..7c8ec6102bbb 100644 --- a/codex-rs/rmcp-client/tests/stdio_message_limits.rs +++ b/codex-rs/rmcp-client/tests/stdio_message_limits.rs @@ -1,5 +1,11 @@ use std::collections::HashMap; use std::ffi::OsString; +#[cfg(windows)] +use std::os::windows::io::AsRawHandle; +#[cfg(windows)] +use std::os::windows::io::FromRawHandle; +#[cfg(windows)] +use std::os::windows::io::OwnedHandle; use std::sync::Arc; use std::time::Duration; @@ -17,6 +23,47 @@ use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; use rmcp::model::ProtocolVersion; +#[cfg(windows)] +#[link(name = "kernel32")] +unsafe extern "system" { + fn OpenProcess( + desired_access: u32, + inherit_handle: i32, + process_id: u32, + ) -> *mut std::ffi::c_void; + fn TerminateProcess(handle: *mut std::ffi::c_void, exit_code: u32) -> i32; + fn WaitForSingleObject(handle: *mut std::ffi::c_void, milliseconds: u32) -> u32; +} + +#[cfg(windows)] +fn open_process_for_wait(process_id: u32) -> std::io::Result { + let handle = unsafe { + OpenProcess( + /*desired_access*/ 0x0010_0001, + /*inherit_handle*/ 0, + process_id, + ) + }; + if handle.is_null() { + return Err(std::io::Error::last_os_error()); + } + Ok(unsafe { OwnedHandle::from_raw_handle(handle.cast()) }) +} + +#[cfg(windows)] +fn wait_for_process_exit(process: &OwnedHandle) -> std::io::Result<()> { + match unsafe { + WaitForSingleObject(process.as_raw_handle().cast(), /*milliseconds*/ 5_000) + } { + 0 => Ok(()), + u32::MAX => Err(std::io::Error::last_os_error()), + _ => Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "process did not exit", + )), + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn stdio_message_limits_preserve_legacy_local_compatibility() -> anyhow::Result<()> { let server = codex_utils_cargo_bin::cargo_bin("test_stdio_server")?; @@ -86,3 +133,94 @@ async fn stdio_message_limits_preserve_legacy_local_compatibility() -> anyhow::R } Ok(()) } + +#[cfg(windows)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn local_stdio_shutdown_terminates_descendants_after_server_exit() -> anyhow::Result<()> { + let server = codex_utils_cargo_bin::cargo_bin("test_stdio_server")?; + + for protocol_mode in [McpProtocolMode::Legacy, McpProtocolMode::V20260728] { + let temp_dir = tempfile::tempdir()?; + let server_pid_file = temp_dir.path().join("server.pid"); + let descendant_pid_file = temp_dir.path().join("descendant.pid"); + let breakaway_denied_file = temp_dir.path().join("breakaway.denied"); + let mut env = HashMap::from([ + ( + OsString::from("MCP_TEST_PID_FILE"), + server_pid_file.clone().into(), + ), + ( + OsString::from("MCP_TEST_DESCENDANT_PID_FILE"), + descendant_pid_file.clone().into(), + ), + ( + OsString::from("MCP_TEST_BREAKAWAY_DENIED_FILE"), + breakaway_denied_file.clone().into(), + ), + ]); + if protocol_mode == McpProtocolMode::V20260728 { + env.insert( + OsString::from("CODEX_MCP_PROTOCOL_VERSION"), + OsString::from("2026-07-28"), + ); + } + + let client = RmcpClient::new_stdio_client_with_protocol_mode( + server.clone().into(), + Vec::new(), + Some(env), + &[], + Some(std::env::current_dir()?.to_string_lossy().into_owned()), + Arc::new(LocalStdioServerLauncher::new(std::env::current_dir()?)), + protocol_mode, + ) + .await?; + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("stdio-cleanup-test", "1.0.0"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(10)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }) + } + .boxed() + }), + ) + .await?; + + assert_eq!( + std::fs::read_to_string(breakaway_denied_file)?, + "denied", + "MCP descendant escaped its Windows job (mode={protocol_mode:?})" + ); + let server_pid = std::fs::read_to_string(server_pid_file)? + .trim() + .parse::()?; + let descendant_pid = std::fs::read_to_string(descendant_pid_file)? + .trim() + .parse::()?; + let server_process = open_process_for_wait(server_pid)?; + let descendant_process = open_process_for_wait(descendant_pid)?; + let terminated = unsafe { + TerminateProcess(server_process.as_raw_handle().cast(), /*exit_code*/ 0) + }; + assert_ne!(terminated, 0, "failed to terminate test MCP server"); + wait_for_process_exit(&server_process)?; + + client.shutdown().await; + + assert!( + wait_for_process_exit(&descendant_process).is_ok(), + "MCP descendant survived its exited server (mode={protocol_mode:?})" + ); + } + Ok(()) +} diff --git a/codex-rs/utils/pty/src/pipe.rs b/codex-rs/utils/pty/src/pipe.rs index b256d2b99756..3715ba72a404 100644 --- a/codex-rs/utils/pty/src/pipe.rs +++ b/codex-rs/utils/pty/src/pipe.rs @@ -62,11 +62,16 @@ impl ChildTerminator for PipeChildTerminator { } fn kill(&mut self) -> io::Result<()> { - #[cfg(unix)] + #[cfg(all(unix, not(target_os = "macos")))] { crate::process_group::kill_process_group(self.process_group_id) } + #[cfg(target_os = "macos")] + { + crate::process_group::kill_process_group_with_member_fallback(self.process_group_id) + } + #[cfg(windows)] { match &self.windows { diff --git a/codex-rs/utils/pty/src/win/job.rs b/codex-rs/utils/pty/src/win/job.rs index f22152eb78a5..b06035f72d22 100644 --- a/codex-rs/utils/pty/src/win/job.rs +++ b/codex-rs/utils/pty/src/win/job.rs @@ -12,12 +12,17 @@ use winapi::um::jobapi2::AssignProcessToJobObject; use winapi::um::jobapi2::CreateJobObjectW; use winapi::um::jobapi2::SetInformationJobObject; use winapi::um::jobapi2::TerminateJobObject; +use winapi::um::processthreadsapi::OpenProcess; +use winapi::um::processthreadsapi::TerminateProcess; use winapi::um::winbase::CREATE_SUSPENDED; use winapi::um::winnt::HANDLE; use winapi::um::winnt::JOB_OBJECT_LIMIT_BREAKAWAY_OK; use winapi::um::winnt::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; use winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION; use winapi::um::winnt::JobObjectExtendedLimitInformation; +use winapi::um::winnt::PROCESS_SET_QUOTA; +use winapi::um::winnt::PROCESS_SUSPEND_RESUME; +use winapi::um::winnt::PROCESS_TERMINATE; #[link(name = "ntdll")] unsafe extern "system" { @@ -53,6 +58,37 @@ impl JobObject { }) } + /// Creates a Job Object whose owned descendants cannot explicitly break away. + pub fn create_without_breakaway() -> io::Result { + let job = Self::create()?; + Self::set_limit_flags(&job.handle, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE)?; + Ok(job) + } + + /// Captures an owned process handle before its numeric identifier can be reused. + pub fn open_process_handle(process_id: u32) -> io::Result { + let handle = unsafe { + OpenProcess(PROCESS_TERMINATE, /*bInheritHandle*/ 0, process_id) + }; + if handle.is_null() { + return Err(io::Error::last_os_error()); + } + + Ok(unsafe { std::os::windows::io::OwnedHandle::from_raw_handle(handle.cast()) }) + } + + /// Terminates the exact process identified by a previously captured handle. + pub fn terminate_process_handle(handle: &std::os::windows::io::OwnedHandle) -> io::Result<()> { + let terminated = unsafe { + TerminateProcess(handle.as_raw_handle().cast(), /*uExitCode*/ 1) + }; + if terminated == 0 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + } + fn set_limit_flags(handle: &OwnedHandle, flags: u32) -> io::Result<()> { let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() }; limits.BasicLimitInformation.LimitFlags = flags; @@ -85,9 +121,53 @@ impl JobObject { } } + /// Prevents a child from running before it can be assigned to this job. + pub fn prepare_suspended_spawn(&self, command: &mut Command) { + command.creation_flags(CREATE_SUSPENDED).kill_on_drop(true); + } + + /// Assigns and resumes a suspended child, returning whether assignment succeeded. + /// + /// Nested jobs can reject assignment. Such a child is resumed without + /// containment so callers can preserve their existing compatibility fallback. + pub fn assign_and_resume_process(&self, process_id: u32) -> io::Result { + let process = unsafe { + OpenProcess( + PROCESS_SET_QUOTA | PROCESS_TERMINATE | PROCESS_SUSPEND_RESUME, + /*bInheritHandle*/ 0, + process_id, + ) + }; + if process.is_null() { + return Err(io::Error::last_os_error()); + } + let process = unsafe { OwnedHandle::from_raw_handle(process.cast()) }; + let assignment = self.assign_process(process.as_raw_handle()); + + let status = unsafe { NtResumeProcess(process.as_raw_handle().cast()) }; + if !NT_SUCCESS(status) { + unsafe { + TerminateProcess(process.as_raw_handle().cast(), /*uExitCode*/ 1) + }; + return Err(io::Error::other(format!( + "failed to resume suspended process: NTSTATUS {status:#x}" + ))); + } + + match assignment { + Ok(()) => Ok(true), + Err(error) => { + log::warn!( + "Windows process job assignment unavailable for pid {process_id}: {error}" + ); + Ok(false) + } + } + } + /// Starts a process only after assigning it to this Job Object. pub fn spawn_contained(&self, command: &mut Command) -> io::Result { - command.creation_flags(CREATE_SUSPENDED).kill_on_drop(true); + self.prepare_suspended_spawn(command); let child = command.spawn()?; let process_handle = child .raw_handle() diff --git a/codex-rs/utils/pty/src/windows_tests.rs b/codex-rs/utils/pty/src/windows_tests.rs index 91322be1cf43..733635b7f3f8 100644 --- a/codex-rs/utils/pty/src/windows_tests.rs +++ b/codex-rs/utils/pty/src/windows_tests.rs @@ -241,6 +241,65 @@ async fn contained_spawn_owns_immediate_descendant() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rejected_job_assignment_resumes_existing_job_member() -> anyhow::Result<()> { + let Some(python) = find_python() else { + eprintln!("python not found; skipping Windows nested-job fallback test"); + return Ok(()); + }; + + let owning_job = crate::JobObject::create()?; + let rejected_job = crate::JobObject::create_without_breakaway()?; + let mut occupied_command = Command::new(&python); + occupied_command + .args(["-c", "import time; time.sleep(60)"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + let mut existing_member = rejected_job.spawn_contained(&mut occupied_command)?; + + let mut command = Command::new(&python); + command + .args([ + "-u", + "-c", + "import time; print('resumed',flush=True); time.sleep(60)", + ]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()); + rejected_job.prepare_suspended_spawn(&mut command); + let mut root = command.spawn()?; + let process_handle = root + .raw_handle() + .ok_or_else(|| anyhow::anyhow!("missing suspended process handle"))?; + owning_job.assign_process(process_handle)?; + let process_id = root + .id() + .ok_or_else(|| anyhow::anyhow!("missing suspended process id"))?; + + assert!( + !rejected_job.assign_and_resume_process(process_id)?, + "unrelated nested job unexpectedly accepted the process" + ); + let stdout = root + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("missing resumed process stdout"))?; + let mut stdout = BufReader::new(stdout); + let mut marker = String::new(); + tokio::time::timeout(Duration::from_secs(10), stdout.read_line(&mut marker)).await??; + assert_eq!(marker.trim(), "resumed"); + + let process_handle = crate::JobObject::open_process_handle(process_id)?; + crate::JobObject::terminate_process_handle(&process_handle)?; + rejected_job.terminate()?; + let status = tokio::time::timeout(Duration::from_secs(10), root.wait()).await??; + assert_eq!(status.code(), Some(1)); + tokio::time::timeout(Duration::from_secs(10), existing_member.wait()).await??; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn conpty_delivers_input_to_foreground_children() -> anyhow::Result<()> { let Some(python) = find_python() else {