more tests and cleanup

Author: brandonskiser

crates/agent/src/agent/agent_loop/mod.rs

@@ -11,7 +11,7 @@ use futures::{
     Stream,
     StreamExt,
 };
-use model::AgentLoopModel;
+use model::Model;
 use protocol::{
     AgentLoopEventKind,
     AgentLoopRequest,
@@ -642,7 +642,7 @@ impl AgentLoopHandle {
         self.loop_event_rx.recv().await
     }
 
-    pub async fn send_request<M: AgentLoopModel>(
+    pub async fn send_request<M: Model>(
         &mut self,
         model: M,
         args: SendRequestArgs,

crates/agent/src/agent/agent_loop/model.rs

@@ -18,7 +18,8 @@ use crate::agent::rts::RtsModel;
 /// Represents a backend implementation for a converse stream compatible API.
 ///
 /// **Important** - implementations should be cancel safe
-pub trait Model {
+pub trait Model: std::fmt::Debug + Send + Sync + 'static {
+    /// Sends a conversation to a model, returning a stream of events as the response.
     fn stream(
         &self,
         messages: Vec<Message>,
@@ -28,12 +29,6 @@ pub trait Model {
     ) -> Pin<Box<dyn Stream<Item = Result<StreamEvent, StreamError>> + Send + 'static>>;
 }
 
-/// Required for defining [Model] with a [Box<dyn Model>] for [super::AgentLoopRequest].
-pub trait AgentLoopModel: Model + std::fmt::Debug + Send + Sync + 'static {}
-
-// Helper blanket impl
-impl<T> AgentLoopModel for T where T: Model + std::fmt::Debug + Send + Sync + 'static {}
-
 /// The supported backends
 #[derive(Debug, Clone)]
 pub enum Models {

crates/agent/src/agent/agent_loop/protocol.rs

@@ -10,7 +10,7 @@ use serde::{
 };
 use tokio::sync::mpsc;
 
-use super::model::AgentLoopModel;
+use super::model::Model;
 use super::types::{
     Message,
     MetadataEvent,
@@ -29,7 +29,7 @@ use super::{
 pub enum AgentLoopRequest {
     GetExecutionState,
     SendRequest {
-        model: Box<dyn AgentLoopModel>,
+        model: Box<dyn Model>,
         args: SendRequestArgs,
     },
     /// Ends the agent loop
@@ -135,8 +135,14 @@ pub enum AgentLoopEventKind {
     /// The agent loop has changed states
     LoopStateChange { from: LoopState, to: LoopState },
     /// Low level event. Generally only useful for [AgentLoop].
+    ///
+    /// This reflects the exact event the agent loop parses from a [Model::stream] response as part
+    /// of executing a user turn.
     StreamEvent(StreamEvent),
     /// Low level event. Generally only useful for [AgentLoop].
+    ///
+    /// This reflects the exact event the agent loop parses from a [Model::stream] response as part
+    /// of executing a user turn.
     StreamError(StreamError),
 }
 

crates/agent/src/agent/mod.rs

@@ -2104,11 +2104,3 @@ pub enum HookStage {
     /// Hooks after executing tool uses
     PostToolUse { tool_results: Vec<ToolExecutorResult> },
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test]
-    async fn test_collect_resources() {}
-}

crates/agent/src/agent/protocol.rs

@@ -29,13 +29,18 @@ use super::types::AgentSnapshot;
 pub enum AgentEvent {
     /// Agent has finished initialization, and is ready to receive requests
     Initialized,
-    /// Events associated with the agent loop
+    /// Events associated with the agent loop.
+    ///
+    /// These events contain information about the model's response, including:
+    /// - Text content
+    /// - Tool uses
+    /// - Metadata about a response stream, and about a complete user turn
     AgentLoop(AgentLoopEvent),
     /// The exact request sent to the backend
     RequestSent(SendRequestArgs),
     /// An unknown error occurred with the model backend that could not be handled by the agent.
     RequestError(LoopError),
-    /// An agent has changed state.
+    /// The agent has changed state.
     StateChange { from: ExecutionState, to: ExecutionState },
     /// A tool use was requested by the model, and the permission was evaluated
     ToolPermissionEvalResult {

crates/agent/src/agent/tools/file_read.rs

@@ -148,7 +148,7 @@ impl FsReadOp {
     async fn execute(&self) -> Result<ToolExecutionOutputItem, ToolExecutionError> {
         let path = PathBuf::from(canonicalize_path(&self.path).map_err(|e| ToolExecutionError::Custom(e.to_string()))?);
 
-        // add line numbers
+        // TODO: add line numbers
         let file_lines = LinesStream::new(
             BufReader::new(
                 fs::File::open(&path)
@@ -157,7 +157,10 @@ impl FsReadOp {
             )
             .lines(),
         );
-        let mut file_lines = file_lines.enumerate().skip(self.offset.unwrap_or_default() as usize);
+        let mut file_lines = file_lines
+            .enumerate()
+            .skip(self.offset.unwrap_or_default() as usize)
+            .take(self.limit.unwrap_or(u32::MAX) as usize);
 
         let mut is_truncated = false;
         let mut content = Vec::new();
@@ -190,10 +193,100 @@ pub struct FileReadContext {}
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::agent::util::test::TestDir;
 
-    #[test]
-    fn test_file_read_tool_schema() {
-        let schema = FsRead::tool_schema();
-        println!("{}", serde_json::to_string_pretty(&schema).unwrap());
+    #[tokio::test]
+    async fn test_fs_read_single_file() {
+        let test_dir = TestDir::new().with_file(("test.txt", "line1\nline2\nline3")).await;
+
+        let tool = FsRead {
+            ops: vec![FsReadOp {
+                path: test_dir.path("test.txt").to_string_lossy().to_string(),
+                limit: None,
+                offset: None,
+            }],
+        };
+
+        assert!(tool.validate().await.is_ok());
+        let result = tool.execute().await.unwrap();
+        assert_eq!(result.items.len(), 1);
+        if let ToolExecutionOutputItem::Text(content) = &result.items[0] {
+            assert_eq!(content, "line1\nline2\nline3");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_fs_read_with_offset_and_limit() {
+        let test_dir = TestDir::new()
+            .with_file(("test.txt", "line1\nline2\nline3\nline4\nline5"))
+            .await;
+
+        let tool = FsRead {
+            ops: vec![FsReadOp {
+                path: test_dir.path("test.txt").to_string_lossy().to_string(),
+                limit: Some(2),
+                offset: Some(1),
+            }],
+        };
+
+        let result = tool.execute().await.unwrap();
+        if let ToolExecutionOutputItem::Text(content) = &result.items[0] {
+            assert_eq!(content, "line2\nline3");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_fs_read_multiple_files() {
+        let test_dir = TestDir::new()
+            .with_file(("file1.txt", "content1"))
+            .await
+            .with_file(("file2.txt", "content2"))
+            .await;
+
+        let tool = FsRead {
+            ops: vec![
+                FsReadOp {
+                    path: test_dir.path("file1.txt").to_string_lossy().to_string(),
+                    limit: None,
+                    offset: None,
+                },
+                FsReadOp {
+                    path: test_dir.path("file2.txt").to_string_lossy().to_string(),
+                    limit: None,
+                    offset: None,
+                },
+            ],
+        };
+
+        let result = tool.execute().await.unwrap();
+        assert_eq!(result.items.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_fs_read_validate_nonexistent_file() {
+        let tool = FsRead {
+            ops: vec![FsReadOp {
+                path: "/nonexistent/file.txt".to_string(),
+                limit: None,
+                offset: None,
+            }],
+        };
+
+        assert!(tool.validate().await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_fs_read_validate_directory_path() {
+        let test_dir = TestDir::new();
+
+        let tool = FsRead {
+            ops: vec![FsReadOp {
+                path: test_dir.path("").to_string_lossy().to_string(),
+                limit: None,
+                offset: None,
+            }],
+        };
+
+        assert!(tool.validate().await.is_err());
     }
 }

crates/agent/src/agent/tools/file_write.rs

@@ -342,3 +342,141 @@ impl FileLineTracker {
         (self.lines_added_by_agent + self.lines_removed_by_agent) as isize
     }
 }
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agent::util::test::TestDir;
+
+    #[tokio::test]
+    async fn test_create_file() {
+        let test_dir = TestDir::new();
+        let tool = FsWrite::Create(FileCreate {
+            path: test_dir.path("new.txt").to_string_lossy().to_string(),
+            content: "hello world".to_string(),
+        });
+
+        assert!(tool.validate().await.is_ok());
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("new.txt")).await.unwrap();
+        assert_eq!(content, "hello world");
+    }
+
+    #[tokio::test]
+    async fn test_create_file_with_parent_dirs() {
+        let test_dir = TestDir::new();
+        let tool = FsWrite::Create(FileCreate {
+            path: test_dir.path("nested/dir/file.txt").to_string_lossy().to_string(),
+            content: "nested content".to_string(),
+        });
+
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("nested/dir/file.txt"))
+            .await
+            .unwrap();
+        assert_eq!(content, "nested content");
+    }
+
+    #[tokio::test]
+    async fn test_str_replace_single_occurrence() {
+        let test_dir = TestDir::new().with_file(("test.txt", "hello world")).await;
+
+        let tool = FsWrite::StrReplace(StrReplace {
+            path: test_dir.path("test.txt").to_string_lossy().to_string(),
+            old_str: "world".to_string(),
+            new_str: "rust".to_string(),
+            replace_all: false,
+        });
+
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("test.txt")).await.unwrap();
+        assert_eq!(content, "hello rust");
+    }
+
+    #[tokio::test]
+    async fn test_str_replace_multiple_occurrences() {
+        let test_dir = TestDir::new().with_file(("test.txt", "foo bar foo")).await;
+
+        let tool = FsWrite::StrReplace(StrReplace {
+            path: test_dir.path("test.txt").to_string_lossy().to_string(),
+            old_str: "foo".to_string(),
+            new_str: "baz".to_string(),
+            replace_all: true,
+        });
+
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("test.txt")).await.unwrap();
+        assert_eq!(content, "baz bar baz");
+    }
+
+    #[tokio::test]
+    async fn test_str_replace_no_match() {
+        let test_dir = TestDir::new().with_file(("test.txt", "hello world")).await;
+
+        let tool = FsWrite::StrReplace(StrReplace {
+            path: test_dir.path("test.txt").to_string_lossy().to_string(),
+            old_str: "missing".to_string(),
+            new_str: "replacement".to_string(),
+            replace_all: false,
+        });
+
+        assert!(tool.execute(None).await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_insert_at_line() {
+        let test_dir = TestDir::new().with_file(("test.txt", "line1\nline2\nline3")).await;
+
+        let tool = FsWrite::Insert(Insert {
+            path: test_dir.path("test.txt").to_string_lossy().to_string(),
+            content: "inserted".to_string(),
+            insert_line: Some(1),
+        });
+
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("test.txt")).await.unwrap();
+        assert_eq!(content, "line1\ninserted\nline2\nline3");
+    }
+
+    #[tokio::test]
+    async fn test_insert_append() {
+        let test_dir = TestDir::new().with_file(("test.txt", "existing")).await;
+
+        let tool = FsWrite::Insert(Insert {
+            path: test_dir.path("test.txt").to_string_lossy().to_string(),
+            content: "appended".to_string(),
+            insert_line: None,
+        });
+
+        assert!(tool.execute(None).await.is_ok());
+
+        let content = tokio::fs::read_to_string(test_dir.path("test.txt")).await.unwrap();
+        assert_eq!(content, "existing\nappended");
+    }
+
+    #[tokio::test]
+    async fn test_fs_write_validate_empty_path() {
+        let tool = FsWrite::Create(FileCreate {
+            path: "".to_string(),
+            content: "content".to_string(),
+        });
+
+        assert!(tool.validate().await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_fs_write_validate_nonexistent_file_for_replace() {
+        let tool = FsWrite::StrReplace(StrReplace {
+            path: "/nonexistent/file.txt".to_string(),
+            old_str: "old".to_string(),
+            new_str: "new".to_string(),
+            replace_all: false,
+        });
+
+        assert!(tool.validate().await.is_err());
+    }
+}