tiny_agent/tasks/
mod.rs

1//! 后台工具任务登记表。
2//!
3//! 声明为 asynchronous 的工具到达让步窗口仍未结束时,运行时会把仍在跑的 `tokio`
4//! task 连同进度缓冲登记进这里,以 id 为键;之后 `check(id)` / `kill(id)` 凭 id
5//! 把它找回来。这一层不关心具体工具类型,bash 只是其中一个使用者。
6
7use serde_json::Value;
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::{Arc, Mutex};
11use tokio::task::JoinHandle;
12
13#[derive(Default)]
14struct BufState {
15    /// 尚未被 `check(id)` 读走的输出(stdout/stderr 合流,贴近终端观感)。
16    data: String,
17    finished: bool,
18    result: Option<Value>,
19    error: Option<String>,
20}
21
22/// 一个后台任务的共享输出缓冲。回调往里 `append`,任务结束时 `finish`。
23#[derive(Clone, Default)]
24pub struct TaskBuffer {
25    inner: Arc<Mutex<BufState>>,
26}
27
28impl TaskBuffer {
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    /// 追加一块输出(在 `execute` 的实时回调里调用)。
34    pub fn append(&self, chunk: &str) {
35        self.inner.lock().unwrap().data.push_str(chunk);
36    }
37
38    /// 标记任务结束:`result` 为工具最终结构化输出,`error` 为工具执行错误。
39    pub fn finish(&self, result: Result<Value, String>) {
40        let mut s = self.inner.lock().unwrap();
41        s.finished = true;
42        match result {
43            Ok(value) => s.result = Some(value),
44            Err(error) => s.error = Some(error),
45        }
46    }
47
48    /// 取出"自上次以来"的新增输出,并释放已消费内容。
49    pub fn drain_new(&self) -> String {
50        let mut s = self.inner.lock().unwrap();
51        std::mem::take(&mut s.data)
52    }
53
54    /// 当前 (是否结束, 结构化结果, 错误)。不消费输出。
55    pub fn status_parts(&self) -> (bool, Option<Value>, Option<String>) {
56        let s = self.inner.lock().unwrap();
57        (s.finished, s.result.clone(), s.error.clone())
58    }
59}
60
61struct BgTask {
62    tool_name: String,
63    handle: JoinHandle<()>,
64    buffer: TaskBuffer,
65}
66
67/// `check(id)` 拿到的一帧快照。
68pub struct TaskSnapshot {
69    /// 自上次以来的新增输出。
70    pub new_output: String,
71    pub tool_name: String,
72    pub finished: bool,
73    pub result: Option<Value>,
74    pub error: Option<String>,
75}
76
77/// 后台任务登记表。`Clone` 即共享同一份(内部 `Arc`),可放进 `ToolCtx` 传给各工具。
78#[derive(Clone, Default)]
79pub struct TaskManager {
80    tasks: Arc<Mutex<HashMap<String, BgTask>>>,
81    counter: Arc<AtomicU64>,
82}
83
84impl TaskManager {
85    pub fn new() -> Self {
86        Self::default()
87    }
88
89    /// 登记一个仍在跑的后台工具任务,返回分配的 id(形如 `task_7`)。
90    pub fn register(
91        &self,
92        tool_name: impl Into<String>,
93        handle: JoinHandle<()>,
94        buffer: TaskBuffer,
95    ) -> String {
96        let n = self.counter.fetch_add(1, Ordering::SeqCst);
97        let id = format!("task_{n}");
98        self.tasks.lock().unwrap().insert(
99            id.clone(),
100            BgTask {
101                tool_name: tool_name.into(),
102                handle,
103                buffer,
104            },
105        );
106        id
107    }
108
109    /// 拿某任务自上次以来的新增输出 + 当前状态。未知 id 返回 `None`。
110    /// 已完成任务在交付这一帧快照后会被移除;之后同一 id 将不可再查询。
111    pub fn poll(&self, id: &str) -> Option<TaskSnapshot> {
112        let mut tasks = self.tasks.lock().unwrap();
113        let task = tasks.get(id)?;
114        let new_output = task.buffer.drain_new();
115        let (finished, result, error) = task.buffer.status_parts();
116        let snapshot = TaskSnapshot {
117            new_output,
118            tool_name: task.tool_name.clone(),
119            finished,
120            result,
121            error,
122        };
123        if snapshot.finished {
124            tasks.remove(id);
125        }
126        Some(snapshot)
127    }
128
129    /// 杀掉某任务:`abort` 那个 tokio task。返回是否确实存在该任务。
130    pub fn kill(&self, id: &str) -> bool {
131        match self.tasks.lock().unwrap().remove(id) {
132            Some(task) => {
133                task.handle.abort();
134                true
135            }
136            None => false,
137        }
138    }
139}
tiny_agent/tasks/mod.rs

tiny_agent/tasks/
mod.rs
