From de5b0a51e2df81883aacdc8f482d0f3667c769bb Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 25 May 2026 18:13:49 +0000 Subject: [PATCH] fix(bash): forcefully terminated canceled shells Added Rust-side descendant termination on shell cancellation paths so aborts and timeouts escalate to SIGKILL even if brush cleanup stalls. Covered SIGTERM-ignoring shell workloads in native tests. Fixes #1347 --- crates/pi-shell/src/shell.rs | 88 +++++++++++----------------- packages/natives/test/native.test.ts | 21 +++++++ 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/crates/pi-shell/src/shell.rs b/crates/pi-shell/src/shell.rs index 55a3dae8a..4efa4f343 100644 --- a/crates/pi-shell/src/shell.rs +++ b/crates/pi-shell/src/shell.rs @@ -1,9 +1,7 @@ //! Runtime-agnostic brush shell execution. -#[cfg(windows)] -use std::collections::HashSet; use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, fs, io::{self, Write}, str, @@ -242,6 +240,7 @@ async fn run_shell_session( ct: &mut CancelToken, ) -> Result { let tokio_cancel = CancellationToken::new(); + let baseline_descendants = process::current_descendant_pids(); let mut run_task = tokio::spawn({ let session = session.clone(); @@ -264,6 +263,7 @@ async fn run_shell_session( res = &mut run_task => res, reason = ct.wait() => { tokio_cancel.cancel(); + terminate_new_descendants(&baseline_descendants).await; let graceful = time::timeout(Duration::from_secs(2), &mut run_task).await; if graceful.is_err() { run_task.abort(); @@ -308,6 +308,7 @@ async fn run_shell_oneshot( ct: CancelToken, ) -> Result { let tokio_cancel = CancellationToken::new(); + let baseline_descendants = process::current_descendant_pids(); let mut task = tokio::spawn({ let tokio_cancel = tokio_cancel.clone(); @@ -321,6 +322,7 @@ async fn run_shell_oneshot( result = &mut task => result, reason = ct.wait() => { tokio_cancel.cancel(); + terminate_new_descendants(&baseline_descendants).await; let graceful = time::timeout(Duration::from_secs(2), &mut task).await; if graceful.is_err() { task.abort(); @@ -353,6 +355,7 @@ async fn run_shell_oneshot_streams( ct: CancelToken, ) -> Result { let tokio_cancel = CancellationToken::new(); + let baseline_descendants = process::current_descendant_pids(); let mut task = tokio::spawn({ let tokio_cancel = tokio_cancel.clone(); @@ -366,6 +369,7 @@ async fn run_shell_oneshot_streams( result = &mut task => result, reason = ct.wait() => { tokio_cancel.cancel(); + terminate_new_descendants(&baseline_descendants).await; let graceful = time::timeout(Duration::from_secs(2), &mut task).await; if graceful.is_err() { task.abort(); @@ -649,34 +653,7 @@ async fn run_shell_command( let baseline_descendants = baseline_descendants.clone(); async move { cancel_token.cancelled().await; - // Rescan-and-signal loop. Each pass picks up grandchildren spawned - // during the previous wave's grace period, then exits early as soon - // as no descendants remain. The first wave is SIGTERM so well-behaved - // programs get a chance to clean up; subsequent waves escalate to - // SIGKILL. Cheaper than the previous 20 Hz tracker loop and avoids - // the constant kernel chatter when no cancellation ever happens. - const WAVES: u32 = 3; - for wave in 0..WAVES { - let mut targets = process::TerminationTargets::new(); - process::add_new_descendants(&mut targets, &baseline_descendants); - if targets.is_empty() { - return; - } - let signal = if wave == 0 { - process::TERM_SIGNAL - } else { - process::KILL_SIGNAL - }; - targets.signal(signal); - if wave + 1 < WAVES { - let pause = if wave == 0 { - Duration::from_millis(75) - } else { - Duration::from_millis(150) - }; - time::sleep(pause).await; - } - } + terminate_new_descendants(&baseline_descendants).await; } }); let source_info = SourceInfo::from("pi-natives:command"); @@ -846,28 +823,7 @@ async fn run_shell_command_streams( let baseline_descendants = baseline_descendants.clone(); async move { cancel_token.cancelled().await; - const WAVES: u32 = 3; - for wave in 0..WAVES { - let mut targets = process::TerminationTargets::new(); - process::add_new_descendants(&mut targets, &baseline_descendants); - if targets.is_empty() { - return; - } - let signal = if wave == 0 { - process::TERM_SIGNAL - } else { - process::KILL_SIGNAL - }; - targets.signal(signal); - if wave + 1 < WAVES { - let pause = if wave == 0 { - Duration::from_millis(75) - } else { - Duration::from_millis(150) - }; - time::sleep(pause).await; - } - } + terminate_new_descendants(&baseline_descendants).await; } }); let source_info = SourceInfo::from("pi-shell:streams"); @@ -1015,6 +971,32 @@ async fn read_output_bytes( } } +// Rescan-and-signal loop for cancellation. Each pass picks up descendants +// spawned during the previous wave's grace period; empty scans keep going so a +// cancellation that wins before `run_string` exposes its child still escalates. +async fn terminate_new_descendants(baseline: &HashSet) { + const WAVES: u32 = 3; + for wave in 0..WAVES { + let mut targets = process::TerminationTargets::new(); + process::add_new_descendants(&mut targets, baseline); + if !targets.is_empty() { + let signal = if wave == 0 { + process::TERM_SIGNAL + } else { + process::KILL_SIGNAL + }; + targets.signal(signal); + } + if wave + 1 < WAVES { + let pause = if wave == 0 { + Duration::from_millis(75) + } else { + Duration::from_millis(150) + }; + time::sleep(pause).await; + } + } +} fn terminate_background_jobs(shell: &BrushShell) { let mut targets = process::TerminationTargets::new(); for job in &shell.jobs().jobs { diff --git a/packages/natives/test/native.test.ts b/packages/natives/test/native.test.ts index bdb1cb14c..4dcb9584e 100644 --- a/packages/natives/test/native.test.ts +++ b/packages/natives/test/native.test.ts @@ -532,6 +532,27 @@ describe("pi-natives", () => { await Bun.sleep(500); expect(await Bun.file(markerPath).exists()).toBe(false); }); + + it("should SIGKILL workloads that ignore SIGTERM on timeout", async () => { + if (process.platform === "win32") { + return; + } + + const markerPath = path.join(testDir, "shell-timeout-sigkill-marker.txt"); + const markerEscaped = markerPath.replace(/'/g, "'\\''"); + await fs.rm(markerPath, { force: true }); + + const result = await executeShell({ + command: `trap '' TERM; sleep 0.3; echo done > '${markerEscaped}'`, + cwd: testDir, + timeoutMs: 50, + }); + + expect(result.timedOut).toBe(true); + + await Bun.sleep(600); + expect(await Bun.file(markerPath).exists()).toBe(false); + }); }); describe("htmlToMarkdown", () => { it("should convert basic HTML to markdown", async () => {