robit_agent/tool/
async_runner.rs

1//! Async tool execution: lets long-running tools return a pending placeholder
2//! and run their actual work in a background task.
3//!
4//! A tool that wants to run asynchronously (decided at runtime inside its
5//! `execute`) calls [`AsyncTaskRunner::submit`] with the real work as a future
6//! and returns a `ToolResult::pending(..)` placeholder. The runner spawns the
7//! work as an independent tokio task; when it finishes (or is cancelled) the
8//! result is sent back to the Agent via the `done` channel, which reinjects it
9//! into the conversation history and wakes the LLM.
10//!
11//! Cancellation is cooperative: each task owns a `CancellationToken`. When the
12//! Agent cancels a task (user `/cancel`, session expiry, or Agent shutdown),
13//! the spawned `select!` favors the cancellation branch and the work future is
14//! dropped (HTTP requests etc. are cancellation-safe).
15
16use std::future::Future;
17use std::panic::AssertUnwindSafe;
18use std::pin::Pin;
19
20use futures_util::future::BoxFuture;
21use futures_util::FutureExt;
22use tokio::sync::mpsc;
23use tokio_util::sync::CancellationToken;
24
25use crate::event::SessionId;
26use crate::tool::ToolResult;
27
28/// A future produced by an async tool representing its background work.
29///
30/// Resolves to the final `ToolResult` that will be reinjected into the
31/// conversation when the task completes.
32pub type AsyncTaskWork = BoxFuture<'static, ToolResult>;
33
34/// Message delivered to the Agent when a background task finishes (success,
35/// failure, or cancellation).
36pub struct AsyncTaskDone {
37    pub task_id: String,
38    pub tool_call_id: String,
39    pub session_id: SessionId,
40    pub tool_name: String,
41    pub result: ToolResult,
42    /// `true` if the task was cancelled (the `cancel` token fired) rather than
43    /// completing on its own. Lets the Agent record `Cancelled` rather than
44    /// `Failed` status.
45    pub cancelled: bool,
46}
47
48/// Handle used by tools to submit background work.
49///
50/// Cheap to clone (just an `mpsc::Sender`); one is placed in every
51/// `ToolContext`. The matching `Receiver` lives on the `Agent`, which `select!`s
52/// on it alongside user input.
53#[derive(Clone)]
54pub struct AsyncTaskRunner {
55    done_tx: mpsc::Sender<AsyncTaskDone>,
56}
57
58impl AsyncTaskRunner {
59    /// Create a runner wired to the Agent's `done` channel.
60    pub fn new(done_tx: mpsc::Sender<AsyncTaskDone>) -> Self {
61        Self { done_tx }
62    }
63
64    /// Spawn `work` as a background task and return its `task_id` immediately.
65    ///
66    /// The caller (the tool's `execute`) should return a `ToolResult::pending`
67    /// placeholder using this `task_id`. When `work` resolves - or `cancel` is
68    /// triggered - the result is sent to the Agent for reinjection.
69    ///
70    /// If the Agent has exited, `done_tx.send` fails silently and the result is
71    /// dropped (side effects like saved files already happened).
72    pub fn submit(
73        &self,
74        tool_call_id: String,
75        session_id: SessionId,
76        tool_name: String,
77        work: AsyncTaskWork,
78        cancel: CancellationToken,
79    ) -> String {
80        let task_id = uuid::Uuid::new_v4().to_string();
81        let done_tx = self.done_tx.clone();
82        let tid = task_id.clone();
83        let cancelled_msg = format!("异步任务 {} 已取消", tid);
84
85        tokio::spawn(async move {
86            // Favor cancellation: if cancelled while still running, return a
87            // cancelled result instead of the work's outcome. Dropping `work`
88            // cancels any in-flight HTTP requests it holds.
89            //
90            // `work` is wrapped in `AssertUnwindSafe(...).catch_unwind()` so a
91            // panic inside the tool's background future is converted into an
92            // error result and delivered to the Agent, instead of aborting the
93            // task silently. Without this, a panicking task would never reach
94            // `done_tx.send`, leaving a permanently-pending task (query_task
95            // stuck, no notification). (Requires panic=unwind, the default;
96            // panic=abort would still terminate the process.)
97            let (result, cancelled) = tokio::select! {
98                biased;
99                _ = cancel.cancelled() => (ToolResult::error(cancelled_msg), true),
100                r = AssertUnwindSafe(work).catch_unwind() => match r {
101                    Ok(tool_result) => (tool_result, false),
102                    Err(panic_payload) => {
103                        let msg = panic_payload_to_string(&panic_payload);
104                        tracing::error!(
105                            "[async] background task panicked: task_id={}, tool={}, session={}, panic={}",
106                            tid, tool_name, session_id, msg
107                        );
108                        (
109                            ToolResult::error(format!("异步任务执行时发生 panic: {}", msg)),
110                            false,
111                        )
112                    }
113                },
114            };
115
116            let done = AsyncTaskDone {
117                task_id: tid,
118                tool_call_id,
119                session_id,
120                tool_name,
121                result,
122                cancelled,
123            };
124            // Agent gone -> drop result. `send` is async but the channel is
125            // bounded; use a blocking-style await (spawned task can wait).
126            if let Err(send_err) = done_tx.send(done).await {
127                // Extract fields from the returned value for logging.
128                // Must convert error to string before destructuring (partial move).
129                let err_msg = send_err.to_string();
130                let lost = send_err.0;
131                tracing::error!(
132                    "[async] FAILED to deliver task result to Agent (Agent likely exited): \
133                     task_id={}, tool={}, session={}, cancelled={}, error={}. \
134                     The task result is permanently lost.",
135                    lost.task_id, lost.tool_name, lost.session_id, lost.cancelled, err_msg
136                );
137            }
138        });
139
140        task_id
141    }
142}
143
144/// Best-effort message extraction from a panic payload.
145///
146/// `panic!("literal")` yields `&'static str`; `panic!("{}", x)` yields `String`.
147/// Anything else falls back to a generic placeholder so the log line is still
148/// useful.
149fn panic_payload_to_string(payload: &Box<dyn std::any::Any + Send>) -> String {
150    if let Some(s) = payload.downcast_ref::<&str>() {
151        (*s).to_string()
152    } else if let Some(s) = payload.downcast_ref::<String>() {
153        s.clone()
154    } else {
155        "<non-string panic payload>".to_string()
156    }
157}
158
159// Keep the raw `Pin<Box<dyn Future>>` alias referenced for documentation;
160// `BoxFuture` is the same type with a Send bound.
161#[allow(dead_code)]
162type _AnyFuture = Pin<Box<dyn Future<Output = ToolResult> + Send>>;
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use std::sync::Arc;
168    use std::sync::Mutex as StdMutex;
169    use tokio::sync::Mutex;
170
171    async fn recv_n(rx: &Mutex<mpsc::Receiver<AsyncTaskDone>>, n: usize) -> Vec<AsyncTaskDone> {
172        let mut out = Vec::new();
173        for _ in 0..n {
174            match rx.lock().await.recv().await {
175                Some(d) => out.push(d),
176                None => break,
177            }
178        }
179        out
180    }
181
182    #[tokio::test]
183    async fn submit_returns_completed_result() {
184        let (tx, rx) = mpsc::channel(8);
185        let runner = AsyncTaskRunner::new(tx);
186        let counter: Arc<StdMutex<u32>> = Arc::new(StdMutex::new(0));
187        let c = counter.clone();
188        let task_id = runner.submit(
189            "tc1".into(),
190            "sess".into(),
191            "test_tool".into(),
192            Box::pin(async move {
193                *c.lock().unwrap() += 1;
194                ToolResult::success("done")
195            }),
196            CancellationToken::new(),
197        );
198        assert!(!task_id.is_empty());
199
200        let results = recv_n(&Mutex::new(rx), 1).await;
201        assert_eq!(results.len(), 1);
202        assert_eq!(results[0].task_id, task_id);
203        assert!(!results[0].result.is_error);
204        assert_eq!(results[0].result.content, "done");
205        assert_eq!(*counter.lock().unwrap(), 1);
206    }
207
208    #[tokio::test]
209    async fn cancel_returns_cancelled_result() {
210        let (tx, rx) = mpsc::channel(8);
211        let runner = AsyncTaskRunner::new(tx);
212        let cancel = CancellationToken::new();
213        let started = Arc::new(tokio::sync::Notify::new());
214        let started2 = started.clone();
215        let task_id = runner.submit(
216            "tc2".into(),
217            "sess".into(),
218            "test_tool".into(),
219            Box::pin(async move {
220                started2.notify_one();
221                // Simulate long work that should be preempted by cancellation.
222                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
223                ToolResult::success("should not reach")
224            }),
225            cancel.clone(),
226        );
227
228        started.notified().await;
229        cancel.cancel();
230
231        let results = recv_n(&Mutex::new(rx), 1).await;
232        assert_eq!(results.len(), 1);
233        assert_eq!(results[0].task_id, task_id);
234        assert!(results[0].result.is_error);
235        assert!(results[0].result.content.contains("已取消"));
236        assert!(results[0].cancelled);
237    }
238
239    #[tokio::test]
240    async fn panicking_work_delivers_error_result() {
241        // A panicking work future must NOT leave the task pending forever.
242        // The panic should be caught, logged, and delivered as an error result
243        // (cancelled=false) so the Agent can report failure to the LLM.
244        let (tx, rx) = mpsc::channel(8);
245        let runner = AsyncTaskRunner::new(tx);
246        let task_id = runner.submit(
247            "tc3".into(),
248            "sess".into(),
249            "test_tool".into(),
250            Box::pin(async {
251                panic!("boom from inside work");
252            }),
253            CancellationToken::new(),
254        );
255
256        let results = recv_n(&Mutex::new(rx), 1).await;
257        assert_eq!(results.len(), 1, "panic should still deliver a result");
258        assert_eq!(results[0].task_id, task_id);
259        assert!(results[0].result.is_error);
260        assert!(
261            results[0].result.content.contains("panic"),
262            "content was: {}",
263            results[0].result.content
264        );
265        assert!(
266            results[0].result.content.contains("boom from inside work"),
267            "content was: {}",
268            results[0].result.content
269        );
270        assert!(!results[0].cancelled);
271    }
272
273    #[test]
274    fn panic_payload_to_string_handles_common_payloads() {
275        let s: Box<dyn std::any::Any + Send> = Box::new("static literal");
276        assert_eq!(panic_payload_to_string(&s), "static literal");
277
278        let s: Box<dyn std::any::Any + Send> = Box::new("owned".to_string());
279        assert_eq!(panic_payload_to_string(&s), "owned");
280
281        let s: Box<dyn std::any::Any + Send> = Box::new(42i32);
282        assert_eq!(panic_payload_to_string(&s), "<non-string panic payload>");
283    }
284}
robit_agent/tool/async_runner.rs

robit_agent/tool/
async_runner.rs
