vtcode_core/exec/
cancellation.rs

1use tokio_util::sync::CancellationToken;
2
3use std::future::Future;
4
5use tokio::task_local;
6
7task_local! {
8    static ACTIVE_TOOL_TOKEN: CancellationToken;
9}
10
11/// Run the provided future with the supplied cancellation token made available to tools.
12pub async fn with_tool_cancellation<F, T>(token: CancellationToken, fut: F) -> T
13where
14    F: Future<Output = T>,
15{
16    ACTIVE_TOOL_TOKEN.scope(token, fut).await
17}
18
19/// Retrieve the currently scoped cancellation token, if any.
20pub fn current_tool_cancellation() -> Option<CancellationToken> {
21    ACTIVE_TOOL_TOKEN.try_with(|token| token.clone()).ok()
22}
23
24#[cfg(test)]
25mod tests {
26    use super::*;
27
28    #[tokio::test]
29    async fn scoped_token_is_accessible() {
30        assert!(current_tool_cancellation().is_none());
31        let token = CancellationToken::new();
32        with_tool_cancellation(token.clone(), async move {
33            let current = current_tool_cancellation().expect("token should be set");
34            assert!(!current.is_cancelled());
35            token.cancel();
36            assert!(current.is_cancelled());
37        })
38        .await;
39        assert!(current_tool_cancellation().is_none());
40    }
41}
vtcode_core/exec/cancellation.rs

vtcode_core/exec/
cancellation.rs
