mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Process-group cleanup for stdio MCP servers to prevent orphan process storms (#10710)
This PR changes stdio MCP child processes to run in their own process group * Add guarded teardown in codex-rmcp-client: send SIGTERM to the group first, then SIGKILL after a short grace period. * Add terminate_process_group helper in process_group.rs. * Add Unix regression test in process_group_cleanup.rs to verify wrapper + grandchild are reaped on client drop. Addresses reported MCP process/thread storm: #10581
This commit is contained in:
committed by
GitHub
Unverified
parent
4d52428fa2
commit
82c981cafc
Generated
+1
@@ -2007,6 +2007,7 @@ dependencies = [
|
||||
"codex-protocol",
|
||||
"codex-utils-cargo-bin",
|
||||
"codex-utils-home-dir",
|
||||
"codex-utils-pty",
|
||||
"futures",
|
||||
"keyring",
|
||||
"oauth2",
|
||||
|
||||
@@ -15,6 +15,7 @@ axum = { workspace = true, default-features = false, features = [
|
||||
] }
|
||||
codex-keyring-store = { workspace = true }
|
||||
codex-protocol = { workspace = true }
|
||||
codex-utils-pty = { workspace = true }
|
||||
codex-utils-home-dir = { workspace = true }
|
||||
futures = { workspace = true, default-features = false, features = ["std"] }
|
||||
keyring = { workspace = true, features = ["crypto-rust"] }
|
||||
|
||||
@@ -60,7 +60,10 @@ use crate::utils::create_env_for_mcp_server;
|
||||
use crate::utils::run_with_timeout;
|
||||
|
||||
enum PendingTransport {
|
||||
ChildProcess(TokioChildProcess),
|
||||
ChildProcess {
|
||||
transport: TokioChildProcess,
|
||||
process_group_guard: Option<ProcessGroupGuard>,
|
||||
},
|
||||
StreamableHttp {
|
||||
transport: StreamableHttpClientTransport<reqwest::Client>,
|
||||
},
|
||||
@@ -75,11 +78,71 @@ enum ClientState {
|
||||
transport: Option<PendingTransport>,
|
||||
},
|
||||
Ready {
|
||||
_process_group_guard: Option<ProcessGroupGuard>,
|
||||
service: Arc<RunningService<RoleClient, LoggingClientHandler>>,
|
||||
oauth: Option<OAuthPersistor>,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
const PROCESS_GROUP_TERM_GRACE_PERIOD: Duration = Duration::from_secs(2);
|
||||
|
||||
#[cfg(unix)]
|
||||
struct ProcessGroupGuard {
|
||||
process_group_id: u32,
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
struct ProcessGroupGuard;
|
||||
|
||||
impl ProcessGroupGuard {
|
||||
fn new(process_group_id: u32) -> Self {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
Self { process_group_id }
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = process_group_id;
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn maybe_terminate_process_group(&self) {
|
||||
let process_group_id = self.process_group_id;
|
||||
let should_escalate =
|
||||
match codex_utils_pty::process_group::terminate_process_group(process_group_id) {
|
||||
Ok(exists) => exists,
|
||||
Err(error) => {
|
||||
warn!("Failed to terminate MCP process group {process_group_id}: {error}");
|
||||
false
|
||||
}
|
||||
};
|
||||
if should_escalate {
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(PROCESS_GROUP_TERM_GRACE_PERIOD);
|
||||
if let Err(error) =
|
||||
codex_utils_pty::process_group::kill_process_group(process_group_id)
|
||||
{
|
||||
warn!("Failed to kill MCP process group {process_group_id}: {error}");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn maybe_terminate_process_group(&self) {}
|
||||
}
|
||||
|
||||
impl Drop for ProcessGroupGuard {
|
||||
fn drop(&mut self) {
|
||||
if cfg!(unix) {
|
||||
self.maybe_terminate_process_group();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub type Elicitation = CreateElicitationRequestParam;
|
||||
pub type ElicitationResponse = CreateElicitationResult;
|
||||
|
||||
@@ -129,6 +192,8 @@ impl RmcpClient {
|
||||
.env_clear()
|
||||
.envs(envs)
|
||||
.args(&args);
|
||||
#[cfg(unix)]
|
||||
command.process_group(0);
|
||||
if let Some(cwd) = cwd {
|
||||
command.current_dir(cwd);
|
||||
}
|
||||
@@ -136,6 +201,7 @@ impl RmcpClient {
|
||||
let (transport, stderr) = TokioChildProcess::builder(command)
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()?;
|
||||
let process_group_guard = transport.id().map(ProcessGroupGuard::new);
|
||||
|
||||
if let Some(stderr) = stderr {
|
||||
tokio::spawn(async move {
|
||||
@@ -157,7 +223,10 @@ impl RmcpClient {
|
||||
|
||||
Ok(Self {
|
||||
state: Mutex::new(ClientState::Connecting {
|
||||
transport: Some(PendingTransport::ChildProcess(transport)),
|
||||
transport: Some(PendingTransport::ChildProcess {
|
||||
transport,
|
||||
process_group_guard,
|
||||
}),
|
||||
}),
|
||||
})
|
||||
}
|
||||
@@ -226,17 +295,22 @@ impl RmcpClient {
|
||||
) -> Result<InitializeResult> {
|
||||
let client_handler = LoggingClientHandler::new(params.clone(), send_elicitation);
|
||||
|
||||
let (transport, oauth_persistor) = {
|
||||
let (transport, oauth_persistor, process_group_guard) = {
|
||||
let mut guard = self.state.lock().await;
|
||||
match &mut *guard {
|
||||
ClientState::Connecting { transport } => match transport.take() {
|
||||
Some(PendingTransport::ChildProcess(transport)) => (
|
||||
Some(PendingTransport::ChildProcess {
|
||||
transport,
|
||||
process_group_guard,
|
||||
}) => (
|
||||
service::serve_client(client_handler.clone(), transport).boxed(),
|
||||
None,
|
||||
process_group_guard,
|
||||
),
|
||||
Some(PendingTransport::StreamableHttp { transport }) => (
|
||||
service::serve_client(client_handler.clone(), transport).boxed(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
Some(PendingTransport::StreamableHttpWithOAuth {
|
||||
transport,
|
||||
@@ -244,6 +318,7 @@ impl RmcpClient {
|
||||
}) => (
|
||||
service::serve_client(client_handler.clone(), transport).boxed(),
|
||||
Some(oauth_persistor),
|
||||
None,
|
||||
),
|
||||
None => return Err(anyhow!("client already initializing")),
|
||||
},
|
||||
@@ -270,6 +345,7 @@ impl RmcpClient {
|
||||
{
|
||||
let mut guard = self.state.lock().await;
|
||||
*guard = ClientState::Ready {
|
||||
_process_group_guard: process_group_guard,
|
||||
service: Arc::new(service),
|
||||
oauth: oauth_persistor.clone(),
|
||||
};
|
||||
@@ -448,7 +524,7 @@ impl RmcpClient {
|
||||
match &*guard {
|
||||
ClientState::Ready {
|
||||
oauth: Some(runtime),
|
||||
service: _,
|
||||
..
|
||||
} => Some(runtime.clone()),
|
||||
_ => None,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
#![cfg(unix)]
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use codex_rmcp_client::RmcpClient;
|
||||
|
||||
fn process_exists(pid: u32) -> bool {
|
||||
std::process::Command::new("kill")
|
||||
.arg("-0")
|
||||
.arg(pid.to_string())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.status()
|
||||
.map(|status| status.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn wait_for_pid_file(path: &Path) -> Result<u32> {
|
||||
for _ in 0..50 {
|
||||
match fs::read_to_string(path) {
|
||||
Ok(content) => {
|
||||
let pid = content
|
||||
.trim()
|
||||
.parse::<u32>()
|
||||
.with_context(|| format!("failed to parse pid from {}", path.display()))?;
|
||||
return Ok(pid);
|
||||
}
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
Err(error) => {
|
||||
return Err(error).with_context(|| format!("failed to read {}", path.display()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
anyhow::bail!("timed out waiting for child pid file at {}", path.display());
|
||||
}
|
||||
|
||||
async fn wait_for_process_exit(pid: u32) -> Result<()> {
|
||||
for _ in 0..50 {
|
||||
if !process_exists(pid) {
|
||||
return Ok(());
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
anyhow::bail!("process {pid} still running after timeout");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn drop_kills_wrapper_process_group() -> Result<()> {
|
||||
let temp_dir = tempfile::tempdir()?;
|
||||
let child_pid_file = temp_dir.path().join("child.pid");
|
||||
let child_pid_file_str = child_pid_file.to_string_lossy().into_owned();
|
||||
|
||||
let client = RmcpClient::new_stdio_client(
|
||||
OsString::from("/bin/sh"),
|
||||
vec![
|
||||
OsString::from("-c"),
|
||||
OsString::from(
|
||||
"sleep 300 & child_pid=$!; echo \"$child_pid\" > \"$CHILD_PID_FILE\"; cat >/dev/null",
|
||||
),
|
||||
],
|
||||
Some(HashMap::from([(
|
||||
"CHILD_PID_FILE".to_string(),
|
||||
child_pid_file_str,
|
||||
)])),
|
||||
&[],
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let grandchild_pid = wait_for_pid_file(&child_pid_file).await?;
|
||||
assert!(
|
||||
process_exists(grandchild_pid),
|
||||
"expected grandchild process {grandchild_pid} to be running before dropping client"
|
||||
);
|
||||
|
||||
drop(client);
|
||||
|
||||
wait_for_process_exit(grandchild_pid).await
|
||||
}
|
||||
@@ -117,6 +117,33 @@ pub fn kill_process_group_by_pid(_pid: u32) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
/// Send SIGTERM to a specific process group ID (best-effort).
|
||||
///
|
||||
/// Returns `Ok(true)` when SIGTERM was delivered to an existing group and
|
||||
/// `Ok(false)` when the group no longer exists.
|
||||
pub fn terminate_process_group(process_group_id: u32) -> io::Result<bool> {
|
||||
use std::io::ErrorKind;
|
||||
|
||||
let pgid = process_group_id as libc::pid_t;
|
||||
let result = unsafe { libc::killpg(pgid, libc::SIGTERM) };
|
||||
if result == -1 {
|
||||
let err = io::Error::last_os_error();
|
||||
if err.kind() == ErrorKind::NotFound {
|
||||
return Ok(false);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
/// No-op on non-Unix platforms.
|
||||
pub fn terminate_process_group(_process_group_id: u32) -> io::Result<bool> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
/// Kill a specific process group ID (best-effort).
|
||||
pub fn kill_process_group(process_group_id: u32) -> io::Result<()> {
|
||||
|
||||
Reference in New Issue
Block a user