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    /// What the model asked for. A build or a full test run needs minutes, while
18    /// most commands need seconds, and only the caller knows which this is.
19    #[serde(default)]
20    pub timeout_seconds: Option<u64>,
21}
22
23#[derive(Debug, Clone)]
24pub struct BashOptions {
25    /// Used when the call does not ask for a timeout.
26    pub timeout: Duration,
27    /// Ceiling for what a call may ask for, so a hung command cannot run
28    /// unbounded.
29    pub max_timeout: Duration,
30    pub max_output_bytes: usize,
31}
32
33impl BashOptions {
34    /// The timeout this call runs under, and the requested value if it had to be
35    /// capped. Callers report the cap so the model can see why its request did
36    /// not take effect.
37    fn resolve_timeout(&self, args: &BashArgs) -> (Duration, Option<u64>) {
38        let Some(requested) = args.timeout_seconds.filter(|seconds| *seconds > 0) else {
39            return (self.timeout.min(self.max_timeout), None);
40        };
41        let requested_duration = Duration::from_secs(requested);
42        if requested_duration > self.max_timeout {
43            (self.max_timeout, Some(requested))
44        } else {
45            (requested_duration, None)
46        }
47    }
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct BashResult {
52    pub output: String,
53    pub exit_code: Option<i32>,
54    pub timed_out: bool,
55    pub cancelled: bool,
56}
57
58pub type OutputSink = Arc<dyn Fn(String) + Send + Sync>;
59
60pub async fn execute_bash(
61    root: &Path,
62    args: &BashArgs,
63    options: &BashOptions,
64    output_sink: Option<OutputSink>,
65) -> Result<BashResult> {
66    execute_bash_cancellable(root, args, options, output_sink, CancellationToken::new()).await
67}
68
69pub async fn execute_bash_cancellable(
70    root: &Path,
71    args: &BashArgs,
72    options: &BashOptions,
73    output_sink: Option<OutputSink>,
74    cancel: CancellationToken,
75) -> Result<BashResult> {
76    let mut command = Command::new("bash");
77    command
78        .arg("-o")
79        .arg("pipefail")
80        .arg("-lc")
81        .arg(&args.command)
82        .current_dir(root)
83        .stdin(Stdio::null())
84        .stdout(Stdio::piped())
85        .stderr(Stdio::piped())
86        .kill_on_drop(true);
87    #[cfg(unix)]
88    command.process_group(0);
89    let mut child = command
90        .spawn()
91        .with_context(|| format!("start bash command: {}", args.command))?;
92
93    let stdout = child.stdout.take().context("capture bash stdout")?;
94    let stderr = child.stderr.take().context("capture bash stderr")?;
95    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(32);
96    tokio::spawn(pump(stdout, tx.clone()));
97    tokio::spawn(pump(stderr, tx));
98
99    let (timeout, capped_from) = options.resolve_timeout(args);
100    let deadline = Instant::now() + timeout;
101    let mut bounded = BoundedOutput::new(options.max_output_bytes);
102    let mut timed_out = false;
103    let mut cancelled = false;
104    let status = loop {
105        tokio::select! {
106            status = child.wait() => break status.context("wait for bash command")?,
107            chunk = rx.recv() => {
108                if let Some(chunk) = chunk {
109                    if let Some(sink) = &output_sink {
110                        sink(String::from_utf8_lossy(&chunk).into_owned());
111                    }
112                    bounded.push(&chunk);
113                }
114            }
115            _ = sleep_until(deadline) => {
116                timed_out = true;
117                break terminate(&mut child).await.context("terminate timed out bash command")?;
118            }
119            _ = cancel.cancelled() => {
120                cancelled = true;
121                break terminate(&mut child).await.context("terminate cancelled bash command")?;
122            }
123        }
124    };
125    while let Some(chunk) = rx.recv().await {
126        if let Some(sink) = &output_sink {
127            sink(String::from_utf8_lossy(&chunk).into_owned());
128        }
129        bounded.push(&chunk);
130    }
131
132    let mut output = bounded.finish();
133    if timed_out {
134        // Say what the limit was and how to raise it: a bare "timed out" leaves
135        // the model guessing whether the command or the limit was wrong.
136        output.push_str(&format!(
137            "\n[bash timed out after {:.1}s",
138            timeout.as_secs_f64()
139        ));
140        match capped_from {
141            Some(requested) => output.push_str(&format!(
142                "; timeout_seconds {requested} was capped at {}]",
143                options.max_timeout.as_secs()
144            )),
145            None => output.push_str(&format!(
146                "; pass timeout_seconds up to {} for slow commands]",
147                options.max_timeout.as_secs()
148            )),
149        }
150    } else if cancelled {
151        output.push_str("\n[bash cancelled]");
152    }
153    Ok(BashResult {
154        output,
155        exit_code: status.code(),
156        timed_out,
157        cancelled,
158    })
159}
160
161async fn terminate(child: &mut tokio::process::Child) -> Result<std::process::ExitStatus> {
162    #[cfg(unix)]
163    if let Some(pid) = child.id() {
164        // SAFETY: this negative PID addresses the dedicated process group created above.
165        let result = unsafe { libc::kill(-(pid as i32), libc::SIGTERM) };
166        if result != 0 {
167            let error = std::io::Error::last_os_error();
168            if error.raw_os_error() != Some(libc::ESRCH) {
169                return Err(error.into());
170            }
171        }
172        match tokio::time::timeout(Duration::from_secs(1), child.wait()).await {
173            Ok(status) => return Ok(status?),
174            Err(_) => {
175                // SAFETY: the same process group is still owned by this child.
176                unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
177                return Ok(child.wait().await?);
178            }
179        }
180    }
181    child.kill().await?;
182    Ok(child.wait().await?)
183}
184
185async fn pump(mut stream: impl tokio::io::AsyncRead + Unpin, tx: mpsc::Sender<Vec<u8>>) {
186    let mut buffer = vec![0_u8; 8192];
187    loop {
188        match stream.read(&mut buffer).await {
189            Ok(0) | Err(_) => break,
190            Ok(count) if tx.send(buffer[..count].to_vec()).await.is_err() => break,
191            Ok(_) => {}
192        }
193    }
194}
195
196struct BoundedOutput {
197    bytes: Vec<u8>,
198    max: usize,
199    dropped: usize,
200}
201
202impl BoundedOutput {
203    fn new(max: usize) -> Self {
204        Self {
205            bytes: Vec::new(),
206            max,
207            dropped: 0,
208        }
209    }
210
211    fn push(&mut self, chunk: &[u8]) {
212        self.bytes.extend_from_slice(chunk);
213        if self.bytes.len() > self.max {
214            let remove = self.bytes.len() - self.max;
215            self.bytes.drain(..remove);
216            self.dropped += remove;
217        }
218    }
219
220    fn finish(self) -> String {
221        let tail = String::from_utf8_lossy(&self.bytes);
222        if self.dropped == 0 {
223            tail.into_owned()
224        } else {
225            format!("[output truncated; dropped {} bytes]\n{tail}", self.dropped)
226        }
227    }
228}
a_agent/tools/bash.rs

a_agent/tools/
bash.rs
