rust_mcp_transport/
message_dispatcher.rs

1use crate::error::{TransportError, TransportResult};
2use crate::schema::{
3    schema_utils::{
4        self, ClientMessage, ClientMessages, McpMessage, RpcMessage, ServerMessage, ServerMessages,
5    },
6    JsonrpcErrorResponse,
7};
8use crate::schema::{RequestId, RpcError};
9use crate::utils::await_timeout;
10use crate::McpDispatch;
11use async_trait::async_trait;
12use futures::future::join_all;
13use std::collections::HashMap;
14use std::pin::Pin;
15use std::sync::Arc;
16use std::time::Duration;
17use tokio::io::AsyncWriteExt;
18use tokio::sync::oneshot::{self};
19use tokio::sync::Mutex;
20
21/// Provides a dispatcher for sending MCP messages and handling responses.
22///
23/// `MessageDispatcher` facilitates MCP communication by managing message sending, request tracking,
24/// and response handling. It supports both client-to-server and server-to-client message flows through
25/// implementations of the `McpDispatch` trait. The dispatcher uses a transport mechanism
26/// (e.g., stdin/stdout) to serialize and send messages, and it tracks pending requests with
27/// a configurable timeout mechanism for asynchronous responses.
28pub struct MessageDispatcher<R> {
29    pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<R>>>>,
30    writable_std: Option<Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>>>,
31    writable_tx: Option<
32        tokio::sync::mpsc::Sender<(
33            String,
34            tokio::sync::oneshot::Sender<crate::error::TransportResult<()>>,
35        )>,
36    >,
37    request_timeout: Duration,
38}
39
40impl<R> MessageDispatcher<R> {
41    /// Creates a new `MessageDispatcher` instance with the given configuration.
42    ///
43    /// # Arguments
44    /// * `pending_requests` - A thread-safe map for storing pending request IDs and their response channels.
45    /// * `writable_std` - A mutex-protected, pinned writer (e.g., stdout) for sending serialized messages.
46    /// * `message_id_counter` - An atomic counter for generating unique request IDs.
47    /// * `request_timeout` - The timeout duration in milliseconds for awaiting responses.
48    ///
49    /// # Returns
50    /// A new `MessageDispatcher` instance configured for MCP message handling.
51    pub fn new(
52        pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<R>>>>,
53        writable_std: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>>,
54        request_timeout: Duration,
55    ) -> Self {
56        Self {
57            pending_requests,
58            writable_std: Some(writable_std),
59            writable_tx: None,
60            request_timeout,
61        }
62    }
63
64    pub fn new_with_acknowledgement(
65        pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<R>>>>,
66        writable_tx: tokio::sync::mpsc::Sender<(
67            String,
68            tokio::sync::oneshot::Sender<crate::error::TransportResult<()>>,
69        )>,
70        request_timeout: Duration,
71    ) -> Self {
72        Self {
73            pending_requests,
74            writable_tx: Some(writable_tx),
75            writable_std: None,
76            request_timeout,
77        }
78    }
79
80    async fn store_pending_request(
81        &self,
82        request_id: RequestId,
83    ) -> tokio::sync::oneshot::Receiver<R> {
84        let (tx_response, rx_response) = oneshot::channel::<R>();
85        let mut pending_requests = self.pending_requests.lock().await;
86        // store request id in the hashmap while waiting for a matching response
87        pending_requests.insert(request_id.clone(), tx_response);
88        rx_response
89    }
90
91    async fn store_pending_request_for_message<M: McpMessage + RpcMessage>(
92        &self,
93        message: &M,
94    ) -> Option<tokio::sync::oneshot::Receiver<R>> {
95        if message.is_request() {
96            if let Some(request_id) = message.request_id() {
97                Some(self.store_pending_request(request_id.clone()).await)
98            } else {
99                None
100            }
101        } else {
102            None
103        }
104    }
105}
106
107// Client side dispatcher
108#[async_trait]
109impl McpDispatch<ServerMessages, ClientMessages, ServerMessage, ClientMessage>
110    for MessageDispatcher<ServerMessage>
111{
112    /// Sends a message from the client to the server and awaits a response if applicable.
113    ///
114    /// Serializes the `ClientMessages` to JSON, writes it to the transport, and waits for a
115    /// `ServerMessages` response if the message is a request. Notifications and responses return
116    /// `Ok(None)`.
117    ///
118    /// # Arguments
119    /// * `messages` - The client message to send, coulld be a single message or batch.
120    ///
121    /// # Returns
122    /// A `TransportResult` containing `Some(ServerMessages)` for requests with a response,
123    /// or `None` for notifications/responses, or an error if the operation fails.
124    ///
125    /// # Errors
126    /// Returns a `TransportError` if serialization, writing, or timeout occurs.
127    async fn send_message(
128        &self,
129        messages: ClientMessages,
130        request_timeout: Option<Duration>,
131    ) -> TransportResult<Option<ServerMessages>> {
132        match messages {
133            ClientMessages::Single(message) => {
134                let rx_response: Option<tokio::sync::oneshot::Receiver<ServerMessage>> =
135                    self.store_pending_request_for_message(&message).await;
136
137                //serialize the message and write it to the writable_std
138                let message_payload = serde_json::to_string(&message).map_err(|_| {
139                    crate::error::TransportError::JsonrpcError(RpcError::parse_error())
140                })?;
141
142                self.write_str(message_payload.as_str(), true).await?;
143
144                if let Some(rx) = rx_response {
145                    // Wait for the response with timeout
146                    match await_timeout(rx, request_timeout.unwrap_or(self.request_timeout)).await {
147                        Ok(response) => Ok(Some(ServerMessages::Single(response))),
148                        Err(error) => match error {
149                            TransportError::ChannelClosed(_) => {
150                                Err(schema_utils::SdkError::connection_closed().into())
151                            }
152                            _ => Err(error),
153                        },
154                    }
155                } else {
156                    Ok(None)
157                }
158            }
159            ClientMessages::Batch(client_messages) => {
160                let (request_ids, pending_tasks): (Vec<_>, Vec<_>) = client_messages
161                    .iter()
162                    .filter(|message| message.is_request())
163                    .map(|message| {
164                        (
165                            message.request_id(),
166                            self.store_pending_request_for_message(message),
167                        )
168                    })
169                    .unzip();
170
171                // Ensure all request IDs are stored before sending the request
172                let tasks = join_all(pending_tasks).await;
173
174                // send the batch messages to the server
175                let message_payload = serde_json::to_string(&client_messages).map_err(|_| {
176                    crate::error::TransportError::JsonrpcError(RpcError::parse_error())
177                })?;
178                self.write_str(message_payload.as_str(), true).await?;
179
180                // no request in the batch, no need to wait for the result
181                if request_ids.is_empty() {
182                    return Ok(None);
183                }
184
185                let timeout_wrapped_futures = tasks.into_iter().filter_map(|rx| {
186                    rx.map(|rx| await_timeout(rx, request_timeout.unwrap_or(self.request_timeout)))
187                });
188
189                let results: Vec<_> = join_all(timeout_wrapped_futures)
190                    .await
191                    .into_iter()
192                    .zip(request_ids)
193                    .map(|(res, request_id)| match res {
194                        Ok(response) => response,
195                        Err(error) => ServerMessage::Error(JsonrpcErrorResponse::new(
196                            RpcError::internal_error().with_message(error.to_string()),
197                            request_id.cloned(),
198                        )),
199                    })
200                    .collect();
201
202                Ok(Some(ServerMessages::Batch(results)))
203            }
204        }
205    }
206
207    async fn send(
208        &self,
209        message: ClientMessage,
210        request_timeout: Option<Duration>,
211    ) -> TransportResult<Option<ServerMessage>> {
212        let response = self.send_message(message.into(), request_timeout).await?;
213        match response {
214            Some(r) => Ok(Some(r.as_single()?)),
215            None => Ok(None),
216        }
217    }
218
219    /// Writes a string payload to the underlying asynchronous writable stream,
220    /// appending a newline character and flushing the stream afterward.
221    ///
222    async fn write_str(&self, payload: &str, _skip_store: bool) -> TransportResult<()> {
223        if let Some(writable_std) = self.writable_std.as_ref() {
224            let mut writable_std = writable_std.lock().await;
225            writable_std.write_all(payload.as_bytes()).await?;
226            writable_std.write_all(b"\n").await?; // new line
227            writable_std.flush().await?;
228            return Ok(());
229        };
230
231        if let Some(writable_tx) = self.writable_tx.as_ref() {
232            let (resp_tx, resp_rx) = oneshot::channel();
233            writable_tx
234                .send((payload.to_string(), resp_tx))
235                .await
236                .map_err(|err| TransportError::Internal(format!("{err}")))?; // Send fails if channel closed
237            return resp_rx.await?; // Await the POST result; propagates the error if POST failed
238        }
239
240        Err(TransportError::Internal("Invalid dispatcher!".to_string()))
241    }
242}
243
244// Server side dispatcher, Sends S and Returns R
245#[async_trait]
246impl McpDispatch<ClientMessages, ServerMessages, ClientMessage, ServerMessage>
247    for MessageDispatcher<ClientMessage>
248{
249    /// Sends a message from the server to the client and awaits a response if applicable.
250    ///
251    /// Serializes the `ServerMessages` to JSON, writes it to the transport, and waits for a
252    /// `ClientMessages` response if the message is a request. Notifications and responses return
253    /// `Ok(None)`.
254    ///
255    /// # Arguments
256    /// * `messages` - The client message to send, coulld be a single message or batch.
257    ///
258    /// # Returns
259    /// A `TransportResult` containing `Some(ClientMessages)` for requests with a response,
260    /// or `None` for notifications/responses, or an error if the operation fails.
261    ///
262    /// # Errors
263    /// Returns a `TransportError` if serialization, writing, or timeout occurs.
264    async fn send_message(
265        &self,
266        messages: ServerMessages,
267        request_timeout: Option<Duration>,
268    ) -> TransportResult<Option<ClientMessages>> {
269        match messages {
270            ServerMessages::Single(message) => {
271                let rx_response: Option<tokio::sync::oneshot::Receiver<ClientMessage>> =
272                    self.store_pending_request_for_message(&message).await;
273
274                let message_payload = serde_json::to_string(&message).map_err(|_| {
275                    crate::error::TransportError::JsonrpcError(RpcError::parse_error())
276                })?;
277
278                self.write_str(message_payload.as_str(), false).await?;
279
280                if let Some(rx) = rx_response {
281                    match await_timeout(rx, request_timeout.unwrap_or(self.request_timeout)).await {
282                        Ok(response) => Ok(Some(ClientMessages::Single(response))),
283                        Err(error) => Err(error),
284                    }
285                } else {
286                    Ok(None)
287                }
288            }
289            ServerMessages::Batch(server_messages) => {
290                let (request_ids, pending_tasks): (Vec<_>, Vec<_>) = server_messages
291                    .iter()
292                    .filter(|message| message.is_request())
293                    .map(|message| {
294                        (
295                            message.request_id(),
296                            self.store_pending_request_for_message(message),
297                        )
298                    })
299                    .unzip();
300
301                // send the batch messages to the client
302                let message_payload = serde_json::to_string(&server_messages).map_err(|_| {
303                    crate::error::TransportError::JsonrpcError(RpcError::parse_error())
304                })?;
305
306                self.write_str(message_payload.as_str(), false).await?;
307
308                // no request in the batch, no need to wait for the result
309                if pending_tasks.is_empty() {
310                    return Ok(None);
311                }
312
313                let tasks = join_all(pending_tasks).await;
314
315                let timeout_wrapped_futures = tasks.into_iter().filter_map(|rx| {
316                    rx.map(|rx| await_timeout(rx, request_timeout.unwrap_or(self.request_timeout)))
317                });
318
319                let results: Vec<_> = join_all(timeout_wrapped_futures)
320                    .await
321                    .into_iter()
322                    .zip(request_ids)
323                    .map(|(res, request_id)| match res {
324                        Ok(response) => response,
325                        Err(error) => ClientMessage::Error(JsonrpcErrorResponse::new(
326                            RpcError::internal_error().with_message(error.to_string()),
327                            request_id.cloned(),
328                        )),
329                    })
330                    .collect();
331
332                Ok(Some(ClientMessages::Batch(results)))
333            }
334        }
335    }
336
337    async fn send(
338        &self,
339        message: ServerMessage,
340        request_timeout: Option<Duration>,
341    ) -> TransportResult<Option<ClientMessage>> {
342        let response = self.send_message(message.into(), request_timeout).await?;
343        match response {
344            Some(r) => Ok(Some(r.as_single()?)),
345            None => Ok(None),
346        }
347    }
348
349    async fn write_str(&self, payload: &str, _skip_store: bool) -> TransportResult<()> {
350        if let Some(writable_std) = self.writable_std.as_ref() {
351            let mut writable_std = writable_std.lock().await;
352            writable_std.write_all(payload.as_bytes()).await?;
353            writable_std.write_all(b"\n").await?; // new line
354            writable_std.flush().await?;
355            return Ok(());
356        };
357
358        if let Some(writable_tx) = self.writable_tx.as_ref() {
359            let (resp_tx, resp_rx) = oneshot::channel();
360            writable_tx
361                .send((payload.to_string(), resp_tx))
362                .await
363                .map_err(|err| TransportError::Internal(err.to_string()))?; // Send fails if channel closed
364            return resp_rx.await?; // Await the POST result; propagates the error if POST failed
365        }
366
367        Err(TransportError::Internal("Invalid dispatcher!".to_string()))
368    }
369}
rust_mcp_transport/message_dispatcher.rs

rust_mcp_transport/
message_dispatcher.rs
