ag_ui/server/
cancel.rs

1//! Cooperative cancellation, without a runtime.
2//!
3//! A transport trips the token when the client disconnects or a deadline
4//! passes; the run notices and unwinds. There is deliberately no
5//! `tokio_util::CancellationToken` here — this crate must build for wasm and
6//! for non-tokio executors, so the token is an [`AtomicBool`] plus a waker list.
7
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Mutex};
12use std::task::{Context, Poll, Waker};
13
14/// A shared "stop now" flag.
15///
16/// Cloning is cheap and every clone refers to the same flag, so a transport can
17/// keep one and hand another to the run.
18#[derive(Clone, Debug, Default)]
19pub struct CancellationToken {
20    inner: Arc<Inner>,
21}
22
23#[derive(Debug, Default)]
24struct Inner {
25    cancelled: AtomicBool,
26    wakers: Mutex<Vec<Waker>>,
27}
28
29impl CancellationToken {
30    /// A fresh, un-cancelled token.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Trips the token and wakes everything waiting on it.
36    ///
37    /// Idempotent — cancelling twice is a no-op.
38    pub fn cancel(&self) {
39        if !self.inner.cancelled.swap(true, Ordering::SeqCst) {
40            let mut wakers = lock(&self.inner.wakers);
41            for waker in wakers.drain(..) {
42                waker.wake();
43            }
44        }
45    }
46
47    /// Whether the token has been tripped.
48    pub fn is_cancelled(&self) -> bool {
49        self.inner.cancelled.load(Ordering::SeqCst)
50    }
51
52    /// Resolves once the token is tripped.
53    ///
54    /// Use it to race an in-flight model call:
55    ///
56    /// ```
57    /// # use ag_ui::server::CancellationToken;
58    /// # let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
59    /// # rt.block_on(async {
60    /// let token = CancellationToken::new();
61    /// token.cancel();
62    /// token.cancelled().await;
63    /// # });
64    /// ```
65    pub fn cancelled(&self) -> Cancelled {
66        Cancelled {
67            token: self.clone(),
68        }
69    }
70}
71
72/// The future returned by [`CancellationToken::cancelled`].
73///
74/// It owns a clone of the token rather than borrowing one — one `Arc` bump, in
75/// exchange for a `'static` future that an agent can hold across an await
76/// without dragging a borrow of the run context along with it.
77#[derive(Clone, Debug)]
78#[must_use = "a future does nothing unless awaited"]
79pub struct Cancelled {
80    token: CancellationToken,
81}
82
83impl Future for Cancelled {
84    type Output = ();
85
86    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
87        if self.token.is_cancelled() {
88            return Poll::Ready(());
89        }
90        let mut wakers = lock(&self.token.inner.wakers);
91        // Re-check under the lock: `cancel` may have drained the list between
92        // the load above and here, and would then never see our waker.
93        if self.token.is_cancelled() {
94            return Poll::Ready(());
95        }
96        if !wakers.iter().any(|waker| waker.will_wake(cx.waker())) {
97            wakers.push(cx.waker().clone());
98        }
99        Poll::Pending
100    }
101}
102
103/// A poisoned waker list is still a perfectly good waker list: the only code
104/// that touches it cannot panic while holding the guard.
105fn lock(mutex: &Mutex<Vec<Waker>>) -> std::sync::MutexGuard<'_, Vec<Waker>> {
106    mutex
107        .lock()
108        .unwrap_or_else(|poisoned| poisoned.into_inner())
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn cancel_is_visible_through_clones() {
117        let token = CancellationToken::new();
118        let clone = token.clone();
119        assert!(!clone.is_cancelled());
120        token.cancel();
121        assert!(clone.is_cancelled());
122    }
123
124    #[tokio::test]
125    async fn cancelled_future_resolves() {
126        let token = CancellationToken::new();
127        let waiter = token.clone();
128        let handle = tokio::spawn(async move { waiter.cancelled().await });
129        token.cancel();
130        handle.await.expect("waiter task panicked");
131    }
132
133    #[tokio::test]
134    async fn cancelled_future_is_ready_when_already_cancelled() {
135        let token = CancellationToken::new();
136        token.cancel();
137        token.cancelled().await;
138    }
139}
ag_ui/server/cancel.rs

ag_ui/server/
cancel.rs
