lc_core/runnables/
cancellation.rs

1// lc-core/src/runnables/cancellation.rs
2//! CancellationToken for aborting long-running operations.
3//!
4//! # Example
5//!
6//! ```rust,ignore
7//! use lc_core::runnables::CancellationToken;
8//! use std::time::Duration;
9//!
10//! let token = CancellationToken::new();
11//!
12//! // In another task, cancel after 30 seconds
13//! let cloned = token.clone();
14//! tokio::spawn(async move {
15//!     tokio::time::sleep(Duration::from_secs(30)).await;
16//!     cloned.cancel();
17//! });
18//!
19//! // In the agent loop, check for cancellation
20//! if token.is_cancelled() {
21//!     return Ok("Agent stopped by cancellation".to_string());
22//! }
23//! ```
24
25use std::sync::atomic::{AtomicBool, Ordering};
26use std::sync::Arc;
27
28/// A token that can be used to signal cancellation of a long-running operation.
29///
30/// Clones share the same underlying cancellation state — cancelling one clone
31/// cancels all of them.
32#[derive(Debug, Clone)]
33pub struct CancellationToken {
34    inner: Arc<AtomicBool>,
35}
36
37impl CancellationToken {
38    /// Creates a new, uncancelled token.
39    pub fn new() -> Self {
40        Self {
41            inner: Arc::new(AtomicBool::new(false)),
42        }
43    }
44
45    /// Signals cancellation.
46    ///
47    /// All clones of this token will become cancelled.
48    pub fn cancel(&self) {
49        self.inner.store(true, Ordering::SeqCst);
50    }
51
52    /// Returns `true` if cancellation has been signaled.
53    pub fn is_cancelled(&self) -> bool {
54        self.inner.load(Ordering::SeqCst)
55    }
56
57    /// Returns a future that resolves when cancellation is signaled.
58    ///
59    /// Uses `tokio::sync::Notify`-style polling. This is a lightweight
60    /// check — it does not block a thread.
61    pub async fn cancelled(&self) {
62        // Simple spin-based wait with yield.
63        // For production use, consider a Notify-based approach.
64        while !self.inner.load(Ordering::SeqCst) {
65            tokio::task::yield_now().await;
66        }
67    }
68}
69
70impl Default for CancellationToken {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn test_new_token_is_not_cancelled() {
82        let token = CancellationToken::new();
83        assert!(!token.is_cancelled());
84    }
85
86    #[test]
87    fn test_cancel_sets_is_cancelled() {
88        let token = CancellationToken::new();
89        token.cancel();
90        assert!(token.is_cancelled());
91    }
92
93    #[test]
94    fn test_clone_shares_cancellation_state() {
95        let token = CancellationToken::new();
96        let clone = token.clone();
97
98        assert!(!token.is_cancelled());
99        assert!(!clone.is_cancelled());
100
101        clone.cancel();
102
103        assert!(token.is_cancelled());
104        assert!(clone.is_cancelled());
105    }
106
107    #[test]
108    fn test_multiple_clones_all_cancelled() {
109        let token = CancellationToken::new();
110        let c1 = token.clone();
111        let c2 = token.clone();
112        let c3 = token.clone();
113
114        token.cancel();
115
116        assert!(c1.is_cancelled());
117        assert!(c2.is_cancelled());
118        assert!(c3.is_cancelled());
119    }
120
121    #[tokio::test]
122    async fn test_cancelled_future_resolves() {
123        let token = CancellationToken::new();
124
125        // Cancel in a background task
126        let cloned = token.clone();
127        tokio::spawn(async move {
128            cloned.cancel();
129        });
130
131        // This should resolve quickly
132        token.cancelled().await;
133        assert!(token.is_cancelled());
134    }
135}
lc_core/runnables/cancellation.rs

lc_core/runnables/
cancellation.rs
