mcp_utils/
transport.rs

1use rmcp::RoleClient;
2use rmcp::RoleServer;
3use rmcp::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
4use rmcp::transport::Transport;
5use std::future::Future;
6use std::sync::Arc;
7use thiserror::Error;
8use tokio::sync::{Mutex, mpsc};
9
10#[derive(Debug, Error)]
11pub enum InMemoryTransportError {
12    #[error("Channel closed")]
13    ChannelClosed,
14}
15
16/// In-memory transport for connecting `McpServer` and `McpClient` in tests
17pub struct InMemoryTransport<R: ServiceRole> {
18    tx: Arc<Mutex<mpsc::Sender<TxJsonRpcMessage<R>>>>,
19    rx: Arc<Mutex<mpsc::Receiver<RxJsonRpcMessage<R>>>>,
20}
21
22impl<R: ServiceRole> InMemoryTransport<R> {
23    fn new(tx: mpsc::Sender<TxJsonRpcMessage<R>>, rx: mpsc::Receiver<RxJsonRpcMessage<R>>) -> Self {
24        Self { tx: Arc::new(Mutex::new(tx)), rx: Arc::new(Mutex::new(rx)) }
25    }
26}
27
28/// Create a pair of transports for client and server
29pub fn create_in_memory_transport() -> (InMemoryTransport<RoleClient>, InMemoryTransport<RoleServer>) {
30    // Client sends ClientRequest/ClientResult, receives ServerRequest/ServerResult
31    // Server sends ServerRequest/ServerResult, receives ClientRequest/ClientResult
32    let (client_tx, server_rx) = mpsc::channel(1000); // Client -> Server
33    let (server_tx, client_rx) = mpsc::channel(1000); // Server -> Client
34
35    let client_transport = InMemoryTransport::new(client_tx, client_rx);
36    let server_transport = InMemoryTransport::new(server_tx, server_rx);
37
38    (client_transport, server_transport)
39}
40
41impl<R: ServiceRole> Transport<R> for InMemoryTransport<R> {
42    type Error = InMemoryTransportError;
43
44    fn send(&mut self, item: TxJsonRpcMessage<R>) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
45        let tx = self.tx.clone();
46        async move {
47            let tx = tx.lock().await;
48            tx.send(item).await.map_err(|_| InMemoryTransportError::ChannelClosed)?;
49            Ok(())
50        }
51    }
52
53    fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<R>>> + Send {
54        let rx = self.rx.clone();
55        async move {
56            let mut rx = rx.lock().await;
57            rx.recv().await
58        }
59    }
60
61    async fn close(&mut self) -> Result<(), Self::Error> {
62        // Channels will be closed when dropped
63        Ok(())
64    }
65}
mcp_utils/transport.rs

mcp_utils/
transport.rs
