Move whole rmcp crate to offcial rust sdk (#44)

* get rid of async-trait

* move rmcp to official-rust-sdk

* replace master with main

* update readme, change license

* fix typo

* Update README.md

Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>

* fix typos

---------

Co-authored-by: = <=>
Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>

crates/rmcp/src/service.rs

@@ -0,0 +1,674 @@
+use crate::error::Error as McpError;
+use crate::model::{
+    CancelledNotification, CancelledNotificationParam, JsonRpcMessage, Message, RequestId,
+};
+use crate::transport::IntoTransport;
+use futures::future::BoxFuture;
+use thiserror::Error;
+#[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;
+#[cfg(feature = "tower")]
+pub use tower::*;
+
+use tokio_util::sync::CancellationToken;
+use tracing::instrument;
+#[derive(Error, Debug)]
+#[non_exhaustive]
+pub enum ServiceError {
+    #[error("Mcp error: {0}")]
+    McpError(McpError),
+    #[error("Transport error: {0}")]
+    Transport(std::io::Error),
+    #[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 },
+}
+
+impl ServiceError {}
+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;
+    type Resp: TransferObject;
+    type Not: TryInto<CancelledNotification, Error = Self::Not>
+        + From<CancelledNotification>
+        + TransferObject;
+    type PeerReq: TransferObject;
+    type PeerResp: TransferObject;
+    type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
+        + From<CancelledNotification>
+        + TransferObject;
+    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,
+>;
+
+pub type TxMessage<R> =
+    Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
+pub type RxMessage<R> =
+    Message<<R as ServiceRole>::PeerReq, <R as ServiceRole>::PeerResp, <R as ServiceRole>::PeerNot>;
+
+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>> + Send + '_;
+    fn handle_notification(
+        &self,
+        notification: R::PeerNot,
+    ) -> impl Future<Output = Result<(), McpError>> + Send + '_;
+    fn get_peer(&self) -> Option<Peer<R>>;
+    fn set_peer(&mut self, peer: Peer<R>);
+    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>, E>> + Send
+    where
+        T: IntoTransport<R, E, A>,
+        E: std::error::Error + From<std::io::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>, E>> + Send
+    where
+        T: IntoTransport<R, E, A>,
+        E: std::error::Error + From<std::io::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>> + Send + '_ {
+        DynService::handle_request(self.as_ref(), request, context)
+    }
+
+    fn handle_notification(
+        &self,
+        notification: R::PeerNot,
+    ) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
+        DynService::handle_notification(self.as_ref(), notification)
+    }
+
+    fn get_peer(&self) -> Option<Peer<R>> {
+        DynService::get_peer(self.as_ref())
+    }
+
+    fn set_peer(&mut self, peer: Peer<R>) {
+        DynService::set_peer(self.as_mut(), peer)
+    }
+
+    fn get_info(&self) -> R::Info {
+        DynService::get_info(self.as_ref())
+    }
+}
+
+pub trait DynService<R: ServiceRole>: Send + Sync {
+    fn handle_request(
+        &self,
+        request: R::PeerReq,
+        context: RequestContext<R>,
+    ) -> BoxFuture<Result<R::Resp, McpError>>;
+    fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>>;
+    fn get_peer(&self) -> Option<Peer<R>>;
+    fn set_peer(&mut self, peer: Peer<R>);
+    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>,
+    ) -> BoxFuture<Result<R::Resp, McpError>> {
+        Box::pin(self.handle_request(request, context))
+    }
+    fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>> {
+        Box::pin(self.handle_notification(notification))
+    }
+    fn get_peer(&self) -> Option<Peer<R>> {
+        self.get_peer()
+    }
+    fn set_peer(&mut self, peer: Peer<R>) {
+        self.set_peer(peer)
+    }
+    fn get_info(&self) -> R::Info {
+        self.get_info()
+    }
+}
+
+use std::collections::HashMap;
+use std::ops::Deref;
+use std::sync::Arc;
+use std::sync::atomic::AtomicU32;
+use std::time::Duration;
+
+use tokio::sync::mpsc;
+
+pub trait RequestIdProvider: Send + Sync + 'static {
+    fn next_request_id(&self) -> RequestId;
+}
+
+#[derive(Debug, Default)]
+pub struct AtomicU32RequestIdProvider {
+    id: AtomicU32,
+}
+
+impl RequestIdProvider for AtomicU32RequestIdProvider {
+    fn next_request_id(&self) -> RequestId {
+        RequestId::Number(self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
+    }
+}
+
+type Responder<T> = tokio::sync::oneshot::Sender<T>;
+
+/// 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)]
+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,
+}
+
+impl<R: ServiceRole> RequestHandle<R> {
+    pub const REQUEST_TIMEOUT_REASON: &str = "request timeout";
+    pub async fn await_response(self) -> Result<R::PeerResp, ServiceError> {
+        if let Some(timeout) = self.options.timeout {
+            let timeout_result = tokio::time::timeout(timeout, async move {
+                self.rx
+                    .await
+                    .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
+            })
+            .await;
+            match timeout_result {
+                Ok(response) => response,
+                Err(_) => {
+                    let error = Err(ServiceError::Timeout { timeout });
+                    // cancel this request
+                    let notification = CancelledNotification {
+                        params: CancelledNotificationParam {
+                            request_id: self.id,
+                            reason: Some(Self::REQUEST_TIMEOUT_REASON.to_owned()),
+                        },
+                        method: crate::model::CancelledNotificationMethod,
+                    };
+                    let _ = self.peer.send_notification(notification.into()).await;
+                    error
+                }
+            }
+        } else {
+            self.rx
+                .await
+                .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
+        }
+    }
+
+    /// Cancel this request
+    pub async fn cancel(self, reason: Option<String>) -> Result<(), ServiceError> {
+        let notification = CancelledNotification {
+            params: CancelledNotificationParam {
+                request_id: self.id,
+                reason,
+            },
+            method: crate::model::CancelledNotificationMethod,
+        };
+        self.peer.send_notification(notification.into()).await?;
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+pub enum PeerSinkMessage<R: ServiceRole> {
+    Request(
+        R::Req,
+        RequestId,
+        Responder<Result<R::PeerResp, ServiceError>>,
+    ),
+    Notification(R::Not, 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_cancellable_request`].
+#[derive(Clone)]
+pub struct Peer<R: ServiceRole> {
+    tx: mpsc::Sender<PeerSinkMessage<R>>,
+    request_id_provider: Arc<dyn RequestIdProvider>,
+    info: 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)]
+pub struct PeerRequestOptions {
+    timeout: Option<Duration>,
+}
+
+impl PeerRequestOptions {
+    pub fn no_options() -> Self {
+        Self::default()
+    }
+}
+
+impl<R: ServiceRole> Peer<R> {
+    const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
+    pub fn new(
+        request_id_provider: Arc<dyn RequestIdProvider>,
+        peer_info: R::PeerInfo,
+    ) -> (Peer<R>, ProxyOutbound<R>) {
+        let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
+        (
+            Self {
+                tx,
+                request_id_provider,
+                info: peer_info.into(),
+            },
+            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::Transport(std::io::Error::other("disconnected")))?;
+        receiver
+            .await
+            .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
+    }
+    pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
+        self.send_cancellable_request(request, PeerRequestOptions::no_options())
+            .await?
+            .await_response()
+            .await
+    }
+    pub async fn send_cancellable_request(
+        &self,
+        request: R::Req,
+        options: PeerRequestOptions,
+    ) -> Result<RequestHandle<R>, ServiceError> {
+        let id = self.request_id_provider.next_request_id();
+        let (responder, receiver) = tokio::sync::oneshot::channel();
+        self.tx
+            .send(PeerSinkMessage::Request(request, id.clone(), responder))
+            .await
+            .map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
+        Ok(RequestHandle {
+            id,
+            rx: receiver,
+            options,
+            peer: self.clone(),
+        })
+    }
+    pub fn peer_info(&self) -> &R::PeerInfo {
+        &self.info
+    }
+}
+
+#[derive(Debug)]
+pub struct RunningService<R: ServiceRole, S: Service<R>> {
+    service: Arc<S>,
+    peer: Peer<R>,
+    handle: tokio::task::JoinHandle<QuitReason>,
+    /// cancellation token
+    ct: CancellationToken,
+}
+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()
+    }
+    pub async fn waiting(self) -> Result<QuitReason, tokio::task::JoinError> {
+        self.handle.await
+    }
+    pub async fn cancel(self) -> Result<QuitReason, tokio::task::JoinError> {
+        self.ct.cancel();
+        self.waiting().await
+    }
+}
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub enum QuitReason {
+    Cancelled,
+    Closed,
+}
+
+/// Request execution context
+#[derive(Debug, Clone)]
+pub struct RequestContext<R: ServiceRole> {
+    /// this token will be cancelled when the [`CancelledNotification`] is received.
+    pub ct: CancellationToken,
+    pub id: RequestId,
+    /// An interface to fetch the remote client or server
+    pub peer: Peer<R>,
+}
+
+/// Use this function to skip initialization process
+pub async fn serve_directly<R, S, T, E, A>(
+    service: S,
+    transport: T,
+    peer_info: R::PeerInfo,
+) -> Result<RunningService<R, S>, E>
+where
+    R: ServiceRole,
+    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()).await
+}
+
+/// Use this function to skip initialization process
+pub async fn serve_directly_with_ct<R, S, T, E, A>(
+    service: S,
+    transport: T,
+    peer_info: R::PeerInfo,
+    ct: CancellationToken,
+) -> Result<RunningService<R, S>, E>
+where
+    R: ServiceRole,
+    S: Service<R>,
+    T: IntoTransport<R, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    serve_inner(service, transport, peer_info, Default::default(), ct).await
+}
+
+#[instrument(skip_all)]
+async fn serve_inner<R, S, T, E, A>(
+    mut service: S,
+    transport: T,
+    peer_info: R::PeerInfo,
+    id_provider: Arc<AtomicU32RequestIdProvider>,
+    ct: CancellationToken,
+) -> Result<RunningService<R, S>, E>
+where
+    R: ServiceRole,
+    S: Service<R>,
+    T: IntoTransport<R, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    use futures::{SinkExt, StreamExt};
+    const SINK_PROXY_BUFFER_SIZE: usize = 64;
+    let (sink_proxy_tx, mut sink_proxy_rx) = tokio::sync::mpsc::channel::<
+        Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>,
+    >(SINK_PROXY_BUFFER_SIZE);
+
+    if R::IS_CLIENT {
+        tracing::info!(?peer_info, "Service initialized as client");
+    } else {
+        tracing::info!(?peer_info, "Service initialized as server");
+    }
+
+    let (peer, mut peer_proxy) = <Peer<R>>::new(id_provider, peer_info);
+    service.set_peer(peer.clone());
+    let mut local_responder_pool = HashMap::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 handle = tokio::spawn(async move {
+        let (mut sink, mut stream) = transport.into_transport();
+        let mut sink = std::pin::pin!(sink);
+        let mut stream = std::pin::pin!(stream);
+        #[derive(Debug)]
+        enum Event<P, R, T> {
+            ProxyMessage(P),
+            PeerMessage(R),
+            ToSink(T),
+        }
+        let quit_reason = loop {
+            let evt = tokio::select! {
+                m = sink_proxy_rx.recv() => {
+                    if let Some(m) = m {
+                        Event::ToSink(m)
+                    } else {
+                        continue
+                    }
+                }
+                m = stream.next() => {
+                    if let Some(m) = m {
+                        Event::PeerMessage(m.into_message())
+                    } else {
+                        // input stream closed
+                        tracing::info!("input stream terminated");
+                        break QuitReason::Closed
+                    }
+                }
+                m = peer_proxy.recv() => {
+                    if let Some(m) = m {
+                        Event::ProxyMessage(m)
+                    } else {
+                        continue
+                    }
+                }
+                _ = serve_loop_ct.cancelled() => {
+                    tracing::info!("task cancelled");
+                    break QuitReason::Cancelled
+                }
+            };
+            tracing::debug!(?evt, "new event");
+            match evt {
+                // response and error
+                Event::ToSink(e) => {
+                    if let Some(id) = match &e {
+                        Message::Response(_, id) => Some(id),
+                        Message::Error(_, id) => Some(id),
+                        _ => None,
+                    } {
+                        if let Some(ct) = local_ct_pool.remove(id) {
+                            ct.cancel();
+                        }
+                        let send_result = sink.send(e.into_json_rpc_message()).await;
+                        if let Err(error) = send_result {
+                            tracing::error!(%error, "fail to response message");
+                        }
+                    }
+                }
+                Event::ProxyMessage(PeerSinkMessage::Request(request, id, responder)) => {
+                    local_responder_pool.insert(id.clone(), responder);
+                    let send_result = sink
+                        .send(Message::Request(request, id.clone()).into_json_rpc_message())
+                        .await;
+                    if let Err(e) = send_result {
+                        if let Some(responder) = local_responder_pool.remove(&id) {
+                            let _ = responder
+                                .send(Err(ServiceError::Transport(std::io::Error::other(e))));
+                        }
+                    }
+                }
+                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_result = sink
+                        .send(Message::Notification(notification).into_json_rpc_message())
+                        .await;
+                    if let Err(e) = send_result {
+                        let _ =
+                            responder.send(Err(ServiceError::Transport(std::io::Error::other(e))));
+                    }
+                    if let Some(param) = cancellation_param {
+                        if let Some(responder) = local_responder_pool.remove(&param.request_id) {
+                            tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
+                            let _response_result = responder.send(Err(ServiceError::Cancelled {
+                                reason: param.reason.clone(),
+                            }));
+                        }
+                    }
+                }
+                Event::PeerMessage(Message::Request(request, id)) => {
+                    tracing::info!(%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 context = RequestContext {
+                            ct: context_ct,
+                            id: id.clone(),
+                            peer: peer.clone(),
+                        };
+                        tokio::spawn(async move {
+                            let result = service.handle_request(request, context).await;
+                            let response = match result {
+                                Ok(result) => {
+                                    tracing::info!(%id, ?result, "response message");
+                                    Message::Response(result, id)
+                                }
+                                Err(error) => {
+                                    tracing::warn!(%id, ?error, "response error");
+                                    Message::Error(error, id)
+                                }
+                            };
+                            let _send_result = sink.send(response).await;
+                        });
+                    }
+                }
+                Event::PeerMessage(Message::Notification(notification)) => {
+                    tracing::info!(?notification, "received notification");
+                    // catch cancelled notification
+                    let notification = match notification.try_into() {
+                        Ok::<CancelledNotification, _>(cancelled) => {
+                            if let Some(ct) = local_ct_pool.remove(&cancelled.params.request_id) {
+                                tracing::info!(id = %cancelled.params.request_id, reason = cancelled.params.reason, "cancelled");
+                                ct.cancel();
+                            }
+                            cancelled.into()
+                        }
+                        Err(notification) => notification,
+                    };
+                    {
+                        let service = shared_service.clone();
+                        tokio::spawn(async move {
+                            let result = service.handle_notification(notification).await;
+                            if let Err(error) = result {
+                                tracing::warn!(%error, "Error sending notification");
+                            }
+                        });
+                    }
+                }
+                Event::PeerMessage(Message::Response(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(Message::Error(error, id)) => {
+                    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");
+                        }
+                    }
+                }
+            }
+        };
+        let sink_close_result = sink.close().await;
+        if let Err(e) = sink_close_result {
+            tracing::error!(%e, "fail to close sink");
+        }
+        tracing::info!(?quit_reason, "serve finished");
+        quit_reason
+    });
+    Ok(RunningService {
+        service,
+        peer: peer_return,
+        handle,
+        ct,
+    })
+}