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::{AsyncToolConfig, ToolRegistry, ToolSpec};
use std::sync::Arc;
use std::time::Duration;

/// 让步窗口:声明为 asynchronous 的内置工具在此时间内跑完就直接返回完整结果,
/// 否则转后台并返回 `running` + id。
const DEFAULT_YIELD_AFTER: Duration = Duration::from_secs(2);

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

    #[derive(Deserialize, JsonSchema)]
    struct EmptyArgs {}

    #[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 共享的:异步工具登记的任务,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 asynchronous_tool_policy_is_not_bash_specific() {
        let mut registry = ToolRegistry::new();
        register_default_tools(&mut registry);
        registry.register_with_ctx_tool(
            ToolSpec::new("slow_report", "generic slow report", false)
                .asynchronous(AsyncToolConfig::new(Duration::from_millis(50))),
            |_sandbox, _args: EmptyArgs, ctx| async move {
                ctx.progress.emit_text("started\n");
                tokio::time::sleep(Duration::from_millis(150)).await;
                ctx.progress.emit_text("finished\n");
                Ok(json!({ "answer": 42 }).into())
            },
        );

        let sb: Arc<dyn Sandbox> = Arc::new(HostSandbox::new());
        let ctx = crate::tools::ToolCtx::default();
        let out = registry
            .call("slow_report", json!({}), sb.clone(), ctx.clone())
            .await
            .unwrap();
        let running = completed(out);
        assert_eq!(running["status"], "running");
        assert_eq!(running["tool_name"], "slow_report");
        assert!(running["output"].as_str().unwrap().contains("started"));
        let id = running["id"].as_str().unwrap().to_string();

        tokio::time::sleep(Duration::from_millis(200)).await;
        let checked = registry
            .call(
                "check",
                json!({ "id": id.clone() }),
                sb.clone(),
                ctx.clone(),
            )
            .await
            .unwrap();
        let completed = completed(checked);
        assert_eq!(completed["status"], "completed");
        assert_eq!(completed["answer"], 42);
        assert!(completed["output"].as_str().unwrap().contains("finished"));

        let checked_again = registry.call("check", json!({ "id": id }), sb, ctx).await;
        assert!(checked_again.is_err());
    }

    #[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());
    }
}