diff --git a/crates/agent/src/execution_env.rs b/crates/agent/src/execution_env.rs index 9e9c9c949..bb66f5ec1 100644 --- a/crates/agent/src/execution_env.rs +++ b/crates/agent/src/execution_env.rs @@ -4,6 +4,85 @@ use std::fmt::Write; use std::sync::Arc; use tokio_util::sync::CancellationToken; +/// Generates an `#[async_trait] impl ExecutionEnvironment` block for a decorator type +/// that wraps an `Arc`. The caller provides custom method +/// implementations; all remaining trait methods delegate to the inner field. +/// +/// # Usage +/// +/// ```ignore +/// delegate_execution_env! { +/// MyDecorator => inner { +/// // Only provide methods with custom logic — the rest delegate automatically. +/// async fn read_file(&self, path: &str, offset: Option, limit: Option) -> Result { +/// // custom logic... +/// } +/// } +/// } +/// ``` +#[macro_export] +macro_rules! delegate_execution_env { + ( + $type:ty => $field:ident { + $($custom:item)* + } + ) => { + #[async_trait::async_trait] + impl $crate::execution_env::ExecutionEnvironment for $type { + $($custom)* + + async fn file_exists(&self, path: &str) -> Result { + self.$field.file_exists(path).await + } + + async fn list_directory( + &self, + path: &str, + depth: Option, + ) -> Result, String> { + self.$field.list_directory(path, depth).await + } + + async fn exec_command( + &self, + command: &str, + timeout_ms: u64, + working_dir: Option<&str>, + env_vars: Option<&std::collections::HashMap>, + cancel_token: Option, + ) -> Result<$crate::execution_env::ExecResult, String> { + self.$field + .exec_command(command, timeout_ms, working_dir, env_vars, cancel_token) + .await + } + + async fn glob(&self, pattern: &str, path: Option<&str>) -> Result, String> { + self.$field.glob(pattern, path).await + } + + async fn initialize(&self) -> Result<(), String> { + self.$field.initialize().await + } + + async fn cleanup(&self) -> Result<(), String> { + self.$field.cleanup().await + } + + fn working_directory(&self) -> &str { + self.$field.working_directory() + } + + fn platform(&self) -> &str { + self.$field.platform() + } + + fn os_version(&self) -> String { + self.$field.os_version() + } + } + }; +} + /// Events emitted during execution environment lifecycle operations. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ExecutionEnvEvent { diff --git a/crates/agent/src/read_before_write_env.rs b/crates/agent/src/read_before_write_env.rs index fd7be07c2..3407a6be4 100644 --- a/crates/agent/src/read_before_write_env.rs +++ b/crates/agent/src/read_before_write_env.rs @@ -1,9 +1,7 @@ use crate::execution_env::*; -use async_trait::async_trait; use std::collections::HashSet; use std::path::{Component, PathBuf}; use std::sync::{Arc, Mutex}; -use tokio_util::sync::CancellationToken; /// Decorator that prevents writing to files the agent hasn't read first. /// @@ -73,93 +71,45 @@ impl ReadBeforeWriteEnvironment { } } -#[async_trait] -impl ExecutionEnvironment for ReadBeforeWriteEnvironment { - async fn read_file( - &self, - path: &str, - offset: Option, - limit: Option, - ) -> Result { - let result = self.inner.read_file(path, offset, limit).await?; - self.mark_read(path); - Ok(result) - } +crate::delegate_execution_env! { + ReadBeforeWriteEnvironment => inner { + async fn read_file( + &self, + path: &str, + offset: Option, + limit: Option, + ) -> Result { + let result = self.inner.read_file(path, offset, limit).await?; + self.mark_read(path); + Ok(result) + } - async fn write_file(&self, path: &str, content: &str) -> Result<(), String> { - self.guard_write(path).await?; - self.inner.write_file(path, content).await - } + async fn write_file(&self, path: &str, content: &str) -> Result<(), String> { + self.guard_write(path).await?; + self.inner.write_file(path, content).await + } - async fn delete_file(&self, path: &str) -> Result<(), String> { - self.guard_write(path).await?; - self.inner.delete_file(path).await - } + async fn delete_file(&self, path: &str) -> Result<(), String> { + self.guard_write(path).await?; + self.inner.delete_file(path).await + } - async fn file_exists(&self, path: &str) -> Result { - self.inner.file_exists(path).await - } - - async fn list_directory( - &self, - path: &str, - depth: Option, - ) -> Result, String> { - self.inner.list_directory(path, depth).await - } - - async fn exec_command( - &self, - command: &str, - timeout_ms: u64, - working_dir: Option<&str>, - env_vars: Option<&std::collections::HashMap>, - cancel_token: Option, - ) -> Result { - self.inner - .exec_command(command, timeout_ms, working_dir, env_vars, cancel_token) - .await - } - - async fn grep( - &self, - pattern: &str, - path: &str, - options: &GrepOptions, - ) -> Result, String> { - let results = self.inner.grep(pattern, path, options).await?; - for line in &results { - if let Some(file_path) = line.split(':').next() { - if !file_path.is_empty() { - self.mark_read(file_path); + async fn grep( + &self, + pattern: &str, + path: &str, + options: &GrepOptions, + ) -> Result, String> { + let results = self.inner.grep(pattern, path, options).await?; + for line in &results { + if let Some(file_path) = line.split(':').next() { + if !file_path.is_empty() { + self.mark_read(file_path); + } } } + Ok(results) } - Ok(results) - } - - async fn glob(&self, pattern: &str, path: Option<&str>) -> Result, String> { - self.inner.glob(pattern, path).await - } - - async fn initialize(&self) -> Result<(), String> { - self.inner.initialize().await - } - - async fn cleanup(&self) -> Result<(), String> { - self.inner.cleanup().await - } - - fn working_directory(&self) -> &str { - self.inner.working_directory() - } - - fn platform(&self) -> &str { - self.inner.platform() - } - - fn os_version(&self) -> String { - self.inner.os_version() } }