diff --git a/.github/workflows/foreground-cancellation.yml b/.github/workflows/foreground-cancellation.yml new file mode 100644 index 00000000..f2b3bef1 --- /dev/null +++ b/.github/workflows/foreground-cancellation.yml @@ -0,0 +1,98 @@ +name: Foreground cancellation qualification +on: + push: + branches: [fix/foreground-command-cancellation] + pull_request: + paths: + - 'libs/mcp/server/**' + - '.github/workflows/foreground-cancellation.yml' +permissions: + contents: read +jobs: + linux-qualification: + timeout-minutes: 30 + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@master + with: + toolchain: '1.94.1' + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + with: + cache-on-failure: true + - name: Generate formatting receipt + run: | + cargo fmt --all + git diff > formatting.patch + - uses: actions/upload-artifact@v4 + with: + name: formatting-${{ github.sha }} + path: formatting.patch + - run: cargo test --locked -p stakpak-mcp-server foreground_ -- --nocapture + - name: Stock executor must fail the cancellation regression + shell: bash + run: | + git fetch origin main --depth=1 + git show FETCH_HEAD:libs/mcp/server/src/local_tools.rs > libs/mcp/server/src/local_tools.rs + python3 - <<'PYTHON' + from pathlib import Path + p = Path('libs/mcp/server/src/local_tools.rs') + s = p.read_text().replace(' fn local_container_with_profile(', ' pub(super) fn local_container_with_profile(') + s += '\n#[cfg(all(test, target_os = "linux"))]\n#[path = "foreground_tests.rs"]\nmod foreground_tests;\n' + p.write_text(s) + PYTHON + cargo test --locked -p stakpak-mcp-server --no-run + if cargo test --locked -p stakpak-mcp-server foreground_cancel_without_timeout -- --nocapture; then + echo 'Stock executor unexpectedly passed; regression does not demonstrate the defect' + exit 1 + fi + git restore libs/mcp/server/src/local_tools.rs + cargo fmt --all + - run: cargo clippy --locked -p stakpak-mcp-server --all-targets -- -D warnings + - run: cargo build --locked --release -p stakpak + - run: git diff --exit-code + - name: Record prerelease artifact provenance + run: | + git rev-parse HEAD > target/release/source-commit.txt + rustc --version > target/release/rust-version.txt + sha256sum target/release/stakpak > target/release/stakpak.sha256 + - uses: actions/upload-artifact@v4 + with: + name: stakpak-linux-prerelease-${{ github.sha }} + path: | + target/release/stakpak + target/release/stakpak.sha256 + target/release/source-commit.txt + target/release/rust-version.txt + + upstream-suites: + timeout-minutes: 45 + runs-on: ubuntu-latest + env: + # Keep assertions/debug behavior; omit symbols to bound linker memory. + CARGO_BUILD_JOBS: '1' + CARGO_PROFILE_DEV_DEBUG: '0' + CARGO_PROFILE_TEST_DEBUG: '0' + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@master + with: + toolchain: '1.94.1' + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + with: + cache-on-failure: true + - run: cargo fmt --all -- --check + - run: cargo clippy --locked --all-targets -- -D warnings + - name: Complete workspace suite + timeout-minutes: 20 + run: cargo test --locked --workspace --verbose + - name: Complete feature suites + timeout-minutes: 15 + run: | + cargo test --locked -p stakpak-shared --features sqlite --verbose + cargo test --locked -p stakpak --features libsql-test --verbose + cargo test --locked -p stakpak-gateway --features libsql-test --verbose + cargo test --locked -p stakai --features network-tests --verbose + - run: cargo build --locked --verbose diff --git a/Cargo.lock b/Cargo.lock index ddd572cc..50e566d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6326,6 +6326,7 @@ dependencies = [ "grep-regex", "grep-searcher", "ignore", + "libc", "rand 0.9.2", "reqwest", "rmcp", diff --git a/libs/mcp/server/Cargo.toml b/libs/mcp/server/Cargo.toml index 976183c7..c40ba6de 100644 --- a/libs/mcp/server/Cargo.toml +++ b/libs/mcp/server/Cargo.toml @@ -35,6 +35,9 @@ grep-searcher = { workspace = true } ignore = { workspace = true } globset = { workspace = true } +[target.'cfg(unix)'.dependencies] +libc = "0.2" + [dev-dependencies] tempfile = "3.8" diff --git a/libs/mcp/server/src/foreground_command.rs b/libs/mcp/server/src/foreground_command.rs new file mode 100644 index 00000000..b67c2643 --- /dev/null +++ b/libs/mcp/server/src/foreground_command.rs @@ -0,0 +1,90 @@ +use std::io; +use std::process::ExitStatus; +use tokio::process::Child; + +/// Owns a foreground child's group until the child is reaped. Commands that +/// deliberately leave this group are outside this cancellation boundary. +pub(crate) struct ForegroundCommand { + child: Child, + owned: bool, +} + +impl ForegroundCommand { + pub(crate) fn new(child: Child) -> Self { + Self { child, owned: true } + } + + pub(crate) fn child_mut(&mut self) -> &mut Child { + &mut self.child + } + + #[cfg(unix)] + fn kill_group(&mut self) -> io::Result<()> { + if !self.owned { + return Ok(()); + } + if let Some(pid) = self.child.id() { + let pid = libc::pid_t::try_from(pid) + .map_err(|_| io::Error::other("invalid foreground process ID"))?; + if pid <= 1 { + return Err(io::Error::other("invalid foreground process group")); + } + // SAFETY: process_group(0) made this owned child the group leader. + // Child::id() becomes None after wait/try_wait successfully reaps it; + // while unreaped (even a zombie), its PID cannot be reused. + if unsafe { libc::kill(-pid, libc::SIGKILL) } == -1 { + let error = io::Error::last_os_error(); + if error.raw_os_error() != Some(libc::ESRCH) { + return Err(error); + } + } + } + Ok(()) + } + + pub(crate) async fn wait(&mut self) -> io::Result { + let result = self.child.wait().await; + match &result { + Ok(_) => self.owned = false, + // Child::id remains Some on a wait error. ECHILD means the + // exclusive ownership assumption was broken by another reaper; + // never signal that numeric PID/group again, including in Drop. + #[cfg(unix)] + Err(error) if error.raw_os_error() == Some(libc::ECHILD) => self.owned = false, + Err(_) => {} + } + result + } + + pub(crate) async fn terminate(&mut self) -> io::Result { + #[cfg(unix)] + self.kill_group()?; + #[cfg(not(unix))] + self.child.start_kill()?; + tokio::time::timeout(std::time::Duration::from_secs(5), self.wait()) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "foreground cleanup timed out"))? + } +} + +impl Drop for ForegroundCommand { + fn drop(&mut self) { + if !self.owned { + return; + } + #[cfg(unix)] + if let Err(error) = self.kill_group() { + tracing::error!("Failed to kill foreground command group: {error}"); + } + // Direct-child fallback is also restricted to a still-owned child. + // Command's default kill_on_drop(false) avoids a deferred stale-PID + // kill after ECHILD; Tokio still handles reaping when Child is dropped. + if self.child.id().is_some() { + let _ = self.child.start_kill(); + } + } +} + +#[cfg(all(test, unix))] +#[path = "foreground_command_tests.rs"] +mod tests; diff --git a/libs/mcp/server/src/foreground_command_tests.rs b/libs/mcp/server/src/foreground_command_tests.rs new file mode 100644 index 00000000..689f778e --- /dev/null +++ b/libs/mcp/server/src/foreground_command_tests.rs @@ -0,0 +1,60 @@ +// Fixture helpers intentionally fail fast on setup/transport errors. +#![allow(clippy::unwrap_used, clippy::expect_used)] + +use super::*; +use tokio::process::Command; + +fn successful_child() -> ForegroundCommand { + ForegroundCommand::new( + Command::new("sh") + .arg("-c") + .arg("exit 7") + .process_group(0) + .spawn() + .unwrap(), + ) +} + +#[tokio::test] +async fn foreground_wait_disarms_group_after_reap() { + let mut owned = successful_child(); + assert_eq!(owned.wait().await.unwrap().code(), Some(7)); + assert!(owned.child_mut().id().is_none()); + // No numeric authority remains even if Drop is delayed. + owned.kill_group().unwrap(); + drop(owned); +} + +#[tokio::test] +async fn foreground_try_wait_disarms_group_after_reap() { + let mut owned = successful_child(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if let Some(status) = owned.child_mut().try_wait().unwrap() { + assert_eq!(status.code(), Some(7)); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(owned.child_mut().id().is_none()); + owned.kill_group().unwrap(); + drop(owned); +} + +#[tokio::test] +async fn foreground_external_reap_disarms_signal_authority() { + let mut owned = successful_child(); + let pid = owned.child_mut().id().unwrap() as libc::pid_t; + let mut status = 0; + // SAFETY: reap only this test's exact owned child, simulating an external + // reaper. It immediately exits; no global waitpid(-1) or PID reuse tricks. + assert_eq!(unsafe { libc::waitpid(pid, &mut status, 0) }, pid); + let error = owned.wait().await.unwrap_err(); + assert_eq!(error.raw_os_error(), Some(libc::ECHILD)); + assert!(!owned.owned); + owned.kill_group().unwrap(); + drop(owned); +} diff --git a/libs/mcp/server/src/foreground_tests.rs b/libs/mcp/server/src/foreground_tests.rs new file mode 100644 index 00000000..11eb667c --- /dev/null +++ b/libs/mcp/server/src/foreground_tests.rs @@ -0,0 +1,247 @@ +// Fixture helpers intentionally fail fast on setup/transport errors. +#![allow(clippy::unwrap_used, clippy::expect_used)] + +use super::tests::local_container_with_profile; +use super::*; +use rmcp::service::PeerRequestOptions; +use rmcp::{ServiceExt, model::ClientRequest}; + +struct ExecutionHarness { + container: ToolContainer, + abort_execution: bool, +} + +impl rmcp::ServerHandler for ExecutionHarness { + fn get_info(&self) -> ServerInfo { + rmcp::ServerHandler::get_info(&self.container) + } + + async fn call_tool( + &self, + request: CallToolRequestParam, + ctx: RequestContext, + ) -> Result { + let command = request + .arguments + .as_ref() + .and_then(|args| args.get("command")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(); + if !self.abort_execution || command == "printf subsequent" { + return rmcp::ServerHandler::call_tool(&self.container, request, ctx).await; + } + let container = self.container.clone(); + let ct = ctx.ct.clone(); + // This case aborts the execution future itself, independently of the + // executor's ordinary request-token cancellation branches. + let mut ctx = ctx; + ctx.ct = Default::default(); + let execution = + tokio::spawn( + async move { container.execute_local_command(&command, None, &ctx).await }, + ); + ct.cancelled().await; + execution.abort(); + assert!(execution.await.unwrap_err().is_cancelled()); + Ok(CallToolResult::cancel(None)) + } +} + +fn process_identity(pid: u32) -> Option { + let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?; + let (_, fields) = stat.rsplit_once(") ")?; + fields.split_whitespace().nth(19).map(str::to_owned) +} + +async fn wait_for_marker(path: &Path) -> u32 { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Ok(text) = fs::read_to_string(path) + && let Ok(pid) = text.trim().parse() + { + return pid; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("nested shell must author its start marker") +} + +async fn assert_gone(pid: u32, identity: &str) { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let stat = fs::read_to_string(format!("/proc/{pid}/stat")); + // A killed orphan may briefly remain a zombie under CI's + // init; it cannot execute or create the delayed artifact. + if process_identity(pid).as_deref() != Some(identity) + || stat.as_ref().map_or(true, |s| { + s.rsplit_once(") ") + .is_some_and(|(_, rest)| rest.starts_with('Z')) + }) + { + break; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("owned shell must stop executing"); +} + +async fn prove_cleanup( + owned: &[(u32, String)], + late: &Path, + control: &mut tokio::process::Child, + identity: &str, +) { + for (pid, identity) in owned { + assert_gone(*pid, identity).await; + } + let control_pid = control.id().unwrap(); + assert!( + control.try_wait().unwrap().is_none(), + "unrelated process survives" + ); + assert_eq!(process_identity(control_pid).as_deref(), Some(identity)); + sleep(Duration::from_millis(2200)).await; + assert!(!late.exists(), "cancelled descendant must not write later"); + assert!( + control.try_wait().unwrap().is_none(), + "control survives deadline" + ); + assert_eq!(process_identity(control_pid).as_deref(), Some(identity)); +} + +async fn exercise(kind: &str, configured_timeout: Option) { + let dir = tempfile::tempdir().unwrap(); + let marker = dir.path().join("started"); + let late = dir.path().join("late"); + let outer = dir.path().join("outer"); + let sleeper = dir.path().join("sleeper"); + let completion = if kind == "leader-exit" { + "exit 0" + } else { + "wait" + }; + let command = format!( + r#"echo $$ > "{}"; sh -c 'echo $$ > "{}"; sleep 2 & echo $! > "{}"; wait; echo late > "{}"' & {completion}"#, + outer.display(), + marker.display(), + sleeper.display(), + late.display() + ); + let mut control = Command::new("sleep") + .arg("30") + .kill_on_drop(true) + .spawn() + .unwrap(); + let control_pid = control.id().unwrap(); + let control_identity = process_identity(control_pid).unwrap(); + let (client_io, server_io) = tokio::io::duplex(65536); + let abort_execution = kind == "abort"; + let server = tokio::spawn(async move { + ExecutionHarness { + container: local_container_with_profile(None), + abort_execution, + } + .serve(server_io) + .await + .unwrap() + }); + let client = ().serve(client_io).await.unwrap(); + let server = server.await.unwrap(); + let request = client + .send_cancellable_request( + ClientRequest::CallToolRequest(Request::new(CallToolRequestParam { + name: "run_command".into(), + arguments: Some( + serde_json::json!({ + "command": command, + "timeout": configured_timeout, + }) + .as_object() + .unwrap() + .clone(), + ), + })), + PeerRequestOptions::default(), + ) + .await + .unwrap(); + let pid = wait_for_marker(&marker).await; + let outer_pid = wait_for_marker(&outer).await; + let sleeper_pid = wait_for_marker(&sleeper).await; + let owned = [outer_pid, pid, sleeper_pid].map(|pid| (pid, process_identity(pid).unwrap())); + if kind == "disconnect" { + client.cancel().await.unwrap(); + server.cancel().await.unwrap(); + prove_cleanup(&owned, &late, &mut control, &control_identity).await; + } else { + if kind == "cancel" || kind == "abort" || kind == "leader-exit" { + client + .notify_cancelled(CancelledNotificationParam { + request_id: request.id.clone(), + reason: Some("regression cancellation".into()), + }) + .await + .unwrap(); + } + let result = tokio::time::timeout(Duration::from_secs(5), request.await_response()) + .await + .expect("executor must settle boundedly"); + if kind == "timeout" { + let result = result.unwrap(); + assert!(format!("{result:?}").contains("timed out")); + } + prove_cleanup(&owned, &late, &mut control, &control_identity).await; + // The same MCP session must still execute a subsequent command. + let next = client + .call_tool(CallToolRequestParam { + name: "run_command".into(), + arguments: Some( + serde_json::json!({"command": "printf subsequent"}) + .as_object() + .unwrap() + .clone(), + ), + }) + .await + .unwrap(); + assert!(format!("{next:?}").contains("subsequent"), "{next:?}"); + client.cancel().await.unwrap(); + server.cancel().await.unwrap(); + } + control.kill().await.unwrap(); +} + +#[tokio::test] +async fn foreground_cancel_without_timeout() { + exercise("cancel", None).await; +} + +#[tokio::test] +async fn foreground_cancel_with_timeout() { + exercise("cancel", Some(30)).await; +} + +#[tokio::test] +async fn foreground_timeout() { + exercise("timeout", Some(1)).await; +} + +#[tokio::test] +async fn foreground_disconnect() { + exercise("disconnect", None).await; +} + +#[tokio::test] +async fn foreground_execution_future_abort() { + exercise("abort", None).await; +} + +#[tokio::test] +async fn foreground_cancel_after_leader_exit_before_pipe_eof() { + exercise("leader-exit", None).await; +} diff --git a/libs/mcp/server/src/lib.rs b/libs/mcp/server/src/lib.rs index d3272326..6b330854 100644 --- a/libs/mcp/server/src/lib.rs +++ b/libs/mcp/server/src/lib.rs @@ -17,6 +17,7 @@ use stakpak_api::AgentProvider; use stakpak_shared::cert_utils::CertificateChain; use stakpak_shared::task_manager::{TaskManager, TaskManagerHandle}; +mod foreground_command; pub mod integrations; pub mod local_tools; pub mod remote_tools; diff --git a/libs/mcp/server/src/local_tools.rs b/libs/mcp/server/src/local_tools.rs index 9e742836..5f3df1c7 100644 --- a/libs/mcp/server/src/local_tools.rs +++ b/libs/mcp/server/src/local_tools.rs @@ -1,3 +1,4 @@ +use crate::foreground_command::ForegroundCommand; use crate::tool_container::ToolContainer; use rmcp::service::RequestContext; use rmcp::{ErrorData as McpError, handler::server::wrapper::Parameters, model::*, schemars, tool}; @@ -1351,7 +1352,7 @@ SAFETY NOTES: cmd.creation_flags(0x00000200); // CREATE_NEW_PROCESS_GROUP } - let mut child = cmd.spawn().map_err(|e| { + let child = cmd.spawn().map_err(|e| { error!("Failed to run command: {}", e); CallToolResult::error(vec![ Content::text("COMMAND_ERROR"), @@ -1359,14 +1360,17 @@ SAFETY NOTES: ]) })?; + let mut child = ForegroundCommand::new(child); #[allow(clippy::unwrap_used)] - let stdout = child.stdout.take().unwrap(); + let stdout = child.child_mut().stdout.take().unwrap(); #[allow(clippy::unwrap_used)] - let stderr = child.stderr.take().unwrap(); - let mut stdout_reader = BufReader::new(stdout); - let mut stderr_reader = BufReader::new(stderr); - let mut stdout_buf = String::new(); - let mut stderr_buf = String::new(); + let stderr = child.child_mut().stderr.take().unwrap(); + // Lines::next_line is cancellation-safe in select!: read_line can + // discard a partial final line when the other pipe or timer wins. + let mut stdout_reader = BufReader::new(stdout).lines(); + let mut stderr_reader = BufReader::new(stderr).lines(); + let mut stdout_done = false; + let mut stderr_done = false; let mut result = String::new(); let progress_id = Uuid::new_v4(); @@ -1382,14 +1386,12 @@ SAFETY NOTES: stall_check_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); macro_rules! handle_output { - ($read_result:expr, $buf:expr) => { + ($read_result:expr, $done:expr) => { match $read_result { - Ok(Ok(0)) => break, // EOF - Ok(Ok(_)) => { + Ok(None) => $done = true, // Drain the other pipe before waiting + Ok(Some(line)) => { last_output_time = std::time::Instant::now(); stall_start_time = None; // Reset stall tracking on output - let line = $buf.trim_end_matches('\n').to_string(); - $buf.clear(); result.push_str(&format!("{}\n", line)); let _ = ctx .peer @@ -1410,8 +1412,7 @@ SAFETY NOTES: }) .await; } - Ok(Err(_)) => break, // Read error - Err(_) => {} // Timeout - continue loop + Err(error) => return Err(error), } }; } @@ -1452,17 +1453,16 @@ SAFETY NOTES: } } - read_result = tokio::time::timeout(Duration::from_millis(100), stderr_reader.read_line(&mut stderr_buf)) => { - handle_output!(read_result, stderr_buf); + read_result = stderr_reader.next_line(), if !stderr_done => { + handle_output!(read_result, stderr_done); } - read_result = tokio::time::timeout(Duration::from_millis(100), stdout_reader.read_line(&mut stdout_buf)) => { - handle_output!(read_result, stdout_buf); + read_result = stdout_reader.next_line(), if !stdout_done => { + handle_output!(read_result, stdout_done); } } - // Check if process has exited - if let Ok(Some(_)) = child.try_wait() { + if stdout_done && stderr_done { break; } } @@ -1479,7 +1479,11 @@ SAFETY NOTES: result = tokio::time::timeout(timeout_duration, stream_and_wait) => result, _ = ctx.ct.cancelled() => { // Cancellation occurred, kill the process - let _ = child.kill().await; + if let Err(error) = child.terminate().await { + return Err(CallToolResult::error(vec![Content::text(format!( + "Foreground command cleanup failed: {error}" + ))])); + } return Err(CallToolResult::cancel(Some(&vec![ Content::text("COMMAND_CANCELLED"), Content::text("Command execution was cancelled"), @@ -1490,7 +1494,11 @@ SAFETY NOTES: tokio::select! { result = stream_and_wait => Ok(result), _ = ctx.ct.cancelled() => { - let _ = child.kill().await; + if let Err(error) = child.terminate().await { + return Err(CallToolResult::error(vec![Content::text(format!( + "Foreground command cleanup failed: {error}" + ))])); + } return Err(CallToolResult::cancel(Some(&vec![ Content::text("COMMAND_CANCELLED"), Content::text("Command execution was cancelled"), @@ -1509,7 +1517,11 @@ SAFETY NOTES: } Err(_) => { // Timeout occurred, kill the process - let _ = child.kill().await; + if let Err(error) = child.terminate().await { + return Err(CallToolResult::error(vec![Content::text(format!( + "Foreground command cleanup failed: {error}" + ))])); + } result.push_str(&format!( "Command timed out after {} seconds\n", timeout.unwrap_or_default() @@ -3370,7 +3382,7 @@ mod tests { ); } - fn local_container_with_profile(profile_name: Option<&str>) -> ToolContainer { + pub(super) fn local_container_with_profile(profile_name: Option<&str>) -> ToolContainer { let task_manager = stakpak_shared::task_manager::TaskManager::new(); ToolContainer::new( None, @@ -4096,3 +4108,7 @@ mod tests { assert_eq!(result, Some(("a***b***c".to_string(), 2))); } } + +#[cfg(all(test, target_os = "linux"))] +#[path = "foreground_tests.rs"] +mod foreground_tests;