1use std::collections::{HashMap, HashSet};
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::{Arc, RwLock};
10use std::task::{Context, Poll};
11
12use tower_service::Service;
13
14use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
15
16use crate::async_task::TaskStore;
17use crate::context::{
18 CancellationToken, ClientRequesterHandle, NotificationSender, RequestContext,
19 ServerNotification,
20};
21use crate::error::{Error, JsonRpcError, Result};
22use crate::filter::{PromptFilter, ResourceFilter, ToolFilter};
23use crate::prompt::Prompt;
24use crate::protocol::*;
25#[cfg(feature = "dynamic-tools")]
26use crate::registry::{
27 DynamicPromptRegistry, DynamicPromptsInner, DynamicResourceRegistry,
28 DynamicResourceTemplateRegistry, DynamicResourceTemplatesInner, DynamicResourcesInner,
29 DynamicToolRegistry, DynamicToolsInner,
30};
31use crate::resource::{Resource, ResourceTemplate};
32use crate::session::SessionState;
33use crate::tool::Tool;
34
35pub(crate) type CompletionHandler = Arc<
37 dyn Fn(CompleteParams) -> Pin<Box<dyn Future<Output = Result<CompleteResult>> + Send>>
38 + Send
39 + Sync,
40>;
41
42fn decode_cursor(cursor: &str) -> Result<usize> {
46 let bytes = BASE64
47 .decode(cursor)
48 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
49 let s = String::from_utf8(bytes)
50 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
51 s.parse::<usize>()
52 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))
53}
54
55fn encode_cursor(offset: usize) -> String {
57 BASE64.encode(offset.to_string())
58}
59
60fn paginate<T>(
64 items: Vec<T>,
65 cursor: Option<&str>,
66 page_size: Option<usize>,
67) -> Result<(Vec<T>, Option<String>)> {
68 let Some(page_size) = page_size else {
69 return Ok((items, None));
70 };
71
72 let offset = match cursor {
73 Some(c) => decode_cursor(c)?,
74 None => 0,
75 };
76
77 if offset >= items.len() {
78 return Ok((Vec::new(), None));
79 }
80
81 let end = (offset + page_size).min(items.len());
82 let next_cursor = if end < items.len() {
83 Some(encode_cursor(end))
84 } else {
85 None
86 };
87
88 let mut items = items;
89 let page = items.drain(offset..end).collect();
90 Ok((page, next_cursor))
91}
92
93#[derive(Clone)]
117pub struct McpRouter {
118 inner: Arc<McpRouterInner>,
119 session: SessionState,
120}
121
122impl std::fmt::Debug for McpRouter {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 f.debug_struct("McpRouter")
125 .field("server_name", &self.inner.server_name)
126 .field("server_version", &self.inner.server_version)
127 .field("tools_count", &self.inner.tools.len())
128 .field("resources_count", &self.inner.resources.len())
129 .field("prompts_count", &self.inner.prompts.len())
130 .field("session_phase", &self.session.phase())
131 .finish()
132 }
133}
134
135#[derive(Clone, Debug)]
137struct AutoInstructionsConfig {
138 prefix: Option<String>,
139 suffix: Option<String>,
140}
141
142#[derive(Clone)]
144struct McpRouterInner {
145 server_name: String,
146 server_version: String,
147 server_title: Option<String>,
149 server_description: Option<String>,
151 server_icons: Option<Vec<ToolIcon>>,
153 server_website_url: Option<String>,
155 instructions: Option<String>,
156 auto_instructions: Option<AutoInstructionsConfig>,
157 tools: HashMap<String, Arc<Tool>>,
158 resources: HashMap<String, Arc<Resource>>,
159 resource_templates: Vec<Arc<ResourceTemplate>>,
161 prompts: HashMap<String, Arc<Prompt>>,
162 in_flight: Arc<RwLock<HashMap<RequestId, CancellationToken>>>,
164 notification_tx: Option<NotificationSender>,
166 client_requester: Option<ClientRequesterHandle>,
168 task_store: TaskStore,
170 subscriptions: Arc<RwLock<HashSet<String>>>,
172 completion_handler: Option<CompletionHandler>,
174 tool_filter: Option<ToolFilter>,
176 resource_filter: Option<ResourceFilter>,
178 prompt_filter: Option<PromptFilter>,
180 extensions: Arc<crate::context::Extensions>,
182 min_log_level: Arc<RwLock<LogLevel>>,
184 page_size: Option<usize>,
186 #[cfg(feature = "dynamic-tools")]
188 dynamic_tools: Option<Arc<DynamicToolsInner>>,
189 #[cfg(feature = "dynamic-tools")]
191 dynamic_prompts: Option<Arc<DynamicPromptsInner>>,
192 #[cfg(feature = "dynamic-tools")]
194 dynamic_resources: Option<Arc<DynamicResourcesInner>>,
195 #[cfg(feature = "dynamic-tools")]
197 dynamic_resource_templates: Option<Arc<DynamicResourceTemplatesInner>>,
198}
199
200impl McpRouterInner {
201 fn generate_instructions(&self, config: &AutoInstructionsConfig) -> String {
203 let mut parts = Vec::new();
204
205 if let Some(prefix) = &config.prefix {
206 parts.push(prefix.clone());
207 }
208
209 if !self.tools.is_empty() {
211 let mut lines = vec!["## Tools".to_string(), String::new()];
212 let mut tools: Vec<_> = self.tools.values().collect();
213 tools.sort_by(|a, b| a.name.cmp(&b.name));
214 for tool in tools {
215 let desc = tool.description.as_deref().unwrap_or("No description");
216 let tags = annotation_tags(tool.annotations.as_ref());
217 if tags.is_empty() {
218 lines.push(format!("- **{}**: {}", tool.name, desc));
219 } else {
220 lines.push(format!("- **{}**: {} [{}]", tool.name, desc, tags));
221 }
222 }
223 parts.push(lines.join("\n"));
224 }
225
226 if !self.resources.is_empty() || !self.resource_templates.is_empty() {
228 let mut lines = vec!["## Resources".to_string(), String::new()];
229 let mut resources: Vec<_> = self.resources.values().collect();
230 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
231 for resource in resources {
232 let desc = resource.description.as_deref().unwrap_or("No description");
233 lines.push(format!("- **{}**: {}", resource.uri, desc));
234 }
235 let mut templates: Vec<_> = self.resource_templates.iter().collect();
236 templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
237 for template in templates {
238 let desc = template.description.as_deref().unwrap_or("No description");
239 lines.push(format!("- **{}**: {}", template.uri_template, desc));
240 }
241 parts.push(lines.join("\n"));
242 }
243
244 if !self.prompts.is_empty() {
246 let mut lines = vec!["## Prompts".to_string(), String::new()];
247 let mut prompts: Vec<_> = self.prompts.values().collect();
248 prompts.sort_by(|a, b| a.name.cmp(&b.name));
249 for prompt in prompts {
250 let desc = prompt.description.as_deref().unwrap_or("No description");
251 lines.push(format!("- **{}**: {}", prompt.name, desc));
252 }
253 parts.push(lines.join("\n"));
254 }
255
256 if let Some(suffix) = &config.suffix {
257 parts.push(suffix.clone());
258 }
259
260 parts.join("\n\n")
261 }
262}
263
264fn annotation_tags(annotations: Option<&crate::protocol::ToolAnnotations>) -> String {
270 let Some(ann) = annotations else {
271 return String::new();
272 };
273 let mut tags = Vec::new();
274 if ann.is_read_only() {
275 tags.push("read-only");
276 }
277 if ann.is_idempotent() {
278 tags.push("idempotent");
279 }
280 tags.join(", ")
281}
282
283impl McpRouter {
284 pub fn new() -> Self {
286 Self {
287 inner: Arc::new(McpRouterInner {
288 server_name: "tower-mcp".to_string(),
289 server_version: env!("CARGO_PKG_VERSION").to_string(),
290 server_title: None,
291 server_description: None,
292 server_icons: None,
293 server_website_url: None,
294 instructions: None,
295 auto_instructions: None,
296 tools: HashMap::new(),
297 resources: HashMap::new(),
298 resource_templates: Vec::new(),
299 prompts: HashMap::new(),
300 in_flight: Arc::new(RwLock::new(HashMap::new())),
301 notification_tx: None,
302 client_requester: None,
303 task_store: TaskStore::new(),
304 subscriptions: Arc::new(RwLock::new(HashSet::new())),
305 extensions: Arc::new(crate::context::Extensions::new()),
306 completion_handler: None,
307 tool_filter: None,
308 resource_filter: None,
309 prompt_filter: None,
310 min_log_level: Arc::new(RwLock::new(LogLevel::Debug)),
311 page_size: None,
312 #[cfg(feature = "dynamic-tools")]
313 dynamic_tools: None,
314 #[cfg(feature = "dynamic-tools")]
315 dynamic_prompts: None,
316 #[cfg(feature = "dynamic-tools")]
317 dynamic_resources: None,
318 #[cfg(feature = "dynamic-tools")]
319 dynamic_resource_templates: None,
320 }),
321 session: SessionState::new(),
322 }
323 }
324
325 pub fn with_fresh_session(&self) -> Self {
333 Self {
334 inner: self.inner.clone(),
335 session: SessionState::new(),
336 }
337 }
338
339 pub fn tool_annotations_map(&self) -> ToolAnnotationsMap {
349 let mut map = HashMap::new();
350 for (name, tool) in &self.inner.tools {
351 if let Some(annotations) = &tool.annotations {
352 map.insert(name.clone(), annotations.clone());
353 }
354 }
355 #[cfg(feature = "dynamic-tools")]
356 if let Some(dynamic) = &self.inner.dynamic_tools {
357 for tool in dynamic.list() {
358 if !map.contains_key(&tool.name)
360 && let Some(ref annotations) = tool.annotations
361 {
362 map.insert(tool.name.clone(), annotations.clone());
363 }
364 }
365 }
366 ToolAnnotationsMap { map: Arc::new(map) }
367 }
368
369 pub fn task_store(&self) -> &TaskStore {
371 &self.inner.task_store
372 }
373
374 #[cfg(feature = "dynamic-tools")]
404 pub fn with_dynamic_tools(mut self) -> (Self, DynamicToolRegistry) {
405 let inner_dyn = Arc::new(DynamicToolsInner::new());
406 Arc::make_mut(&mut self.inner).dynamic_tools = Some(inner_dyn.clone());
407 (self, DynamicToolRegistry::new(inner_dyn))
408 }
409
410 #[cfg(feature = "dynamic-tools")]
433 pub fn with_dynamic_prompts(mut self) -> (Self, DynamicPromptRegistry) {
434 let inner_dyn = Arc::new(DynamicPromptsInner::new());
435 Arc::make_mut(&mut self.inner).dynamic_prompts = Some(inner_dyn.clone());
436 (self, DynamicPromptRegistry::new(inner_dyn))
437 }
438
439 #[cfg(feature = "dynamic-tools")]
462 pub fn with_dynamic_resources(mut self) -> (Self, DynamicResourceRegistry) {
463 let inner_dyn = Arc::new(DynamicResourcesInner::new());
464 Arc::make_mut(&mut self.inner).dynamic_resources = Some(inner_dyn.clone());
465 (self, DynamicResourceRegistry::new(inner_dyn))
466 }
467
468 #[cfg(feature = "dynamic-tools")]
490 pub fn with_dynamic_resource_templates(mut self) -> (Self, DynamicResourceTemplateRegistry) {
491 let inner_dyn = Arc::new(DynamicResourceTemplatesInner::new());
492 Arc::make_mut(&mut self.inner).dynamic_resource_templates = Some(inner_dyn.clone());
493 (self, DynamicResourceTemplateRegistry::new(inner_dyn))
494 }
495
496 pub fn with_notification_sender(mut self, tx: NotificationSender) -> Self {
500 let inner = Arc::make_mut(&mut self.inner);
501 #[cfg(feature = "dynamic-tools")]
504 if let Some(ref dynamic_tools) = inner.dynamic_tools {
505 dynamic_tools.add_notification_sender(tx.clone());
506 }
507 #[cfg(feature = "dynamic-tools")]
508 if let Some(ref dynamic_prompts) = inner.dynamic_prompts {
509 dynamic_prompts.add_notification_sender(tx.clone());
510 }
511 #[cfg(feature = "dynamic-tools")]
512 if let Some(ref dynamic_resources) = inner.dynamic_resources {
513 dynamic_resources.add_notification_sender(tx.clone());
514 }
515 #[cfg(feature = "dynamic-tools")]
516 if let Some(ref dynamic_resource_templates) = inner.dynamic_resource_templates {
517 dynamic_resource_templates.add_notification_sender(tx.clone());
518 }
519 inner.notification_tx = Some(tx);
520 self
521 }
522
523 pub fn notification_sender(&self) -> Option<&NotificationSender> {
525 self.inner.notification_tx.as_ref()
526 }
527
528 pub fn with_client_requester(mut self, requester: ClientRequesterHandle) -> Self {
533 Arc::make_mut(&mut self.inner).client_requester = Some(requester);
534 self
535 }
536
537 pub fn client_requester(&self) -> Option<&ClientRequesterHandle> {
539 self.inner.client_requester.as_ref()
540 }
541
542 pub fn with_state<T: Clone + Send + Sync + 'static>(mut self, state: T) -> Self {
585 let inner = Arc::make_mut(&mut self.inner);
586 Arc::make_mut(&mut inner.extensions).insert(state);
587 self
588 }
589
590 pub fn with_extension<T: Clone + Send + Sync + 'static>(self, value: T) -> Self {
595 self.with_state(value)
596 }
597
598 pub fn extensions(&self) -> &crate::context::Extensions {
600 &self.inner.extensions
601 }
602
603 pub fn create_context(
608 &self,
609 request_id: RequestId,
610 progress_token: Option<ProgressToken>,
611 ) -> RequestContext {
612 let ctx = RequestContext::new(request_id.clone());
613
614 let ctx = if let Some(token) = progress_token {
616 ctx.with_progress_token(token)
617 } else {
618 ctx
619 };
620
621 let ctx = if let Some(tx) = &self.inner.notification_tx {
623 ctx.with_notification_sender(tx.clone())
624 } else {
625 ctx
626 };
627
628 let ctx = if let Some(requester) = &self.inner.client_requester {
630 ctx.with_client_requester(requester.clone())
631 } else {
632 ctx
633 };
634
635 let ctx = ctx.with_extensions(self.inner.extensions.clone());
637
638 let ctx = ctx.with_min_log_level(self.inner.min_log_level.clone());
640
641 let token = ctx.cancellation_token();
643 if let Ok(mut in_flight) = self.inner.in_flight.write() {
644 in_flight.insert(request_id, token);
645 }
646
647 ctx
648 }
649
650 pub fn complete_request(&self, request_id: &RequestId) {
652 if let Ok(mut in_flight) = self.inner.in_flight.write() {
653 in_flight.remove(request_id);
654 }
655 }
656
657 fn cancel_request(&self, request_id: &RequestId) -> bool {
659 let Ok(in_flight) = self.inner.in_flight.read() else {
660 return false;
661 };
662 let Some(token) = in_flight.get(request_id) else {
663 return false;
664 };
665 token.cancel();
666 true
667 }
668
669 pub fn server_info(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
671 let inner = Arc::make_mut(&mut self.inner);
672 inner.server_name = name.into();
673 inner.server_version = version.into();
674 self
675 }
676
677 pub fn page_size(mut self, size: usize) -> Self {
684 Arc::make_mut(&mut self.inner).page_size = Some(size);
685 self
686 }
687
688 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
690 Arc::make_mut(&mut self.inner).instructions = Some(instructions.into());
691 self
692 }
693
694 pub fn auto_instructions(mut self) -> Self {
726 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
727 prefix: None,
728 suffix: None,
729 });
730 self
731 }
732
733 pub fn auto_instructions_with(
750 mut self,
751 prefix: Option<impl Into<String>>,
752 suffix: Option<impl Into<String>>,
753 ) -> Self {
754 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
755 prefix: prefix.map(Into::into),
756 suffix: suffix.map(Into::into),
757 });
758 self
759 }
760
761 pub fn server_title(mut self, title: impl Into<String>) -> Self {
763 Arc::make_mut(&mut self.inner).server_title = Some(title.into());
764 self
765 }
766
767 pub fn server_description(mut self, description: impl Into<String>) -> Self {
769 Arc::make_mut(&mut self.inner).server_description = Some(description.into());
770 self
771 }
772
773 pub fn server_icons(mut self, icons: Vec<ToolIcon>) -> Self {
775 Arc::make_mut(&mut self.inner).server_icons = Some(icons);
776 self
777 }
778
779 pub fn server_website_url(mut self, url: impl Into<String>) -> Self {
781 Arc::make_mut(&mut self.inner).server_website_url = Some(url.into());
782 self
783 }
784
785 pub fn tool(mut self, tool: Tool) -> Self {
787 Arc::make_mut(&mut self.inner)
788 .tools
789 .insert(tool.name.clone(), Arc::new(tool));
790 self
791 }
792
793 pub fn tool_if(self, condition: bool, tool: Tool) -> Self {
819 if condition { self.tool(tool) } else { self }
820 }
821
822 pub fn resource(mut self, resource: Resource) -> Self {
824 Arc::make_mut(&mut self.inner)
825 .resources
826 .insert(resource.uri.clone(), Arc::new(resource));
827 self
828 }
829
830 pub fn resource_if(self, condition: bool, resource: Resource) -> Self {
849 if condition {
850 self.resource(resource)
851 } else {
852 self
853 }
854 }
855
856 pub fn resource_template(mut self, template: ResourceTemplate) -> Self {
889 Arc::make_mut(&mut self.inner)
890 .resource_templates
891 .push(Arc::new(template));
892 self
893 }
894
895 pub fn prompt(mut self, prompt: Prompt) -> Self {
897 Arc::make_mut(&mut self.inner)
898 .prompts
899 .insert(prompt.name.clone(), Arc::new(prompt));
900 self
901 }
902
903 pub fn prompt_if(self, condition: bool, prompt: Prompt) -> Self {
922 if condition { self.prompt(prompt) } else { self }
923 }
924
925 pub fn tools(self, tools: impl IntoIterator<Item = Tool>) -> Self {
951 tools
952 .into_iter()
953 .fold(self, |router, tool| router.tool(tool))
954 }
955
956 pub fn tools_if(self, condition: bool, tools: impl IntoIterator<Item = Tool>) -> Self {
960 if condition { self.tools(tools) } else { self }
961 }
962
963 pub fn resources(self, resources: impl IntoIterator<Item = Resource>) -> Self {
982 resources
983 .into_iter()
984 .fold(self, |router, resource| router.resource(resource))
985 }
986
987 pub fn resources_if(
991 self,
992 condition: bool,
993 resources: impl IntoIterator<Item = Resource>,
994 ) -> Self {
995 if condition {
996 self.resources(resources)
997 } else {
998 self
999 }
1000 }
1001
1002 pub fn prompts(self, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1021 prompts
1022 .into_iter()
1023 .fold(self, |router, prompt| router.prompt(prompt))
1024 }
1025
1026 pub fn prompts_if(self, condition: bool, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1030 if condition {
1031 self.prompts(prompts)
1032 } else {
1033 self
1034 }
1035 }
1036
1037 pub fn merge(mut self, other: McpRouter) -> Self {
1082 let inner = Arc::make_mut(&mut self.inner);
1083 let other_inner = other.inner;
1084
1085 for (name, tool) in &other_inner.tools {
1087 inner.tools.insert(name.clone(), tool.clone());
1088 }
1089
1090 for (uri, resource) in &other_inner.resources {
1092 inner.resources.insert(uri.clone(), resource.clone());
1093 }
1094
1095 for template in &other_inner.resource_templates {
1098 inner.resource_templates.push(template.clone());
1099 }
1100
1101 for (name, prompt) in &other_inner.prompts {
1103 inner.prompts.insert(name.clone(), prompt.clone());
1104 }
1105
1106 self
1107 }
1108
1109 pub fn nest(mut self, prefix: impl Into<String>, other: McpRouter) -> Self {
1149 let prefix = prefix.into();
1150 let inner = Arc::make_mut(&mut self.inner);
1151 let other_inner = other.inner;
1152
1153 for tool in other_inner.tools.values() {
1155 let prefixed_tool = tool.with_name_prefix(&prefix);
1156 inner
1157 .tools
1158 .insert(prefixed_tool.name.clone(), Arc::new(prefixed_tool));
1159 }
1160
1161 for (uri, resource) in &other_inner.resources {
1163 inner.resources.insert(uri.clone(), resource.clone());
1164 }
1165
1166 for template in &other_inner.resource_templates {
1168 inner.resource_templates.push(template.clone());
1169 }
1170
1171 for (name, prompt) in &other_inner.prompts {
1173 inner.prompts.insert(name.clone(), prompt.clone());
1174 }
1175
1176 self
1177 }
1178
1179 pub fn completion_handler<F, Fut>(mut self, handler: F) -> Self
1207 where
1208 F: Fn(CompleteParams) -> Fut + Send + Sync + 'static,
1209 Fut: Future<Output = Result<CompleteResult>> + Send + 'static,
1210 {
1211 Arc::make_mut(&mut self.inner).completion_handler =
1212 Some(Arc::new(move |params| Box::pin(handler(params))));
1213 self
1214 }
1215
1216 pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
1251 Arc::make_mut(&mut self.inner).tool_filter = Some(filter);
1252 self
1253 }
1254
1255 pub fn resource_filter(mut self, filter: ResourceFilter) -> Self {
1286 Arc::make_mut(&mut self.inner).resource_filter = Some(filter);
1287 self
1288 }
1289
1290 pub fn prompt_filter(mut self, filter: PromptFilter) -> Self {
1319 Arc::make_mut(&mut self.inner).prompt_filter = Some(filter);
1320 self
1321 }
1322
1323 pub fn session(&self) -> &SessionState {
1325 &self.session
1326 }
1327
1328 pub fn log(&self, params: LoggingMessageParams) -> bool {
1350 let Some(tx) = &self.inner.notification_tx else {
1351 return false;
1352 };
1353 tx.try_send(ServerNotification::LogMessage(params)).is_ok()
1354 }
1355
1356 pub fn log_info(&self, message: &str) -> bool {
1360 self.log(LoggingMessageParams::new(
1361 LogLevel::Info,
1362 serde_json::json!({ "message": message }),
1363 ))
1364 }
1365
1366 pub fn log_warning(&self, message: &str) -> bool {
1368 self.log(LoggingMessageParams::new(
1369 LogLevel::Warning,
1370 serde_json::json!({ "message": message }),
1371 ))
1372 }
1373
1374 pub fn log_error(&self, message: &str) -> bool {
1376 self.log(LoggingMessageParams::new(
1377 LogLevel::Error,
1378 serde_json::json!({ "message": message }),
1379 ))
1380 }
1381
1382 pub fn log_debug(&self, message: &str) -> bool {
1384 self.log(LoggingMessageParams::new(
1385 LogLevel::Debug,
1386 serde_json::json!({ "message": message }),
1387 ))
1388 }
1389
1390 pub fn is_subscribed(&self, uri: &str) -> bool {
1392 if let Ok(subs) = self.inner.subscriptions.read() {
1393 return subs.contains(uri);
1394 }
1395 false
1396 }
1397
1398 pub fn subscribed_uris(&self) -> Vec<String> {
1400 if let Ok(subs) = self.inner.subscriptions.read() {
1401 return subs.iter().cloned().collect();
1402 }
1403 Vec::new()
1404 }
1405
1406 fn subscribe(&self, uri: &str) -> bool {
1408 if let Ok(mut subs) = self.inner.subscriptions.write() {
1409 return subs.insert(uri.to_string());
1410 }
1411 false
1412 }
1413
1414 fn unsubscribe(&self, uri: &str) -> bool {
1416 if let Ok(mut subs) = self.inner.subscriptions.write() {
1417 return subs.remove(uri);
1418 }
1419 false
1420 }
1421
1422 pub fn notify_resource_updated(&self, uri: &str) -> bool {
1427 if !self.is_subscribed(uri) {
1429 return false;
1430 }
1431
1432 let Some(tx) = &self.inner.notification_tx else {
1433 return false;
1434 };
1435 tx.try_send(ServerNotification::ResourceUpdated {
1436 uri: uri.to_string(),
1437 })
1438 .is_ok()
1439 }
1440
1441 pub fn notify_resources_list_changed(&self) -> bool {
1445 let Some(tx) = &self.inner.notification_tx else {
1446 return false;
1447 };
1448 tx.try_send(ServerNotification::ResourcesListChanged)
1449 .is_ok()
1450 }
1451
1452 pub fn notify_tools_list_changed(&self) -> bool {
1456 let Some(tx) = &self.inner.notification_tx else {
1457 return false;
1458 };
1459 tx.try_send(ServerNotification::ToolsListChanged).is_ok()
1460 }
1461
1462 pub fn notify_prompts_list_changed(&self) -> bool {
1466 let Some(tx) = &self.inner.notification_tx else {
1467 return false;
1468 };
1469 tx.try_send(ServerNotification::PromptsListChanged).is_ok()
1470 }
1471
1472 fn capabilities(&self) -> ServerCapabilities {
1474 let has_resources =
1475 !self.inner.resources.is_empty() || !self.inner.resource_templates.is_empty();
1476 let has_notifications = self.inner.notification_tx.is_some();
1477
1478 #[cfg(feature = "dynamic-tools")]
1479 let has_dynamic_tools = self.inner.dynamic_tools.is_some();
1480 #[cfg(not(feature = "dynamic-tools"))]
1481 let has_dynamic_tools = false;
1482
1483 #[cfg(feature = "dynamic-tools")]
1484 let has_dynamic_prompts = self.inner.dynamic_prompts.is_some();
1485 #[cfg(not(feature = "dynamic-tools"))]
1486 let has_dynamic_prompts = false;
1487
1488 #[cfg(feature = "dynamic-tools")]
1489 let has_dynamic_resources = self.inner.dynamic_resources.is_some()
1490 || self.inner.dynamic_resource_templates.is_some();
1491 #[cfg(not(feature = "dynamic-tools"))]
1492 let has_dynamic_resources = false;
1493
1494 ServerCapabilities {
1495 tools: if self.inner.tools.is_empty() && !has_dynamic_tools {
1496 None
1497 } else {
1498 Some(ToolsCapability {
1499 list_changed: has_notifications,
1500 })
1501 },
1502 resources: if has_resources || has_dynamic_resources {
1503 Some(ResourcesCapability {
1504 subscribe: true,
1505 list_changed: has_notifications,
1506 })
1507 } else {
1508 None
1509 },
1510 prompts: if self.inner.prompts.is_empty() && !has_dynamic_prompts {
1511 None
1512 } else {
1513 Some(PromptsCapability {
1514 list_changed: has_notifications,
1515 })
1516 },
1517 logging: if self.inner.notification_tx.is_some() {
1519 Some(LoggingCapability::default())
1520 } else {
1521 None
1522 },
1523 tasks: {
1525 let has_task_support = self
1526 .inner
1527 .tools
1528 .values()
1529 .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1530 if has_task_support {
1531 Some(TasksCapability {
1532 list: Some(TasksListCapability {}),
1533 cancel: Some(TasksCancelCapability {}),
1534 requests: Some(TasksRequestsCapability {
1535 tools: Some(TasksToolsRequestsCapability {
1536 call: Some(TasksToolsCallCapability {}),
1537 }),
1538 }),
1539 })
1540 } else {
1541 None
1542 }
1543 },
1544 completions: if self.inner.completion_handler.is_some() {
1546 Some(CompletionsCapability::default())
1547 } else {
1548 None
1549 },
1550 experimental: None,
1551 extensions: None,
1552 }
1553 }
1554
1555 async fn handle(&self, request_id: RequestId, request: McpRequest) -> Result<McpResponse> {
1557 let method = request.method_name();
1559 if !self.session.is_request_allowed(method) {
1560 tracing::warn!(
1561 method = %method,
1562 phase = ?self.session.phase(),
1563 "Request rejected: session not initialized"
1564 );
1565 return Err(Error::JsonRpc(JsonRpcError::invalid_request(format!(
1566 "Session not initialized. Only 'initialize' and 'ping' are allowed before initialization. Got: {}",
1567 method
1568 ))));
1569 }
1570
1571 match request {
1572 McpRequest::Initialize(params) => {
1573 tracing::info!(
1574 client = %params.client_info.name,
1575 version = %params.client_info.version,
1576 "Client initializing"
1577 );
1578
1579 let protocol_version = if crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1582 .contains(¶ms.protocol_version.as_str())
1583 {
1584 params.protocol_version
1585 } else {
1586 crate::protocol::LATEST_PROTOCOL_VERSION.to_string()
1587 };
1588
1589 self.session.mark_initializing();
1591
1592 Ok(McpResponse::Initialize(InitializeResult {
1593 protocol_version,
1594 capabilities: self.capabilities(),
1595 server_info: Implementation {
1596 name: self.inner.server_name.clone(),
1597 version: self.inner.server_version.clone(),
1598 title: self.inner.server_title.clone(),
1599 description: self.inner.server_description.clone(),
1600 icons: self.inner.server_icons.clone(),
1601 website_url: self.inner.server_website_url.clone(),
1602 meta: None,
1603 },
1604 instructions: if let Some(config) = &self.inner.auto_instructions {
1605 Some(self.inner.generate_instructions(config))
1606 } else {
1607 self.inner.instructions.clone()
1608 },
1609 meta: None,
1610 }))
1611 }
1612
1613 McpRequest::ListTools(params) => {
1614 let filter = self.inner.tool_filter.as_ref();
1615 let is_visible = |t: &Tool| {
1616 filter
1617 .map(|f| f.is_visible(&self.session, t))
1618 .unwrap_or(true)
1619 };
1620
1621 let mut tools: Vec<ToolDefinition> = self
1623 .inner
1624 .tools
1625 .values()
1626 .filter(|t| is_visible(t))
1627 .map(|t| t.definition())
1628 .collect();
1629
1630 #[cfg(feature = "dynamic-tools")]
1632 if let Some(ref dynamic) = self.inner.dynamic_tools {
1633 let static_names: HashSet<String> =
1634 tools.iter().map(|t| t.name.clone()).collect();
1635 for t in dynamic.list() {
1636 if !static_names.contains(&t.name) && is_visible(&t) {
1637 tools.push(t.definition());
1638 }
1639 }
1640 }
1641
1642 tools.sort_by(|a, b| a.name.cmp(&b.name));
1643
1644 let (tools, next_cursor) =
1645 paginate(tools, params.cursor.as_deref(), self.inner.page_size)?;
1646
1647 Ok(McpResponse::ListTools(ListToolsResult {
1648 tools,
1649 next_cursor,
1650 meta: None,
1651 }))
1652 }
1653
1654 McpRequest::CallTool(params) => {
1655 let tool = self.inner.tools.get(¶ms.name).cloned();
1657 #[cfg(feature = "dynamic-tools")]
1658 let tool = tool.or_else(|| {
1659 self.inner
1660 .dynamic_tools
1661 .as_ref()
1662 .and_then(|d| d.get(¶ms.name))
1663 });
1664
1665 let tool = tool
1666 .ok_or_else(|| Error::JsonRpc(JsonRpcError::method_not_found(¶ms.name)))?;
1667
1668 if let Some(filter) = &self.inner.tool_filter
1670 && !filter.is_visible(&self.session, &tool)
1671 {
1672 return Err(filter.denial_error(¶ms.name));
1673 }
1674
1675 if let Some(task_params) = params.task {
1676 if matches!(tool.task_support, TaskSupportMode::Forbidden) {
1678 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1679 "Tool '{}' does not support async tasks",
1680 params.name
1681 ))));
1682 }
1683
1684 let (task_id, cancellation_token) = self.inner.task_store.create_task(
1686 ¶ms.name,
1687 params.arguments.clone(),
1688 task_params.ttl,
1689 );
1690
1691 tracing::info!(task_id = %task_id, tool = %params.name, "Created async task");
1692
1693 let progress_token = params.meta.and_then(|m| m.progress_token);
1695 let ctx = self.create_context(request_id, progress_token);
1696
1697 let task_store = self.inner.task_store.clone();
1699 let tool = tool.clone();
1700 let arguments = params.arguments;
1701 let task_id_clone = task_id.clone();
1702
1703 tokio::spawn(async move {
1704 if cancellation_token.is_cancelled() {
1706 tracing::debug!(task_id = %task_id_clone, "Task cancelled before execution");
1707 return;
1708 }
1709
1710 let result = tool.call_with_context(ctx, arguments).await;
1712
1713 if cancellation_token.is_cancelled() {
1714 tracing::debug!(task_id = %task_id_clone, "Task cancelled during execution");
1715 } else if result.is_error {
1716 let error_msg = result.first_text().unwrap_or("Tool execution failed");
1718 task_store.fail_task(&task_id_clone, error_msg);
1719 tracing::warn!(task_id = %task_id_clone, error = %error_msg, "Task failed");
1720 } else {
1721 task_store.complete_task(&task_id_clone, result);
1722 tracing::debug!(task_id = %task_id_clone, "Task completed successfully");
1723 }
1724 });
1725
1726 let task = self.inner.task_store.get_task(&task_id).ok_or_else(|| {
1727 Error::JsonRpc(JsonRpcError::internal_error(
1728 "Failed to retrieve created task",
1729 ))
1730 })?;
1731
1732 Ok(McpResponse::CreateTask(CreateTaskResult {
1733 task,
1734 meta: None,
1735 }))
1736 } else {
1737 if matches!(tool.task_support, TaskSupportMode::Required) {
1739 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1740 "Tool '{}' requires async task execution (include 'task' in params)",
1741 params.name
1742 ))));
1743 }
1744
1745 let progress_token = params.meta.and_then(|m| m.progress_token);
1747 let ctx = self.create_context(request_id, progress_token);
1748
1749 tracing::debug!(tool = %params.name, "Calling tool");
1750 let result = tool.call_with_context(ctx, params.arguments).await;
1751
1752 Ok(McpResponse::CallTool(result))
1753 }
1754 }
1755
1756 McpRequest::ListResources(params) => {
1757 let is_visible = |r: &Resource| -> bool {
1758 self.inner
1759 .resource_filter
1760 .as_ref()
1761 .map(|f| f.is_visible(&self.session, r))
1762 .unwrap_or(true)
1763 };
1764
1765 let mut resources: Vec<ResourceDefinition> = self
1766 .inner
1767 .resources
1768 .values()
1769 .filter(|r| is_visible(r))
1770 .map(|r| r.definition())
1771 .collect();
1772
1773 #[cfg(feature = "dynamic-tools")]
1775 if let Some(ref dynamic) = self.inner.dynamic_resources {
1776 let static_uris: HashSet<String> =
1777 resources.iter().map(|r| r.uri.clone()).collect();
1778 for r in dynamic.list() {
1779 if !static_uris.contains(&r.uri) && is_visible(&r) {
1780 resources.push(r.definition());
1781 }
1782 }
1783 }
1784
1785 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
1786
1787 let (resources, next_cursor) =
1788 paginate(resources, params.cursor.as_deref(), self.inner.page_size)?;
1789
1790 Ok(McpResponse::ListResources(ListResourcesResult {
1791 resources,
1792 next_cursor,
1793 meta: None,
1794 }))
1795 }
1796
1797 McpRequest::ListResourceTemplates(params) => {
1798 let mut resource_templates: Vec<ResourceTemplateDefinition> = self
1799 .inner
1800 .resource_templates
1801 .iter()
1802 .map(|t| t.definition())
1803 .collect();
1804
1805 #[cfg(feature = "dynamic-tools")]
1807 if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
1808 let static_patterns: HashSet<String> = resource_templates
1809 .iter()
1810 .map(|t| t.uri_template.clone())
1811 .collect();
1812 for t in dynamic.list() {
1813 if !static_patterns.contains(&t.uri_template) {
1814 resource_templates.push(t.definition());
1815 }
1816 }
1817 }
1818
1819 resource_templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
1820
1821 let (resource_templates, next_cursor) = paginate(
1822 resource_templates,
1823 params.cursor.as_deref(),
1824 self.inner.page_size,
1825 )?;
1826
1827 Ok(McpResponse::ListResourceTemplates(
1828 ListResourceTemplatesResult {
1829 resource_templates,
1830 next_cursor,
1831 meta: None,
1832 },
1833 ))
1834 }
1835
1836 McpRequest::ReadResource(params) => {
1837 if let Some(resource) = self.inner.resources.get(¶ms.uri) {
1839 if let Some(filter) = &self.inner.resource_filter
1841 && !filter.is_visible(&self.session, resource)
1842 {
1843 return Err(filter.denial_error(¶ms.uri));
1844 }
1845
1846 tracing::debug!(uri = %params.uri, "Reading static resource");
1847 let result = resource.read().await;
1848 return Ok(McpResponse::ReadResource(result));
1849 }
1850
1851 #[cfg(feature = "dynamic-tools")]
1853 #[allow(clippy::collapsible_if)]
1854 if let Some(ref dynamic) = self.inner.dynamic_resources {
1855 if let Some(resource) = dynamic.get(¶ms.uri) {
1856 if let Some(filter) = &self.inner.resource_filter
1857 && !filter.is_visible(&self.session, &resource)
1858 {
1859 return Err(filter.denial_error(¶ms.uri));
1860 }
1861 tracing::debug!(uri = %params.uri, "Reading dynamic resource");
1862 let result = resource.read().await;
1863 return Ok(McpResponse::ReadResource(result));
1864 }
1865 }
1866
1867 for template in &self.inner.resource_templates {
1869 if let Some(variables) = template.match_uri(¶ms.uri) {
1870 tracing::debug!(
1871 uri = %params.uri,
1872 template = %template.uri_template,
1873 "Reading resource via template"
1874 );
1875 let result = template.read(¶ms.uri, variables).await?;
1876 return Ok(McpResponse::ReadResource(result));
1877 }
1878 }
1879
1880 #[cfg(feature = "dynamic-tools")]
1882 #[allow(clippy::collapsible_if)]
1883 if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
1884 if let Some((template, variables)) = dynamic.match_uri(¶ms.uri) {
1885 tracing::debug!(
1886 uri = %params.uri,
1887 template = %template.uri_template,
1888 "Reading resource via dynamic template"
1889 );
1890 let result = template.read(¶ms.uri, variables).await?;
1891 return Ok(McpResponse::ReadResource(result));
1892 }
1893 }
1894
1895 Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1897 ¶ms.uri,
1898 )))
1899 }
1900
1901 McpRequest::SubscribeResource(params) => {
1902 if !self.inner.resources.contains_key(¶ms.uri) {
1904 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1905 ¶ms.uri,
1906 )));
1907 }
1908
1909 tracing::debug!(uri = %params.uri, "Subscribing to resource");
1910 self.subscribe(¶ms.uri);
1911
1912 Ok(McpResponse::SubscribeResource(EmptyResult {}))
1913 }
1914
1915 McpRequest::UnsubscribeResource(params) => {
1916 if !self.inner.resources.contains_key(¶ms.uri) {
1918 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1919 ¶ms.uri,
1920 )));
1921 }
1922
1923 tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
1924 self.unsubscribe(¶ms.uri);
1925
1926 Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
1927 }
1928
1929 McpRequest::ListPrompts(params) => {
1930 let is_visible = |p: &Prompt| -> bool {
1931 self.inner
1932 .prompt_filter
1933 .as_ref()
1934 .map(|f| f.is_visible(&self.session, p))
1935 .unwrap_or(true)
1936 };
1937
1938 let mut prompts: Vec<PromptDefinition> = self
1939 .inner
1940 .prompts
1941 .values()
1942 .filter(|p| is_visible(p))
1943 .map(|p| p.definition())
1944 .collect();
1945
1946 #[cfg(feature = "dynamic-tools")]
1948 if let Some(ref dynamic) = self.inner.dynamic_prompts {
1949 let static_names: HashSet<String> =
1950 prompts.iter().map(|p| p.name.clone()).collect();
1951 for p in dynamic.list() {
1952 if !static_names.contains(&p.name) && is_visible(&p) {
1953 prompts.push(p.definition());
1954 }
1955 }
1956 }
1957
1958 prompts.sort_by(|a, b| a.name.cmp(&b.name));
1959
1960 let (prompts, next_cursor) =
1961 paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
1962
1963 Ok(McpResponse::ListPrompts(ListPromptsResult {
1964 prompts,
1965 next_cursor,
1966 meta: None,
1967 }))
1968 }
1969
1970 McpRequest::GetPrompt(params) => {
1971 let prompt = self.inner.prompts.get(¶ms.name).cloned();
1973 #[cfg(feature = "dynamic-tools")]
1974 let prompt = prompt.or_else(|| {
1975 self.inner
1976 .dynamic_prompts
1977 .as_ref()
1978 .and_then(|d| d.get(¶ms.name))
1979 });
1980 let prompt = prompt.ok_or_else(|| {
1981 Error::JsonRpc(JsonRpcError::method_not_found(&format!(
1982 "Prompt not found: {}",
1983 params.name
1984 )))
1985 })?;
1986
1987 if let Some(filter) = &self.inner.prompt_filter
1989 && !filter.is_visible(&self.session, &prompt)
1990 {
1991 return Err(filter.denial_error(¶ms.name));
1992 }
1993
1994 tracing::debug!(name = %params.name, "Getting prompt");
1995 let result = prompt.get(params.arguments).await?;
1996
1997 Ok(McpResponse::GetPrompt(result))
1998 }
1999
2000 McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
2001
2002 McpRequest::ListTasks(params) => {
2003 let tasks = self.inner.task_store.list_tasks(params.status);
2004
2005 let (tasks, next_cursor) =
2006 paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
2007
2008 Ok(McpResponse::ListTasks(ListTasksResult {
2009 tasks,
2010 next_cursor,
2011 }))
2012 }
2013
2014 McpRequest::GetTaskInfo(params) => {
2015 let task = self
2016 .inner
2017 .task_store
2018 .get_task(¶ms.task_id)
2019 .ok_or_else(|| {
2020 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2021 "Task not found: {}",
2022 params.task_id
2023 )))
2024 })?;
2025
2026 Ok(McpResponse::GetTaskInfo(task))
2027 }
2028
2029 McpRequest::GetTaskResult(params) => {
2030 let (task_obj, result, error) = self
2032 .inner
2033 .task_store
2034 .wait_for_completion(¶ms.task_id)
2035 .await
2036 .ok_or_else(|| {
2037 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2038 "Task not found: {}",
2039 params.task_id
2040 )))
2041 })?;
2042
2043 let meta = serde_json::json!({
2045 "io.modelcontextprotocol/related-task": task_obj
2046 });
2047
2048 match task_obj.status {
2049 TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
2050 format!("Task {} was cancelled", params.task_id),
2051 ))),
2052 TaskStatus::Failed => {
2053 let mut call_result = CallToolResult::error(
2054 error.unwrap_or_else(|| "Task failed".to_string()),
2055 );
2056 call_result.meta = Some(meta);
2057 Ok(McpResponse::GetTaskResult(call_result))
2058 }
2059 _ => {
2060 let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
2061 call_result.meta = Some(meta);
2062 Ok(McpResponse::GetTaskResult(call_result))
2063 }
2064 }
2065 }
2066
2067 McpRequest::CancelTask(params) => {
2068 let current = self
2070 .inner
2071 .task_store
2072 .get_task(¶ms.task_id)
2073 .ok_or_else(|| {
2074 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2075 "Task not found: {}",
2076 params.task_id
2077 )))
2078 })?;
2079
2080 if current.status.is_terminal() {
2081 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
2082 "Task {} is already in terminal state: {}",
2083 params.task_id, current.status
2084 ))));
2085 }
2086
2087 let task_obj = self
2088 .inner
2089 .task_store
2090 .cancel_task(¶ms.task_id, params.reason.as_deref())
2091 .ok_or_else(|| {
2092 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2093 "Task not found: {}",
2094 params.task_id
2095 )))
2096 })?;
2097
2098 Ok(McpResponse::CancelTask(task_obj))
2099 }
2100
2101 McpRequest::SetLoggingLevel(params) => {
2102 tracing::debug!(level = ?params.level, "Client set logging level");
2103 if let Ok(mut level) = self.inner.min_log_level.write() {
2104 *level = params.level;
2105 }
2106 Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
2107 }
2108
2109 McpRequest::Complete(params) => {
2110 tracing::debug!(
2111 reference = ?params.reference,
2112 argument = %params.argument.name,
2113 "Completion request"
2114 );
2115
2116 if let Some(ref handler) = self.inner.completion_handler {
2118 let result = handler(params).await?;
2119 Ok(McpResponse::Complete(result))
2120 } else {
2121 Ok(McpResponse::Complete(CompleteResult::new(vec![])))
2123 }
2124 }
2125
2126 McpRequest::Unknown { method, .. } => {
2127 Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
2128 }
2129 _ => Err(Error::JsonRpc(JsonRpcError::method_not_found(
2130 "unknown method",
2131 ))),
2132 }
2133 }
2134
2135 pub fn handle_notification(&self, notification: McpNotification) {
2137 match notification {
2138 McpNotification::Initialized => {
2139 let phase_before = self.session.phase();
2140 if self.session.mark_initialized() {
2141 if phase_before == crate::session::SessionPhase::Uninitialized {
2142 tracing::info!(
2143 "Session initialized from uninitialized state (race resolved)"
2144 );
2145 } else {
2146 tracing::info!("Session initialized, entering operation phase");
2147 }
2148 } else {
2149 tracing::warn!(
2150 phase = ?self.session.phase(),
2151 "Received initialized notification in unexpected state"
2152 );
2153 }
2154 }
2155 McpNotification::Cancelled(params) => {
2156 if let Some(ref request_id) = params.request_id {
2157 if self.cancel_request(request_id) {
2158 tracing::info!(
2159 request_id = ?request_id,
2160 reason = ?params.reason,
2161 "Request cancelled"
2162 );
2163 } else {
2164 tracing::debug!(
2165 request_id = ?request_id,
2166 reason = ?params.reason,
2167 "Cancellation requested for unknown request"
2168 );
2169 }
2170 } else {
2171 tracing::debug!(
2172 reason = ?params.reason,
2173 "Cancellation notification received without request_id"
2174 );
2175 }
2176 }
2177 McpNotification::Progress(params) => {
2178 tracing::trace!(
2179 token = ?params.progress_token,
2180 progress = params.progress,
2181 total = ?params.total,
2182 "Progress notification"
2183 );
2184 }
2186 McpNotification::RootsListChanged => {
2187 tracing::info!("Client roots list changed");
2188 }
2191 McpNotification::Unknown { method, .. } => {
2192 tracing::debug!(method = %method, "Unknown notification received");
2193 }
2194 _ => {
2195 tracing::debug!("Unrecognized notification variant received");
2196 }
2197 }
2198 }
2199}
2200
2201impl Default for McpRouter {
2202 fn default() -> Self {
2203 Self::new()
2204 }
2205}
2206
2207pub use crate::context::Extensions;
2213
2214#[derive(Debug, Clone)]
2239pub struct ToolAnnotationsMap {
2240 map: Arc<HashMap<String, ToolAnnotations>>,
2241}
2242
2243impl ToolAnnotationsMap {
2244 pub fn get(&self, tool_name: &str) -> Option<&ToolAnnotations> {
2248 self.map.get(tool_name)
2249 }
2250
2251 pub fn is_read_only(&self, tool_name: &str) -> bool {
2256 self.map.get(tool_name).is_some_and(|a| a.read_only_hint)
2257 }
2258
2259 pub fn is_destructive(&self, tool_name: &str) -> bool {
2264 self.map.get(tool_name).is_none_or(|a| a.destructive_hint)
2265 }
2266
2267 pub fn is_idempotent(&self, tool_name: &str) -> bool {
2272 self.map.get(tool_name).is_some_and(|a| a.idempotent_hint)
2273 }
2274}
2275
2276#[derive(Debug, Clone)]
2278pub struct RouterRequest {
2279 pub id: RequestId,
2281 pub inner: McpRequest,
2283 pub extensions: Extensions,
2285}
2286
2287#[derive(Debug, Clone)]
2289pub struct RouterResponse {
2290 pub id: RequestId,
2292 pub inner: std::result::Result<McpResponse, JsonRpcError>,
2294}
2295
2296impl RouterResponse {
2297 pub fn into_jsonrpc(self) -> JsonRpcResponse {
2299 match self.inner {
2300 Ok(response) => match serde_json::to_value(response) {
2301 Ok(result) => JsonRpcResponse::result(self.id, result),
2302 Err(e) => {
2303 tracing::error!(error = %e, "Failed to serialize response");
2304 JsonRpcResponse::error(
2305 Some(self.id),
2306 JsonRpcError::internal_error(format!("Serialization error: {}", e)),
2307 )
2308 }
2309 },
2310 Err(error) => JsonRpcResponse::error(Some(self.id), error),
2311 }
2312 }
2313}
2314
2315impl Service<RouterRequest> for McpRouter {
2316 type Response = RouterResponse;
2317 type Error = std::convert::Infallible; type Future =
2319 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
2320
2321 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
2322 Poll::Ready(Ok(()))
2323 }
2324
2325 fn call(&mut self, req: RouterRequest) -> Self::Future {
2326 let router = self.clone();
2327 let request_id = req.id.clone();
2328 Box::pin(async move {
2329 let result = router.handle(req.id, req.inner).await;
2330 router.complete_request(&request_id);
2332 Ok(RouterResponse {
2333 id: request_id,
2334 inner: result.map_err(|e| match e {
2339 Error::JsonRpc(err) => err,
2340 Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
2341 e => JsonRpcError::internal_error(e.to_string()),
2342 }),
2343 })
2344 })
2345 }
2346}
2347
2348#[cfg(test)]
2349mod tests {
2350 use super::*;
2351 use crate::extract::{Context, Json};
2352 use crate::jsonrpc::JsonRpcService;
2353 use crate::tool::ToolBuilder;
2354 use schemars::JsonSchema;
2355 use serde::Deserialize;
2356 use tower::ServiceExt;
2357
2358 #[derive(Debug, Deserialize, JsonSchema)]
2359 struct AddInput {
2360 a: i64,
2361 b: i64,
2362 }
2363
2364 async fn init_router(router: &mut McpRouter) {
2366 let init_req = RouterRequest {
2368 id: RequestId::Number(0),
2369 inner: McpRequest::Initialize(InitializeParams {
2370 protocol_version: "2025-11-25".to_string(),
2371 capabilities: ClientCapabilities {
2372 roots: None,
2373 sampling: None,
2374 elicitation: None,
2375 tasks: None,
2376 experimental: None,
2377 extensions: None,
2378 },
2379 client_info: Implementation {
2380 name: "test".to_string(),
2381 version: "1.0".to_string(),
2382 ..Default::default()
2383 },
2384 meta: None,
2385 }),
2386 extensions: Extensions::new(),
2387 };
2388 let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
2389 router.handle_notification(McpNotification::Initialized);
2391 }
2392
2393 #[tokio::test]
2394 async fn test_router_list_tools() {
2395 let add_tool = ToolBuilder::new("add")
2396 .description("Add two numbers")
2397 .handler(|input: AddInput| async move {
2398 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2399 })
2400 .build();
2401
2402 let mut router = McpRouter::new().tool(add_tool);
2403
2404 init_router(&mut router).await;
2406
2407 let req = RouterRequest {
2408 id: RequestId::Number(1),
2409 inner: McpRequest::ListTools(ListToolsParams::default()),
2410 extensions: Extensions::new(),
2411 };
2412
2413 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2414
2415 match resp.inner {
2416 Ok(McpResponse::ListTools(result)) => {
2417 assert_eq!(result.tools.len(), 1);
2418 assert_eq!(result.tools[0].name, "add");
2419 }
2420 _ => panic!("Expected ListTools response"),
2421 }
2422 }
2423
2424 #[tokio::test]
2425 async fn test_router_call_tool() {
2426 let add_tool = ToolBuilder::new("add")
2427 .description("Add two numbers")
2428 .handler(|input: AddInput| async move {
2429 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2430 })
2431 .build();
2432
2433 let mut router = McpRouter::new().tool(add_tool);
2434
2435 init_router(&mut router).await;
2437
2438 let req = RouterRequest {
2439 id: RequestId::Number(1),
2440 inner: McpRequest::CallTool(CallToolParams {
2441 name: "add".to_string(),
2442 arguments: serde_json::json!({"a": 2, "b": 3}),
2443 meta: None,
2444 task: None,
2445 }),
2446 extensions: Extensions::new(),
2447 };
2448
2449 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2450
2451 match resp.inner {
2452 Ok(McpResponse::CallTool(result)) => {
2453 assert!(!result.is_error);
2454 match &result.content[0] {
2456 Content::Text { text, .. } => assert_eq!(text, "5"),
2457 _ => panic!("Expected text content"),
2458 }
2459 }
2460 _ => panic!("Expected CallTool response"),
2461 }
2462 }
2463
2464 async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2466 let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2467 "protocolVersion": "2025-11-25",
2468 "capabilities": {},
2469 "clientInfo": { "name": "test", "version": "1.0" }
2470 }));
2471 let _ = service.call_single(init_req).await.unwrap();
2472 router.handle_notification(McpNotification::Initialized);
2473 }
2474
2475 #[tokio::test]
2476 async fn test_jsonrpc_service() {
2477 let add_tool = ToolBuilder::new("add")
2478 .description("Add two numbers")
2479 .handler(|input: AddInput| async move {
2480 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2481 })
2482 .build();
2483
2484 let router = McpRouter::new().tool(add_tool);
2485 let mut service = JsonRpcService::new(router.clone());
2486
2487 init_jsonrpc_service(&mut service, &router).await;
2489
2490 let req = JsonRpcRequest::new(1, "tools/list");
2491
2492 let resp = service.call_single(req).await.unwrap();
2493
2494 match resp {
2495 JsonRpcResponse::Result(r) => {
2496 assert_eq!(r.id, RequestId::Number(1));
2497 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2498 assert_eq!(tools.len(), 1);
2499 }
2500 JsonRpcResponse::Error(_) => panic!("Expected success response"),
2501 _ => panic!("unexpected response variant"),
2502 }
2503 }
2504
2505 #[tokio::test]
2506 async fn test_batch_request() {
2507 let add_tool = ToolBuilder::new("add")
2508 .description("Add two numbers")
2509 .handler(|input: AddInput| async move {
2510 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2511 })
2512 .build();
2513
2514 let router = McpRouter::new().tool(add_tool);
2515 let mut service = JsonRpcService::new(router.clone());
2516
2517 init_jsonrpc_service(&mut service, &router).await;
2519
2520 let requests = vec![
2522 JsonRpcRequest::new(1, "tools/list"),
2523 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2524 "name": "add",
2525 "arguments": {"a": 10, "b": 20}
2526 })),
2527 JsonRpcRequest::new(3, "ping"),
2528 ];
2529
2530 let responses = service.call_batch(requests).await.unwrap();
2531
2532 assert_eq!(responses.len(), 3);
2533
2534 match &responses[0] {
2536 JsonRpcResponse::Result(r) => {
2537 assert_eq!(r.id, RequestId::Number(1));
2538 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2539 assert_eq!(tools.len(), 1);
2540 }
2541 JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2542 _ => panic!("unexpected response variant"),
2543 }
2544
2545 match &responses[1] {
2547 JsonRpcResponse::Result(r) => {
2548 assert_eq!(r.id, RequestId::Number(2));
2549 let content = r.result.get("content").unwrap().as_array().unwrap();
2550 let text = content[0].get("text").unwrap().as_str().unwrap();
2551 assert_eq!(text, "30");
2552 }
2553 JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2554 _ => panic!("unexpected response variant"),
2555 }
2556
2557 match &responses[2] {
2559 JsonRpcResponse::Result(r) => {
2560 assert_eq!(r.id, RequestId::Number(3));
2561 }
2562 JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2563 _ => panic!("unexpected response variant"),
2564 }
2565 }
2566
2567 #[tokio::test]
2568 async fn test_empty_batch_error() {
2569 let router = McpRouter::new();
2570 let mut service = JsonRpcService::new(router);
2571
2572 let result = service.call_batch(vec![]).await;
2573 assert!(result.is_err());
2574 }
2575
2576 #[tokio::test]
2581 async fn test_progress_token_extraction() {
2582 use crate::context::{ServerNotification, notification_channel};
2583 use crate::protocol::ProgressToken;
2584 use std::sync::Arc;
2585 use std::sync::atomic::{AtomicBool, Ordering};
2586
2587 let progress_reported = Arc::new(AtomicBool::new(false));
2589 let progress_ref = progress_reported.clone();
2590
2591 let tool = ToolBuilder::new("progress_tool")
2593 .description("Tool that reports progress")
2594 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2595 let reported = progress_ref.clone();
2596 async move {
2597 ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
2599 .await;
2600 reported.store(true, Ordering::SeqCst);
2601 Ok(CallToolResult::text("done"))
2602 }
2603 })
2604 .build();
2605
2606 let (tx, mut rx) = notification_channel(10);
2608 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2609 let mut service = JsonRpcService::new(router.clone());
2610
2611 init_jsonrpc_service(&mut service, &router).await;
2613
2614 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2616 "name": "progress_tool",
2617 "arguments": {"a": 1, "b": 2},
2618 "_meta": {
2619 "progressToken": "test-token-123"
2620 }
2621 }));
2622
2623 let resp = service.call_single(req).await.unwrap();
2624
2625 match resp {
2627 JsonRpcResponse::Result(_) => {}
2628 JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
2629 _ => panic!("unexpected response variant"),
2630 }
2631
2632 assert!(progress_reported.load(Ordering::SeqCst));
2634
2635 let notification = rx.try_recv().expect("Expected progress notification");
2637 match notification {
2638 ServerNotification::Progress(params) => {
2639 assert_eq!(
2640 params.progress_token,
2641 ProgressToken::String("test-token-123".to_string())
2642 );
2643 assert_eq!(params.progress, 50.0);
2644 assert_eq!(params.total, Some(100.0));
2645 assert_eq!(params.message.as_deref(), Some("Halfway"));
2646 }
2647 _ => panic!("Expected Progress notification"),
2648 }
2649 }
2650
2651 #[tokio::test]
2652 async fn test_tool_call_without_progress_token() {
2653 use crate::context::notification_channel;
2654 use std::sync::Arc;
2655 use std::sync::atomic::{AtomicBool, Ordering};
2656
2657 let progress_attempted = Arc::new(AtomicBool::new(false));
2658 let progress_ref = progress_attempted.clone();
2659
2660 let tool = ToolBuilder::new("no_token_tool")
2661 .description("Tool that tries to report progress without token")
2662 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2663 let attempted = progress_ref.clone();
2664 async move {
2665 ctx.report_progress(50.0, Some(100.0), None).await;
2667 attempted.store(true, Ordering::SeqCst);
2668 Ok(CallToolResult::text("done"))
2669 }
2670 })
2671 .build();
2672
2673 let (tx, mut rx) = notification_channel(10);
2674 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2675 let mut service = JsonRpcService::new(router.clone());
2676
2677 init_jsonrpc_service(&mut service, &router).await;
2678
2679 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2681 "name": "no_token_tool",
2682 "arguments": {"a": 1, "b": 2}
2683 }));
2684
2685 let resp = service.call_single(req).await.unwrap();
2686 assert!(matches!(resp, JsonRpcResponse::Result(_)));
2687
2688 assert!(progress_attempted.load(Ordering::SeqCst));
2690
2691 assert!(rx.try_recv().is_err());
2693 }
2694
2695 #[tokio::test]
2696 async fn test_batch_errors_returned_not_dropped() {
2697 let add_tool = ToolBuilder::new("add")
2698 .description("Add two numbers")
2699 .handler(|input: AddInput| async move {
2700 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2701 })
2702 .build();
2703
2704 let router = McpRouter::new().tool(add_tool);
2705 let mut service = JsonRpcService::new(router.clone());
2706
2707 init_jsonrpc_service(&mut service, &router).await;
2708
2709 let requests = vec![
2711 JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2713 "name": "add",
2714 "arguments": {"a": 10, "b": 20}
2715 })),
2716 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2718 "name": "nonexistent_tool",
2719 "arguments": {}
2720 })),
2721 JsonRpcRequest::new(3, "ping"),
2723 ];
2724
2725 let responses = service.call_batch(requests).await.unwrap();
2726
2727 assert_eq!(responses.len(), 3);
2729
2730 match &responses[0] {
2732 JsonRpcResponse::Result(r) => {
2733 assert_eq!(r.id, RequestId::Number(1));
2734 }
2735 JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
2736 _ => panic!("unexpected response variant"),
2737 }
2738
2739 match &responses[1] {
2741 JsonRpcResponse::Error(e) => {
2742 assert_eq!(e.id, Some(RequestId::Number(2)));
2743 assert!(e.error.message.contains("not found") || e.error.code == -32601);
2745 }
2746 JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
2747 _ => panic!("unexpected response variant"),
2748 }
2749
2750 match &responses[2] {
2752 JsonRpcResponse::Result(r) => {
2753 assert_eq!(r.id, RequestId::Number(3));
2754 }
2755 JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
2756 _ => panic!("unexpected response variant"),
2757 }
2758 }
2759
2760 #[tokio::test]
2765 async fn test_list_resource_templates() {
2766 use crate::resource::ResourceTemplateBuilder;
2767 use std::collections::HashMap;
2768
2769 let template = ResourceTemplateBuilder::new("file:///{path}")
2770 .name("Project Files")
2771 .description("Access project files")
2772 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2773 Ok(ReadResourceResult {
2774 contents: vec![ResourceContent {
2775 uri,
2776 mime_type: None,
2777 text: None,
2778 blob: None,
2779 meta: None,
2780 }],
2781 meta: None,
2782 })
2783 });
2784
2785 let mut router = McpRouter::new().resource_template(template);
2786
2787 init_router(&mut router).await;
2789
2790 let req = RouterRequest {
2791 id: RequestId::Number(1),
2792 inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
2793 extensions: Extensions::new(),
2794 };
2795
2796 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2797
2798 match resp.inner {
2799 Ok(McpResponse::ListResourceTemplates(result)) => {
2800 assert_eq!(result.resource_templates.len(), 1);
2801 assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
2802 assert_eq!(result.resource_templates[0].name, "Project Files");
2803 }
2804 _ => panic!("Expected ListResourceTemplates response"),
2805 }
2806 }
2807
2808 #[tokio::test]
2809 async fn test_read_resource_via_template() {
2810 use crate::resource::ResourceTemplateBuilder;
2811 use std::collections::HashMap;
2812
2813 let template = ResourceTemplateBuilder::new("db://users/{id}")
2814 .name("User Records")
2815 .handler(|uri: String, vars: HashMap<String, String>| async move {
2816 let id = vars.get("id").unwrap().clone();
2817 Ok(ReadResourceResult {
2818 contents: vec![ResourceContent {
2819 uri,
2820 mime_type: Some("application/json".to_string()),
2821 text: Some(format!(r#"{{"id": "{}"}}"#, id)),
2822 blob: None,
2823 meta: None,
2824 }],
2825 meta: None,
2826 })
2827 });
2828
2829 let mut router = McpRouter::new().resource_template(template);
2830
2831 init_router(&mut router).await;
2833
2834 let req = RouterRequest {
2836 id: RequestId::Number(1),
2837 inner: McpRequest::ReadResource(ReadResourceParams {
2838 uri: "db://users/123".to_string(),
2839 meta: None,
2840 }),
2841 extensions: Extensions::new(),
2842 };
2843
2844 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2845
2846 match resp.inner {
2847 Ok(McpResponse::ReadResource(result)) => {
2848 assert_eq!(result.contents.len(), 1);
2849 assert_eq!(result.contents[0].uri, "db://users/123");
2850 assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
2851 }
2852 _ => panic!("Expected ReadResource response"),
2853 }
2854 }
2855
2856 #[tokio::test]
2857 async fn test_static_resource_takes_precedence_over_template() {
2858 use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
2859 use std::collections::HashMap;
2860
2861 let template = ResourceTemplateBuilder::new("file:///{path}")
2863 .name("Files Template")
2864 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2865 Ok(ReadResourceResult {
2866 contents: vec![ResourceContent {
2867 uri,
2868 mime_type: None,
2869 text: Some("from template".to_string()),
2870 blob: None,
2871 meta: None,
2872 }],
2873 meta: None,
2874 })
2875 });
2876
2877 let static_resource = ResourceBuilder::new("file:///README.md")
2879 .name("README")
2880 .text("from static resource");
2881
2882 let mut router = McpRouter::new()
2883 .resource_template(template)
2884 .resource(static_resource);
2885
2886 init_router(&mut router).await;
2888
2889 let req = RouterRequest {
2891 id: RequestId::Number(1),
2892 inner: McpRequest::ReadResource(ReadResourceParams {
2893 uri: "file:///README.md".to_string(),
2894 meta: None,
2895 }),
2896 extensions: Extensions::new(),
2897 };
2898
2899 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2900
2901 match resp.inner {
2902 Ok(McpResponse::ReadResource(result)) => {
2903 assert_eq!(
2905 result.contents[0].text.as_deref(),
2906 Some("from static resource")
2907 );
2908 }
2909 _ => panic!("Expected ReadResource response"),
2910 }
2911 }
2912
2913 #[tokio::test]
2914 async fn test_resource_not_found_when_no_match() {
2915 use crate::resource::ResourceTemplateBuilder;
2916 use std::collections::HashMap;
2917
2918 let template = ResourceTemplateBuilder::new("db://users/{id}")
2919 .name("Users")
2920 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2921 Ok(ReadResourceResult {
2922 contents: vec![ResourceContent {
2923 uri,
2924 mime_type: None,
2925 text: None,
2926 blob: None,
2927 meta: None,
2928 }],
2929 meta: None,
2930 })
2931 });
2932
2933 let mut router = McpRouter::new().resource_template(template);
2934
2935 init_router(&mut router).await;
2937
2938 let req = RouterRequest {
2940 id: RequestId::Number(1),
2941 inner: McpRequest::ReadResource(ReadResourceParams {
2942 uri: "db://posts/123".to_string(),
2943 meta: None,
2944 }),
2945 extensions: Extensions::new(),
2946 };
2947
2948 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2949
2950 match resp.inner {
2951 Err(err) => {
2952 assert!(err.message.contains("not found"));
2953 }
2954 Ok(_) => panic!("Expected error for non-matching URI"),
2955 }
2956 }
2957
2958 #[tokio::test]
2959 async fn test_capabilities_include_resources_with_only_templates() {
2960 use crate::resource::ResourceTemplateBuilder;
2961 use std::collections::HashMap;
2962
2963 let template = ResourceTemplateBuilder::new("file:///{path}")
2964 .name("Files")
2965 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2966 Ok(ReadResourceResult {
2967 contents: vec![ResourceContent {
2968 uri,
2969 mime_type: None,
2970 text: None,
2971 blob: None,
2972 meta: None,
2973 }],
2974 meta: None,
2975 })
2976 });
2977
2978 let mut router = McpRouter::new().resource_template(template);
2979
2980 let init_req = RouterRequest {
2982 id: RequestId::Number(0),
2983 inner: McpRequest::Initialize(InitializeParams {
2984 protocol_version: "2025-11-25".to_string(),
2985 capabilities: ClientCapabilities {
2986 roots: None,
2987 sampling: None,
2988 elicitation: None,
2989 tasks: None,
2990 experimental: None,
2991 extensions: None,
2992 },
2993 client_info: Implementation {
2994 name: "test".to_string(),
2995 version: "1.0".to_string(),
2996 ..Default::default()
2997 },
2998 meta: None,
2999 }),
3000 extensions: Extensions::new(),
3001 };
3002 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3003
3004 match resp.inner {
3005 Ok(McpResponse::Initialize(result)) => {
3006 assert!(result.capabilities.resources.is_some());
3008 }
3009 _ => panic!("Expected Initialize response"),
3010 }
3011 }
3012
3013 #[tokio::test]
3018 async fn test_log_sends_notification() {
3019 use crate::context::notification_channel;
3020
3021 let (tx, mut rx) = notification_channel(10);
3022 let router = McpRouter::new().with_notification_sender(tx);
3023
3024 let sent = router.log_info("Test message");
3026 assert!(sent);
3027
3028 let notification = rx.try_recv().unwrap();
3030 match notification {
3031 ServerNotification::LogMessage(params) => {
3032 assert_eq!(params.level, LogLevel::Info);
3033 let data = params.data;
3034 assert_eq!(
3035 data.get("message").unwrap().as_str().unwrap(),
3036 "Test message"
3037 );
3038 }
3039 _ => panic!("Expected LogMessage notification"),
3040 }
3041 }
3042
3043 #[tokio::test]
3044 async fn test_log_with_custom_params() {
3045 use crate::context::notification_channel;
3046
3047 let (tx, mut rx) = notification_channel(10);
3048 let router = McpRouter::new().with_notification_sender(tx);
3049
3050 let params = LoggingMessageParams::new(
3052 LogLevel::Error,
3053 serde_json::json!({
3054 "error": "Connection failed",
3055 "host": "localhost"
3056 }),
3057 )
3058 .with_logger("database");
3059
3060 let sent = router.log(params);
3061 assert!(sent);
3062
3063 let notification = rx.try_recv().unwrap();
3064 match notification {
3065 ServerNotification::LogMessage(params) => {
3066 assert_eq!(params.level, LogLevel::Error);
3067 assert_eq!(params.logger.as_deref(), Some("database"));
3068 let data = params.data;
3069 assert_eq!(
3070 data.get("error").unwrap().as_str().unwrap(),
3071 "Connection failed"
3072 );
3073 }
3074 _ => panic!("Expected LogMessage notification"),
3075 }
3076 }
3077
3078 #[tokio::test]
3079 async fn test_log_without_channel_returns_false() {
3080 let router = McpRouter::new();
3082
3083 assert!(!router.log_info("Test"));
3085 assert!(!router.log_warning("Test"));
3086 assert!(!router.log_error("Test"));
3087 assert!(!router.log_debug("Test"));
3088 }
3089
3090 #[tokio::test]
3091 async fn test_logging_capability_with_channel() {
3092 use crate::context::notification_channel;
3093
3094 let (tx, _rx) = notification_channel(10);
3095 let mut router = McpRouter::new().with_notification_sender(tx);
3096
3097 let init_req = RouterRequest {
3099 id: RequestId::Number(0),
3100 inner: McpRequest::Initialize(InitializeParams {
3101 protocol_version: "2025-11-25".to_string(),
3102 capabilities: ClientCapabilities {
3103 roots: None,
3104 sampling: None,
3105 elicitation: None,
3106 tasks: None,
3107 experimental: None,
3108 extensions: None,
3109 },
3110 client_info: Implementation {
3111 name: "test".to_string(),
3112 version: "1.0".to_string(),
3113 ..Default::default()
3114 },
3115 meta: None,
3116 }),
3117 extensions: Extensions::new(),
3118 };
3119 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3120
3121 match resp.inner {
3122 Ok(McpResponse::Initialize(result)) => {
3123 assert!(result.capabilities.logging.is_some());
3125 }
3126 _ => panic!("Expected Initialize response"),
3127 }
3128 }
3129
3130 #[tokio::test]
3131 async fn test_no_logging_capability_without_channel() {
3132 let mut router = McpRouter::new();
3133
3134 let init_req = RouterRequest {
3136 id: RequestId::Number(0),
3137 inner: McpRequest::Initialize(InitializeParams {
3138 protocol_version: "2025-11-25".to_string(),
3139 capabilities: ClientCapabilities {
3140 roots: None,
3141 sampling: None,
3142 elicitation: None,
3143 tasks: None,
3144 experimental: None,
3145 extensions: None,
3146 },
3147 client_info: Implementation {
3148 name: "test".to_string(),
3149 version: "1.0".to_string(),
3150 ..Default::default()
3151 },
3152 meta: None,
3153 }),
3154 extensions: Extensions::new(),
3155 };
3156 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3157
3158 match resp.inner {
3159 Ok(McpResponse::Initialize(result)) => {
3160 assert!(result.capabilities.logging.is_none());
3162 }
3163 _ => panic!("Expected Initialize response"),
3164 }
3165 }
3166
3167 #[tokio::test]
3172 async fn test_create_task_via_call_tool() {
3173 let add_tool = ToolBuilder::new("add")
3174 .description("Add two numbers")
3175 .task_support(TaskSupportMode::Optional)
3176 .handler(|input: AddInput| async move {
3177 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3178 })
3179 .build();
3180
3181 let mut router = McpRouter::new().tool(add_tool);
3182 init_router(&mut router).await;
3183
3184 let req = RouterRequest {
3185 id: RequestId::Number(1),
3186 inner: McpRequest::CallTool(CallToolParams {
3187 name: "add".to_string(),
3188 arguments: serde_json::json!({"a": 5, "b": 10}),
3189 meta: None,
3190 task: Some(TaskRequestParams { ttl: None }),
3191 }),
3192 extensions: Extensions::new(),
3193 };
3194
3195 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3196
3197 match resp.inner {
3198 Ok(McpResponse::CreateTask(result)) => {
3199 assert!(result.task.task_id.starts_with("task-"));
3200 assert_eq!(result.task.status, TaskStatus::Working);
3201 }
3202 _ => panic!("Expected CreateTask response"),
3203 }
3204 }
3205
3206 #[tokio::test]
3207 async fn test_list_tasks_empty() {
3208 let mut router = McpRouter::new();
3209 init_router(&mut router).await;
3210
3211 let req = RouterRequest {
3212 id: RequestId::Number(1),
3213 inner: McpRequest::ListTasks(ListTasksParams::default()),
3214 extensions: Extensions::new(),
3215 };
3216
3217 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3218
3219 match resp.inner {
3220 Ok(McpResponse::ListTasks(result)) => {
3221 assert!(result.tasks.is_empty());
3222 }
3223 _ => panic!("Expected ListTasks response"),
3224 }
3225 }
3226
3227 #[tokio::test]
3228 async fn test_task_lifecycle_complete() {
3229 let add_tool = ToolBuilder::new("add")
3230 .description("Add two numbers")
3231 .task_support(TaskSupportMode::Optional)
3232 .handler(|input: AddInput| async move {
3233 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3234 })
3235 .build();
3236
3237 let mut router = McpRouter::new().tool(add_tool);
3238 init_router(&mut router).await;
3239
3240 let req = RouterRequest {
3242 id: RequestId::Number(1),
3243 inner: McpRequest::CallTool(CallToolParams {
3244 name: "add".to_string(),
3245 arguments: serde_json::json!({"a": 7, "b": 8}),
3246 meta: None,
3247 task: Some(TaskRequestParams { ttl: None }),
3248 }),
3249 extensions: Extensions::new(),
3250 };
3251
3252 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3253 let task_id = match resp.inner {
3254 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3255 _ => panic!("Expected CreateTask response"),
3256 };
3257
3258 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3260
3261 let req = RouterRequest {
3263 id: RequestId::Number(2),
3264 inner: McpRequest::GetTaskResult(GetTaskResultParams {
3265 task_id: task_id.clone(),
3266 meta: None,
3267 }),
3268 extensions: Extensions::new(),
3269 };
3270
3271 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3272
3273 match resp.inner {
3274 Ok(McpResponse::GetTaskResult(result)) => {
3275 assert!(result.meta.is_some());
3277 match &result.content[0] {
3279 Content::Text { text, .. } => assert_eq!(text, "15"),
3280 _ => panic!("Expected text content"),
3281 }
3282 }
3283 _ => panic!("Expected GetTaskResult response"),
3284 }
3285 }
3286
3287 #[tokio::test]
3288 async fn test_task_cancellation() {
3289 let slow_tool = ToolBuilder::new("slow")
3291 .description("Slow tool")
3292 .task_support(TaskSupportMode::Optional)
3293 .handler(|_input: serde_json::Value| async move {
3294 tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
3295 Ok(CallToolResult::text("done"))
3296 })
3297 .build();
3298
3299 let mut router = McpRouter::new().tool(slow_tool);
3300 init_router(&mut router).await;
3301
3302 let req = RouterRequest {
3304 id: RequestId::Number(1),
3305 inner: McpRequest::CallTool(CallToolParams {
3306 name: "slow".to_string(),
3307 arguments: serde_json::json!({}),
3308 meta: None,
3309 task: Some(TaskRequestParams { ttl: None }),
3310 }),
3311 extensions: Extensions::new(),
3312 };
3313
3314 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3315 let task_id = match resp.inner {
3316 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3317 _ => panic!("Expected CreateTask response"),
3318 };
3319
3320 let req = RouterRequest {
3322 id: RequestId::Number(2),
3323 inner: McpRequest::CancelTask(CancelTaskParams {
3324 task_id: task_id.clone(),
3325 reason: Some("Test cancellation".to_string()),
3326 meta: None,
3327 }),
3328 extensions: Extensions::new(),
3329 };
3330
3331 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3332
3333 match resp.inner {
3334 Ok(McpResponse::CancelTask(task_obj)) => {
3335 assert_eq!(task_obj.status, TaskStatus::Cancelled);
3336 }
3337 _ => panic!("Expected CancelTask response"),
3338 }
3339 }
3340
3341 #[tokio::test]
3342 async fn test_get_task_info() {
3343 let add_tool = ToolBuilder::new("add")
3344 .description("Add two numbers")
3345 .task_support(TaskSupportMode::Optional)
3346 .handler(|input: AddInput| async move {
3347 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3348 })
3349 .build();
3350
3351 let mut router = McpRouter::new().tool(add_tool);
3352 init_router(&mut router).await;
3353
3354 let req = RouterRequest {
3356 id: RequestId::Number(1),
3357 inner: McpRequest::CallTool(CallToolParams {
3358 name: "add".to_string(),
3359 arguments: serde_json::json!({"a": 1, "b": 2}),
3360 meta: None,
3361 task: Some(TaskRequestParams { ttl: Some(600_000) }),
3362 }),
3363 extensions: Extensions::new(),
3364 };
3365
3366 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3367 let task_id = match resp.inner {
3368 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3369 _ => panic!("Expected CreateTask response"),
3370 };
3371
3372 let req = RouterRequest {
3374 id: RequestId::Number(2),
3375 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3376 task_id: task_id.clone(),
3377 meta: None,
3378 }),
3379 extensions: Extensions::new(),
3380 };
3381
3382 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3383
3384 match resp.inner {
3385 Ok(McpResponse::GetTaskInfo(info)) => {
3386 assert_eq!(info.task_id, task_id);
3387 assert!(info.created_at.contains('T')); assert_eq!(info.ttl, Some(600_000));
3389 }
3390 _ => panic!("Expected GetTaskInfo response"),
3391 }
3392 }
3393
3394 #[tokio::test]
3395 async fn test_task_forbidden_tool_rejects_task_params() {
3396 let tool = ToolBuilder::new("sync_only")
3397 .description("Sync only tool")
3398 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3399 .build();
3400
3401 let mut router = McpRouter::new().tool(tool);
3402 init_router(&mut router).await;
3403
3404 let req = RouterRequest {
3406 id: RequestId::Number(1),
3407 inner: McpRequest::CallTool(CallToolParams {
3408 name: "sync_only".to_string(),
3409 arguments: serde_json::json!({}),
3410 meta: None,
3411 task: Some(TaskRequestParams { ttl: None }),
3412 }),
3413 extensions: Extensions::new(),
3414 };
3415
3416 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3417
3418 match resp.inner {
3419 Err(e) => {
3420 assert!(e.message.contains("does not support async tasks"));
3421 }
3422 _ => panic!("Expected error response"),
3423 }
3424 }
3425
3426 #[tokio::test]
3427 async fn test_get_nonexistent_task() {
3428 let mut router = McpRouter::new();
3429 init_router(&mut router).await;
3430
3431 let req = RouterRequest {
3432 id: RequestId::Number(1),
3433 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3434 task_id: "task-999".to_string(),
3435 meta: None,
3436 }),
3437 extensions: Extensions::new(),
3438 };
3439
3440 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3441
3442 match resp.inner {
3443 Err(e) => {
3444 assert!(e.message.contains("not found"));
3445 }
3446 _ => panic!("Expected error response"),
3447 }
3448 }
3449
3450 #[tokio::test]
3455 async fn test_subscribe_to_resource() {
3456 use crate::resource::ResourceBuilder;
3457
3458 let resource = ResourceBuilder::new("file:///test.txt")
3459 .name("Test File")
3460 .text("Hello");
3461
3462 let mut router = McpRouter::new().resource(resource);
3463 init_router(&mut router).await;
3464
3465 let req = RouterRequest {
3467 id: RequestId::Number(1),
3468 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3469 uri: "file:///test.txt".to_string(),
3470 meta: None,
3471 }),
3472 extensions: Extensions::new(),
3473 };
3474
3475 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3476
3477 match resp.inner {
3478 Ok(McpResponse::SubscribeResource(_)) => {
3479 assert!(router.is_subscribed("file:///test.txt"));
3481 }
3482 _ => panic!("Expected SubscribeResource response"),
3483 }
3484 }
3485
3486 #[tokio::test]
3487 async fn test_unsubscribe_from_resource() {
3488 use crate::resource::ResourceBuilder;
3489
3490 let resource = ResourceBuilder::new("file:///test.txt")
3491 .name("Test File")
3492 .text("Hello");
3493
3494 let mut router = McpRouter::new().resource(resource);
3495 init_router(&mut router).await;
3496
3497 let req = RouterRequest {
3499 id: RequestId::Number(1),
3500 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3501 uri: "file:///test.txt".to_string(),
3502 meta: None,
3503 }),
3504 extensions: Extensions::new(),
3505 };
3506 let _ = router.ready().await.unwrap().call(req).await.unwrap();
3507 assert!(router.is_subscribed("file:///test.txt"));
3508
3509 let req = RouterRequest {
3511 id: RequestId::Number(2),
3512 inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3513 uri: "file:///test.txt".to_string(),
3514 meta: None,
3515 }),
3516 extensions: Extensions::new(),
3517 };
3518
3519 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3520
3521 match resp.inner {
3522 Ok(McpResponse::UnsubscribeResource(_)) => {
3523 assert!(!router.is_subscribed("file:///test.txt"));
3525 }
3526 _ => panic!("Expected UnsubscribeResource response"),
3527 }
3528 }
3529
3530 #[tokio::test]
3531 async fn test_subscribe_nonexistent_resource() {
3532 let mut router = McpRouter::new();
3533 init_router(&mut router).await;
3534
3535 let req = RouterRequest {
3536 id: RequestId::Number(1),
3537 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3538 uri: "file:///nonexistent.txt".to_string(),
3539 meta: None,
3540 }),
3541 extensions: Extensions::new(),
3542 };
3543
3544 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3545
3546 match resp.inner {
3547 Err(e) => {
3548 assert!(e.message.contains("not found"));
3549 }
3550 _ => panic!("Expected error response"),
3551 }
3552 }
3553
3554 #[tokio::test]
3555 async fn test_notify_resource_updated() {
3556 use crate::context::notification_channel;
3557 use crate::resource::ResourceBuilder;
3558
3559 let (tx, mut rx) = notification_channel(10);
3560
3561 let resource = ResourceBuilder::new("file:///test.txt")
3562 .name("Test File")
3563 .text("Hello");
3564
3565 let router = McpRouter::new()
3566 .resource(resource)
3567 .with_notification_sender(tx);
3568
3569 router.subscribe("file:///test.txt");
3571
3572 let sent = router.notify_resource_updated("file:///test.txt");
3574 assert!(sent);
3575
3576 let notification = rx.try_recv().unwrap();
3578 match notification {
3579 ServerNotification::ResourceUpdated { uri } => {
3580 assert_eq!(uri, "file:///test.txt");
3581 }
3582 _ => panic!("Expected ResourceUpdated notification"),
3583 }
3584 }
3585
3586 #[tokio::test]
3587 async fn test_notify_resource_updated_not_subscribed() {
3588 use crate::context::notification_channel;
3589 use crate::resource::ResourceBuilder;
3590
3591 let (tx, mut rx) = notification_channel(10);
3592
3593 let resource = ResourceBuilder::new("file:///test.txt")
3594 .name("Test File")
3595 .text("Hello");
3596
3597 let router = McpRouter::new()
3598 .resource(resource)
3599 .with_notification_sender(tx);
3600
3601 let sent = router.notify_resource_updated("file:///test.txt");
3603 assert!(!sent); assert!(rx.try_recv().is_err());
3607 }
3608
3609 #[tokio::test]
3610 async fn test_notify_resources_list_changed() {
3611 use crate::context::notification_channel;
3612
3613 let (tx, mut rx) = notification_channel(10);
3614 let router = McpRouter::new().with_notification_sender(tx);
3615
3616 let sent = router.notify_resources_list_changed();
3617 assert!(sent);
3618
3619 let notification = rx.try_recv().unwrap();
3620 match notification {
3621 ServerNotification::ResourcesListChanged => {}
3622 _ => panic!("Expected ResourcesListChanged notification"),
3623 }
3624 }
3625
3626 #[tokio::test]
3627 async fn test_subscribed_uris() {
3628 use crate::resource::ResourceBuilder;
3629
3630 let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
3631
3632 let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
3633
3634 let router = McpRouter::new().resource(resource1).resource(resource2);
3635
3636 router.subscribe("file:///a.txt");
3638 router.subscribe("file:///b.txt");
3639
3640 let uris = router.subscribed_uris();
3641 assert_eq!(uris.len(), 2);
3642 assert!(uris.contains(&"file:///a.txt".to_string()));
3643 assert!(uris.contains(&"file:///b.txt".to_string()));
3644 }
3645
3646 #[tokio::test]
3647 async fn test_subscription_capability_advertised() {
3648 use crate::resource::ResourceBuilder;
3649
3650 let resource = ResourceBuilder::new("file:///test.txt")
3651 .name("Test")
3652 .text("Hello");
3653
3654 let mut router = McpRouter::new().resource(resource);
3655
3656 let init_req = RouterRequest {
3658 id: RequestId::Number(0),
3659 inner: McpRequest::Initialize(InitializeParams {
3660 protocol_version: "2025-11-25".to_string(),
3661 capabilities: ClientCapabilities {
3662 roots: None,
3663 sampling: None,
3664 elicitation: None,
3665 tasks: None,
3666 experimental: None,
3667 extensions: None,
3668 },
3669 client_info: Implementation {
3670 name: "test".to_string(),
3671 version: "1.0".to_string(),
3672 ..Default::default()
3673 },
3674 meta: None,
3675 }),
3676 extensions: Extensions::new(),
3677 };
3678 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3679
3680 match resp.inner {
3681 Ok(McpResponse::Initialize(result)) => {
3682 let resources_cap = result.capabilities.resources.unwrap();
3684 assert!(resources_cap.subscribe);
3685 }
3686 _ => panic!("Expected Initialize response"),
3687 }
3688 }
3689
3690 #[tokio::test]
3691 async fn test_completion_handler() {
3692 let router = McpRouter::new()
3693 .server_info("test", "1.0")
3694 .completion_handler(|params: CompleteParams| async move {
3695 let prefix = ¶ms.argument.value;
3697 let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
3698 .into_iter()
3699 .filter(|s| s.starts_with(prefix))
3700 .map(String::from)
3701 .collect();
3702 Ok(CompleteResult::new(suggestions))
3703 });
3704
3705 let init_req = RouterRequest {
3707 id: RequestId::Number(0),
3708 inner: McpRequest::Initialize(InitializeParams {
3709 protocol_version: "2025-11-25".to_string(),
3710 capabilities: ClientCapabilities::default(),
3711 client_info: Implementation {
3712 name: "test".to_string(),
3713 version: "1.0".to_string(),
3714 ..Default::default()
3715 },
3716 meta: None,
3717 }),
3718 extensions: Extensions::new(),
3719 };
3720 let resp = router
3721 .clone()
3722 .ready()
3723 .await
3724 .unwrap()
3725 .call(init_req)
3726 .await
3727 .unwrap();
3728
3729 match resp.inner {
3731 Ok(McpResponse::Initialize(result)) => {
3732 assert!(result.capabilities.completions.is_some());
3733 }
3734 _ => panic!("Expected Initialize response"),
3735 }
3736
3737 router.handle_notification(McpNotification::Initialized);
3739
3740 let complete_req = RouterRequest {
3742 id: RequestId::Number(1),
3743 inner: McpRequest::Complete(CompleteParams {
3744 reference: CompletionReference::prompt("test-prompt"),
3745 argument: CompletionArgument::new("query", "al"),
3746 context: None,
3747 meta: None,
3748 }),
3749 extensions: Extensions::new(),
3750 };
3751 let resp = router
3752 .clone()
3753 .ready()
3754 .await
3755 .unwrap()
3756 .call(complete_req)
3757 .await
3758 .unwrap();
3759
3760 match resp.inner {
3761 Ok(McpResponse::Complete(result)) => {
3762 assert_eq!(result.completion.values, vec!["alpha"]);
3763 }
3764 _ => panic!("Expected Complete response"),
3765 }
3766 }
3767
3768 #[tokio::test]
3769 async fn test_completion_without_handler_returns_empty() {
3770 let router = McpRouter::new().server_info("test", "1.0");
3771
3772 let init_req = RouterRequest {
3774 id: RequestId::Number(0),
3775 inner: McpRequest::Initialize(InitializeParams {
3776 protocol_version: "2025-11-25".to_string(),
3777 capabilities: ClientCapabilities::default(),
3778 client_info: Implementation {
3779 name: "test".to_string(),
3780 version: "1.0".to_string(),
3781 ..Default::default()
3782 },
3783 meta: None,
3784 }),
3785 extensions: Extensions::new(),
3786 };
3787 let resp = router
3788 .clone()
3789 .ready()
3790 .await
3791 .unwrap()
3792 .call(init_req)
3793 .await
3794 .unwrap();
3795
3796 match resp.inner {
3798 Ok(McpResponse::Initialize(result)) => {
3799 assert!(result.capabilities.completions.is_none());
3800 }
3801 _ => panic!("Expected Initialize response"),
3802 }
3803
3804 router.handle_notification(McpNotification::Initialized);
3806
3807 let complete_req = RouterRequest {
3809 id: RequestId::Number(1),
3810 inner: McpRequest::Complete(CompleteParams {
3811 reference: CompletionReference::prompt("test-prompt"),
3812 argument: CompletionArgument::new("query", "al"),
3813 context: None,
3814 meta: None,
3815 }),
3816 extensions: Extensions::new(),
3817 };
3818 let resp = router
3819 .clone()
3820 .ready()
3821 .await
3822 .unwrap()
3823 .call(complete_req)
3824 .await
3825 .unwrap();
3826
3827 match resp.inner {
3828 Ok(McpResponse::Complete(result)) => {
3829 assert!(result.completion.values.is_empty());
3830 }
3831 _ => panic!("Expected Complete response"),
3832 }
3833 }
3834
3835 #[tokio::test]
3836 async fn test_tool_filter_list() {
3837 use crate::filter::CapabilityFilter;
3838 use crate::tool::Tool;
3839
3840 let public_tool = ToolBuilder::new("public")
3841 .description("Public tool")
3842 .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
3843 .build();
3844
3845 let admin_tool = ToolBuilder::new("admin")
3846 .description("Admin tool")
3847 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3848 .build();
3849
3850 let mut router = McpRouter::new()
3851 .tool(public_tool)
3852 .tool(admin_tool)
3853 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
3854
3855 init_router(&mut router).await;
3857
3858 let req = RouterRequest {
3859 id: RequestId::Number(1),
3860 inner: McpRequest::ListTools(ListToolsParams::default()),
3861 extensions: Extensions::new(),
3862 };
3863
3864 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3865
3866 match resp.inner {
3867 Ok(McpResponse::ListTools(result)) => {
3868 assert_eq!(result.tools.len(), 1);
3870 assert_eq!(result.tools[0].name, "public");
3871 }
3872 _ => panic!("Expected ListTools response"),
3873 }
3874 }
3875
3876 #[tokio::test]
3877 async fn test_tool_filter_call_denied() {
3878 use crate::filter::CapabilityFilter;
3879 use crate::tool::Tool;
3880
3881 let admin_tool = ToolBuilder::new("admin")
3882 .description("Admin tool")
3883 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3884 .build();
3885
3886 let mut router = McpRouter::new()
3887 .tool(admin_tool)
3888 .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); init_router(&mut router).await;
3892
3893 let req = RouterRequest {
3894 id: RequestId::Number(1),
3895 inner: McpRequest::CallTool(CallToolParams {
3896 name: "admin".to_string(),
3897 arguments: serde_json::json!({"a": 1, "b": 2}),
3898 meta: None,
3899 task: None,
3900 }),
3901 extensions: Extensions::new(),
3902 };
3903
3904 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3905
3906 match resp.inner {
3908 Err(e) => {
3909 assert_eq!(e.code, -32601); }
3911 _ => panic!("Expected JsonRpc error"),
3912 }
3913 }
3914
3915 #[tokio::test]
3916 async fn test_tool_filter_call_allowed() {
3917 use crate::filter::CapabilityFilter;
3918 use crate::tool::Tool;
3919
3920 let public_tool = ToolBuilder::new("public")
3921 .description("Public tool")
3922 .handler(|input: AddInput| async move {
3923 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3924 })
3925 .build();
3926
3927 let mut router = McpRouter::new()
3928 .tool(public_tool)
3929 .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); init_router(&mut router).await;
3933
3934 let req = RouterRequest {
3935 id: RequestId::Number(1),
3936 inner: McpRequest::CallTool(CallToolParams {
3937 name: "public".to_string(),
3938 arguments: serde_json::json!({"a": 1, "b": 2}),
3939 meta: None,
3940 task: None,
3941 }),
3942 extensions: Extensions::new(),
3943 };
3944
3945 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3946
3947 match resp.inner {
3948 Ok(McpResponse::CallTool(result)) => {
3949 assert!(!result.is_error);
3950 }
3951 _ => panic!("Expected CallTool response"),
3952 }
3953 }
3954
3955 #[tokio::test]
3956 async fn test_tool_filter_custom_denial() {
3957 use crate::filter::{CapabilityFilter, DenialBehavior};
3958 use crate::tool::Tool;
3959
3960 let admin_tool = ToolBuilder::new("admin")
3961 .description("Admin tool")
3962 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3963 .build();
3964
3965 let mut router = McpRouter::new().tool(admin_tool).tool_filter(
3966 CapabilityFilter::new(|_, _: &Tool| false)
3967 .denial_behavior(DenialBehavior::Unauthorized),
3968 );
3969
3970 init_router(&mut router).await;
3972
3973 let req = RouterRequest {
3974 id: RequestId::Number(1),
3975 inner: McpRequest::CallTool(CallToolParams {
3976 name: "admin".to_string(),
3977 arguments: serde_json::json!({"a": 1, "b": 2}),
3978 meta: None,
3979 task: None,
3980 }),
3981 extensions: Extensions::new(),
3982 };
3983
3984 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3985
3986 match resp.inner {
3988 Err(e) => {
3989 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
3991 }
3992 _ => panic!("Expected JsonRpc error"),
3993 }
3994 }
3995
3996 #[tokio::test]
3997 async fn test_resource_filter_list() {
3998 use crate::filter::CapabilityFilter;
3999 use crate::resource::{Resource, ResourceBuilder};
4000
4001 let public_resource = ResourceBuilder::new("file:///public.txt")
4002 .name("Public File")
4003 .text("public content");
4004
4005 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4006 .name("Secret File")
4007 .text("secret content");
4008
4009 let mut router = McpRouter::new()
4010 .resource(public_resource)
4011 .resource(secret_resource)
4012 .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
4013 !r.name.contains("Secret")
4014 }));
4015
4016 init_router(&mut router).await;
4018
4019 let req = RouterRequest {
4020 id: RequestId::Number(1),
4021 inner: McpRequest::ListResources(ListResourcesParams::default()),
4022 extensions: Extensions::new(),
4023 };
4024
4025 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4026
4027 match resp.inner {
4028 Ok(McpResponse::ListResources(result)) => {
4029 assert_eq!(result.resources.len(), 1);
4031 assert_eq!(result.resources[0].name, "Public File");
4032 }
4033 _ => panic!("Expected ListResources response"),
4034 }
4035 }
4036
4037 #[tokio::test]
4038 async fn test_resource_filter_read_denied() {
4039 use crate::filter::CapabilityFilter;
4040 use crate::resource::{Resource, ResourceBuilder};
4041
4042 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4043 .name("Secret File")
4044 .text("secret content");
4045
4046 let mut router = McpRouter::new()
4047 .resource(secret_resource)
4048 .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); init_router(&mut router).await;
4052
4053 let req = RouterRequest {
4054 id: RequestId::Number(1),
4055 inner: McpRequest::ReadResource(ReadResourceParams {
4056 uri: "file:///secret.txt".to_string(),
4057 meta: None,
4058 }),
4059 extensions: Extensions::new(),
4060 };
4061
4062 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4063
4064 match resp.inner {
4066 Err(e) => {
4067 assert_eq!(e.code, -32601); }
4069 _ => panic!("Expected JsonRpc error"),
4070 }
4071 }
4072
4073 #[tokio::test]
4074 async fn test_resource_filter_read_allowed() {
4075 use crate::filter::CapabilityFilter;
4076 use crate::resource::{Resource, ResourceBuilder};
4077
4078 let public_resource = ResourceBuilder::new("file:///public.txt")
4079 .name("Public File")
4080 .text("public content");
4081
4082 let mut router = McpRouter::new()
4083 .resource(public_resource)
4084 .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); init_router(&mut router).await;
4088
4089 let req = RouterRequest {
4090 id: RequestId::Number(1),
4091 inner: McpRequest::ReadResource(ReadResourceParams {
4092 uri: "file:///public.txt".to_string(),
4093 meta: None,
4094 }),
4095 extensions: Extensions::new(),
4096 };
4097
4098 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4099
4100 match resp.inner {
4101 Ok(McpResponse::ReadResource(result)) => {
4102 assert_eq!(result.contents.len(), 1);
4103 assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
4104 }
4105 _ => panic!("Expected ReadResource response"),
4106 }
4107 }
4108
4109 #[tokio::test]
4110 async fn test_resource_filter_custom_denial() {
4111 use crate::filter::{CapabilityFilter, DenialBehavior};
4112 use crate::resource::{Resource, ResourceBuilder};
4113
4114 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4115 .name("Secret File")
4116 .text("secret content");
4117
4118 let mut router = McpRouter::new().resource(secret_resource).resource_filter(
4119 CapabilityFilter::new(|_, _: &Resource| false)
4120 .denial_behavior(DenialBehavior::Unauthorized),
4121 );
4122
4123 init_router(&mut router).await;
4125
4126 let req = RouterRequest {
4127 id: RequestId::Number(1),
4128 inner: McpRequest::ReadResource(ReadResourceParams {
4129 uri: "file:///secret.txt".to_string(),
4130 meta: None,
4131 }),
4132 extensions: Extensions::new(),
4133 };
4134
4135 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4136
4137 match resp.inner {
4139 Err(e) => {
4140 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4142 }
4143 _ => panic!("Expected JsonRpc error"),
4144 }
4145 }
4146
4147 #[tokio::test]
4148 async fn test_prompt_filter_list() {
4149 use crate::filter::CapabilityFilter;
4150 use crate::prompt::{Prompt, PromptBuilder};
4151
4152 let public_prompt = PromptBuilder::new("greeting")
4153 .description("A greeting")
4154 .user_message("Hello!");
4155
4156 let admin_prompt = PromptBuilder::new("system_debug")
4157 .description("Admin prompt")
4158 .user_message("Debug");
4159
4160 let mut router = McpRouter::new()
4161 .prompt(public_prompt)
4162 .prompt(admin_prompt)
4163 .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
4164 !p.name.contains("system")
4165 }));
4166
4167 init_router(&mut router).await;
4169
4170 let req = RouterRequest {
4171 id: RequestId::Number(1),
4172 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4173 extensions: Extensions::new(),
4174 };
4175
4176 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4177
4178 match resp.inner {
4179 Ok(McpResponse::ListPrompts(result)) => {
4180 assert_eq!(result.prompts.len(), 1);
4182 assert_eq!(result.prompts[0].name, "greeting");
4183 }
4184 _ => panic!("Expected ListPrompts response"),
4185 }
4186 }
4187
4188 #[tokio::test]
4189 async fn test_prompt_filter_get_denied() {
4190 use crate::filter::CapabilityFilter;
4191 use crate::prompt::{Prompt, PromptBuilder};
4192 use std::collections::HashMap;
4193
4194 let admin_prompt = PromptBuilder::new("system_debug")
4195 .description("Admin prompt")
4196 .user_message("Debug");
4197
4198 let mut router = McpRouter::new()
4199 .prompt(admin_prompt)
4200 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); init_router(&mut router).await;
4204
4205 let req = RouterRequest {
4206 id: RequestId::Number(1),
4207 inner: McpRequest::GetPrompt(GetPromptParams {
4208 name: "system_debug".to_string(),
4209 arguments: HashMap::new(),
4210 meta: None,
4211 }),
4212 extensions: Extensions::new(),
4213 };
4214
4215 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4216
4217 match resp.inner {
4219 Err(e) => {
4220 assert_eq!(e.code, -32601); }
4222 _ => panic!("Expected JsonRpc error"),
4223 }
4224 }
4225
4226 #[tokio::test]
4227 async fn test_prompt_filter_get_allowed() {
4228 use crate::filter::CapabilityFilter;
4229 use crate::prompt::{Prompt, PromptBuilder};
4230 use std::collections::HashMap;
4231
4232 let public_prompt = PromptBuilder::new("greeting")
4233 .description("A greeting")
4234 .user_message("Hello!");
4235
4236 let mut router = McpRouter::new()
4237 .prompt(public_prompt)
4238 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); init_router(&mut router).await;
4242
4243 let req = RouterRequest {
4244 id: RequestId::Number(1),
4245 inner: McpRequest::GetPrompt(GetPromptParams {
4246 name: "greeting".to_string(),
4247 arguments: HashMap::new(),
4248 meta: None,
4249 }),
4250 extensions: Extensions::new(),
4251 };
4252
4253 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4254
4255 match resp.inner {
4256 Ok(McpResponse::GetPrompt(result)) => {
4257 assert_eq!(result.messages.len(), 1);
4258 }
4259 _ => panic!("Expected GetPrompt response"),
4260 }
4261 }
4262
4263 #[tokio::test]
4264 async fn test_prompt_filter_custom_denial() {
4265 use crate::filter::{CapabilityFilter, DenialBehavior};
4266 use crate::prompt::{Prompt, PromptBuilder};
4267 use std::collections::HashMap;
4268
4269 let admin_prompt = PromptBuilder::new("system_debug")
4270 .description("Admin prompt")
4271 .user_message("Debug");
4272
4273 let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
4274 CapabilityFilter::new(|_, _: &Prompt| false)
4275 .denial_behavior(DenialBehavior::Unauthorized),
4276 );
4277
4278 init_router(&mut router).await;
4280
4281 let req = RouterRequest {
4282 id: RequestId::Number(1),
4283 inner: McpRequest::GetPrompt(GetPromptParams {
4284 name: "system_debug".to_string(),
4285 arguments: HashMap::new(),
4286 meta: None,
4287 }),
4288 extensions: Extensions::new(),
4289 };
4290
4291 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4292
4293 match resp.inner {
4295 Err(e) => {
4296 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4298 }
4299 _ => panic!("Expected JsonRpc error"),
4300 }
4301 }
4302
4303 #[derive(Debug, Deserialize, JsonSchema)]
4308 struct StringInput {
4309 value: String,
4310 }
4311
4312 #[tokio::test]
4313 async fn test_router_merge_tools() {
4314 let tool_a = ToolBuilder::new("tool_a")
4316 .description("Tool A")
4317 .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
4318 .build();
4319
4320 let router_a = McpRouter::new().tool(tool_a);
4321
4322 let tool_b = ToolBuilder::new("tool_b")
4324 .description("Tool B")
4325 .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
4326 .build();
4327 let tool_c = ToolBuilder::new("tool_c")
4328 .description("Tool C")
4329 .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
4330 .build();
4331
4332 let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
4333
4334 let mut merged = McpRouter::new()
4336 .server_info("merged", "1.0")
4337 .merge(router_a)
4338 .merge(router_b);
4339
4340 init_router(&mut merged).await;
4341
4342 let req = RouterRequest {
4344 id: RequestId::Number(1),
4345 inner: McpRequest::ListTools(ListToolsParams::default()),
4346 extensions: Extensions::new(),
4347 };
4348
4349 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4350
4351 match resp.inner {
4352 Ok(McpResponse::ListTools(result)) => {
4353 assert_eq!(result.tools.len(), 3);
4354 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4355 assert!(names.contains(&"tool_a"));
4356 assert!(names.contains(&"tool_b"));
4357 assert!(names.contains(&"tool_c"));
4358 }
4359 _ => panic!("Expected ListTools response"),
4360 }
4361 }
4362
4363 #[tokio::test]
4364 async fn test_router_merge_overwrites_duplicates() {
4365 let tool_v1 = ToolBuilder::new("shared")
4367 .description("Version 1")
4368 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
4369 .build();
4370
4371 let router_a = McpRouter::new().tool(tool_v1);
4372
4373 let tool_v2 = ToolBuilder::new("shared")
4375 .description("Version 2")
4376 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
4377 .build();
4378
4379 let router_b = McpRouter::new().tool(tool_v2);
4380
4381 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4383
4384 init_router(&mut merged).await;
4385
4386 let req = RouterRequest {
4387 id: RequestId::Number(1),
4388 inner: McpRequest::ListTools(ListToolsParams::default()),
4389 extensions: Extensions::new(),
4390 };
4391
4392 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4393
4394 match resp.inner {
4395 Ok(McpResponse::ListTools(result)) => {
4396 assert_eq!(result.tools.len(), 1);
4397 assert_eq!(result.tools[0].name, "shared");
4398 assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
4399 }
4400 _ => panic!("Expected ListTools response"),
4401 }
4402 }
4403
4404 #[tokio::test]
4405 async fn test_router_merge_resources() {
4406 use crate::resource::ResourceBuilder;
4407
4408 let router_a = McpRouter::new().resource(
4410 ResourceBuilder::new("file:///a.txt")
4411 .name("File A")
4412 .text("content a"),
4413 );
4414
4415 let router_b = McpRouter::new().resource(
4416 ResourceBuilder::new("file:///b.txt")
4417 .name("File B")
4418 .text("content b"),
4419 );
4420
4421 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4422
4423 init_router(&mut merged).await;
4424
4425 let req = RouterRequest {
4426 id: RequestId::Number(1),
4427 inner: McpRequest::ListResources(ListResourcesParams::default()),
4428 extensions: Extensions::new(),
4429 };
4430
4431 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4432
4433 match resp.inner {
4434 Ok(McpResponse::ListResources(result)) => {
4435 assert_eq!(result.resources.len(), 2);
4436 let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4437 assert!(uris.contains(&"file:///a.txt"));
4438 assert!(uris.contains(&"file:///b.txt"));
4439 }
4440 _ => panic!("Expected ListResources response"),
4441 }
4442 }
4443
4444 #[tokio::test]
4445 async fn test_router_merge_prompts() {
4446 use crate::prompt::PromptBuilder;
4447
4448 let router_a =
4449 McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4450
4451 let router_b =
4452 McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4453
4454 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4455
4456 init_router(&mut merged).await;
4457
4458 let req = RouterRequest {
4459 id: RequestId::Number(1),
4460 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4461 extensions: Extensions::new(),
4462 };
4463
4464 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4465
4466 match resp.inner {
4467 Ok(McpResponse::ListPrompts(result)) => {
4468 assert_eq!(result.prompts.len(), 2);
4469 let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4470 assert!(names.contains(&"prompt_a"));
4471 assert!(names.contains(&"prompt_b"));
4472 }
4473 _ => panic!("Expected ListPrompts response"),
4474 }
4475 }
4476
4477 #[tokio::test]
4478 async fn test_router_nest_prefixes_tools() {
4479 let tool_query = ToolBuilder::new("query")
4481 .description("Query the database")
4482 .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4483 .build();
4484 let tool_insert = ToolBuilder::new("insert")
4485 .description("Insert into database")
4486 .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4487 .build();
4488
4489 let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4490
4491 let mut router = McpRouter::new()
4493 .server_info("nested", "1.0")
4494 .nest("db", db_router);
4495
4496 init_router(&mut router).await;
4497
4498 let req = RouterRequest {
4499 id: RequestId::Number(1),
4500 inner: McpRequest::ListTools(ListToolsParams::default()),
4501 extensions: Extensions::new(),
4502 };
4503
4504 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4505
4506 match resp.inner {
4507 Ok(McpResponse::ListTools(result)) => {
4508 assert_eq!(result.tools.len(), 2);
4509 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4510 assert!(names.contains(&"db.query"));
4511 assert!(names.contains(&"db.insert"));
4512 }
4513 _ => panic!("Expected ListTools response"),
4514 }
4515 }
4516
4517 #[tokio::test]
4518 async fn test_router_nest_call_prefixed_tool() {
4519 let tool = ToolBuilder::new("echo")
4520 .description("Echo input")
4521 .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4522 .build();
4523
4524 let nested_router = McpRouter::new().tool(tool);
4525
4526 let mut router = McpRouter::new().nest("api", nested_router);
4527
4528 init_router(&mut router).await;
4529
4530 let req = RouterRequest {
4532 id: RequestId::Number(1),
4533 inner: McpRequest::CallTool(CallToolParams {
4534 name: "api.echo".to_string(),
4535 arguments: serde_json::json!({"value": "hello world"}),
4536 meta: None,
4537 task: None,
4538 }),
4539 extensions: Extensions::new(),
4540 };
4541
4542 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4543
4544 match resp.inner {
4545 Ok(McpResponse::CallTool(result)) => {
4546 assert!(!result.is_error);
4547 match &result.content[0] {
4548 Content::Text { text, .. } => assert_eq!(text, "hello world"),
4549 _ => panic!("Expected text content"),
4550 }
4551 }
4552 _ => panic!("Expected CallTool response"),
4553 }
4554 }
4555
4556 #[tokio::test]
4557 async fn test_router_multiple_nests() {
4558 let db_tool = ToolBuilder::new("query")
4559 .description("Database query")
4560 .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4561 .build();
4562
4563 let api_tool = ToolBuilder::new("fetch")
4564 .description("API fetch")
4565 .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4566 .build();
4567
4568 let db_router = McpRouter::new().tool(db_tool);
4569 let api_router = McpRouter::new().tool(api_tool);
4570
4571 let mut router = McpRouter::new()
4572 .nest("db", db_router)
4573 .nest("api", api_router);
4574
4575 init_router(&mut router).await;
4576
4577 let req = RouterRequest {
4578 id: RequestId::Number(1),
4579 inner: McpRequest::ListTools(ListToolsParams::default()),
4580 extensions: Extensions::new(),
4581 };
4582
4583 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4584
4585 match resp.inner {
4586 Ok(McpResponse::ListTools(result)) => {
4587 assert_eq!(result.tools.len(), 2);
4588 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4589 assert!(names.contains(&"db.query"));
4590 assert!(names.contains(&"api.fetch"));
4591 }
4592 _ => panic!("Expected ListTools response"),
4593 }
4594 }
4595
4596 #[tokio::test]
4597 async fn test_router_merge_and_nest_combined() {
4598 let tool_a = ToolBuilder::new("local")
4600 .description("Local tool")
4601 .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
4602 .build();
4603
4604 let nested_tool = ToolBuilder::new("remote")
4605 .description("Remote tool")
4606 .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
4607 .build();
4608
4609 let nested_router = McpRouter::new().tool(nested_tool);
4610
4611 let mut router = McpRouter::new()
4612 .tool(tool_a)
4613 .nest("external", nested_router);
4614
4615 init_router(&mut router).await;
4616
4617 let req = RouterRequest {
4618 id: RequestId::Number(1),
4619 inner: McpRequest::ListTools(ListToolsParams::default()),
4620 extensions: Extensions::new(),
4621 };
4622
4623 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4624
4625 match resp.inner {
4626 Ok(McpResponse::ListTools(result)) => {
4627 assert_eq!(result.tools.len(), 2);
4628 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4629 assert!(names.contains(&"local"));
4630 assert!(names.contains(&"external.remote"));
4631 }
4632 _ => panic!("Expected ListTools response"),
4633 }
4634 }
4635
4636 #[tokio::test]
4637 async fn test_router_merge_preserves_server_info() {
4638 let child_router = McpRouter::new()
4639 .server_info("child", "2.0")
4640 .instructions("Child instructions");
4641
4642 let mut router = McpRouter::new()
4643 .server_info("parent", "1.0")
4644 .instructions("Parent instructions")
4645 .merge(child_router);
4646
4647 init_router(&mut router).await;
4648
4649 let init_req = RouterRequest {
4651 id: RequestId::Number(99),
4652 inner: McpRequest::Initialize(InitializeParams {
4653 protocol_version: "2025-11-25".to_string(),
4654 capabilities: ClientCapabilities::default(),
4655 client_info: Implementation {
4656 name: "test".to_string(),
4657 version: "1.0".to_string(),
4658 ..Default::default()
4659 },
4660 meta: None,
4661 }),
4662 extensions: Extensions::new(),
4663 };
4664
4665 let child_router2 = McpRouter::new().server_info("child", "2.0");
4667 let mut fresh_router = McpRouter::new()
4668 .server_info("parent", "1.0")
4669 .merge(child_router2);
4670
4671 let resp = fresh_router
4672 .ready()
4673 .await
4674 .unwrap()
4675 .call(init_req)
4676 .await
4677 .unwrap();
4678
4679 match resp.inner {
4680 Ok(McpResponse::Initialize(result)) => {
4681 assert_eq!(result.server_info.name, "parent");
4682 assert_eq!(result.server_info.version, "1.0");
4683 }
4684 _ => panic!("Expected Initialize response"),
4685 }
4686 }
4687
4688 #[tokio::test]
4693 async fn test_auto_instructions_tools_only() {
4694 let tool_a = ToolBuilder::new("alpha")
4695 .description("Alpha tool")
4696 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4697 .build();
4698 let tool_b = ToolBuilder::new("beta")
4699 .description("Beta tool")
4700 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4701 .build();
4702
4703 let mut router = McpRouter::new()
4704 .auto_instructions()
4705 .tool(tool_a)
4706 .tool(tool_b);
4707
4708 let resp = send_initialize(&mut router).await;
4709 let instructions = resp.instructions.expect("should have instructions");
4710
4711 assert!(instructions.contains("## Tools"));
4712 assert!(instructions.contains("- **alpha**: Alpha tool"));
4713 assert!(instructions.contains("- **beta**: Beta tool"));
4714 assert!(!instructions.contains("## Resources"));
4716 assert!(!instructions.contains("## Prompts"));
4717 }
4718
4719 #[tokio::test]
4720 async fn test_auto_instructions_with_annotations() {
4721 let read_only_tool = ToolBuilder::new("query")
4722 .description("Run a query")
4723 .read_only()
4724 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4725 .build();
4726 let destructive_tool = ToolBuilder::new("delete")
4727 .description("Delete a record")
4728 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4729 .build();
4730 let idempotent_tool = ToolBuilder::new("upsert")
4731 .description("Upsert a record")
4732 .non_destructive()
4733 .idempotent()
4734 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4735 .build();
4736
4737 let mut router = McpRouter::new()
4738 .auto_instructions()
4739 .tool(read_only_tool)
4740 .tool(destructive_tool)
4741 .tool(idempotent_tool);
4742
4743 let resp = send_initialize(&mut router).await;
4744 let instructions = resp.instructions.unwrap();
4745
4746 assert!(instructions.contains("- **query**: Run a query [read-only]"));
4747 assert!(instructions.contains("- **delete**: Delete a record\n"));
4749 assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
4750 }
4751
4752 #[tokio::test]
4753 async fn test_auto_instructions_with_resources() {
4754 use crate::resource::ResourceBuilder;
4755
4756 let resource = ResourceBuilder::new("file:///schema.sql")
4757 .name("Schema")
4758 .description("Database schema")
4759 .text("CREATE TABLE ...");
4760
4761 let mut router = McpRouter::new().auto_instructions().resource(resource);
4762
4763 let resp = send_initialize(&mut router).await;
4764 let instructions = resp.instructions.unwrap();
4765
4766 assert!(instructions.contains("## Resources"));
4767 assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
4768 assert!(!instructions.contains("## Tools"));
4769 }
4770
4771 #[tokio::test]
4772 async fn test_auto_instructions_with_resource_templates() {
4773 use crate::resource::ResourceTemplateBuilder;
4774
4775 let template = ResourceTemplateBuilder::new("file:///{path}")
4776 .name("File")
4777 .description("Read a file by path")
4778 .handler(
4779 |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
4780 Ok(crate::ReadResourceResult::text("content", "text/plain"))
4781 },
4782 );
4783
4784 let mut router = McpRouter::new()
4785 .auto_instructions()
4786 .resource_template(template);
4787
4788 let resp = send_initialize(&mut router).await;
4789 let instructions = resp.instructions.unwrap();
4790
4791 assert!(instructions.contains("## Resources"));
4792 assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
4793 }
4794
4795 #[tokio::test]
4796 async fn test_auto_instructions_with_prompts() {
4797 use crate::prompt::PromptBuilder;
4798
4799 let prompt = PromptBuilder::new("write_query")
4800 .description("Help write a SQL query")
4801 .user_message("Write a query for: {task}");
4802
4803 let mut router = McpRouter::new().auto_instructions().prompt(prompt);
4804
4805 let resp = send_initialize(&mut router).await;
4806 let instructions = resp.instructions.unwrap();
4807
4808 assert!(instructions.contains("## Prompts"));
4809 assert!(instructions.contains("- **write_query**: Help write a SQL query"));
4810 assert!(!instructions.contains("## Tools"));
4811 }
4812
4813 #[tokio::test]
4814 async fn test_auto_instructions_all_sections() {
4815 use crate::prompt::PromptBuilder;
4816 use crate::resource::ResourceBuilder;
4817
4818 let tool = ToolBuilder::new("query")
4819 .description("Execute SQL")
4820 .read_only()
4821 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4822 .build();
4823 let resource = ResourceBuilder::new("db://schema")
4824 .name("Schema")
4825 .description("Full database schema")
4826 .text("schema");
4827 let prompt = PromptBuilder::new("write_query")
4828 .description("Help write a SQL query")
4829 .user_message("Write a query");
4830
4831 let mut router = McpRouter::new()
4832 .auto_instructions()
4833 .tool(tool)
4834 .resource(resource)
4835 .prompt(prompt);
4836
4837 let resp = send_initialize(&mut router).await;
4838 let instructions = resp.instructions.unwrap();
4839
4840 assert!(instructions.contains("## Tools"));
4842 assert!(instructions.contains("## Resources"));
4843 assert!(instructions.contains("## Prompts"));
4844
4845 let tools_pos = instructions.find("## Tools").unwrap();
4847 let resources_pos = instructions.find("## Resources").unwrap();
4848 let prompts_pos = instructions.find("## Prompts").unwrap();
4849 assert!(tools_pos < resources_pos);
4850 assert!(resources_pos < prompts_pos);
4851 }
4852
4853 #[tokio::test]
4854 async fn test_auto_instructions_with_prefix_and_suffix() {
4855 let tool = ToolBuilder::new("echo")
4856 .description("Echo input")
4857 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4858 .build();
4859
4860 let mut router = McpRouter::new()
4861 .auto_instructions_with(
4862 Some("This server provides echo capabilities."),
4863 Some("Contact admin@example.com for support."),
4864 )
4865 .tool(tool);
4866
4867 let resp = send_initialize(&mut router).await;
4868 let instructions = resp.instructions.unwrap();
4869
4870 assert!(instructions.starts_with("This server provides echo capabilities."));
4871 assert!(instructions.ends_with("Contact admin@example.com for support."));
4872 assert!(instructions.contains("## Tools"));
4873 assert!(instructions.contains("- **echo**: Echo input"));
4874 }
4875
4876 #[tokio::test]
4877 async fn test_auto_instructions_prefix_only() {
4878 let tool = ToolBuilder::new("echo")
4879 .description("Echo input")
4880 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4881 .build();
4882
4883 let mut router = McpRouter::new()
4884 .auto_instructions_with(Some("My server intro."), None::<String>)
4885 .tool(tool);
4886
4887 let resp = send_initialize(&mut router).await;
4888 let instructions = resp.instructions.unwrap();
4889
4890 assert!(instructions.starts_with("My server intro."));
4891 assert!(instructions.contains("- **echo**: Echo input"));
4892 }
4893
4894 #[tokio::test]
4895 async fn test_auto_instructions_empty_router() {
4896 let mut router = McpRouter::new().auto_instructions();
4897
4898 let resp = send_initialize(&mut router).await;
4899 let instructions = resp.instructions.expect("should have instructions");
4900
4901 assert!(!instructions.contains("## Tools"));
4903 assert!(!instructions.contains("## Resources"));
4904 assert!(!instructions.contains("## Prompts"));
4905 assert!(instructions.is_empty());
4906 }
4907
4908 #[tokio::test]
4909 async fn test_auto_instructions_overrides_manual() {
4910 let tool = ToolBuilder::new("echo")
4911 .description("Echo input")
4912 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4913 .build();
4914
4915 let mut router = McpRouter::new()
4916 .instructions("This will be overridden")
4917 .auto_instructions()
4918 .tool(tool);
4919
4920 let resp = send_initialize(&mut router).await;
4921 let instructions = resp.instructions.unwrap();
4922
4923 assert!(!instructions.contains("This will be overridden"));
4924 assert!(instructions.contains("- **echo**: Echo input"));
4925 }
4926
4927 #[tokio::test]
4928 async fn test_no_auto_instructions_returns_manual() {
4929 let tool = ToolBuilder::new("echo")
4930 .description("Echo input")
4931 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4932 .build();
4933
4934 let mut router = McpRouter::new()
4935 .instructions("Manual instructions here")
4936 .tool(tool);
4937
4938 let resp = send_initialize(&mut router).await;
4939 let instructions = resp.instructions.unwrap();
4940
4941 assert_eq!(instructions, "Manual instructions here");
4942 }
4943
4944 #[tokio::test]
4945 async fn test_auto_instructions_no_description_fallback() {
4946 let tool = ToolBuilder::new("mystery")
4947 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4948 .build();
4949
4950 let mut router = McpRouter::new().auto_instructions().tool(tool);
4951
4952 let resp = send_initialize(&mut router).await;
4953 let instructions = resp.instructions.unwrap();
4954
4955 assert!(instructions.contains("- **mystery**: No description"));
4956 }
4957
4958 #[tokio::test]
4959 async fn test_auto_instructions_sorted_alphabetically() {
4960 let tool_z = ToolBuilder::new("zebra")
4961 .description("Z tool")
4962 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4963 .build();
4964 let tool_a = ToolBuilder::new("alpha")
4965 .description("A tool")
4966 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4967 .build();
4968 let tool_m = ToolBuilder::new("middle")
4969 .description("M tool")
4970 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4971 .build();
4972
4973 let mut router = McpRouter::new()
4974 .auto_instructions()
4975 .tool(tool_z)
4976 .tool(tool_a)
4977 .tool(tool_m);
4978
4979 let resp = send_initialize(&mut router).await;
4980 let instructions = resp.instructions.unwrap();
4981
4982 let alpha_pos = instructions.find("**alpha**").unwrap();
4983 let middle_pos = instructions.find("**middle**").unwrap();
4984 let zebra_pos = instructions.find("**zebra**").unwrap();
4985 assert!(alpha_pos < middle_pos);
4986 assert!(middle_pos < zebra_pos);
4987 }
4988
4989 #[tokio::test]
4990 async fn test_auto_instructions_read_only_and_idempotent_tags() {
4991 let tool = ToolBuilder::new("safe_update")
4992 .description("Safe update operation")
4993 .idempotent()
4994 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4995 .build();
4996
4997 let mut router = McpRouter::new().auto_instructions().tool(tool);
4998
4999 let resp = send_initialize(&mut router).await;
5000 let instructions = resp.instructions.unwrap();
5001
5002 assert!(
5003 instructions.contains("[idempotent]"),
5004 "got: {}",
5005 instructions
5006 );
5007 }
5008
5009 #[tokio::test]
5010 async fn test_auto_instructions_lazy_generation() {
5011 let mut router = McpRouter::new().auto_instructions();
5014
5015 let tool = ToolBuilder::new("late_tool")
5016 .description("Added after auto_instructions")
5017 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5018 .build();
5019
5020 router = router.tool(tool);
5021
5022 let resp = send_initialize(&mut router).await;
5023 let instructions = resp.instructions.unwrap();
5024
5025 assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
5026 }
5027
5028 #[tokio::test]
5029 async fn test_auto_instructions_multiple_annotation_tags() {
5030 let tool = ToolBuilder::new("update")
5031 .description("Update a record")
5032 .annotations(ToolAnnotations {
5033 read_only_hint: true,
5034 idempotent_hint: true,
5035 ..Default::default()
5036 })
5037 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5038 .build();
5039
5040 let mut router = McpRouter::new().auto_instructions().tool(tool);
5041
5042 let resp = send_initialize(&mut router).await;
5043 let instructions = resp.instructions.unwrap();
5044
5045 assert!(
5046 instructions.contains("[read-only, idempotent]"),
5047 "got: {}",
5048 instructions
5049 );
5050 }
5051
5052 #[tokio::test]
5053 async fn test_auto_instructions_no_annotations_no_tags() {
5054 let tool = ToolBuilder::new("fetch")
5056 .description("Fetch data")
5057 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5058 .build();
5059
5060 let mut router = McpRouter::new().auto_instructions().tool(tool);
5061
5062 let resp = send_initialize(&mut router).await;
5063 let instructions = resp.instructions.unwrap();
5064
5065 assert!(
5067 !instructions.contains('['),
5068 "should have no tags, got: {}",
5069 instructions
5070 );
5071 assert!(instructions.contains("- **fetch**: Fetch data"));
5072 }
5073
5074 async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
5076 let init_req = RouterRequest {
5077 id: RequestId::Number(0),
5078 inner: McpRequest::Initialize(InitializeParams {
5079 protocol_version: "2025-11-25".to_string(),
5080 capabilities: ClientCapabilities {
5081 roots: None,
5082 sampling: None,
5083 elicitation: None,
5084 tasks: None,
5085 experimental: None,
5086 extensions: None,
5087 },
5088 client_info: Implementation {
5089 name: "test".to_string(),
5090 version: "1.0".to_string(),
5091 ..Default::default()
5092 },
5093 meta: None,
5094 }),
5095 extensions: Extensions::new(),
5096 };
5097 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5098 match resp.inner {
5099 Ok(McpResponse::Initialize(result)) => result,
5100 other => panic!("Expected Initialize response, got {:?}", other),
5101 }
5102 }
5103
5104 #[tokio::test]
5105 async fn test_notify_tools_list_changed() {
5106 let (tx, mut rx) = crate::context::notification_channel(16);
5107
5108 let router = McpRouter::new()
5109 .server_info("test", "1.0")
5110 .with_notification_sender(tx);
5111
5112 assert!(router.notify_tools_list_changed());
5113
5114 let notification = rx.recv().await.unwrap();
5115 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5116 }
5117
5118 #[tokio::test]
5119 async fn test_notify_prompts_list_changed() {
5120 let (tx, mut rx) = crate::context::notification_channel(16);
5121
5122 let router = McpRouter::new()
5123 .server_info("test", "1.0")
5124 .with_notification_sender(tx);
5125
5126 assert!(router.notify_prompts_list_changed());
5127
5128 let notification = rx.recv().await.unwrap();
5129 assert!(matches!(
5130 notification,
5131 ServerNotification::PromptsListChanged
5132 ));
5133 }
5134
5135 #[tokio::test]
5136 async fn test_notify_without_sender_returns_false() {
5137 let router = McpRouter::new().server_info("test", "1.0");
5138
5139 assert!(!router.notify_tools_list_changed());
5140 assert!(!router.notify_prompts_list_changed());
5141 assert!(!router.notify_resources_list_changed());
5142 }
5143
5144 #[tokio::test]
5145 async fn test_list_changed_capabilities_with_notification_sender() {
5146 let (tx, _rx) = crate::context::notification_channel(16);
5147 let tool = ToolBuilder::new("test")
5148 .description("test")
5149 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5150 .build();
5151
5152 let mut router = McpRouter::new()
5153 .server_info("test", "1.0")
5154 .tool(tool)
5155 .with_notification_sender(tx);
5156
5157 init_router(&mut router).await;
5158
5159 let caps = router.capabilities();
5160 let tools_cap = caps.tools.expect("tools capability should be present");
5161 assert!(
5162 tools_cap.list_changed,
5163 "tools.listChanged should be true when notification sender is configured"
5164 );
5165 }
5166
5167 #[tokio::test]
5168 async fn test_list_changed_capabilities_without_notification_sender() {
5169 let tool = ToolBuilder::new("test")
5170 .description("test")
5171 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5172 .build();
5173
5174 let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
5175
5176 init_router(&mut router).await;
5177
5178 let caps = router.capabilities();
5179 let tools_cap = caps.tools.expect("tools capability should be present");
5180 assert!(
5181 !tools_cap.list_changed,
5182 "tools.listChanged should be false without notification sender"
5183 );
5184 }
5185
5186 #[tokio::test]
5187 async fn test_set_logging_level_filters_messages() {
5188 let (tx, mut rx) = crate::context::notification_channel(16);
5189
5190 let mut router = McpRouter::new()
5191 .server_info("test", "1.0")
5192 .with_notification_sender(tx);
5193
5194 init_router(&mut router).await;
5195
5196 let set_level_req = RouterRequest {
5198 id: RequestId::Number(99),
5199 inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
5200 level: LogLevel::Warning,
5201 meta: None,
5202 }),
5203 extensions: crate::context::Extensions::new(),
5204 };
5205 let resp = router
5206 .ready()
5207 .await
5208 .unwrap()
5209 .call(set_level_req)
5210 .await
5211 .unwrap();
5212 assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
5213
5214 let ctx = router.create_context(RequestId::Number(100), None);
5216
5217 ctx.send_log(LoggingMessageParams::new(
5219 LogLevel::Error,
5220 serde_json::Value::Null,
5221 ));
5222 assert!(
5223 rx.try_recv().is_ok(),
5224 "Error should pass through Warning filter"
5225 );
5226
5227 ctx.send_log(LoggingMessageParams::new(
5229 LogLevel::Info,
5230 serde_json::Value::Null,
5231 ));
5232 assert!(
5233 rx.try_recv().is_err(),
5234 "Info should be filtered at Warning level"
5235 );
5236 }
5237
5238 #[test]
5239 fn test_paginate_no_page_size() {
5240 let items = vec![1, 2, 3, 4, 5];
5241 let (page, cursor) = paginate(items.clone(), None, None).unwrap();
5242 assert_eq!(page, items);
5243 assert!(cursor.is_none());
5244 }
5245
5246 #[test]
5247 fn test_paginate_first_page() {
5248 let items = vec![1, 2, 3, 4, 5];
5249 let (page, cursor) = paginate(items, None, Some(2)).unwrap();
5250 assert_eq!(page, vec![1, 2]);
5251 assert!(cursor.is_some());
5252 }
5253
5254 #[test]
5255 fn test_paginate_middle_page() {
5256 let items = vec![1, 2, 3, 4, 5];
5257 let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
5258 assert_eq!(page1, vec![1, 2]);
5259
5260 let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
5261 assert_eq!(page2, vec![3, 4]);
5262 assert!(cursor2.is_some());
5263 }
5264
5265 #[test]
5266 fn test_paginate_last_page() {
5267 let items = vec![1, 2, 3, 4, 5];
5268 let cursor = encode_cursor(4);
5270 let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
5271 assert_eq!(page, vec![5]);
5272 assert!(next.is_none());
5273 }
5274
5275 #[test]
5276 fn test_paginate_exact_boundary() {
5277 let items = vec![1, 2, 3, 4];
5278 let (page, cursor) = paginate(items, None, Some(4)).unwrap();
5279 assert_eq!(page, vec![1, 2, 3, 4]);
5280 assert!(cursor.is_none());
5281 }
5282
5283 #[test]
5284 fn test_paginate_invalid_cursor() {
5285 let items = vec![1, 2, 3];
5286 let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
5287 assert!(result.is_err());
5288 }
5289
5290 #[test]
5291 fn test_cursor_round_trip() {
5292 let offset = 42;
5293 let encoded = encode_cursor(offset);
5294 let decoded = decode_cursor(&encoded).unwrap();
5295 assert_eq!(decoded, offset);
5296 }
5297
5298 #[tokio::test]
5299 async fn test_list_tools_pagination() {
5300 let tool_a = ToolBuilder::new("alpha")
5301 .description("a")
5302 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5303 .build();
5304 let tool_b = ToolBuilder::new("beta")
5305 .description("b")
5306 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5307 .build();
5308 let tool_c = ToolBuilder::new("gamma")
5309 .description("c")
5310 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5311 .build();
5312
5313 let mut router = McpRouter::new()
5314 .server_info("test", "1.0")
5315 .page_size(2)
5316 .tool(tool_a)
5317 .tool(tool_b)
5318 .tool(tool_c);
5319
5320 init_router(&mut router).await;
5321
5322 let req = RouterRequest {
5324 id: RequestId::Number(1),
5325 inner: McpRequest::ListTools(ListToolsParams {
5326 cursor: None,
5327 meta: None,
5328 }),
5329 extensions: Extensions::new(),
5330 };
5331 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5332 let (tools, next_cursor) = match resp.inner {
5333 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5334 other => panic!("Expected ListTools, got {:?}", other),
5335 };
5336 assert_eq!(tools.len(), 2);
5337 assert_eq!(tools[0].name, "alpha");
5338 assert_eq!(tools[1].name, "beta");
5339 assert!(next_cursor.is_some());
5340
5341 let req = RouterRequest {
5343 id: RequestId::Number(2),
5344 inner: McpRequest::ListTools(ListToolsParams {
5345 cursor: next_cursor,
5346 meta: None,
5347 }),
5348 extensions: Extensions::new(),
5349 };
5350 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5351 let (tools, next_cursor) = match resp.inner {
5352 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5353 other => panic!("Expected ListTools, got {:?}", other),
5354 };
5355 assert_eq!(tools.len(), 1);
5356 assert_eq!(tools[0].name, "gamma");
5357 assert!(next_cursor.is_none());
5358 }
5359
5360 #[tokio::test]
5361 async fn test_list_tools_no_pagination_by_default() {
5362 let tool_a = ToolBuilder::new("alpha")
5363 .description("a")
5364 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5365 .build();
5366 let tool_b = ToolBuilder::new("beta")
5367 .description("b")
5368 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5369 .build();
5370
5371 let mut router = McpRouter::new()
5372 .server_info("test", "1.0")
5373 .tool(tool_a)
5374 .tool(tool_b);
5375
5376 init_router(&mut router).await;
5377
5378 let req = RouterRequest {
5379 id: RequestId::Number(1),
5380 inner: McpRequest::ListTools(ListToolsParams {
5381 cursor: None,
5382 meta: None,
5383 }),
5384 extensions: Extensions::new(),
5385 };
5386 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5387 match resp.inner {
5388 Ok(McpResponse::ListTools(result)) => {
5389 assert_eq!(result.tools.len(), 2);
5390 assert!(result.next_cursor.is_none());
5391 }
5392 other => panic!("Expected ListTools, got {:?}", other),
5393 }
5394 }
5395
5396 #[cfg(feature = "dynamic-tools")]
5401 mod dynamic_tools_tests {
5402 use super::*;
5403
5404 #[tokio::test]
5405 async fn test_dynamic_tools_register_and_list() {
5406 let (router, registry) = McpRouter::new()
5407 .server_info("test", "1.0")
5408 .with_dynamic_tools();
5409
5410 let tool = ToolBuilder::new("dynamic_echo")
5411 .description("Dynamic echo")
5412 .handler(|input: AddInput| async move {
5413 Ok(CallToolResult::text(format!("{}", input.a)))
5414 })
5415 .build();
5416
5417 registry.register(tool);
5418
5419 let mut router = router;
5420 init_router(&mut router).await;
5421
5422 let req = RouterRequest {
5423 id: RequestId::Number(1),
5424 inner: McpRequest::ListTools(ListToolsParams::default()),
5425 extensions: Extensions::new(),
5426 };
5427
5428 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5429 match resp.inner {
5430 Ok(McpResponse::ListTools(result)) => {
5431 assert_eq!(result.tools.len(), 1);
5432 assert_eq!(result.tools[0].name, "dynamic_echo");
5433 }
5434 _ => panic!("Expected ListTools response"),
5435 }
5436 }
5437
5438 #[tokio::test]
5439 async fn test_dynamic_tools_unregister() {
5440 let (router, registry) = McpRouter::new()
5441 .server_info("test", "1.0")
5442 .with_dynamic_tools();
5443
5444 let tool = ToolBuilder::new("temp")
5445 .description("Temporary")
5446 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5447 .build();
5448
5449 registry.register(tool);
5450 assert!(registry.contains("temp"));
5451
5452 let removed = registry.unregister("temp");
5453 assert!(removed);
5454 assert!(!registry.contains("temp"));
5455
5456 assert!(!registry.unregister("temp"));
5458
5459 let mut router = router;
5460 init_router(&mut router).await;
5461
5462 let req = RouterRequest {
5463 id: RequestId::Number(1),
5464 inner: McpRequest::ListTools(ListToolsParams::default()),
5465 extensions: Extensions::new(),
5466 };
5467
5468 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5469 match resp.inner {
5470 Ok(McpResponse::ListTools(result)) => {
5471 assert_eq!(result.tools.len(), 0);
5472 }
5473 _ => panic!("Expected ListTools response"),
5474 }
5475 }
5476
5477 #[tokio::test]
5478 async fn test_dynamic_tools_merged_with_static() {
5479 let static_tool = ToolBuilder::new("static_tool")
5480 .description("Static")
5481 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5482 .build();
5483
5484 let (router, registry) = McpRouter::new()
5485 .server_info("test", "1.0")
5486 .tool(static_tool)
5487 .with_dynamic_tools();
5488
5489 let dynamic_tool = ToolBuilder::new("dynamic_tool")
5490 .description("Dynamic")
5491 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5492 .build();
5493
5494 registry.register(dynamic_tool);
5495
5496 let mut router = router;
5497 init_router(&mut router).await;
5498
5499 let req = RouterRequest {
5500 id: RequestId::Number(1),
5501 inner: McpRequest::ListTools(ListToolsParams::default()),
5502 extensions: Extensions::new(),
5503 };
5504
5505 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5506 match resp.inner {
5507 Ok(McpResponse::ListTools(result)) => {
5508 assert_eq!(result.tools.len(), 2);
5509 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5510 assert!(names.contains(&"static_tool"));
5511 assert!(names.contains(&"dynamic_tool"));
5512 }
5513 _ => panic!("Expected ListTools response"),
5514 }
5515 }
5516
5517 #[tokio::test]
5518 async fn test_static_tools_shadow_dynamic() {
5519 let static_tool = ToolBuilder::new("shared")
5520 .description("Static version")
5521 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5522 .build();
5523
5524 let (router, registry) = McpRouter::new()
5525 .server_info("test", "1.0")
5526 .tool(static_tool)
5527 .with_dynamic_tools();
5528
5529 let dynamic_tool = ToolBuilder::new("shared")
5530 .description("Dynamic version")
5531 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5532 .build();
5533
5534 registry.register(dynamic_tool);
5535
5536 let mut router = router;
5537 init_router(&mut router).await;
5538
5539 let req = RouterRequest {
5541 id: RequestId::Number(1),
5542 inner: McpRequest::ListTools(ListToolsParams::default()),
5543 extensions: Extensions::new(),
5544 };
5545
5546 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5547 match resp.inner {
5548 Ok(McpResponse::ListTools(result)) => {
5549 assert_eq!(result.tools.len(), 1);
5550 assert_eq!(result.tools[0].name, "shared");
5551 assert_eq!(
5552 result.tools[0].description.as_deref(),
5553 Some("Static version")
5554 );
5555 }
5556 _ => panic!("Expected ListTools response"),
5557 }
5558
5559 let req = RouterRequest {
5561 id: RequestId::Number(2),
5562 inner: McpRequest::CallTool(CallToolParams {
5563 name: "shared".to_string(),
5564 arguments: serde_json::json!({"a": 1, "b": 2}),
5565 meta: None,
5566 task: None,
5567 }),
5568 extensions: Extensions::new(),
5569 };
5570
5571 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5572 match resp.inner {
5573 Ok(McpResponse::CallTool(result)) => {
5574 assert!(!result.is_error);
5575 match &result.content[0] {
5576 Content::Text { text, .. } => assert_eq!(text, "static"),
5577 _ => panic!("Expected text content"),
5578 }
5579 }
5580 _ => panic!("Expected CallTool response"),
5581 }
5582 }
5583
5584 #[tokio::test]
5585 async fn test_dynamic_tools_call() {
5586 let (router, registry) = McpRouter::new()
5587 .server_info("test", "1.0")
5588 .with_dynamic_tools();
5589
5590 let tool = ToolBuilder::new("add")
5591 .description("Add two numbers")
5592 .handler(|input: AddInput| async move {
5593 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
5594 })
5595 .build();
5596
5597 registry.register(tool);
5598
5599 let mut router = router;
5600 init_router(&mut router).await;
5601
5602 let req = RouterRequest {
5603 id: RequestId::Number(1),
5604 inner: McpRequest::CallTool(CallToolParams {
5605 name: "add".to_string(),
5606 arguments: serde_json::json!({"a": 3, "b": 4}),
5607 meta: None,
5608 task: None,
5609 }),
5610 extensions: Extensions::new(),
5611 };
5612
5613 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5614 match resp.inner {
5615 Ok(McpResponse::CallTool(result)) => {
5616 assert!(!result.is_error);
5617 match &result.content[0] {
5618 Content::Text { text, .. } => assert_eq!(text, "7"),
5619 _ => panic!("Expected text content"),
5620 }
5621 }
5622 _ => panic!("Expected CallTool response"),
5623 }
5624 }
5625
5626 #[tokio::test]
5627 async fn test_dynamic_tools_notification_on_register() {
5628 let (tx, mut rx) = crate::context::notification_channel(16);
5629 let (router, registry) = McpRouter::new()
5630 .server_info("test", "1.0")
5631 .with_dynamic_tools();
5632 let _router = router.with_notification_sender(tx);
5633
5634 let tool = ToolBuilder::new("notified")
5635 .description("Test")
5636 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5637 .build();
5638
5639 registry.register(tool);
5640
5641 let notification = rx.recv().await.unwrap();
5642 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5643 }
5644
5645 #[tokio::test]
5646 async fn test_dynamic_tools_notification_on_unregister() {
5647 let (tx, mut rx) = crate::context::notification_channel(16);
5648 let (router, registry) = McpRouter::new()
5649 .server_info("test", "1.0")
5650 .with_dynamic_tools();
5651 let _router = router.with_notification_sender(tx);
5652
5653 let tool = ToolBuilder::new("notified")
5654 .description("Test")
5655 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5656 .build();
5657
5658 registry.register(tool);
5659 let _ = rx.recv().await.unwrap();
5661
5662 registry.unregister("notified");
5663 let notification = rx.recv().await.unwrap();
5664 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5665 }
5666
5667 #[tokio::test]
5668 async fn test_dynamic_tools_no_notification_on_empty_unregister() {
5669 let (tx, mut rx) = crate::context::notification_channel(16);
5670 let (router, registry) = McpRouter::new()
5671 .server_info("test", "1.0")
5672 .with_dynamic_tools();
5673 let _router = router.with_notification_sender(tx);
5674
5675 assert!(!registry.unregister("nonexistent"));
5677
5678 assert!(rx.try_recv().is_err());
5680 }
5681
5682 #[tokio::test]
5683 async fn test_dynamic_tools_filter_applies() {
5684 use crate::filter::CapabilityFilter;
5685
5686 let (router, registry) = McpRouter::new()
5687 .server_info("test", "1.0")
5688 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
5689 tool.name != "hidden"
5690 }))
5691 .with_dynamic_tools();
5692
5693 let visible = ToolBuilder::new("visible")
5694 .description("Visible")
5695 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5696 .build();
5697
5698 let hidden = ToolBuilder::new("hidden")
5699 .description("Hidden")
5700 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5701 .build();
5702
5703 registry.register(visible);
5704 registry.register(hidden);
5705
5706 let mut router = router;
5707 init_router(&mut router).await;
5708
5709 let req = RouterRequest {
5711 id: RequestId::Number(1),
5712 inner: McpRequest::ListTools(ListToolsParams::default()),
5713 extensions: Extensions::new(),
5714 };
5715
5716 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5717 match resp.inner {
5718 Ok(McpResponse::ListTools(result)) => {
5719 assert_eq!(result.tools.len(), 1);
5720 assert_eq!(result.tools[0].name, "visible");
5721 }
5722 _ => panic!("Expected ListTools response"),
5723 }
5724
5725 let req = RouterRequest {
5727 id: RequestId::Number(2),
5728 inner: McpRequest::CallTool(CallToolParams {
5729 name: "hidden".to_string(),
5730 arguments: serde_json::json!({"a": 1, "b": 2}),
5731 meta: None,
5732 task: None,
5733 }),
5734 extensions: Extensions::new(),
5735 };
5736
5737 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5738 match resp.inner {
5739 Err(e) => {
5740 assert_eq!(e.code, -32601); }
5742 _ => panic!("Expected JsonRpc error"),
5743 }
5744 }
5745
5746 #[tokio::test]
5747 async fn test_dynamic_tools_capabilities_advertised() {
5748 let (mut router, _registry) = McpRouter::new()
5750 .server_info("test", "1.0")
5751 .with_dynamic_tools();
5752
5753 let init_req = RouterRequest {
5754 id: RequestId::Number(1),
5755 inner: McpRequest::Initialize(InitializeParams {
5756 protocol_version: "2025-11-25".to_string(),
5757 capabilities: ClientCapabilities::default(),
5758 client_info: Implementation {
5759 name: "test".to_string(),
5760 version: "1.0".to_string(),
5761 ..Default::default()
5762 },
5763 meta: None,
5764 }),
5765 extensions: Extensions::new(),
5766 };
5767
5768 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5769 match resp.inner {
5770 Ok(McpResponse::Initialize(result)) => {
5771 assert!(result.capabilities.tools.is_some());
5772 }
5773 _ => panic!("Expected Initialize response"),
5774 }
5775 }
5776
5777 #[tokio::test]
5778 async fn test_dynamic_tools_multi_session_notification() {
5779 let (tx1, mut rx1) = crate::context::notification_channel(16);
5780 let (tx2, mut rx2) = crate::context::notification_channel(16);
5781
5782 let (router, registry) = McpRouter::new()
5783 .server_info("test", "1.0")
5784 .with_dynamic_tools();
5785
5786 let _session1 = router.clone().with_notification_sender(tx1);
5788 let _session2 = router.clone().with_notification_sender(tx2);
5789
5790 let tool = ToolBuilder::new("broadcast")
5791 .description("Test")
5792 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5793 .build();
5794
5795 registry.register(tool);
5796
5797 let n1 = rx1.recv().await.unwrap();
5799 let n2 = rx2.recv().await.unwrap();
5800 assert!(matches!(n1, ServerNotification::ToolsListChanged));
5801 assert!(matches!(n2, ServerNotification::ToolsListChanged));
5802 }
5803
5804 #[tokio::test]
5805 async fn test_dynamic_tools_call_not_found() {
5806 let (router, _registry) = McpRouter::new()
5807 .server_info("test", "1.0")
5808 .with_dynamic_tools();
5809
5810 let mut router = router;
5811 init_router(&mut router).await;
5812
5813 let req = RouterRequest {
5814 id: RequestId::Number(1),
5815 inner: McpRequest::CallTool(CallToolParams {
5816 name: "nonexistent".to_string(),
5817 arguments: serde_json::json!({}),
5818 meta: None,
5819 task: None,
5820 }),
5821 extensions: Extensions::new(),
5822 };
5823
5824 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5825 match resp.inner {
5826 Err(e) => {
5827 assert_eq!(e.code, -32601);
5828 }
5829 _ => panic!("Expected method not found error"),
5830 }
5831 }
5832
5833 #[tokio::test]
5834 async fn test_dynamic_tools_registry_list() {
5835 let (_, registry) = McpRouter::new()
5836 .server_info("test", "1.0")
5837 .with_dynamic_tools();
5838
5839 assert!(registry.list().is_empty());
5840
5841 let tool = ToolBuilder::new("tool_a")
5842 .description("A")
5843 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5844 .build();
5845 registry.register(tool);
5846
5847 let tool = ToolBuilder::new("tool_b")
5848 .description("B")
5849 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5850 .build();
5851 registry.register(tool);
5852
5853 let tools = registry.list();
5854 assert_eq!(tools.len(), 2);
5855 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
5856 assert!(names.contains(&"tool_a"));
5857 assert!(names.contains(&"tool_b"));
5858 }
5859 } #[tokio::test]
5862 async fn test_tool_if_true_registers() {
5863 let tool = ToolBuilder::new("conditional")
5864 .description("Conditional tool")
5865 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5866 .build();
5867
5868 let mut router = McpRouter::new().tool_if(true, tool);
5869 init_router(&mut router).await;
5870
5871 let req = RouterRequest {
5872 id: RequestId::Number(1),
5873 inner: McpRequest::ListTools(ListToolsParams::default()),
5874 extensions: Extensions::new(),
5875 };
5876 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5877 match resp.inner {
5878 Ok(McpResponse::ListTools(result)) => {
5879 assert_eq!(result.tools.len(), 1);
5880 assert_eq!(result.tools[0].name, "conditional");
5881 }
5882 _ => panic!("Expected ListTools response"),
5883 }
5884 }
5885
5886 #[tokio::test]
5887 async fn test_tool_if_false_skips() {
5888 let tool = ToolBuilder::new("conditional")
5889 .description("Conditional tool")
5890 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5891 .build();
5892
5893 let mut router = McpRouter::new().tool_if(false, tool);
5894 init_router(&mut router).await;
5895
5896 let req = RouterRequest {
5897 id: RequestId::Number(1),
5898 inner: McpRequest::ListTools(ListToolsParams::default()),
5899 extensions: Extensions::new(),
5900 };
5901 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5902 match resp.inner {
5903 Ok(McpResponse::ListTools(result)) => {
5904 assert_eq!(result.tools.len(), 0);
5905 }
5906 _ => panic!("Expected ListTools response"),
5907 }
5908 }
5909
5910 #[tokio::test]
5911 async fn test_tools_if_batch_conditional() {
5912 let tools = vec![
5913 ToolBuilder::new("a")
5914 .description("Tool A")
5915 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5916 .build(),
5917 ToolBuilder::new("b")
5918 .description("Tool B")
5919 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5920 .build(),
5921 ];
5922
5923 let mut router = McpRouter::new().tools_if(false, tools);
5924 init_router(&mut router).await;
5925
5926 let req = RouterRequest {
5927 id: RequestId::Number(1),
5928 inner: McpRequest::ListTools(ListToolsParams::default()),
5929 extensions: Extensions::new(),
5930 };
5931 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5932 match resp.inner {
5933 Ok(McpResponse::ListTools(result)) => {
5934 assert_eq!(result.tools.len(), 0);
5935 }
5936 _ => panic!("Expected ListTools response"),
5937 }
5938 }
5939
5940 #[test]
5941 fn test_resource_if_true_registers() {
5942 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
5943 .name("test")
5944 .text("hello");
5945
5946 let router = McpRouter::new().resource_if(true, resource);
5947 assert_eq!(router.inner.resources.len(), 1);
5948 }
5949
5950 #[test]
5951 fn test_resource_if_false_skips() {
5952 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
5953 .name("test")
5954 .text("hello");
5955
5956 let router = McpRouter::new().resource_if(false, resource);
5957 assert_eq!(router.inner.resources.len(), 0);
5958 }
5959
5960 #[test]
5961 fn test_prompt_if_true_registers() {
5962 let prompt = crate::prompt::PromptBuilder::new("greet")
5963 .description("Greeting")
5964 .user_message("Hello!");
5965
5966 let router = McpRouter::new().prompt_if(true, prompt);
5967 assert_eq!(router.inner.prompts.len(), 1);
5968 }
5969
5970 #[test]
5971 fn test_prompt_if_false_skips() {
5972 let prompt = crate::prompt::PromptBuilder::new("greet")
5973 .description("Greeting")
5974 .user_message("Hello!");
5975
5976 let router = McpRouter::new().prompt_if(false, prompt);
5977 assert_eq!(router.inner.prompts.len(), 0);
5978 }
5979}