Add tests

Author: brandonskiser

Cargo.lock

@@ -64,7 +64,7 @@ shellexpand.workspace = true
 strum.workspace = true
 syntect = "5.2.0"
 sysinfo.workspace = true
-textwrap = "0.16.2"
+tempfile.workspace = true
 thiserror.workspace = true
 time.workspace = true
 tokio.workspace = true

crates/agent/Cargo.toml

@@ -6,23 +6,33 @@ use std::str::FromStr;
 use crate::agent::agent_loop::types::ToolUseBlock;
 use crate::agent::protocol::AgentError;
 use crate::agent::tools::BuiltInToolName;
-use crate::agent::util::path::canonicalize_path;
+use crate::agent::util::path::canonicalize_path_impl;
+use crate::agent::util::providers::{
+    RealProvider,
+    SystemProvider,
+};
 
 /// Represents a value from the `resources` array in the agent config.
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ResourceKind<'a> {
     File { original: &'a str, file_path: &'a str },
     FileGlob { original: &'a str, pattern: glob::Pattern },
 }
 
 impl<'a> ResourceKind<'a> {
     pub fn parse(value: &'a str) -> Result<Self, String> {
+        let sys = RealProvider;
+        Self::parse_impl(value, &sys)
+    }
+
+    fn parse_impl(value: &'a str, sys: &impl SystemProvider) -> Result<Self, String> {
         if !value.starts_with("file://") {
             return Err("Only file schemes are currently supported".to_string());
         }
 
         let file_path = value.trim_start_matches("file://");
         if file_path.contains('*') || file_path.contains('?') {
-            let canon = canonicalize_path(file_path)
+            let canon = canonicalize_path_impl(file_path, sys, sys, sys)
                 .map_err(|err| format!("Failed to canonicalize path for {}: {}", file_path, err))?;
             let pattern = glob::Pattern::new(canon.as_str())
                 .map_err(|err| format!("Failed to create glob for {}: {}", canon, err))?;
@@ -244,3 +254,37 @@ impl FromStr for CanonicalToolName {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agent::util::test::TestSystem;
+
+    #[test]
+    fn test_resource_kind_parse_nonfile() {
+        assert!(
+            ResourceKind::parse("https://google.com").is_err(),
+            "non-file scheme should be an error"
+        );
+    }
+
+    #[test]
+    fn test_resource_kind_parse_file_scheme() {
+        let sys = TestSystem::new();
+
+        let resource = "file://project/README.md";
+        assert_eq!(ResourceKind::parse_impl(resource, &sys).unwrap(), ResourceKind::File {
+            original: resource,
+            file_path: "project/README.md"
+        });
+
+        let resource = "file://~/project/**/*.rs";
+        assert_eq!(
+            ResourceKind::parse_impl(resource, &sys).unwrap(),
+            ResourceKind::FileGlob {
+                original: resource,
+                pattern: glob::Pattern::new("/home/testuser/project/**/*.rs").unwrap()
+            }
+        );
+    }
+}

crates/agent/src/agent/agent_config/parse.rs

@@ -2104,3 +2104,11 @@ 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/mod.rs

@@ -49,6 +49,10 @@ FEATURES:
 LIMITATIONS:
 - Maximum file size is 250KB
 - Cannot display binary files or images
+
+TIPS:
+- Read multiple files in one go if you know you want to read more than one file
+- Dont use limit and offset for small files
 "#;
 
 // TODO - migrate from JsonSchema, it's not very configurable and prone to breaking changes in the

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

@@ -44,25 +44,25 @@ const FS_WRITE_SCHEMA: &str = r#"
             "type": "string",
             "enum": [
                 "create",
-                "str_replace",
+                "strReplace",
                 "insert"
             ],
-            "description": "The commands to run. Allowed options are: `create`, `str_replace`, `insert`"
+            "description": "The commands to run. Allowed options are: `create`, `strReplace`, `insert`"
         },
         "content": {
             "description": "Required parameter of `create` and `insert` commands.",
             "type": "string"
         },
-        "insert_line": {
+        "insertLine": {
             "description": "Optional parameter of `insert` command. Line is 0-indexed. `content` will be inserted at the provided line. If not provided, content will be inserted at the end of the file on a new line, inserting a newline at the end of the file if it is missing.",
             "type": "integer"
         },
-        "new_str": {
-            "description": "Required parameter of `str_replace` command containing the new string.",
+        "newStr": {
+            "description": "Required parameter of `strReplace` command containing the new string.",
             "type": "string"
         },
-        "old_str": {
-            "description": "Required parameter of `str_replace` command containing the string in `path` to replace.",
+        "oldStr": {
+            "description": "Required parameter of `strReplace` command containing the string in `path` to replace.",
             "type": "string"
         },
         "path": {

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

@@ -5,6 +5,8 @@ pub mod glob;
 pub mod path;
 pub mod providers;
 pub mod request_channel;
+#[cfg(test)]
+pub mod test;
 
 use std::collections::HashMap;
 use std::env::VarError;
@@ -100,6 +102,15 @@ pub async fn read_file_with_max_limit(
         .await
         .with_context(|| format!("Failed to query file metadata at '{}'", path.to_string_lossy()))?;
 
+    // Read only the max supported length.
+    let mut reader = BufReader::new(file).take(max_file_length);
+    let mut content = Vec::new();
+    reader
+        .read_to_end(&mut content)
+        .await
+        .with_context(|| format!("Failed to read from file at '{}'", path.to_string_lossy()))?;
+    let mut content = content.to_str_lossy().to_string();
+
     let truncated_amount = if md.size() > max_file_length {
         // Edge case check to ensure the suffix is less than max file length.
         if suffix.len() as u64 > max_file_length {
@@ -110,18 +121,11 @@ pub async fn read_file_with_max_limit(
         0
     };
 
-    // Read only the max supported length.
-    let mut reader = BufReader::new(file).take(max_file_length);
-    let mut content = Vec::new();
-    reader
-        .read_to_end(&mut content)
-        .await
-        .with_context(|| format!("Failed to read from file at '{}'", path.to_string_lossy()))?;
-
-    // Truncate content safely.
-    let mut content = content.to_str_lossy().to_string();
-    truncate_safe_in_place(&mut content, max_file_length as usize, suffix);
+    if truncated_amount == 0 {
+        return Ok((content, 0));
+    }
 
+    content.replace_range((content.len().saturating_sub(suffix.len())).., suffix);
     Ok((content, truncated_amount))
 }
 
@@ -132,6 +136,7 @@ pub fn is_integ_test() -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::agent::util::test::TestDir;
 
     #[test]
     fn test_truncate_safe() {
@@ -185,4 +190,26 @@ mod tests {
         assert_eq!(env_vars.get("KEY1").unwrap(), "Value is test_value");
         assert_eq!(env_vars.get("KEY2").unwrap(), "No substitution");
     }
+
+    #[tokio::test]
+    async fn test_read_file_with_max_limit() {
+        // Test file with 30 bytes in length
+        let test_file = "123456789\n".repeat(3);
+        let d = TestDir::new().with_file(("test.txt", &test_file)).await;
+
+        // Test not truncated
+        let (content, bytes_truncated) = read_file_with_max_limit(d.path("test.txt"), 100, "...").await.unwrap();
+        assert_eq!(content, test_file);
+        assert_eq!(bytes_truncated, 0);
+
+        // Test truncated
+        let (content, bytes_truncated) = read_file_with_max_limit(d.path("test.txt"), 10, "...").await.unwrap();
+        assert_eq!(content, "1234567...");
+        assert_eq!(bytes_truncated, 23);
+
+        // Test suffix greater than max length
+        let (content, bytes_truncated) = read_file_with_max_limit(d.path("test.txt"), 1, "...").await.unwrap();
+        assert_eq!(content, "");
+        assert_eq!(bytes_truncated, 30);
+    }
 }

crates/agent/src/agent/util/mod.rs

@@ -12,12 +12,12 @@ use super::providers::{
     CwdProvider,
     EnvProvider,
     HomeProvider,
-    SystemProvider,
+    RealProvider,
 };
 
 /// Performs tilde and environment variable expansion on the provided input.
 pub fn expand_path(input: &str) -> Result<Cow<'_, str>, UtilError> {
-    let sys = SystemProvider;
+    let sys = RealProvider;
     Ok(shellexpand::full_with_context(
         input,
         sys.shellexpand_home(),
@@ -32,7 +32,7 @@ pub fn expand_path(input: &str) -> Result<Cow<'_, str>, UtilError> {
 /// - Performs env var expansion
 /// - Resolves `.` and `..` path components
 pub fn canonicalize_path(path: impl AsRef<str>) -> Result<String, UtilError> {
-    let sys = SystemProvider;
+    let sys = RealProvider;
     canonicalize_path_impl(path, &sys, &sys, &sys)
 }
 
@@ -97,7 +97,7 @@ fn normalize_path(path: &Path) -> PathBuf {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::agent::util::providers::TestSystem;
+    use crate::agent::util::test::TestSystem;
 
     #[test]
     fn test_canonicalize_path() {

crates/agent/src/agent/util/path.rs

@@ -3,6 +3,12 @@ use std::path::PathBuf;
 
 use super::directories;
 
+/// A trait for accessing system and process context (env vars, home dir, current working dir,
+/// etc.).
+pub trait SystemProvider: EnvProvider + HomeProvider + CwdProvider {}
+
+impl<T> SystemProvider for T where T: EnvProvider + HomeProvider + CwdProvider {}
+
 /// A trait for accessing environment variables.
 ///
 /// This provides unit tests the capability to fake system context.
@@ -36,84 +42,22 @@ pub trait CwdProvider {
 
 /// Provides real implementations for [EnvProvider], [HomeProvider], and [CwdProvider].
 #[derive(Clone, Copy)]
-pub struct SystemProvider;
+pub struct RealProvider;
 
-impl EnvProvider for SystemProvider {
+impl EnvProvider for RealProvider {
     fn var(&self, input: &str) -> Result<String, VarError> {
         std::env::var(input)
     }
 }
 
-impl HomeProvider for SystemProvider {
+impl HomeProvider for RealProvider {
     fn home(&self) -> Option<PathBuf> {
         directories::home_dir().ok()
     }
 }
 
-impl CwdProvider for SystemProvider {
+impl CwdProvider for RealProvider {
     fn cwd(&self) -> Result<PathBuf, std::io::Error> {
         std::env::current_dir()
     }
 }
-
-#[cfg(test)]
-#[derive(Debug, Clone)]
-pub struct TestSystem {
-    env: std::collections::HashMap<String, String>,
-    home: Option<PathBuf>,
-    cwd: Option<PathBuf>,
-}
-
-#[cfg(test)]
-impl TestSystem {
-    pub fn new() -> Self {
-        let mut env = std::collections::HashMap::new();
-        env.insert("HOME".to_string(), "/home/testuser".to_string());
-        Self {
-            env,
-            home: Some(PathBuf::from("/home/testuser")),
-            cwd: Some(PathBuf::from("/home/testuser")),
-        }
-    }
-
-    pub fn with_var(mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Self {
-        self.env.insert(key.as_ref().to_string(), value.as_ref().to_string());
-        self
-    }
-
-    pub fn with_cwd(mut self, cwd: impl AsRef<std::path::Path>) -> Self {
-        self.cwd = Some(PathBuf::from(cwd.as_ref()));
-        self
-    }
-}
-
-#[cfg(test)]
-impl Default for TestSystem {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-#[cfg(test)]
-impl EnvProvider for TestSystem {
-    fn var(&self, input: &str) -> Result<String, VarError> {
-        self.env.get(input).cloned().ok_or(VarError::NotPresent)
-    }
-}
-
-#[cfg(test)]
-impl HomeProvider for TestSystem {
-    fn home(&self) -> Option<PathBuf> {
-        self.home.as_ref().cloned()
-    }
-}
-
-#[cfg(test)]
-impl CwdProvider for TestSystem {
-    fn cwd(&self) -> Result<PathBuf, std::io::Error> {
-        self.cwd.as_ref().cloned().ok_or(std::io::Error::new(
-            std::io::ErrorKind::NotFound,
-            eyre::eyre!("not found"),
-        ))
-    }
-}