tiny-agent 0.3.0

use super::{ExecOutput, OutStream, OutputCallback, Sandbox, SandboxError};
use async_trait::async_trait;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
use tokio::process::Command;

/// 直接在宿主机上执行命令、读写真实文件系统的 sandbox。
///
/// 这是默认实现,行为等价于"不沙箱",不提供任何隔离。
#[derive(Default)]
pub struct HostSandbox {
    /// 命令执行的工作目录;`None` = 继承当前进程的 cwd(原行为)。
    root: Option<PathBuf>,
}

impl HostSandbox {
    pub fn new() -> Self {
        Self { root: None }
    }

    /// 把命令限定在 `root` 目录下执行(以它为工作目录)。
    ///
    /// 注意:这是**工作目录**级别的限定——相对路径会落在 `root` 内,但**不阻止**绝对路径
    /// (如 `/etc/...`)或 `cd ..` 逃逸,因此**不是安全沙箱**,只是"默认在这个目录里干活"。
    pub fn with_root(root: impl Into<PathBuf>) -> Self {
        Self {
            root: Some(root.into()),
        }
    }
}

#[async_trait]
impl Sandbox for HostSandbox {
    async fn open(self: Arc<Self>, _session_id: &str) -> Result<Arc<dyn Sandbox>, SandboxError> {
        Ok(self)
    }

    async fn execute(
        &self,
        command: &str,
        on_output: Option<OutputCallback>,
    ) -> Result<ExecOutput, SandboxError> {
        let mut cmd = Command::new("bash");
        cmd.arg("-c")
            .arg(command)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            // 被 drop(取消/超时杀任务)时连带杀掉直接子进程。
            .kill_on_drop(true);
        // 限定了 root 就以它为工作目录,命令里的相对路径都落在目录内。
        if let Some(root) = &self.root {
            cmd.current_dir(root);
        }
        let mut child = cmd.spawn()?;

        let stdout = child.stdout.take().expect("stdout piped");
        let stderr = child.stderr.take().expect("stderr piped");

        // 两个流并发逐行读:每行既回调(实时)又累积(最终返回)。
        let (out_acc, err_acc) = tokio::join!(
            pump(stdout, OutStream::Stdout, on_output.clone()),
            pump(stderr, OutStream::Stderr, on_output.clone()),
        );

        let status = child.wait().await?;
        Ok(ExecOutput {
            stdout: out_acc,
            stderr: err_acc,
            exit_code: status.code(),
        })
    }
}

/// 逐行读一个流:每行调一次回调(若有)并累积,EOF 后返回累积内容。
async fn pump<R: AsyncRead + Unpin>(
    reader: R,
    stream: OutStream,
    on_output: Option<OutputCallback>,
) -> String {
    let mut lines = BufReader::new(reader).lines();
    let mut acc = String::new();
    while let Ok(Some(line)) = lines.next_line().await {
        if let Some(cb) = &on_output {
            cb(stream, &line);
        }
        acc.push_str(&line);
        acc.push('\n');
    }
    acc
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn execute_captures_stdout_and_exit_code() {
        let sb = HostSandbox::new();
        let out = sb.execute("echo hello", None).await.unwrap();
        assert_eq!(out.stdout, "hello\n");
        assert!(out.is_success());
    }

    #[tokio::test]
    async fn execute_reports_nonzero_exit() {
        let sb = HostSandbox::new();
        let out = sb.execute("exit 3", None).await.unwrap();
        assert_eq!(out.exit_code, Some(3));
        assert!(!out.is_success());
    }

    #[tokio::test]
    async fn with_root_runs_in_that_directory() {
        let dir = std::env::temp_dir();
        let sb = HostSandbox::with_root(&dir);
        let out = sb.execute("pwd", None).await.unwrap();
        // macOS 上 /var/folders 通过软链到 /private/var,两边都 canonicalize 再比。
        let got = std::fs::canonicalize(out.stdout.trim()).unwrap();
        let want = std::fs::canonicalize(&dir).unwrap();
        assert_eq!(got, want);
        // 相对路径写文件会落在 root 内。
        let probe = format!("host_root_probe_{}.txt", std::process::id());
        sb.execute(&format!("echo hi > {probe}"), None).await.unwrap();
        let landed = dir.join(&probe);
        assert!(landed.exists(), "相对路径应落在 root: {}", landed.display());
        let _ = std::fs::remove_file(landed);
    }
}