rust_mcp_transport/
client_streamable_http.rs

1use crate::error::TransportError;
2use crate::mcp_stream::MCPStream;
3
4use crate::schema::{
5    schema_utils::{
6        ClientMessage, ClientMessages, McpMessage, MessageFromClient, SdkError, ServerMessage,
7        ServerMessages,
8    },
9    RequestId,
10};
11use crate::utils::{
12    http_delete, http_post, CancellationTokenSource, ReadableChannel, StreamableHttpStream,
13    WritableChannel,
14};
15use crate::{error::TransportResult, IoStream, McpDispatch, MessageDispatcher, Transport};
16use crate::{SessionId, TransportDispatcher, TransportOptions};
17use async_trait::async_trait;
18use bytes::Bytes;
19use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
20use reqwest::Client;
21use std::collections::HashMap;
22use std::pin::Pin;
23use std::{sync::Arc, time::Duration};
24use tokio::io::{BufReader, BufWriter};
25use tokio::sync::oneshot::Sender;
26use tokio::sync::{mpsc, oneshot, Mutex};
27use tokio::task::JoinHandle;
28
29const DEFAULT_CHANNEL_CAPACITY: usize = 64;
30const DEFAULT_MAX_RETRY: usize = 5;
31const DEFAULT_RETRY_TIME_SECONDS: u64 = 1;
32const SHUTDOWN_TIMEOUT_SECONDS: u64 = 5;
33
34pub struct StreamableTransportOptions {
35    pub mcp_url: String,
36    pub request_options: RequestOptions,
37}
38
39impl StreamableTransportOptions {
40    pub async fn terminate_session(&self, session_id: Option<&SessionId>) {
41        let client = Client::new();
42        match http_delete(&client, &self.mcp_url, session_id, None).await {
43            Ok(_) => {}
44            Err(TransportError::Http(status_code)) => {
45                tracing::info!("Session termination failed with status code {status_code}",);
46            }
47            Err(error) => {
48                tracing::info!("Session termination failed with error :{error}");
49            }
50        };
51    }
52}
53
54pub struct RequestOptions {
55    pub request_timeout: Duration,
56    pub max_line_length: usize,
57    pub channel_capacity: usize,
58    pub retry_delay: Option<Duration>,
59    pub max_retries: Option<usize>,
60    pub custom_headers: Option<HashMap<String, String>>,
61}
62
63impl Default for RequestOptions {
64    fn default() -> Self {
65        Self {
66            request_timeout: TransportOptions::default().timeout,
67            max_line_length: TransportOptions::default().max_line_length,
68            channel_capacity: TransportOptions::default().channel_capacity,
69            retry_delay: None,
70            max_retries: None,
71            custom_headers: None,
72        }
73    }
74}
75
76pub struct ClientStreamableTransport<R>
77where
78    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
79{
80    /// Optional cancellation token source for shutting down the transport
81    shutdown_source: tokio::sync::RwLock<Option<CancellationTokenSource>>,
82    /// Flag indicating if the transport is shut down
83    is_shut_down: Mutex<bool>,
84    /// Timeout duration for MCP messages
85    request_timeout: Duration,
86    /// Maximum line length for incoming messages
87    max_line_length: usize,
88    /// Capacity of the incoming-message channel buffer
89    channel_capacity: usize,
90    /// HTTP client for making requests
91    client: Client,
92    /// URL for the SSE endpoint
93    mcp_server_url: String,
94    /// Delay between retry attempts
95    retry_delay: Duration,
96    /// Maximum number of retry attempts
97    max_retries: usize,
98    /// Optional custom HTTP headers
99    custom_headers: Option<HeaderMap>,
100    sse_task: tokio::sync::RwLock<Option<tokio::task::JoinHandle<()>>>,
101    post_task: tokio::sync::RwLock<Option<tokio::task::JoinHandle<()>>>,
102    message_sender: Arc<tokio::sync::RwLock<Option<MessageDispatcher<R>>>>,
103    error_stream: tokio::sync::RwLock<Option<IoStream>>,
104    pending_requests: Arc<Mutex<HashMap<RequestId, tokio::sync::oneshot::Sender<R>>>>,
105    session_id: Arc<tokio::sync::RwLock<Option<SessionId>>>,
106    standalone: bool,
107}
108
109impl<R> ClientStreamableTransport<R>
110where
111    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
112{
113    pub fn new(
114        options: &StreamableTransportOptions,
115        session_id: Option<SessionId>,
116        standalone: bool,
117    ) -> TransportResult<Self> {
118        let client = Client::new();
119
120        let headers = match &options.request_options.custom_headers {
121            Some(h) => Some(Self::validate_headers(h)?),
122            None => None,
123        };
124
125        let mcp_server_url = options.mcp_url.to_owned();
126        Ok(Self {
127            shutdown_source: tokio::sync::RwLock::new(None),
128            is_shut_down: Mutex::new(false),
129            request_timeout: options.request_options.request_timeout,
130            max_line_length: options.request_options.max_line_length,
131            channel_capacity: options.request_options.channel_capacity,
132            client,
133            mcp_server_url,
134            retry_delay: options
135                .request_options
136                .retry_delay
137                .unwrap_or(Duration::from_secs(DEFAULT_RETRY_TIME_SECONDS)),
138            max_retries: options
139                .request_options
140                .max_retries
141                .unwrap_or(DEFAULT_MAX_RETRY),
142            sse_task: tokio::sync::RwLock::new(None),
143            post_task: tokio::sync::RwLock::new(None),
144            custom_headers: headers,
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            session_id: Arc::new(tokio::sync::RwLock::new(session_id)),
149            standalone,
150        })
151    }
152
153    fn validate_headers(headers: &HashMap<String, String>) -> TransportResult<HeaderMap> {
154        let mut header_map = HeaderMap::new();
155        for (key, value) in headers {
156            let header_name =
157                key.parse::<HeaderName>()
158                    .map_err(|e| TransportError::Configuration {
159                        message: format!("Invalid header name: {e}"),
160                    })?;
161            let header_value =
162                HeaderValue::from_str(value).map_err(|e| TransportError::Configuration {
163                    message: format!("Invalid header value: {e}"),
164                })?;
165            header_map.insert(header_name, header_value);
166        }
167        Ok(header_map)
168    }
169
170    pub(crate) async fn set_message_sender(&self, sender: MessageDispatcher<R>) {
171        let mut lock = self.message_sender.write().await;
172        *lock = Some(sender);
173    }
174
175    pub(crate) async fn set_error_stream(
176        &self,
177        error_stream: Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>>,
178    ) {
179        let mut lock = self.error_stream.write().await;
180        *lock = Some(IoStream::Readable(error_stream));
181    }
182}
183
184#[async_trait]
185impl<R, S, M, OR, OM> Transport<R, S, M, OR, OM> for ClientStreamableTransport<M>
186where
187    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
188    S: McpMessage + Clone + Send + Sync + serde::Serialize + 'static,
189    M: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
190    OR: Clone + Send + Sync + serde::Serialize + 'static,
191    OM: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
192{
193    async fn start(&self) -> TransportResult<tokio_stream::wrappers::ReceiverStream<R>>
194    where
195        MessageDispatcher<M>: McpDispatch<R, OR, M, OM>,
196    {
197        if self.standalone {
198            // Create CancellationTokenSource and token
199            let (cancellation_source, cancellation_token) = CancellationTokenSource::new();
200            let mut lock = self.shutdown_source.write().await;
201            *lock = Some(cancellation_source);
202
203            let (write_tx, mut write_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
204            let (read_tx, read_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
205
206            let max_retries = self.max_retries;
207            let retry_delay = self.retry_delay;
208
209            let post_url = self.mcp_server_url.clone();
210            let custom_headers = self.custom_headers.clone();
211            let cancellation_token_post = cancellation_token.clone();
212            let cancellation_token_sse = cancellation_token.clone();
213
214            let session_id_clone = self.session_id.clone();
215
216            let mut streamable_http = StreamableHttpStream {
217                client: self.client.clone(),
218                mcp_url: post_url,
219                max_retries,
220                retry_delay,
221                read_tx,
222                session_id: session_id_clone, //Arc<RwLock<Option<String>>>
223            };
224
225            let session_id = self.session_id.read().await.to_owned();
226
227            let sse_response = streamable_http
228                .make_standalone_stream_connection(&cancellation_token_sse, &custom_headers, None)
229                .await?;
230
231            let sse_task_handle = tokio::spawn(async move {
232                if let Err(error) = streamable_http
233                    .run_standalone(&cancellation_token_sse, &custom_headers, sse_response)
234                    .await
235                {
236                    if !matches!(error, TransportError::Cancelled(_)) {
237                        tracing::warn!("{error}");
238                    }
239                }
240            });
241
242            let mut sse_task_lock = self.sse_task.write().await;
243            *sse_task_lock = Some(sse_task_handle);
244
245            let post_url = self.mcp_server_url.clone();
246            let client = self.client.clone();
247            let custom_headers = self.custom_headers.clone();
248
249            // Initiate a task to process POST requests from messages received via the writable stream.
250            let post_task_handle = tokio::spawn(async move {
251                loop {
252                    tokio::select! {
253                    _ = cancellation_token_post.cancelled() =>
254                    {
255                            break;
256                    },
257                    data = write_rx.recv() => {
258                        match data{
259                          Some(data) => {
260                              // trim the trailing \n before making a request
261                              let payload = String::from_utf8_lossy(&data).trim().to_string();
262
263                             if let Err(e) = http_post(
264                                  &client,
265                                  &post_url,
266                                  payload.to_string(),
267                                  session_id.as_ref(),
268                                  custom_headers.as_ref(),
269                              )
270                              .await{
271                                tracing::error!("Failed to POST message: {e}")
272                          }
273                        },
274                        None => break, // Exit if channel is closed
275                        }
276                       }
277                    }
278                }
279            });
280            let mut post_task_lock = self.post_task.write().await;
281            *post_task_lock = Some(post_task_handle);
282
283            // Create writable stream
284            let writable: Mutex<Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>> =
285                Mutex::new(Box::pin(BufWriter::new(WritableChannel { write_tx })));
286
287            // Create readable stream
288            let readable: Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>> =
289                Box::pin(BufReader::new(ReadableChannel {
290                    read_rx,
291                    buffer: Bytes::new(),
292                }));
293
294            let (stream, sender, error_stream) = MCPStream::create(
295                readable,
296                writable,
297                IoStream::Writable(Box::pin(tokio::io::stderr())),
298                self.pending_requests.clone(),
299                self.request_timeout,
300                self.max_line_length,
301                cancellation_token,
302                self.channel_capacity,
303            );
304
305            self.set_message_sender(sender).await;
306
307            if let IoStream::Readable(error_stream) = error_stream {
308                self.set_error_stream(error_stream).await;
309            }
310            Ok(stream)
311        } else {
312            // Create CancellationTokenSource and token
313            let (cancellation_source, cancellation_token) = CancellationTokenSource::new();
314            let mut lock = self.shutdown_source.write().await;
315            *lock = Some(cancellation_source);
316
317            // let (write_tx, mut write_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
318            let (write_tx, mut write_rx): (
319                tokio::sync::mpsc::Sender<(
320                    String,
321                    tokio::sync::oneshot::Sender<crate::error::TransportResult<()>>,
322                )>,
323                tokio::sync::mpsc::Receiver<(
324                    String,
325                    tokio::sync::oneshot::Sender<crate::error::TransportResult<()>>,
326                )>,
327            ) = tokio::sync::mpsc::channel(DEFAULT_CHANNEL_CAPACITY); // Buffer size as needed
328            let (read_tx, read_rx) = mpsc::channel::<Bytes>(DEFAULT_CHANNEL_CAPACITY);
329
330            let max_retries = self.max_retries;
331            let retry_delay = self.retry_delay;
332
333            let post_url = self.mcp_server_url.clone();
334            let custom_headers = self.custom_headers.clone();
335            let cancellation_token_post = cancellation_token.clone();
336            let cancellation_token_sse = cancellation_token.clone();
337
338            let session_id_clone = self.session_id.clone();
339
340            let mut streamable_http = StreamableHttpStream {
341                client: self.client.clone(),
342                mcp_url: post_url,
343                max_retries,
344                retry_delay,
345                read_tx,
346                session_id: session_id_clone, //Arc<RwLock<Option<String>>>
347            };
348
349            // Initiate a task to process POST requests from messages received via the writable stream.
350            let post_task_handle = tokio::spawn(async move {
351                loop {
352                    tokio::select! {
353                    _ = cancellation_token_post.cancelled() =>
354                    {
355                            break;
356                    },
357                    data = write_rx.recv() => {
358                        match data{
359                          Some((data, ack_tx)) => {
360                            // trim the trailing \n before making a request
361                            let payload = data.trim().to_string();
362                            let result = streamable_http.run(payload, &cancellation_token_sse, &custom_headers).await;
363                            let _ = ack_tx.send(result);// Ignore error if receiver dropped
364                        },
365                        None => break, // Exit if channel is closed
366                        }
367                       }
368                    }
369                }
370            });
371            let mut post_task_lock = self.post_task.write().await;
372            *post_task_lock = Some(post_task_handle);
373
374            // Create readable stream
375            let readable: Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>> =
376                Box::pin(BufReader::new(ReadableChannel {
377                    read_rx,
378                    buffer: Bytes::new(),
379                }));
380
381            let (stream, sender, error_stream) = MCPStream::create_with_ack(
382                readable,
383                write_tx,
384                IoStream::Writable(Box::pin(tokio::io::stderr())),
385                self.pending_requests.clone(),
386                self.request_timeout,
387                self.max_line_length,
388                cancellation_token,
389                self.channel_capacity,
390            );
391
392            self.set_message_sender(sender).await;
393
394            if let IoStream::Readable(error_stream) = error_stream {
395                self.set_error_stream(error_stream).await;
396            }
397
398            Ok(stream)
399        }
400    }
401
402    fn message_sender(&self) -> Arc<tokio::sync::RwLock<Option<MessageDispatcher<M>>>> {
403        self.message_sender.clone() as _
404    }
405
406    fn error_stream(&self) -> &tokio::sync::RwLock<Option<IoStream>> {
407        &self.error_stream as _
408    }
409    async fn shut_down(&self) -> TransportResult<()> {
410        // Trigger cancellation
411        let mut cancellation_lock = self.shutdown_source.write().await;
412        if let Some(source) = cancellation_lock.as_ref() {
413            source.cancel()?;
414        }
415        *cancellation_lock = None; // Clear cancellation_source
416
417        // Mark as shut down
418        let mut is_shut_down_lock = self.is_shut_down.lock().await;
419        *is_shut_down_lock = true;
420
421        // Get task handle
422        let post_task = self.post_task.write().await.take();
423
424        // // Wait for tasks to complete with a timeout
425        let timeout = Duration::from_secs(SHUTDOWN_TIMEOUT_SECONDS);
426        let shutdown_future = async {
427            if let Some(post_handle) = post_task {
428                let _ = post_handle.await;
429            }
430            Ok::<(), TransportError>(())
431        };
432
433        tokio::select! {
434            result = shutdown_future => {
435                result // result of task completion
436            }
437            _ = tokio::time::sleep(timeout) => {
438                tracing::warn!("Shutdown timed out after {:?}", timeout);
439                Err(TransportError::ShutdownTimeout)
440            }
441        }
442    }
443    async fn is_shut_down(&self) -> bool {
444        let result = self.is_shut_down.lock().await;
445        *result
446    }
447    async fn consume_string_payload(&self, _: &str) -> TransportResult<()> {
448        Err(TransportError::Internal(
449            "Invalid invocation of consume_string_payload() function for ClientStreamableTransport"
450                .to_string(),
451        ))
452    }
453
454    async fn pending_request_tx(&self, request_id: &RequestId) -> Option<Sender<M>> {
455        let mut pending_requests = self.pending_requests.lock().await;
456        pending_requests.remove(request_id)
457    }
458
459    async fn keep_alive(
460        &self,
461        _: Duration,
462        _: oneshot::Sender<()>,
463    ) -> TransportResult<JoinHandle<()>> {
464        Err(TransportError::Internal(
465            "Invalid invocation of keep_alive() function for ClientStreamableTransport".to_string(),
466        ))
467    }
468
469    async fn session_id(&self) -> Option<SessionId> {
470        let guard = self.session_id.read().await;
471        guard.clone()
472    }
473}
474
475#[async_trait]
476impl McpDispatch<ServerMessages, ClientMessages, ServerMessage, ClientMessage>
477    for ClientStreamableTransport<ServerMessage>
478{
479    async fn send_message(
480        &self,
481        message: ClientMessages,
482        request_timeout: Option<Duration>,
483    ) -> TransportResult<Option<ServerMessages>> {
484        let sender = self.message_sender.read().await;
485
486        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
487
488        sender.send_message(message, request_timeout).await
489    }
490
491    async fn send(
492        &self,
493        message: ClientMessage,
494        request_timeout: Option<Duration>,
495    ) -> TransportResult<Option<ServerMessage>> {
496        let sender = self.message_sender.read().await;
497
498        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
499
500        sender.send(message, request_timeout).await
501    }
502
503    async fn send_batch(
504        &self,
505        message: Vec<ClientMessage>,
506        request_timeout: Option<Duration>,
507    ) -> TransportResult<Option<Vec<ServerMessage>>> {
508        let sender = self.message_sender.read().await;
509        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
510        sender.send_batch(message, request_timeout).await
511    }
512
513    async fn write_str(&self, payload: &str, skip_store: bool) -> TransportResult<()> {
514        let sender = self.message_sender.read().await;
515        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
516        sender.write_str(payload, skip_store).await
517    }
518}
519
520impl
521    TransportDispatcher<
522        ServerMessages,
523        MessageFromClient,
524        ServerMessage,
525        ClientMessages,
526        ClientMessage,
527    > for ClientStreamableTransport<ServerMessage>
528{
529}
rust_mcp_transport/client_streamable_http.rs

rust_mcp_transport/
client_streamable_http.rs
