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

mcp_utils/testing/
fake_mcp.rs
