agent_tools/base/
cmd.rs

1use crate::error::{Error, Result};
2use std::collections::HashMap;
3use std::fs::File;
4use std::io::{self, Read, Write};
5use std::path::PathBuf;
6use std::process::{Command, Stdio};
7use std::thread;
8use std::time::Duration;
9use wait_timeout::ChildExt;
10
11pub struct CmdTool;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum CmdStdin {
15    Text(String),
16    Bytes(Vec<u8>),
17    File(PathBuf),
18    Null,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct CmdRequest {
23    pub program: String,
24    pub args: Vec<String>,
25    pub cwd: Option<String>,
26    pub env: Option<HashMap<String, String>>,
27    pub timeout_ms: Option<u64>,
28    pub fail_on_non_zero: bool,
29    pub stdin: Option<CmdStdin>,
30    pub background: bool,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct ShellCmdRequest {
35    pub command: String,
36    pub cwd: Option<String>,
37    pub env: Option<HashMap<String, String>>,
38    pub timeout_ms: Option<u64>,
39    pub fail_on_non_zero: bool,
40    pub stdin: Option<CmdStdin>,
41    pub background: bool,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct CmdOutput {
46    pub stdout: String,
47    pub stderr: String,
48    pub exit_code: i32,
49    pub pid: Option<u32>,
50}
51
52impl CmdTool {
53    pub fn run(req: CmdRequest) -> Result<CmdOutput> {
54        let mut cmd = Command::new(&req.program);
55        cmd.args(&req.args);
56
57        run_inner(
58            &mut cmd,
59            req.cwd,
60            req.env,
61            req.timeout_ms,
62            req.fail_on_non_zero,
63            req.stdin,
64            req.background,
65        )
66    }
67
68    pub fn run_shell(req: ShellCmdRequest) -> Result<CmdOutput> {
69        let mut cmd = build_shell_command(&req.command);
70
71        run_inner(
72            &mut cmd,
73            req.cwd,
74            req.env,
75            req.timeout_ms,
76            req.fail_on_non_zero,
77            req.stdin,
78            req.background,
79        )
80    }
81}
82
83fn run_inner(
84    cmd: &mut Command,
85    cwd: Option<String>,
86    env: Option<HashMap<String, String>>,
87    timeout_ms: Option<u64>,
88    fail_on_non_zero: bool,
89    stdin: Option<CmdStdin>,
90    background: bool,
91) -> Result<CmdOutput> {
92    configure_command(cmd, cwd, env, stdin.as_ref(), background)?;
93    let mut child = spawn_child(cmd)?;
94    let stdout_handle = take_output_reader(&mut child.stdout);
95    let stderr_handle = take_output_reader(&mut child.stderr);
96
97    write_stdin(&mut child, stdin.as_ref())?;
98
99    if background {
100        return Ok(background_output(&child));
101    }
102
103    let status = match wait_for_child(&mut child, timeout_ms) {
104        Ok(status) => status,
105        Err(err) => {
106            let _ = collect_output(stdout_handle);
107            let _ = collect_output(stderr_handle);
108            return Err(err);
109        }
110    };
111
112    build_foreground_output(status, stdout_handle, stderr_handle, fail_on_non_zero)
113}
114
115fn configure_command(
116    cmd: &mut Command,
117    cwd: Option<String>,
118    env: Option<HashMap<String, String>>,
119    stdin: Option<&CmdStdin>,
120    background: bool,
121) -> Result<()> {
122    if let Some(cwd) = cwd {
123        cmd.current_dir(cwd);
124    }
125
126    if let Some(env) = env {
127        cmd.envs(env);
128    }
129
130    configure_stdin(cmd, stdin, background)?;
131    configure_output(cmd, background);
132    Ok(())
133}
134
135fn configure_stdin(cmd: &mut Command, stdin: Option<&CmdStdin>, background: bool) -> Result<()> {
136    match stdin {
137        Some(CmdStdin::File(path)) => {
138            let file = File::open(path).map_err(Error::tool_io)?;
139            cmd.stdin(file);
140        }
141        Some(CmdStdin::Null) => {
142            cmd.stdin(Stdio::null());
143        }
144        Some(_) => {
145            cmd.stdin(Stdio::piped());
146        }
147        None if background => {
148            cmd.stdin(Stdio::null());
149        }
150        None => {}
151    }
152
153    Ok(())
154}
155
156fn configure_output(cmd: &mut Command, background: bool) {
157    if background {
158        cmd.stdout(Stdio::null());
159        cmd.stderr(Stdio::null());
160        return;
161    }
162
163    cmd.stdout(Stdio::piped());
164    cmd.stderr(Stdio::piped());
165}
166
167fn spawn_child(cmd: &mut Command) -> Result<std::process::Child> {
168    cmd.spawn().map_err(Error::tool_io)
169}
170
171fn take_output_reader<R>(pipe: &mut Option<R>) -> Option<thread::JoinHandle<io::Result<Vec<u8>>>>
172where
173    R: Read + Send + 'static,
174{
175    pipe.take().map(spawn_reader)
176}
177
178fn background_output(child: &std::process::Child) -> CmdOutput {
179    CmdOutput {
180        stdout: String::new(),
181        stderr: String::new(),
182        exit_code: 0,
183        pid: Some(child.id()),
184    }
185}
186
187fn wait_for_child(
188    child: &mut std::process::Child,
189    timeout_ms: Option<u64>,
190) -> Result<std::process::ExitStatus> {
191    match timeout_ms {
192        Some(timeout_ms) => wait_with_timeout(child, timeout_ms),
193        None => child.wait().map_err(Error::tool_io),
194    }
195}
196
197fn wait_with_timeout(
198    child: &mut std::process::Child,
199    timeout_ms: u64,
200) -> Result<std::process::ExitStatus> {
201    let duration = Duration::from_millis(timeout_ms);
202    match child.wait_timeout(duration).map_err(Error::tool_io)? {
203        Some(status) => Ok(status),
204        None => {
205            kill_timed_out_child(child)?;
206            Err(Error::tool_timeout())
207        }
208    }
209}
210
211fn kill_timed_out_child(child: &mut std::process::Child) -> Result<()> {
212    child.kill().map_err(Error::tool_io)?;
213    child.wait().map_err(Error::tool_io)?;
214    Ok(())
215}
216
217fn write_stdin(child: &mut std::process::Child, stdin_content: Option<&CmdStdin>) -> Result<()> {
218    let Some(stdin_content) = stdin_content else {
219        return Ok(());
220    };
221
222    let Some(mut stdin) = child.stdin.take() else {
223        if matches!(stdin_content, CmdStdin::Text(_) | CmdStdin::Bytes(_)) {
224            return Err(Error::tool_io(io::Error::new(
225                io::ErrorKind::BrokenPipe,
226                "stdin pipe not available",
227            )));
228        }
229        return Ok(());
230    };
231
232    match stdin_content {
233        CmdStdin::Text(text) => stdin.write_all(text.as_bytes()).map_err(Error::tool_io)?,
234        CmdStdin::Bytes(bytes) => stdin.write_all(bytes).map_err(Error::tool_io)?,
235        CmdStdin::File(_) | CmdStdin::Null => {}
236    }
237
238    stdin.flush().map_err(Error::tool_io)?;
239    drop(stdin);
240    Ok(())
241}
242
243fn spawn_reader<R>(mut reader: R) -> thread::JoinHandle<io::Result<Vec<u8>>>
244where
245    R: Read + Send + 'static,
246{
247    thread::spawn(move || {
248        let mut buf = Vec::new();
249        reader.read_to_end(&mut buf)?;
250        Ok(buf)
251    })
252}
253
254fn collect_output(handle: Option<thread::JoinHandle<io::Result<Vec<u8>>>>) -> Result<String> {
255    let Some(handle) = handle else {
256        return Ok(String::new());
257    };
258
259    let bytes = handle
260        .join()
261        .map_err(|_| Error::tool_io(io::Error::other("output reader thread panicked")))?
262        .map_err(Error::tool_io)?;
263
264    Ok(String::from_utf8_lossy(&bytes).into_owned())
265}
266
267fn build_foreground_output(
268    status: std::process::ExitStatus,
269    stdout_handle: Option<thread::JoinHandle<io::Result<Vec<u8>>>>,
270    stderr_handle: Option<thread::JoinHandle<io::Result<Vec<u8>>>>,
271    fail_on_non_zero: bool,
272) -> Result<CmdOutput> {
273    let stdout = collect_output(stdout_handle)?;
274    let stderr = collect_output(stderr_handle)?;
275    let exit_code = status.code().unwrap_or(-1);
276
277    if fail_on_non_zero && exit_code != 0 {
278        return Err(Error::tool_cmd_failed(exit_code));
279    }
280
281    Ok(CmdOutput {
282        stdout,
283        stderr,
284        exit_code,
285        pid: None,
286    })
287}
288
289fn build_shell_command(command: &str) -> Command {
290    if cfg!(target_os = "windows") {
291        let mut cmd = Command::new("cmd.exe");
292        cmd.arg("/c").arg(command);
293        cmd
294    } else {
295        let mut cmd = Command::new("sh");
296        cmd.arg("-c").arg(command);
297        cmd
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use std::io::Write;
305    use tempfile::NamedTempFile;
306
307    #[test]
308    fn test_successful_command() {
309        let req = CmdRequest {
310            program: "echo".to_string(),
311            args: vec!["hello".to_string()],
312            cwd: None,
313            env: None,
314            timeout_ms: None,
315            fail_on_non_zero: false,
316            stdin: None,
317            background: false,
318        };
319        let result = CmdTool::run(req);
320        assert!(result.is_ok());
321        let output = result.unwrap();
322        assert_eq!(output.exit_code, 0);
323        assert_eq!(output.stdout.trim(), "hello");
324        assert!(output.pid.is_none());
325    }
326
327    #[test]
328    fn test_shell_command() {
329        let req = ShellCmdRequest {
330            command: "echo 'hello from shell'".to_string(),
331            cwd: None,
332            env: None,
333            timeout_ms: None,
334            fail_on_non_zero: false,
335            stdin: None,
336            background: false,
337        };
338        let result = CmdTool::run_shell(req);
339        assert!(result.is_ok());
340        let output = result.unwrap();
341        assert_eq!(output.exit_code, 0);
342        assert_eq!(output.stdout.trim(), "hello from shell");
343        assert!(output.pid.is_none());
344    }
345
346    #[test]
347    fn test_timeout_command() {
348        let req = CmdRequest {
349            program: "sleep".to_string(),
350            args: vec!["2".to_string()],
351            cwd: None,
352            env: None,
353            timeout_ms: Some(100),
354            fail_on_non_zero: false,
355            stdin: None,
356            background: false,
357        };
358        let result = CmdTool::run(req);
359        assert!(result.is_err());
360        let err_msg = result.unwrap_err().to_string().to_lowercase();
361        assert!(err_msg.contains("timed out"));
362    }
363
364    #[test]
365    fn test_non_existent_command() {
366        let req = CmdRequest {
367            program: "this_command_does_not_exist_12345".to_string(),
368            args: vec![],
369            cwd: None,
370            env: None,
371            timeout_ms: None,
372            fail_on_non_zero: false,
373            stdin: None,
374            background: false,
375        };
376        let result = CmdTool::run(req);
377        assert!(result.is_err());
378        let err_msg = result.unwrap_err().to_string().to_lowercase();
379        assert!(err_msg.contains("no such file") || err_msg.contains("not found"));
380    }
381
382    #[test]
383    fn test_stdin_text() {
384        let req = CmdRequest {
385            program: "cat".to_string(),
386            args: vec![],
387            cwd: None,
388            env: None,
389            timeout_ms: None,
390            fail_on_non_zero: false,
391            stdin: Some(CmdStdin::Text("hello stdin text".to_string())),
392            background: false,
393        };
394        let result = CmdTool::run(req);
395        assert!(result.is_ok());
396        let output = result.unwrap();
397        assert_eq!(output.exit_code, 0);
398        assert_eq!(output.stdout, "hello stdin text");
399        assert!(output.pid.is_none());
400    }
401
402    #[test]
403    fn test_stdin_bytes() {
404        let req = CmdRequest {
405            program: "cat".to_string(),
406            args: vec![],
407            cwd: None,
408            env: None,
409            timeout_ms: None,
410            fail_on_non_zero: false,
411            stdin: Some(CmdStdin::Bytes(b"hello stdin bytes".to_vec())),
412            background: false,
413        };
414        let result = CmdTool::run(req);
415        assert!(result.is_ok());
416        let output = result.unwrap();
417        assert_eq!(output.exit_code, 0);
418        assert_eq!(output.stdout, "hello stdin bytes");
419        assert!(output.pid.is_none());
420    }
421
422    #[test]
423    fn test_stdin_file() {
424        let mut temp_file = NamedTempFile::new().unwrap();
425        write!(temp_file, "hello stdin file").unwrap();
426
427        let req = CmdRequest {
428            program: "cat".to_string(),
429            args: vec![],
430            cwd: None,
431            env: None,
432            timeout_ms: None,
433            fail_on_non_zero: false,
434            stdin: Some(CmdStdin::File(temp_file.path().to_path_buf())),
435            background: false,
436        };
437        let result = CmdTool::run(req);
438        assert!(result.is_ok());
439        let output = result.unwrap();
440        assert_eq!(output.exit_code, 0);
441        assert_eq!(output.stdout, "hello stdin file");
442        assert!(output.pid.is_none());
443    }
444
445    #[test]
446    fn test_stdin_null() {
447        let req = CmdRequest {
448            program: "cat".to_string(),
449            args: vec![],
450            cwd: None,
451            env: None,
452            timeout_ms: Some(1_000),
453            fail_on_non_zero: false,
454            stdin: Some(CmdStdin::Null),
455            background: false,
456        };
457        let result = CmdTool::run(req);
458        assert!(result.is_ok());
459        let output = result.unwrap();
460        assert_eq!(output.exit_code, 0);
461        assert!(output.stdout.is_empty());
462        assert!(output.pid.is_none());
463    }
464
465    #[test]
466    fn test_background() {
467        let req = CmdRequest {
468            program: "sleep".to_string(),
469            args: vec!["1".to_string()],
470            cwd: None,
471            env: None,
472            timeout_ms: None,
473            fail_on_non_zero: false,
474            stdin: None,
475            background: true,
476        };
477        let result = CmdTool::run(req);
478        assert!(result.is_ok());
479        let output = result.unwrap();
480        assert_eq!(output.exit_code, 0);
481        assert!(output.stdout.is_empty());
482        assert!(output.pid.is_some());
483        assert!(output.pid.unwrap() > 0);
484    }
485
486    #[test]
487    fn test_shell_pipe() {
488        let command = if cfg!(target_os = "windows") {
489            "echo hello pipe | findstr pipe"
490        } else {
491            "echo 'hello pipe' | grep pipe"
492        };
493
494        let req = ShellCmdRequest {
495            command: command.to_string(),
496            cwd: None,
497            env: None,
498            timeout_ms: None,
499            fail_on_non_zero: false,
500            stdin: None,
501            background: false,
502        };
503        let result = CmdTool::run_shell(req);
504        assert!(result.is_ok());
505        let output = result.unwrap();
506        assert_eq!(output.exit_code, 0);
507        assert!(output.stdout.contains("hello pipe"));
508        assert!(output.pid.is_none());
509    }
510
511    #[test]
512    fn test_non_zero_exit_can_fail() {
513        let req = if cfg!(target_os = "windows") {
514            ShellCmdRequest {
515                command: "cmd /c exit 7".to_string(),
516                cwd: None,
517                env: None,
518                timeout_ms: None,
519                fail_on_non_zero: true,
520                stdin: None,
521                background: false,
522            }
523        } else {
524            ShellCmdRequest {
525                command: "sh -c 'exit 7'".to_string(),
526                cwd: None,
527                env: None,
528                timeout_ms: None,
529                fail_on_non_zero: true,
530                stdin: None,
531                background: false,
532            }
533        };
534
535        let result = CmdTool::run_shell(req);
536        assert!(result.is_err());
537        let err_msg = result.unwrap_err().to_string().to_lowercase();
538        assert!(err_msg.contains("exit code 7"));
539    }
540
541    #[test]
542    fn test_non_zero_exit_can_be_observed_without_error() {
543        let req = if cfg!(target_os = "windows") {
544            ShellCmdRequest {
545                command: "cmd /c exit 9".to_string(),
546                cwd: None,
547                env: None,
548                timeout_ms: None,
549                fail_on_non_zero: false,
550                stdin: None,
551                background: false,
552            }
553        } else {
554            ShellCmdRequest {
555                command: "sh -c 'exit 9'".to_string(),
556                cwd: None,
557                env: None,
558                timeout_ms: None,
559                fail_on_non_zero: false,
560                stdin: None,
561                background: false,
562            }
563        };
564
565        let result = CmdTool::run_shell(req).unwrap();
566        assert_eq!(result.exit_code, 9);
567    }
568
569    #[cfg(not(target_os = "windows"))]
570    #[test]
571    fn test_non_utf8_stdout_is_preserved_lossily() {
572        let req = ShellCmdRequest {
573            command: "printf '\\377\\376abc'".to_string(),
574            cwd: None,
575            env: None,
576            timeout_ms: None,
577            fail_on_non_zero: false,
578            stdin: None,
579            background: false,
580        };
581
582        let result = CmdTool::run_shell(req).unwrap();
583        assert!(result.stdout.contains("abc"));
584        assert!(!result.stdout.is_empty());
585    }
586}
agent_tools/base/cmd.rs

agent_tools/base/
cmd.rs
