tiny-agent 0.2.0

use crate::messages::MessageKind;
use crate::sandbox::{OutStream, Sandbox};
use crate::tasks::TaskBuffer;
use crate::tools::{ToolCtx, ToolOutcome};
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;

/// 让步窗口:命令在此时间内跑完就直接返回完整结果,否则转后台并返回 `running` + id。
const YIELD_WINDOW: Duration = Duration::from_secs(2);

#[derive(Deserialize, JsonSchema)]
pub struct BashArgs {
    /// 要执行的 bash 命令
    pub command: String,
}

pub async fn bash(
    sandbox: Arc<dyn Sandbox>,
    args: BashArgs,
    ctx: ToolCtx,
) -> Result<ToolOutcome, String> {
    let buffer = TaskBuffer::new();

    // 实时回调:每行输出 ① 进 buffer(供 check 拿增量)② 推 messages(实时刷屏)。
    let on_output = {
        let buffer = buffer.clone();
        let messages = ctx.messages.clone();
        let session = ctx.session_id.clone();
        Arc::new(move |_stream: OutStream, chunk: &str| {
            let line = format!("{chunk}\n");
            buffer.append(&line);
            messages.send(&session, MessageKind::ToolOutput { chunk: line });
        })
    };

    // 先 spawn 出去脱离,这样超时走 timeout 分支时进程**不会**被 drop 掉,继续后台跑。
    let mut handle = tokio::spawn({
        let sandbox = sandbox.clone();
        let buffer = buffer.clone();
        let command = args.command.clone();
        async move {
            match sandbox.execute(&command, Some(on_output)).await {
                Ok(out) => buffer.finish(out.exit_code, None),
                Err(e) => buffer.finish(None, Some(e.to_string())),
            }
        }
    });

    tokio::select! {
        _ = &mut handle => {
            // 2s 内跑完:返回完整结果。
            let output = buffer.drain_new();
            let (_finished, exit_code, error) = buffer.status_parts();
            if let Some(err) = error {
                return Err(err);
            }
            Ok(json!({
                "status": "completed",
                "exit_code": exit_code,
                "output": output,
            })
            .into())
        }
        _ = tokio::time::sleep(YIELD_WINDOW) => {
            // 没跑完:进程在后台继续(handle 没被 drop,move 进登记表),返回这 2s 的增量 + id。
            let partial = buffer.drain_new();
            let id = ctx.tasks.register(handle, buffer);
            Ok(json!({
                "id": id,
                "status": "running",
                "output": partial,
                "note": "命令仍在后台运行;用 check 工具按 id 回来查看新增输出与最终状态,用 kill 终止。",
            })
            .into())
        }
    }
}