server: Add error enum while deal server 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 <hu.dingyang@zte.com.cn>

Author: jokemanfire

crates/rmcp/src/service/server.rs

@@ -1,13 +1,7 @@
 use crate::model::{
-    CancelledNotification, CancelledNotificationParam, ClientInfo, ClientNotification,
-    ClientRequest, ClientResult, CreateMessageRequest, CreateMessageRequestParam,
-    CreateMessageResult, ListRootsRequest, ListRootsResult, LoggingMessageNotification,
-    LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam,
-    PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification,
-    ResourceUpdatedNotificationParam, ServerInfo, ServerMessage, ServerNotification, ServerRequest,
-    ServerResult, ToolListChangedNotification,
+    CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientMessage, ClientNotification, ClientRequest, ClientResult, CreateMessageRequest, CreateMessageRequestParam, CreateMessageResult, ListRootsRequest, ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam, PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerInfo, ServerMessage, ServerNotification, ServerRequest, ServerResult, ToolListChangedNotification
 };
-
+use thiserror::Error;
 use super::*;
 use futures::{SinkExt, StreamExt};
 
@@ -26,6 +20,24 @@ impl ServiceRole for RoleServer {
     const IS_CLIENT: bool = false;
 }
 
+/// It represents the error that may occur when serving the server.
+///
+/// if you want to handle the error, you can use `serve_server_with_ct` or `serve_server` with `Result<RunningService<RoleServer, S>, ServerError>`
+#[derive(Error, Debug)]
+pub enum ServerError {
+    #[error("expect initialize request, but received: {0:?}")]
+    ExpectedInitRequest(Option<ClientMessage>),
+    
+    #[error("expect initialize notification, but received: {0:?}")]
+    ExpectedInitNotification(Option<ClientMessage>),
+    
+    #[error("connection closed: {0}")]
+    ConnectionClosed(String),
+    
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+}
+
 pub type ClientSink = Peer<RoleServer>;
 
 impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
@@ -55,6 +67,46 @@ where
     serve_server_with_ct(service, transport, CancellationToken::new()).await
 }
 
+/// Helper function to get the next message from the stream
+async fn expect_next_message<S>(
+    stream: &mut S, 
+    context: &str
+) -> Result<ClientMessage, ServerError> 
+where S: StreamExt<Item = ClientJsonRpcMessage> + Unpin
+{
+    Ok(stream
+        .next()
+        .await
+        .ok_or_else(|| ServerError::ConnectionClosed(context.to_string()))?
+        .into_message())
+}
+
+/// Helper function to expect a request from the stream
+async fn expect_request<S>(
+    stream: &mut S, 
+    context: &str
+) -> Result<(ClientRequest, RequestId), ServerError> 
+where S: StreamExt<Item = ClientJsonRpcMessage> + Unpin
+{
+    let msg = expect_next_message(stream, context).await?;
+    let msg_clone = msg.clone();
+    msg.into_request()
+        .ok_or_else(|| ServerError::ExpectedInitRequest(Some(msg_clone)))
+}
+
+/// Helper function to expect a notification from the stream
+async fn expect_notification<S>(
+    stream: &mut S, 
+    context: &str
+) -> Result<ClientNotification, ServerError> 
+where S: StreamExt<Item = ClientJsonRpcMessage> + Unpin
+{
+    let msg = expect_next_message(stream, context).await?;
+    let msg_clone = msg.clone();
+    msg.into_notification()
+        .ok_or_else(|| ServerError::ExpectedInitNotification(Some(msg_clone)))
+}
+
 pub async fn serve_server_with_ct<S, T, E, A>(
     service: S,
     transport: T,
@@ -70,54 +122,46 @@ where
     let mut stream = Box::pin(stream);
     let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
 
-    // service
-    let (request, id) = stream
-        .next()
-        .await
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::UnexpectedEof,
-            "expect initialize request",
-        ))?
-        .into_message()
-        .into_request()
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize request",
-        ))?;
+    // Convert ServerError to std::io::Error, then to E
+    let handle_server_error = |e: ServerError| -> E {
+        match e {
+            ServerError::Io(io_err) => io_err.into(),
+            other => std::io::Error::new(
+                std::io::ErrorKind::Other, 
+                format!("{}", other)
+            ).into()
+        }
+    };
+
+    // Get initialize request
+    let (request, id) = expect_request(&mut stream, "initialize request")
+        .await.map_err(handle_server_error)?;
+    
     let ClientRequest::InitializeRequest(peer_info) = request else {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize request",
-        )
-        .into());
+        return Err(handle_server_error(ServerError::ExpectedInitRequest(
+            Some(ClientMessage::Request(request, id))
+        )));
     };
+    
+    // Send initialize response
     let init_response = service.get_info();
     sink.send(
         ServerMessage::Response(ServerResult::InitializeResult(init_response), id)
             .into_json_rpc_message(),
     )
     .await?;
-    // waiting for notification
-    let notification = stream
-        .next()
-        .await
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::UnexpectedEof,
-            "expect initialize notification",
-        ))?
-        .into_message()
-        .into_notification()
-        .ok_or(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize notification",
-        ))?;
+    
+    // Wait for initialize notification
+    let notification = expect_notification(&mut stream, "initialize notification")
+        .await.map_err(handle_server_error)?;
+    
     let ClientNotification::InitializedNotification(_) = notification else {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            "expect initialize notification",
-        )
-        .into());
+        return Err(handle_server_error(ServerError::ExpectedInitNotification(
+            Some(ClientMessage::Notification(notification))
+        )));
     };
+    
+    // Continue processing service
     serve_inner(service, (sink, stream), peer_info.params, id_provider, ct).await
 }