use crate::agent::WorkResult;
use crate::types::{Action, WorkerId};
pub trait Environment: Send + Sync {
fn step(&self, worker_id: WorkerId, action: &Action) -> WorkResult;
fn reset(&self);
fn name(&self) -> &str;
}
pub type EnvironmentBox = Box<dyn Environment>;
use std::fs;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::process::Command;
pub struct DefaultEnvironment {
working_dir: PathBuf,
}
impl DefaultEnvironment {
pub fn new() -> Self {
Self {
working_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
}
}
pub fn with_working_dir(working_dir: impl Into<PathBuf>) -> Self {
Self {
working_dir: working_dir.into(),
}
}
fn handle_bash(&self, action: &Action) -> WorkResult {
let command = action.params.target.as_deref().unwrap_or("");
let mut cmd = Command::new("sh");
cmd.arg("-c").arg(command);
cmd.current_dir(&self.working_dir);
match cmd.output() {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
if output.status.success() {
WorkResult::env_success_with_data("Command executed successfully", stdout)
} else {
WorkResult::env_failure(format!(
"Exit code: {:?}\nstderr: {}",
output.status.code(),
stderr
))
}
}
Err(e) => WorkResult::env_failure(format!("Failed to execute: {}", e)),
}
}
fn handle_read(&self, action: &Action) -> WorkResult {
let path = action.params.target.as_deref().unwrap_or("");
let full_path = self.resolve_path(path);
match fs::read_to_string(&full_path) {
Ok(content) => WorkResult::env_success_with_data("File read successfully", content),
Err(e) => WorkResult::env_failure(format!("Failed to read {}: {}", path, e)),
}
}
fn handle_write(&self, action: &Action) -> WorkResult {
let path = action.params.target.as_deref().unwrap_or("");
let content = action
.params
.args
.get("content")
.map(|s| s.as_str())
.unwrap_or("");
let full_path = self.resolve_path(path);
if let Some(parent) = full_path.parent() {
if !parent.exists() {
if let Err(e) = fs::create_dir_all(parent) {
return WorkResult::env_failure(format!("Failed to create directory: {}", e));
}
}
}
match fs::write(&full_path, content) {
Ok(()) => WorkResult::env_success(format!("Written to {}", path)),
Err(e) => WorkResult::env_failure(format!("Failed to write {}: {}", path, e)),
}
}
fn handle_grep(&self, action: &Action) -> WorkResult {
let target_str = action.params.target.as_deref().unwrap_or("");
let (pattern, search_path) = if let Some(p) = action.params.args.get("pattern") {
let path = if target_str.is_empty() {
"."
} else {
target_str
};
(p.as_str(), path)
} else if !target_str.is_empty()
&& !target_str.contains('/')
&& !target_str.contains('\\')
&& !target_str.ends_with(".rs")
&& !target_str.ends_with(".txt")
&& !target_str.ends_with(".toml")
{
(target_str, ".")
} else {
("", target_str)
};
let full_path = self.resolve_path(search_path);
if full_path.is_dir() {
return self.grep_directory(&full_path, pattern);
}
let file = match fs::File::open(&full_path) {
Ok(f) => f,
Err(e) => {
return WorkResult::env_failure(format!("Failed to open {}: {}", search_path, e))
}
};
let reader = BufReader::new(file);
let mut matches = Vec::new();
for (line_num, line) in reader.lines().enumerate() {
if let Ok(line) = line {
if line.contains(pattern) {
matches.push(format!("{}:{}:{}", search_path, line_num + 1, line));
}
}
}
WorkResult::env_success_with_data(
format!("Found {} matches", matches.len()),
matches.join("\n"),
)
}
fn grep_directory(&self, dir: &Path, pattern: &str) -> WorkResult {
let mut matches = Vec::new();
fn search_dir(dir: &Path, pattern: &str, matches: &mut Vec<String>) {
if let Ok(entries) = fs::read_dir(dir) {
for entry in entries.filter_map(|e| e.ok()) {
let path = entry.path();
if path.is_dir() {
if let Some(name) = path.file_name() {
let name = name.to_string_lossy();
if !name.starts_with('.') && name != "target" && name != "node_modules"
{
search_dir(&path, pattern, matches);
}
}
} else if path.extension().map(|e| e == "rs").unwrap_or(false) {
if let Ok(content) = fs::read_to_string(&path) {
for (line_num, line) in content.lines().enumerate() {
if line.contains(pattern) {
matches.push(format!(
"{}:{}:{}",
path.display(),
line_num + 1,
line
));
}
}
}
}
}
}
}
search_dir(dir, pattern, &mut matches);
WorkResult::env_success_with_data(
format!("Found {} matches", matches.len()),
matches.join("\n"),
)
}
fn handle_glob(&self, action: &Action) -> WorkResult {
let target_str = action.params.target.as_deref().unwrap_or(".");
let (pattern, search_dir) = if let Some(p) = action.params.args.get("pattern") {
(p.as_str(), target_str)
} else if target_str.contains('*') {
(target_str, ".")
} else {
("*", target_str)
};
let full_path = self.resolve_path(search_dir);
if pattern.contains("**") {
return self.glob_recursive(&full_path, pattern);
}
match fs::read_dir(&full_path) {
Ok(entries) => {
let files: Vec<String> = entries
.filter_map(|e| e.ok())
.filter(|e| {
if pattern == "*" {
return true;
}
if let Some(ext) = pattern.strip_prefix("*.") {
return e.path().extension().map(|x| x == ext).unwrap_or(false);
}
e.file_name().to_string_lossy().contains(pattern)
})
.map(|e| e.path().display().to_string())
.collect();
WorkResult::env_success_with_data(
format!("Found {} files", files.len()),
files.join("\n"),
)
}
Err(e) => WorkResult::env_failure(format!("Failed to read directory: {}", e)),
}
}
fn glob_recursive(&self, dir: &Path, pattern: &str) -> WorkResult {
let mut files = Vec::new();
let ext = if pattern.contains("*.") {
pattern.rsplit("*.").next()
} else {
None
};
fn collect_files(dir: &Path, ext: Option<&str>, files: &mut Vec<String>) {
if let Ok(entries) = fs::read_dir(dir) {
for entry in entries.filter_map(|e| e.ok()) {
let path = entry.path();
if path.is_dir() {
if let Some(name) = path.file_name() {
let name = name.to_string_lossy();
if !name.starts_with('.') && name != "target" && name != "node_modules"
{
collect_files(&path, ext, files);
}
}
} else if let Some(ext) = ext {
if path.extension().map(|e| e == ext).unwrap_or(false) {
files.push(path.display().to_string());
}
} else {
files.push(path.display().to_string());
}
}
}
}
collect_files(dir, ext, &mut files);
WorkResult::env_success_with_data(format!("Found {} files", files.len()), files.join("\n"))
}
fn handle_answer(&self, action: &Action) -> WorkResult {
let answer = action.params.target.as_deref().unwrap_or("");
WorkResult::done_success(format!("Answer: {}", answer))
}
fn handle_continue(&self, _action: &Action) -> WorkResult {
WorkResult::env_success("Continuing...")
}
fn resolve_path(&self, path: &str) -> PathBuf {
let p = Path::new(path);
if p.is_absolute() {
p.to_path_buf()
} else {
self.working_dir.join(p)
}
}
}
impl Default for DefaultEnvironment {
fn default() -> Self {
Self::new()
}
}
impl Environment for DefaultEnvironment {
fn step(&self, _worker_id: WorkerId, action: &Action) -> WorkResult {
match action.name.as_str() {
"Bash" => self.handle_bash(action),
"Read" => self.handle_read(action),
"Write" => self.handle_write(action),
"Grep" => self.handle_grep(action),
"Glob" => self.handle_glob(action),
"Answer" => self.handle_answer(action),
"Continue" => self.handle_continue(action),
_ => WorkResult::unsupported(&action.name),
}
}
fn reset(&self) {
}
fn name(&self) -> &str {
"DefaultEnvironment"
}
}