aither_command/
lib.rs

1use std::{borrow::Cow, path::PathBuf};
2
3use aither_core::llm::ToolContext;
4use aither_core::llm::{Tool, ToolResult, tool::json};
5use anyhow::{Context, Result, bail};
6use async_process::Command;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10/// Execute a shell command with arguments.
11///
12/// Runs a program with the specified arguments and returns stdout, stderr,
13/// and exit code. Use this for executing system commands when you need
14/// fine-grained control over arguments.
15///
16/// For complex scripts with pipelines, use `terminal` instead.
17#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
18pub struct CommandArgs {
19    /// Program name to execute (e.g., "ls", "grep", "git", "cargo").
20    pub program: String,
21    /// Command-line arguments as separate strings (e.g., ["-la"] for ls, ["-r", "TODO", "src/"] for grep).
22    #[serde(default)]
23    pub args: Vec<String>,
24    /// Working directory for command execution. Omit to use default.
25    pub cwd: Option<PathBuf>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
29pub struct CommandOutput {
30    pub program: String,
31    pub status: i32,
32    pub stdout: String,
33    pub stderr: String,
34}
35
36#[derive(Debug, Clone)]
37pub struct CommandTool {
38    allowed: Option<Vec<String>>,
39    default_cwd: PathBuf,
40    max_output: usize,
41    name: String,
42}
43
44impl CommandTool {
45    pub fn new(default_cwd: impl Into<PathBuf>) -> Self {
46        let default_cwd = default_cwd.into();
47        Self {
48            allowed: None,
49            default_cwd,
50            max_output: 16 * 1024,
51            name: "command".into(),
52        }
53    }
54
55    pub fn restrict_to(mut self, commands: impl IntoIterator<Item = impl Into<String>>) -> Self {
56        self.allowed = Some(commands.into_iter().map(Into::into).collect());
57        self
58    }
59
60    pub fn max_output(mut self, bytes: usize) -> Self {
61        self.max_output = bytes;
62        self
63    }
64
65    pub fn named(mut self, name: impl Into<String>) -> Self {
66        self.name = name.into();
67        self
68    }
69
70    fn ensure_allowed(&self, program: &str) -> Result<()> {
71        if let Some(allowed) = &self.allowed
72            && !allowed.iter().any(|entry| entry == program)
73        {
74            bail!(
75                "Program '{program}' is not allowed. Allowed commands: {}",
76                allowed.join(", ")
77            );
78        }
79        Ok(())
80    }
81
82    fn truncate(&self, text: String) -> String {
83        if text.len() <= self.max_output {
84            return text;
85        }
86
87        let mut truncated = text;
88        truncated.truncate(self.max_output);
89        truncated.push_str("\n...output truncated...");
90        truncated
91    }
92}
93
94impl Tool for CommandTool {
95    fn name(&self) -> Cow<'static, str> {
96        Cow::Owned(self.name.clone())
97    }
98
99    type Arguments = CommandArgs;
100    type Res = ToolResult;
101
102    async fn call(
103        &self,
104        arguments: Self::Arguments,
105        _cx: ToolContext,
106    ) -> aither_core::Result<Self::Res> {
107        self.ensure_allowed(&arguments.program)?;
108
109        let working_dir = arguments.cwd.unwrap_or_else(|| self.default_cwd.clone());
110        let output = Command::new(&arguments.program)
111            .args(&arguments.args)
112            .current_dir(&working_dir)
113            .output()
114            .await
115            .with_context(|| {
116                format!(
117                    "failed to execute '{}' in {}",
118                    arguments.program,
119                    working_dir.display()
120                )
121            })?;
122
123        let response = CommandOutput {
124            program: arguments.program,
125            status: output.status.code().unwrap_or_default(),
126            stdout: self.truncate(String::from_utf8_lossy(&output.stdout).to_string()),
127            stderr: self.truncate(String::from_utf8_lossy(&output.stderr).to_string()),
128        };
129
130        Ok(ToolResult::text(json(&response)?))
131    }
132}
aither_command/lib.rs

aither_command/
lib.rs
