yapi_core/
bash_executor.rs

1//! User `!` commands: run in the shell, stream sanitized output, keep the tail.
2//!
3//! Port of `packages/coding-agent/src/core/bash-executor.ts` in pi `v1.0.0`.
4
5use std::io::Write;
6use std::path::{Path, PathBuf};
7use std::sync::LazyLock;
8
9use regex_lite::Regex;
10use tokio_util::sync::CancellationToken;
11
12use crate::tools::bash::{Event, create_log, exit_code_of, run_child, shell_command, shell_with};
13use crate::tools::truncate::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncate_tail};
14
15pub use yapi_types::rpc::BashResult;
16
17static ANSI: LazyLock<Option<Regex>> = LazyLock::new(|| {
18    // OSC up to the first string terminator (BEL, ESC \ or 0x9c), or CSI.
19    Regex::new(concat!(
20        r"(?:\x1b\][\s\S]*?(?:\x07|\x1b\\|\x{9c}))",
21        r"|[\x1b\x{9b}][\[\]()#;?]*(?:\d{1,4}(?:[;:]\d{0,4})*)?[\dA-PR-TZcf-nq-uy=><~]",
22    ))
23    .ok()
24});
25
26/// pi's `stripAnsi`: removes OSC (operating system command) and CSI (control
27/// sequence introducer) sequences.
28pub fn strip_ansi(text: &str) -> String {
29    if !text.contains(['\x1b', '\u{9b}']) {
30        return text.to_owned();
31    }
32    match ANSI.as_ref() {
33        Some(regex) => regex.replace_all(text, "").into_owned(),
34        None => text.to_owned(),
35    }
36}
37
38/// pi's `sanitizeBinaryOutput`: drops control characters other than tab,
39/// newline and carriage return, and Unicode interlinear annotations.
40pub fn sanitize_binary(text: &str) -> String {
41    text.chars()
42        .filter(|c| {
43            !matches!(c, '\x00'..='\x08' | '\x0b' | '\x0c' | '\x0e'..='\x1f' | '\u{fff9}'..='\u{fffb}')
44        })
45        .collect()
46}
47
48#[derive(Default)]
49struct Output {
50    chunks: Vec<String>,
51    kept: usize,
52    total_bytes: usize,
53    pending: Vec<u8>,
54    file: Option<(PathBuf, std::fs::File)>,
55}
56
57impl Output {
58    fn ensure_file(&mut self) {
59        if self.file.is_some() {
60            return;
61        }
62        if let Some((path, mut file)) = create_log() {
63            for chunk in &self.chunks {
64                let _ = file.write_all(chunk.as_bytes());
65            }
66            self.file = Some((path, file));
67        }
68    }
69
70    fn push(&mut self, data: &[u8], on_chunk: &mut impl FnMut(&str)) {
71        self.total_bytes += data.len();
72        let decoded = yapi_types::js::decode_utf8_stream(&mut self.pending, data);
73        self.text(&decoded, on_chunk);
74    }
75
76    fn text(&mut self, decoded: &str, on_chunk: &mut impl FnMut(&str)) {
77        let text = sanitize_binary(&strip_ansi(decoded)).replace('\r', "");
78        if self.total_bytes > DEFAULT_MAX_BYTES {
79            self.ensure_file();
80        }
81        if let Some((_, file)) = &mut self.file {
82            let _ = file.write_all(text.as_bytes());
83        }
84        self.kept += text.encode_utf16().count();
85        self.chunks.push(text.clone());
86        while self.kept > DEFAULT_MAX_BYTES * 2 && self.chunks.len() > 1 {
87            let removed = self.chunks.remove(0);
88            self.kept -= removed.encode_utf16().count();
89        }
90        on_chunk(&text);
91    }
92
93    fn finish(mut self, exit_code: Option<i32>, cancelled: bool) -> BashResult {
94        let full = self.chunks.concat();
95        let truncation = truncate_tail(&full, DEFAULT_MAX_LINES, DEFAULT_MAX_BYTES);
96        if truncation.truncated {
97            self.ensure_file();
98        }
99        BashResult {
100            output: if truncation.truncated {
101                truncation.content
102            } else {
103                full
104            },
105            exit_code: if cancelled { None } else { exit_code },
106            cancelled,
107            truncated: truncation.truncated,
108            full_output_path: self.file.map(|(path, _)| path.display().to_string()),
109        }
110    }
111}
112
113/// Runs `command` in `cwd` with the configured shell, calling `on_chunk` with
114/// each sanitized piece of output, until it exits or `cancel` fires.
115pub async fn execute(
116    command: &str,
117    cwd: &Path,
118    shell_path: Option<&str>,
119    bin_dir: &Path,
120    cancel: CancellationToken,
121    mut on_chunk: impl FnMut(&str),
122) -> Result<BashResult, String> {
123    if cancel.is_cancelled() {
124        return Err("aborted".into());
125    }
126    let shell = shell_with(shell_path)?;
127    if !cwd.exists() {
128        return Err(format!(
129            "Working directory does not exist: {}\nCannot execute bash commands.",
130            cwd.display()
131        ));
132    }
133    let child = shell_command(shell, command, cwd, bin_dir)
134        .spawn()
135        .map_err(|err| err.to_string())?;
136    let mut output = Output::default();
137    let (status, stop) = run_child(child, &cancel, None, |event| {
138        if let Event::Output(bytes) = event {
139            output.push(bytes, &mut on_chunk);
140        }
141        None
142    })
143    .await;
144    let exit_code = status.map(|status| exit_code_of(status).unwrap_or(1));
145    // pi never flushes its decoder, so an incomplete UTF-8 sequence at the
146    // end of the output is dropped.
147    Ok(output.finish(exit_code, stop.is_some()))
148}
149
150/// pi's `executeBashWithOperations`: runs a command through `exec`, which
151/// sends its output to the sender it gets and returns the exit code, and
152/// collects the output as [`execute`] does. A run that fails or ends after
153/// `cancel` fired is cancelled.
154pub async fn execute_with<F>(
155    exec: impl FnOnce(tokio::sync::mpsc::UnboundedSender<Vec<u8>>) -> F,
156    cancel: &CancellationToken,
157    mut on_chunk: impl FnMut(&str),
158) -> Result<BashResult, String>
159where
160    F: std::future::Future<Output = Result<Option<i32>, String>>,
161{
162    let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
163    let mut output = Output::default();
164    let run = exec(sender);
165    tokio::pin!(run);
166    let outcome = loop {
167        tokio::select! {
168            biased;
169            Some(bytes) = receiver.recv() => output.push(&bytes, &mut on_chunk),
170            outcome = &mut run => break outcome,
171        }
172    };
173    while let Ok(bytes) = receiver.try_recv() {
174        output.push(&bytes, &mut on_chunk);
175    }
176    let cancelled = cancel.is_cancelled();
177    match outcome {
178        Ok(exit_code) => Ok(output.finish(exit_code, cancelled)),
179        Err(_) if cancelled => Ok(output.finish(None, true)),
180        Err(error) => Err(error),
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn strips_escapes_and_controls() {
190        assert_eq!(strip_ansi("\x1b[31mred\x1b[0m"), "red");
191        assert_eq!(strip_ansi("\x1b]8;;http://x\x07link\x1b]8;;\x07"), "link");
192        assert_eq!(sanitize_binary("a\x00b\tc\x07"), "ab\tc");
193    }
194
195    #[cfg(unix)]
196    #[tokio::test]
197    async fn runs_and_cancels() {
198        let dir = std::env::temp_dir();
199        let mut seen = String::new();
200        // An empty `shellPath` means the default shell.
201        let result = execute(
202            "printf 'one\\r\\ntwo\\n'; exit 3",
203            &dir,
204            Some(""),
205            &dir,
206            CancellationToken::new(),
207            |chunk| seen.push_str(chunk),
208        )
209        .await
210        .unwrap_or_default();
211        assert_eq!(result.output, "one\ntwo\n");
212        assert_eq!(result.exit_code, Some(3));
213        assert_eq!(seen, "one\ntwo\n");
214
215        let cancel = CancellationToken::new();
216        let trigger = cancel.clone();
217        tokio::spawn(async move {
218            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
219            trigger.cancel();
220        });
221        let result = execute("sleep 5", &dir, None, &dir, cancel, |_| {})
222            .await
223            .unwrap_or_default();
224        assert!(result.cancelled);
225        assert_eq!(result.exit_code, None);
226    }
227
228    #[cfg(unix)]
229    #[tokio::test]
230    async fn drops_trailing_partial_char() {
231        let dir = std::env::temp_dir();
232        let mut seen = String::new();
233        // The first two bytes of "€".
234        let result = execute(
235            "printf 'ok \\342\\202'",
236            &dir,
237            None,
238            &dir,
239            CancellationToken::new(),
240            |chunk| seen.push_str(chunk),
241        )
242        .await
243        .unwrap_or_default();
244        assert_eq!(result.output, "ok ");
245        assert_eq!(seen, "ok ");
246    }
247}
yapi_core/bash_executor.rs

yapi_core/
bash_executor.rs
