1use std::borrow::Cow;
4use std::collections::VecDeque;
5use std::ffi::{OsStr, OsString};
6use std::fs::File;
7use std::io::Write as _;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11
12use rig_core::completion::ToolAnnotations;
13use rig_core::redact;
14use rig_core::tool::{
15 ExitStatus, KeptEnd, OutputCut, Tool, ToolContext, ToolExecutionError, args_schema,
16};
17use schemars::JsonSchema;
18use serde::Deserialize;
19
20use crate::process::{Ended, KillOnDrop, Launcher, OutputLine, Process, Spawn};
21use crate::workspace::{Local, Workspace};
22use crate::{Live, OutputLimit, Spill, blocking, clip};
23
24const DEFAULT_TIMEOUT: u64 = 120;
25const OUTPUT_LIMIT: u64 = 10 * 1024 * 1024;
29
30#[derive(Clone, Debug)]
40pub struct Shell {
41 pub workspace: Arc<dyn Workspace>,
43 pub launcher: Launcher,
45 pub live: Option<Live>,
48 pub spill: Option<Spill>,
50 pub cut: Option<OutputLimit>,
54 pub confine: bool,
58 pub max_timeout: u64,
61 pub quiet: Duration,
65}
66
67impl Default for Shell {
68 fn default() -> Self {
73 Self {
74 workspace: Arc::new(Local),
75 launcher: Launcher::default(),
76 live: None,
77 spill: None,
78 cut: Some(OutputLimit::default()),
79 confine: false,
80 max_timeout: 600,
81 quiet: Duration::from_secs(60),
82 }
83 }
84}
85
86#[derive(Clone, Debug, PartialEq, Eq, Deserialize, JsonSchema)]
88#[serde(deny_unknown_fields)]
89pub struct ShellArgs {
90 pub command: String,
92 pub timeout_secs: Option<u64>,
94 pub cwd: Option<String>,
97}
98
99impl ShellArgs {
100 pub fn new(command: impl Into<String>) -> Self {
102 Self {
103 command: command.into(),
104 timeout_secs: None,
105 cwd: None,
106 }
107 }
108
109 pub fn timeout_secs(mut self, secs: u64) -> Self {
111 self.timeout_secs = Some(secs);
112 self
113 }
114
115 pub fn in_dir(mut self, dir: impl Into<String>) -> Self {
117 self.cwd = Some(dir.into());
118 self
119 }
120}
121
122#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct ShellRun {
125 pub output: String,
129 pub status: ExitStatus,
131}
132
133impl Shell {
134 pub async fn run(&self, args: ShellArgs) -> Result<ShellRun, ToolExecutionError> {
139 self.execute(args).await.map(|(ran, _)| ran)
140 }
141
142 async fn execute(
144 &self,
145 args: ShellArgs,
146 ) -> Result<(ShellRun, Option<OutputCut>), ToolExecutionError> {
147 let timeout = Duration::from_secs(
148 args.timeout_secs
149 .unwrap_or(DEFAULT_TIMEOUT.min(self.max_timeout))
150 .clamp(1, self.max_timeout.max(1)),
151 );
152 let quiet = match args.timeout_secs {
153 None => Some(self.quiet).filter(|quiet| *quiet < timeout),
154 Some(_) => None,
155 };
156 let process = self.start(&args)?;
157 let _stop = KillOnDrop(process.group());
158 let shell = self.clone();
159 blocking(move || shell.finish(process, timeout, quiet)).await
160 }
161}
162
163impl Tool for Shell {
164 const NAME: &'static str = "shell";
165 type Args = ShellArgs;
166 type Output = String;
167 type Error = ToolExecutionError;
168
169 fn description(&self) -> String {
170 let kept = match self.cut {
171 Some(cut) => format!(
172 "; past {} lines or {} KB, its first and last lines and the lines with errors \
173 or failures from between",
174 cut.lines,
175 cut.bytes / 1024
176 ),
177 None => String::new(),
178 };
179 format!(
180 "Run a shell command with `sh -c` in the current directory, or in cwd. stdin is \
181 empty. Returns the exit code and stdout and stderr interleaved line by line{kept}. \
182 The command is killed after timeout_secs (default {}, at most {}), or once it \
183 printed more than {} MB; with the default, also once it printed nothing for {}s, \
184 so name timeout_secs for a build or wait that is quiet for long; a quick command \
185 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 \
186 carry over to the next call. Do not start programs that wait for input or never end, such as editors, \
187 pagers or servers in the foreground. Do not assume a git repository: a failed \
188 `git` command ends a `&&` chain. Check formatting only of files you changed.",
189 DEFAULT_TIMEOUT.min(self.max_timeout).max(1),
190 self.max_timeout.max(1),
191 OUTPUT_LIMIT / 1024 / 1024,
192 self.quiet.as_secs()
193 )
194 }
195
196 fn parameters(&self) -> serde_json::Value {
197 args_schema::<ShellArgs>()
198 }
199
200 fn annotations(&self) -> ToolAnnotations {
201 ToolAnnotations::default().with_title("Shell")
202 }
203
204 async fn call(
205 &self,
206 context: &mut ToolContext,
207 args: ShellArgs,
208 ) -> Result<String, ToolExecutionError> {
209 let (ran, cut) = self.execute(args).await?;
210 context.insert_result(ran.status)?;
211 if let Some(cut) = cut {
212 context.insert_result(cut)?;
213 }
214 Ok(ran.output)
215 }
216}
217
218impl Shell {
219 fn start(&self, args: &ShellArgs) -> Result<Process, ToolExecutionError> {
223 let base = self.launcher.cwd.as_deref();
224 let dir = match &args.cwd {
225 Some(dir) => {
226 if self.confine {
227 self.workspace.place(base, true, dir)?;
228 }
229 Some(
230 self.workspace
231 .join(base.unwrap_or(Path::new("")), Path::new(dir)),
232 )
233 }
234 None => base.map(Path::to_path_buf),
235 };
236 let how = Spawn {
237 stdin: None,
238 merge_stderr: false,
239 };
240 let line = [OsString::from("-c"), OsString::from(&args.command)];
241 self.workspace
242 .spawn(&self.launcher, dir.as_deref(), OsStr::new("sh"), &line, how)
243 .map_err(|error| ToolExecutionError::other(format!("could not run sh: {error}")))
244 }
245
246 fn finish(
249 &self,
250 mut process: Process,
251 timeout: Duration,
252 quiet: Option<Duration>,
253 ) -> Result<(ShellRun, Option<OutputCut>), ToolExecutionError> {
254 let mut kept = Kept {
255 spill: self.spill.clone(),
256 cut: self.cut,
257 live: self.live.clone(),
258 ..Kept::default()
259 };
260 let deadline = Instant::now().checked_add(timeout);
261 let end = process
262 .run(deadline, quiet, |line| {
263 kept.push(&line);
264 kept.written <= OUTPUT_LIMIT
265 })
266 .map_err(|error| {
267 ToolExecutionError::other(format!("waiting for the command failed: {error}"))
268 })?;
269 let silent = kept.written == 0;
270 let (mut output, cut) = kept.text();
271 let status = match end {
272 Ended::Exited(status) => match status.code() {
273 Some(code) => {
274 output.push_str(&format!("[exit code {code}]"));
275 ExitStatus::Exited(code)
276 }
277 None => {
278 output.push_str(&killed(status));
279 ExitStatus::Signaled
280 }
281 },
282 Ended::TimedOut => {
283 let secs = timeout.as_secs();
284 output.push_str(&format!(
285 "[timed out after {secs}s and stopped: {}",
286 stopped(&process)
287 ));
288 if silent {
289 output.push_str(SILENT);
290 }
291 if secs < self.max_timeout {
292 output.push_str(&format!(
293 ". Pass timeout_secs up to {} to wait longer",
294 self.max_timeout
295 ));
296 }
297 output.push(']');
298 ExitStatus::TimedOut
299 }
300 Ended::Quiet => {
301 output.push_str(&format!(
302 "[printed nothing for {}s, so it was stopped: {}",
303 quiet.unwrap_or_default().as_secs(),
304 stopped(&process)
305 ));
306 if silent {
307 output.push_str(SILENT);
308 }
309 output.push_str(
310 ". A command that is stuck, such as one waiting for input or looping \
311 forever, needs fixing, not a longer wait; only one known to be quiet for \
312 long, such as a build, needs timeout_secs]",
313 );
314 ExitStatus::TimedOut
315 }
316 Ended::Stopped => {
317 output.push_str(&format!(
318 "[killed after printing more than {} MB: {}]",
319 OUTPUT_LIMIT / 1024 / 1024,
320 stopped(&process)
321 ));
322 ExitStatus::Killed
323 }
324 };
325 Ok((ShellRun { output, status }, cut))
326 }
327}
328
329const SILENT: &str = ". It printed nothing: a program whose output is not a terminal may hold \
331 it back until it ends, so make it print as it goes (such as \
332 `python3 -u` or `stdbuf -oL`) to see how far it got";
333
334fn stopped(process: &Process) -> String {
337 let grace = crate::process::STOP_GRACE.as_secs();
338 let how = if !cfg!(unix) {
339 "it was killed".to_owned()
340 } else if process.forced() {
341 format!("it got SIGTERM, then SIGKILL {grace}s later as it still ran")
342 } else {
343 "it got SIGTERM and ended".to_owned()
344 };
345 match process.remote() {
346 Some(host) => format!(
347 "{how}. That was the transport on this machine; whether the command on {host} \
348 stopped too is the transport's doing, so check there before running it again"
349 ),
350 None => how,
351 }
352}
353
354const STATUS_STARTS: [&str; 5] = [
357 "[exit code ",
358 "[killed by ",
359 "[timed out after ",
360 "[printed nothing for ",
361 "[killed after printing more than ",
362];
363
364fn killed(status: std::process::ExitStatus) -> String {
368 #[cfg(unix)]
369 if let Some(signal) = std::os::unix::process::ExitStatusExt::signal(&status) {
370 let (name, why) = sent(signal);
371 return format!(
372 "[killed by {name} (signal {signal}), which this tool did not send: {why}. It was \
373 the process this call started here, `sh` or the last program it ran; when that \
374 program was a client of another machine, such as a remote shell, the command \
375 there may still be running or may have finished, so check there before running \
376 it again]"
377 );
378 }
379 #[cfg(not(unix))]
380 let _ = status;
381 "[killed by a signal]".to_owned()
382}
383
384#[cfg(unix)]
386fn sent(signal: i32) -> (&'static str, &'static str) {
387 match signal {
388 libc::SIGHUP => ("SIGHUP", "its terminal or session went away"),
389 libc::SIGINT => ("SIGINT", "something on this machine interrupted it"),
390 libc::SIGQUIT => ("SIGQUIT", "something on this machine asked it to quit"),
391 libc::SIGILL => ("SIGILL", "it crashed"),
392 libc::SIGTRAP => ("SIGTRAP", "it crashed, or a debugger stopped it"),
393 libc::SIGABRT => ("SIGABRT", "it aborted itself, as on a failed assertion"),
394 libc::SIGBUS => ("SIGBUS", "it crashed"),
395 libc::SIGFPE => ("SIGFPE", "it crashed on an arithmetic error"),
396 libc::SIGKILL => (
397 "SIGKILL",
398 "something on this machine killed it, such as the out-of-memory killer or a `kill`",
399 ),
400 libc::SIGSEGV => ("SIGSEGV", "it crashed"),
401 libc::SIGPIPE => ("SIGPIPE", "it wrote to a pipe that nothing read any more"),
402 libc::SIGTERM => ("SIGTERM", "something on this machine asked it to stop"),
403 _ => ("a signal", "something on this machine sent it"),
404 }
405}
406
407pub fn split_status(text: &str) -> (&str, Option<&str>) {
420 let start = text.rfind('\n').map_or(0, |newline| newline + 1);
421 match text.split_at_checked(start) {
422 Some((printed, last))
423 if last.ends_with(']') && STATUS_STARTS.iter().any(|at| last.starts_with(at)) =>
424 {
425 (printed, Some(last))
426 }
427 _ => (text, None),
428 }
429}
430
431const MAX_HELD: usize = 1024 * 1024;
434const CONTEXT: usize = 3;
437const MAX_PICKED: usize = 300;
439
440fn matters(line: &str) -> bool {
444 let start = line.trim_start();
445 start.starts_with("error")
446 || start.starts_with("Error")
447 || [
448 "FAILED",
449 "FAIL:",
450 "failed",
451 "panicked",
452 "Traceback",
453 "Error:",
454 ]
455 .iter()
456 .any(|word| line.contains(word))
457}
458
459fn plain(line: &str) -> Cow<'_, str> {
463 if !line.contains('\x1b') {
464 return Cow::Borrowed(line);
465 }
466 let mut text = String::with_capacity(line.len());
467 let mut chars = line.chars().peekable();
468 while let Some(c) = chars.next() {
469 if c != '\x1b' {
470 text.push(c);
471 continue;
472 }
473 match chars.next() {
474 Some('[') => {
476 let _ = chars.by_ref().find(|c| ('\x40'..='\x7e').contains(c));
477 }
478 Some(']' | 'P' | 'X' | '^' | '_') => {
480 while let Some(c) = chars.next() {
481 if c == '\x07' || (c == '\x1b' && chars.next_if_eq(&'\\').is_some()) {
482 break;
483 }
484 }
485 }
486 Some(' '..='/') => {
488 let _ = chars.by_ref().find(|c| !(' '..='/').contains(c));
489 }
490 _ => {}
491 }
492 }
493 Cow::Owned(text)
494}
495
496#[derive(Default)]
503struct Kept {
504 written: u64,
506 cut: Option<OutputLimit>,
508 held: [String; 2],
511 lines: usize,
513 head: Vec<(usize, String)>,
515 head_bytes: usize,
517 picked: Vec<(usize, String)>,
519 picked_bytes: usize,
521 later: VecDeque<(usize, String, bool)>,
525 later_bytes: usize,
527 missed: bool,
529 context: usize,
531 tail: VecDeque<(usize, String)>,
533 tail_bytes: usize,
535 cutting: bool,
537 spill: Option<Spill>,
539 spilled: Option<(PathBuf, File)>,
541 live: Option<Live>,
543 in_key: bool,
546}
547
548impl Kept {
549 fn push(&mut self, piece: &OutputLine) {
553 self.written += piece.text().len() as u64;
554 let held = match piece {
555 OutputLine::Out(_) => 0,
556 OutputLine::Err(_) => 1,
557 };
558 let Some(line) = self.held.get_mut(held) else {
559 return;
560 };
561 line.push_str(piece.text());
562 if line.ends_with('\n') || line.len() >= MAX_HELD {
563 let line = std::mem::take(line);
564 self.line(line);
565 }
566 }
567
568 fn flush(&mut self) {
571 for held in 0..self.held.len() {
572 if let Some(line) = self.held.get_mut(held).map(std::mem::take)
573 && !line.is_empty()
574 {
575 self.line(line);
576 }
577 }
578 }
579
580 fn line(&mut self, line: String) {
585 let line = plain(&line);
586 let mut line = self.redact(&line).into_owned();
587 if line.is_empty() {
588 return;
589 }
590 if !line.ends_with('\n') {
591 line.push('\n');
592 }
593 if let Some(live) = &self.live {
594 live.line(&line);
595 }
596 if let Some((_, file)) = &mut self.spilled
597 && file.write_all(line.as_bytes()).is_err()
598 {
599 self.spilled = None;
600 }
601 self.tail_bytes += line.len();
602 self.tail.push_back((self.lines, line));
603 self.lines += 1;
604 let Some(limit) = self.cut else {
605 return;
606 };
607 if !self.cutting {
608 if self.tail.len() <= limit.lines && self.tail_bytes <= limit.bytes {
609 return;
610 }
611 self.cutting = true;
612 self.spill();
613 while self.head.len() < limit.lines / 2
614 && let Some((at, line)) = self.tail.pop_front()
615 {
616 if self.head_bytes + line.len() > limit.bytes / 2 {
617 self.tail.push_front((at, line));
618 break;
619 }
620 self.head_bytes += line.len();
621 self.tail_bytes -= line.len();
622 self.head.push((at, line));
623 }
624 }
625 while self.tail.len()
627 > limit
628 .lines
629 .saturating_sub(self.head.len() + self.picked.len() + self.later.len())
630 || self.tail_bytes
631 > limit
632 .bytes
633 .saturating_sub(self.head_bytes + self.picked_bytes + self.later_bytes)
634 {
635 let Some((at, line)) = self.tail.pop_front() else {
636 break;
637 };
638 self.tail_bytes -= line.len();
639 let starts = matters(&line);
640 if starts {
641 self.context = CONTEXT + 1;
642 }
643 if self.context > 0 {
644 self.context -= 1;
645 let mut line = clip(line.trim_end_matches('\n'), MAX_PICKED).to_owned();
646 line.push('\n');
647 self.pick(limit, at, line, starts);
648 }
649 }
650 }
651
652 fn pick(&mut self, limit: OutputLimit, at: usize, line: String, starts: bool) {
656 let (lines, bytes) = (limit.lines / 8, limit.bytes / 8);
657 if self.later.is_empty()
658 && self.picked.len() < lines
659 && self.picked_bytes + line.len() <= bytes
660 {
661 self.picked_bytes += line.len();
662 self.picked.push((at, line));
663 return;
664 }
665 self.later_bytes += line.len();
666 self.later.push_back((at, line, starts));
667 while self.later.len() > lines || self.later_bytes > bytes {
668 self.missed = true;
669 while let Some((_, line, _)) = self.later.pop_front() {
670 self.later_bytes -= line.len();
671 if self.later.front().is_none_or(|(_, _, starts)| *starts) {
672 break;
673 }
674 }
675 }
676 }
677
678 fn redact<'a>(&mut self, piece: &'a str) -> Cow<'a, str> {
682 if self.in_key {
683 self.in_key = !piece.contains("-----END ");
684 return Cow::Borrowed("");
685 }
686 let mut redacted = redact::text(piece);
687 if let Cow::Owned(line) = &mut redacted {
688 self.in_key = redact::find(piece).any(|secret| {
689 secret.kind == redact::Kind::PrivateKey
690 && piece
691 .get(secret.range)
692 .is_some_and(|key| !key.contains("-----END "))
693 });
694 if piece.ends_with('\n') && !line.ends_with('\n') {
696 line.push('\n');
697 }
698 }
699 redacted
700 }
701
702 fn spill(&mut self) {
705 let Some(spill) = self.spill.take() else {
706 return;
707 };
708 self.spilled = spill.create("shell").ok().filter(|(_, file)| {
709 let mut file = file;
710 self.tail
711 .iter()
712 .try_for_each(|(_, line)| file.write_all(line.as_bytes()))
713 .is_ok()
714 });
715 }
716
717 fn text(mut self) -> (String, Option<OutputCut>) {
722 self.flush();
723 let kept_lines = self.head.len() + self.picked.len() + self.later.len() + self.tail.len();
724 let errors = match self.missed {
725 true => "the first and last lines with errors or failures",
726 false => "lines with errors or failures",
727 };
728 let cut = self.lines - kept_lines;
729 let mut text = String::new();
730 let mut fact = None;
731 if cut > 0 {
732 match &self.spilled {
733 Some((path, _)) => {
734 let first = self.head.len() + 1;
736 let next = self
737 .picked
738 .first()
739 .map(|(at, _)| *at)
740 .or(self.later.front().map(|(at, _, _)| *at))
741 .or(self.tail.front().map(|(at, _)| *at));
742 let gap = next.map_or(cut, |at| at + 1 - first);
743 text.push_str(&format!(
744 "[{cut} of {} lines cut, but for {errors}; all of the output is in {}: \
745 `read` it there with offset {first} and limit {} for the first lines \
746 cut, rather than running the command again]\n",
747 self.lines,
748 path.display(),
749 gap.min(crate::MAX_LINES),
750 ));
751 }
752 None => text.push_str(&format!(
753 "[{cut} of {} lines cut, but for {errors}]\n",
754 self.lines
755 )),
756 }
757 fact = Some(OutputCut {
758 kept: KeptEnd::Ends,
759 kept_lines,
760 total_lines: self.lines,
761 spill: self
762 .spilled
763 .as_ref()
764 .map(|(path, _)| path.to_string_lossy().into_owned()),
765 });
766 }
767 let mut next = 0;
768 let kept = (self.head.iter().chain(&self.picked))
769 .map(|(at, line)| (at, line))
770 .chain(self.later.iter().map(|(at, line, _)| (at, line)))
771 .chain(self.tail.iter().map(|(at, line)| (at, line)));
772 for (at, line) in kept {
773 if *at > next {
774 text.push_str(&format!(
775 "[… {} lines cut (lines {}-{at}) …]\n",
776 at - next,
777 next + 1
778 ));
779 }
780 text.push_str(line);
781 next = at + 1;
782 }
783 (text, fact)
784 }
785}
786
787#[cfg(test)]
788mod tests;