use std::borrow::Cow;
use std::collections::VecDeque;
use std::ffi::{OsStr, OsString};
use std::fs::File;
use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use rig_core::completion::ToolAnnotations;
use rig_core::redact;
use rig_core::tool::{
ExitStatus, KeptEnd, OutputCut, Tool, ToolContext, ToolExecutionError, args_schema,
};
use schemars::JsonSchema;
use serde::Deserialize;
use crate::process::{Ended, KillOnDrop, Launcher, OutputLine, Process, Spawn};
use crate::workspace::{Local, Workspace};
use crate::{Live, OutputLimit, Spill, blocking, clip};
const DEFAULT_TIMEOUT: u64 = 120;
const OUTPUT_LIMIT: u64 = 10 * 1024 * 1024;
#[derive(Clone, Debug)]
pub struct Shell {
pub workspace: Arc<dyn Workspace>,
pub launcher: Launcher,
pub live: Option<Live>,
pub spill: Option<Spill>,
pub cut: Option<OutputLimit>,
pub confine: bool,
pub max_timeout: u64,
pub quiet: Duration,
}
impl Default for Shell {
fn default() -> Self {
Self {
workspace: Arc::new(Local),
launcher: Launcher::default(),
live: None,
spill: None,
cut: Some(OutputLimit::default()),
confine: false,
max_timeout: 600,
quiet: Duration::from_secs(60),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ShellArgs {
pub command: String,
pub timeout_secs: Option<u64>,
pub cwd: Option<String>,
}
impl ShellArgs {
pub fn new(command: impl Into<String>) -> Self {
Self {
command: command.into(),
timeout_secs: None,
cwd: None,
}
}
pub fn timeout_secs(mut self, secs: u64) -> Self {
self.timeout_secs = Some(secs);
self
}
pub fn in_dir(mut self, dir: impl Into<String>) -> Self {
self.cwd = Some(dir.into());
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ShellRun {
pub output: String,
pub status: ExitStatus,
}
impl Shell {
pub async fn run(&self, args: ShellArgs) -> Result<ShellRun, ToolExecutionError> {
self.execute(args).await.map(|(ran, _)| ran)
}
async fn execute(
&self,
args: ShellArgs,
) -> Result<(ShellRun, Option<OutputCut>), ToolExecutionError> {
let timeout = Duration::from_secs(
args.timeout_secs
.unwrap_or(DEFAULT_TIMEOUT.min(self.max_timeout))
.clamp(1, self.max_timeout.max(1)),
);
let quiet = match args.timeout_secs {
None => Some(self.quiet).filter(|quiet| *quiet < timeout),
Some(_) => None,
};
let process = self.start(&args)?;
let _stop = KillOnDrop(process.group());
let shell = self.clone();
blocking(move || shell.finish(process, timeout, quiet)).await
}
}
impl Tool for Shell {
const NAME: &'static str = "shell";
type Args = ShellArgs;
type Output = String;
type Error = ToolExecutionError;
fn description(&self) -> String {
let kept = match self.cut {
Some(cut) => format!(
"; past {} lines or {} KB, its first and last lines and the lines with errors \
or failures from between",
cut.lines,
cut.bytes / 1024
),
None => String::new(),
};
format!(
"Run a shell command with `sh -c` in the current directory, or in cwd. stdin is \
empty. Returns the exit code and stdout and stderr interleaved line by line{kept}. \
The command is killed after timeout_secs (default {}, at most {}), or once it \
printed more than {} MB; with the default, also once it printed nothing for {}s, \
so name timeout_secs for a build or wait that is quiet for long; a quick command \
needs none. Use it to build, test, run programs and use git. Each call starts a fresh `sh`: a `cd` or a variable does not \
carry over to the next call. Do not start programs that wait for input or never end, such as editors, \
pagers or servers in the foreground. Do not assume a git repository: a failed \
`git` command ends a `&&` chain. Check formatting only of files you changed.",
DEFAULT_TIMEOUT.min(self.max_timeout).max(1),
self.max_timeout.max(1),
OUTPUT_LIMIT / 1024 / 1024,
self.quiet.as_secs()
)
}
fn parameters(&self) -> serde_json::Value {
args_schema::<ShellArgs>()
}
fn annotations(&self) -> ToolAnnotations {
ToolAnnotations::default().with_title("Shell")
}
async fn call(
&self,
context: &mut ToolContext,
args: ShellArgs,
) -> Result<String, ToolExecutionError> {
let (ran, cut) = self.execute(args).await?;
context.insert_result(ran.status)?;
if let Some(cut) = cut {
context.insert_result(cut)?;
}
Ok(ran.output)
}
}
impl Shell {
fn start(&self, args: &ShellArgs) -> Result<Process, ToolExecutionError> {
let base = self.launcher.cwd.as_deref();
let dir = match &args.cwd {
Some(dir) => {
if self.confine {
self.workspace.place(base, true, dir)?;
}
Some(
self.workspace
.join(base.unwrap_or(Path::new("")), Path::new(dir)),
)
}
None => base.map(Path::to_path_buf),
};
let how = Spawn {
stdin: None,
merge_stderr: false,
};
let line = [OsString::from("-c"), OsString::from(&args.command)];
self.workspace
.spawn(&self.launcher, dir.as_deref(), OsStr::new("sh"), &line, how)
.map_err(|error| ToolExecutionError::other(format!("could not run sh: {error}")))
}
fn finish(
&self,
mut process: Process,
timeout: Duration,
quiet: Option<Duration>,
) -> Result<(ShellRun, Option<OutputCut>), ToolExecutionError> {
let mut kept = Kept {
spill: self.spill.clone(),
cut: self.cut,
live: self.live.clone(),
..Kept::default()
};
let deadline = Instant::now().checked_add(timeout);
let end = process
.run(deadline, quiet, |line| {
kept.push(&line);
kept.written <= OUTPUT_LIMIT
})
.map_err(|error| {
ToolExecutionError::other(format!("waiting for the command failed: {error}"))
})?;
let silent = kept.written == 0;
let (mut output, cut) = kept.text();
let status = match end {
Ended::Exited(status) => match status.code() {
Some(code) => {
output.push_str(&format!("[exit code {code}]"));
ExitStatus::Exited(code)
}
None => {
output.push_str(&killed(status));
ExitStatus::Signaled
}
},
Ended::TimedOut => {
let secs = timeout.as_secs();
output.push_str(&format!(
"[timed out after {secs}s and stopped: {}",
stopped(&process)
));
if silent {
output.push_str(SILENT);
}
if secs < self.max_timeout {
output.push_str(&format!(
". Pass timeout_secs up to {} to wait longer",
self.max_timeout
));
}
output.push(']');
ExitStatus::TimedOut
}
Ended::Quiet => {
output.push_str(&format!(
"[printed nothing for {}s, so it was stopped: {}",
quiet.unwrap_or_default().as_secs(),
stopped(&process)
));
if silent {
output.push_str(SILENT);
}
output.push_str(
". A command that is stuck, such as one waiting for input or looping \
forever, needs fixing, not a longer wait; only one known to be quiet for \
long, such as a build, needs timeout_secs]",
);
ExitStatus::TimedOut
}
Ended::Stopped => {
output.push_str(&format!(
"[killed after printing more than {} MB: {}]",
OUTPUT_LIMIT / 1024 / 1024,
stopped(&process)
));
ExitStatus::Killed
}
};
Ok((ShellRun { output, status }, cut))
}
}
const SILENT: &str = ". It printed nothing: a program whose output is not a terminal may hold \
it back until it ends, so make it print as it goes (such as \
`python3 -u` or `stdbuf -oL`) to see how far it got";
fn stopped(process: &Process) -> String {
let grace = crate::process::STOP_GRACE.as_secs();
let how = if !cfg!(unix) {
"it was killed".to_owned()
} else if process.forced() {
format!("it got SIGTERM, then SIGKILL {grace}s later as it still ran")
} else {
"it got SIGTERM and ended".to_owned()
};
match process.remote() {
Some(host) => format!(
"{how}. That was the transport on this machine; whether the command on {host} \
stopped too is the transport's doing, so check there before running it again"
),
None => how,
}
}
const STATUS_STARTS: [&str; 5] = [
"[exit code ",
"[killed by ",
"[timed out after ",
"[printed nothing for ",
"[killed after printing more than ",
];
fn killed(status: std::process::ExitStatus) -> String {
#[cfg(unix)]
if let Some(signal) = std::os::unix::process::ExitStatusExt::signal(&status) {
let (name, why) = sent(signal);
return format!(
"[killed by {name} (signal {signal}), which this tool did not send: {why}. It was \
the process this call started here, `sh` or the last program it ran; when that \
program was a client of another machine, such as a remote shell, the command \
there may still be running or may have finished, so check there before running \
it again]"
);
}
#[cfg(not(unix))]
let _ = status;
"[killed by a signal]".to_owned()
}
#[cfg(unix)]
fn sent(signal: i32) -> (&'static str, &'static str) {
match signal {
libc::SIGHUP => ("SIGHUP", "its terminal or session went away"),
libc::SIGINT => ("SIGINT", "something on this machine interrupted it"),
libc::SIGQUIT => ("SIGQUIT", "something on this machine asked it to quit"),
libc::SIGILL => ("SIGILL", "it crashed"),
libc::SIGTRAP => ("SIGTRAP", "it crashed, or a debugger stopped it"),
libc::SIGABRT => ("SIGABRT", "it aborted itself, as on a failed assertion"),
libc::SIGBUS => ("SIGBUS", "it crashed"),
libc::SIGFPE => ("SIGFPE", "it crashed on an arithmetic error"),
libc::SIGKILL => (
"SIGKILL",
"something on this machine killed it, such as the out-of-memory killer or a `kill`",
),
libc::SIGSEGV => ("SIGSEGV", "it crashed"),
libc::SIGPIPE => ("SIGPIPE", "it wrote to a pipe that nothing read any more"),
libc::SIGTERM => ("SIGTERM", "something on this machine asked it to stop"),
_ => ("a signal", "something on this machine sent it"),
}
}
pub fn split_status(text: &str) -> (&str, Option<&str>) {
let start = text.rfind('\n').map_or(0, |newline| newline + 1);
match text.split_at_checked(start) {
Some((printed, last))
if last.ends_with(']') && STATUS_STARTS.iter().any(|at| last.starts_with(at)) =>
{
(printed, Some(last))
}
_ => (text, None),
}
}
const MAX_HELD: usize = 1024 * 1024;
const CONTEXT: usize = 3;
const MAX_PICKED: usize = 300;
fn matters(line: &str) -> bool {
let start = line.trim_start();
start.starts_with("error")
|| start.starts_with("Error")
|| [
"FAILED",
"FAIL:",
"failed",
"panicked",
"Traceback",
"Error:",
]
.iter()
.any(|word| line.contains(word))
}
fn plain(line: &str) -> Cow<'_, str> {
if !line.contains('\x1b') {
return Cow::Borrowed(line);
}
let mut text = String::with_capacity(line.len());
let mut chars = line.chars().peekable();
while let Some(c) = chars.next() {
if c != '\x1b' {
text.push(c);
continue;
}
match chars.next() {
Some('[') => {
let _ = chars.by_ref().find(|c| ('\x40'..='\x7e').contains(c));
}
Some(']' | 'P' | 'X' | '^' | '_') => {
while let Some(c) = chars.next() {
if c == '\x07' || (c == '\x1b' && chars.next_if_eq(&'\\').is_some()) {
break;
}
}
}
Some(' '..='/') => {
let _ = chars.by_ref().find(|c| !(' '..='/').contains(c));
}
_ => {}
}
}
Cow::Owned(text)
}
#[derive(Default)]
struct Kept {
written: u64,
cut: Option<OutputLimit>,
held: [String; 2],
lines: usize,
head: Vec<(usize, String)>,
head_bytes: usize,
picked: Vec<(usize, String)>,
picked_bytes: usize,
later: VecDeque<(usize, String, bool)>,
later_bytes: usize,
missed: bool,
context: usize,
tail: VecDeque<(usize, String)>,
tail_bytes: usize,
cutting: bool,
spill: Option<Spill>,
spilled: Option<(PathBuf, File)>,
live: Option<Live>,
in_key: bool,
}
impl Kept {
fn push(&mut self, piece: &OutputLine) {
self.written += piece.text().len() as u64;
let held = match piece {
OutputLine::Out(_) => 0,
OutputLine::Err(_) => 1,
};
let Some(line) = self.held.get_mut(held) else {
return;
};
line.push_str(piece.text());
if line.ends_with('\n') || line.len() >= MAX_HELD {
let line = std::mem::take(line);
self.line(line);
}
}
fn flush(&mut self) {
for held in 0..self.held.len() {
if let Some(line) = self.held.get_mut(held).map(std::mem::take)
&& !line.is_empty()
{
self.line(line);
}
}
}
fn line(&mut self, line: String) {
let line = plain(&line);
let mut line = self.redact(&line).into_owned();
if line.is_empty() {
return;
}
if !line.ends_with('\n') {
line.push('\n');
}
if let Some(live) = &self.live {
live.line(&line);
}
if let Some((_, file)) = &mut self.spilled
&& file.write_all(line.as_bytes()).is_err()
{
self.spilled = None;
}
self.tail_bytes += line.len();
self.tail.push_back((self.lines, line));
self.lines += 1;
let Some(limit) = self.cut else {
return;
};
if !self.cutting {
if self.tail.len() <= limit.lines && self.tail_bytes <= limit.bytes {
return;
}
self.cutting = true;
self.spill();
while self.head.len() < limit.lines / 2
&& let Some((at, line)) = self.tail.pop_front()
{
if self.head_bytes + line.len() > limit.bytes / 2 {
self.tail.push_front((at, line));
break;
}
self.head_bytes += line.len();
self.tail_bytes -= line.len();
self.head.push((at, line));
}
}
while self.tail.len()
> limit
.lines
.saturating_sub(self.head.len() + self.picked.len() + self.later.len())
|| self.tail_bytes
> limit
.bytes
.saturating_sub(self.head_bytes + self.picked_bytes + self.later_bytes)
{
let Some((at, line)) = self.tail.pop_front() else {
break;
};
self.tail_bytes -= line.len();
let starts = matters(&line);
if starts {
self.context = CONTEXT + 1;
}
if self.context > 0 {
self.context -= 1;
let mut line = clip(line.trim_end_matches('\n'), MAX_PICKED).to_owned();
line.push('\n');
self.pick(limit, at, line, starts);
}
}
}
fn pick(&mut self, limit: OutputLimit, at: usize, line: String, starts: bool) {
let (lines, bytes) = (limit.lines / 8, limit.bytes / 8);
if self.later.is_empty()
&& self.picked.len() < lines
&& self.picked_bytes + line.len() <= bytes
{
self.picked_bytes += line.len();
self.picked.push((at, line));
return;
}
self.later_bytes += line.len();
self.later.push_back((at, line, starts));
while self.later.len() > lines || self.later_bytes > bytes {
self.missed = true;
while let Some((_, line, _)) = self.later.pop_front() {
self.later_bytes -= line.len();
if self.later.front().is_none_or(|(_, _, starts)| *starts) {
break;
}
}
}
}
fn redact<'a>(&mut self, piece: &'a str) -> Cow<'a, str> {
if self.in_key {
self.in_key = !piece.contains("-----END ");
return Cow::Borrowed("");
}
let mut redacted = redact::text(piece);
if let Cow::Owned(line) = &mut redacted {
self.in_key = redact::find(piece).any(|secret| {
secret.kind == redact::Kind::PrivateKey
&& piece
.get(secret.range)
.is_some_and(|key| !key.contains("-----END "))
});
if piece.ends_with('\n') && !line.ends_with('\n') {
line.push('\n');
}
}
redacted
}
fn spill(&mut self) {
let Some(spill) = self.spill.take() else {
return;
};
self.spilled = spill.create("shell").ok().filter(|(_, file)| {
let mut file = file;
self.tail
.iter()
.try_for_each(|(_, line)| file.write_all(line.as_bytes()))
.is_ok()
});
}
fn text(mut self) -> (String, Option<OutputCut>) {
self.flush();
let kept_lines = self.head.len() + self.picked.len() + self.later.len() + self.tail.len();
let errors = match self.missed {
true => "the first and last lines with errors or failures",
false => "lines with errors or failures",
};
let cut = self.lines - kept_lines;
let mut text = String::new();
let mut fact = None;
if cut > 0 {
match &self.spilled {
Some((path, _)) => {
let first = self.head.len() + 1;
let next = self
.picked
.first()
.map(|(at, _)| *at)
.or(self.later.front().map(|(at, _, _)| *at))
.or(self.tail.front().map(|(at, _)| *at));
let gap = next.map_or(cut, |at| at + 1 - first);
text.push_str(&format!(
"[{cut} of {} lines cut, but for {errors}; all of the output is in {}: \
`read` it there with offset {first} and limit {} for the first lines \
cut, rather than running the command again]\n",
self.lines,
path.display(),
gap.min(crate::MAX_LINES),
));
}
None => text.push_str(&format!(
"[{cut} of {} lines cut, but for {errors}]\n",
self.lines
)),
}
fact = Some(OutputCut {
kept: KeptEnd::Ends,
kept_lines,
total_lines: self.lines,
spill: self
.spilled
.as_ref()
.map(|(path, _)| path.to_string_lossy().into_owned()),
});
}
let mut next = 0;
let kept = (self.head.iter().chain(&self.picked))
.map(|(at, line)| (at, line))
.chain(self.later.iter().map(|(at, line, _)| (at, line)))
.chain(self.tail.iter().map(|(at, line)| (at, line)));
for (at, line) in kept {
if *at > next {
text.push_str(&format!(
"[… {} lines cut (lines {}-{at}) …]\n",
at - next,
next + 1
));
}
text.push_str(line);
next = at + 1;
}
(text, fact)
}
}
#[cfg(test)]
mod tests;