Skip to main content

ai_agents_tools/builtin/
command.rs

1use async_trait::async_trait;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::collections::HashMap;
6
7use ai_agents_core::{
8    CommandPolicyBinding, PathAccessMode, PathBindingKind, PathPolicyBinding, ResultLimitBinding,
9    ResultLimitKind, Tool, ToolCallClassification, ToolExecutionContext, ToolOperationKind,
10    ToolPolicyBindings, ToolResult, ToolSafetyMetadata, ToolSideEffectLevel,
11};
12
13use crate::generate_schema;
14use crate::types::{CommandRequest, CommandResponse, CommandRunnerSlot};
15
16const DEFAULT_TIMEOUT_MS: u64 = 30_000;
17const DEFAULT_MAX_OUTPUT_CHARS: usize = 20_000;
18
19/// Runs allowlisted non-interactive commands through a host runner.
20pub struct CommandTool {
21    runner: CommandRunnerSlot,
22}
23
24impl CommandTool {
25    /// Create a command tool backed by a shared runner slot.
26    pub fn new(runner: CommandRunnerSlot) -> Self {
27        Self { runner }
28    }
29}
30
31#[derive(Debug, Deserialize, JsonSchema)]
32struct CommandInput {
33    /// Preferred command vector including executable and arguments.
34    #[serde(default)]
35    argv: Vec<String>,
36    /// Compatibility command string. Shell syntax is rejected by default.
37    #[serde(default)]
38    command: Option<String>,
39    /// Working directory. Defaults to current directory.
40    #[serde(default)]
41    cwd: Option<String>,
42    /// Environment variables to pass when policy allows them.
43    #[serde(default)]
44    env: HashMap<String, String>,
45    /// Timeout in milliseconds.
46    #[serde(default)]
47    timeout_ms: Option<u64>,
48    /// Maximum combined output characters.
49    #[serde(default)]
50    max_output_chars: Option<usize>,
51    /// User-visible reason for running the command.
52    #[serde(default)]
53    reason: Option<String>,
54}
55
56#[derive(Debug, Serialize)]
57struct CommandToolOutput {
58    success: bool,
59    exit_code: Option<i32>,
60    termination: String,
61    stdout: String,
62    stderr: String,
63    combined_output: String,
64    truncated: bool,
65    timed_out: bool,
66    cwd: String,
67    argv: Vec<String>,
68    reason: Option<String>,
69}
70
71#[async_trait]
72impl Tool for CommandTool {
73    fn id(&self) -> &str {
74        "command"
75    }
76
77    fn name(&self) -> &str {
78        "Command"
79    }
80
81    fn description(&self) -> &str {
82        "Run exact allowlisted non-interactive argv commands with timeout and bounded output."
83    }
84
85    fn input_schema(&self) -> Value {
86        generate_schema::<CommandInput>()
87    }
88
89    fn safety_metadata(&self) -> ToolSafetyMetadata {
90        ToolSafetyMetadata {
91            read_only: false,
92            concurrency_safe: false,
93            operation: ToolOperationKind::Command,
94            side_effect_level: ToolSideEffectLevel::LocalWrite,
95            requires_network: false,
96            destructive: false,
97            open_world: false,
98            host_dependent: true,
99            requires_user_interaction: false,
100            supports_cancellation: true,
101            default_requires_approval: true,
102            should_defer_schema: false,
103            max_output_chars: Some(DEFAULT_MAX_OUTPUT_CHARS),
104            max_result_size_chars: Some(DEFAULT_MAX_OUTPUT_CHARS),
105        }
106    }
107
108    fn classify_call(&self, _args: &Value) -> ToolCallClassification {
109        let mut classification = ToolCallClassification::from_metadata(&self.safety_metadata());
110        classification.safely_retryable = false;
111        classification
112    }
113
114    fn policy_bindings(&self) -> ToolPolicyBindings {
115        ToolPolicyBindings {
116            command_fields: vec![
117                CommandPolicyBinding::argv("argv"),
118                CommandPolicyBinding::command("command"),
119                CommandPolicyBinding::env("env"),
120            ],
121            path_fields: vec![
122                PathPolicyBinding::new("cwd", PathAccessMode::ReadWrite, PathBindingKind::Cwd)
123                    .with_default_path("."),
124            ],
125            result_limit_fields: vec![ResultLimitBinding::new(
126                "max_output_chars",
127                ResultLimitKind::MaxOutputChars,
128            )],
129            ..Default::default()
130        }
131    }
132
133    async fn execute(&self, args: Value, ctx: ToolExecutionContext) -> ToolResult {
134        let input: CommandInput = match serde_json::from_value(args) {
135            Ok(input) => input,
136            Err(error) => return ToolResult::error(format!("Invalid input: {}", error)),
137        };
138        let argv = match command_argv(&input) {
139            Ok(argv) => argv,
140            Err(error) => return ToolResult::error(error),
141        };
142        if argv.is_empty() {
143            return ToolResult::error("argv must not be empty");
144        }
145        if input.command.is_some() && contains_shell_syntax(&argv.join(" ")) {
146            return ToolResult::error("command string contains shell syntax denied by default");
147        }
148        let policy = CommandPolicySnapshot::from_context(&ctx.policy_snapshot);
149        let env = filter_env(input.env, &policy);
150        let timeout_ms = input
151            .timeout_ms
152            .unwrap_or(DEFAULT_TIMEOUT_MS)
153            .min(ctx.limits.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS));
154        let max_output_chars = input
155            .max_output_chars
156            .unwrap_or(DEFAULT_MAX_OUTPUT_CHARS)
157            .min(
158                ctx.limits
159                    .max_output_chars
160                    .unwrap_or(DEFAULT_MAX_OUTPUT_CHARS),
161            );
162        let cwd = input.cwd.unwrap_or_else(|| ".".to_string());
163        let request = CommandRequest {
164            argv: argv.clone(),
165            cwd: Some(cwd.clone()),
166            env,
167            timeout_ms: Some(timeout_ms),
168            max_output_chars: Some(max_output_chars),
169            reason: input.reason.clone(),
170        };
171        let runner = self.runner.read().clone();
172        let response = runner.run_command(request, ctx).await;
173        command_response_to_result(response, cwd, argv, input.reason)
174    }
175}
176
177#[derive(Debug, Default)]
178struct CommandPolicySnapshot {
179    env_passthrough: Vec<String>,
180    redact_env: Vec<String>,
181}
182
183impl CommandPolicySnapshot {
184    fn from_context(value: &Value) -> Self {
185        let mut snapshot = Self::default();
186        snapshot
187            .env_passthrough
188            .extend(strings_at(value, "env_passthrough"));
189        snapshot.redact_env.extend(strings_at(value, "redact_env"));
190        if let Some(commands) = value.get("commands") {
191            snapshot
192                .env_passthrough
193                .extend(strings_at(commands, "env_passthrough"));
194        }
195        snapshot
196    }
197}
198
199fn command_argv(input: &CommandInput) -> Result<Vec<String>, String> {
200    if !input.argv.is_empty() {
201        return Ok(input.argv.clone());
202    }
203    let Some(command) = input.command.as_deref() else {
204        return Err("either argv or command is required".to_string());
205    };
206    if contains_shell_syntax(command) {
207        return Err("command string contains shell syntax denied by default".to_string());
208    }
209    parse_command_words(command).ok_or_else(|| "command string could not be parsed".to_string())
210}
211
212fn filter_env(
213    env: HashMap<String, String>,
214    policy: &CommandPolicySnapshot,
215) -> HashMap<String, String> {
216    if policy.env_passthrough.is_empty() {
217        return HashMap::new();
218    }
219    env.into_iter()
220        .filter(|(key, _)| policy.env_passthrough.iter().any(|allowed| allowed == key))
221        .filter(|(key, _)| !policy.redact_env.iter().any(|redacted| redacted == key))
222        .collect()
223}
224
225fn strings_at(value: &Value, field: &str) -> Vec<String> {
226    value
227        .get(field)
228        .and_then(Value::as_array)
229        .into_iter()
230        .flatten()
231        .filter_map(Value::as_str)
232        .map(str::to_string)
233        .collect()
234}
235
236fn command_response_to_result(
237    response: CommandResponse,
238    cwd: String,
239    argv: Vec<String>,
240    reason: Option<String>,
241) -> ToolResult {
242    let output = CommandToolOutput {
243        success: response.success,
244        exit_code: response.exit_code,
245        termination: response.termination,
246        stdout: response.stdout,
247        stderr: response.stderr,
248        combined_output: response.combined_output,
249        truncated: response.truncated,
250        timed_out: response.timed_out,
251        cwd,
252        argv: if response.argv_redacted.is_empty() {
253            redact_argv(&argv)
254        } else {
255            response.argv_redacted
256        },
257        reason,
258    };
259    let json = match serde_json::to_string(&output) {
260        Ok(json) => json,
261        Err(error) => return ToolResult::error(format!("Serialization error: {}", error)),
262    };
263    let mut metadata = HashMap::new();
264    metadata.insert("truncated".to_string(), Value::Bool(output.truncated));
265    metadata.insert("timed_out".to_string(), Value::Bool(output.timed_out));
266    metadata.insert(
267        "timeout_cleanup".to_string(),
268        Value::String(if output.timed_out {
269            "kill_on_drop".to_string()
270        } else {
271            "not_needed".to_string()
272        }),
273    );
274    metadata.insert("argv".to_string(), serde_json::json!(output.argv));
275    ToolResult::ok_with_metadata(json, metadata)
276}
277
278fn redact_argv(argv: &[String]) -> Vec<String> {
279    let mut redacted = Vec::with_capacity(argv.len());
280    let mut redact_next = false;
281    for arg in argv {
282        let lower = arg.to_ascii_lowercase();
283        let sensitive = lower.contains("token")
284            || lower.contains("secret")
285            || lower.contains("password")
286            || lower.contains("apikey")
287            || lower.contains("api-key");
288        if redact_next || sensitive {
289            redacted.push("[redacted]".to_string());
290        } else {
291            redacted.push(arg.clone());
292        }
293        redact_next = matches!(
294            lower.as_str(),
295            "--token" | "--secret" | "--password" | "--api-key"
296        );
297    }
298    redacted
299}
300
301fn contains_shell_syntax(value: &str) -> bool {
302    const DENIED: &[char] = &[';', '&', '|', '<', '>', '`', '$', '\n', '\r'];
303    value.chars().any(|ch| DENIED.contains(&ch))
304        || value.contains("$(")
305        || value.contains("${")
306        || value.contains("<(")
307        || value.contains(">(")
308}
309
310fn parse_command_words(value: &str) -> Option<Vec<String>> {
311    let mut words = Vec::new();
312    let mut current = String::new();
313    let mut quote: Option<char> = None;
314    for ch in value.chars() {
315        match (quote, ch) {
316            (Some(q), c) if c == q => quote = None,
317            (Some(_), c) => current.push(c),
318            (None, '\'' | '"') => quote = Some(ch),
319            (None, c) if c.is_whitespace() => {
320                if !current.is_empty() {
321                    words.push(std::mem::take(&mut current));
322                }
323            }
324            (None, c) => current.push(c),
325        }
326    }
327    if quote.is_some() {
328        return None;
329    }
330    if !current.is_empty() {
331        words.push(current);
332    }
333    (!words.is_empty()).then_some(words)
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use crate::types::{CommandResponse, StaticCommandRunner};
340    use parking_lot::RwLock;
341    use std::sync::Arc;
342
343    #[tokio::test]
344    async fn static_runner_executes_allowed_argv() {
345        let mut responses = HashMap::new();
346        responses.insert(
347            vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
348            CommandResponse {
349                success: true,
350                exit_code: Some(0),
351                termination: "exited".to_string(),
352                stdout: "ok".to_string(),
353                combined_output: "ok".to_string(),
354                argv_redacted: vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
355                ..CommandResponse::default()
356            },
357        );
358        let runner = Arc::new(RwLock::new(
359            Arc::new(StaticCommandRunner::new(responses)) as Arc<_>
360        ));
361        let tool = CommandTool::new(runner);
362        let result = tool
363            .execute(
364                serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
365                ToolExecutionContext::test("command"),
366            )
367            .await;
368        assert!(result.success);
369    }
370
371    #[tokio::test]
372    async fn command_metadata_uses_redacted_argv() {
373        let mut responses = HashMap::new();
374        responses.insert(
375            vec![
376                "deploy".to_string(),
377                "--token".to_string(),
378                "abc123".to_string(),
379            ],
380            CommandResponse {
381                success: true,
382                termination: "exited".to_string(),
383                combined_output: "ok".to_string(),
384                ..CommandResponse::default()
385            },
386        );
387        let runner = Arc::new(RwLock::new(
388            Arc::new(StaticCommandRunner::new(responses)) as Arc<_>
389        ));
390        let tool = CommandTool::new(runner);
391        let result = tool
392            .execute(
393                serde_json::json!({"argv": ["deploy", "--token", "abc123"]}),
394                ToolExecutionContext::test("command"),
395            )
396            .await;
397        assert!(result.success);
398        let metadata = result.metadata.unwrap();
399        assert_eq!(
400            metadata["argv"],
401            serde_json::json!(["deploy", "[redacted]", "[redacted]"])
402        );
403    }
404
405    #[tokio::test]
406    async fn command_string_rejects_shell_syntax() {
407        let runner = Arc::new(RwLock::new(
408            Arc::new(StaticCommandRunner::default()) as Arc<_>
409        ));
410        let tool = CommandTool::new(runner);
411        let result = tool
412            .execute(
413                serde_json::json!({"command": "cargo test && rm -rf target"}),
414                ToolExecutionContext::test("command"),
415            )
416            .await;
417        assert!(!result.success);
418    }
419}