rust_mcp_transport/
client_sse.rs

1use crate::error::{TransportError, TransportResult};
2use crate::mcp_stream::MCPStream;
3use crate::message_dispatcher::MessageDispatcher;
4use crate::transport::Transport;
5use crate::utils::{
6    extract_origin, http_post, CancellationTokenSource, ReadableChannel, SseStream, WritableChannel,
7};
8use crate::{IoStream, McpDispatch, TransportDispatcher, TransportOptions};
9use async_trait::async_trait;
10use bytes::Bytes;
11use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
12use reqwest::Client;
13use tokio::sync::oneshot::Sender;
14use tokio::task::JoinHandle;
15
16use crate::schema::{
17    schema_utils::{
18        ClientMessage, ClientMessages, McpMessage, MessageFromClient, SdkError, ServerMessage,
19        ServerMessages,
20    },
21    RequestId,
22};
23use std::cmp::Ordering;
24use std::collections::HashMap;
25use std::pin::Pin;
26use std::sync::Arc;
27use std::time::Duration;
28use tokio::io::{BufReader, BufWriter};
29use tokio::sync::{mpsc, oneshot, Mutex};
30
31const DEFAULT_CHANNEL_CAPACITY: usize = 64;
32const DEFAULT_MAX_RETRY: usize = 5;
33const DEFAULT_RETRY_TIME_SECONDS: u64 = 1;
34const SHUTDOWN_TIMEOUT_SECONDS: u64 = 5;
35
36/// Configuration options for the Client SSE Transport
37///
38/// Defines settings for request timeouts, retry behavior, and custom HTTP headers.
39pub struct ClientSseTransportOptions {
40    pub request_timeout: Duration,
41    pub max_line_length: usize,
42    pub channel_capacity: usize,
43    pub retry_delay: Option<Duration>,
44    pub max_retries: Option<usize>,
45    pub custom_headers: Option<HashMap<String, String>>,
46}
47
48/// Provides default values for ClientSseTransportOptions
49impl Default for ClientSseTransportOptions {
50    fn default() -> Self {
51        Self {
52            request_timeout: TransportOptions::default().timeout,
53            max_line_length: TransportOptions::default().max_line_length,
54            channel_capacity: TransportOptions::default().channel_capacity,
55            retry_delay: None,
56            max_retries: None,
57            custom_headers: None,
58        }
59    }
60}
61
62/// Client-side Server-Sent Events (SSE) transport implementation
63///
64/// Manages SSE connections, HTTP POST requests, and message streaming for client-server communication.
65pub struct ClientSseTransport<R>
66where
67    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
68{
69    /// Optional cancellation token source for shutting down the transport
70    shutdown_source: tokio::sync::RwLock<Option<CancellationTokenSource>>,
71    /// Flag indicating if the transport is shut down
72    is_shut_down: Mutex<bool>,
73    /// Timeout duration for MCP messages
74    request_timeout: Duration,
75    /// Maximum line length for incoming messages
76    max_line_length: usize,
77    /// Capacity of the incoming-message channel buffer
78    channel_capacity: usize,
79    /// HTTP client for making requests
80    client: Client,
81    /// URL for the SSE endpoint
82    sse_url: String,
83    /// Base URL extracted from the server URL
84    base_url: String,
85    /// Delay between retry attempts
86    retry_delay: Duration,
87    /// Maximum number of retry attempts
88    max_retries: usize,
89    /// Optional custom HTTP headers
90    custom_headers: Option<HeaderMap>,
91    sse_task: tokio::sync::RwLock<Option<tokio::task::JoinHandle<()>>>,
92    post_task: tokio::sync::RwLock<Option<tokio::task::JoinHandle<()>>>,
93    message_sender: Arc<tokio::sync::RwLock<Option<MessageDispatcher<R>>>>,
94    error_stream: tokio::sync::RwLock<Option<IoStream>>,
95    pending_requests: Arc<Mutex<HashMap<RequestId, tokio::sync::oneshot::Sender<R>>>>,
96}
97
98impl<R> ClientSseTransport<R>
99where
100    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
101{
102    /// Creates a new ClientSseTransport instance
103    ///
104    /// Initializes the transport with the provided server URL and options.
105    ///
106    /// # Arguments
107    /// * `server_url` - The URL of the SSE server
108    /// * `options` - Configuration options for the transport
109    ///
110    /// # Returns
111    /// * `TransportResult<Self>` - The initialized transport or an error
112    pub fn new(server_url: &str, options: ClientSseTransportOptions) -> TransportResult<Self> {
113        let client = Client::new();
114
115        let base_url = match extract_origin(server_url) {
116            Some(url) => url,
117            None => {
118                let message = format!("Failed to extract origin from server URL: {server_url}");
119                tracing::error!(message);
120                return Err(TransportError::Configuration { message });
121            }
122        };
123
124        let headers = match &options.custom_headers {
125            Some(h) => Some(Self::validate_headers(h)?),
126            None => None,
127        };
128
129        Ok(Self {
130            client,
131            base_url,
132            sse_url: server_url.to_string(),
133            max_retries: options.max_retries.unwrap_or(DEFAULT_MAX_RETRY),
134            retry_delay: options
135                .retry_delay
136                .unwrap_or(Duration::from_secs(DEFAULT_RETRY_TIME_SECONDS)),
137            shutdown_source: tokio::sync::RwLock::new(None),
138            is_shut_down: Mutex::new(false),
139            request_timeout: options.request_timeout,
140            max_line_length: options.max_line_length,
141            channel_capacity: options.channel_capacity,
142            custom_headers: headers,
143            sse_task: tokio::sync::RwLock::new(None),
144            post_task: tokio::sync::RwLock::new(None),
145            message_sender: Arc::new(tokio::sync::RwLock::new(None)),
146            error_stream: tokio::sync::RwLock::new(None),
147            pending_requests: Arc::new(Mutex::new(HashMap::new())),
148        })
149    }
150
151    /// Validates and converts a HashMap of headers into a HeaderMap
152    ///
153    /// # Arguments
154    /// * `headers` - The HashMap of header names and values
155    ///
156    /// # Returns
157    /// * `TransportResult<HeaderMap>` - The validated HeaderMap or an error
158    fn validate_headers(headers: &HashMap<String, String>) -> TransportResult<HeaderMap> {
159        let mut header_map = HeaderMap::new();
160
161        for (key, value) in headers {
162            let header_name =
163                key.parse::<HeaderName>()
164                    .map_err(|e| TransportError::Configuration {
165                        message: format!("Invalid header name: {e}"),
166                    })?;
167            let header_value =
168                HeaderValue::from_str(value).map_err(|e| TransportError::Configuration {
169                    message: format!("Invalid header value: {e}"),
170                })?;
171            header_map.insert(header_name, header_value);
172        }
173
174        Ok(header_map)
175    }
176
177    /// Validates the message endpoint URL
178    ///
179    /// Ensures the endpoint is either relative to the base URL or matches the base URL's origin.
180    ///
181    /// # Arguments
182    /// * `endpoint` - The endpoint URL to validate
183    ///
184    /// # Returns
185    /// * `TransportResult<String>` - The validated endpoint URL or an error
186    pub fn validate_message_endpoint(&self, endpoint: String) -> TransportResult<String> {
187        if endpoint.starts_with("/") {
188            return Ok(format!("{}{}", self.base_url, endpoint));
189        }
190        if let Some(endpoint_origin) = extract_origin(&endpoint) {
191            if endpoint_origin.cmp(&self.base_url) != Ordering::Equal {
192                return Err(TransportError::Configuration {
193                    message: format!(
194                    "Endpoint origin does not match connection origin. expected: {} , received: {}",
195                    self.base_url, endpoint_origin
196                ),
197                });
198            }
199            return Ok(endpoint);
200        }
201        Ok(endpoint)
202    }
203
204    pub(crate) async fn set_message_sender(&self, sender: MessageDispatcher<R>) {
205        let mut lock = self.message_sender.write().await;
206        *lock = Some(sender);
207    }
208
209    pub(crate) async fn set_error_stream(
210        &self,
211        error_stream: Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>>,
212    ) {
213        let mut lock = self.error_stream.write().await;
214        *lock = Some(IoStream::Readable(error_stream));
215    }
216}
217
218#[async_trait]
219impl<R, S, M, OR, OM> Transport<R, S, M, OR, OM> for ClientSseTransport<M>
220where
221    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
222    S: McpMessage + Clone + Send + Sync + serde::Serialize + 'static,
223    M: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
224    OR: Clone + Send + Sync + serde::Serialize + 'static,
225    OM: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
226{
227    /// Starts the transport, initializing SSE and POST tasks
228    ///
229    /// Sets up the SSE stream, POST request handler, and message streams for communication.
230    ///
231    /// # Returns
232    /// * `TransportResult<(Pin<Box<dyn Stream<Item = R> + Send>>, MessageDispatcher<R>, IoStream)>`
233    ///   - The message stream, dispatcher, and error stream
234    async fn start(&self) -> TransportResult<tokio_stream::wrappers::ReceiverStream<R>>
235    where
236        MessageDispatcher<M>: McpDispatch<R, OR, M, OM>,
237    {
238        // Create CancellationTokenSource and token
239        let (cancellation_source, cancellation_token) = CancellationTokenSource::new();
240        let mut lock = self.shutdown_source.write().await;
241        *lock = Some(cancellation_source);
242
243        let (write_tx, mut write_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
244        let (read_tx, read_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
245
246        // Create oneshot channel for signaling SSE endpoint event message
247        let (endpoint_event_tx, endpoint_event_rx) = oneshot::channel::<Option<String>>();
248        let endpoint_event_tx = Some(endpoint_event_tx);
249
250        let sse_client = self.client.clone();
251        let sse_url = self.sse_url.clone();
252
253        let max_retries = self.max_retries;
254        let retry_delay = self.retry_delay;
255
256        let custom_headers = self.custom_headers.clone();
257
258        let read_stream = SseStream {
259            sse_client,
260            sse_url,
261            max_retries,
262            retry_delay,
263            read_tx,
264        };
265
266        // Spawn task to handle SSE stream with reconnection
267        let cancellation_token_sse = cancellation_token.clone();
268        let sse_task_handle = tokio::spawn(async move {
269            read_stream
270                .run(endpoint_event_tx, cancellation_token_sse, &custom_headers)
271                .await;
272        });
273        let mut sse_task_lock = self.sse_task.write().await;
274        *sse_task_lock = Some(sse_task_handle);
275
276        // Await the first SSE message, expected to receive messages endpoint from he server
277        let err =
278            || std::io::Error::other("Failed to receive 'messages' endpoint from the server.");
279        let post_url = endpoint_event_rx
280            .await
281            .map_err(|_| err())?
282            .ok_or_else(err)?;
283
284        let post_url = self.validate_message_endpoint(post_url)?;
285
286        let client_clone = self.client.clone();
287
288        let custom_headers = self.custom_headers.clone();
289
290        let cancellation_token_post = cancellation_token.clone();
291        // Spawn task to handle POST requests from writable stream
292        let post_task_handle = tokio::spawn(async move {
293            loop {
294                tokio::select! {
295
296                _ = cancellation_token_post.cancelled() =>
297                {
298                        break;
299                },
300
301                data = write_rx.recv() => {
302                    match data{
303                      Some(data) => {
304                        // trim the trailing \n before making a request
305                        let body = String::from_utf8_lossy(&data).trim().to_string();
306                          if let Err(e) = http_post(&client_clone, &post_url, body,None, custom_headers.as_ref()).await {
307                            tracing::error!("Failed to POST message: {e}");
308                      }
309                    },
310                    None => break, // Exit if channel is closed
311                    }
312                   }
313                }
314            }
315        });
316        let mut post_task_lock = self.post_task.write().await;
317        *post_task_lock = Some(post_task_handle);
318
319        // Create writable stream
320        let writable: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>> =
321            Mutex::new(Box::pin(BufWriter::new(WritableChannel { write_tx })));
322
323        // Create readable stream
324        let readable: Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>> =
325            Box::pin(BufReader::new(ReadableChannel {
326                read_rx,
327                buffer: Bytes::new(),
328            }));
329
330        let (stream, sender, error_stream) = MCPStream::create(
331            readable,
332            writable,
333            IoStream::Writable(Box::pin(tokio::io::stderr())),
334            self.pending_requests.clone(),
335            self.request_timeout,
336            self.max_line_length,
337            cancellation_token,
338            self.channel_capacity,
339        );
340
341        self.set_message_sender(sender).await;
342
343        if let IoStream::Readable(error_stream) = error_stream {
344            self.set_error_stream(error_stream).await;
345        }
346
347        Ok(stream)
348    }
349
350    fn message_sender(&self) -> Arc<tokio::sync::RwLock<Option<MessageDispatcher<M>>>> {
351        self.message_sender.clone() as _
352    }
353
354    fn error_stream(&self) -> &tokio::sync::RwLock<Option<IoStream>> {
355        &self.error_stream as _
356    }
357
358    async fn consume_string_payload(&self, _payload: &str) -> TransportResult<()> {
359        Err(TransportError::Internal(
360            "Invalid invocation of consume_string_payload() function for ClientSseTransport"
361                .to_string(),
362        ))
363    }
364
365    async fn keep_alive(
366        &self,
367        _: Duration,
368        _: oneshot::Sender<()>,
369    ) -> TransportResult<JoinHandle<()>> {
370        Err(TransportError::Internal(
371            "Invalid invocation of keep_alive() function for ClientSseTransport".to_string(),
372        ))
373    }
374
375    /// Checks if the transport has been shut down
376    ///
377    /// # Returns
378    /// * `bool` - True if the transport is shut down, false otherwise
379    async fn is_shut_down(&self) -> bool {
380        let result = self.is_shut_down.lock().await;
381        *result
382    }
383
384    // Shuts down the transport, terminating any subprocess and signaling closure.
385    ///
386    /// Sends a shutdown signal via the watch channel and kills the subprocess if present.
387    ///
388    /// # Returns
389    /// A `TransportResult` indicating success or failure.
390    ///
391    /// # Errors
392    /// Returns a `TransportError` if the shutdown signal fails or the process cannot be killed.
393    async fn shut_down(&self) -> TransportResult<()> {
394        // Trigger cancellation
395        let mut cancellation_lock = self.shutdown_source.write().await;
396        if let Some(source) = cancellation_lock.as_ref() {
397            source.cancel()?;
398        }
399        *cancellation_lock = None; // Clear cancellation_source
400
401        // Mark as shut down
402        let mut is_shut_down_lock = self.is_shut_down.lock().await;
403        *is_shut_down_lock = true;
404
405        // Get task handles
406        let sse_task = self.sse_task.write().await.take();
407        let post_task = self.post_task.write().await.take();
408
409        // Wait for tasks to complete with a timeout
410        let timeout = Duration::from_secs(SHUTDOWN_TIMEOUT_SECONDS);
411        let shutdown_future = async {
412            if let Some(post_handle) = post_task {
413                let _ = post_handle.await;
414            }
415            if let Some(sse_handle) = sse_task {
416                let _ = sse_handle.await;
417            }
418            Ok::<(), TransportError>(())
419        };
420
421        tokio::select! {
422            result = shutdown_future => {
423                result // result of task completion
424            }
425            _ = tokio::time::sleep(timeout) => {
426                tracing::warn!("Shutdown timed out after {:?}", timeout);
427                Err(TransportError::ShutdownTimeout)
428            }
429        }
430    }
431
432    async fn pending_request_tx(&self, request_id: &RequestId) -> Option<Sender<M>> {
433        let mut pending_requests = self.pending_requests.lock().await;
434        pending_requests.remove(request_id)
435    }
436}
437
438#[async_trait]
439impl McpDispatch<ServerMessages, ClientMessages, ServerMessage, ClientMessage>
440    for ClientSseTransport<ServerMessage>
441{
442    async fn send_message(
443        &self,
444        message: ClientMessages,
445        request_timeout: Option<Duration>,
446    ) -> TransportResult<Option<ServerMessages>> {
447        let sender = self.message_sender.read().await;
448        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
449        sender.send_message(message, request_timeout).await
450    }
451
452    async fn send(
453        &self,
454        message: ClientMessage,
455        request_timeout: Option<Duration>,
456    ) -> TransportResult<Option<ServerMessage>> {
457        let sender = self.message_sender.read().await;
458        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
459        sender.send(message, request_timeout).await
460    }
461
462    async fn send_batch(
463        &self,
464        message: Vec<ClientMessage>,
465        request_timeout: Option<Duration>,
466    ) -> TransportResult<Option<Vec<ServerMessage>>> {
467        let sender = self.message_sender.read().await;
468        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
469        sender.send_batch(message, request_timeout).await
470    }
471
472    async fn write_str(&self, payload: &str, skip_store: bool) -> TransportResult<()> {
473        let sender = self.message_sender.read().await;
474        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
475        sender.write_str(payload, skip_store).await
476    }
477}
478
479impl
480    TransportDispatcher<
481        ServerMessages,
482        MessageFromClient,
483        ServerMessage,
484        ClientMessages,
485        ClientMessage,
486    > for ClientSseTransport<ServerMessage>
487{
488}
rust_mcp_transport/client_sse.rs

rust_mcp_transport/
client_sse.rs
