tiny_agent/tasks/
mod.rs

1//! 后台任务登记表。
2//!
3//! `bash` 的"2s 让步窗口"到点还没跑完时,会把那个还在后台跑的 `tokio` task 连同它的
4//! 输出缓冲登记进这里,以 id 为键;之后 `check(id)` / `kill(id)` 凭 id 把它找回来。
5//!
6//! 这一层**完全在工具层之上**,只用到 `sandbox.execute` 这个原语,sandbox 本身对
7//! "任务"一无所知。
8
9use std::collections::HashMap;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::{Arc, Mutex};
12use tokio::task::JoinHandle;
13
14#[derive(Default)]
15struct BufState {
16    /// 至今为止的全部输出(stdout/stderr 合流,贴近终端观感)。
17    data: String,
18    /// 模型上次读到的位置 —— `drain_new` 只返回这之后的内容。
19    cursor: usize,
20    finished: bool,
21    exit_code: Option<i32>,
22    /// execute 自身报错(如 sandbox 起进程失败),非命令退出码。
23    error: Option<String>,
24}
25
26/// 一个后台任务的共享输出缓冲。回调往里 `append`,任务结束时 `finish`。
27#[derive(Clone, Default)]
28pub struct TaskBuffer {
29    inner: Arc<Mutex<BufState>>,
30}
31
32impl TaskBuffer {
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// 追加一块输出(在 `execute` 的实时回调里调用)。
38    pub fn append(&self, chunk: &str) {
39        self.inner.lock().unwrap().data.push_str(chunk);
40    }
41
42    /// 标记任务结束:`exit_code` 为命令退出码,`error` 为 execute 自身的错误。
43    pub fn finish(&self, exit_code: Option<i32>, error: Option<String>) {
44        let mut s = self.inner.lock().unwrap();
45        s.finished = true;
46        s.exit_code = exit_code;
47        s.error = error;
48    }
49
50    /// 取出"自上次以来"的新增输出,游标推到末尾。
51    pub fn drain_new(&self) -> String {
52        let mut s = self.inner.lock().unwrap();
53        let new = s.data[s.cursor..].to_string();
54        s.cursor = s.data.len();
55        new
56    }
57
58    /// 当前 (是否结束, 退出码, execute 错误)。不消费输出。
59    pub fn status_parts(&self) -> (bool, Option<i32>, Option<String>) {
60        let s = self.inner.lock().unwrap();
61        (s.finished, s.exit_code, s.error.clone())
62    }
63}
64
65struct BgTask {
66    handle: JoinHandle<()>,
67    buffer: TaskBuffer,
68}
69
70/// `check(id)` 拿到的一帧快照。
71pub struct TaskSnapshot {
72    /// 自上次以来的新增输出。
73    pub new_output: String,
74    pub finished: bool,
75    pub exit_code: Option<i32>,
76    pub error: Option<String>,
77}
78
79/// 后台任务登记表。`Clone` 即共享同一份(内部 `Arc`),可放进 `ToolCtx` 传给各工具。
80#[derive(Clone, Default)]
81pub struct TaskManager {
82    tasks: Arc<Mutex<HashMap<String, BgTask>>>,
83    counter: Arc<AtomicU64>,
84}
85
86impl TaskManager {
87    pub fn new() -> Self {
88        Self::default()
89    }
90
91    /// 登记一个仍在跑的后台任务,返回分配的 id(形如 `task_7`)。
92    pub fn register(&self, handle: JoinHandle<()>, buffer: TaskBuffer) -> String {
93        let n = self.counter.fetch_add(1, Ordering::SeqCst);
94        let id = format!("task_{n}");
95        self.tasks
96            .lock()
97            .unwrap()
98            .insert(id.clone(), BgTask { handle, buffer });
99        id
100    }
101
102    /// 拿某任务自上次以来的新增输出 + 当前状态。未知 id 返回 `None`。
103    pub fn poll(&self, id: &str) -> Option<TaskSnapshot> {
104        let tasks = self.tasks.lock().unwrap();
105        let task = tasks.get(id)?;
106        let new_output = task.buffer.drain_new();
107        let (finished, exit_code, error) = task.buffer.status_parts();
108        Some(TaskSnapshot {
109            new_output,
110            finished,
111            exit_code,
112            error,
113        })
114    }
115
116    /// 杀掉某任务:`abort` 那个 tokio task → `execute` future 被 drop → `kill_on_drop`
117    /// 杀掉子进程。返回是否确实存在该任务。
118    pub fn kill(&self, id: &str) -> bool {
119        match self.tasks.lock().unwrap().remove(id) {
120            Some(task) => {
121                task.handle.abort();
122                true
123            }
124            None => false,
125        }
126    }
127}
tiny_agent/tasks/mod.rs

tiny_agent/tasks/
mod.rs
