diff --git a/Cargo.lock b/Cargo.lock index be372a555..73ede88c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2505,6 +2505,7 @@ dependencies = [ "brush-builtins", "brush-core", "brush-parser 0.3.0", + "bytes", "clap", "libc", "os_pipe", diff --git a/Cargo.toml b/Cargo.toml index 3753692c4..517a527fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -189,6 +189,7 @@ toml = "1.1" # Error Handling # ────────────────────────────────────────────────────────────────────────────── anyhow = "1.0" +bytes = "1" # ────────────────────────────────────────────────────────────────────────────── # CLI & Configuration diff --git a/crates/pi-shell/Cargo.toml b/crates/pi-shell/Cargo.toml index de3104ba9..9bae65ae1 100644 --- a/crates/pi-shell/Cargo.toml +++ b/crates/pi-shell/Cargo.toml @@ -11,6 +11,7 @@ workspace = true [dependencies] anyhow.workspace = true +bytes.workspace = true brush-builtins.workspace = true brush-core.workspace = true brush-parser.workspace = true diff --git a/crates/pi-shell/src/lib.rs b/crates/pi-shell/src/lib.rs index e1571845d..9d9f90635 100644 --- a/crates/pi-shell/src/lib.rs +++ b/crates/pi-shell/src/lib.rs @@ -8,5 +8,5 @@ pub mod windows; pub use brush_core::commands::{ChildSessionAction, child_session_action}; pub use shell::{ MinimizerResult, Shell, ShellExecuteOptions, ShellExecuteResult, ShellOptions, ShellRunOptions, - ShellRunResult, execute_shell, + ShellRunResult, StreamSinks, execute_shell, execute_shell_streams, }; diff --git a/crates/pi-shell/src/shell.rs b/crates/pi-shell/src/shell.rs index 2b93a65f9..acde96d3e 100644 --- a/crates/pi-shell/src/shell.rs +++ b/crates/pi-shell/src/shell.rs @@ -12,6 +12,7 @@ use std::{ }; use anyhow::{Error, Result}; +use bytes::Bytes; use brush_builtins::{BuiltinSet, default_builtins}; use brush_core::{ ExecutionContext, ExecutionControlFlow, ExecutionExitCode, ExecutionResult, ProcessGroupPolicy, @@ -196,6 +197,40 @@ pub async fn execute_shell( run_shell_oneshot(config, run_config, on_chunk, cancel_token).await } +/// Optional per-stream raw byte sinks for [`execute_shell_streams`]. +/// +/// When a sink is `Some`, that stream's pipe is drained directly into the +/// channel with no UTF-8 decoding and no merging. When `None`, the +/// corresponding pipe is still drained (to avoid blocking the child) but +/// its bytes are dropped. +#[derive(Default)] +pub struct StreamSinks { + pub stdout: Option>, + pub stderr: Option>, +} + +/// One-shot execution that delivers stdout and stderr as raw byte chunks on +/// separate channels. The minimizer is intentionally disabled — its +/// `MinimizerResult.text` contract presumes a single merged transcript. +pub async fn execute_shell_streams( + options: ShellExecuteOptions, + streams: StreamSinks, + cancel_token: CancelToken, +) -> Result { + let config = ShellConfig { + session_env: options.session_env, + snapshot_path: options.snapshot_path, + minimizer: None, + }; + let run_config = ShellRunConfig { + command: options.command, + cwd: options.cwd, + env: options.env, + minimizer: None, + }; + run_shell_oneshot_streams(config, run_config, streams, cancel_token).await +} + async fn run_shell_session( session: Arc>>, abort_state: ShellAbortState, @@ -309,6 +344,51 @@ async fn run_shell_oneshot( }) } +async fn run_shell_oneshot_streams( + config: ShellConfig, + run_config: ShellRunConfig, + streams: StreamSinks, + ct: CancelToken, +) -> Result { + let tokio_cancel = CancellationToken::new(); + + let mut task = tokio::spawn({ + let tokio_cancel = tokio_cancel.clone(); + async move { + let mut session = create_session(&config).await?; + run_shell_command_streams(&mut session, &run_config, streams, tokio_cancel).await + } + }); + + let run_result = tokio::select! { + result = &mut task => result, + reason = ct.wait() => { + tokio_cancel.cancel(); + let graceful = time::timeout(Duration::from_secs(2), &mut task).await; + if graceful.is_err() { + task.abort(); + let _ = task.await; + } + return Ok(ShellExecuteResult { + exit_code: None, + cancelled: matches!(reason, AbortReason::Signal), + timed_out: matches!(reason, AbortReason::Timeout), + minimized: None, + }); + }, + }; + + let res = run_result + .unwrap_or_else(|err| Err(Error::msg(format!("Shell execution task failed: {err}")))); + let exec = res?; + Ok(ShellExecuteResult { + exit_code: Some(exit_code(&exec)), + cancelled: false, + timed_out: false, + minimized: None, + }) +} + fn null_file() -> Result { openfiles::null().map_err(|err| Error::msg(format!("Failed to create null file: {err}"))) } @@ -729,6 +809,259 @@ async fn run_shell_command( Ok((result, minimized_out)) } +async fn run_shell_command_streams( + session: &mut ShellSessionCore, + options: &ShellRunConfig, + streams: StreamSinks, + cancel_token: CancellationToken, +) -> Result { + if let Some(cwd) = options.cwd.as_deref() { + session + .shell + .set_working_dir(cwd) + .map_err(|err| Error::msg(format!("Failed to set cwd: {err}")))?; + } + + let mut env_scope_pushed = false; + if let Some(env) = options.env.as_ref() { + session + .shell + .env_mut() + .push_scope(EnvironmentScope::Command); + env_scope_pushed = true; + for (key, value) in env { + let normalized_key = normalize_env_key(key); + if should_skip_env_var(normalized_key) { + continue; + } + let mut var = ShellVariable::new(ShellValue::String(value.clone())); + var.export(); + if let Err(err) = + session + .shell + .env_mut() + .add(normalized_key, var, EnvironmentScope::Command) + { + let _ = session.shell.env_mut().pop_scope(EnvironmentScope::Command); + return Err(Error::msg(format!("Failed to set env: {err}"))); + } + } + } + + let (stdout_reader, stdout_writer) = pipe_to_files("stdout")?; + let (stderr_reader, stderr_writer) = pipe_to_files("stderr")?; + + let stdout_file = OpenFile::from(stdout_writer); + let stderr_file = OpenFile::from(stderr_writer); + + let mut params = session.shell.default_exec_params(); + params.set_fd(OpenFiles::STDIN_FD, null_file()?); + params.set_fd(OpenFiles::STDOUT_FD, stdout_file); + params.set_fd(OpenFiles::STDERR_FD, stderr_file); + params.process_group_policy = ProcessGroupPolicy::NewProcessGroup; + params.set_cancel_token(cancel_token.clone()); + let baseline_descendants = process::current_descendant_pids(); + let reader_cancel = CancellationToken::new(); + let (activity_tx, mut activity_rx) = mpsc::channel::<()>(1); + + let StreamSinks { stdout: stdout_sink, stderr: stderr_sink } = streams; + let mut stdout_handle = tokio::spawn(Box::pin(read_output_bytes( + stdout_reader, + stdout_sink, + reader_cancel.clone(), + activity_tx.clone(), + ))); + let mut stderr_handle = tokio::spawn(Box::pin(read_output_bytes( + stderr_reader, + stderr_sink, + reader_cancel.clone(), + activity_tx, + ))); + + let cancel_bridge = tokio::spawn({ + let cancel_token = cancel_token.clone(); + let reader_cancel = reader_cancel.clone(); + async move { + cancel_token.cancelled().await; + reader_cancel.cancel(); + } + }); + let process_cancel_bridge = tokio::spawn({ + let cancel_token = cancel_token.clone(); + 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; + } + } + } + }); + let source_info = SourceInfo::from("pi-shell:streams"); + let result = session + .shell + .run_string(options.command.clone(), &source_info, ¶ms) + .await; + + if cancel_token.is_cancelled() { + terminate_background_jobs(&session.shell); + } + + if env_scope_pushed { + session + .shell + .env_mut() + .pop_scope(EnvironmentScope::Command) + .map_err(|err| Error::msg(format!("Failed to pop env scope: {err}")))?; + } + + drop(params); + + const POST_EXIT_IDLE: Duration = Duration::from_millis(250); + const POST_EXIT_MAX: Duration = Duration::from_secs(2); + const READER_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(250); + + let mut stdout_finished = false; + let mut stderr_finished = false; + let mut idle_timer = Box::pin(time::sleep(POST_EXIT_IDLE)); + let mut max_timer = Box::pin(time::sleep(POST_EXIT_MAX)); + + loop { + if stdout_finished && stderr_finished { + break; + } + tokio::select! { + res = &mut stdout_handle, if !stdout_finished => { + let _ = res; + stdout_finished = true; + } + res = &mut stderr_handle, if !stderr_finished => { + let _ = res; + stderr_finished = true; + } + msg = activity_rx.recv() => { + if msg.is_none() { + break; + } + idle_timer.as_mut().reset(time::Instant::now() + POST_EXIT_IDLE); + } + () = &mut idle_timer => break, + () = &mut max_timer => break, + } + } + + if !stdout_finished || !stderr_finished { + reader_cancel.cancel(); + } + if !stdout_finished { + if time::timeout(READER_SHUTDOWN_TIMEOUT, &mut stdout_handle) + .await + .is_err() + { + stdout_handle.abort(); + let _ = stdout_handle.await; + } + } + if !stderr_finished { + if time::timeout(READER_SHUTDOWN_TIMEOUT, &mut stderr_handle) + .await + .is_err() + { + stderr_handle.abort(); + let _ = stderr_handle.await; + } + } + cancel_bridge.abort(); + let _ = cancel_bridge.await; + if cancel_token.is_cancelled() { + // Let the kill-wave bridge finish all three signal passes so stragglers + // have a chance to receive SIGKILL. + let _ = process_cancel_bridge.await; + } else { + process_cancel_bridge.abort(); + let _ = process_cancel_bridge.await; + } + + let result = result.map_err(|err| Error::msg(format!("Shell execution failed: {err}")))?; + Ok(result) +} + +async fn read_output_bytes( + reader: fs::File, + sink: Option>, + cancel_token: CancellationToken, + activity: mpsc::Sender<()>, +) { + const BUF: usize = 65536; + + #[cfg(unix)] + let Ok(reader) = register_nonblocking_pipe(reader) else { + return; + }; + #[cfg(not(unix))] + let mut reader = tokio::fs::File::from_std(reader); + + loop { + let mut buf = vec![0u8; BUF]; + #[cfg(unix)] + let n = { + let Ok(mut readiness) = (tokio::select! { + ready = reader.readable() => ready, + () = cancel_token.cancelled() => break, + }) else { + break; + }; + match readiness.try_io(|inner| read_nonblocking(inner.get_ref(), &mut buf)) { + Ok(Ok(0)) => break, + Ok(Ok(n)) => n, + Ok(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue, + Ok(Err(_)) => break, + Err(_would_block) => continue, + } + }; + #[cfg(not(unix))] + let n = { + let read_future = reader.read(&mut buf); + tokio::pin!(read_future); + match tokio::select! { + res = &mut read_future => res, + () = cancel_token.cancelled() => break, + } { + Ok(0) => break, + Ok(n) => n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => break, + } + }; + let _ = activity.try_send(()); + buf.truncate(n); + if let Some(sink) = sink.as_ref() + && sink.send(Bytes::from(buf)).is_err() + { + // Receiver dropped — stop forwarding and let the pipe close. + break; + } + } +} + fn terminate_background_jobs(shell: &BrushShell) { let mut targets = process::TerminationTargets::new(); for job in &shell.jobs().jobs { @@ -1505,4 +1838,46 @@ mod tests { .expect("reader task should stop after cancellation") .expect("reader task should not panic"); } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread")] + async fn execute_shell_streams_separates_stdout_and_stderr() { + let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); + let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); + let options = ShellExecuteOptions { + command: "echo out; echo err 1>&2".to_string(), + ..Default::default() + }; + let streams = StreamSinks { stdout: Some(stdout_tx), stderr: Some(stderr_tx) }; + let result = execute_shell_streams(options, streams, CancelToken::default()) + .await + .expect("execute should succeed"); + assert_eq!(result.exit_code, Some(0)); + assert!(!result.cancelled); + + let mut stdout = Vec::new(); + while let Some(chunk) = stdout_rx.recv().await { + stdout.extend_from_slice(&chunk); + } + let mut stderr = Vec::new(); + while let Some(chunk) = stderr_rx.recv().await { + stderr.extend_from_slice(&chunk); + } + assert_eq!(stdout, b"out\n"); + assert_eq!(stderr, b"err\n"); + } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread")] + async fn execute_shell_streams_works_when_sinks_are_none() { + // Both sinks `None` — pipes must still drain so the child can exit. + let options = ShellExecuteOptions { + command: "yes done | head -n 100 1>&2; echo final".to_string(), + ..Default::default() + }; + let result = execute_shell_streams(options, StreamSinks::default(), CancelToken::default()) + .await + .expect("execute should succeed"); + assert_eq!(result.exit_code, Some(0)); + } }