rust_mcp_transport/
stdio.rs

1use crate::schema::schema_utils::{
2    ClientMessage, ClientMessages, MessageFromClient, MessageFromServer, SdkError, ServerMessage,
3    ServerMessages,
4};
5use crate::schema::RequestId;
6use async_trait::async_trait;
7use serde::de::DeserializeOwned;
8use std::collections::HashMap;
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::process::Command;
12use tokio::sync::oneshot::Sender;
13use tokio::sync::{oneshot, Mutex};
14use tokio::task::JoinHandle;
15
16use crate::error::{TransportError, TransportResult};
17use crate::mcp_stream::MCPStream;
18use crate::message_dispatcher::MessageDispatcher;
19use crate::transport::Transport;
20use crate::utils::CancellationTokenSource;
21use crate::{IoStream, McpDispatch, TransportDispatcher, TransportOptions};
22
23/// Implements a standard I/O transport for MCP communication.
24///
25/// This module provides the `StdioTransport` struct, which serves as a transport layer for the
26/// Model Context Protocol (MCP) using standard input/output (stdio). It supports both client-side
27/// and server-side communication by optionally launching a subprocess or using the current
28/// process's stdio streams. The transport handles message streaming, dispatching, and shutdown
29/// operations, integrating with the MCP runtime ecosystem.
30pub struct StdioTransport<R>
31where
32    R: Clone + Send + Sync + DeserializeOwned + 'static,
33{
34    command: Option<String>,
35    args: Option<Vec<String>>,
36    env: Option<HashMap<String, String>>,
37    options: TransportOptions,
38    shutdown_source: tokio::sync::RwLock<Option<CancellationTokenSource>>,
39    is_shut_down: Mutex<bool>,
40    message_sender: Arc<tokio::sync::RwLock<Option<MessageDispatcher<R>>>>,
41    error_stream: tokio::sync::RwLock<Option<IoStream>>,
42    pending_requests: Arc<Mutex<HashMap<RequestId, tokio::sync::oneshot::Sender<R>>>>,
43}
44
45impl<R> StdioTransport<R>
46where
47    R: Clone + Send + Sync + DeserializeOwned + 'static,
48{
49    /// Creates a new `StdioTransport` instance for MCP Server.
50    ///
51    /// This constructor configures the transport to use the current process's stdio streams,
52    ///
53    /// # Arguments
54    /// * `options` - Configuration options for the transport, including timeout settings.
55    ///
56    /// # Returns
57    /// A `TransportResult` containing the initialized `StdioTransport` instance.
58    ///
59    /// # Errors
60    /// Currently, this method does not fail, but it returns a `TransportResult` for API consistency.
61    pub fn new(options: TransportOptions) -> TransportResult<Self> {
62        Ok(Self {
63            // when transport is used for MCP Server, we do not need a command
64            args: None,
65            command: None,
66            env: None,
67            options,
68            shutdown_source: tokio::sync::RwLock::new(None),
69            is_shut_down: Mutex::new(false),
70            message_sender: Arc::new(tokio::sync::RwLock::new(None)),
71            error_stream: tokio::sync::RwLock::new(None),
72            pending_requests: Arc::new(Mutex::new(HashMap::new())),
73        })
74    }
75
76    /// Creates a new `StdioTransport` instance with a subprocess for MCP Client use.
77    ///
78    /// This constructor configures the transport to launch a MCP Server with a specified command
79    /// arguments and optional environment variables
80    ///
81    /// # Arguments
82    /// * `command` - The command to execute (e.g., "rust-mcp-filesystem").
83    /// * `args` - Arguments to pass to the command. (e.g., "~/Documents").
84    /// * `env` - Optional environment variables for the subprocess.
85    /// * `options` - Configuration options for the transport, including timeout settings.
86    ///
87    /// # Returns
88    /// A `TransportResult` containing the initialized `StdioTransport` instance, ready to launch
89    /// the MCP server on `start`.
90    pub fn create_with_server_launch<C: Into<String>>(
91        command: C,
92        args: Vec<String>,
93        env: Option<HashMap<String, String>>,
94        options: TransportOptions,
95    ) -> TransportResult<Self> {
96        Ok(Self {
97            // when transport is used for MCP Server, we do not need a command
98            args: Some(args),
99            command: Some(command.into()),
100            env,
101            options,
102            shutdown_source: tokio::sync::RwLock::new(None),
103            is_shut_down: Mutex::new(false),
104            message_sender: Arc::new(tokio::sync::RwLock::new(None)),
105            error_stream: tokio::sync::RwLock::new(None),
106            pending_requests: Arc::new(Mutex::new(HashMap::new())),
107        })
108    }
109
110    /// Retrieves the command and arguments for launching the subprocess.
111    ///
112    /// Adjusts the command based on the platform: on Windows, wraps it with `cmd.exe /c`.
113    ///
114    /// # Returns
115    /// A tuple of the command string and its arguments.
116    fn launch_commands(&self) -> (String, Vec<std::string::String>) {
117        #[cfg(windows)]
118        {
119            let command = "cmd.exe".to_string();
120            let mut command_args = vec!["/c".to_string(), self.command.clone().unwrap_or_default()];
121            command_args.extend(self.args.clone().unwrap_or_default());
122            (command, command_args)
123        }
124
125        #[cfg(unix)]
126        {
127            let command = self.command.clone().unwrap_or_default();
128            let command_args = self.args.clone().unwrap_or_default();
129            (command, command_args)
130        }
131    }
132
133    pub(crate) async fn set_message_sender(&self, sender: MessageDispatcher<R>) {
134        let mut lock = self.message_sender.write().await;
135        *lock = Some(sender);
136    }
137
138    pub(crate) async fn set_error_stream(&self, error_stream: IoStream) {
139        let mut lock = self.error_stream.write().await;
140        *lock = Some(error_stream);
141    }
142}
143
144#[async_trait]
145impl<R, S, M, OR, OM> Transport<R, S, M, OR, OM> for StdioTransport<M>
146where
147    R: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
148    S: Clone + Send + Sync + serde::Serialize + 'static,
149    M: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
150    OR: Clone + Send + Sync + serde::Serialize + 'static,
151    OM: Clone + Send + Sync + serde::de::DeserializeOwned + 'static,
152{
153    /// Starts the transport, initializing streams and the message dispatcher.
154    ///
155    /// If configured with a command (MCP Client), launches the MCP server and connects its stdio streams.
156    /// Otherwise, uses the current process's stdio for server-side communication.
157    ///
158    /// # Returns
159    /// A `TransportResult` containing:
160    /// - A pinned stream of incoming messages.
161    /// - A `MessageDispatcher<R>` for sending messages.
162    /// - An `IoStream` for stderr (readable) or stdout (writable) depending on the mode.
163    ///
164    /// # Errors
165    /// Returns a `TransportError` if the subprocess fails to spawn or stdio streams cannot be accessed.
166    async fn start(&self) -> TransportResult<tokio_stream::wrappers::ReceiverStream<R>>
167    where
168        MessageDispatcher<M>: McpDispatch<R, OR, M, OM>,
169    {
170        // Create CancellationTokenSource and token
171        let (cancellation_source, cancellation_token) = CancellationTokenSource::new();
172        let mut lock = self.shutdown_source.write().await;
173        *lock = Some(cancellation_source);
174
175        if self.command.is_some() {
176            let (command_name, command_args) = self.launch_commands();
177
178            let mut command = Command::new(command_name);
179            command
180                .envs(self.env.as_ref().unwrap_or(&HashMap::new()))
181                .args(&command_args)
182                .stdout(std::process::Stdio::piped())
183                .stdin(std::process::Stdio::piped())
184                .stderr(std::process::Stdio::piped())
185                .kill_on_drop(true);
186
187            #[cfg(windows)]
188            command.creation_flags(0x08000000); // https://learn.microsoft.com/en-us/windows/win32/procthread/process-creation-flags
189
190            #[cfg(unix)]
191            command.process_group(0);
192
193            let mut process = command.spawn().map_err(TransportError::Io)?;
194
195            let stdin = process
196                .stdin
197                .take()
198                .ok_or_else(|| TransportError::Internal("Unable to retrieve stdin.".into()))?;
199
200            let stdout = process
201                .stdout
202                .take()
203                .ok_or_else(|| TransportError::Internal("Unable to retrieve stdout.".into()))?;
204
205            let stderr = process
206                .stderr
207                .take()
208                .ok_or_else(|| TransportError::Internal("Unable to retrieve stderr.".into()))?;
209
210            let pending_requests_clone = self.pending_requests.clone();
211
212            tokio::spawn(async move {
213                let _ = process.wait().await;
214                // clean up pending requests to cancel waiting tasks
215                let mut pending_requests = pending_requests_clone.lock().await;
216                pending_requests.clear();
217            });
218
219            let (stream, sender, error_stream) = MCPStream::create(
220                Box::pin(stdout),
221                Mutex::new(Box::pin(stdin)),
222                IoStream::Readable(Box::pin(stderr)),
223                self.pending_requests.clone(),
224                self.options.timeout,
225                self.options.max_line_length,
226                cancellation_token,
227                self.options.channel_capacity,
228            );
229
230            self.set_message_sender(sender).await;
231            self.set_error_stream(error_stream).await;
232
233            Ok(stream)
234        } else {
235            let (stream, sender, error_stream) = MCPStream::create(
236                Box::pin(tokio::io::stdin()),
237                Mutex::new(Box::pin(tokio::io::stdout())),
238                IoStream::Writable(Box::pin(tokio::io::stderr())),
239                self.pending_requests.clone(),
240                self.options.timeout,
241                self.options.max_line_length,
242                cancellation_token,
243                self.options.channel_capacity,
244            );
245
246            self.set_message_sender(sender).await;
247            self.set_error_stream(error_stream).await;
248            Ok(stream)
249        }
250    }
251
252    async fn pending_request_tx(&self, request_id: &RequestId) -> Option<Sender<M>> {
253        let mut pending_requests = self.pending_requests.lock().await;
254        pending_requests.remove(request_id)
255    }
256
257    /// Checks if the transport has been shut down.
258    async fn is_shut_down(&self) -> bool {
259        let result = self.is_shut_down.lock().await;
260        *result
261    }
262
263    fn message_sender(&self) -> Arc<tokio::sync::RwLock<Option<MessageDispatcher<M>>>> {
264        self.message_sender.clone() as _
265    }
266
267    fn error_stream(&self) -> &tokio::sync::RwLock<Option<IoStream>> {
268        &self.error_stream as _
269    }
270
271    async fn consume_string_payload(&self, _payload: &str) -> TransportResult<()> {
272        Err(TransportError::Internal(
273            "Invalid invocation of consume_string_payload() function in StdioTransport".to_string(),
274        ))
275    }
276
277    async fn keep_alive(
278        &self,
279        _interval: Duration,
280        _disconnect_tx: oneshot::Sender<()>,
281    ) -> TransportResult<JoinHandle<()>> {
282        Err(TransportError::Internal(
283            "Invalid invocation of keep_alive() function for StdioTransport".to_string(),
284        ))
285    }
286
287    // Shuts down the transport, terminating any subprocess and signaling closure.
288    ///
289    /// Sends a shutdown signal via the watch channel and kills the subprocess if present.
290    ///
291    /// # Returns
292    /// A `TransportResult` indicating success or failure.
293    ///
294    /// # Errors
295    /// Returns a `TransportError` if the shutdown signal fails or the process cannot be killed.
296    async fn shut_down(&self) -> TransportResult<()> {
297        // Trigger cancellation
298        let mut cancellation_lock = self.shutdown_source.write().await;
299        if let Some(source) = cancellation_lock.as_ref() {
300            source.cancel()?;
301        }
302        *cancellation_lock = None; // Clear cancellation_source
303
304        // Mark as shut down
305        let mut is_shut_down_lock = self.is_shut_down.lock().await;
306        *is_shut_down_lock = true;
307        Ok(())
308    }
309}
310
311#[async_trait]
312impl McpDispatch<ClientMessages, ServerMessages, ClientMessage, ServerMessage>
313    for StdioTransport<ClientMessage>
314{
315    async fn send_message(
316        &self,
317        message: ServerMessages,
318        request_timeout: Option<Duration>,
319    ) -> TransportResult<Option<ClientMessages>> {
320        let sender = self.message_sender.read().await;
321        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
322        sender.send_message(message, request_timeout).await
323    }
324
325    async fn send(
326        &self,
327        message: ServerMessage,
328        request_timeout: Option<Duration>,
329    ) -> TransportResult<Option<ClientMessage>> {
330        let sender = self.message_sender.read().await;
331        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
332        sender.send(message, request_timeout).await
333    }
334
335    async fn send_batch(
336        &self,
337        message: Vec<ServerMessage>,
338        request_timeout: Option<Duration>,
339    ) -> TransportResult<Option<Vec<ClientMessage>>> {
340        let sender = self.message_sender.read().await;
341        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
342        sender.send_batch(message, request_timeout).await
343    }
344
345    async fn write_str(&self, payload: &str, skip_store: bool) -> TransportResult<()> {
346        let sender = self.message_sender.read().await;
347        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
348        sender.write_str(payload, skip_store).await
349    }
350}
351
352impl
353    TransportDispatcher<
354        ClientMessages,
355        MessageFromServer,
356        ClientMessage,
357        ServerMessages,
358        ServerMessage,
359    > for StdioTransport<ClientMessage>
360{
361}
362
363#[async_trait]
364impl McpDispatch<ServerMessages, ClientMessages, ServerMessage, ClientMessage>
365    for StdioTransport<ServerMessage>
366{
367    async fn send_message(
368        &self,
369        message: ClientMessages,
370        request_timeout: Option<Duration>,
371    ) -> TransportResult<Option<ServerMessages>> {
372        let sender = self.message_sender.read().await;
373        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
374        sender.send_message(message, request_timeout).await
375    }
376
377    async fn send(
378        &self,
379        message: ClientMessage,
380        request_timeout: Option<Duration>,
381    ) -> TransportResult<Option<ServerMessage>> {
382        let sender = self.message_sender.read().await;
383        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
384        sender.send(message, request_timeout).await
385    }
386
387    async fn send_batch(
388        &self,
389        message: Vec<ClientMessage>,
390        request_timeout: Option<Duration>,
391    ) -> TransportResult<Option<Vec<ServerMessage>>> {
392        let sender = self.message_sender.read().await;
393        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
394        sender.send_batch(message, request_timeout).await
395    }
396
397    async fn write_str(&self, payload: &str, skip_store: bool) -> TransportResult<()> {
398        let sender = self.message_sender.read().await;
399        let sender = sender.as_ref().ok_or(SdkError::connection_closed())?;
400        sender.write_str(payload, skip_store).await
401    }
402}
403
404impl
405    TransportDispatcher<
406        ServerMessages,
407        MessageFromClient,
408        ServerMessage,
409        ClientMessages,
410        ClientMessage,
411    > for StdioTransport<ServerMessage>
412{
413}
rust_mcp_transport/stdio.rs

rust_mcp_transport/
stdio.rs
