fix(client): add error enum while deal client info

1. wrap the error type for more standardized
2. add more information in error for debug trace
3. wrap helper func for more user-friendly code

Signed-off-by: jokemanfire <[email protected]>

Author: jokemanfire

crates/rmcp/src/service/client.rs

@@ -1,20 +1,73 @@
-use futures::{SinkExt, StreamExt};
+use futures::{SinkExt, Stream, StreamExt};
+use thiserror::Error;
 
 use super::*;
 use crate::model::{
     CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
     CancelledNotificationParam, ClientInfo, ClientMessage, ClientNotification, ClientRequest,
     ClientResult, CompleteRequest, CompleteRequestParam, CompleteResult, GetPromptRequest,
     GetPromptRequestParam, GetPromptResult, InitializeRequest, InitializedNotification,
-    ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
+    JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
     ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
     ListToolsResult, PaginatedRequestParam, PaginatedRequestParamInner, ProgressNotification,
     ProgressNotificationParam, ReadResourceRequest, ReadResourceRequestParam, ReadResourceResult,
-    RootsListChangedNotification, ServerInfo, ServerNotification, ServerRequest, ServerResult,
-    SetLevelRequest, SetLevelRequestParam, SubscribeRequest, SubscribeRequestParam,
-    UnsubscribeRequest, UnsubscribeRequestParam,
+    RequestId, RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification,
+    ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParam, SubscribeRequest,
+    SubscribeRequestParam, UnsubscribeRequest, UnsubscribeRequestParam,
 };
 
+/// It represents the error that may occur when serving the client.
+///
+/// if you want to handle the error, you can use `serve_client_with_ct` or `serve_client` with `Result<RunningService<RoleClient, S>, ClientError>`
+#[derive(Error, Debug)]
+pub enum ClientError {
+    #[error("expect initialized response, but received: {0:?}")]
+    ExpectedInitResponse(Option<ServerJsonRpcMessage>),
+
+    #[error("expect initialized result, but received: {0:?}")]
+    ExpectedInitResult(Option<ServerResult>),
+
+    #[error("conflict initialized response id: expected {0}, got {1}")]
+    ConflictInitResponseId(RequestId, RequestId),
+
+    #[error("connection closed: {0}")]
+    ConnectionClosed(String),
+
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+}
+
+/// Helper function to get the next message from the stream
+async fn expect_next_message<S>(
+    stream: &mut S,
+    context: &str,
+) -> Result<ServerJsonRpcMessage, ClientError>
+where
+    S: Stream<Item = ServerJsonRpcMessage> + Unpin,
+{
+    stream
+        .next()
+        .await
+        .ok_or_else(|| ClientError::ConnectionClosed(context.to_string()))
+        .map_err(|e| ClientError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))
+}
+
+/// Helper function to expect a response from the stream
+async fn expect_response<S>(
+    stream: &mut S,
+    context: &str,
+) -> Result<(ServerResult, RequestId), ClientError>
+where
+    S: Stream<Item = ServerJsonRpcMessage> + Unpin,
+{
+    let msg = expect_next_message(stream, context).await?;
+
+    match msg {
+        ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => Ok((result, id)),
+        _ => Err(ClientError::ExpectedInitResponse(Some(msg))),
+    }
+}
+
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
 pub struct RoleClient;
 
@@ -74,6 +127,15 @@ where
     let mut sink = Box::pin(sink);
     let mut stream = Box::pin(stream);
     let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
+
+    // Convert ClientError to std::io::Error, then to E
+    let handle_client_error = |e: ClientError| -> E {
+        match e {
+            ClientError::Io(io_err) => io_err.into(),
+            other => std::io::Error::new(std::io::ErrorKind::Other, format!("{}", other)).into(),
+        }
+    };
+
     // service
     let id = id_provider.next_request_id();
     let init_request = InitializeRequest {
@@ -85,34 +147,24 @@ where
             .into_json_rpc_message(),
     )
     .await?;
-    let (response, response_id) = stream
-        .next()
+
+    let (response, response_id) = expect_response(&mut stream, "initialize response")
         .await
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::UnexpectedEof,
-            "expect initialize response",
-        ))?
-        .into_message()
-        .into_result()
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize result",
-        ))?;
+        .map_err(handle_client_error)?;
+
     if id != response_id {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "conflict initialize response id",
-        )
-        .into());
+        return Err(handle_client_error(ClientError::ConflictInitResponseId(
+            id,
+            response_id,
+        )));
     }
-    let response = response.map_err(std::io::Error::other)?;
+
     let ServerResult::InitializeResult(initialize_result) = response else {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize result",
-        )
-        .into());
+        return Err(handle_client_error(ClientError::ExpectedInitResult(Some(
+            response,
+        ))));
     };
+
     // send notification
     let notification = ClientMessage::Notification(ClientNotification::InitializedNotification(
         InitializedNotification {