tiny_agent/tools/default_tools/

mod.rs

mod ask_followup_question;
mod bash;
mod read_file;
mod read_skill;
mod spawn_agent;
mod write_file;

pub use spawn_agent::register_spawn_agent;

use crate::skills::SkillManager;
use crate::tools::ToolRegistry;
use std::sync::Arc;

/// 把内置工具一次性注册进 registry。
///
/// 工具函数签名都是 `fn(Arc<dyn Sandbox>, Args)`,sandbox 在 `ToolRegistry::call`
/// 时按会话注入,所以这里直接注册函数本身,无需包装闭包。
pub fn register_default_tools(registry: &mut ToolRegistry) {
    registry.register(
        "askFollowupQuestion",
        "向用户提出一个追问,并让 agent 暂停到 WaitingForUser 状态",
        false,
        ask_followup_question::ask_followup_question,
    );
    registry.register(
        "read_file",
        "读取指定文件的内容",
        true,
        read_file::read_file,
    );
    registry.register(
        "write_file",
        "写入或局部编辑文件,支持 write/replace/insert_after/insert_before/append",
        false,
        write_file::write_file,
    );
    registry.register("bash", "执行一条 bash 命令并返回输出", false, bash::bash);
}

pub fn register_skill_tools(registry: &mut ToolRegistry, skills: Arc<SkillManager>) {
    read_skill::register_read_skill(registry, skills);
}

/// 用单引号包裹字符串,使其能安全地作为单个 shell 参数传递。
pub fn sh_quote(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('\'');
    for ch in s.chars() {
        if ch == '\'' {
            out.push_str("'\\''");
        } else {
            out.push(ch);
        }
    }
    out.push('\'');
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sandbox::{HostSandbox, Sandbox};
    use crate::tools::ToolOutcome;
    use read_file::ReadFileArgs;
    use serde_json::{Value, json};
    use std::sync::Arc;
    use write_file::{WriteFileArgs, WriteFileMode};

    /// 取出普通完成态的 `Value`,测试里断言工具输出时用。
    fn completed(outcome: ToolOutcome) -> Value {
        match outcome {
            ToolOutcome::Completed(value) => value,
            other => panic!("expected Completed, got {other:?}"),
        }
    }

    #[test]
    fn sh_quote_wraps_plain_string() {
        assert_eq!(sh_quote("foo.txt"), "'foo.txt'");
    }

    #[test]
    fn sh_quote_escapes_single_quote() {
        assert_eq!(sh_quote("a'b"), "'a'\\''b'");
    }

    #[tokio::test]
    async fn write_then_read_round_trips_through_sandbox() {
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let dir = std::env::temp_dir();
        let path = dir.join(format!("tiny_agent_test_{}.txt", uuid::Uuid::new_v4()));
        let path_str = path.to_string_lossy().to_string();
        let content = "line one\nline 'two'\n";

        let write_args = WriteFileArgs {
            path: path_str.clone(),
            mode: Some(WriteFileMode::Write),
            content: Some(content.to_string()),
            old_string: None,
            new_string: None,
            replace_all: None,
            expected_replacements: None,
        };
        write_file::write_file(sb.clone(), write_args)
            .await
            .unwrap();

        let read_args = ReadFileArgs {
            path: path_str.clone(),
        };
        let out = read_file::read_file(sb.clone(), read_args).await.unwrap();
        assert_eq!(completed(out)["content"], content);

        let _ = tokio::fs::remove_file(&path).await;
    }

    #[tokio::test]
    async fn ask_followup_question_returns_needs_user_interaction() {
        let mut registry = ToolRegistry::new();
        register_default_tools(&mut registry);
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());

        let outcome = registry
            .call(
                "askFollowupQuestion",
                json!({
                    "question": "Which file should I edit?",
                    "options": ["src/main.rs", "Cargo.toml"]
                }),
                sb,
            )
            .await
            .unwrap();

        assert!(matches!(
            outcome,
            ToolOutcome::NeedsUserInteraction(interaction)
                if interaction.kind == "followup_question"
                    && interaction.payload["question"] == "Which file should I edit?"
                    && interaction.payload["options"] == json!(["src/main.rs", "Cargo.toml"])
        ));
    }

    #[tokio::test]
    async fn write_file_replaces_unique_old_string() {
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let dir = std::env::temp_dir();
        let path = dir.join(format!("tiny_agent_test_{}.txt", uuid::Uuid::new_v4()));
        let path_str = path.to_string_lossy().to_string();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Write),
                content: Some("alpha\nbeta\ngamma\n".to_string()),
                old_string: None,
                new_string: None,
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Replace),
                content: None,
                old_string: Some("beta".to_string()),
                new_string: Some("BETA".to_string()),
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        let out = read_file::read_file(
            sb.clone(),
            ReadFileArgs {
                path: path_str.clone(),
            },
        )
        .await
        .unwrap();
        assert_eq!(completed(out)["content"], "alpha\nBETA\ngamma\n");

        let _ = tokio::fs::remove_file(&path).await;
    }

    #[tokio::test]
    async fn write_file_rejects_ambiguous_replace_by_default() {
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let dir = std::env::temp_dir();
        let path = dir.join(format!("tiny_agent_test_{}.txt", uuid::Uuid::new_v4()));
        let path_str = path.to_string_lossy().to_string();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Write),
                content: Some("same\nsame\n".to_string()),
                old_string: None,
                new_string: None,
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        let result = write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Replace),
                content: None,
                old_string: Some("same".to_string()),
                new_string: Some("changed".to_string()),
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await;

        assert!(result.is_err());
        let _ = tokio::fs::remove_file(&path).await;
    }

    #[tokio::test]
    async fn write_file_can_insert_and_append() {
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let dir = std::env::temp_dir();
        let path = dir.join(format!("tiny_agent_test_{}.txt", uuid::Uuid::new_v4()));
        let path_str = path.to_string_lossy().to_string();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Write),
                content: Some("fn main() {\n}\n".to_string()),
                old_string: None,
                new_string: None,
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::InsertAfter),
                content: Some("\n    println!(\"hi\");".to_string()),
                old_string: Some("fn main() {".to_string()),
                new_string: None,
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        write_file::write_file(
            sb.clone(),
            WriteFileArgs {
                path: path_str.clone(),
                mode: Some(WriteFileMode::Append),
                content: Some("// end\n".to_string()),
                old_string: None,
                new_string: None,
                replace_all: None,
                expected_replacements: None,
            },
        )
        .await
        .unwrap();

        let out = read_file::read_file(
            sb.clone(),
            ReadFileArgs {
                path: path_str.clone(),
            },
        )
        .await
        .unwrap();
        assert_eq!(
            completed(out)["content"],
            "fn main() {\n    println!(\"hi\");\n}\n// end\n"
        );

        let _ = tokio::fs::remove_file(&path).await;
    }

    #[tokio::test]
    async fn read_missing_file_returns_err() {
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let args = ReadFileArgs {
            path: "/no/such/path/really_missing_12345".to_string(),
        };
        assert!(read_file::read_file(sb, args).await.is_err());
    }
}