rmcp 3.0.0-beta.1

use std::sync::OnceLock;

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, ClientCapabilities, Extensions,
        GetExtensions, GetMeta, Implementation, JsonRpcError, JsonRpcMessage, JsonRpcNotification,
        JsonRpcRequest, JsonRpcResponse, NotificationMetaObject, NumberOrString, ProgressToken,
        ProtocolVersion, RequestId, RequestMetaObject,
    },
    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("subscription consumer lagged behind its {capacity}-message buffer")]
    SubscriptionLagged { capacity: usize },
    #[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 },
    /// The peer kept returning `input_required` beyond the configured round cap.
    #[error("input_required did not complete within {max_rounds} MRTR rounds")]
    InputRequiredRoundsExceeded { max_rounds: usize },
}

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<Metadata = RequestMetaObject> + GetExtensions;
    type Resp: TransferObject;
    type Not: TryInto<CancelledNotification, Error = Self::Not>
        + From<CancelledNotification>
        + TransferObject;
    type PeerReq: TransferObject + GetMeta<Metadata = RequestMetaObject> + GetExtensions;
    type PeerResp: TransferObject;
    type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
        + From<CancelledNotification>
        + TransferObject
        + GetMeta<Metadata = NotificationMetaObject>
        + GetExtensions;
    type InitializeError;
    const IS_CLIENT: bool;
    type Info: TransferObject;
    type PeerInfo: TransferObject;
    #[doc(hidden)]
    fn configure_direct_peer(_peer: &Peer<Self>, _info: &Self::Info) {}
    #[doc(hidden)]
    fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> {
        None
    }
    /// Invalidate any response cache affected by an inbound peer notification.
    ///
    /// The serve loop calls this for every notification *before* subscription
    /// routing, so cache invalidation still runs when a notification is
    /// delivered through a `listen` subscription channel rather than the
    /// [`Service::handle_notification`] callbacks.
    #[doc(hidden)]
    fn invalidate_response_cache(
        _peer: &Peer<Self>,
        _notification: &Self::PeerNot,
    ) -> impl Future<Output = ()> + MaybeSendFuture {
        async {}
    }

    #[doc(hidden)]
    fn enforce_request_association(
        _request: &Self::Req,
        _peer_info: Option<&Self::PeerInfo>,
        _in_request_handler_scope: bool,
    ) -> Result<(), ServiceError> {
        Ok(())
    }

    /// Receive-side counterpart of [`Self::enforce_request_association`]:
    /// SEP-2260 says clients receiving a server-to-client request with no
    /// associated outbound request should reject it with invalid params. An
    /// error return is sent back to the peer instead of dispatching to the
    /// handler.
    #[doc(hidden)]
    fn enforce_peer_request_association(
        _peer_request: &Self::PeerReq,
        _peer_info: Option<&Self::PeerInfo>,
        _has_pending_outbound_request: bool,
    ) -> Result<(), McpError> {
        Ok(())
    }
}

pub(crate) fn uses_legacy_lifecycle(
    protocol_version: Option<&ProtocolVersion>,
    uses_discover_lifecycle: bool,
) -> bool {
    !uses_discover_lifecycle
        && protocol_version.is_none_or(|version| version < &ProtocolVersion::V_2026_07_28)
}

tokio::task_local! {
    pub(crate) static ORIGINATING_REQUEST: RequestId;
}

pub(crate) fn in_request_handler_scope() -> bool {
    ORIGINATING_REQUEST.try_with(|_| ()).is_ok()
}

/// Marker in an outbound request's non-serialized [`Extensions`] identifying
/// the in-flight peer request it was issued from (SEP-2260). Attached for both
/// roles whenever a request is sent from within a request handler; the
/// streamable HTTP server reads it to deliver server-initiated requests on the
/// originating request's SSE stream. Never on the wire (SEP-2260 defines no
/// wire field), so session managers that serialize messages between processes
/// lose it and such requests fall back to the standalone stream with a warning.
///
/// # Caller requirements
///
/// From protocol version `2026-07-28`, server-to-client sampling, roots, and
/// elicitation requests must be issued while handling a client request;
/// outside a handler they return an `invalid_request` error. The association
/// is task-local and does not cross `tokio::spawn`, so use the task manager
/// for long-running work.
#[derive(Debug, Clone, PartialEq, Eq)]
#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
pub struct OriginatingRequestId(pub RequestId);

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;

pub(crate) fn remove_pending_request<T>(
    pending_requests: &mut HashMap<RequestId, T>,
    response_id: &RequestId,
) -> Option<T> {
    pending_requests.remove(response_id).or_else(|| {
        response_id
            .numeric_string_value()
            .and_then(|id| pending_requests.remove(&RequestId::Number(id)))
    })
}

#[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<()>>>>;
type SubscriptionChannel<N> = (mpsc::Sender<N>, usize);
type SubscriptionChannelMap<N> = HashMap<RequestId, SubscriptionChannel<N>>;

/// 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;
        self.peer.unregister_subscription(&self.id);
        result
    }

    async fn send_timeout_cancel_notification(&self, reason: &str) {
        self.peer.unregister_subscription(&self.id);
        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;
        self.peer.unregister_subscription(&self.id);
        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>>,
    },
}

#[derive(Debug, Clone)]
pub(crate) struct ClientRequestMetadata {
    pub protocol_version: ProtocolVersion,
    pub client_info: Implementation,
    pub client_capabilities: ClientCapabilities,
}

/// 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`].
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>>>>,
    client_request_metadata: Arc<OnceLock<ClientRequestMetadata>>,
    request_metadata_required: Arc<std::sync::atomic::AtomicBool>,
    subscription_channels: Arc<std::sync::RwLock<SubscriptionChannelMap<R::PeerNot>>>,
    #[cfg(feature = "client")]
    response_cache: client::cache::PeerResponseCache<R>,
}

impl<R: Clone + ServiceRole> Clone for Peer<R>
where
    R::PeerInfo: Clone,
{
    fn clone(&self) -> Peer<R> {
        Self {
            tx: self.tx.clone(),
            request_id_provider: self.request_id_provider.clone(),
            progress_token_provider: self.progress_token_provider.clone(),
            progress_timeout_watchers: self.progress_timeout_watchers.clone(),
            info: self.info.clone(),
            client_request_metadata: self.client_request_metadata.clone(),
            request_metadata_required: self.request_metadata_required.clone(),
            subscription_channels: self.subscription_channels.clone(),
            #[cfg(feature = "client")]
            response_cache: self.response_cache.clone(),
        }
    }
}

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<RequestMetaObject>,
    /// 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()
        }
    }

    /// Adds request metadata while preserving any other configured options.
    ///
    /// Explicit values take precedence over discover-lifecycle metadata defaults.
    pub fn with_meta(mut self, meta: RequestMetaObject) -> Self {
        self.meta = Some(meta);
        self
    }

    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))),
                client_request_metadata: Default::default(),
                request_metadata_required: Default::default(),
                subscription_channels: Default::default(),
                #[cfg(feature = "client")]
                response_cache: Default::default(),
            },
            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,
        request: R::Req,
        options: PeerRequestOptions,
    ) -> Result<RequestHandle<R>, ServiceError> {
        self.send_request_with_option_and_subscription(request, options, None)
            .await
    }

    async fn send_request_with_option_and_subscription(
        &self,
        mut request: R::Req,
        options: PeerRequestOptions,
        subscription_sender: Option<SubscriptionChannel<R::PeerNot>>,
    ) -> Result<RequestHandle<R>, ServiceError> {
        R::enforce_request_association(
            &request,
            self.peer_info().as_deref(),
            in_request_handler_scope(),
        )?;
        if let Ok(originating) = ORIGINATING_REQUEST.try_with(|id| id.clone()) {
            request
                .extensions_mut()
                .insert(OriginatingRequestId(originating));
        }
        let id = self.request_id_provider.next_request_id();
        let progress_token = self.progress_token_provider.next_progress_token();
        if let Some(metadata) = self.client_request_metadata.get() {
            let meta = request.get_meta_mut();
            meta.set_protocol_version(metadata.protocol_version.clone());
            meta.set_client_info(metadata.client_info.clone());
            meta.set_client_capabilities(metadata.client_capabilities.clone());
        }
        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 let Some(channel) = subscription_sender {
            self.subscription_channels_write()
                .insert(id.clone(), channel);
        }
        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);
            }
            self.unregister_subscription(&id);
            return Err(ServiceError::TransportClosed);
        }
        Ok(RequestHandle {
            id,
            rx: receiver,
            progress_token,
            options,
            peer: self.clone(),
            progress_reset_rx,
        })
    }

    #[cfg(feature = "client")]
    pub(crate) async fn send_subscription_request(
        &self,
        request: R::Req,
        options: PeerRequestOptions,
        channel_capacity: usize,
    ) -> Result<(RequestHandle<R>, mpsc::Receiver<R::PeerNot>), ServiceError> {
        let (sender, receiver) = mpsc::channel(channel_capacity);
        let handle = self
            .send_request_with_option_and_subscription(
                request,
                options,
                Some((sender, channel_capacity)),
            )
            .await?;
        Ok((handle, receiver))
    }

    fn subscription_sender(&self, id: &RequestId) -> Option<SubscriptionChannel<R::PeerNot>> {
        self.subscription_channels_read().get(id).cloned()
    }

    pub(crate) fn unregister_subscription(&self, id: &RequestId) {
        self.subscription_channels_write().remove(id);
    }

    fn subscription_channels_read(
        &self,
    ) -> std::sync::RwLockReadGuard<'_, SubscriptionChannelMap<R::PeerNot>> {
        match self.subscription_channels.read() {
            Ok(channels) => channels,
            Err(poisoned) => poisoned.into_inner(),
        }
    }

    fn subscription_channels_write(
        &self,
    ) -> std::sync::RwLockWriteGuard<'_, SubscriptionChannelMap<R::PeerNot>> {
        match self.subscription_channels.write() {
            Ok(channels) => channels,
            Err(poisoned) => poisoned.into_inner(),
        }
    }

    pub(crate) fn try_cancel_request(&self, id: RequestId, reason: Option<String>) {
        let notification = CancelledNotification {
            params: CancelledNotificationParam {
                request_id: Some(id),
                reason,
                meta: None,
            },
            method: crate::model::CancelledNotificationMethod,
            extensions: Default::default(),
        };
        let (responder, _receiver) = tokio::sync::oneshot::channel();
        let _ = self.tx.try_send(PeerSinkMessage::Notification {
            notification: notification.into(),
            responder,
        });
    }

    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(crate) fn set_client_request_metadata(&self, metadata: ClientRequestMetadata) {
        let result = self.client_request_metadata.set(metadata);
        debug_assert!(result.is_ok(), "client request metadata set more than once");
    }

    pub(crate) fn require_request_metadata(&self) {
        self.request_metadata_required
            .store(true, std::sync::atomic::Ordering::Release);
    }

    pub(crate) fn request_metadata_required(&self) -> bool {
        self.request_metadata_required
            .load(std::sync::atomic::Ordering::Acquire)
    }

    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: RequestMetaObject,
    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: RequestMetaObject::default(),
            extensions: Extensions::default(),
            peer,
        }
    }
}

#[cfg(feature = "server")]
impl RequestContext<RoleServer> {
    /// The current request's protocol version, falling back to legacy handshake state.
    pub fn protocol_version(&self) -> Option<crate::model::ProtocolVersion> {
        self.meta.protocol_version().or_else(|| {
            self.peer
                .peer_info()
                .map(|info| info.protocol_version.clone())
        })
    }

    /// The current request's client implementation, falling back only for legacy sessions.
    pub fn client_info(&self) -> Option<Implementation> {
        if self.peer.request_metadata_required() {
            self.meta.client_info()
        } else {
            self.meta
                .client_info()
                .or_else(|| self.peer.peer_info().map(|info| info.client_info.clone()))
        }
    }

    /// The current request's client capabilities, falling back only for legacy sessions.
    pub fn client_capabilities(&self) -> Option<ClientCapabilities> {
        if self.peer.request_metadata_required() {
            self.meta.client_capabilities()
        } else {
            self.meta
                .client_capabilities()
                .or_else(|| self.peer.peer_info().map(|info| info.capabilities.clone()))
        }
    }
}

/// Request execution context
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct NotificationContext<R: ServiceRole> {
    pub meta: NotificationMetaObject,
    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);
    R::configure_direct_peer(&peer, &service.get_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),
            ResponseSendTaskResult(Result<(), tokio::task::JoinError>),
        }

        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)
                            }
                        }
                    }
                    result = response_send_tasks.join_next(), if !response_send_tasks.is_empty() => {
                        Event::ResponseSendTaskResult(
                            result.expect("non-empty response send task set")
                        )
                    }
                    _ = 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(),
                                }));
                            }
                        }
                    }
                }
                Event::ResponseSendTaskResult(result) => {
                    if let Err(error) = result {
                        tracing::error!(%error, "response send task failed");
                    }
                }
                // 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,
                    } {
                        let Some(ct) = local_ct_pool.remove(id) else {
                            tracing::debug!(%id, "dropping response for cancelled request");
                            continue;
                        };
                        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");
                    if let Err(error) = R::enforce_peer_request_association(
                        &request,
                        peer.peer_info().as_deref(),
                        !local_responder_pool.is_empty(),
                    ) {
                        tracing::warn!(%id, message = %error.message, "rejected peer request");
                        // send directly: the sink proxy path would drop the
                        // error since the request was never registered in
                        // local_ct_pool
                        let send = transport.send(JsonRpcMessage::error(error, Some(id)));
                        let current_span = tracing::Span::current();
                        response_send_tasks.spawn(async move {
                            if let Err(error) = send.await {
                                tracing::error!(%error, "fail to send rejection error");
                            }
                        }.instrument(current_span));
                        continue;
                    }
                    {
                        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 = RequestMetaObject::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();
                        let handler_id = id.clone();
                        spawn_service_task(async move {
                            let result = ORIGINATING_REQUEST
                                .scope(handler_id, 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");
                    R::invalidate_response_cache(&peer, &notification).await;
                    let cancellation_request_id =
                        if let Some(cancelled) = R::peer_cancelled_params(&notification) {
                            let request_id = cancelled.request_id.clone();
                            if let Some(request_id) = request_id.as_ref() {
                                if R::IS_CLIENT {
                                    if let Some(responder) =
                                        local_responder_pool.remove(request_id)
                                    {
                                        let _ = responder.send(Err(ServiceError::Cancelled {
                                            reason: cancelled.reason.clone(),
                                        }));
                                    }
                                } else if let Some(ct) = local_ct_pool.remove(request_id) {
                                    tracing::info!(id = %request_id, reason = cancelled.reason, "cancelled");
                                    ct.cancel();
                                }
                            }
                            request_id
                        } else {
                            None
                        };
                    let subscription_id = notification
                        .get_meta()
                        .subscription_id()
                        .or(cancellation_request_id);
                    if let Some(subscription_id) = subscription_id
                        && let Some((sender, capacity)) =
                            peer.subscription_sender(&subscription_id)
                    {
                        match sender.try_send(notification) {
                            Ok(()) => {}
                            Err(mpsc::error::TrySendError::Full(_)) => {
                                tracing::warn!(
                                    id = %subscription_id,
                                    capacity,
                                    "subscription notification buffer full"
                                );
                                if R::IS_CLIENT
                                    && let Some(responder) =
                                        local_responder_pool.remove(&subscription_id)
                                {
                                    let _ = responder
                                        .send(Err(ServiceError::SubscriptionLagged { capacity }));
                                }
                                peer.unregister_subscription(&subscription_id);
                                peer.try_cancel_request(
                                    subscription_id,
                                    Some("subscription notification buffer full".to_owned()),
                                );
                            }
                            Err(mpsc::error::TrySendError::Closed(_)) => {
                                peer.unregister_subscription(&subscription_id);
                                peer.try_cancel_request(
                                    subscription_id,
                                    Some("subscription notification receiver closed".to_owned()),
                                );
                            }
                        }
                        continue;
                    }
                    let mut 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 = NotificationMetaObject::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) =
                        remove_pending_request(&mut local_responder_pool, &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) =
                        remove_pending_request(&mut local_responder_pool, &id)
                    {
                        let service_error = if error.is_transport_closed() {
                            ServiceError::TransportClosed
                        } else {
                            ServiceError::McpError(error)
                        };
                        let _response_result = responder.send(Err(service_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(),
    }
}

#[cfg(all(test, feature = "server"))]
mod sep2260_marker_tests {
    use std::sync::Arc;

    use super::*;
    use crate::model::{PingRequest, RequestId, ServerRequest};

    fn ping() -> ServerRequest {
        ServerRequest::PingRequest(PingRequest {
            method: Default::default(),
            extensions: Default::default(),
        })
    }

    async fn send_and_capture(scope: Option<RequestId>) -> <RoleServer as ServiceRole>::Req {
        // peer_info None keeps enforcement non-strict; only the sink message matters.
        let (peer, mut rx) =
            Peer::<RoleServer>::new(Arc::new(AtomicU32RequestIdProvider::default()), None);
        let send = peer.send_request_with_option(ping(), PeerRequestOptions::no_options());
        let _handle = match scope {
            Some(id) => ORIGINATING_REQUEST.scope(id, send).await.unwrap(),
            None => send.await.unwrap(),
        };
        let PeerSinkMessage::Request { request, .. } = rx.recv().await.expect("sink message")
        else {
            panic!("expected a request sink message");
        };
        request
    }

    #[tokio::test]
    async fn outbound_request_carries_originating_id_when_in_scope() {
        let request = send_and_capture(Some(RequestId::Number(7))).await;
        let marker = request
            .extensions()
            .get::<OriginatingRequestId>()
            .expect("marker attached");
        assert_eq!(marker.0, RequestId::Number(7));
    }

    #[tokio::test]
    async fn outbound_request_has_no_marker_outside_scope() {
        let request = send_and_capture(None).await;
        assert!(request.extensions().get::<OriginatingRequestId>().is_none());
    }
}