vtcode_bash_runner/
process.rs

1//! Unified process handle types for PTY and pipe backends.
2//!
3//! This module provides abstractions for interacting with spawned processes
4//! regardless of whether they use a PTY or regular pipes.
5//!
6//! Inspired by [codex-rs] PTY process handle patterns (Apache-2.0).
7//! Copyright 2025 OpenAI. See the repository `THIRD-PARTY-NOTICES` file for
8//! full attribution.
9//!
10//! [codex-rs]: https://github.com/openai/codex
11//!
12//! ## Async-drop pattern
13//!
14//! `Drop` cannot be `async`, but some cleanup requires async operations
15//! (e.g. asking a runtime to stop a process, removing a container, closing
16//! a network connection). The pattern borrowed from `testcontainers-rs` is to
17//! spin a dedicated thread inside `Drop`, create a temporary Tokio runtime on
18//! that thread, and block the `Drop` call until the future resolves. This is
19//! heavier than a true `async drop`, but it lets us bridge sync `Drop` into
20//! async cleanup without requiring nightly Rust.
21
22use std::fmt;
23use std::io;
24use std::sync::Arc;
25use std::sync::Mutex as StdMutex;
26use std::sync::atomic::{AtomicBool, Ordering};
27
28use bytes::Bytes;
29use tokio::sync::{broadcast, mpsc, oneshot};
30use tokio::task::{AbortHandle, JoinHandle};
31
32/// Run an async cleanup future from synchronous `Drop`.
33///
34/// This bridges the gap between sync `Drop` and async resource cleanup.
35/// A dedicated thread is spawned with its own Tokio runtime so the future
36/// can make full use of async APIs. The `Drop` call blocks until the runtime
37/// shuts down, giving us deterministic cleanup semantics similar to RAII.
38///
39/// Borrowed from the `testcontainers-rs` pattern for async-drop in Rust
40/// (where true `async drop` is still nightly-only).
41pub(crate) fn async_drop<F, Fut>(f: F)
42where
43    F: FnOnce() -> Fut + Send + 'static,
44    Fut: Future<Output = ()> + Send + 'static,
45{
46    let handle = std::thread::spawn(move || {
47        let rt = match tokio::runtime::Runtime::new() {
48            Ok(rt) => rt,
49            Err(_) => return,
50        };
51        rt.block_on(f());
52    });
53    let _ = handle.join();
54}
55
56/// Trait for process termination strategies.
57///
58/// Different backends (PTY vs pipe) may need different termination approaches.
59pub trait ChildTerminator: Send + Sync {
60    /// Kill the child process.
61    fn kill(&mut self) -> io::Result<()>;
62}
63
64/// Optional PTY-specific handles that must be preserved.
65///
66/// For PTY processes, the slave handle must be kept alive because the process
67/// will receive SIGHUP if it's closed.
68pub struct PtyHandles {
69    /// The slave PTY handle (kept alive to prevent SIGHUP).
70    pub _slave: Option<Box<dyn Send>>,
71    /// The master PTY handle.
72    pub _master: Box<dyn Send>,
73}
74
75impl fmt::Debug for PtyHandles {
76    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77        f.debug_struct("PtyHandles").finish()
78    }
79}
80
81/// Handle for driving an interactive or non-interactive process.
82///
83/// This provides a unified interface for both PTY and pipe-based processes:
84/// - Write to stdin via `writer_sender()`
85/// - Read merged stdout/stderr via `output_receiver()`
86/// - Check exit status via `has_exited()` and `exit_code()`
87/// - Clean up via `terminate()`
88pub struct ProcessHandle {
89    writer_tx: mpsc::Sender<Vec<u8>>,
90    output_tx: broadcast::Sender<Bytes>,
91    killer: StdMutex<Option<Box<dyn ChildTerminator>>>,
92    reader_handle: StdMutex<Option<JoinHandle<()>>>,
93    reader_abort_handles: StdMutex<Vec<AbortHandle>>,
94    writer_handle: StdMutex<Option<JoinHandle<()>>>,
95    wait_handle: StdMutex<Option<JoinHandle<()>>>,
96    exit_status: Arc<AtomicBool>,
97    exit_code: Arc<StdMutex<Option<i32>>>,
98    // PTY handles must be preserved to prevent the process from receiving Control+C
99    _pty_handles: StdMutex<Option<PtyHandles>>,
100}
101
102impl fmt::Debug for ProcessHandle {
103    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104        f.debug_struct("ProcessHandle")
105            .field("has_exited", &self.has_exited())
106            .field("exit_code", &self.exit_code())
107            .finish()
108    }
109}
110
111impl ProcessHandle {
112    /// Create a new process handle with all required components.
113    #[allow(clippy::too_many_arguments)]
114    pub(crate) fn new(
115        writer_tx: mpsc::Sender<Vec<u8>>,
116        output_tx: broadcast::Sender<Bytes>,
117        initial_output_rx: broadcast::Receiver<Bytes>,
118        killer: Box<dyn ChildTerminator>,
119        reader_handle: JoinHandle<()>,
120        reader_abort_handles: Vec<AbortHandle>,
121        writer_handle: JoinHandle<()>,
122        wait_handle: JoinHandle<()>,
123        exit_status: Arc<AtomicBool>,
124        exit_code: Arc<StdMutex<Option<i32>>>,
125        pty_handles: Option<PtyHandles>,
126    ) -> (Self, broadcast::Receiver<Bytes>) {
127        (
128            Self {
129                writer_tx,
130                output_tx,
131                killer: StdMutex::new(Some(killer)),
132                reader_handle: StdMutex::new(Some(reader_handle)),
133                reader_abort_handles: StdMutex::new(reader_abort_handles),
134                writer_handle: StdMutex::new(Some(writer_handle)),
135                wait_handle: StdMutex::new(Some(wait_handle)),
136                exit_status,
137                exit_code,
138                _pty_handles: StdMutex::new(pty_handles),
139            },
140            initial_output_rx,
141        )
142    }
143
144    /// Returns a channel sender for writing raw bytes to the child stdin.
145    ///
146    /// # Example
147    /// ```ignore
148    /// let writer = handle.writer_sender();
149    /// writer.send(b"input\n".to_vec()).await?;
150    /// ```
151    #[inline]
152    pub fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
153        self.writer_tx.clone()
154    }
155
156    /// Returns a broadcast receiver that yields stdout/stderr chunks.
157    ///
158    /// Multiple receivers can be created; each receives all output from the
159    /// point of subscription.
160    #[inline]
161    pub fn output_receiver(&self) -> broadcast::Receiver<Bytes> {
162        self.output_tx.subscribe()
163    }
164
165    /// True if the child process has exited.
166    #[inline]
167    pub fn has_exited(&self) -> bool {
168        self.exit_status.load(Ordering::SeqCst)
169    }
170
171    /// Returns the exit code if the process has exited.
172    #[inline]
173    pub fn exit_code(&self) -> Option<i32> {
174        *self.exit_code.lock().unwrap_or_else(|e| e.into_inner())
175    }
176
177    /// True once the stdout/stderr reader task has drained the child streams.
178    #[inline]
179    pub fn is_output_drained(&self) -> bool {
180        self.reader_handle
181            .lock()
182            .ok()
183            .and_then(|guard| guard.as_ref().map(JoinHandle::is_finished))
184            .unwrap_or(true)
185    }
186
187    /// Attempts to kill the child and abort helper tasks.
188    ///
189    /// This is idempotent and safe to call multiple times.
190    pub fn terminate(&self) {
191        self.terminate_internal();
192    }
193
194    /// Internal termination that aborts all tasks.
195    fn terminate_internal(&self) {
196        // Kill the child process
197        if let Ok(mut killer_opt) = self.killer.lock()
198            && let Some(mut killer) = killer_opt.take()
199        {
200            let _ = killer.kill();
201        }
202
203        self.abort_tasks();
204    }
205
206    /// Abort all background tasks associated with this process.
207    fn abort_tasks(&self) {
208        // Abort reader handle
209        if let Ok(mut h) = self.reader_handle.lock()
210            && let Some(handle) = h.take()
211        {
212            handle.abort();
213        }
214
215        // Abort individual reader abort handles
216        if let Ok(mut handles) = self.reader_abort_handles.lock() {
217            for handle in handles.drain(..) {
218                handle.abort();
219            }
220        }
221
222        // Abort writer handle
223        if let Ok(mut h) = self.writer_handle.lock()
224            && let Some(handle) = h.take()
225        {
226            handle.abort();
227        }
228
229        // Abort wait handle
230        if let Ok(mut h) = self.wait_handle.lock()
231            && let Some(handle) = h.take()
232        {
233            handle.abort();
234        }
235    }
236
237    /// Check if the process is still running.
238    #[inline]
239    pub fn is_running(&self) -> bool {
240        !self.has_exited() && !self.is_writer_closed()
241    }
242
243    /// Send bytes to the process stdin.
244    ///
245    /// Returns an error if the stdin channel is closed.
246    pub async fn write(&self, bytes: impl Into<Vec<u8>>) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
247        self.writer_tx.send(bytes.into()).await
248    }
249
250    /// Check if the writer channel is closed.
251    #[inline]
252    pub fn is_writer_closed(&self) -> bool {
253        self.writer_tx.is_closed()
254    }
255}
256
257impl Drop for ProcessHandle {
258    fn drop(&mut self) {
259        // Use the async-drop pattern so cleanup can block on async waits
260        // (e.g. waiting for the OS to reap the child) without blocking the
261        // caller's thread. This mirrors testcontainers-rs's approach for
262        // async resource cleanup from synchronous Drop.
263        //
264        // We must take ownership of the inner values here because the async
265        // block needs to own everything it captures.
266        let killer = self.killer.lock().ok().and_then(|mut g| g.take());
267        let mut reader_handle = self.reader_handle.lock().ok().and_then(|mut g| g.take());
268        let reader_abort_handles = self
269            .reader_abort_handles
270            .lock()
271            .ok()
272            .map(|mut g| g.drain(..).collect::<Vec<_>>());
273        let mut writer_handle = self.writer_handle.lock().ok().and_then(|mut g| g.take());
274        let mut wait_handle = self.wait_handle.lock().ok().and_then(|mut g| g.take());
275
276        async_drop(move || async move {
277            if let Some(mut killer) = killer {
278                let _ = killer.kill();
279            }
280            if let Some(handle) = reader_handle.take() {
281                handle.abort();
282            }
283            if let Some(handle) = writer_handle.take() {
284                handle.abort();
285            }
286            if let Some(handle) = wait_handle.take() {
287                handle.abort();
288            }
289            if let Some(handles) = reader_abort_handles {
290                for handle in handles {
291                    handle.abort();
292                }
293            }
294        });
295    }
296}
297
298/// Return value from spawn helpers (PTY or pipe).
299///
300/// Bundles the process handle with receivers for output and exit notification.
301#[derive(Debug)]
302pub struct SpawnedProcess {
303    /// Handle for interacting with the process.
304    pub session: ProcessHandle,
305    /// Receiver for stdout/stderr output chunks.
306    pub output_rx: broadcast::Receiver<Bytes>,
307    /// Receiver for exit code (receives once when process exits).
308    pub exit_rx: oneshot::Receiver<i32>,
309}
310
311impl SpawnedProcess {
312    /// Convenience method to wait for the process to exit and collect output.
313    ///
314    /// Returns (collected_output, exit_code).
315    pub async fn wait_with_output(self, timeout_ms: u64) -> (Vec<u8>, i32) {
316        collect_output_until_exit(self.output_rx, self.exit_rx, timeout_ms).await
317    }
318}
319
320/// Collect output from a process until it exits or times out.
321///
322/// This is useful for tests and simple use cases where you want all output.
323pub async fn collect_output_until_exit(
324    mut output_rx: broadcast::Receiver<Bytes>,
325    exit_rx: oneshot::Receiver<i32>,
326    timeout_ms: u64,
327) -> (Vec<u8>, i32) {
328    let mut collected = Vec::new();
329    let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(timeout_ms);
330    tokio::pin!(exit_rx);
331
332    loop {
333        tokio::select! {
334            res = output_rx.recv() => {
335                if let Ok(chunk) = res {
336                    collected.extend_from_slice(&chunk);
337                }
338            }
339            res = &mut exit_rx => {
340                let code = res.unwrap_or(-1);
341                // Drain remaining output briefly after exit
342                let quiet = tokio::time::Duration::from_millis(50);
343                let max_deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(500);
344
345                while tokio::time::Instant::now() < max_deadline {
346                    match tokio::time::timeout(quiet, output_rx.recv()).await {
347                        Ok(Ok(chunk)) => collected.extend_from_slice(&chunk),
348                        Ok(Err(broadcast::error::RecvError::Lagged(count))) => {
349                            eprintln!("[vtcode] output stream lagged ({count} dropped)");
350                            continue;
351                        }
352                        Ok(Err(broadcast::error::RecvError::Closed)) => break,
353                        Err(_) => break, // Timeout - quiet period reached
354                    }
355                }
356                return (collected, code);
357            }
358            _ = tokio::time::sleep_until(deadline) => {
359                return (collected, -1);
360            }
361        }
362    }
363}
364
365/// Backwards-compatible alias for ProcessHandle.
366pub type ExecCommandSession = ProcessHandle;
367
368/// Backwards-compatible alias for SpawnedProcess.
369pub type SpawnedPty = SpawnedProcess;
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    struct NoopTerminator;
376    impl ChildTerminator for NoopTerminator {
377        fn kill(&mut self) -> io::Result<()> {
378            Ok(())
379        }
380    }
381
382    #[tokio::test]
383    async fn test_process_handle_debug() {
384        // Just verify Debug impl doesn't panic
385        let exit_status = Arc::new(AtomicBool::new(false));
386        let exit_code = Arc::new(StdMutex::new(None));
387
388        let (writer_tx, _) = mpsc::channel(1);
389        let (output_tx, initial_rx) = broadcast::channel(1);
390
391        let (handle, _) = ProcessHandle::new(
392            writer_tx,
393            output_tx,
394            initial_rx,
395            Box::new(NoopTerminator),
396            tokio::spawn(async {}),
397            vec![],
398            tokio::spawn(async {}),
399            tokio::spawn(async {}),
400            exit_status,
401            exit_code,
402            None,
403        );
404
405        let debug_str = format!("{handle:?}");
406        assert!(debug_str.contains("ProcessHandle"));
407    }
408
409    #[tokio::test]
410    async fn test_has_exited() {
411        let exit_status = Arc::new(AtomicBool::new(false));
412        let exit_code = Arc::new(StdMutex::new(None));
413
414        let (writer_tx, _) = mpsc::channel(1);
415        let (output_tx, initial_rx) = broadcast::channel(1);
416
417        let (handle, _) = ProcessHandle::new(
418            writer_tx,
419            output_tx,
420            initial_rx,
421            Box::new(NoopTerminator),
422            tokio::spawn(async {}),
423            vec![],
424            tokio::spawn(async {}),
425            tokio::spawn(async {}),
426            Arc::clone(&exit_status),
427            exit_code,
428            None,
429        );
430
431        assert!(!handle.has_exited());
432        exit_status.store(true, Ordering::SeqCst);
433        assert!(handle.has_exited());
434    }
435}
vtcode_bash_runner/process.rs

vtcode_bash_runner/
process.rs
