a_agent/tools/
bash.rs

1use std::path::Path;
2use std::process::Stdio;
3use std::sync::Arc;
4use std::time::Duration;
5
6use anyhow::{Context, Result};
7use serde::{Deserialize, Serialize};
8use tokio::io::AsyncReadExt;
9use tokio::process::Command;
10use tokio::sync::mpsc;
11use tokio::time::{Instant, sleep_until};
12use tokio_util::sync::CancellationToken;
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct BashArgs {
16    pub command: String,
17}
18
19#[derive(Debug, Clone)]
20pub struct BashOptions {
21    pub timeout: Duration,
22    pub max_output_bytes: usize,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct BashResult {
27    pub output: String,
28    pub exit_code: Option<i32>,
29    pub timed_out: bool,
30    pub cancelled: bool,
31}
32
33pub type OutputSink = Arc<dyn Fn(String) + Send + Sync>;
34
35pub async fn execute_bash(
36    root: &Path,
37    args: &BashArgs,
38    options: &BashOptions,
39    output_sink: Option<OutputSink>,
40) -> Result<BashResult> {
41    execute_bash_cancellable(root, args, options, output_sink, CancellationToken::new()).await
42}
43
44pub async fn execute_bash_cancellable(
45    root: &Path,
46    args: &BashArgs,
47    options: &BashOptions,
48    output_sink: Option<OutputSink>,
49    cancel: CancellationToken,
50) -> Result<BashResult> {
51    let mut command = Command::new("bash");
52    command
53        .arg("-o")
54        .arg("pipefail")
55        .arg("-lc")
56        .arg(&args.command)
57        .current_dir(root)
58        .stdin(Stdio::null())
59        .stdout(Stdio::piped())
60        .stderr(Stdio::piped())
61        .kill_on_drop(true);
62    #[cfg(unix)]
63    command.process_group(0);
64    let mut child = command
65        .spawn()
66        .with_context(|| format!("start bash command: {}", args.command))?;
67
68    let stdout = child.stdout.take().context("capture bash stdout")?;
69    let stderr = child.stderr.take().context("capture bash stderr")?;
70    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(32);
71    tokio::spawn(pump(stdout, tx.clone()));
72    tokio::spawn(pump(stderr, tx));
73
74    let deadline = Instant::now() + options.timeout;
75    let mut bounded = BoundedOutput::new(options.max_output_bytes);
76    let mut timed_out = false;
77    let mut cancelled = false;
78    let status = loop {
79        tokio::select! {
80            status = child.wait() => break status.context("wait for bash command")?,
81            chunk = rx.recv() => {
82                if let Some(chunk) = chunk {
83                    if let Some(sink) = &output_sink {
84                        sink(String::from_utf8_lossy(&chunk).into_owned());
85                    }
86                    bounded.push(&chunk);
87                }
88            }
89            _ = sleep_until(deadline) => {
90                timed_out = true;
91                break terminate(&mut child).await.context("terminate timed out bash command")?;
92            }
93            _ = cancel.cancelled() => {
94                cancelled = true;
95                break terminate(&mut child).await.context("terminate cancelled bash command")?;
96            }
97        }
98    };
99    while let Some(chunk) = rx.recv().await {
100        if let Some(sink) = &output_sink {
101            sink(String::from_utf8_lossy(&chunk).into_owned());
102        }
103        bounded.push(&chunk);
104    }
105
106    let mut output = bounded.finish();
107    if timed_out {
108        output.push_str(&format!(
109            "\n[bash timed out after {:.1}s]",
110            options.timeout.as_secs_f64()
111        ));
112    } else if cancelled {
113        output.push_str("\n[bash cancelled]");
114    }
115    Ok(BashResult {
116        output,
117        exit_code: status.code(),
118        timed_out,
119        cancelled,
120    })
121}
122
123async fn terminate(child: &mut tokio::process::Child) -> Result<std::process::ExitStatus> {
124    #[cfg(unix)]
125    if let Some(pid) = child.id() {
126        // SAFETY: this negative PID addresses the dedicated process group created above.
127        let result = unsafe { libc::kill(-(pid as i32), libc::SIGTERM) };
128        if result != 0 {
129            let error = std::io::Error::last_os_error();
130            if error.raw_os_error() != Some(libc::ESRCH) {
131                return Err(error.into());
132            }
133        }
134        match tokio::time::timeout(Duration::from_secs(1), child.wait()).await {
135            Ok(status) => return Ok(status?),
136            Err(_) => {
137                // SAFETY: the same process group is still owned by this child.
138                unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
139                return Ok(child.wait().await?);
140            }
141        }
142    }
143    child.kill().await?;
144    Ok(child.wait().await?)
145}
146
147async fn pump(mut stream: impl tokio::io::AsyncRead + Unpin, tx: mpsc::Sender<Vec<u8>>) {
148    let mut buffer = vec![0_u8; 8192];
149    loop {
150        match stream.read(&mut buffer).await {
151            Ok(0) | Err(_) => break,
152            Ok(count) if tx.send(buffer[..count].to_vec()).await.is_err() => break,
153            Ok(_) => {}
154        }
155    }
156}
157
158struct BoundedOutput {
159    bytes: Vec<u8>,
160    max: usize,
161    dropped: usize,
162}
163
164impl BoundedOutput {
165    fn new(max: usize) -> Self {
166        Self {
167            bytes: Vec::new(),
168            max,
169            dropped: 0,
170        }
171    }
172
173    fn push(&mut self, chunk: &[u8]) {
174        self.bytes.extend_from_slice(chunk);
175        if self.bytes.len() > self.max {
176            let remove = self.bytes.len() - self.max;
177            self.bytes.drain(..remove);
178            self.dropped += remove;
179        }
180    }
181
182    fn finish(self) -> String {
183        let tail = String::from_utf8_lossy(&self.bytes);
184        if self.dropped == 0 {
185            tail.into_owned()
186        } else {
187            format!("[output truncated; dropped {} bytes]\n{tail}", self.dropped)
188        }
189    }
190}
a_agent/tools/bash.rs

a_agent/tools/
bash.rs
