robit_agent/tool/

async_runner.rs

//! Async tool execution: lets long-running tools return a pending placeholder
//! and run their actual work in a background task.
//!
//! A tool that wants to run asynchronously (decided at runtime inside its
//! `execute`) calls [`AsyncTaskRunner::submit`] with the real work as a future
//! and returns a `ToolResult::pending(..)` placeholder. The runner spawns the
//! work as an independent tokio task; when it finishes (or is cancelled) the
//! result is sent back to the Agent via the `done` channel, which reinjects it
//! into the conversation history and wakes the LLM.
//!
//! Cancellation is cooperative: each task owns a `CancellationToken`. When the
//! Agent cancels a task (user `/cancel`, session expiry, or Agent shutdown),
//! the spawned `select!` favors the cancellation branch and the work future is
//! dropped (HTTP requests etc. are cancellation-safe).

use std::future::Future;
use std::pin::Pin;

use futures_util::future::BoxFuture;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

use crate::event::SessionId;
use crate::tool::ToolResult;

/// A future produced by an async tool representing its background work.
///
/// Resolves to the final `ToolResult` that will be reinjected into the
/// conversation when the task completes.
pub type AsyncTaskWork = BoxFuture<'static, ToolResult>;

/// Message delivered to the Agent when a background task finishes (success,
/// failure, or cancellation).
pub struct AsyncTaskDone {
    pub task_id: String,
    pub tool_call_id: String,
    pub session_id: SessionId,
    pub tool_name: String,
    pub result: ToolResult,
    /// `true` if the task was cancelled (the `cancel` token fired) rather than
    /// completing on its own. Lets the Agent record `Cancelled` rather than
    /// `Failed` status.
    pub cancelled: bool,
}

/// Handle used by tools to submit background work.
///
/// Cheap to clone (just an `mpsc::Sender`); one is placed in every
/// `ToolContext`. The matching `Receiver` lives on the `Agent`, which `select!`s
/// on it alongside user input.
#[derive(Clone)]
pub struct AsyncTaskRunner {
    done_tx: mpsc::Sender<AsyncTaskDone>,
}

impl AsyncTaskRunner {
    /// Create a runner wired to the Agent's `done` channel.
    pub fn new(done_tx: mpsc::Sender<AsyncTaskDone>) -> Self {
        Self { done_tx }
    }

    /// Spawn `work` as a background task and return its `task_id` immediately.
    ///
    /// The caller (the tool's `execute`) should return a `ToolResult::pending`
    /// placeholder using this `task_id`. When `work` resolves - or `cancel` is
    /// triggered - the result is sent to the Agent for reinjection.
    ///
    /// If the Agent has exited, `done_tx.send` fails silently and the result is
    /// dropped (side effects like saved files already happened).
    pub fn submit(
        &self,
        tool_call_id: String,
        session_id: SessionId,
        tool_name: String,
        work: AsyncTaskWork,
        cancel: CancellationToken,
    ) -> String {
        let task_id = uuid::Uuid::new_v4().to_string();
        let done_tx = self.done_tx.clone();
        let tid = task_id.clone();
        let cancelled_msg = format!("异步任务 {} 已取消", tid);

        tokio::spawn(async move {
            // Favor cancellation: if cancelled while still running, return a
            // cancelled result instead of the work's outcome. Dropping `work`
            // cancels any in-flight HTTP requests it holds.
            let (result, cancelled) = tokio::select! {
                biased;
                _ = cancel.cancelled() => (ToolResult::error(cancelled_msg), true),
                r = work => (r, false),
            };

            let done = AsyncTaskDone {
                task_id: tid,
                tool_call_id,
                session_id,
                tool_name,
                result,
                cancelled,
            };
            // Agent gone -> drop result. `send` is async but the channel is
            // bounded; use a blocking-style await (spawned task can wait).
            let _ = done_tx.send(done).await;
        });

        task_id
    }
}

// Keep the raw `Pin<Box<dyn Future>>` alias referenced for documentation;
// `BoxFuture` is the same type with a Send bound.
#[allow(dead_code)]
type _AnyFuture = Pin<Box<dyn Future<Output = ToolResult> + Send>>;

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::sync::Mutex as StdMutex;
    use tokio::sync::Mutex;

    async fn recv_n(rx: &Mutex<mpsc::Receiver<AsyncTaskDone>>, n: usize) -> Vec<AsyncTaskDone> {
        let mut out = Vec::new();
        for _ in 0..n {
            match rx.lock().await.recv().await {
                Some(d) => out.push(d),
                None => break,
            }
        }
        out
    }

    #[tokio::test]
    async fn submit_returns_completed_result() {
        let (tx, rx) = mpsc::channel(8);
        let runner = AsyncTaskRunner::new(tx);
        let counter: Arc<StdMutex<u32>> = Arc::new(StdMutex::new(0));
        let c = counter.clone();
        let task_id = runner.submit(
            "tc1".into(),
            "sess".into(),
            "test_tool".into(),
            Box::pin(async move {
                *c.lock().unwrap() += 1;
                ToolResult::success("done")
            }),
            CancellationToken::new(),
        );
        assert!(!task_id.is_empty());

        let results = recv_n(&Mutex::new(rx), 1).await;
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].task_id, task_id);
        assert!(!results[0].result.is_error);
        assert_eq!(results[0].result.content, "done");
        assert_eq!(*counter.lock().unwrap(), 1);
    }

    #[tokio::test]
    async fn cancel_returns_cancelled_result() {
        let (tx, rx) = mpsc::channel(8);
        let runner = AsyncTaskRunner::new(tx);
        let cancel = CancellationToken::new();
        let started = Arc::new(tokio::sync::Notify::new());
        let started2 = started.clone();
        let task_id = runner.submit(
            "tc2".into(),
            "sess".into(),
            "test_tool".into(),
            Box::pin(async move {
                started2.notify_one();
                // Simulate long work that should be preempted by cancellation.
                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
                ToolResult::success("should not reach")
            }),
            cancel.clone(),
        );

        started.notified().await;
        cancel.cancel();

        let results = recv_n(&Mutex::new(rx), 1).await;
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].task_id, task_id);
        assert!(results[0].result.is_error);
        assert!(results[0].result.content.contains("已取消"));
        assert!(results[0].cancelled);
    }
}