tiny-agent 0.2.0

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

/// 直接在宿主机上执行命令、读写真实文件系统的 sandbox。
///
/// 这是默认实现,行为等价于"不沙箱",不提供任何隔离。
#[derive(Default)]
pub struct HostSandbox;

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

#[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 child = Command::new("bash")
            .arg("-c")
            .arg(command)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            // 被 drop(取消/超时杀任务)时连带杀掉直接子进程。
            .kill_on_drop(true)
            .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());
    }
}