rust_mcp_transport/transport.rs
1use std::{pin::Pin, time::Duration};
2
3use async_trait::async_trait;
4use rust_mcp_schema::{schema_utils::McpMessage, RequestId};
5
6use futures::Stream;
7
8use crate::{error::TransportResult, message_dispatcher::MessageDispatcher};
9
10/// Default Timeout in milliseconds
11const DEFAULT_TIMEOUT_MSEC: u64 = 60_000;
12
13/// Enum representing a stream that can either be readable or writable.
14/// This allows the reuse of the same traits for both MCP Server and MCP Client,
15/// where the data direction is reversed.
16///
17/// It encapsulates two types of I/O streams:
18/// - `Readable`: A stream that implements the `AsyncRead` trait for reading data asynchronously.
19/// - `Writable`: A stream that implements the `AsyncWrite` trait for writing data asynchronously.
20///
21pub enum IoStream {
22 Readable(Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>>),
23 Writable(Pin<Box<dyn tokio::io::AsyncWrite + Send + Sync>>),
24}
25
26/// Configuration for the transport layer
27#[derive(Debug, Clone)]
28pub struct TransportOptions {
29 /// The timeout in milliseconds for requests.
30 ///
31 /// This value defines the maximum amount of time to wait for a response before
32 /// considering the request as timed out.
33 pub timeout: Duration,
34}
35impl Default for TransportOptions {
36 fn default() -> Self {
37 Self {
38 timeout: Duration::from_millis(DEFAULT_TIMEOUT_MSEC),
39 }
40 }
41}
42
43/// A trait for sending MCP messages.
44///
45///It is intended to be implemented by types that send messages in the MCP protocol, such as servers or clients.
46///
47/// The `McpDispatch` trait requires two associated types:
48/// - `R`: The type of the response, which must implement the `McpMessage` trait and be capable of deserialization.
49/// - `S`: The type of the message to send, which must be serializable and cloneable.
50///
51/// Both associated types `R` and `S` must be `Send`, `Sync`, and `'static` to ensure they can be used
52/// safely in an asynchronous context and across threads.
53///
54/// # Associated Types
55///
56/// - `R`: The response type, which must implement the `McpMessage` trait, be `Clone`, `Send`, `Sync`, and
57/// be deserializable (`DeserializeOwned`).
58/// - `S`: The type of the message to send, which must be `Clone`, `Send`, `Sync`, and serializable (`Serialize`).
59///
60/// # Methods
61///
62/// ### `send`
63///
64/// Sends a raw message represented by type `S` and optionally includes a `request_id`.
65/// The method returns a `TransportResult<Option<R>>`, where:
66/// - `Option<R>`: The response, which can be `None` or contain the response of type `R`.
67/// - `TransportResult`: Represents the result of the operation, which can include success or failure.
68///
69/// # Arguments
70/// - `message`: The message to send, of type `S`, which will be serialized before transmission.
71/// - `request_id`: An optional `RequestId` to associate with this message. It can be used for tracking
72/// or correlating the request with its response.
73///
74/// # Example
75///
76/// let sender: Box<dyn McpDispatch<MyResponse, MyMessage>> = ...;
77/// let result = sender.send(my_message, Some(request_id)).await;
78///
79#[async_trait]
80pub trait McpDispatch<R, S>: Send + Sync + 'static
81where
82 R: McpMessage + Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
83 S: Clone + Send + Sync + serde::Serialize + 'static,
84{
85 /// Sends a raw message represented by type `S` and optionally includes a `request_id`.
86 /// The `request_id` is used when sending a message in response to an MCP request.
87 /// It should match the `request_id` of the original request.
88 async fn send(
89 &self,
90 message: S,
91 request_id: Option<RequestId>,
92 request_timeout: Option<Duration>,
93 ) -> TransportResult<Option<R>>;
94}
95
96/// A trait representing the transport layer for MCP.
97///
98/// This trait is designed for handling the transport of messages within an MCP protocol system. It
99/// provides a method to start the transport process, which involves setting up a stream, a message sender,
100/// and handling I/O operations.
101///
102/// The `Transport` trait requires three associated types:
103/// - `R`: The message type to send, which must implement the `McpMessage` trait.
104/// - `S`: The message type to send.
105/// - `M`: The type of message that we expect to receive as a response to the sent message.
106///
107#[async_trait]
108pub trait Transport<R, S>: Send + Sync + 'static
109where
110 R: McpMessage + Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
111 S: Clone + Send + Sync + serde::Serialize + 'static,
112{
113 async fn start(
114 &self,
115 ) -> TransportResult<(
116 Pin<Box<dyn Stream<Item = R> + Send>>,
117 MessageDispatcher<R>,
118 IoStream,
119 )>
120 where
121 MessageDispatcher<R>: McpDispatch<R, S>;
122 async fn shut_down(&self) -> TransportResult<()>;
123 async fn is_shut_down(&self) -> bool;
124}