mcp_utils/testing/
fake_mcp.rs

1use crate::client::{RuntimeMcpServer, RuntimeMcpTransport, ToolExposure};
2use rmcp::{
3    ErrorData as McpError, Peer, RoleServer, ServerHandler,
4    model::{
5        CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ClientCapabilities,
6        ContentBlock, CreateTaskResult, DetailedTask, DiscoverResult, GetTaskParams, GetTaskResult, Implementation,
7        ListToolsResult, PaginatedRequestParams, ProgressNotificationParam, ProtocolVersion, ResultType,
8        ServerCapabilities, ServerInfo, Tool, UpdateTaskParams,
9    },
10    service::{DynService, RequestContext},
11};
12use serde_json::json;
13use std::collections::{BTreeMap, HashMap, VecDeque};
14use std::future::Future;
15use std::sync::{Arc, Mutex};
16use std::time::Duration;
17
18pub fn fake_mcp(name: &str, server: FakeMcpServer) -> RuntimeMcpServer {
19    RuntimeMcpServer::new(name, RuntimeMcpTransport::InMemory { server: server.into_dyn() }, ToolExposure::ModelVisible)
20}
21
22/// A fake MCP server preloaded with the classic math tools (`add_numbers`,
23/// `divide_numbers`, `slow_tool`); add scripted tools with [`Self::with_tool`].
24#[derive(Clone)]
25pub struct FakeMcpServer {
26    state: FakeMcpState,
27}
28
29#[derive(Clone, Default)]
30pub struct FakeMcpState {
31    inner: Arc<Mutex<FakeMcpStateInner>>,
32}
33
34#[derive(Clone)]
35pub struct CapturedToolCall {
36    pub request: CallToolRequestParams,
37    pub context_meta: serde_json::Map<String, serde_json::Value>,
38}
39
40#[derive(Clone)]
41pub struct CapturedTaskUpdate {
42    pub task_id: String,
43    pub input_responses: rmcp::model::InputResponses,
44}
45
46#[derive(Clone)]
47pub struct FakeTool {
48    definition: Tool,
49    responses: HashMap<Option<String>, FakeToolResponse>,
50    handler: Option<ToolHandler>,
51}
52
53#[derive(Clone)]
54pub struct FakeToolResponse {
55    response: CallToolResponse,
56    delay: Duration,
57    progress: Vec<(f64, Option<f64>)>,
58    task_progress: Vec<(f64, Option<f64>)>,
59}
60
61impl FakeMcpServer {
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    pub fn with_tool(self, tool: FakeTool) -> Self {
67        self.state.add_tool(tool);
68        self
69    }
70
71    pub fn with_task(self, task_id: impl Into<String>, states: impl IntoIterator<Item = DetailedTask>) -> Self {
72        self.state.script_task(task_id, states);
73        self
74    }
75
76    pub fn with_task_get_failures(self, failures: usize) -> Self {
77        self.state.lock().task_get_failures = failures;
78        self
79    }
80
81    pub fn with_task_update_failures(self, failures: usize) -> Self {
82        self.state.lock().task_update_failures = failures;
83        self
84    }
85
86    pub fn state(&self) -> FakeMcpState {
87        self.state.clone()
88    }
89
90    pub fn into_dyn(self) -> Box<dyn DynService<RoleServer>> {
91        Box::new(self)
92    }
93}
94
95impl FakeMcpState {
96    pub fn calls_for(&self, tool: &str) -> Vec<CapturedToolCall> {
97        self.lock().calls.iter().filter(|call| call.request.name.as_ref() == tool).cloned().collect()
98    }
99
100    pub fn task_get_ids(&self) -> Vec<String> {
101        self.lock().task_get_ids.clone()
102    }
103
104    pub fn task_updates(&self) -> Vec<CapturedTaskUpdate> {
105        self.lock().task_updates.clone()
106    }
107
108    pub fn task_cancel_ids(&self) -> Vec<String> {
109        self.lock().task_cancel_ids.clone()
110    }
111
112    pub fn client_capabilities(&self) -> Option<ClientCapabilities> {
113        self.lock().client_capabilities.clone()
114    }
115
116    pub fn script_task(&self, task_id: impl Into<String>, states: impl IntoIterator<Item = DetailedTask>) {
117        self.lock().tasks.insert(task_id.into(), states.into_iter().collect());
118    }
119
120    fn task_for(&self, task_id: &str) -> Result<Option<DetailedTask>, ()> {
121        let mut inner = self.lock();
122        inner.task_get_ids.push(task_id.to_string());
123        if inner.task_get_failures > 0 {
124            inner.task_get_failures -= 1;
125            return Err(());
126        }
127        let Some(states) = inner.tasks.get_mut(task_id) else {
128            return Ok(None);
129        };
130        Ok(if states.len() > 1 { states.pop_front() } else { states.front().cloned() })
131    }
132
133    fn record_task_update(&self, request: UpdateTaskParams) -> bool {
134        let mut inner = self.lock();
135        inner
136            .task_updates
137            .push(CapturedTaskUpdate { task_id: request.task_id, input_responses: request.input_responses });
138        if inner.task_update_failures > 0 {
139            inner.task_update_failures -= 1;
140            false
141        } else {
142            true
143        }
144    }
145
146    fn record_task_cancel(&self, request: CancelTaskParams) {
147        self.lock().task_cancel_ids.push(request.task_id);
148    }
149
150    pub fn add_tool(&self, tool: FakeTool) {
151        self.lock().tools.insert(tool.definition.name.to_string(), tool);
152    }
153
154    pub async fn add_tool_and_notify(&self, tool: FakeTool) {
155        let peers = {
156            let mut inner = self.lock();
157            inner.tools.insert(tool.definition.name.to_string(), tool);
158            inner.peers.clone()
159        };
160        for peer in peers {
161            let _ = peer.notify_tool_list_changed().await;
162        }
163    }
164
165    pub async fn clear_tools_and_notify(&self) {
166        let peers = {
167            let mut inner = self.lock();
168            inner.tools.clear();
169            inner.peers.clone()
170        };
171        for peer in peers {
172            let _ = peer.notify_tool_list_changed().await;
173        }
174    }
175
176    pub fn fail_next_tool_list(&self) {
177        self.lock().tool_list_failures += 1;
178    }
179
180    fn definitions(&self) -> Vec<Tool> {
181        self.lock().tools.values().map(|tool| tool.definition.clone()).collect()
182    }
183
184    fn response_for(
185        &self,
186        request: &CallToolRequestParams,
187        context_meta: serde_json::Map<String, serde_json::Value>,
188    ) -> Option<FakeToolResponse> {
189        let mut inner = self.lock();
190        inner.calls.push(CapturedToolCall { request: request.clone(), context_meta });
191        inner.tools.get(request.name.as_ref()).and_then(|tool| tool.response_for(request))
192    }
193
194    fn lock(&self) -> std::sync::MutexGuard<'_, FakeMcpStateInner> {
195        self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
196    }
197}
198
199impl FakeTool {
200    pub fn new(name: impl Into<String>) -> Self {
201        let name = name.into();
202        let schema = serde_json::from_value(json!({ "type": "object", "properties": {} }))
203            .expect("empty object schema is valid");
204        Self {
205            definition: Tool::new(name, "Fake MCP tool", Arc::new(schema)),
206            responses: HashMap::new(),
207            handler: None,
208        }
209    }
210
211    pub fn description(mut self, description: impl Into<String>) -> Self {
212        self.definition.description = Some(description.into().into());
213        self
214    }
215
216    pub fn responds(mut self, response: impl Into<FakeToolResponse>) -> Self {
217        self.responses.insert(None, response.into());
218        self
219    }
220
221    pub fn when_state(mut self, state: impl Into<String>, response: impl Into<FakeToolResponse>) -> Self {
222        self.responses.insert(Some(state.into()), response.into());
223        self
224    }
225
226    /// Compute the response from the request, for tools whose output depends
227    /// on their arguments. Scripted responses take precedence.
228    pub fn responds_with(
229        mut self,
230        handler: impl Fn(&CallToolRequestParams) -> FakeToolResponse + Send + Sync + 'static,
231    ) -> Self {
232        self.handler = Some(Arc::new(handler));
233        self
234    }
235
236    fn response_for(&self, request: &CallToolRequestParams) -> Option<FakeToolResponse> {
237        self.responses
238            .get(&request.request_state.as_deref().map(str::to_string))
239            .cloned()
240            .or_else(|| self.handler.as_ref().map(|handler| handler(request)))
241    }
242}
243
244impl FakeToolResponse {
245    pub fn new(response: impl Into<CallToolResponse>) -> Self {
246        Self { response: response.into(), delay: Duration::ZERO, progress: Vec::new(), task_progress: Vec::new() }
247    }
248
249    pub fn text(text: impl Into<String>) -> Self {
250        Self::new(CallToolResult::success(vec![ContentBlock::text(text.into())]))
251    }
252
253    pub fn task(task: CreateTaskResult) -> Self {
254        Self::new(CallToolResponse::Task(task))
255    }
256
257    pub fn delay(mut self, delay: Duration) -> Self {
258        self.delay = delay;
259        self
260    }
261
262    pub fn progress(mut self, progress: f64, total: Option<f64>) -> Self {
263        self.progress.push((progress, total));
264        self
265    }
266
267    pub fn task_progress(mut self, progress: f64, total: Option<f64>) -> Self {
268        self.task_progress.push((progress, total));
269        self
270    }
271}
272
273impl<T> From<T> for FakeToolResponse
274where
275    T: Into<CallToolResponse>,
276{
277    fn from(response: T) -> Self {
278        Self::new(response)
279    }
280}
281
282impl Default for FakeMcpServer {
283    fn default() -> Self {
284        Self { state: FakeMcpState::default() }
285            .with_tool(add_numbers())
286            .with_tool(divide_numbers())
287            .with_tool(slow_tool())
288    }
289}
290
291impl ServerHandler for FakeMcpServer {
292    fn discover(
293        &self,
294        context: RequestContext<RoleServer>,
295    ) -> impl Future<Output = Result<DiscoverResult, McpError>> + Send + '_ {
296        self.state.lock().client_capabilities = context.meta.client_capabilities();
297        std::future::ready(Ok(DiscoverResult::from_server_info(
298            ServerHandler::supported_protocol_versions(self).into_owned(),
299            ServerHandler::get_info(self),
300        )))
301    }
302
303    fn get_info(&self) -> ServerInfo {
304        ServerInfo::new(ServerCapabilities::builder().enable_tools().enable_tasks().build())
305            .with_server_info(
306                Implementation::new("fake-mcp-server", "0.1.0").with_description("A fake MCP server for testing"),
307            )
308            .with_instructions("A fake MCP server for testing")
309    }
310
311    fn get_task(
312        &self,
313        request: GetTaskParams,
314        _context: RequestContext<RoleServer>,
315    ) -> impl Future<Output = Result<GetTaskResult, McpError>> + Send + '_ {
316        let result = match self.state.task_for(&request.task_id) {
317            Ok(Some(task)) => Ok(GetTaskResult::new(task)),
318            Ok(None) => Err(McpError::invalid_params(format!("unknown task: {}", request.task_id), None)),
319            Err(()) => Err(McpError::internal_error("scripted tasks/get failure", None)),
320        };
321        std::future::ready(result)
322    }
323
324    fn update_task(
325        &self,
326        request: UpdateTaskParams,
327        _context: RequestContext<RoleServer>,
328    ) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
329        std::future::ready(
330            self.state
331                .record_task_update(request)
332                .then_some(())
333                .ok_or_else(|| McpError::internal_error("scripted tasks/update failure", None)),
334        )
335    }
336
337    fn cancel_task(
338        &self,
339        request: CancelTaskParams,
340        _context: RequestContext<RoleServer>,
341    ) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
342        self.state.record_task_cancel(request);
343        std::future::ready(Ok(()))
344    }
345
346    fn list_tools(
347        &self,
348        _request: Option<PaginatedRequestParams>,
349        context: RequestContext<RoleServer>,
350    ) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
351        let supports_cache_hints =
352            context.protocol_version().is_some_and(|version| version >= ProtocolVersion::V_2026_07_28);
353        let tools = {
354            let mut inner = self.state.lock();
355            if inner.tool_list_failures > 0 {
356                inner.tool_list_failures -= 1;
357                return std::future::ready(Err(McpError::internal_error("scripted tools/list failure", None)));
358            }
359            if inner.peers.is_empty() {
360                inner.peers.push(context.peer);
361            }
362            inner.tools.values().map(|tool| tool.definition.clone()).collect()
363        };
364        std::future::ready(Ok(ListToolsResult {
365            result_type: Some(ResultType::COMPLETE),
366            tools,
367            meta: None,
368            next_cursor: None,
369            ttl_ms: supports_cache_hints.then_some(0),
370            cache_scope: supports_cache_hints.then_some(CacheScope::Public),
371        }))
372    }
373
374    fn get_tool(&self, name: &str) -> Option<Tool> {
375        self.state.definitions().into_iter().find(|tool| tool.name == name)
376    }
377
378    async fn call_tool(
379        &self,
380        request: CallToolRequestParams,
381        context: RequestContext<RoleServer>,
382    ) -> Result<CallToolResponse, McpError> {
383        let response = self.state.response_for(&request, context.meta.0.0.clone());
384        let Some(response) = response else {
385            return Err(McpError::invalid_params(format!("unknown tool: {}", request.name), None));
386        };
387
388        if !response.delay.is_zero() {
389            tokio::time::sleep(response.delay).await;
390        }
391        if let Some(token) = context.meta.get_progress_token() {
392            for (progress, total) in response.progress {
393                let mut notification = ProgressNotificationParam::new(token.clone(), progress);
394                if let Some(total) = total {
395                    notification = notification.with_total(total);
396                }
397                let _ = context.peer.notify_progress(notification).await;
398            }
399            if !response.task_progress.is_empty() {
400                let peer = context.peer.clone();
401                let token = token.clone();
402                tokio::spawn(async move {
403                    tokio::task::yield_now().await;
404                    for (progress, total) in response.task_progress {
405                        let mut notification = ProgressNotificationParam::new(token.clone(), progress);
406                        if let Some(total) = total {
407                            notification = notification.with_total(total);
408                        }
409                        let _ = peer.notify_progress(notification).await;
410                    }
411                });
412            }
413        }
414        Ok(response.response)
415    }
416}
417
418type ToolHandler = Arc<dyn Fn(&CallToolRequestParams) -> FakeToolResponse + Send + Sync>;
419
420#[derive(Default)]
421struct FakeMcpStateInner {
422    tools: BTreeMap<String, FakeTool>,
423    calls: Vec<CapturedToolCall>,
424    client_capabilities: Option<ClientCapabilities>,
425    tasks: HashMap<String, VecDeque<DetailedTask>>,
426    task_get_ids: Vec<String>,
427    task_updates: Vec<CapturedTaskUpdate>,
428    task_cancel_ids: Vec<String>,
429    task_get_failures: usize,
430    task_update_failures: usize,
431    tool_list_failures: usize,
432    peers: Vec<Peer<RoleServer>>,
433}
434
435fn add_numbers() -> FakeTool {
436    FakeTool::new("add_numbers").description("Adds two numbers together").responds_with(|request| {
437        let sum = int_arg(request, "a") + int_arg(request, "b");
438        FakeToolResponse::new(CallToolResult::structured(json!({ "sum": sum })))
439    })
440}
441
442fn divide_numbers() -> FakeTool {
443    FakeTool::new("divide_numbers").description("Divides two numbers").responds_with(|request| {
444        let (a, b) = (int_arg(request, "a"), int_arg(request, "b"));
445        if b == 0 {
446            return FakeToolResponse::new(CallToolResult::error(vec![ContentBlock::text("Division by zero")]));
447        }
448        FakeToolResponse::new(CallToolResult::structured(json!({ "quotient": a / b })))
449    })
450}
451
452fn slow_tool() -> FakeTool {
453    FakeTool::new("slow_tool")
454        .description("A tool that sleeps for a specified duration (for testing timeouts)")
455        .responds_with(|request| {
456            let sleep_ms = int_arg(request, "sleep_ms").unsigned_abs();
457            FakeToolResponse::new(CallToolResult::structured(json!({ "message": format!("Slept for {sleep_ms}ms") })))
458                .delay(Duration::from_millis(sleep_ms))
459        })
460}
461
462fn int_arg(request: &CallToolRequestParams, name: &str) -> i64 {
463    request.arguments.as_ref().and_then(|args| args.get(name)).and_then(serde_json::Value::as_i64).unwrap_or_default()
464}
mcp_utils/testing/fake_mcp.rs

mcp_utils/testing/
fake_mcp.rs
