fix: ensure server-initiated requests include a valid request_id (#80)

Author: hashemix

crates/rust-mcp-sdk/src/mcp_handlers/mcp_server_handler.rs

@@ -51,6 +51,7 @@ pub trait ServerHandler: Send + Sync + 'static {
 
         runtime
             .set_client_details(initialize_request.params.clone())
+            .await
             .map_err(|err| RpcError::internal_error().with_message(format!("{err}")))?;
 
         Ok(server_info)

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

@@ -1,30 +1,27 @@
 pub mod mcp_server_runtime;
 pub mod mcp_server_runtime_core;
+use crate::error::SdkResult;
+use crate::mcp_traits::mcp_handler::McpServerHandler;
+use crate::mcp_traits::mcp_server::McpServer;
 use crate::schema::{
     schema_utils::{
-        ClientMessage, ClientMessages, FromMessage, MessageFromServer, SdkError, ServerMessage,
-        ServerMessages,
+        ClientMessage, ClientMessages, FromMessage, McpMessage, MessageFromServer, SdkError,
+        ServerMessage, ServerMessages,
     },
     InitializeRequestParams, InitializeResult, RequestId, RpcError,
 };
-
 use async_trait::async_trait;
 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 std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use tokio::io::AsyncWriteExt;
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, watch};
 
-use crate::error::SdkResult;
-use crate::mcp_traits::mcp_handler::McpServerHandler;
-use crate::mcp_traits::mcp_server::McpServer;
-#[cfg(feature = "hyper-server")]
-use rust_mcp_transport::SessionId;
 pub const DEFAULT_STREAM_ID: &str = "STANDALONE-STREAM";
 
 // Define a type alias for the TransportDispatcher trait object
@@ -49,21 +46,32 @@ pub struct ServerRuntime {
     #[cfg(feature = "hyper-server")]
     session_id: Option<SessionId>,
     transport_map: tokio::sync::RwLock<HashMap<String, TransportType>>,
+    client_details_tx: watch::Sender<Option<InitializeRequestParams>>,
+    client_details_rx: watch::Receiver<Option<InitializeRequestParams>>,
 }
 
 #[async_trait]
 impl McpServer for ServerRuntime {
     /// Set the client details, storing them in client_details
-    fn set_client_details(&self, client_details: InitializeRequestParams) -> SdkResult<()> {
-        match self.client_details.write() {
-            Ok(mut details) => {
-                *details = Some(client_details);
-                Ok(())
+    async fn set_client_details(&self, client_details: InitializeRequestParams) -> SdkResult<()> {
+        self.handler.on_server_started(self).await;
+
+        self.client_details_tx
+            .send(Some(client_details))
+            .map_err(|_| {
+                RpcError::internal_error()
+                    .with_message("Failed to set client details".to_string())
+                    .into()
+            })
+    }
+
+    async fn wait_for_initialization(&self) {
+        loop {
+            if self.client_details_rx.borrow().is_some() {
+                return;
             }
-            // Failed to acquire read lock, likely due to PoisonError from a thread panic. Returning None.
-            Err(_) => Err(RpcError::internal_error()
-                .with_message("Internal Error: Failed to acquire write lock.".to_string())
-                .into()),
+            let mut rx = self.client_details_rx.clone();
+            rx.changed().await.ok();
         }
     }
 
@@ -79,7 +87,19 @@ impl McpServer for ServerRuntime {
                 .with_message("transport stream does not exists or is closed!".to_string()),
         )?;
 
-        let mcp_message = ServerMessage::from_message(message, request_id)?;
+        // generate a new request_id for request messages
+        let outgoing_request_id = if message.is_request() {
+            match request_id {
+                Some(_) => Err(RpcError::internal_error().with_message(
+                    "request_id should not have a value when sending a new request".to_string(),
+                )),
+                None => Ok(self.next_request_id(transport).await),
+            }
+        } else {
+            Ok(request_id)
+        }?;
+
+        let mcp_message = ServerMessage::from_message(message, outgoing_request_id)?;
         transport
             .send_message(ServerMessages::Single(mcp_message), request_timeout)
             .map_err(|err| err.into())
@@ -130,8 +150,6 @@ impl McpServer for ServerRuntime {
 
         let mut stream = transport.start().await?;
 
-        self.handler.on_server_started(self).await;
-
         // Process incoming messages from the client
         while let Some(mcp_messages) = stream.next().await {
             match mcp_messages {
@@ -207,6 +225,25 @@ impl ServerRuntime {
         Ok(())
     }
 
+    pub(crate) async fn next_request_id(
+        &self,
+        transport: &Arc<
+            dyn TransportDispatcher<
+                ClientMessages,
+                MessageFromServer,
+                ClientMessage,
+                ServerMessages,
+                ServerMessage,
+            >,
+        >,
+    ) -> Option<RequestId> {
+        let message_sender = transport.message_sender();
+        let guard = message_sender.read().await;
+        guard
+            .as_ref()
+            .map(|dispatcher| dispatcher.next_request_id())
+    }
+
     pub(crate) async fn handle_message(
         &self,
         message: ClientMessage,
@@ -416,12 +453,16 @@ impl ServerRuntime {
         handler: Arc<dyn McpServerHandler>,
         session_id: SessionId,
     ) -> Self {
+        let (client_details_tx, client_details_rx) =
+            watch::channel::<Option<InitializeRequestParams>>(None);
         Self {
             server_details,
             client_details: Arc::new(RwLock::new(None)),
             handler,
             session_id: Some(session_id),
             transport_map: tokio::sync::RwLock::new(HashMap::new()),
+            client_details_tx,
+            client_details_rx,
         }
     }
 
@@ -438,13 +479,17 @@ impl ServerRuntime {
     ) -> Self {
         let mut map: HashMap<String, TransportType> = HashMap::new();
         map.insert(DEFAULT_STREAM_ID.to_string(), Arc::new(transport));
+        let (client_details_tx, client_details_rx) =
+            watch::channel::<Option<InitializeRequestParams>>(None);
         Self {
             server_details: Arc::new(server_details),
             client_details: Arc::new(RwLock::new(None)),
             handler,
             #[cfg(feature = "hyper-server")]
             session_id: None,
             transport_map: tokio::sync::RwLock::new(map),
+            client_details_tx,
+            client_details_rx,
         }
     }
 }

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

@@ -76,6 +76,7 @@ impl McpServerHandler for RuntimeCoreInternalHandler<Box<dyn ServerHandlerCore>>
             // keep a copy of the InitializeRequestParams which includes client_info and capabilities
             runtime
                 .set_client_details(initialize_request.params.clone())
+                .await
                 .map_err(|err| RpcError::internal_error().with_message(format!("{err}")))?;
         }
 

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

@@ -23,10 +23,12 @@ use crate::{error::SdkResult, utils::format_assertion_message};
 #[async_trait]
 pub trait McpServer: Sync + Send {
     async fn start(&self) -> SdkResult<()>;
-    fn set_client_details(&self, client_details: InitializeRequestParams) -> SdkResult<()>;
+    async fn set_client_details(&self, client_details: InitializeRequestParams) -> SdkResult<()>;
     fn server_info(&self) -> &InitializeResult;
     fn client_info(&self) -> Option<InitializeRequestParams>;
 
+    async fn wait_for_initialization(&self);
+
     #[deprecated(since = "0.2.0", note = "Use `client_info()` instead.")]
     fn get_client_info(&self) -> Option<InitializeRequestParams> {
         self.client_info()