fix: more friendly interface to get service error (#190)

crates/rmcp/src/handler/client.rs

@@ -53,14 +53,6 @@ impl<H: ClientHandler> Service<RoleClient> for H {
         Ok(())
     }
 
-    fn get_peer(&self) -> Option<Peer<RoleClient>> {
-        self.get_peer()
-    }
-
-    fn set_peer(&mut self, peer: Peer<RoleClient>) {
-        self.set_peer(peer);
-    }
-
     fn get_info(&self) -> <RoleClient as ServiceRole>::Info {
         self.get_info()
     }

crates/rmcp/src/handler/server.rs

@@ -89,14 +89,6 @@ impl<H: ServerHandler> Service<RoleServer> for H {
         Ok(())
     }
 
-    fn get_peer(&self) -> Option<Peer<RoleServer>> {
-        self.get_peer()
-    }
-
-    fn set_peer(&mut self, peer: Peer<RoleServer>) {
-        self.set_peer(peer);
-    }
-
     fn get_info(&self) -> <RoleServer as ServiceRole>::Info {
         self.get_info()
     }

crates/rmcp/src/service.rs

@@ -36,8 +36,10 @@ use tracing::instrument;
 pub enum ServiceError {
     #[error("Mcp error: {0}")]
     McpError(McpError),
-    #[error("Transport error: {0}")]
+    #[error("Transport send error: {0}")]
-    Transport(std::io::Error),
+    TransportSend(Box<dyn std::error::Error + Send + Sync>),
+    #[error("Transport closed")]
+    TransportClosed,
     #[error("Unexpected response type")]
     UnexpectedResponse,
     #[error("task cancelled for reason {}", reason.as_deref().unwrap_or("<unknown>"))]
@@ -98,8 +100,6 @@ pub trait Service<R: ServiceRole>: Send + Sync + 'static {
         &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;
 }
 
@@ -148,14 +148,6 @@ impl<R: ServiceRole> Service<R> for Box<dyn DynService<R>> {
         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())
     }
@@ -168,8 +160,6 @@ pub trait DynService<R: ServiceRole>: Send + Sync {
         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;
 }
 
@@ -184,12 +174,6 @@ impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
     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()
     }
@@ -255,9 +239,7 @@ impl<R: ServiceRole> RequestHandle<R> {
     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
+                self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
-                    .await
-                    .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
             })
             .await;
             match timeout_result {
@@ -278,9 +260,7 @@ impl<R: ServiceRole> RequestHandle<R> {
                 }
             }
         } else {
-            self.rx
+            self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
-                .await
-                .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
         }
     }
 
@@ -373,12 +353,8 @@ impl<R: ServiceRole> Peer<R> {
                 responder,
             })
             .await
-            .map_err(|_m| {
+            .map_err(|_m| ServiceError::TransportClosed)?;
-                ServiceError::Transport(std::io::Error::other("disconnected: receiver dropped"))
+        receiver.await.map_err(|_e| ServiceError::TransportClosed)?
-            })?;
-        receiver.await.map_err(|_e| {
-            ServiceError::Transport(std::io::Error::other("disconnected: responder dropped"))
-        })?
     }
     pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
         self.send_request_with_option(request, PeerRequestOptions::no_options())
@@ -416,7 +392,7 @@ impl<R: ServiceRole> Peer<R> {
                 responder,
             })
             .await
-            .map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
+            .map_err(|_m| ServiceError::TransportClosed)?;
         Ok(RequestHandle {
             id,
             rx: receiver,
@@ -428,6 +404,10 @@ impl<R: ServiceRole> Peer<R> {
     pub fn peer_info(&self) -> &R::PeerInfo {
         &self.info
     }
+
+    pub fn is_transport_closed(&self) -> bool {
+        self.tx.is_closed()
+    }
 }
 
 #[derive(Debug)]
@@ -518,7 +498,7 @@ where
 
 #[instrument(skip_all)]
 async fn serve_inner<R, S, T, E, A>(
-    mut service: S,
+    service: S,
     transport: T,
     peer: Peer<R>,
     mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
@@ -540,7 +520,6 @@ where
         tracing::info!(?peer_info, "Service initialized as server");
     }
 
-    service.set_peer(peer.clone());
     let mut local_responder_pool =
         HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::new();
     let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
@@ -631,8 +610,7 @@ where
                 Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
                     if let Err(e) = result {
                         if let Some(responder) = local_responder_pool.remove(&id) {
-                            let _ = responder
+                            let _ = responder.send(Err(ServiceError::TransportSend(Box::new(e))));
-                                .send(Err(ServiceError::Transport(std::io::Error::other(e))));
                         }
                     }
                 }
@@ -642,7 +620,7 @@ where
                     cancellation_param,
                 }) => {
                     let response = if let Err(e) = result {
-                        Err(ServiceError::Transport(std::io::Error::other(e)))
+                        Err(ServiceError::TransportSend(Box::new(e)))
                     } else {
                         Ok(())
                     };

crates/rmcp/src/service/tower.rs

@@ -2,12 +2,11 @@ use std::{future::poll_fn, marker::PhantomData};
 
 use tower_service::Service as TowerService;
 
-use crate::service::{Peer, RequestContext, Service, ServiceRole};
+use crate::service::{RequestContext, Service, ServiceRole};
 
 pub struct TowerHandler<S, R: ServiceRole> {
     pub service: S,
     pub info: R::Info,
-    pub peer: Option<Peer<R>>,
     role: PhantomData<R>,
 }
 
@@ -17,7 +16,6 @@ impl<S, R: ServiceRole> TowerHandler<S, R> {
             service,
             role: PhantomData,
             info,
-            peer: None,
         }
     }
 }
@@ -48,14 +46,6 @@ where
         std::future::ready(Ok(()))
     }
 
-    fn get_peer(&self) -> Option<Peer<R>> {
-        self.peer.clone()
-    }
-
-    fn set_peer(&mut self, peer: Peer<R>) {
-        self.peer = Some(peer);
-    }
-
     fn get_info(&self) -> R::Info {
         self.info.clone()
     }