rmcp/

service.rs

use futures::FutureExt;
#[cfg(not(feature = "local"))]
use futures::future::BoxFuture;
#[cfg(feature = "local")]
use futures::future::LocalBoxFuture;
use thiserror::Error;

// ---------------------------------------------------------------------------
// Conditional Send helpers
//
// `MaybeSend`       – supertrait alias: `Send + Sync` without `local`, empty with `local`
// `MaybeSendFuture` – future bound alias: `Send` without `local`, empty with `local`
// `MaybeBoxFuture`  – boxed future type: `BoxFuture` without `local`, `LocalBoxFuture` with `local`
// ---------------------------------------------------------------------------

#[cfg(not(feature = "local"))]
#[doc(hidden)]
pub trait MaybeSend: Send + Sync {}
#[cfg(not(feature = "local"))]
impl<T: Send + Sync> MaybeSend for T {}

#[cfg(feature = "local")]
#[doc(hidden)]
pub trait MaybeSend {}
#[cfg(feature = "local")]
impl<T> MaybeSend for T {}

#[cfg(not(feature = "local"))]
#[doc(hidden)]
pub trait MaybeSendFuture: Send {}
#[cfg(not(feature = "local"))]
impl<T: Send> MaybeSendFuture for T {}

#[cfg(feature = "local")]
#[doc(hidden)]
pub trait MaybeSendFuture {}
#[cfg(feature = "local")]
impl<T> MaybeSendFuture for T {}

#[cfg(not(feature = "local"))]
pub(crate) type MaybeBoxFuture<'a, T> = BoxFuture<'a, T>;
#[cfg(feature = "local")]
pub(crate) type MaybeBoxFuture<'a, T> = LocalBoxFuture<'a, T>;

#[cfg(feature = "server")]
use crate::model::ClientNotification;
#[cfg(feature = "server")]
use crate::model::ServerJsonRpcMessage;
#[cfg(feature = "client")]
use crate::model::ServerNotification;
use crate::{
    error::ErrorData as McpError,
    model::{
        CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta,
        JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta,
        NumberOrString, ProgressToken, RequestId,
    },
    transport::{DynamicTransportError, IntoTransport, Transport},
};
#[cfg(feature = "client")]
mod client;
#[cfg(feature = "client")]
pub use client::*;
#[cfg(feature = "server")]
mod server;
#[cfg(feature = "server")]
pub use server::*;
#[cfg(feature = "tower")]
mod tower;
use tokio_util::sync::{CancellationToken, DropGuard};
#[cfg(feature = "tower")]
pub use tower::*;
use tracing::{Instrument as _, instrument};
#[derive(Error, Debug)]
#[non_exhaustive]
pub enum ServiceError {
    #[error("Mcp error: {0}")]
    McpError(McpError),
    #[error("Transport send error: {0}")]
    TransportSend(DynamicTransportError),
    #[error("Transport closed")]
    TransportClosed,
    #[error("Unexpected response type")]
    UnexpectedResponse,
    #[error("task cancelled for reason {}", reason.as_deref().unwrap_or("<unknown>"))]
    Cancelled { reason: Option<String> },
    #[error("request timeout after {}", chrono::Duration::from_std(*timeout).unwrap_or_default())]
    Timeout { timeout: Duration },
}

trait TransferObject:
    std::fmt::Debug + Clone + serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static
{
}

impl<T> TransferObject for T where
    T: std::fmt::Debug
        + serde::Serialize
        + serde::de::DeserializeOwned
        + Send
        + Sync
        + 'static
        + Clone
{
}

#[allow(private_bounds, reason = "there's no the third implementation")]
pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
    type Req: TransferObject + GetMeta + GetExtensions;
    type Resp: TransferObject;
    type Not: TryInto<CancelledNotification, Error = Self::Not>
        + From<CancelledNotification>
        + TransferObject;
    type PeerReq: TransferObject + GetMeta + GetExtensions;
    type PeerResp: TransferObject;
    type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
        + From<CancelledNotification>
        + TransferObject
        + GetMeta
        + GetExtensions;
    type InitializeError;
    const IS_CLIENT: bool;
    type Info: TransferObject;
    type PeerInfo: TransferObject;
}

pub type TxJsonRpcMessage<R> =
    JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
pub type RxJsonRpcMessage<R> = JsonRpcMessage<
    <R as ServiceRole>::PeerReq,
    <R as ServiceRole>::PeerResp,
    <R as ServiceRole>::PeerNot,
>;

#[cfg(not(feature = "local"))]
pub trait Service<R: ServiceRole>: Send + Sync + 'static {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> impl Future<Output = Result<R::Resp, McpError>> + MaybeSendFuture + '_;
    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_;
    fn get_info(&self) -> R::Info;
}

#[cfg(feature = "local")]
pub trait Service<R: ServiceRole>: 'static {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> impl Future<Output = Result<R::Resp, McpError>> + MaybeSendFuture + '_;
    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_;
    fn get_info(&self) -> R::Info;
}

pub trait ServiceExt<R: ServiceRole>: Service<R> + Sized {
    /// Convert this service to a dynamic boxed service
    ///
    /// This could be very helpful when you want to store the services in a collection
    fn into_dyn(self) -> Box<dyn DynService<R>> {
        Box::new(self)
    }
    fn serve<T, E, A>(
        self,
        transport: T,
    ) -> impl Future<Output = Result<RunningService<R, Self>, R::InitializeError>> + MaybeSendFuture
    where
        T: IntoTransport<R, E, A>,
        E: std::error::Error + Send + Sync + 'static,
        Self: Sized,
    {
        Self::serve_with_ct(self, transport, Default::default())
    }
    fn serve_with_ct<T, E, A>(
        self,
        transport: T,
        ct: CancellationToken,
    ) -> impl Future<Output = Result<RunningService<R, Self>, R::InitializeError>> + MaybeSendFuture
    where
        T: IntoTransport<R, E, A>,
        E: std::error::Error + Send + Sync + 'static,
        Self: Sized;
}

impl<R: ServiceRole> Service<R> for Box<dyn DynService<R>> {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> impl Future<Output = Result<R::Resp, McpError>> + MaybeSendFuture + '_ {
        DynService::handle_request(self.as_ref(), request, context)
    }

    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_ {
        DynService::handle_notification(self.as_ref(), notification, context)
    }

    fn get_info(&self) -> R::Info {
        DynService::get_info(self.as_ref())
    }
}

#[cfg(not(feature = "local"))]
pub trait DynService<R: ServiceRole>: Send + Sync {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>>;
    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> MaybeBoxFuture<'_, Result<(), McpError>>;
    fn get_info(&self) -> R::Info;
}

#[cfg(feature = "local")]
pub trait DynService<R: ServiceRole> {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>>;
    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> MaybeBoxFuture<'_, Result<(), McpError>>;
    fn get_info(&self) -> R::Info;
}

impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
    fn handle_request(
        &self,
        request: R::PeerReq,
        context: RequestContext<R>,
    ) -> MaybeBoxFuture<'_, Result<R::Resp, McpError>> {
        Box::pin(self.handle_request(request, context))
    }
    fn handle_notification(
        &self,
        notification: R::PeerNot,
        context: NotificationContext<R>,
    ) -> MaybeBoxFuture<'_, Result<(), McpError>> {
        Box::pin(self.handle_notification(notification, context))
    }
    fn get_info(&self) -> R::Info {
        self.get_info()
    }
}

use std::{
    collections::{HashMap, VecDeque},
    ops::Deref,
    sync::{Arc, atomic::AtomicU64},
    time::Duration,
};

use tokio::sync::mpsc;

pub trait RequestIdProvider: Send + Sync + 'static {
    fn next_request_id(&self) -> RequestId;
}

pub trait ProgressTokenProvider: Send + Sync + 'static {
    fn next_progress_token(&self) -> ProgressToken;
}

pub type AtomicU32RequestIdProvider = AtomicU32Provider;
pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;

#[derive(Debug, Default)]
pub struct AtomicU32Provider {
    id: AtomicU64,
}

impl RequestIdProvider for AtomicU32Provider {
    fn next_request_id(&self) -> RequestId {
        let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        // Safe conversion: we start at 0 and increment by 1, so we won't overflow i64::MAX in practice
        RequestId::Number(id as i64)
    }
}

impl ProgressTokenProvider for AtomicU32Provider {
    fn next_progress_token(&self) -> ProgressToken {
        let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        ProgressToken(NumberOrString::Number(id as i64))
    }
}

#[doc(hidden)]
pub trait ProgressNotificationToken {
    fn progress_token(&self) -> Option<&ProgressToken>;
}

#[cfg(feature = "server")]
impl ProgressNotificationToken for ClientNotification {
    fn progress_token(&self) -> Option<&ProgressToken> {
        match self {
            ClientNotification::ProgressNotification(notification) => {
                Some(&notification.params.progress_token)
            }
            _ => None,
        }
    }
}

#[cfg(feature = "client")]
impl ProgressNotificationToken for ServerNotification {
    fn progress_token(&self) -> Option<&ProgressToken> {
        match self {
            ServerNotification::ProgressNotification(notification) => {
                Some(&notification.params.progress_token)
            }
            _ => None,
        }
    }
}

type Responder<T> = tokio::sync::oneshot::Sender<T>;
type ProgressTimeoutWatchers = Arc<tokio::sync::RwLock<HashMap<ProgressToken, mpsc::Sender<()>>>>;

/// A handle to a remote request
///
/// You can cancel it by call [`RequestHandle::cancel`] with a reason,
///
/// or wait for response by call [`RequestHandle::await_response`]
#[derive(Debug)]
#[non_exhaustive]
pub struct RequestHandle<R: ServiceRole> {
    pub rx: tokio::sync::oneshot::Receiver<Result<R::PeerResp, ServiceError>>,
    pub options: PeerRequestOptions,
    pub peer: Peer<R>,
    pub id: RequestId,
    pub progress_token: ProgressToken,
    progress_reset_rx: Option<mpsc::Receiver<()>>,
}

impl<R: ServiceRole> RequestHandle<R> {
    pub const REQUEST_TIMEOUT_REASON: &str = "request timeout";
    pub const REQUEST_MAX_TOTAL_TIMEOUT_REASON: &str = "maximum total timeout exceeded";

    pub async fn await_response(mut self) -> Result<R::PeerResp, ServiceError> {
        let timeout = self.options.timeout;
        let max_total_timeout = self.options.max_total_timeout;
        let reset_timeout_on_progress = self.options.reset_timeout_on_progress;

        let has_progress_reset_rx = self.progress_reset_rx.is_some();
        let progress_token = self.progress_token.clone();

        let result = match (timeout, max_total_timeout, reset_timeout_on_progress) {
            (Some(timeout), None, false) => match tokio::time::timeout(timeout, &mut self.rx).await
            {
                Ok(response) => response.map_err(|_e| ServiceError::TransportClosed)?,
                Err(_) => {
                    let error = Err(ServiceError::Timeout { timeout });
                    // cancel this request
                    self.send_timeout_cancel_notification(Self::REQUEST_TIMEOUT_REASON)
                        .await;
                    error
                }
            },
            (None, None, _) => (&mut self.rx)
                .await
                .map_err(|_e| ServiceError::TransportClosed)?,
            _ => {
                self.await_response_with_progress_timeout(
                    timeout,
                    max_total_timeout,
                    reset_timeout_on_progress,
                )
                .await
            }
        };

        Self::cleanup_progress_timeout_watcher(
            &self.peer.progress_timeout_watchers,
            &progress_token,
            has_progress_reset_rx,
        )
        .await;
        result
    }

    async fn send_timeout_cancel_notification(&self, reason: &str) {
        let notification = CancelledNotification {
            params: CancelledNotificationParam {
                request_id: Some(self.id.clone()),
                reason: Some(reason.to_owned()),
                meta: None,
            },
            method: crate::model::CancelledNotificationMethod,
            extensions: Default::default(),
        };
        let _ = self.peer.send_notification(notification.into()).await;
    }

    async fn await_response_with_progress_timeout(
        &mut self,
        timeout: Option<Duration>,
        max_total_timeout: Option<Duration>,
        reset_timeout_on_progress: bool,
    ) -> Result<R::PeerResp, ServiceError> {
        let mut idle_sleep =
            timeout.map(|timeout| (timeout, Box::pin(tokio::time::sleep(timeout))));
        let mut max_total_sleep =
            max_total_timeout.map(|timeout| (timeout, Box::pin(tokio::time::sleep(timeout))));

        loop {
            tokio::select! {
                biased;

                response = &mut self.rx => {
                    return response.map_err(|_e| ServiceError::TransportClosed)?;
                }
                _ = async {
                    if let Some((_, sleep)) = idle_sleep.as_mut() {
                        sleep.as_mut().await;
                    }
                }, if idle_sleep.is_some() => {
                    if let Some((timeout, _)) = idle_sleep.as_ref() {
                        self.send_timeout_cancel_notification(Self::REQUEST_TIMEOUT_REASON).await;
                        return Err(ServiceError::Timeout { timeout: *timeout });
                    }
                }
                _ = async {
                    if let Some((_, sleep)) = max_total_sleep.as_mut() {
                        sleep.as_mut().await;
                    }
                }, if max_total_sleep.is_some() => {
                    if let Some((timeout, _)) = max_total_sleep.as_ref() {
                        self.send_timeout_cancel_notification(Self::REQUEST_MAX_TOTAL_TIMEOUT_REASON).await;
                        return Err(ServiceError::Timeout { timeout: *timeout });
                    }
                }
                progress = async {
                    match self.progress_reset_rx.as_mut() {
                        Some(rx) => rx.recv().await,
                        None => None,
                    }
                }, if reset_timeout_on_progress && idle_sleep.is_some() && self.progress_reset_rx.is_some() => {
                    if progress.is_some() {
                        if let Some((timeout, sleep)) = idle_sleep.as_mut() {
                            sleep.as_mut().reset(tokio::time::Instant::now() + *timeout);
                        }
                    }
                }
            }
        }
    }

    /// Cancel this request
    pub async fn cancel(self, reason: Option<String>) -> Result<(), ServiceError> {
        Self::cleanup_progress_timeout_watcher(
            &self.peer.progress_timeout_watchers,
            &self.progress_token,
            self.progress_reset_rx.is_some(),
        )
        .await;
        let notification = CancelledNotification {
            params: CancelledNotificationParam {
                request_id: Some(self.id),
                reason,
                meta: None,
            },
            method: crate::model::CancelledNotificationMethod,
            extensions: Default::default(),
        };
        self.peer.send_notification(notification.into()).await?;
        Ok(())
    }

    async fn cleanup_progress_timeout_watcher(
        progress_timeout_watchers: &ProgressTimeoutWatchers,
        progress_token: &ProgressToken,
        has_progress_reset_rx: bool,
    ) {
        if has_progress_reset_rx {
            progress_timeout_watchers
                .write()
                .await
                .remove(progress_token);
        }
    }
}

#[derive(Debug)]
pub(crate) enum PeerSinkMessage<R: ServiceRole> {
    Request {
        request: R::Req,
        id: RequestId,
        responder: Responder<Result<R::PeerResp, ServiceError>>,
    },
    Notification {
        notification: R::Not,
        responder: Responder<Result<(), ServiceError>>,
    },
}

/// An interface to fetch the remote client or server
///
/// For general purpose, call [`Peer::send_request`] or [`Peer::send_notification`] to send message to remote peer.
///
/// To create a cancellable request, call [`Peer::send_request_with_option`].
#[derive(Clone)]
pub struct Peer<R: ServiceRole> {
    tx: mpsc::Sender<PeerSinkMessage<R>>,
    request_id_provider: Arc<dyn RequestIdProvider>,
    progress_token_provider: Arc<dyn ProgressTokenProvider>,
    progress_timeout_watchers: ProgressTimeoutWatchers,
    info: Arc<std::sync::RwLock<Option<Arc<R::PeerInfo>>>>,
}

impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PeerSink")
            .field("tx", &self.tx)
            .field("is_client", &R::IS_CLIENT)
            .finish()
    }
}

type ProxyOutbound<R> = mpsc::Receiver<PeerSinkMessage<R>>;

#[derive(Debug, Default)]
#[non_exhaustive]
pub struct PeerRequestOptions {
    pub timeout: Option<Duration>,
    pub meta: Option<Meta>,
    /// Reset the request timeout when a matching progress notification is received.
    pub reset_timeout_on_progress: bool,
    /// Maximum total time to wait for the request, regardless of progress notifications.
    pub max_total_timeout: Option<Duration>,
}

impl PeerRequestOptions {
    pub fn no_options() -> Self {
        Self::default()
    }

    pub fn with_timeout(timeout: Duration) -> Self {
        Self {
            timeout: Some(timeout),
            ..Self::default()
        }
    }

    pub fn reset_timeout_on_progress(mut self) -> Self {
        self.reset_timeout_on_progress = true;
        self
    }

    pub fn with_max_total_timeout(mut self, timeout: Duration) -> Self {
        self.max_total_timeout = Some(timeout);
        self
    }
}

impl<R: ServiceRole> Peer<R> {
    const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
    pub(crate) fn new(
        request_id_provider: Arc<dyn RequestIdProvider>,
        peer_info: Option<R::PeerInfo>,
    ) -> (Peer<R>, ProxyOutbound<R>) {
        let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
        (
            Self {
                tx,
                request_id_provider,
                progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
                progress_timeout_watchers: Default::default(),
                info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))),
            },
            rx,
        )
    }
    pub async fn send_notification(&self, notification: R::Not) -> Result<(), ServiceError> {
        let (responder, receiver) = tokio::sync::oneshot::channel();
        self.tx
            .send(PeerSinkMessage::Notification {
                notification,
                responder,
            })
            .await
            .map_err(|_m| ServiceError::TransportClosed)?;
        receiver.await.map_err(|_e| ServiceError::TransportClosed)?
    }
    pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
        self.send_request_with_option(request, PeerRequestOptions::no_options())
            .await?
            .await_response()
            .await
    }

    pub async fn send_cancellable_request(
        &self,
        request: R::Req,
        options: PeerRequestOptions,
    ) -> Result<RequestHandle<R>, ServiceError> {
        self.send_request_with_option(request, options).await
    }

    pub async fn send_request_with_option(
        &self,
        mut request: R::Req,
        options: PeerRequestOptions,
    ) -> Result<RequestHandle<R>, ServiceError> {
        let id = self.request_id_provider.next_request_id();
        let progress_token = self.progress_token_provider.next_progress_token();
        if let Some(meta) = options.meta.clone() {
            request.get_meta_mut().extend(meta);
        }
        request
            .get_meta_mut()
            .set_progress_token(progress_token.clone());
        let (responder, receiver) = tokio::sync::oneshot::channel();
        let progress_reset_rx = if options.reset_timeout_on_progress && options.timeout.is_some() {
            let (sender, receiver) = mpsc::channel(1);
            self.progress_timeout_watchers
                .write()
                .await
                .insert(progress_token.clone(), sender);
            Some(receiver)
        } else {
            None
        };
        if self
            .tx
            .send(PeerSinkMessage::Request {
                request,
                id: id.clone(),
                responder,
            })
            .await
            .is_err()
        {
            if progress_reset_rx.is_some() {
                self.progress_timeout_watchers
                    .write()
                    .await
                    .remove(&progress_token);
            }
            return Err(ServiceError::TransportClosed);
        }
        Ok(RequestHandle {
            id,
            rx: receiver,
            progress_token,
            options,
            peer: self.clone(),
            progress_reset_rx,
        })
    }

    async fn notify_progress_timeout_watcher(&self, progress_token: &ProgressToken) {
        let sender = self
            .progress_timeout_watchers
            .read()
            .await
            .get(progress_token)
            .cloned();
        if let Some(sender) = sender {
            match sender.try_send(()) {
                Ok(()) => {}
                Err(mpsc::error::TrySendError::Full(_)) => {
                    tracing::trace!(?progress_token, "progress timeout watcher channel is full");
                }
                Err(mpsc::error::TrySendError::Closed(_)) => {
                    self.progress_timeout_watchers
                        .write()
                        .await
                        .remove(progress_token);
                }
            }
        }
    }

    /// Snapshot of the peer's handshake info.
    pub fn peer_info(&self) -> Option<Arc<R::PeerInfo>> {
        self.info.read().expect("peer info lock poisoned").clone()
    }

    /// Stores the peer's handshake info, overwriting any previous value.
    pub fn set_peer_info(&self, info: R::PeerInfo) {
        *self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info));
    }

    pub fn is_transport_closed(&self) -> bool {
        self.tx.is_closed()
    }
}

#[derive(Debug)]
pub struct RunningService<R: ServiceRole, S: Service<R>> {
    service: Arc<S>,
    peer: Peer<R>,
    handle: Option<tokio::task::JoinHandle<QuitReason>>,
    cancellation_token: CancellationToken,
    dg: DropGuard,
}
impl<R: ServiceRole, S: Service<R>> Deref for RunningService<R, S> {
    type Target = Peer<R>;

    fn deref(&self) -> &Self::Target {
        &self.peer
    }
}

impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
    #[inline]
    pub fn peer(&self) -> &Peer<R> {
        &self.peer
    }
    #[inline]
    pub fn service(&self) -> &S {
        self.service.as_ref()
    }
    #[inline]
    pub fn cancellation_token(&self) -> RunningServiceCancellationToken {
        RunningServiceCancellationToken(self.cancellation_token.clone())
    }

    /// Returns true if the service has been closed or cancelled.
    #[inline]
    pub fn is_closed(&self) -> bool {
        self.handle.is_none() || self.cancellation_token.is_cancelled()
    }

    /// Wait for the service to complete.
    ///
    /// This will block until the service loop terminates (either due to
    /// cancellation, transport closure, or an error).
    #[inline]
    pub async fn waiting(mut self) -> Result<QuitReason, tokio::task::JoinError> {
        match self.handle.take() {
            Some(handle) => handle.await,
            None => Ok(QuitReason::Closed),
        }
    }

    /// Gracefully close the connection and wait for cleanup to complete.
    ///
    /// This method cancels the service, waits for the background task to finish
    /// (which includes calling `transport.close()`), and ensures all cleanup
    /// operations complete before returning.
    ///
    /// Unlike [`cancel`](Self::cancel), this method takes `&mut self` and can be
    /// called without consuming the `RunningService`. After calling this method,
    /// the service is considered closed and subsequent operations will fail.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let mut client = ().serve(transport).await?;
    /// // ... use the client ...
    /// client.close().await?;
    /// ```
    pub async fn close(&mut self) -> Result<QuitReason, tokio::task::JoinError> {
        if let Some(handle) = self.handle.take() {
            // Disarm the drop guard so it doesn't try to cancel again
            // We need to cancel manually and wait for completion
            self.cancellation_token.cancel();
            handle.await
        } else {
            // Already closed
            Ok(QuitReason::Closed)
        }
    }

    /// Gracefully close the connection with a timeout.
    ///
    /// Similar to [`close`](Self::close), but returns after the specified timeout
    /// if the cleanup doesn't complete in time. This is useful for ensuring
    /// a bounded shutdown time.
    ///
    /// Returns `Ok(Some(reason))` if shutdown completed within the timeout,
    /// `Ok(None)` if the timeout was reached, or `Err` if there was a join error.
    pub async fn close_with_timeout(
        &mut self,
        timeout: Duration,
    ) -> Result<Option<QuitReason>, tokio::task::JoinError> {
        if let Some(handle) = self.handle.take() {
            self.cancellation_token.cancel();
            match tokio::time::timeout(timeout, handle).await {
                Ok(result) => result.map(Some),
                Err(_elapsed) => {
                    tracing::warn!(
                        "close_with_timeout: cleanup did not complete within {:?}",
                        timeout
                    );
                    Ok(None)
                }
            }
        } else {
            Ok(Some(QuitReason::Closed))
        }
    }

    /// Cancel the service and wait for cleanup to complete.
    ///
    /// This consumes the `RunningService` and ensures the connection is properly
    /// closed. For a non-consuming alternative, see [`close`](Self::close).
    pub async fn cancel(mut self) -> Result<QuitReason, tokio::task::JoinError> {
        // Disarm the drop guard since we're handling cancellation explicitly
        let _ = std::mem::replace(&mut self.dg, self.cancellation_token.clone().drop_guard());
        self.close().await
    }
}

impl<R: ServiceRole, S: Service<R>> Drop for RunningService<R, S> {
    fn drop(&mut self) {
        if self.handle.is_some() && !self.cancellation_token.is_cancelled() {
            tracing::debug!(
                "RunningService dropped without explicit close(). \
                 The connection will be closed asynchronously. \
                 For guaranteed cleanup, call close() or cancel() before dropping."
            );
        }
        // The DropGuard will handle cancellation
    }
}

// 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)]
#[non_exhaustive]
pub enum QuitReason {
    Cancelled,
    Closed,
    JoinError(tokio::task::JoinError),
}

/// Request execution context
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct RequestContext<R: ServiceRole> {
    /// this token will be cancelled when the [`CancelledNotification`] is received.
    pub ct: CancellationToken,
    pub id: RequestId,
    pub meta: Meta,
    pub extensions: Extensions,
    /// An interface to fetch the remote client or server
    pub peer: Peer<R>,
}

impl<R: ServiceRole> RequestContext<R> {
    /// Create a new RequestContext.
    pub fn new(id: RequestId, peer: Peer<R>) -> Self {
        Self {
            ct: CancellationToken::new(),
            id,
            meta: Meta::default(),
            extensions: Extensions::default(),
            peer,
        }
    }
}

#[cfg(feature = "server")]
impl RequestContext<RoleServer> {
    /// The protocol version the client negotiated, or `None` before peer info is recorded.
    pub fn protocol_version(&self) -> Option<crate::model::ProtocolVersion> {
        self.peer
            .peer_info()
            .map(|info| info.protocol_version.clone())
    }
}

/// Request execution context
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct NotificationContext<R: ServiceRole> {
    pub meta: Meta,
    pub extensions: Extensions,
    /// An interface to fetch the remote client or server
    pub peer: Peer<R>,
}

/// Use this function to skip initialization process
pub fn serve_directly<R, S, T, E, A>(
    service: S,
    transport: T,
    peer_info: Option<R::PeerInfo>,
) -> RunningService<R, S>
where
    R: ServiceRole,
    R::PeerNot: ProgressNotificationToken,
    S: Service<R>,
    T: IntoTransport<R, E, A>,
    E: std::error::Error + Send + Sync + 'static,
{
    serve_directly_with_ct(service, transport, peer_info, Default::default())
}

/// Use this function to skip initialization process
pub fn serve_directly_with_ct<R, S, T, E, A>(
    service: S,
    transport: T,
    peer_info: Option<R::PeerInfo>,
    ct: CancellationToken,
) -> RunningService<R, S>
where
    R: ServiceRole,
    R::PeerNot: ProgressNotificationToken,
    S: Service<R>,
    T: IntoTransport<R, E, A>,
    E: std::error::Error + Send + Sync + 'static,
{
    let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
    serve_inner(service, transport.into_transport(), peer, peer_rx, ct)
}

/// Spawn a task that may hold `!Send` state when the `local` feature is active.
///
/// Without the `local` feature this is `tokio::spawn` (requires `Future: Send + 'static`).
/// With `local` it uses `tokio::task::spawn_local` (requires only `Future: 'static`).
#[cfg(not(feature = "local"))]
fn spawn_service_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    tokio::spawn(future)
}

#[cfg(feature = "local")]
fn spawn_service_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + 'static,
    F::Output: 'static,
{
    tokio::task::spawn_local(future)
}

#[instrument(skip_all)]
fn serve_inner<R, S, T>(
    service: S,
    transport: T,
    peer: Peer<R>,
    mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
    ct: CancellationToken,
) -> RunningService<R, S>
where
    R: ServiceRole,
    R::PeerNot: ProgressNotificationToken,
    S: Service<R>,
    T: Transport<R> + 'static,
{
    const SINK_PROXY_BUFFER_SIZE: usize = 64;
    let (sink_proxy_tx, mut sink_proxy_rx) =
        tokio::sync::mpsc::channel::<TxJsonRpcMessage<R>>(SINK_PROXY_BUFFER_SIZE);
    let peer_info = peer.peer_info();
    if R::IS_CLIENT {
        tracing::info!(?peer_info, "Service initialized as client");
    } else {
        tracing::info!(?peer_info, "Service initialized as server");
    }

    let mut local_responder_pool =
        HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::new();
    let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
    let shared_service = Arc::new(service);
    // for return
    let service = shared_service.clone();

    // let message_sink = tokio::sync::
    // let mut stream = std::pin::pin!(stream);
    let serve_loop_ct = ct.child_token();
    let peer_return: Peer<R> = peer.clone();
    let current_span = tracing::Span::current();
    let handle = spawn_service_task(async move {
        let mut transport = transport.into_transport();
        let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
        let mut send_task_set = tokio::task::JoinSet::<SendTaskResult>::new();
        let mut response_send_tasks = tokio::task::JoinSet::<()>::new();
        #[derive(Debug)]
        enum SendTaskResult {
            Request {
                id: RequestId,
                result: Result<(), DynamicTransportError>,
            },
            Notification {
                responder: Responder<Result<(), ServiceError>>,
                cancellation_param: Option<CancelledNotificationParam>,
                result: Result<(), DynamicTransportError>,
            },
        }
        #[derive(Debug)]
        enum Event<R: ServiceRole> {
            ProxyMessage(PeerSinkMessage<R>),
            PeerMessage(RxJsonRpcMessage<R>),
            ToSink(TxJsonRpcMessage<R>),
            SendTaskResult(SendTaskResult),
        }

        let quit_reason = loop {
            let evt = if let Some(m) = batch_messages.pop_front() {
                Event::PeerMessage(m)
            } else {
                tokio::select! {
                    m = sink_proxy_rx.recv(), if !sink_proxy_rx.is_closed() => {
                        if let Some(m) = m {
                            Event::ToSink(m)
                        } else {
                            continue
                        }
                    }
                    m = transport.receive() => {
                        if let Some(m) = m {
                            Event::PeerMessage(m)
                        } else {
                            // input stream closed
                            tracing::info!("input stream terminated");
                            break QuitReason::Closed
                        }
                    }
                    m = peer_rx.recv(), if !peer_rx.is_closed() => {
                        if let Some(m) = m {
                            Event::ProxyMessage(m)
                        } else {
                            continue
                        }
                    }
                    m = send_task_set.join_next(), if !send_task_set.is_empty() => {
                        let Some(result) = m else {
                            continue
                        };
                        match result {
                            Err(e) => {
                                // join error, which is serious, we should quit.
                                tracing::error!(%e, "send request task encounter a tokio join error");
                                break QuitReason::JoinError(e)
                            }
                            Ok(result) => {
                                Event::SendTaskResult(result)
                            }
                        }
                    }
                    _ = serve_loop_ct.cancelled() => {
                        tracing::info!("task cancelled");
                        break QuitReason::Cancelled
                    }
                }
            };

            tracing::trace!(?evt, "new event");
            match evt {
                Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
                    if let Err(e) = result {
                        if let Some(responder) = local_responder_pool.remove(&id) {
                            let _ = responder.send(Err(ServiceError::TransportSend(e)));
                        }
                    }
                }
                Event::SendTaskResult(SendTaskResult::Notification {
                    responder,
                    result,
                    cancellation_param,
                }) => {
                    let response = if let Err(e) = result {
                        Err(ServiceError::TransportSend(e))
                    } else {
                        Ok(())
                    };
                    let _ = responder.send(response);
                    if let Some(param) = cancellation_param {
                        if let Some(request_id) = &param.request_id {
                            if let Some(responder) = local_responder_pool.remove(request_id) {
                                tracing::info!(id = %request_id, reason = param.reason, "cancelled");
                                let _response_result = responder.send(Err(ServiceError::Cancelled {
                                    reason: param.reason.clone(),
                                }));
                            }
                        }
                    }
                }
                // response and error
                Event::ToSink(m) => {
                    if let Some(id) = match &m {
                        JsonRpcMessage::Response(response) => Some(&response.id),
                        JsonRpcMessage::Error(error) => error.id.as_ref(),
                        _ => None,
                    } {
                        if let Some(ct) = local_ct_pool.remove(id) {
                            ct.cancel();
                        }
                        let send = transport.send(m);
                        let current_span = tracing::Span::current();
                        response_send_tasks.spawn(async move {
                            let send_result = send.await;
                            if let Err(error) = send_result {
                                tracing::error!(%error, "fail to response message");
                            }
                        }.instrument(current_span));
                    }
                }
                Event::ProxyMessage(PeerSinkMessage::Request {
                    request,
                    id,
                    responder,
                }) => {
                    local_responder_pool.insert(id.clone(), responder);
                    let send = transport.send(JsonRpcMessage::request(request, id.clone()));
                    {
                        let id = id.clone();
                        let current_span = tracing::Span::current();
                        send_task_set.spawn(send.map(move |r| SendTaskResult::Request {
                            id,
                            result: r.map_err(DynamicTransportError::new::<T, R>),
                        }).instrument(current_span));
                    }
                }
                Event::ProxyMessage(PeerSinkMessage::Notification {
                    notification,
                    responder,
                }) => {
                    // catch cancellation notification
                    let mut cancellation_param = None;
                    let notification = match notification.try_into() {
                        Ok::<CancelledNotification, _>(cancelled) => {
                            cancellation_param.replace(cancelled.params.clone());
                            cancelled.into()
                        }
                        Err(notification) => notification,
                    };
                    let send = transport.send(JsonRpcMessage::notification(notification));
                    let current_span = tracing::Span::current();
                    send_task_set.spawn(send.map(move |result| SendTaskResult::Notification {
                        responder,
                        cancellation_param,
                        result: result.map_err(DynamicTransportError::new::<T, R>),
                    }).instrument(current_span));
                }
                Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest {
                    id,
                    mut request,
                    ..
                })) => {
                    tracing::debug!(%id, ?request, "received request");
                    {
                        let service = shared_service.clone();
                        let sink = sink_proxy_tx.clone();
                        let request_ct = serve_loop_ct.child_token();
                        let context_ct = request_ct.child_token();
                        local_ct_pool.insert(id.clone(), request_ct);
                        let mut extensions = Extensions::new();
                        let mut meta = Meta::new();
                        // avoid clone
                        // swap meta firstly, otherwise progress token will be lost
                        std::mem::swap(&mut meta, request.get_meta_mut());
                        std::mem::swap(&mut extensions, request.extensions_mut());
                        let context = RequestContext {
                            ct: context_ct,
                            id: id.clone(),
                            peer: peer.clone(),
                            meta,
                            extensions,
                        };
                        let current_span = tracing::Span::current();
                        spawn_service_task(async move {
                            let result = service
                                .handle_request(request, context)
                                .await;
                            let response = match result {
                                Ok(result) => {
                                    tracing::debug!(%id, ?result, "response message");
                                    JsonRpcMessage::response(result, id)
                                }
                                Err(error) => {
                                    tracing::warn!(%id, ?error, "response error");
                                    JsonRpcMessage::error(error, Some(id))
                                }
                            };
                            let _send_result = sink.send(response).await;
                        }.instrument(current_span));
                    }
                }
                Event::PeerMessage(JsonRpcMessage::Notification(JsonRpcNotification {
                    notification,
                    ..
                })) => {
                    tracing::info!(?notification, "received notification");
                    // catch cancelled notification
                    let mut notification = match notification.try_into() {
                        Ok::<CancelledNotification, _>(cancelled) => {
                            if let Some(request_id) = &cancelled.params.request_id {
                                if let Some(ct) = local_ct_pool.remove(request_id) {
                                    tracing::info!(id = %request_id, reason = cancelled.params.reason, "cancelled");
                                    ct.cancel();
                                }
                            }
                            cancelled.into()
                        }
                        Err(notification) => notification,
                    };
                    if let Some(progress_token) = notification.progress_token() {
                        peer.notify_progress_timeout_watcher(progress_token).await;
                    }
                    {
                        let service = shared_service.clone();
                        let mut extensions = Extensions::new();
                        let mut meta = Meta::new();
                        // avoid clone
                        std::mem::swap(&mut extensions, notification.extensions_mut());
                        std::mem::swap(&mut meta, notification.get_meta_mut());
                        let context = NotificationContext {
                            peer: peer.clone(),
                            meta,
                            extensions,
                        };
                        let current_span = tracing::Span::current();
                        spawn_service_task(async move {
                            let result = service.handle_notification(notification, context).await;
                            if let Err(error) = result {
                                tracing::warn!(%error, "Error sending notification");
                            }
                        }.instrument(current_span));
                    }
                }
                Event::PeerMessage(JsonRpcMessage::Response(JsonRpcResponse {
                    result,
                    id,
                    ..
                })) => {
                    if let Some(responder) = local_responder_pool.remove(&id) {
                        let response_result = responder.send(Ok(result));
                        if let Err(_error) = response_result {
                            tracing::warn!(%id, "Error sending response");
                        }
                    }
                }
                Event::PeerMessage(JsonRpcMessage::Error(JsonRpcError { error, id, .. })) => {
                    let Some(id) = id else {
                        // MCP error responses without an id (e.g. Parse error / Invalid Request)
                        // can't be routed back to a pending request — log and drop.
                        tracing::debug!(?error, "received id-less peer error");
                        continue;
                    };
                    if let Some(responder) = local_responder_pool.remove(&id) {
                        let _response_result = responder.send(Err(ServiceError::McpError(error)));
                        if let Err(_error) = _response_result {
                            tracing::warn!(%id, "Error sending response");
                        }
                    }
                }
            }
        };

        // Drain in-flight handler responses before closing the transport.
        // When stdin EOF or cancellation arrives, spawned handler tasks may still
        // be finishing. We need to:
        // 1. Wait for response sends that were already spawned in the main loop
        // 2. Drain any remaining handler responses from the channel
        let drain_timeout = match &quit_reason {
            QuitReason::Closed => Some(Duration::from_secs(5)),
            QuitReason::Cancelled => Some(Duration::from_secs(2)),
            _ => None,
        };
        if let Some(timeout_duration) = drain_timeout {
            // Drop our sender so the channel closes once all handler task
            // clones finish sending their responses (or are dropped).
            drop(sink_proxy_tx);
            let drain_result = tokio::time::timeout(timeout_duration, async {
                // First, wait for any response sends already dispatched by the
                // main loop (these hold transport write futures).
                while let Some(result) = response_send_tasks.join_next().await {
                    if let Err(error) = result {
                        tracing::error!(%error, "response send task failed during drain");
                    }
                }
                // Then drain any handler responses still in the channel
                // (handlers that finished after the loop broke).
                while let Some(m) = sink_proxy_rx.recv().await {
                    if let Err(error) = transport.send(m).await {
                        tracing::error!(%error, "failed to send pending response during drain");
                        break;
                    }
                }
            })
            .await;
            if drain_result.is_err() {
                tracing::warn!("timed out draining in-flight responses");
            }
        }

        let sink_close_result = transport.close().await;
        if let Err(e) = sink_close_result {
            tracing::error!(%e, "fail to close sink");
        }
        tracing::info!(?quit_reason, "serve finished");
        quit_reason
    }.instrument(current_span));
    RunningService {
        service,
        peer: peer_return,
        handle: Some(handle),
        cancellation_token: ct.clone(),
        dg: ct.drop_guard(),
    }
}