command_stream/stream.rs
1//! Streaming and async iteration support
2//!
3//! This module provides async streaming capabilities similar to JavaScript's
4//! async iterators and stream handling in `$.stream-utils.mjs`.
5//!
6//! It mirrors the JavaScript implementation's behavior for issue #155:
7//!
8//! 1. The stream yields an explicit `OutputChunk::Exit(code)` when the
9//! process exits, so consumers can observe the exit code from inside the
10//! loop.
11//! 2. The stream does not hang forever when the process has exited but a
12//! grandchild keeps the stdio pipes open (the readers are drained with a
13//! grace period and then aborted).
14//! 3. The process can be stopped from inside the loop via
15//! [`OutputStream::kill`] / [`OutputStream::kill_with`], and abandoning the
16//! stream (e.g. `break`) also stops the process.
17//! 4. The stop signal is configurable via
18//! [`StreamingRunner::kill_signal`] (default `SIGTERM`), just like the
19//! JavaScript `killSignal` option.
20//!
21//! ## Usage
22//!
23//! ```rust,no_run
24//! use command_stream::{StreamingRunner, OutputChunk};
25//!
26//! #[tokio::main]
27//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
28//! let runner = StreamingRunner::new("yes hello");
29//!
30//! // Stream output as it arrives
31//! let mut stream = runner.stream();
32//! let mut count = 0;
33//! while let Some(chunk) = stream.next().await {
34//! match chunk {
35//! OutputChunk::Stdout(data) => {
36//! print!("{}", String::from_utf8_lossy(&data));
37//! count += 1;
38//! if count >= 5 {
39//! // Stop the process from inside the loop.
40//! stream.kill();
41//! }
42//! }
43//! OutputChunk::Stderr(data) => {
44//! eprint!("{}", String::from_utf8_lossy(&data));
45//! }
46//! OutputChunk::Exit(code) => {
47//! println!("Process exited with code: {}", code);
48//! break;
49//! }
50//! }
51//! }
52//!
53//! Ok(())
54//! }
55//! ```
56
57use std::collections::HashMap;
58use std::ffi::OsString;
59use std::path::PathBuf;
60use std::time::Duration;
61use tokio::sync::{mpsc, watch};
62use tokio::task::JoinHandle;
63
64use crate::signal::{DEFAULT_KILL_GRACE_MS, DEFAULT_KILL_SIGNAL};
65use crate::stream_process::{
66 run_streaming_process, ChildEnv, GraceWindows, ProcessSetup, StreamChannels, StreamingCommand,
67};
68use crate::stream_stdin::{StdinRequest, StdinSender};
69use crate::trace::trace_lazy;
70use crate::{CommandResult, Result};
71
72/// Default grace period (in milliseconds) to keep draining the stdio pipes
73/// after the process has exited before aborting any lingering readers. Mirrors
74/// the JavaScript `exitPumpGrace` default.
75pub(crate) const DEFAULT_EXIT_PUMP_GRACE_MS: u64 = 100;
76
77/// A chunk of output from a streaming process
78#[derive(Debug, Clone)]
79pub enum OutputChunk {
80 /// Stdout data
81 Stdout(Vec<u8>),
82 /// Stderr data
83 Stderr(Vec<u8>),
84 /// Process exit code
85 Exit(i32),
86}
87
88/// A streaming process runner that allows async iteration over output
89pub struct StreamingRunner {
90 command: StreamingCommand,
91 cwd: Option<PathBuf>,
92 env: Option<HashMap<String, String>>,
93 env_clear: bool,
94 env_remove: Vec<String>,
95 prefer_local: crate::PreferLocal,
96 stdin_content: Option<Vec<u8>>,
97 stdin_pipe: bool,
98 timeout: Option<Duration>,
99 kill_signal: String,
100 kill_grace_ms: u64,
101 exit_pump_grace_ms: u64,
102}
103
104impl StreamingRunner {
105 /// Create a streaming runner for a command string interpreted by the
106 /// platform shell.
107 pub fn new(command: impl Into<String>) -> Self {
108 Self::with_command(StreamingCommand::Shell(command.into()))
109 }
110
111 /// Create a streaming runner for an executable and exact argument vector.
112 ///
113 /// Unlike [`StreamingRunner::new`], this constructor bypasses the platform
114 /// shell. Argument boundaries are therefore preserved on every platform,
115 /// including Windows, without requiring shell-specific quoting.
116 pub fn from_argv<P, I, S>(program: P, args: I) -> Self
117 where
118 P: Into<OsString>,
119 I: IntoIterator<Item = S>,
120 S: Into<OsString>,
121 {
122 Self::with_command(StreamingCommand::Argv {
123 program: program.into(),
124 args: args.into_iter().map(Into::into).collect(),
125 })
126 }
127
128 pub(crate) fn with_command(command: StreamingCommand) -> Self {
129 StreamingRunner {
130 command,
131 cwd: None,
132 env: None,
133 env_clear: false,
134 env_remove: Vec::new(),
135 prefer_local: crate::PreferLocal::Off,
136 stdin_content: None,
137 stdin_pipe: false,
138 timeout: None,
139 kill_signal: DEFAULT_KILL_SIGNAL.to_string(),
140 kill_grace_ms: DEFAULT_KILL_GRACE_MS,
141 exit_pump_grace_ms: DEFAULT_EXIT_PUMP_GRACE_MS,
142 }
143 }
144
145 /// Set the working directory
146 pub fn cwd(mut self, path: impl Into<PathBuf>) -> Self {
147 self.cwd = Some(path.into());
148 self
149 }
150
151 /// Set environment variables
152 pub fn env(mut self, env: HashMap<String, String>) -> Self {
153 self.env = Some(env);
154 self
155 }
156
157 /// Prefer executables from the command's working directory or explicit directories.
158 pub fn prefer_local(mut self, preference: crate::PreferLocal) -> Self {
159 self.prefer_local = preference;
160 self
161 }
162
163 /// Set stdin content
164 pub fn stdin(mut self, content: impl Into<String>) -> Self {
165 self.stdin_content = Some(content.into().into_bytes());
166 self
167 }
168
169 /// Set binary stdin without UTF-8 conversion.
170 pub fn stdin_bytes(mut self, content: impl Into<Vec<u8>>) -> Self {
171 self.stdin_content = Some(content.into());
172 self
173 }
174
175 /// Clear inherited environment variables before applying `env`.
176 pub fn clear_env(mut self, clear: bool) -> Self {
177 self.env_clear = clear;
178 self
179 }
180
181 /// Remove one variable from the child's environment, the equivalent of
182 /// `env -u NAME` or [`std::process::Command::env_remove`].
183 ///
184 /// Removals are applied after [`env`](Self::env) and
185 /// [`clear_env`](Self::clear_env), so a removed variable is absent even
186 /// when the `env` map also sets it.
187 ///
188 /// ```no_run
189 /// use command_stream::StreamingRunner;
190 ///
191 /// # async fn example() -> command_stream::Result<()> {
192 /// let result = StreamingRunner::from_argv("printenv", ["CLAUDECODE"])
193 /// .env_remove("CLAUDECODE")
194 /// .collect()
195 /// .await?;
196 /// assert_eq!(result.code, 1); // printenv reports the variable as missing
197 /// # Ok(())
198 /// # }
199 /// ```
200 pub fn env_remove(mut self, name: impl Into<String>) -> Self {
201 self.env_remove.push(name.into());
202 self
203 }
204
205 /// Remove several variables from the child's environment. See
206 /// [`env_remove`](Self::env_remove).
207 pub fn env_remove_all<I, S>(mut self, names: I) -> Self
208 where
209 I: IntoIterator<Item = S>,
210 S: Into<String>,
211 {
212 self.env_remove.extend(names.into_iter().map(Into::into));
213 self
214 }
215
216 /// Keep the child's stdin open for the lifetime of the stream, so that
217 /// [`OutputStream::write_stdin`] can send input in reaction to output.
218 ///
219 /// Without it stdin is the null device (or, with [`stdin`](Self::stdin),
220 /// a fixed buffer that is closed once written). Content set with `stdin`
221 /// or `stdin_bytes` is still written first. Finish with
222 /// [`OutputStream::close_stdin`] so the child reads end-of-file; dropping
223 /// the stream closes stdin as well.
224 ///
225 /// ```no_run
226 /// use command_stream::{OutputChunk, StreamingRunner};
227 ///
228 /// # async fn example() -> command_stream::Result<()> {
229 /// let mut stream = StreamingRunner::from_argv("cat", Vec::<String>::new())
230 /// .stdin_pipe()
231 /// .stream();
232 /// stream.write_stdin(b"ping\n").await?;
233 /// while let Some(chunk) = stream.next().await {
234 /// if let OutputChunk::Stdout(data) = chunk {
235 /// assert_eq!(data, b"ping\n");
236 /// stream.close_stdin().await?;
237 /// }
238 /// }
239 /// # Ok(())
240 /// # }
241 /// ```
242 pub fn stdin_pipe(mut self) -> Self {
243 self.stdin_pipe = true;
244 self
245 }
246
247 /// Bound the run: once `limit` has passed since the spawn, the child is
248 /// stopped with the [`kill_signal`](Self::kill_signal) (then `SIGKILL`
249 /// after the grace period), exactly as [`OutputStream::kill`] would.
250 ///
251 /// Mirrors the JavaScript `timeout` option. The `Exit` chunk carries the
252 /// `128 + signal` code (143 for the default `SIGTERM`) and
253 /// [`OutputStream::timed_out`] reports `true`.
254 pub fn timeout(mut self, limit: Duration) -> Self {
255 self.timeout = Some(limit);
256 self
257 }
258
259 /// Configure the signal used to stop the process when it is killed without
260 /// an explicit signal — i.e. [`OutputStream::kill`] or abandoning the
261 /// stream. Mirrors the JavaScript `killSignal` option (default `SIGTERM`).
262 ///
263 /// The reported exit code follows the conventional `128 + signal` mapping
264 /// (e.g. `SIGTERM` => 143, `SIGINT` => 130, `SIGKILL` => 137).
265 pub fn kill_signal(mut self, signal: impl Into<String>) -> Self {
266 self.kill_signal = signal.into();
267 self
268 }
269
270 /// Configure how long (in milliseconds) the child is given to handle the
271 /// kill signal before `SIGKILL` is sent. Mirrors the JavaScript `killGrace`
272 /// option (default 100ms).
273 ///
274 /// This is the window in which a child running its own `SIGTERM` handler
275 /// can shut down on its own terms. Set it to `0` to escalate immediately.
276 pub fn kill_grace_ms(mut self, ms: u64) -> Self {
277 self.kill_grace_ms = ms;
278 self
279 }
280
281 /// Configure the grace period (in milliseconds) to keep draining the stdio
282 /// pipes after the process exits before aborting lingering readers. Mirrors
283 /// the JavaScript `exitPumpGrace` option (default 100ms).
284 pub fn exit_pump_grace_ms(mut self, ms: u64) -> Self {
285 self.exit_pump_grace_ms = ms;
286 self
287 }
288
289 /// Take the working directory and the environment the child is spawned
290 /// with, including the `prefer_local` PATH.
291 fn take_child_env(&mut self) -> (Option<PathBuf>, ChildEnv) {
292 let cwd = self.cwd.take();
293 let mut vars = self.env.take();
294 let local_cwd = cwd
295 .clone()
296 .or_else(|| std::env::current_dir().ok())
297 .unwrap_or_else(|| PathBuf::from("."));
298 if let Some((key, path)) =
299 crate::local_bin::preferred_path(vars.as_ref(), &local_cwd, &self.prefer_local)
300 {
301 vars.get_or_insert_with(|| std::env::vars().collect())
302 .insert(key, path);
303 }
304 let env = ChildEnv {
305 vars,
306 clear: self.env_clear,
307 remove: std::mem::take(&mut self.env_remove),
308 };
309 (cwd, env)
310 }
311
312 /// Start the command detached from the caller and return its pid, without
313 /// waiting for it (issue #232).
314 ///
315 /// The child gets null stdio, leads a process group of its own on Unix
316 /// (`CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS` on Windows) and is not
317 /// killed when anything is dropped, so it keeps running after the caller
318 /// returns or exits. [`cwd`](Self::cwd), [`env`](Self::env),
319 /// [`clear_env`](Self::clear_env), [`env_remove`](Self::env_remove) and
320 /// [`prefer_local`](Self::prefer_local) apply as usual; stdin content,
321 /// [`stdin_pipe`](Self::stdin_pipe), [`timeout`](Self::timeout) and the
322 /// kill settings do not, since nothing supervises the child. A background
323 /// thread reaps it when it exits. No Tokio runtime is needed.
324 ///
325 /// ```no_run
326 /// use command_stream::StreamingRunner;
327 ///
328 /// let pid = StreamingRunner::from_argv("sleep", ["60"]).cwd("/tmp").detach()?;
329 /// println!("started in the background as pid {pid:?}");
330 /// # Ok::<(), command_stream::Error>(())
331 /// ```
332 pub fn detach(mut self) -> Result<Option<u32>> {
333 let (cwd, env) = self.take_child_env();
334 crate::detach::spawn_detached(&self.command, cwd.as_deref(), &env)
335 }
336
337 pub(crate) fn spawn(mut self) -> (OutputStream, JoinHandle<Result<()>>) {
338 let (tx, rx) = mpsc::channel(1024);
339 // Unbounded so a synchronous Drop can request a kill without awaiting.
340 let (kill_tx, kill_rx) = mpsc::unbounded_channel::<String>();
341 // The child is spawned inside the task below, so its id is not known
342 // when this returns. The task publishes it here as soon as the spawn
343 // succeeds; `OutputStream::pid` reads the latest value (issue #18).
344 let (pid_tx, pid_rx) = watch::channel(None);
345 let (exit_signal_tx, exit_signal_rx) = watch::channel(None);
346
347 // Spawn the process handling task
348 let command = self.command.clone();
349 let (cwd, env) = self.take_child_env();
350 let (timed_out_tx, timed_out_rx) = watch::channel(false);
351 let (stdin_tx, stdin_requests) = if self.stdin_pipe {
352 let (sender, receiver) = crate::stream_stdin::channel();
353 (Some(sender), Some(receiver))
354 } else {
355 (None, None)
356 };
357 let setup = ProcessSetup {
358 command,
359 cwd,
360 env,
361 stdin_content: self.stdin_content.take(),
362 stdin_requests,
363 timeout: self.timeout,
364 kill_signal: self.kill_signal.clone(),
365 };
366 let grace = GraceWindows {
367 exit_pump_ms: self.exit_pump_grace_ms,
368 kill_ms: self.kill_grace_ms,
369 };
370 let kill_signal = self.kill_signal.clone();
371
372 let task = tokio::spawn(async move {
373 let channels = StreamChannels {
374 output_tx: tx,
375 kill_rx,
376 pid_tx,
377 exit_signal_tx,
378 timed_out_tx,
379 };
380 let result = run_streaming_process(setup, grace, channels).await;
381 if let Err(error) = &result {
382 trace_lazy("StreamingRunner", || format!("Error: {error}"));
383 }
384 result
385 });
386
387 (
388 OutputStream {
389 rx,
390 kill_tx,
391 kill_signal,
392 killed: false,
393 pid_rx,
394 exit_signal_rx,
395 timed_out_rx,
396 stdin_tx,
397 },
398 task,
399 )
400 }
401
402 /// Start the process and return a stream of output chunks
403 pub fn stream(self) -> OutputStream {
404 self.spawn().0
405 }
406
407 /// Run to completion and collect all output
408 pub async fn collect(self) -> Result<CommandResult> {
409 let stdin_content = self.stdin_content.clone();
410 let mut stdout = Vec::new();
411 let mut stderr = Vec::new();
412 let mut exit_code = 0;
413
414 let (mut stream, task) = self.spawn();
415 while let Some(chunk) = stream.rx.recv().await {
416 match chunk {
417 OutputChunk::Stdout(data) => stdout.extend(data),
418 OutputChunk::Stderr(data) => stderr.extend(data),
419 OutputChunk::Exit(code) => exit_code = code,
420 }
421 }
422
423 task.await.map_err(|error| {
424 std::io::Error::other(format!("streaming process task failed: {error}"))
425 })??;
426
427 let mut result = CommandResult::new(
428 String::from_utf8_lossy(&stdout).to_string(),
429 String::from_utf8_lossy(&stderr).to_string(),
430 exit_code,
431 );
432 if let Some(content) = stdin_content {
433 result.stdin = crate::result_streams::CapturedInput::new(content);
434 }
435 Ok(result)
436 }
437
438 /// Run an exact-argument command to completion from synchronous code.
439 ///
440 /// Build the command with [`Self::from_argv`] and configure it with the
441 /// same `cwd`, `env`, and `stdin` methods used by [`Self::collect`]. Call
442 /// this outside a Tokio runtime; async callers should use `collect().await`.
443 pub fn collect_blocking(self) -> Result<CommandResult> {
444 if tokio::runtime::Handle::try_current().is_ok() {
445 return Err(std::io::Error::other(
446 "collect_blocking cannot run inside a Tokio runtime; use collect().await",
447 )
448 .into());
449 }
450 let runtime = tokio::runtime::Runtime::new()?;
451 runtime.block_on(self.collect())
452 }
453}
454
455/// Stream of output chunks from a process
456pub struct OutputStream {
457 rx: mpsc::Receiver<OutputChunk>,
458 kill_tx: mpsc::UnboundedSender<String>,
459 kill_signal: String,
460 killed: bool,
461 pid_rx: watch::Receiver<Option<u32>>,
462 exit_signal_rx: watch::Receiver<Option<String>>,
463 timed_out_rx: watch::Receiver<bool>,
464 stdin_tx: Option<StdinSender>,
465}
466
467impl OutputStream {
468 /// Receive the next chunk
469 pub async fn next(&mut self) -> Option<OutputChunk> {
470 self.rx.recv().await
471 }
472
473 /// Signal reported by the native child, available once it exits.
474 /// Ordinary exit codes above 128 are not mistaken for signals.
475 pub fn exit_signal(&self) -> Option<String> {
476 self.exit_signal_rx.borrow().clone()
477 }
478
479 /// Whether the [`StreamingRunner::timeout`] elapsed and stopped the child.
480 ///
481 /// It becomes `true` before the terminating `Exit` chunk is delivered, so
482 /// it can be read as soon as that chunk arrives. An explicit kill, or a
483 /// child exiting on its own, leaves it `false`.
484 pub fn timed_out(&self) -> bool {
485 *self.timed_out_rx.borrow()
486 }
487
488 /// Write bytes to the child's stdin while it runs, and flush them.
489 ///
490 /// Requires [`StreamingRunner::stdin_pipe`]. The write is carried out by
491 /// the task that owns the pipe, so the output readers keep draining
492 /// stdout and stderr meanwhile. Fails with a broken-pipe error when stdin
493 /// was never piped, has been closed, or the child has exited.
494 pub async fn write_stdin(&self, data: impl AsRef<[u8]>) -> Result<()> {
495 let data = data.as_ref().to_vec();
496 crate::stream_stdin::request(self.stdin_tx.as_ref(), |ack| StdinRequest::Write(data, ack))
497 .await?;
498 Ok(())
499 }
500
501 /// Close the child's stdin so it reads end-of-file.
502 ///
503 /// Later [`write_stdin`](Self::write_stdin) calls fail. Closing twice, or
504 /// after the child has exited, is a no-op.
505 pub async fn close_stdin(&mut self) -> Result<()> {
506 let Some(sender) = self.stdin_tx.take() else {
507 return Ok(());
508 };
509 match crate::stream_stdin::request(Some(&sender), StdinRequest::Close).await {
510 Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
511 other => Ok(other?),
512 }
513 }
514
515 /// Process id of the streamed command, as currently known.
516 ///
517 /// The child is spawned by a background task, so this is `None` for the
518 /// short window between [`StreamingRunner::stream`] returning and the spawn
519 /// completing, and stays `None` if the spawn failed. From the first
520 /// delivered chunk onwards it is set, and it remains readable after the
521 /// process has exited. Use [`wait_for_pid`](Self::wait_for_pid) to avoid
522 /// the startup window.
523 pub fn pid(&self) -> Option<u32> {
524 *self.pid_rx.borrow()
525 }
526
527 /// Process id of the streamed command, waiting for the spawn to complete.
528 ///
529 /// Resolves as soon as the child exists, and returns `None` if the process
530 /// could never be spawned. This is the streaming counterpart of awaiting a
531 /// stream before reading `runner.pid` in JavaScript.
532 pub async fn wait_for_pid(&mut self) -> Option<u32> {
533 // `wait_for` checks the current value first, so an already-published id
534 // returns without waiting. An error means the sending task is gone,
535 // which only happens when the spawn failed.
536 match self.pid_rx.wait_for(|pid| pid.is_some()).await {
537 Ok(pid) => *pid,
538 Err(_) => None,
539 }
540 }
541
542 /// Stop the process using the configured kill signal (default `SIGTERM`).
543 ///
544 /// This can be called from inside the consumption loop to stop a
545 /// long-running or endless process; a terminating `OutputChunk::Exit` is
546 /// still delivered afterwards.
547 pub fn kill(&mut self) {
548 let signal = self.kill_signal.clone();
549 self.kill_with(&signal);
550 }
551
552 /// Stop the process using an explicit signal, overriding the configured
553 /// kill signal for this call.
554 pub fn kill_with(&mut self, signal: &str) {
555 if self.killed {
556 return;
557 }
558 self.killed = true;
559 trace_lazy("OutputStream", || format!("kill | signal={}", signal));
560 // Best effort: the task may have already finished, in which case the
561 // receiver is gone and the send fails harmlessly.
562 let _ = self.kill_tx.send(signal.to_string());
563 }
564
565 /// Collect all remaining output into vectors
566 pub async fn collect(mut self) -> (Vec<u8>, Vec<u8>, i32) {
567 let mut stdout = Vec::new();
568 let mut stderr = Vec::new();
569 let mut exit_code = 0;
570
571 while let Some(chunk) = self.rx.recv().await {
572 match chunk {
573 OutputChunk::Stdout(data) => stdout.extend(data),
574 OutputChunk::Stderr(data) => stderr.extend(data),
575 OutputChunk::Exit(code) => exit_code = code,
576 }
577 }
578
579 (stdout, stderr, exit_code)
580 }
581
582 /// Collect stdout only, discarding stderr
583 pub async fn collect_stdout(mut self) -> Vec<u8> {
584 let mut stdout = Vec::new();
585
586 while let Some(chunk) = self.rx.recv().await {
587 if let OutputChunk::Stdout(data) = chunk {
588 stdout.extend(data);
589 }
590 }
591
592 stdout
593 }
594}
595
596impl Drop for OutputStream {
597 fn drop(&mut self) {
598 // Abandoning the stream (e.g. `break`-ing out of the loop) must stop the
599 // process, matching the JavaScript iterator's `finally` cleanup. If the
600 // process already finished this is a harmless no-op.
601 if !self.killed {
602 let _ = self.kill_tx.send(self.kill_signal.clone());
603 }
604 }
605}
606
607/// Async iterator trait for output streams
608#[async_trait::async_trait]
609pub trait AsyncIterator {
610 type Item;
611
612 /// Get the next item from the iterator
613 async fn next(&mut self) -> Option<Self::Item>;
614}
615
616#[async_trait::async_trait]
617impl AsyncIterator for OutputStream {
618 type Item = OutputChunk;
619
620 async fn next(&mut self) -> Option<Self::Item> {
621 self.rx.recv().await
622 }
623}
624
625/// Extension trait to convert ProcessRunner into a stream
626pub trait IntoStream {
627 /// Convert into an output stream
628 fn into_stream(self) -> OutputStream;
629}
630
631impl IntoStream for crate::ProcessRunner {
632 fn into_stream(self) -> OutputStream {
633 let streaming = StreamingRunner::new(self.command().to_string());
634 streaming.stream()
635 }
636}