use crate::error::Result;
use crate::tools::{ToolExecutor, ToolResult};
use async_trait::async_trait;
use std::path::Path;
use std::sync::Arc;
#[async_trait]
pub trait ShellPreprocessor: Send + Sync {
async fn process(
&self,
content: &str,
workspace: &Path,
executor: &Arc<ToolExecutor>,
) -> Result<String>;
}
pub struct DefaultShellPreprocessor {
max_output_length: usize,
}
impl DefaultShellPreprocessor {
pub fn new() -> Self {
Self {
max_output_length: 10 * 1024, }
}
fn process_inline(&self, content: &str) -> Vec<(String, String)> {
let mut results = Vec::new();
let re = regex::Regex::new(r"(?:^|\s)(!`([^`]+)`)").unwrap();
for cap in re.captures_iter(content) {
let full_match = cap.get(1).unwrap().as_str();
let command = cap.get(2).unwrap().as_str().trim();
results.push((full_match.to_string(), command.to_string()));
}
results
}
fn process_blocks(&self, content: &str) -> Vec<(String, String)> {
let mut results = Vec::new();
let re = regex::Regex::new(r"```!\s*\n?([\s\S]*?)\n?```").unwrap();
for cap in re.captures_iter(content) {
let full_match = cap.get(0).unwrap().as_str();
let command = cap.get(1).unwrap().as_str().trim();
if !command.is_empty() {
results.push((full_match.to_string(), command.to_string()));
}
}
results
}
}
impl Default for DefaultShellPreprocessor {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ShellPreprocessor for DefaultShellPreprocessor {
async fn process(
&self,
content: &str,
workspace: &Path,
executor: &Arc<ToolExecutor>,
) -> Result<String> {
let mut result = content.to_string();
for (pattern, command) in self.process_blocks(content) {
let output = self.execute_command(&command, workspace, executor).await?;
result = result.replace(&pattern, &output);
}
if result.contains("!`") {
for (pattern, command) in self.process_inline(&result) {
let output = self.execute_command(&command, workspace, executor).await?;
result = result.replace(&pattern, &output);
}
}
Ok(result)
}
}
impl DefaultShellPreprocessor {
async fn execute_command(
&self,
command: &str,
workspace: &Path,
executor: &Arc<ToolExecutor>,
) -> Result<String> {
let args = serde_json::json!({ "command": command });
let result: ToolResult = executor
.execute_with_context(
"bash",
&args,
&crate::tools::ToolContext {
workspace: workspace.to_path_buf(),
session_id: None,
event_tx: None,
agent_event_tx: None,
search_config: None,
sandbox: None,
command_env: None,
},
)
.await?;
if result.exit_code != 0 {
return Ok(format!("[Error: {}]", result.output));
}
let output = result.output.trim().to_string();
if output.len() > self.max_output_length {
Ok(format!(
"{}... [truncated {} bytes]",
&output[..self.max_output_length],
output.len() - self.max_output_length
))
} else {
Ok(output)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_process_inline_pattern() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("Run !`echo hello` now");
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].0, "!`echo hello`");
assert_eq!(matches[0].1, "echo hello");
}
#[test]
fn test_process_block_pattern() {
let preprocessor = DefaultShellPreprocessor::new();
let content = r#"```!
ls -la
```"#;
let matches = preprocessor.process_blocks(content);
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].1, "ls -la");
}
#[test]
fn test_process_block_pattern_with_newlines() {
let preprocessor = DefaultShellPreprocessor::new();
let content = r#"```!
ls -la
```"#;
let matches = preprocessor.process_blocks(content);
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].1, "ls -la");
}
#[test]
fn test_no_false_positives() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("echo $! and !`echo hi`");
assert_eq!(matches.len(), 1); assert_eq!(matches[0].1, "echo hi");
}
#[test]
fn test_process_inline_at_start_of_line() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("!`pwd`");
assert_eq!(matches.len(), 1);
assert_eq!(matches[0].1, "pwd");
}
#[test]
fn test_process_inline_multiple() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("First !`cmd1` then !`cmd2`");
assert_eq!(matches.len(), 2);
assert_eq!(matches[0].1, "cmd1");
assert_eq!(matches[1].1, "cmd2");
}
#[test]
fn test_process_inline_no_matches() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("No shell commands here");
assert!(matches.is_empty());
}
#[test]
fn test_process_inline_backtick_in_string() {
let preprocessor = DefaultShellPreprocessor::new();
let matches = preprocessor.process_inline("`echo hello`");
assert!(matches.is_empty());
}
#[test]
fn test_process_blocks_multiple() {
let preprocessor = DefaultShellPreprocessor::new();
let content = r#"```!
ls -la
```
Some text
```!
pwd
```"#;
let matches = preprocessor.process_blocks(content);
assert_eq!(matches.len(), 2);
assert_eq!(matches[0].1, "ls -la");
assert_eq!(matches[1].1, "pwd");
}
#[test]
fn test_process_blocks_no_matches() {
let preprocessor = DefaultShellPreprocessor::new();
let content = "```\nsome code\n```";
let matches = preprocessor.process_blocks(content);
assert!(matches.is_empty());
}
#[test]
fn test_process_blocks_empty_command() {
let preprocessor = DefaultShellPreprocessor::new();
let content = "```!\n```";
let matches = preprocessor.process_blocks(content);
assert!(matches.is_empty());
}
#[test]
fn test_default_shell_preprocessor_new() {
let preprocessor = DefaultShellPreprocessor::new();
assert_eq!(preprocessor.max_output_length, 10 * 1024);
}
}