Skip to main content

mcp_utils/
testing.rs

1use std::future::Future;
2
3use crate::protocol::client_lifecycle_mode;
4use crate::transport::create_in_memory_transport;
5use rmcp::{
6    RoleClient, RoleServer, Service, serve_client_with_lifecycle, serve_server,
7    service::{ClientInitializeError, RunningService, ServerInitializeError},
8};
9
10#[cfg(feature = "client")]
11pub use elicitation_script::{CapturedElicitation, ElicitationScript};
12#[cfg(all(feature = "client", any(test, feature = "testing")))]
13pub use fake_mcp::{
14    CapturedTaskUpdate, CapturedToolCall, FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, fake_mcp,
15};
16
17#[cfg(all(feature = "client", any(test, feature = "testing")))]
18mod fake_mcp;
19
20pub type ConnectedServices<T, U> = (RunningService<RoleServer, T>, RunningService<RoleClient, U>);
21
22/// Helper function to connect an MCP server and client via in-memory transport
23/// This handles the dual-era discovery/initialization handshake by running both concurrently
24pub fn connect<T, U>(server: T, client: U) -> impl Future<Output = Result<ConnectedServices<T, U>, ConnectError>>
25where
26    T: Service<RoleServer>,
27    U: Service<RoleClient>,
28{
29    Box::pin(async move {
30        let (client_transport, server_transport) = create_in_memory_transport();
31
32        let (server_result, client_result) = tokio::join!(
33            serve_server(server, server_transport),
34            serve_client_with_lifecycle(client, client_transport, client_lifecycle_mode())
35        );
36
37        let server = server_result.map_err(ConnectError::ServerInit)?;
38        let client = client_result.map_err(ConnectError::ClientInit)?;
39
40        Ok((server, client))
41    })
42}
43
44#[derive(Debug, thiserror::Error)]
45pub enum ConnectError {
46    #[error("Server initialization failed: {0}")]
47    ServerInit(ServerInitializeError),
48    #[error("Client initialization failed: {0}")]
49    ClientInit(ClientInitializeError),
50}
51
52#[cfg(feature = "client")]
53mod elicitation_script {
54    use crate::client::McpClientEvent;
55    use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction};
56    use std::collections::VecDeque;
57    use std::sync::{Arc, Mutex, PoisonError};
58    use tokio::sync::mpsc;
59    use tokio::task::JoinHandle;
60
61    /// Scripts the user's side of elicitation round trips: answers each
62    /// incoming request with the next queued response (Cancel once the queue
63    /// is empty) and records what arrived for assertions.
64    pub struct ElicitationScript {
65        captured: Arc<Mutex<Vec<CapturedElicitation>>>,
66        task: JoinHandle<()>,
67    }
68
69    #[derive(Clone)]
70    pub struct CapturedElicitation {
71        pub server_name: String,
72        pub request: ElicitRequestParams,
73    }
74
75    impl ElicitationScript {
76        pub fn spawn(
77            mut event_rx: mpsc::Receiver<McpClientEvent>,
78            responses: impl IntoIterator<Item = ElicitResult>,
79        ) -> Self {
80            let mut responses = responses.into_iter().collect::<VecDeque<_>>();
81            let captured = Arc::new(Mutex::new(Vec::new()));
82            let recorder = Arc::clone(&captured);
83            let task = tokio::spawn(async move {
84                while let Some(event) = event_rx.recv().await {
85                    if let McpClientEvent::Elicitation(event) = event {
86                        recorder
87                            .lock()
88                            .unwrap_or_else(PoisonError::into_inner)
89                            .push(CapturedElicitation { server_name: event.server_name, request: event.request });
90                        let response =
91                            responses.pop_front().unwrap_or_else(|| ElicitResult::new(ElicitationAction::Cancel));
92                        let _ = event.response_sender.send(response);
93                    }
94                }
95            });
96            Self { captured, task }
97        }
98
99        pub fn captured(&self) -> Vec<CapturedElicitation> {
100            self.captured.lock().unwrap_or_else(PoisonError::into_inner).clone()
101        }
102    }
103
104    impl Drop for ElicitationScript {
105        fn drop(&mut self) {
106            self.task.abort();
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::connect;
114    use rmcp::{
115        ClientHandler, ServerHandler,
116        model::{ErrorData, Implementation, InitializeRequestParams, ProtocolVersion, ServerCapabilities, ServerInfo},
117        service::RequestContext,
118    };
119    use std::borrow::Cow;
120
121    #[tokio::test]
122    async fn connect_prefers_stateless_discovery_for_modern_servers() {
123        let (_server, client) = connect(McpServer728, TestClient).await.expect("connect");
124        assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2026_07_28);
125        client.list_tools(None).await.expect("list tools");
126        client.cancel().await.expect("cancel client");
127    }
128
129    #[tokio::test]
130    async fn connect_selects_an_older_mutually_supported_revision() {
131        let (_server, client) = connect(McpServer618, TestClient).await.expect("connect");
132
133        assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2025_06_18);
134        client.list_tools(None).await.expect("list tools");
135        client.cancel().await.expect("cancel client");
136    }
137
138    #[tokio::test]
139    async fn connect_falls_back_to_legacy_initialization() {
140        let (_server, client) = connect(McpServer1125, TestClient).await.expect("connect");
141
142        assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2025_11_25);
143        client.cancel().await.expect("cancel client");
144    }
145
146    #[derive(Clone, Default)]
147    struct TestClient;
148
149    impl ClientHandler for TestClient {}
150
151    #[derive(Clone, Default)]
152    struct McpServer728;
153
154    impl ServerHandler for McpServer728 {
155        fn get_info(&self) -> ServerInfo {
156            ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
157                .with_server_info(Implementation::new("modern-only", "1.0.0"))
158                .with_protocol_version(ProtocolVersion::V_2026_07_28)
159        }
160
161        fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
162            Cow::Owned(vec![ProtocolVersion::V_2026_07_28])
163        }
164
165        fn initialize(
166            &self,
167            _request: InitializeRequestParams,
168            _context: RequestContext<rmcp::RoleServer>,
169        ) -> impl std::future::Future<Output = Result<rmcp::model::InitializeResult, ErrorData>> + Send + '_ {
170            std::future::ready(Err(ErrorData::new(
171                rmcp::model::ErrorCode::METHOD_NOT_FOUND,
172                "initialize is not supported",
173                None,
174            )))
175        }
176    }
177
178    #[derive(Clone, Default)]
179    struct McpServer618;
180
181    impl ServerHandler for McpServer618 {
182        fn get_info(&self) -> ServerInfo {
183            ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
184                .with_server_info(Implementation::new("older-revision", "1.0.0"))
185                .with_protocol_version(ProtocolVersion::V_2025_06_18)
186        }
187
188        fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
189            Cow::Owned(vec![ProtocolVersion::V_2025_06_18])
190        }
191    }
192
193    #[derive(Clone, Default)]
194    struct McpServer1125;
195
196    impl ServerHandler for McpServer1125 {
197        fn get_info(&self) -> ServerInfo {
198            ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
199                .with_server_info(Implementation::new("legacy", "1.0.0"))
200                .with_protocol_version(ProtocolVersion::V_2025_11_25)
201        }
202
203        fn discover(
204            &self,
205            _context: RequestContext<rmcp::RoleServer>,
206        ) -> impl Future<Output = Result<rmcp::model::DiscoverResult, ErrorData>> + Send + '_ {
207            std::future::ready(Err(ErrorData::new(
208                rmcp::model::ErrorCode::METHOD_NOT_FOUND,
209                "server/discover is not supported",
210                None,
211            )))
212        }
213    }
214}