rust_mcp_transport/
transport.rs

1use std::pin::Pin;
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
27pub struct TransportOptions {
28    /// The timeout in milliseconds for requests.
29    ///
30    /// This value defines the maximum amount of time to wait for a response before
31    /// considering the request as timed out.
32    pub timeout: u64,
33}
34impl Default for TransportOptions {
35    fn default() -> Self {
36        Self {
37            timeout: DEFAULT_TIMEOUT_MSEC,
38        }
39    }
40}
41
42/// A trait for sending MCP messages.
43///
44///It is intended to be implemented by types that send messages in the MCP protocol, such as servers or clients.
45///
46/// The `McpDispatch` trait requires two associated types:
47/// - `R`: The type of the response, which must implement the `MCPMessage` trait and be capable of deserialization.
48/// - `S`: The type of the message to send, which must be serializable and cloneable.
49///
50/// Both associated types `R` and `S` must be `Send`, `Sync`, and `'static` to ensure they can be used
51/// safely in an asynchronous context and across threads.
52///
53/// # Associated Types
54///
55/// - `R`: The response type, which must implement the `MCPMessage` trait, be `Clone`, `Send`, `Sync`, and
56///   be deserializable (`DeserializeOwned`).
57/// - `S`: The type of the message to send, which must be `Clone`, `Send`, `Sync`, and serializable (`Serialize`).
58///
59/// # Methods
60///
61/// ### `send`
62///
63/// Sends a raw message represented by type `S` and optionally includes a `request_id`.
64/// The method returns a `TransportResult<Option<R>>`, where:
65/// - `Option<R>`: The response, which can be `None` or contain the response of type `R`.
66/// - `TransportResult`: Represents the result of the operation, which can include success or failure.
67///
68/// # Arguments
69/// - `message`: The message to send, of type `S`, which will be serialized before transmission.
70/// - `request_id`: An optional `RequestId` to associate with this message. It can be used for tracking
71///   or correlating the request with its response.
72///
73/// # Example
74///
75/// let sender: Box<dyn McpDispatch<MyResponse, MyMessage>> = ...;
76/// let result = sender.send(my_message, Some(request_id)).await;
77///
78#[async_trait]
79pub trait McpDispatch<R, S>: Send + Sync + 'static
80where
81    R: MCPMessage + Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
82    S: Clone + Send + Sync + serde::Serialize + 'static,
83{
84    /// Sends a raw message represented by type `S` and optionally includes a `request_id`.
85    /// The `request_id` is used when sending a message in response to an MCP request.
86    /// It should match the `request_id` of the original request.
87    async fn send(&self, message: S, request_id: Option<RequestId>) -> TransportResult<Option<R>>;
88}
89
90/// A trait representing the transport layer for MCP.
91///
92/// This trait is designed for handling the transport of messages within an MCP protocol system. It
93/// provides a method to start the transport process, which involves setting up a stream, a message sender,
94/// and handling I/O operations.
95///
96/// The `Transport` trait requires three associated types:
97/// - `R`: The message type to send, which must implement the `MCPMessage` trait.
98/// - `S`: The message type to send.
99/// - `M`: The type of message that we expect to receive as a response to the sent message.
100///
101#[async_trait]
102pub trait Transport<R, S>: Send + Sync + 'static
103where
104    R: MCPMessage + Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
105    S: Clone + Send + Sync + serde::Serialize + 'static,
106{
107    async fn start(
108        &self,
109    ) -> TransportResult<(
110        Pin<Box<dyn Stream<Item = R> + Send>>,
111        MessageDispatcher<R>,
112        IoStream,
113    )>
114    where
115        MessageDispatcher<R>: McpDispatch<R, S>;
116    async fn shut_down(&self) -> TransportResult<()>;
117    async fn is_shut_down(&self) -> bool;
118}