fix: cancellable initialization process (#280)

Author: 4t145

crates/rmcp/src/service/client.rs

@@ -39,6 +39,9 @@ pub enum ClientInitializeError<E> {
         error: E,
         context: Cow<'static, str>,
     },
+
+    #[error("Cancelled")]
+    Cancelled,
 }
 
 /// Helper function to get the next message from the stream
@@ -121,6 +124,24 @@ pub async fn serve_client_with_ct<S, T, E, A>(
     transport: T,
     ct: CancellationToken,
 ) -> Result<RunningService<RoleClient, S>, ClientInitializeError<E>>
+where
+    S: Service<RoleClient>,
+    T: IntoTransport<RoleClient, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    tokio::select! {
+        result = serve_client_with_ct_inner(service, transport, ct.clone()) => { result }
+        _ = ct.cancelled() => {
+            Err(ClientInitializeError::Cancelled)
+        }
+    }
+}
+
+async fn serve_client_with_ct_inner<S, T, E, A>(
+    service: S,
+    transport: T,
+    ct: CancellationToken,
+) -> Result<RunningService<RoleClient, S>, ClientInitializeError<E>>
 where
     S: Service<RoleClient>,
     T: IntoTransport<RoleClient, E, A>,

crates/rmcp/src/service/server.rs

@@ -58,6 +58,9 @@ pub enum ServerInitializeError<E> {
         error: E,
         context: Cow<'static, str>,
     },
+
+    #[error("Cancelled")]
+    Cancelled,
 }
 
 pub type ClientSink = Peer<RoleServer>;
@@ -140,6 +143,24 @@ pub async fn serve_server_with_ct<S, T, E, A>(
     transport: T,
     ct: CancellationToken,
 ) -> Result<RunningService<RoleServer, S>, ServerInitializeError<E>>
+where
+    S: Service<RoleServer>,
+    T: IntoTransport<RoleServer, E, A>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    tokio::select! {
+        result = serve_server_with_ct_inner(service, transport, ct.clone()) => { result }
+        _ = ct.cancelled() => {
+            Err(ServerInitializeError::Cancelled)
+        }
+    }
+}
+
+async fn serve_server_with_ct_inner<S, T, E, A>(
+    service: S,
+    transport: T,
+    ct: CancellationToken,
+) -> Result<RunningService<RoleServer, S>, ServerInitializeError<E>>
 where
     S: Service<RoleServer>,
     T: IntoTransport<RoleServer, E, A>,