refactor: request-id generation and messaging functions

Author: hashemix

crates/rust-mcp-sdk/src/mcp_runtimes/client_runtime.rs

@@ -7,14 +7,19 @@ use crate::schema::{
         ServerMessages,
     },
     InitializeRequest, InitializeRequestParams, InitializeResult, InitializedNotification,
-    RpcError, ServerResult,
+    RequestId, RpcError, ServerResult,
 };
 use async_trait::async_trait;
 use futures::future::{join_all, try_join_all};
 use futures::StreamExt;
 
-use rust_mcp_transport::{IoStream, McpDispatch, MessageDispatcher, Transport};
-use std::sync::{Arc, RwLock};
+use rust_mcp_transport::{
+    IoStream, McpDispatch, MessageDispatcher, RequestIdGen, RequestIdGenNumeric, Transport,
+};
+use std::{
+    sync::{Arc, RwLock},
+    time::Duration,
+};
 use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::sync::Mutex;
 
@@ -41,6 +46,7 @@ pub struct ClientRuntime {
     // Details about the connected server
     server_details: Arc<RwLock<Option<InitializeResult>>>,
     handlers: Mutex<Vec<tokio::task::JoinHandle<Result<(), McpSdkError>>>>,
+    request_id_gen: Box<dyn RequestIdGen>,
 }
 
 impl ClientRuntime {
@@ -61,6 +67,7 @@ impl ClientRuntime {
             client_details,
             server_details: Arc::new(RwLock::new(None)),
             handlers: Mutex::new(vec![]),
+            request_id_gen: Box::new(RequestIdGenNumeric::new(None)),
         }
     }
 
@@ -284,6 +291,33 @@ impl McpClient for ClientRuntime {
         }
     }
 
+    async fn send(
+        &self,
+        message: MessageFromClient,
+        request_id: Option<RequestId>,
+        timeout: Option<Duration>,
+    ) -> SdkResult<Option<ServerMessage>> {
+        let sender = self.sender();
+        let sender = sender.read().await;
+        let sender = sender
+            .as_ref()
+            .ok_or(schema_utils::SdkError::connection_closed())?;
+
+        let outgoing_request_id = self
+            .request_id_gen
+            .request_id_for_message(&message, request_id);
+
+        let mcp_message = ClientMessage::from_message(message, outgoing_request_id)?;
+
+        let response = sender
+            .send_message(ClientMessages::Single(mcp_message), timeout)
+            .await?
+            .map(|res| res.as_single())
+            .transpose()?;
+
+        Ok(response)
+    }
+
     async fn is_shut_down(&self) -> bool {
         self.transport.is_shut_down().await
     }

crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs

@@ -16,7 +16,7 @@ use futures::future::try_join_all;
 use futures::{StreamExt, TryFutureExt};
 #[cfg(feature = "hyper-server")]
 use rust_mcp_transport::SessionId;
-use rust_mcp_transport::{IoStream, TransportDispatcher};
+use rust_mcp_transport::{IoStream, RequestIdGen, RequestIdGenNumeric, TransportDispatcher};
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
@@ -45,6 +45,7 @@ pub struct ServerRuntime {
     #[cfg(feature = "hyper-server")]
     session_id: Option<SessionId>,
     transport_map: tokio::sync::RwLock<HashMap<String, TransportType>>,
+    request_id_gen: Box<dyn RequestIdGen>,
     client_details_tx: watch::Sender<Option<InitializeRequestParams>>,
     client_details_rx: watch::Receiver<Option<InitializeRequestParams>>,
 }
@@ -79,22 +80,26 @@ impl McpServer for ServerRuntime {
         message: MessageFromServer,
         request_id: Option<RequestId>,
         request_timeout: Option<Duration>,
-    ) -> SdkResult<Option<ClientMessages>> {
+    ) -> SdkResult<Option<ClientMessage>> {
         let transport_map = self.transport_map.read().await;
         let transport = transport_map.get(DEFAULT_STREAM_ID).ok_or(
             RpcError::internal_error()
                 .with_message("transport stream does not exists or is closed!".to_string()),
         )?;
 
         let outgoing_request_id = self
-            .request_id_for_message(transport, &message, request_id)
-            .await;
+            .request_id_gen
+            .request_id_for_message(&message, request_id);
 
         let mcp_message = ServerMessage::from_message(message, outgoing_request_id)?;
-        transport
+
+        let response = transport
             .send_message(ServerMessages::Single(mcp_message), request_timeout)
-            .map_err(|err| err.into())
-            .await
+            .await?
+            .map(|res| res.as_single())
+            .transpose()?;
+
+        Ok(response)
     }
 
     async fn send_batch(
@@ -211,40 +216,6 @@ impl ServerRuntime {
         Ok(())
     }
 
-    /// Determines the request ID for an outgoing MCP message.
-    ///
-    /// For requests, generates a new ID using the internal counter. For responses or errors,
-    /// uses the provided `request_id`. Notifications receive no ID.
-    ///
-    /// # Arguments
-    /// * `message` - The MCP message to evaluate.
-    /// * `request_id` - An optional existing request ID (required for responses/errors).
-    ///
-    /// # Returns
-    /// An `Option<RequestId>`: `Some` for requests or responses/errors, `None` for notifications.
-    pub(crate) async fn request_id_for_message(
-        &self,
-        transport: &Arc<
-            dyn TransportDispatcher<
-                ClientMessages,
-                MessageFromServer,
-                ClientMessage,
-                ServerMessages,
-                ServerMessage,
-            >,
-        >,
-        message: &MessageFromServer,
-        request_id: Option<RequestId>,
-    ) -> Option<RequestId> {
-        let message_sender = transport.message_sender();
-        let guard = message_sender.read().await;
-        if let Some(dispatcher) = guard.as_ref() {
-            dispatcher.request_id_for_message(message, request_id)
-        } else {
-            None
-        }
-    }
-
     pub(crate) async fn handle_message(
         &self,
         message: ClientMessage,
@@ -471,6 +442,7 @@ impl ServerRuntime {
             transport_map: tokio::sync::RwLock::new(HashMap::new()),
             client_details_tx,
             client_details_rx,
+            request_id_gen: Box::new(RequestIdGenNumeric::new(None)),
         }
     }
 
@@ -497,6 +469,7 @@ impl ServerRuntime {
             transport_map: tokio::sync::RwLock::new(map),
             client_details_tx,
             client_details_rx,
+            request_id_gen: Box::new(RequestIdGenNumeric::new(None)),
         }
     }
 }

crates/rust-mcp-sdk/src/mcp_traits/mcp_client.rs

@@ -10,7 +10,7 @@ use crate::schema::{
     InitializeRequestParams, InitializeResult, ListPromptsRequest, ListPromptsRequestParams,
     ListResourceTemplatesRequest, ListResourceTemplatesRequestParams, ListResourcesRequest,
     ListResourcesRequestParams, ListRootsRequest, ListToolsRequest, ListToolsRequestParams,
-    LoggingLevel, PingRequest, ReadResourceRequest, ReadResourceRequestParams,
+    LoggingLevel, PingRequest, ReadResourceRequest, ReadResourceRequestParams, RequestId,
     RootsListChangedNotification, RootsListChangedNotificationParams, RpcError, ServerCapabilities,
     SetLevelRequest, SetLevelRequestParams, SubscribeRequest, SubscribeRequestParams,
     UnsubscribeRequest, UnsubscribeRequestParams,
@@ -175,27 +175,15 @@ pub trait McpClient: Sync + Send {
         request: RequestFromClient,
         timeout: Option<Duration>,
     ) -> SdkResult<ResultFromServer> {
-        let sender = self.sender();
-        let sender = sender.read().await;
-        let sender = sender
-            .as_ref()
-            .ok_or(schema_utils::SdkError::connection_closed())?;
-
-        let request_id = sender.next_request_id();
-
-        let mcp_message =
-            ClientMessage::from_message(MessageFromClient::from(request), Some(request_id))?;
-        let response = sender
-            .send_message(ClientMessages::Single(mcp_message), timeout)
+        let response = self
+            .send(MessageFromClient::RequestFromClient(request), None, timeout)
             .await?;
 
         let server_message = response.ok_or_else(|| {
             RpcError::internal_error()
-                .with_message("An empty response was received from the server.".to_string())
+                .with_message("An empty response was received from the client.".to_string())
         })?;
 
-        let server_message = server_message.as_single()?;
-
         if server_message.is_error() {
             return Err(server_message.as_error()?.error.into());
         }
@@ -205,27 +193,10 @@ pub trait McpClient: Sync + Send {
 
     async fn send(
         &self,
-        message: ClientMessage,
+        message: MessageFromClient,
+        request_id: Option<RequestId>,
         timeout: Option<Duration>,
-    ) -> SdkResult<Option<ServerMessage>> {
-        let sender = self.sender();
-        let sender = sender.read().await;
-        let sender = sender
-            .as_ref()
-            .ok_or(schema_utils::SdkError::connection_closed())?;
-
-        let response = sender
-            .send_message(ClientMessages::Single(message), timeout)
-            .await?;
-
-        match response {
-            Some(res) => {
-                let server_results = res.as_single()?;
-                Ok(Some(server_results))
-            }
-            None => Ok(None),
-        }
-    }
+    ) -> SdkResult<Option<ServerMessage>>;
 
     async fn send_batch(
         &self,

crates/rust-mcp-sdk/src/mcp_traits/mcp_server.rs

@@ -2,8 +2,8 @@ use std::time::Duration;
 
 use crate::schema::{
     schema_utils::{
-        ClientMessage, ClientMessages, McpMessage, MessageFromServer, NotificationFromServer,
-        RequestFromServer, ResultFromClient, ServerMessage,
+        ClientMessage, McpMessage, MessageFromServer, NotificationFromServer, RequestFromServer,
+        ResultFromClient, ServerMessage,
     },
     CallToolRequest, CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult,
     GetPromptRequest, Implementation, InitializeRequestParams, InitializeResult,
@@ -44,7 +44,7 @@ pub trait McpServer: Sync + Send {
         message: MessageFromServer,
         request_id: Option<RequestId>,
         request_timeout: Option<Duration>,
-    ) -> SdkResult<Option<ClientMessages>>;
+    ) -> SdkResult<Option<ClientMessage>>;
 
     async fn send_batch(
         &self,
@@ -84,13 +84,11 @@ pub trait McpServer: Sync + Send {
             .send(MessageFromServer::RequestFromServer(request), None, timeout)
             .await?;
 
-        let client_messages = response.ok_or_else(|| {
+        let client_message = response.ok_or_else(|| {
             RpcError::internal_error()
                 .with_message("An empty response was received from the client.".to_string())
         })?;
 
-        let client_message = client_messages.as_single()?;
-
         if client_message.is_error() {
             return Err(client_message.as_error()?.error.into());
         }

crates/rust-mcp-sdk/tests/common/common.rs

@@ -128,8 +128,10 @@ use futures::stream::Stream;
 // stream: &mut impl Stream<Item = Result<hyper::body::Bytes, hyper::Error>>,
 pub async fn read_sse_event_from_stream(
     stream: &mut (impl Stream<Item = Result<hyper::body::Bytes, reqwest::Error>> + Unpin),
-) -> Option<String> {
+    event_count: usize,
+) -> Option<Vec<String>> {
     let mut buffer = String::new();
+    let mut events = vec![];
 
     while let Some(item) = stream.next().await {
         match item {
@@ -158,7 +160,10 @@ pub async fn read_sse_event_from_stream(
 
                     // Return if data was found
                     if let Some(data) = data {
-                        return Some(data);
+                        events.push(data);
+                        if events.len().eq(&event_count) {
+                            return Some(events);
+                        }
                     }
                 }
             }
@@ -171,9 +176,9 @@ pub async fn read_sse_event_from_stream(
     None
 }
 
-pub async fn read_sse_event(response: Response) -> Option<String> {
+pub async fn read_sse_event(response: Response, event_count: usize) -> Option<Vec<String>> {
     let mut stream = response.bytes_stream();
-    read_sse_event_from_stream(&mut stream).await
+    read_sse_event_from_stream(&mut stream, event_count).await
 }
 
 pub fn test_client_info() -> InitializeRequestParams {