Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 3 additions & 10 deletions crates/goose/src/agents/extension_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,10 @@ use rmcp::service::{ClientInitializeError, ServiceError};
use rmcp::transport::streamable_http_client::{
StreamableHttpClientTransportConfig, StreamableHttpError,
};
use rmcp::transport::{
ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport, TokioChildProcess,
};
use rmcp::transport::{ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport};
use std::collections::HashMap;
use std::path::PathBuf;
use std::pin::Pin;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Weak};
use std::task::{Context, Poll};
Expand Down Expand Up @@ -47,7 +44,7 @@ use crate::config::search_path::SearchPaths;
use crate::config::{get_all_extensions, Config};
use crate::oauth::{oauth_flow, GooseCredentialStore, StaticOAuthClientConfig};
use crate::prompt_template;
use crate::subprocess::configure_subprocess;
use crate::subprocess::spawn_long_lived_mcp_subprocess;
use rmcp::model::{
CallToolRequestParams, CallToolResult, ContentBlock, ErrorCode, ErrorData, GetPromptResult,
MetaObject, Prompt, Resource, ResourceContents, ServerInfo, ServerNotification, Tool,
Expand Down Expand Up @@ -420,8 +417,6 @@ async fn child_process_client(
action_required: Arc<ActionRequiredManager>,
extension_manager: Weak<ExtensionManager>,
) -> ExtensionResult<McpClient> {
configure_subprocess(&mut command);

if let Ok(path) = SearchPaths::builder().path() {
command.env("PATH", path);
}
Expand All @@ -436,9 +431,7 @@ async fn child_process_client(
);
}

let (transport, mut stderr) = TokioChildProcess::builder(command)
.stderr(Stdio::piped())
.spawn()?;
let (transport, mut stderr) = spawn_long_lived_mcp_subprocess(command).await?;
let mut stderr = stderr.take().ok_or_else(|| {
ExtensionError::SetupError("failed to attach child process stderr".to_owned())
})?;
Expand Down
78 changes: 75 additions & 3 deletions crates/goose/src/subprocess.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
use rmcp::transport::TokioChildProcess;
use std::io;
#[cfg(target_os = "linux")]
use std::sync::{mpsc, OnceLock};
use tokio::process::ChildStderr;
use tokio::process::Command;

#[cfg(windows)]
Expand Down Expand Up @@ -60,13 +65,80 @@ impl SubprocessExt for std::process::Command {
}
}

#[allow(unused_variables)]
pub fn configure_subprocess(command: &mut Command) {
fn configure_common_subprocess(command: &mut Command) {
// Isolate subprocess into its own process group so it does not receive
// SIGINT when the user presses Ctrl+C in the terminal.
#[cfg(unix)]
command.process_group(0);
command.set_no_window();
}

#[allow(unused_variables)]
pub fn configure_subprocess(command: &mut Command) {
configure_common_subprocess(command);
#[cfg(target_os = "linux")]
configure_parent_death_signal(command);
command.set_no_window();
}

#[cfg(target_os = "linux")]
struct LongLivedSpawnRequest {
command: Command,
runtime: tokio::runtime::Handle,
response: tokio::sync::oneshot::Sender<io::Result<(TokioChildProcess, Option<ChildStderr>)>>,
}

#[cfg(target_os = "linux")]
fn long_lived_spawn_sender() -> io::Result<mpsc::Sender<LongLivedSpawnRequest>> {
static SENDER: OnceLock<io::Result<mpsc::Sender<LongLivedSpawnRequest>>> = OnceLock::new();

match SENDER.get_or_init(|| {
let (sender, receiver) = mpsc::channel::<LongLivedSpawnRequest>();
std::thread::Builder::new()
.name("goose-extension-spawner".to_owned())
.spawn(move || {
while let Ok(mut request) = receiver.recv() {
let _runtime_guard = request.runtime.enter();
configure_subprocess(&mut request.command);
let result = TokioChildProcess::builder(request.command)
.stderr(std::process::Stdio::piped())
.spawn();
let _ = request.response.send(result);
}
})
.map(|_| sender)
}) {
Ok(sender) => Ok(sender.clone()),
Err(error) => Err(io::Error::new(error.kind(), error.to_string())),
}
}

/// Spawn a long-lived MCP subprocess without tying Linux parent-death cleanup
/// to the Tokio worker that happened to request it.
pub async fn spawn_long_lived_mcp_subprocess(
command: Command,
) -> io::Result<(TokioChildProcess, Option<ChildStderr>)> {
#[cfg(target_os = "linux")]
{
let runtime = tokio::runtime::Handle::try_current().map_err(io::Error::other)?;
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
long_lived_spawn_sender()?
.send(LongLivedSpawnRequest {
command,
runtime,
response: response_tx,
})
.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "extension spawner exited"))?;
response_rx
.await
.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "extension spawner exited"))?
}

#[cfg(not(target_os = "linux"))]
{
let mut command = command;
configure_subprocess(&mut command);
TokioChildProcess::builder(command)
.stderr(std::process::Stdio::piped())
.spawn()
}
}
127 changes: 125 additions & 2 deletions crates/goose/tests/subprocess_cleanup.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,30 @@
#![cfg(target_os = "linux")]

use goose::subprocess::configure_subprocess;
use goose::subprocess::{configure_subprocess, spawn_long_lived_mcp_subprocess};
use std::io::{BufRead, BufReader, Write};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::process::{Child, Command, Stdio};
use std::sync::mpsc;
use std::time::{Duration, Instant};

const HELPER_ENV: &str = "GOOSE_SUBPROCESS_PARENT_DEATH_HELPER";
const THREAD_HELPER_ENV: &str = "GOOSE_SUBPROCESS_THREAD_DEATH_HELPER";

struct HelperProcess(Child);

impl Drop for HelperProcess {
fn drop(&mut self) {
let _ = self.0.kill();
let _ = self.0.wait();
}
}

#[ctor::ctor]
unsafe fn maybe_run_helper() {
if std::env::var_os(THREAD_HELPER_ENV).is_some() {
run_thread_death_helper();
}

if std::env::var_os(HELPER_ENV).is_none() {
return;
}
Expand Down Expand Up @@ -41,6 +56,49 @@ unsafe fn maybe_run_helper() {
}
}

fn run_thread_death_helper() {
let (tx, rx) = mpsc::channel();
let spawn_thread = std::thread::spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime");

let pid = runtime.block_on(async {
let mut command = tokio::process::Command::new("sleep");
command.arg("30");
command.stdin(Stdio::null());
command.stdout(Stdio::null());
command.stderr(Stdio::null());
let (child, _) = spawn_long_lived_mcp_subprocess(command)
.await
.expect("spawn child");
let pid = child.id().expect("child pid");
std::mem::forget(child);
pid
});

tx.send(pid).expect("send child pid");
});

spawn_thread.join().expect("spawn thread");
let child_pid = rx.recv().expect("child pid");
std::thread::sleep(Duration::from_millis(500));

if !process_is_running(child_pid) {
eprintln!("child process {child_pid} exited after spawning thread exit");
unsafe {
libc::_exit(1);
}
}

println!("{child_pid}");
std::io::stdout().flush().expect("flush pid");
loop {
std::thread::park();
}
}

#[test]
fn child_process_exits_when_parent_process_dies() {
let current_exe = std::env::current_exe().expect("current test binary");
Expand Down Expand Up @@ -74,6 +132,71 @@ fn child_process_exits_when_parent_process_dies() {
}
}

#[test]
fn long_lived_child_process_survives_spawning_thread_exit() {
let current_exe = std::env::current_exe().expect("current test binary");
let mut helper = HelperProcess(
Command::new(current_exe)
.env(THREAD_HELPER_ENV, "1")
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.expect("spawn helper"),
);

let stdout = helper.0.stdout.take().expect("helper stdout");
let (pid_tx, pid_rx) = mpsc::channel();
std::thread::spawn(move || {
let pid = BufReader::new(stdout)
.lines()
.next()
.ok_or("helper exited without reporting a child pid")
.and_then(|line| line.map_err(|_| "failed to read helper child pid"))
.and_then(|line| line.parse::<u32>().map_err(|_| "invalid helper child pid"));
let _ = pid_tx.send(pid);
});
let child_pid = pid_rx
.recv_timeout(Duration::from_secs(5))
.expect("timed out waiting for helper child pid")
.expect("helper child pid");

assert!(
process_is_running(child_pid),
"child process {child_pid} exited after spawning thread exit"
);

unsafe {
libc::kill(helper.0.id() as libc::pid_t, libc::SIGKILL);
}
let status = helper.0.wait().expect("wait for helper");
assert!(
!status.success(),
"helper should have been killed: {status}"
);

let deadline = Instant::now() + Duration::from_secs(5);
while process_is_running(child_pid) && Instant::now() < deadline {
std::thread::sleep(Duration::from_millis(100));
}
assert!(
!process_is_running(child_pid),
"child process {child_pid} survived parent process death"
);
}

fn process_exists(pid: u32) -> bool {
PathBuf::from(format!("/proc/{pid}")).exists()
}

fn process_is_running(pid: u32) -> bool {
match process_state(pid) {
Some('Z') | None => false,
Some(_) => true,
}
}

fn process_state(pid: u32) -> Option<char> {
let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
let (_, after_name) = stat.rsplit_once(") ")?;
after_name.chars().next()
}