tiny_agent/tools/default_tools/

mod.rs

mod ask_followup_question;
mod bash;
mod check;
mod kill;
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_with_ctx(
        "bash",
        "执行一条 bash 命令。命令在 2 秒内跑完则直接返回 {status:\"completed\", exit_code, output};\
         超过 2 秒未结束则转入后台,返回 {id, status:\"running\", output:<已有输出>}——此时进程仍在跑,\
         请稍后用 check 工具凭 id 回来查看新增输出和最终状态(模型不会被自动通知完成),或用 kill 终止。",
        false,
        bash::bash,
    );
    registry.register_with_ctx(
        "check",
        "查看某个后台任务(bash 返回的 running 任务)自上次以来的新增输出与当前状态。\
         可传 wait_secs 让 check 先等待指定秒数再返回,避免过于频繁轮询。\
         返回 {status:\"running\"|\"completed\", exit_code, output:<新增>}。",
        true,
        check::check,
    );
    registry.register_with_ctx(
        "kill",
        "终止一个后台任务(按 bash 返回的 id)。",
        false,
        kill::kill,
    );
}

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 std::time::Duration;
    use tokio::time::Instant;
    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,
                crate::tools::ToolCtx::default(),
            )
            .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 bash_completes_fast_command_inline() {
        let mut registry = ToolRegistry::new();
        register_default_tools(&mut registry);
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());

        let out = registry
            .call(
                "bash",
                json!({ "command": "echo hi" }),
                sb,
                crate::tools::ToolCtx::default(),
            )
            .await
            .unwrap();
        let v = completed(out);
        assert_eq!(v["status"], "completed");
        assert_eq!(v["exit_code"], 0);
        assert!(v["output"].as_str().unwrap().contains("hi"));
    }

    #[tokio::test]
    async fn bash_yields_long_command_then_check_sees_output_and_kill_stops_it() {
        let mut registry = ToolRegistry::new();
        register_default_tools(&mut registry);
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        // 同一个 ToolCtx 的 TaskManager 是 Arc 共享的:bash 登记的任务,check/kill 能查到。
        let ctx = crate::tools::ToolCtx::default();

        // 先打印一行,再 sleep 超过 2s 让步窗口。
        let out = registry
            .call(
                "bash",
                json!({ "command": "echo first; sleep 5" }),
                sb.clone(),
                ctx.clone(),
            )
            .await
            .unwrap();
        let v = completed(out);
        assert_eq!(v["status"], "running");
        assert!(v["output"].as_str().unwrap().contains("first"));
        let id = v["id"].as_str().unwrap().to_string();

        // check:任务仍在跑(还没到 exit)。
        let cout = registry
            .call("check", json!({ "id": id }), sb.clone(), ctx.clone())
            .await
            .unwrap();
        assert_eq!(completed(cout)["status"], "running");

        // kill:能找到并终止。
        let kout = registry
            .call("kill", json!({ "id": id }), sb.clone(), ctx.clone())
            .await
            .unwrap();
        assert_eq!(completed(kout)["status"], "killed");

        // kill 之后再 check:任务已不在登记表。
        let err = registry.call("check", json!({ "id": id }), sb, ctx).await;
        assert!(err.is_err());
    }

    #[tokio::test]
    async fn check_wait_secs_sleeps_before_polling() {
        let mut registry = ToolRegistry::new();
        register_default_tools(&mut registry);
        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let ctx = crate::tools::ToolCtx::default();

        let out = registry
            .call(
                "bash",
                json!({ "command": "echo first; sleep 4" }),
                sb.clone(),
                ctx.clone(),
            )
            .await
            .unwrap();
        let id = completed(out)["id"].as_str().unwrap().to_string();

        let started = Instant::now();
        let cout = registry
            .call(
                "check",
                json!({ "id": id, "wait_secs": 1 }),
                sb.clone(),
                ctx.clone(),
            )
            .await
            .unwrap();
        assert!(started.elapsed() >= Duration::from_millis(900));
        assert_eq!(completed(cout)["status"], "running");

        let _ = registry.call("kill", json!({ "id": id }), sb, ctx).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());
    }
}