File size: 2,475 Bytes
52a9af3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | use std::sync::Arc;
use std::sync::Mutex;
use tokio::sync::watch;
/// Coordinates user elicitations that pause tool-result delivery for a session.
///
/// Registrations are counted so concurrent elicitations keep the session paused until all of them
/// finish. Consumers can subscribe to pause timeout progress or wait before returning an already
/// captured result.
#[derive(Clone)]
pub(crate) struct ElicitationService {
inner: Arc<Inner>,
}
struct Inner {
state: Mutex<State>,
paused: watch::Sender<bool>,
}
#[derive(Default)]
struct State {
outstanding: i64,
}
pub(crate) struct ElicitationRegistration {
service: ElicitationService,
}
impl ElicitationService {
pub(crate) fn new() -> Self {
let (paused, _paused_rx) = watch::channel(false);
Self {
inner: Arc::new(Inner {
state: Mutex::new(State::default()),
paused,
}),
}
}
pub(crate) fn register(&self) -> ElicitationRegistration {
self.increment();
ElicitationRegistration {
service: self.clone(),
}
}
fn increment(&self) {
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let was_clear = state.outstanding == 0;
assert_ne!(
state.outstanding,
i64::MAX,
"outstanding elicitation count overflowed"
);
state.outstanding += 1;
if was_clear {
self.inner.paused.send_replace(true);
}
}
pub(crate) fn subscribe(&self) -> watch::Receiver<bool> {
self.inner.paused.subscribe()
}
pub(crate) async fn wait_until_clear(&self) {
let mut paused = self.subscribe();
let _ = paused.wait_for(|paused| !*paused).await;
}
fn decrement(&self) {
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
assert!(
state.outstanding > 0,
"elicitation registration count underflowed"
);
state.outstanding -= 1;
if state.outstanding == 0 {
self.inner.paused.send_replace(false);
}
}
}
impl Drop for ElicitationRegistration {
fn drop(&mut self) {
self.service.decrement();
}
}
#[cfg(test)]
#[path = "elicitation_tests.rs"]
mod tests;
|