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

4t145 · web-flow · commit c1e282e50db3 · 2025-05-18T18:54:44.000+08:00
diff --git a/crates/rmcp/src/handler/client.rs b/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()
     }
diff --git a/crates/rmcp/src/handler/server.rs b/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()
     }
diff --git a/crates/rmcp/src/service.rs b/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}")]
-    Transport(std::io::Error),
+    #[error("Transport send error: {0}")]
+    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
-                    .await
-                    .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
+                self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
             })
             .await;
             match timeout_result {
@@ -278,9 +260,7 @@ impl<R: ServiceRole> RequestHandle<R> {
                 }
             }
         } else {
-            self.rx
-                .await
-                .map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
+            self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
         }
     }
 
@@ -373,12 +353,8 @@ impl<R: ServiceRole> Peer<R> {
                 responder,
             })
             .await
-            .map_err(|_m| {
-                ServiceError::Transport(std::io::Error::other("disconnected: receiver dropped"))
-            })?;
-        receiver.await.map_err(|_e| {
-            ServiceError::Transport(std::io::Error::other("disconnected: responder dropped"))
-        })?
+            .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())
@@ -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
-                                .send(Err(ServiceError::Transport(std::io::Error::other(e))));
+                            let _ = responder.send(Err(ServiceError::TransportSend(Box::new(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(())
                     };
diff --git a/crates/rmcp/src/service/tower.rs b/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()
     }

Original file line number	Diff line number	Diff line change
`@@ -53,14 +53,6 @@ impl<H: ClientHandler> Service<RoleClient> for H {`
`53`	`53`	`Ok(())`
`54`	`54`	`}`
`55`	`55`
`56`		`- fn get_peer(&self) -> Option<Peer<RoleClient>> {`
`57`		`- self.get_peer()`
`58`		`- }`
`59`		`-`
`60`		`- fn set_peer(&mut self, peer: Peer<RoleClient>) {`
`61`		`- self.set_peer(peer);`
`62`		`- }`
`63`		`-`
`64`	`56`	`fn get_info(&self) -> <RoleClient as ServiceRole>::Info {`
`65`	`57`	`self.get_info()`
`66`	`58`	`}`
Original file line number	Diff line number	Diff line change
`@@ -89,14 +89,6 @@ impl<H: ServerHandler> Service<RoleServer> for H {`
`89`	`89`	`Ok(())`
`90`	`90`	`}`
`91`	`91`
`92`		`- fn get_peer(&self) -> Option<Peer<RoleServer>> {`
`93`		`- self.get_peer()`
`94`		`- }`
`95`		`-`
`96`		`- fn set_peer(&mut self, peer: Peer<RoleServer>) {`
`97`		`- self.set_peer(peer);`
`98`		`- }`
`99`		`-`
`100`	`92`	`fn get_info(&self) -> <RoleServer as ServiceRole>::Info {`
`101`	`93`	`self.get_info()`
`102`	`94`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2,12 +2,11 @@ use std::{future::poll_fn, marker::PhantomData};`
`2`	`2`
`3`	`3`	`use tower_service::Service as TowerService;`
`4`	`4`
`5`		`-use crate::service::{Peer, RequestContext, Service, ServiceRole};`
	`5`	`+use crate::service::{RequestContext, Service, ServiceRole};`
`6`	`6`
`7`	`7`	`pub struct TowerHandler<S, R: ServiceRole> {`
`8`	`8`	`pub service: S,`
`9`	`9`	`pub info: R::Info,`
`10`		`- pub peer: Option<Peer<R>>,`
`11`	`10`	`role: PhantomData<R>,`
`12`	`11`	`}`
`13`	`12`
`@@ -17,7 +16,6 @@ impl<S, R: ServiceRole> TowerHandler<S, R> {`
`17`	`16`	`service,`
`18`	`17`	`role: PhantomData,`
`19`	`18`	`info,`
`20`		`- peer: None,`
`21`	`19`	`}`
`22`	`20`	`}`
`23`	`21`	`}`
`@@ -48,14 +46,6 @@ where`
`48`	`46`	`std::future::ready(Ok(()))`
`49`	`47`	`}`
`50`	`48`
`51`		`- fn get_peer(&self) -> Option<Peer<R>> {`
`52`		`- self.peer.clone()`
`53`		`- }`
`54`		`-`
`55`		`- fn set_peer(&mut self, peer: Peer<R>) {`
`56`		`- self.peer = Some(peer);`
`57`		`- }`
`58`		`-`
`59`	`49`	`fn get_info(&self) -> R::Info {`
`60`	`50`	`self.info.clone()`
`61`	`51`	`}`