tiny-agent 0.2.0

//! 后台任务登记表。
//!
//! `bash` 的"2s 让步窗口"到点还没跑完时,会把那个还在后台跑的 `tokio` task 连同它的
//! 输出缓冲登记进这里,以 id 为键;之后 `check(id)` / `kill(id)` 凭 id 把它找回来。
//!
//! 这一层**完全在工具层之上**,只用到 `sandbox.execute` 这个原语,sandbox 本身对
//! "任务"一无所知。

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tokio::task::JoinHandle;

#[derive(Default)]
struct BufState {
    /// 至今为止的全部输出(stdout/stderr 合流,贴近终端观感)。
    data: String,
    /// 模型上次读到的位置 —— `drain_new` 只返回这之后的内容。
    cursor: usize,
    finished: bool,
    exit_code: Option<i32>,
    /// execute 自身报错(如 sandbox 起进程失败),非命令退出码。
    error: Option<String>,
}

/// 一个后台任务的共享输出缓冲。回调往里 `append`,任务结束时 `finish`。
#[derive(Clone, Default)]
pub struct TaskBuffer {
    inner: Arc<Mutex<BufState>>,
}

impl TaskBuffer {
    pub fn new() -> Self {
        Self::default()
    }

    /// 追加一块输出(在 `execute` 的实时回调里调用)。
    pub fn append(&self, chunk: &str) {
        self.inner.lock().unwrap().data.push_str(chunk);
    }

    /// 标记任务结束:`exit_code` 为命令退出码,`error` 为 execute 自身的错误。
    pub fn finish(&self, exit_code: Option<i32>, error: Option<String>) {
        let mut s = self.inner.lock().unwrap();
        s.finished = true;
        s.exit_code = exit_code;
        s.error = error;
    }

    /// 取出"自上次以来"的新增输出,游标推到末尾。
    pub fn drain_new(&self) -> String {
        let mut s = self.inner.lock().unwrap();
        let new = s.data[s.cursor..].to_string();
        s.cursor = s.data.len();
        new
    }

    /// 当前 (是否结束, 退出码, execute 错误)。不消费输出。
    pub fn status_parts(&self) -> (bool, Option<i32>, Option<String>) {
        let s = self.inner.lock().unwrap();
        (s.finished, s.exit_code, s.error.clone())
    }
}

struct BgTask {
    handle: JoinHandle<()>,
    buffer: TaskBuffer,
}

/// `check(id)` 拿到的一帧快照。
pub struct TaskSnapshot {
    /// 自上次以来的新增输出。
    pub new_output: String,
    pub finished: bool,
    pub exit_code: Option<i32>,
    pub error: Option<String>,
}

/// 后台任务登记表。`Clone` 即共享同一份(内部 `Arc`),可放进 `ToolCtx` 传给各工具。
#[derive(Clone, Default)]
pub struct TaskManager {
    tasks: Arc<Mutex<HashMap<String, BgTask>>>,
    counter: Arc<AtomicU64>,
}

impl TaskManager {
    pub fn new() -> Self {
        Self::default()
    }

    /// 登记一个仍在跑的后台任务,返回分配的 id(形如 `task_7`)。
    pub fn register(&self, handle: JoinHandle<()>, buffer: TaskBuffer) -> String {
        let n = self.counter.fetch_add(1, Ordering::SeqCst);
        let id = format!("task_{n}");
        self.tasks
            .lock()
            .unwrap()
            .insert(id.clone(), BgTask { handle, buffer });
        id
    }

    /// 拿某任务自上次以来的新增输出 + 当前状态。未知 id 返回 `None`。
    pub fn poll(&self, id: &str) -> Option<TaskSnapshot> {
        let tasks = self.tasks.lock().unwrap();
        let task = tasks.get(id)?;
        let new_output = task.buffer.drain_new();
        let (finished, exit_code, error) = task.buffer.status_parts();
        Some(TaskSnapshot {
            new_output,
            finished,
            exit_code,
            error,
        })
    }

    /// 杀掉某任务:`abort` 那个 tokio task → `execute` future 被 drop → `kill_on_drop`
    /// 杀掉子进程。返回是否确实存在该任务。
    pub fn kill(&self, id: &str) -> bool {
        match self.tasks.lock().unwrap().remove(id) {
            Some(task) => {
                task.handle.abort();
                true
            }
            None => false,
        }
    }
}