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
19pub struct CommandTool {
21 runner: CommandRunnerSlot,
22}
23
24impl CommandTool {
25 pub fn new(runner: CommandRunnerSlot) -> Self {
27 Self { runner }
28 }
29}
30
31#[derive(Debug, Deserialize, JsonSchema)]
32struct CommandInput {
33 #[serde(default)]
35 argv: Vec<String>,
36 #[serde(default)]
38 command: Option<String>,
39 #[serde(default)]
41 cwd: Option<String>,
42 #[serde(default)]
44 env: HashMap<String, String>,
45 #[serde(default)]
47 timeout_ms: Option<u64>,
48 #[serde(default)]
50 max_output_chars: Option<usize>,
51 #[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}