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

mcp_utils/testing/
fake_mcp.rs
