tiny-agent 0.3.0

//! 后台工具任务登记表。
//!
//! 声明为 asynchronous 的工具到达让步窗口仍未结束时,运行时会把仍在跑的 `tokio`
//! task 连同进度缓冲登记进这里,以 id 为键;之后 `check(id)` / `kill(id)` 凭 id
//! 把它找回来。这一层不关心具体工具类型,bash 只是其中一个使用者。

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

#[derive(Default)]
struct BufState {
    /// 尚未被 `check(id)` 读走的输出(stdout/stderr 合流,贴近终端观感)。
    data: String,
    finished: bool,
    result: Option<Value>,
    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);
    }

    /// 标记任务结束:`result` 为工具最终结构化输出,`error` 为工具执行错误。
    pub fn finish(&self, result: Result<Value, String>) {
        let mut s = self.inner.lock().unwrap();
        s.finished = true;
        match result {
            Ok(value) => s.result = Some(value),
            Err(error) => s.error = Some(error),
        }
    }

    /// 取出"自上次以来"的新增输出,并释放已消费内容。
    pub fn drain_new(&self) -> String {
        let mut s = self.inner.lock().unwrap();
        std::mem::take(&mut s.data)
    }

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

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

/// `check(id)` 拿到的一帧快照。
pub struct TaskSnapshot {
    /// 自上次以来的新增输出。
    pub new_output: String,
    pub tool_name: String,
    pub finished: bool,
    pub result: Option<Value>,
    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,
        tool_name: impl Into<String>,
        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 {
                tool_name: tool_name.into(),
                handle,
                buffer,
            },
        );
        id
    }

    /// 拿某任务自上次以来的新增输出 + 当前状态。未知 id 返回 `None`。
    /// 已完成任务在交付这一帧快照后会被移除;之后同一 id 将不可再查询。
    pub fn poll(&self, id: &str) -> Option<TaskSnapshot> {
        let mut tasks = self.tasks.lock().unwrap();
        let task = tasks.get(id)?;
        let new_output = task.buffer.drain_new();
        let (finished, result, error) = task.buffer.status_parts();
        let snapshot = TaskSnapshot {
            new_output,
            tool_name: task.tool_name.clone(),
            finished,
            result,
            error,
        };
        if snapshot.finished {
            tasks.remove(id);
        }
        Some(snapshot)
    }

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