diff --git a/crates/brush-core-vendored/src/jobs.rs b/crates/brush-core-vendored/src/jobs.rs index 02f4f08b9..0a575b23c 100644 --- a/crates/brush-core-vendored/src/jobs.rs +++ b/crates/brush-core-vendored/src/jobs.rs @@ -437,6 +437,26 @@ impl Job { } } + /// Aborts shell-internal background tasks and drops their join handles. + /// + /// External process jobs are intentionally left alone; callers that abort + /// internal tasks are still responsible for signalling any process trees + /// those tasks may have spawned. + pub fn abort_internal_tasks(&mut self) { + let mut aborted = false; + self.tasks.retain_mut(|task| { + if let JobTask::Internal(handle) = task { + handle.abort(); + aborted = true; + return false; + } + true + }); + if aborted && self.tasks.is_empty() { + self.state = JobState::Done; + } + } + /// Tries to retrieve a "representative" pid for the job. pub fn representative_pid(&self) -> Option { for task in &self.tasks { diff --git a/crates/pi-shell/src/shell.rs b/crates/pi-shell/src/shell.rs index f844fc83f..4de375c1c 100644 --- a/crates/pi-shell/src/shell.rs +++ b/crates/pi-shell/src/shell.rs @@ -663,7 +663,7 @@ async fn run_shell_command( .await; if cancel_token.is_cancelled() { - terminate_background_jobs(&session.shell); + terminate_background_jobs(&mut session.shell); } if env_scope_pushed { @@ -833,7 +833,7 @@ async fn run_shell_command_streams( .await; if cancel_token.is_cancelled() { - terminate_background_jobs(&session.shell); + terminate_background_jobs(&mut session.shell); } if env_scope_pushed { @@ -998,9 +998,10 @@ async fn terminate_new_descendants(baseline: & } } } -fn terminate_background_jobs(shell: &BrushShell) { +fn terminate_background_jobs(shell: &mut BrushShell) { let mut targets = process::TerminationTargets::new(); - for job in &shell.jobs().jobs { + for job in &mut shell.jobs_mut().jobs { + job.abort_internal_tasks(); if let Some(pgid) = job.process_group_id() { targets.add_pgid(pgid); } @@ -1009,11 +1010,9 @@ fn terminate_background_jobs(shell: &BrushShell) { } } if targets.is_empty() { - // Pure descendant cleanup is handled by `process_cancel_bridge` while - // the cancel was still in flight. Here we only signal brush's own - // job-tracked targets — pgids of background-group leaders that may have - // already exited (so the descendant walk would no longer find them as - // new descendants, but their group still holds live grandchildren). + // Shell-internal jobs were aborted above. Pure descendant cleanup is + // handled by `process_cancel_bridge` while the cancel was in flight; + // without job-tracked pgids or pids there is nothing else to signal here. return; } @@ -1934,6 +1933,66 @@ mod tests { assert!(matches!(reason, AbortReason::Signal)); } + #[tokio::test(flavor = "multi_thread")] + async fn cancellation_aborts_internal_background_jobs() { + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos(); + let dir = + std::env::temp_dir().join(format!("pi-shell-bg-cancel-{}-{unique}", std::process::id())); + std::fs::create_dir(&dir).expect("create temp dir"); + let started = dir.join("started"); + let release = dir.join("release"); + let marker = dir.join("marker"); + + let config = ShellConfig { session_env: None, snapshot_path: None, minimizer: None }; + let mut session = create_session(&config).await.expect("create session"); + session + .shell + .set_working_dir(dir.to_string_lossy().as_ref()) + .expect("set cwd"); + + let mut params = session.shell.default_exec_params(); + params.set_fd(OpenFiles::STDIN_FD, null_file().expect("null stdin")); + params.set_fd(OpenFiles::STDOUT_FD, null_file().expect("null stdout")); + params.set_fd(OpenFiles::STDERR_FD, null_file().expect("null stderr")); + + let source_info = SourceInfo::from("pi-shell:test"); + let result = session + .shell + .run_string( + "{ echo started > started; while [ ! -f release ]; do sleep 0.05; done; echo done > \ + marker; } &", + &source_info, + ¶ms, + ) + .await + .expect("spawn background job"); + assert_eq!(exit_code(&result), 0); + + let mut background_started = false; + for _ in 0..200 { + if started.exists() { + background_started = true; + break; + } + time::sleep(Duration::from_millis(10)).await; + } + assert!(background_started, "background job did not reach its wait loop"); + + terminate_background_jobs(&mut session.shell); + std::fs::write(&release, b"").expect("release marker"); + time::sleep(Duration::from_millis(250)).await; + let marker_exists = marker.exists(); + std::fs::remove_dir_all(&dir).expect("cleanup temp dir"); + + assert!( + !marker_exists, + "internal background job survived cancellation and wrote marker after release", + ); + } + #[cfg(unix)] #[tokio::test] async fn read_output_stops_when_cancelled_before_pipe_eof() { diff --git a/packages/coding-agent/test/bash-executor.test.ts b/packages/coding-agent/test/bash-executor.test.ts index ce93ef450..96d92f5b8 100644 --- a/packages/coding-agent/test/bash-executor.test.ts +++ b/packages/coding-agent/test/bash-executor.test.ts @@ -570,11 +570,15 @@ describe("executeBash", () => { if (process.platform === "win32") return; const marker = path.join(tempDir, "marker-bg-abort.txt"); + const release = path.join(tempDir, "marker-bg-abort.release"); + const started = path.join(tempDir, "marker-bg-abort.started"); const markerEscaped = marker.replace(/'/g, "'\\''"); + const releaseEscaped = release.replace(/'/g, "'\\''"); + const startedEscaped = started.replace(/'/g, "'\\''"); const controller = new AbortController(); const promise = executeBash( - `{ sleep ${KILL_MARKER_DELAY_SECONDS}; echo done > '${markerEscaped}'; } & sleep 10`, + `{ touch '${startedEscaped}'; while [ ! -f '${releaseEscaped}' ]; do sleep 0.05; done; echo done > '${markerEscaped}'; } & sleep 10`, { cwd: tempDir, timeout: 10000, @@ -582,7 +586,11 @@ describe("executeBash", () => { }, ); - await Bun.sleep(100); + const startDeadline = Date.now() + 4000; + while (!fs.existsSync(started) && Date.now() < startDeadline) { + await Bun.sleep(2); + } + expect(fs.existsSync(started)).toBe(true); controller.abort(); const result = await promise; @@ -590,6 +598,8 @@ describe("executeBash", () => { expect(result.output).toContain("Command cancelled"); await Bun.sleep(KILL_MARKER_ASSERTION_WAIT_MS); + fs.writeFileSync(release, ""); + await Bun.sleep(150); expect(fs.existsSync(marker)).toBe(false); });