diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index c1a6ecc0..11b7e67d 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -424,14 +424,14 @@ pub struct RunningService> { service: Arc, peer: Peer, handle: tokio::task::JoinHandle, - /// cancellation token with drop guard + cancellation_token: CancellationToken, dg: DropGuard, } impl> Deref for RunningService { type Target = Peer; fn deref(&self) -> &Self::Target { - self.peer() + &self.peer } } @@ -444,6 +444,11 @@ impl> RunningService { pub fn service(&self) -> &S { self.service.as_ref() } + #[inline] + pub fn cancellation_token(&self) -> RunningServiceCancellationToken { + RunningServiceCancellationToken(self.cancellation_token.clone()) + } + #[inline] pub async fn waiting(self) -> Result { self.handle.await } @@ -454,6 +459,15 @@ impl> RunningService { } } +// use a wrapper type so we can tweak the implementation if needed +pub struct RunningServiceCancellationToken(CancellationToken); + +impl RunningServiceCancellationToken { + pub fn cancel(self) { + self.0.cancel(); + } +} + #[derive(Debug)] pub enum QuitReason { Cancelled, @@ -801,6 +815,7 @@ where service, peer: peer_return, handle, + cancellation_token: ct.clone(), dg: ct.drop_guard(), } }