metrale_bench/benchmarks/agentic/
agent_shell.rs1use std::collections::VecDeque;
15use std::process::Stdio;
16use std::sync::Arc;
17use std::time::Duration;
18
19use anyhow::Result;
20use parking_lot::Mutex;
21use tokio::io::AsyncReadExt;
22
23use super::{AgentConfig, DRAIN_GRACE, MAX_TOOL_OUTPUT, norm};
24
25const CAPTURE_END: usize = 16 * MAX_TOOL_OUTPUT;
29
30const ELISION_NOTE: usize = 96;
34
35#[cfg(test)]
36pub(super) const TEST_ELISION_NOTE: usize = ELISION_NOTE;
37
38pub(crate) async fn run_shell(cfg: &AgentConfig, command: &str, limit: Duration) -> Result<String> {
39 let mut cmd = tokio::process::Command::new("sh");
40 cmd.arg("-c")
41 .arg(command)
42 .current_dir(&cfg.sandbox)
43 .stdin(Stdio::null())
44 .stdout(Stdio::piped())
45 .stderr(Stdio::piped())
46 .kill_on_drop(true)
49 .env("METRALE_AGENT_SHELL", "1");
54 #[cfg(unix)]
58 cmd.process_group(0);
59 if let Some(dir) = &cfg.cargo_target_dir {
60 cmd.env("CARGO_TARGET_DIR", dir);
61 }
62 let mut child = cmd.spawn()?;
63 let pid = child.id();
64 let (out, err) = (Arc::default(), Arc::default());
65 let pumps = (
68 tokio::spawn(pump(child.stdout.take(), Arc::clone(&out))),
69 tokio::spawn(pump(child.stderr.take(), Arc::clone(&err))),
70 );
71 let status = match tokio::time::timeout(limit, child.wait()).await {
74 Ok(s) => Some(s?),
75 Err(_) => {
76 if let Some(pid) = pid {
80 kill_group(pid).await;
81 }
82 let _ = child.kill().await;
83 None
84 }
85 };
86 let aborts = (pumps.0.abort_handle(), pumps.1.abort_handle());
87 let _ = tokio::time::timeout(DRAIN_GRACE, async {
88 let _ = pumps.0.await;
89 let _ = pumps.1.await;
90 })
91 .await;
92 aborts.0.abort();
95 aborts.1.abort();
96 let mut text = out.lock().text();
97 let stderr = err.lock().text();
98 if !stderr.trim().is_empty() {
99 text.push_str("\n[stderr]\n");
102 text.push_str(&stderr);
103 }
104 match status {
105 Some(s) if !s.success() => text.push_str(&format!("\n[exit {s}]")),
106 Some(_) => {}
107 None => text.push_str(&format!(
108 "\n[timed out after {}s and was killed; the output above is what it had produced. \
109 If this was a server, start it detached with its output redirected to a file.]",
110 limit.as_secs()
111 )),
112 }
113 Ok(truncate(&norm::normalize(&text)))
117}
118
119async fn pump<R: AsyncReadExt + Unpin>(reader: Option<R>, sink: Arc<Mutex<Capture>>) {
122 let Some(mut reader) = reader else { return };
123 let mut buf = [0u8; 8192];
124 while let Ok(n) = reader.read(&mut buf).await {
125 if n == 0 {
126 return;
127 }
128 sink.lock().push(&buf[..n]);
129 }
130}
131
132#[derive(Default)]
137pub(crate) struct Capture {
138 head: Vec<u8>,
139 tail: VecDeque<u8>,
140 dropped: usize,
141}
142
143impl Capture {
144 fn push(&mut self, mut bytes: &[u8]) {
145 if self.head.len() < CAPTURE_END {
146 let n = (CAPTURE_END - self.head.len()).min(bytes.len());
147 self.head.extend_from_slice(&bytes[..n]);
148 bytes = &bytes[n..];
149 }
150 if bytes.len() >= CAPTURE_END {
151 self.dropped += self.tail.len() + bytes.len() - CAPTURE_END;
152 self.tail.clear();
153 self.tail.extend(&bytes[bytes.len() - CAPTURE_END..]);
154 return;
155 }
156 let overflow = (self.tail.len() + bytes.len()).saturating_sub(CAPTURE_END);
157 self.dropped += overflow;
158 drop(self.tail.drain(..overflow));
159 self.tail.extend(bytes);
160 }
161
162 fn text(&self) -> String {
165 let mut bytes = self.head.clone();
166 if self.dropped > 0 {
167 bytes.extend_from_slice(
168 format!("\n… [{} bytes dropped from the middle] …\n", self.dropped).as_bytes(),
169 );
170 }
171 bytes.extend(self.tail.iter().copied());
172 String::from_utf8_lossy(&bytes).into_owned()
173 }
174
175 #[cfg(test)]
176 fn held(&self) -> usize {
177 self.head.len() + self.tail.len()
178 }
179}
180
181#[cfg(unix)]
186async fn kill_group(pid: u32) {
187 let _ = tokio::process::Command::new("kill")
188 .arg("-9")
189 .arg("--")
190 .arg(format!("-{pid}"))
191 .stdout(Stdio::null())
192 .stderr(Stdio::null())
193 .status()
194 .await;
195}
196
197#[cfg(not(unix))]
198async fn kill_group(_pid: u32) {}
199
200pub fn truncate(text: &str) -> String {
204 if text.len() <= MAX_TOOL_OUTPUT {
205 return text.to_string();
206 }
207 let keep = MAX_TOOL_OUTPUT - ELISION_NOTE;
208 let (mut cut, mut from) = (keep / 2, text.len() - keep / 2);
209 while !text.is_char_boundary(cut) {
210 cut -= 1;
211 }
212 while !text.is_char_boundary(from) {
213 from += 1;
214 }
215 let (head, tail) = (&text[..cut], &text[from..]);
216 let elided = text.len() - head.len() - tail.len();
217 format!("{head}\n… [{elided} characters elided from the middle] …\n{tail}")
218}
219
220#[cfg(test)]
221#[path = "agent_shell_tests.rs"]
222mod tests;