agentic_tools_core/
context.rs

1//! Tool execution context.
2
3use crate::ToolError;
4use std::future::Future;
5use tokio_util::sync::CancellationToken;
6use tokio_util::sync::WaitForCancellationFutureOwned;
7
8/// Context passed to tool executions.
9///
10/// MCP-backed tool calls receive a request-scoped cancellation token through this
11/// context. Non-MCP callers that construct [`ToolContext::default`] or
12/// [`ToolContext::new`] receive a never-cancelled context so existing direct,
13/// native, and NAPI entrypoints remain source-compatible.
14///
15/// Tool authors should prefer these helpers instead of wiring cancellation by hand:
16/// - [`ToolContext::run_cancellable`] for a single async operation that should abort
17///   promptly when the request is cancelled.
18/// - [`ToolContext::cancelled`] when cancellation must participate in a larger
19///   `tokio::select!`.
20/// - [`ToolContext::is_cancelled`] for quick boundary checks before starting more work.
21/// - [`ToolContext::cancellation_token`] when an owned token clone must cross `.await`
22///   boundaries inside a `BoxFuture<'static>` implementation.
23///
24/// Cancellation maps to [`ToolError::Cancelled`]. Returning that error means the tool
25/// stopped because the caller cancelled the request, not because the tool failed.
26///
27/// For subprocess-managing tools, request cancellation should trigger explicit cleanup
28/// before returning. Dropping a future remains a backstop, not the primary cooperative
29/// cleanup path.
30#[derive(Clone, Debug)]
31pub struct ToolContext {
32    cancel: CancellationToken,
33}
34
35impl Default for ToolContext {
36    fn default() -> Self {
37        Self::with_cancel(CancellationToken::new())
38    }
39}
40
41impl ToolContext {
42    /// Create a new default context.
43    pub fn new() -> Self {
44        Self::default()
45    }
46
47    /// Create a context backed by the supplied cancellation token.
48    pub fn with_cancel(cancel: CancellationToken) -> Self {
49        Self { cancel }
50    }
51
52    /// Clone the request cancellation token for use across `.await` boundaries.
53    pub fn cancellation_token(&self) -> CancellationToken {
54        self.cancel.clone()
55    }
56
57    /// Return an owned future that resolves when the request is cancelled.
58    pub fn cancelled(&self) -> WaitForCancellationFutureOwned {
59        self.cancel.clone().cancelled_owned()
60    }
61
62    /// Check whether the request has already been cancelled.
63    pub fn is_cancelled(&self) -> bool {
64        self.cancel.is_cancelled()
65    }
66
67    /// Run an async operation that should stop promptly on request cancellation.
68    pub async fn run_cancellable<F, T, E>(&self, fut: F) -> Result<T, ToolError>
69    where
70        F: Future<Output = Result<T, E>>,
71        E: Into<ToolError>,
72    {
73        tokio::select! {
74            _ = self.cancelled() => {
75                tracing::info!("tool request cancelled during run_cancellable");
76                Err(ToolError::cancelled(None))
77            }
78            result = fut => result.map_err(Into::into),
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use tokio::time::Duration;
87    use tokio::time::sleep;
88    use tokio::time::timeout;
89
90    #[tokio::test]
91    async fn default_context_is_never_cancelled() {
92        let ctx = ToolContext::default();
93
94        assert!(!ctx.is_cancelled());
95        assert!(
96            timeout(Duration::from_millis(25), ctx.cancelled())
97                .await
98                .is_err()
99        );
100    }
101
102    #[tokio::test]
103    async fn with_cancel_propagates_cancellation() {
104        let cancel = CancellationToken::new();
105        let ctx = ToolContext::with_cancel(cancel.clone());
106
107        cancel.cancel();
108        ctx.cancelled().await;
109
110        assert!(ctx.is_cancelled());
111        assert!(ctx.cancellation_token().is_cancelled());
112    }
113
114    #[tokio::test]
115    async fn run_cancellable_returns_inner_success() {
116        let ctx = ToolContext::default();
117
118        let result = ctx
119            .run_cancellable(async { Ok::<_, ToolError>("done") })
120            .await;
121
122        assert!(matches!(result, Ok("done")));
123    }
124
125    #[tokio::test]
126    async fn run_cancellable_returns_cancelled_when_request_is_cancelled() {
127        let cancel = CancellationToken::new();
128        let ctx = ToolContext::with_cancel(cancel.clone());
129
130        let canceller = tokio::spawn(async move {
131            sleep(Duration::from_millis(25)).await;
132            cancel.cancel();
133        });
134
135        let result = ctx
136            .run_cancellable(async {
137                sleep(Duration::from_secs(5)).await;
138                Ok::<(), ToolError>(())
139            })
140            .await;
141
142        let join_result = canceller.await;
143        assert!(join_result.is_ok());
144        assert!(matches!(result, Err(ToolError::Cancelled { reason: None })));
145    }
146}
agentic_tools_core/context.rs

agentic_tools_core/
context.rs
