metrale_bench/benchmarks/agentic/
agent_shell.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Running the shell command the model wrote, and bounding its time,
4//! memory and output.
5//!
6//! Owner: bench, agentic.
7//! Invariants:
8//! - Each command runs as `sh -c` with the sandbox as its working directory
9//!   and, on unix, in its own process group.
10//! - While a command runs, at most `2 * CAPTURE_END` bytes of each stream are
11//!   held.
12//! - [`truncate`] returns at most `MAX_TOOL_OUTPUT` bytes.
13
14use 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
25/// 2026-09-26: Bytes of one stream kept at each end while the command runs
26/// (see [`Capture`]); the middle is counted and dropped. [`truncate`] later
27/// cuts the result to `MAX_TOOL_OUTPUT`.
28const CAPTURE_END: usize = 16 * MAX_TOOL_OUTPUT;
29
30/// 2026-09-26: Room reserved for the elision note, so [`truncate`] returns at
31/// most `MAX_TOOL_OUTPUT` bytes and a second `truncate` (the agent loop
32/// truncates every tool result) returns its input unchanged.
33const 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        // 2026-09-26: The child is killed if this future is dropped before it
47        // exits.
48        .kill_on_drop(true)
49        // 2026-09-26: `run_tier.sh` sets the same variable for opencode.
50        // `bench/fp8_dgx2_drift/harness/cargo-shim/cargo`, when it is on PATH,
51        // reads it to detach `cargo run`. It is set on this child only, so the
52        // scorer's own `cargo run` is not detached.
53        .env("METRALE_AGENT_SHELL", "1");
54    // 2026-09-26: Its own process group, so the timeout can kill everything the
55    // command started. A `setsid` server leaves the group; `agent::reap`
56    // handles those.
57    #[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    // 2026-09-26: Drain concurrently with the wait: output past the pipe buffer
66    // blocks the writer until someone reads it.
67    let pumps = (
68        tokio::spawn(pump(child.stdout.take(), Arc::clone(&out))),
69        tokio::spawn(pump(child.stderr.take(), Arc::clone(&err))),
70    );
71    // 2026-09-26: Wait for the process, not for end-of-pipe: a `setsid cmd &`
72    // that inherits stdout keeps the pipe open after `sh` has exited.
73    let status = match tokio::time::timeout(limit, child.wait()).await {
74        Ok(s) => Some(s?),
75        Err(_) => {
76            // 2026-09-26: Kill the whole group, so a child the command forked
77            // does not keep running until `agent::reap` at the end of the
78            // iteration.
79            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    // 2026-09-26: A pump still running after the grace is reading a pipe that
93    // something not killed holds open; abort it so it does not live on.
94    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        // 2026-09-26: Non-empty stderr is appended on every path, including the
100        // timeout.
101        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    // 2026-09-26: Normalise before truncating: `truncate` cuts at byte offsets
114    // and reports an elided count, so both would vary with the text that
115    // [`norm`] rewrites.
116    Ok(truncate(&norm::normalize(&text)))
117}
118
119/// 2026-09-26: Copies raw bytes into the capture, which decodes them once, so a
120/// UTF-8 sequence split across two reads is not mangled.
121async 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/// 2026-09-26: A bounded window over one stream: the first [`CAPTURE_END`]
133/// bytes and the last [`CAPTURE_END`], with the bytes between counted and
134/// dropped. Reading never stops to enforce the bound, because a writer blocked
135/// on a full pipe would be reported as a timeout.
136#[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    /// 2026-09-26: Decoded once, over the concatenation, so a multi-byte
163    /// character split across the head/tail seam is not mangled.
164    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/// 2026-09-26: SIGKILL the process group led by `pid` (the child was spawned
182/// with `process_group(0)`). The `--` ends option parsing, so `-<pid>` is read
183/// as a group. `a_timed_out_command_takes_the_children_it_forked_with_it`
184/// (agent_shell_tests.rs) checks that the group dies.
185#[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
200/// 2026-09-26: Keep the head and tail of text longer than `MAX_TOOL_OUTPUT`
201/// bytes, with a note giving the elided length in bytes. The result is at
202/// most `MAX_TOOL_OUTPUT` bytes, so truncating it again returns it unchanged.
203pub 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;