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 list_ttl_ms: Option<u64>,
190 logging_deprecated: Option<tower_mcp_types::protocol::DeprecationInfo>,
193 disabled_tools: Arc<RwLock<HashSet<String>>>,
195 disabled_resources: Arc<RwLock<HashSet<String>>>,
197 disabled_prompts: Arc<RwLock<HashSet<String>>>,
199 #[cfg(feature = "dynamic-tools")]
201 dynamic_tools: Option<Arc<DynamicToolsInner>>,
202 #[cfg(feature = "dynamic-tools")]
204 dynamic_prompts: Option<Arc<DynamicPromptsInner>>,
205 #[cfg(feature = "dynamic-tools")]
207 dynamic_resources: Option<Arc<DynamicResourcesInner>>,
208 #[cfg(feature = "dynamic-tools")]
210 dynamic_resource_templates: Option<Arc<DynamicResourceTemplatesInner>>,
211}
212
213impl McpRouterInner {
214 fn generate_instructions(&self, config: &AutoInstructionsConfig) -> String {
216 let mut parts = Vec::new();
217
218 if let Some(prefix) = &config.prefix {
219 parts.push(prefix.clone());
220 }
221
222 if !self.tools.is_empty() {
224 let mut lines = vec!["## Tools".to_string(), String::new()];
225 let mut tools: Vec<_> = self.tools.values().collect();
226 tools.sort_by(|a, b| a.name.cmp(&b.name));
227 for tool in tools {
228 let desc = tool.description.as_deref().unwrap_or("No description");
229 let tags = annotation_tags(tool.annotations.as_ref());
230 if tags.is_empty() {
231 lines.push(format!("- **{}**: {}", tool.name, desc));
232 } else {
233 lines.push(format!("- **{}**: {} [{}]", tool.name, desc, tags));
234 }
235 }
236 parts.push(lines.join("\n"));
237 }
238
239 if !self.resources.is_empty() || !self.resource_templates.is_empty() {
241 let mut lines = vec!["## Resources".to_string(), String::new()];
242 let mut resources: Vec<_> = self.resources.values().collect();
243 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
244 for resource in resources {
245 let desc = resource.description.as_deref().unwrap_or("No description");
246 lines.push(format!("- **{}**: {}", resource.uri, desc));
247 }
248 let mut templates: Vec<_> = self.resource_templates.iter().collect();
249 templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
250 for template in templates {
251 let desc = template.description.as_deref().unwrap_or("No description");
252 lines.push(format!("- **{}**: {}", template.uri_template, desc));
253 }
254 parts.push(lines.join("\n"));
255 }
256
257 if !self.prompts.is_empty() {
259 let mut lines = vec!["## Prompts".to_string(), String::new()];
260 let mut prompts: Vec<_> = self.prompts.values().collect();
261 prompts.sort_by(|a, b| a.name.cmp(&b.name));
262 for prompt in prompts {
263 let desc = prompt.description.as_deref().unwrap_or("No description");
264 lines.push(format!("- **{}**: {}", prompt.name, desc));
265 }
266 parts.push(lines.join("\n"));
267 }
268
269 if let Some(suffix) = &config.suffix {
270 parts.push(suffix.clone());
271 }
272
273 parts.join("\n\n")
274 }
275}
276
277fn annotation_tags(annotations: Option<&crate::protocol::ToolAnnotations>) -> String {
283 let Some(ann) = annotations else {
284 return String::new();
285 };
286 let mut tags = Vec::new();
287 if ann.is_read_only() {
288 tags.push("read-only");
289 }
290 if ann.is_idempotent() {
291 tags.push("idempotent");
292 }
293 tags.join(", ")
294}
295
296impl McpRouter {
297 pub fn new() -> Self {
299 Self {
300 inner: Arc::new(McpRouterInner {
301 server_name: "tower-mcp".to_string(),
302 server_version: env!("CARGO_PKG_VERSION").to_string(),
303 server_title: None,
304 server_description: None,
305 server_icons: None,
306 server_website_url: None,
307 instructions: None,
308 auto_instructions: None,
309 tools: HashMap::new(),
310 resources: HashMap::new(),
311 resource_templates: Vec::new(),
312 prompts: HashMap::new(),
313 in_flight: Arc::new(RwLock::new(HashMap::new())),
314 notification_tx: None,
315 client_requester: None,
316 task_store: TaskStore::new(),
317 subscriptions: Arc::new(RwLock::new(HashSet::new())),
318 extensions: Arc::new(crate::context::Extensions::new()),
319 completion_handler: None,
320 tool_filter: None,
321 resource_filter: None,
322 prompt_filter: None,
323 min_log_level: Arc::new(RwLock::new(LogLevel::Debug)),
324 page_size: None,
325 list_ttl_ms: None,
326 logging_deprecated: None,
327 disabled_tools: Arc::new(RwLock::new(HashSet::new())),
328 disabled_resources: Arc::new(RwLock::new(HashSet::new())),
329 disabled_prompts: Arc::new(RwLock::new(HashSet::new())),
330 #[cfg(feature = "dynamic-tools")]
331 dynamic_tools: None,
332 #[cfg(feature = "dynamic-tools")]
333 dynamic_prompts: None,
334 #[cfg(feature = "dynamic-tools")]
335 dynamic_resources: None,
336 #[cfg(feature = "dynamic-tools")]
337 dynamic_resource_templates: None,
338 }),
339 session: SessionState::new(),
340 }
341 }
342
343 pub fn with_fresh_session(&self) -> Self {
351 Self {
352 inner: self.inner.clone(),
353 session: SessionState::new(),
354 }
355 }
356
357 pub fn tool_annotations_map(&self) -> ToolAnnotationsMap {
367 let disabled = self.inner.disabled_tools.read().unwrap();
368 let mut map = HashMap::new();
369 for (name, tool) in &self.inner.tools {
370 if disabled.contains(name) {
371 continue;
372 }
373 if let Some(annotations) = &tool.annotations {
374 map.insert(name.clone(), annotations.clone());
375 }
376 }
377 #[cfg(feature = "dynamic-tools")]
378 if let Some(dynamic) = &self.inner.dynamic_tools {
379 for tool in dynamic.list() {
380 if disabled.contains(&tool.name) {
381 continue;
382 }
383 if !map.contains_key(&tool.name)
385 && let Some(ref annotations) = tool.annotations
386 {
387 map.insert(tool.name.clone(), annotations.clone());
388 }
389 }
390 }
391 ToolAnnotationsMap { map: Arc::new(map) }
392 }
393
394 pub fn task_store(&self) -> &TaskStore {
396 &self.inner.task_store
397 }
398
399 #[cfg(feature = "dynamic-tools")]
429 pub fn with_dynamic_tools(mut self) -> (Self, DynamicToolRegistry) {
430 let inner_dyn = Arc::new(DynamicToolsInner::new());
431 Arc::make_mut(&mut self.inner).dynamic_tools = Some(inner_dyn.clone());
432 (self, DynamicToolRegistry::new(inner_dyn))
433 }
434
435 #[cfg(feature = "dynamic-tools")]
458 pub fn with_dynamic_prompts(mut self) -> (Self, DynamicPromptRegistry) {
459 let inner_dyn = Arc::new(DynamicPromptsInner::new());
460 Arc::make_mut(&mut self.inner).dynamic_prompts = Some(inner_dyn.clone());
461 (self, DynamicPromptRegistry::new(inner_dyn))
462 }
463
464 #[cfg(feature = "dynamic-tools")]
487 pub fn with_dynamic_resources(mut self) -> (Self, DynamicResourceRegistry) {
488 let inner_dyn = Arc::new(DynamicResourcesInner::new());
489 Arc::make_mut(&mut self.inner).dynamic_resources = Some(inner_dyn.clone());
490 (self, DynamicResourceRegistry::new(inner_dyn))
491 }
492
493 #[cfg(feature = "dynamic-tools")]
515 pub fn with_dynamic_resource_templates(mut self) -> (Self, DynamicResourceTemplateRegistry) {
516 let inner_dyn = Arc::new(DynamicResourceTemplatesInner::new());
517 Arc::make_mut(&mut self.inner).dynamic_resource_templates = Some(inner_dyn.clone());
518 (self, DynamicResourceTemplateRegistry::new(inner_dyn))
519 }
520
521 pub fn with_notification_sender(mut self, tx: NotificationSender) -> Self {
525 let inner = Arc::make_mut(&mut self.inner);
526 #[cfg(feature = "dynamic-tools")]
529 if let Some(ref dynamic_tools) = inner.dynamic_tools {
530 dynamic_tools.add_notification_sender(tx.clone());
531 }
532 #[cfg(feature = "dynamic-tools")]
533 if let Some(ref dynamic_prompts) = inner.dynamic_prompts {
534 dynamic_prompts.add_notification_sender(tx.clone());
535 }
536 #[cfg(feature = "dynamic-tools")]
537 if let Some(ref dynamic_resources) = inner.dynamic_resources {
538 dynamic_resources.add_notification_sender(tx.clone());
539 }
540 #[cfg(feature = "dynamic-tools")]
541 if let Some(ref dynamic_resource_templates) = inner.dynamic_resource_templates {
542 dynamic_resource_templates.add_notification_sender(tx.clone());
543 }
544 inner.notification_tx = Some(tx);
545 self
546 }
547
548 pub fn notification_sender(&self) -> Option<&NotificationSender> {
550 self.inner.notification_tx.as_ref()
551 }
552
553 pub fn with_client_requester(mut self, requester: ClientRequesterHandle) -> Self {
558 Arc::make_mut(&mut self.inner).client_requester = Some(requester);
559 self
560 }
561
562 pub fn client_requester(&self) -> Option<&ClientRequesterHandle> {
564 self.inner.client_requester.as_ref()
565 }
566
567 pub fn with_state<T: Clone + Send + Sync + 'static>(mut self, state: T) -> Self {
610 let inner = Arc::make_mut(&mut self.inner);
611 Arc::make_mut(&mut inner.extensions).insert(state);
612 self
613 }
614
615 pub fn with_extension<T: Clone + Send + Sync + 'static>(self, value: T) -> Self {
620 self.with_state(value)
621 }
622
623 pub fn extensions(&self) -> &crate::context::Extensions {
625 &self.inner.extensions
626 }
627
628 pub fn create_context(
633 &self,
634 request_id: RequestId,
635 progress_token: Option<ProgressToken>,
636 ) -> RequestContext {
637 self.create_context_with_extensions(request_id, progress_token, &Extensions::new())
638 }
639
640 pub(crate) fn create_context_with_extensions(
645 &self,
646 request_id: RequestId,
647 progress_token: Option<ProgressToken>,
648 per_request: &Extensions,
649 ) -> RequestContext {
650 let ctx = RequestContext::new(request_id.clone());
651
652 let ctx = if let Some(token) = progress_token {
654 ctx.with_progress_token(token)
655 } else {
656 ctx
657 };
658
659 let ctx = if let Some(tx) = &self.inner.notification_tx {
661 ctx.with_notification_sender(tx.clone())
662 } else {
663 ctx
664 };
665
666 let ctx = if let Some(requester) = &self.inner.client_requester {
668 ctx.with_client_requester(requester.clone())
669 } else {
670 ctx
671 };
672
673 let mut merged = (*self.inner.extensions).clone();
677 merged.merge(per_request);
678 let ctx = ctx.with_extensions(Arc::new(merged));
679
680 let ctx = ctx.with_min_log_level(self.inner.min_log_level.clone());
682
683 let token = ctx.cancellation_token();
685 if let Ok(mut in_flight) = self.inner.in_flight.write() {
686 in_flight.insert(request_id, token);
687 }
688
689 ctx
690 }
691
692 pub fn complete_request(&self, request_id: &RequestId) {
694 if let Ok(mut in_flight) = self.inner.in_flight.write() {
695 in_flight.remove(request_id);
696 }
697 }
698
699 fn cancel_request(&self, request_id: &RequestId) -> bool {
701 let Ok(in_flight) = self.inner.in_flight.read() else {
702 return false;
703 };
704 let Some(token) = in_flight.get(request_id) else {
705 return false;
706 };
707 token.cancel();
708 true
709 }
710
711 pub fn server_info(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
713 let inner = Arc::make_mut(&mut self.inner);
714 inner.server_name = name.into();
715 inner.server_version = version.into();
716 self
717 }
718
719 pub fn page_size(mut self, size: usize) -> Self {
726 Arc::make_mut(&mut self.inner).page_size = Some(size);
727 self
728 }
729
730 pub fn list_ttl(mut self, ms: u64) -> Self {
736 Arc::make_mut(&mut self.inner).list_ttl_ms = Some(ms);
737 self
738 }
739
740 pub fn logging_deprecated(mut self, info: tower_mcp_types::protocol::DeprecationInfo) -> Self {
746 Arc::make_mut(&mut self.inner).logging_deprecated = Some(info);
747 self
748 }
749
750 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
752 Arc::make_mut(&mut self.inner).instructions = Some(instructions.into());
753 self
754 }
755
756 pub fn auto_instructions(mut self) -> Self {
788 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
789 prefix: None,
790 suffix: None,
791 });
792 self
793 }
794
795 pub fn auto_instructions_with(
812 mut self,
813 prefix: Option<impl Into<String>>,
814 suffix: Option<impl Into<String>>,
815 ) -> Self {
816 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
817 prefix: prefix.map(Into::into),
818 suffix: suffix.map(Into::into),
819 });
820 self
821 }
822
823 pub fn server_title(mut self, title: impl Into<String>) -> Self {
825 Arc::make_mut(&mut self.inner).server_title = Some(title.into());
826 self
827 }
828
829 pub fn server_description(mut self, description: impl Into<String>) -> Self {
831 Arc::make_mut(&mut self.inner).server_description = Some(description.into());
832 self
833 }
834
835 pub fn server_icons(mut self, icons: Vec<ToolIcon>) -> Self {
837 Arc::make_mut(&mut self.inner).server_icons = Some(icons);
838 self
839 }
840
841 pub fn server_website_url(mut self, url: impl Into<String>) -> Self {
843 Arc::make_mut(&mut self.inner).server_website_url = Some(url.into());
844 self
845 }
846
847 pub fn tool(mut self, tool: Tool) -> Self {
849 Arc::make_mut(&mut self.inner)
850 .tools
851 .insert(tool.name.clone(), Arc::new(tool));
852 self
853 }
854
855 pub fn tool_if(self, condition: bool, tool: Tool) -> Self {
881 if condition { self.tool(tool) } else { self }
882 }
883
884 pub fn resource(mut self, resource: Resource) -> Self {
886 Arc::make_mut(&mut self.inner)
887 .resources
888 .insert(resource.uri.clone(), Arc::new(resource));
889 self
890 }
891
892 pub fn resource_if(self, condition: bool, resource: Resource) -> Self {
911 if condition {
912 self.resource(resource)
913 } else {
914 self
915 }
916 }
917
918 pub fn resource_template(mut self, template: ResourceTemplate) -> Self {
951 Arc::make_mut(&mut self.inner)
952 .resource_templates
953 .push(Arc::new(template));
954 self
955 }
956
957 pub fn prompt(mut self, prompt: Prompt) -> Self {
959 Arc::make_mut(&mut self.inner)
960 .prompts
961 .insert(prompt.name.clone(), Arc::new(prompt));
962 self
963 }
964
965 pub fn prompt_if(self, condition: bool, prompt: Prompt) -> Self {
984 if condition { self.prompt(prompt) } else { self }
985 }
986
987 pub fn tools(self, tools: impl IntoIterator<Item = Tool>) -> Self {
1013 tools
1014 .into_iter()
1015 .fold(self, |router, tool| router.tool(tool))
1016 }
1017
1018 pub fn tools_if(self, condition: bool, tools: impl IntoIterator<Item = Tool>) -> Self {
1022 if condition { self.tools(tools) } else { self }
1023 }
1024
1025 pub fn resources(self, resources: impl IntoIterator<Item = Resource>) -> Self {
1044 resources
1045 .into_iter()
1046 .fold(self, |router, resource| router.resource(resource))
1047 }
1048
1049 pub fn resources_if(
1053 self,
1054 condition: bool,
1055 resources: impl IntoIterator<Item = Resource>,
1056 ) -> Self {
1057 if condition {
1058 self.resources(resources)
1059 } else {
1060 self
1061 }
1062 }
1063
1064 pub fn prompts(self, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1083 prompts
1084 .into_iter()
1085 .fold(self, |router, prompt| router.prompt(prompt))
1086 }
1087
1088 pub fn prompts_if(self, condition: bool, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1092 if condition {
1093 self.prompts(prompts)
1094 } else {
1095 self
1096 }
1097 }
1098
1099 pub fn merge(mut self, other: McpRouter) -> Self {
1144 let inner = Arc::make_mut(&mut self.inner);
1145 let other_inner = other.inner;
1146
1147 for (name, tool) in &other_inner.tools {
1149 inner.tools.insert(name.clone(), tool.clone());
1150 }
1151
1152 for (uri, resource) in &other_inner.resources {
1154 inner.resources.insert(uri.clone(), resource.clone());
1155 }
1156
1157 for template in &other_inner.resource_templates {
1160 inner.resource_templates.push(template.clone());
1161 }
1162
1163 for (name, prompt) in &other_inner.prompts {
1165 inner.prompts.insert(name.clone(), prompt.clone());
1166 }
1167
1168 self
1169 }
1170
1171 pub fn nest(mut self, prefix: impl Into<String>, other: McpRouter) -> Self {
1211 let prefix = prefix.into();
1212 let inner = Arc::make_mut(&mut self.inner);
1213 let other_inner = other.inner;
1214
1215 for tool in other_inner.tools.values() {
1217 let prefixed_tool = tool.with_name_prefix(&prefix);
1218 inner
1219 .tools
1220 .insert(prefixed_tool.name.clone(), Arc::new(prefixed_tool));
1221 }
1222
1223 for (uri, resource) in &other_inner.resources {
1225 inner.resources.insert(uri.clone(), resource.clone());
1226 }
1227
1228 for template in &other_inner.resource_templates {
1230 inner.resource_templates.push(template.clone());
1231 }
1232
1233 for (name, prompt) in &other_inner.prompts {
1235 inner.prompts.insert(name.clone(), prompt.clone());
1236 }
1237
1238 self
1239 }
1240
1241 pub fn completion_handler<F, Fut>(mut self, handler: F) -> Self
1269 where
1270 F: Fn(CompleteParams) -> Fut + Send + Sync + 'static,
1271 Fut: Future<Output = Result<CompleteResult>> + Send + 'static,
1272 {
1273 Arc::make_mut(&mut self.inner).completion_handler =
1274 Some(Arc::new(move |params| Box::pin(handler(params))));
1275 self
1276 }
1277
1278 pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
1313 Arc::make_mut(&mut self.inner).tool_filter = Some(filter);
1314 self
1315 }
1316
1317 pub fn resource_filter(mut self, filter: ResourceFilter) -> Self {
1348 Arc::make_mut(&mut self.inner).resource_filter = Some(filter);
1349 self
1350 }
1351
1352 pub fn prompt_filter(mut self, filter: PromptFilter) -> Self {
1381 Arc::make_mut(&mut self.inner).prompt_filter = Some(filter);
1382 self
1383 }
1384
1385 pub fn session(&self) -> &SessionState {
1387 &self.session
1388 }
1389
1390 pub fn log(&self, params: LoggingMessageParams) -> bool {
1412 let Some(tx) = &self.inner.notification_tx else {
1413 return false;
1414 };
1415 tx.try_send(ServerNotification::LogMessage(params)).is_ok()
1416 }
1417
1418 pub fn log_info(&self, message: &str) -> bool {
1422 self.log(LoggingMessageParams::new(
1423 LogLevel::Info,
1424 serde_json::json!({ "message": message }),
1425 ))
1426 }
1427
1428 pub fn log_warning(&self, message: &str) -> bool {
1430 self.log(LoggingMessageParams::new(
1431 LogLevel::Warning,
1432 serde_json::json!({ "message": message }),
1433 ))
1434 }
1435
1436 pub fn log_error(&self, message: &str) -> bool {
1438 self.log(LoggingMessageParams::new(
1439 LogLevel::Error,
1440 serde_json::json!({ "message": message }),
1441 ))
1442 }
1443
1444 pub fn log_debug(&self, message: &str) -> bool {
1446 self.log(LoggingMessageParams::new(
1447 LogLevel::Debug,
1448 serde_json::json!({ "message": message }),
1449 ))
1450 }
1451
1452 pub fn is_subscribed(&self, uri: &str) -> bool {
1454 if let Ok(subs) = self.inner.subscriptions.read() {
1455 return subs.contains(uri);
1456 }
1457 false
1458 }
1459
1460 pub fn subscribed_uris(&self) -> Vec<String> {
1462 if let Ok(subs) = self.inner.subscriptions.read() {
1463 return subs.iter().cloned().collect();
1464 }
1465 Vec::new()
1466 }
1467
1468 fn subscribe(&self, uri: &str) -> bool {
1470 if let Ok(mut subs) = self.inner.subscriptions.write() {
1471 return subs.insert(uri.to_string());
1472 }
1473 false
1474 }
1475
1476 fn unsubscribe(&self, uri: &str) -> bool {
1478 if let Ok(mut subs) = self.inner.subscriptions.write() {
1479 return subs.remove(uri);
1480 }
1481 false
1482 }
1483
1484 pub fn notify_resource_updated(&self, uri: &str) -> bool {
1489 if !self.is_subscribed(uri) {
1491 return false;
1492 }
1493
1494 let Some(tx) = &self.inner.notification_tx else {
1495 return false;
1496 };
1497 tx.try_send(ServerNotification::ResourceUpdated {
1498 uri: uri.to_string(),
1499 })
1500 .is_ok()
1501 }
1502
1503 pub fn notify_resources_list_changed(&self) -> bool {
1507 let Some(tx) = &self.inner.notification_tx else {
1508 return false;
1509 };
1510 tx.try_send(ServerNotification::ResourcesListChanged)
1511 .is_ok()
1512 }
1513
1514 pub fn notify_tools_list_changed(&self) -> bool {
1518 let Some(tx) = &self.inner.notification_tx else {
1519 return false;
1520 };
1521 tx.try_send(ServerNotification::ToolsListChanged).is_ok()
1522 }
1523
1524 pub fn notify_prompts_list_changed(&self) -> bool {
1528 let Some(tx) = &self.inner.notification_tx else {
1529 return false;
1530 };
1531 tx.try_send(ServerNotification::PromptsListChanged).is_ok()
1532 }
1533
1534 pub fn disable_tool(&self, name: impl Into<String>) {
1545 let mut set = self.inner.disabled_tools.write().unwrap();
1546 set.insert(name.into());
1547 }
1548
1549 pub fn enable_tool(&self, name: &str) {
1552 let mut set = self.inner.disabled_tools.write().unwrap();
1553 set.remove(name);
1554 }
1555
1556 pub fn is_tool_enabled(&self, name: &str) -> bool {
1560 !self.inner.disabled_tools.read().unwrap().contains(name)
1561 }
1562
1563 pub fn disable_resource(&self, uri: impl Into<String>) {
1566 let mut set = self.inner.disabled_resources.write().unwrap();
1567 set.insert(uri.into());
1568 }
1569
1570 pub fn enable_resource(&self, uri: &str) {
1572 let mut set = self.inner.disabled_resources.write().unwrap();
1573 set.remove(uri);
1574 }
1575
1576 pub fn is_resource_enabled(&self, uri: &str) -> bool {
1578 !self.inner.disabled_resources.read().unwrap().contains(uri)
1579 }
1580
1581 pub fn disable_prompt(&self, name: impl Into<String>) {
1584 let mut set = self.inner.disabled_prompts.write().unwrap();
1585 set.insert(name.into());
1586 }
1587
1588 pub fn enable_prompt(&self, name: &str) {
1590 let mut set = self.inner.disabled_prompts.write().unwrap();
1591 set.remove(name);
1592 }
1593
1594 pub fn is_prompt_enabled(&self, name: &str) -> bool {
1596 !self.inner.disabled_prompts.read().unwrap().contains(name)
1597 }
1598
1599 fn capabilities(&self) -> ServerCapabilities {
1601 let has_resources =
1602 !self.inner.resources.is_empty() || !self.inner.resource_templates.is_empty();
1603 let has_notifications = self.inner.notification_tx.is_some();
1604
1605 #[cfg(feature = "dynamic-tools")]
1606 let has_dynamic_tools = self.inner.dynamic_tools.is_some();
1607 #[cfg(not(feature = "dynamic-tools"))]
1608 let has_dynamic_tools = false;
1609
1610 #[cfg(feature = "dynamic-tools")]
1611 let has_dynamic_prompts = self.inner.dynamic_prompts.is_some();
1612 #[cfg(not(feature = "dynamic-tools"))]
1613 let has_dynamic_prompts = false;
1614
1615 #[cfg(feature = "dynamic-tools")]
1616 let has_dynamic_resources = self.inner.dynamic_resources.is_some()
1617 || self.inner.dynamic_resource_templates.is_some();
1618 #[cfg(not(feature = "dynamic-tools"))]
1619 let has_dynamic_resources = false;
1620
1621 ServerCapabilities {
1622 tools: if self.inner.tools.is_empty() && !has_dynamic_tools {
1623 None
1624 } else {
1625 Some(ToolsCapability {
1626 list_changed: has_notifications,
1627 })
1628 },
1629 resources: if has_resources || has_dynamic_resources {
1630 Some(ResourcesCapability {
1631 subscribe: true,
1632 list_changed: has_notifications,
1633 })
1634 } else {
1635 None
1636 },
1637 prompts: if self.inner.prompts.is_empty() && !has_dynamic_prompts {
1638 None
1639 } else {
1640 Some(PromptsCapability {
1641 list_changed: has_notifications,
1642 })
1643 },
1644 logging: if self.inner.notification_tx.is_some() {
1646 Some(LoggingCapability {
1647 deprecated: self.inner.logging_deprecated.clone(),
1648 })
1649 } else {
1650 None
1651 },
1652 tasks: {
1658 let has_task_support = self
1659 .inner
1660 .tools
1661 .values()
1662 .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1663 if has_task_support {
1664 Some(TasksCapability {
1665 list: Some(TasksListCapability {}),
1666 cancel: Some(TasksCancelCapability {}),
1667 requests: Some(TasksRequestsCapability {
1668 tools: Some(TasksToolsRequestsCapability {
1669 call: Some(TasksToolsCallCapability {}),
1670 }),
1671 }),
1672 })
1673 } else {
1674 None
1675 }
1676 },
1677 completions: if self.inner.completion_handler.is_some() {
1679 Some(CompletionsCapability::default())
1680 } else {
1681 None
1682 },
1683 experimental: None,
1684 extensions: {
1685 let has_task_support = self
1686 .inner
1687 .tools
1688 .values()
1689 .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1690 if has_task_support {
1691 let mut map = std::collections::HashMap::new();
1692 map.insert(
1693 tower_mcp_types::protocol::TASKS_EXTENSION_ID.to_string(),
1694 serde_json::json!({}),
1695 );
1696 Some(map)
1697 } else {
1698 None
1699 }
1700 },
1701 }
1702 }
1703
1704 async fn handle(
1706 &self,
1707 request_id: RequestId,
1708 request: McpRequest,
1709 extensions: Extensions,
1710 ) -> Result<McpResponse> {
1711 let method = request.method_name();
1713 if !self.session.is_request_allowed(method) {
1714 tracing::warn!(
1715 method = %method,
1716 phase = ?self.session.phase(),
1717 "Request rejected: session not initialized"
1718 );
1719 return Err(Error::JsonRpc(JsonRpcError::invalid_request(format!(
1720 "Session not initialized. Only 'initialize' and 'ping' are allowed before initialization. Got: {}",
1721 method
1722 ))));
1723 }
1724
1725 match request {
1726 McpRequest::Initialize(params) => {
1727 tracing::info!(
1728 client = %params.client_info.name,
1729 version = %params.client_info.version,
1730 "Client initializing"
1731 );
1732
1733 let protocol_version = if crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1736 .contains(¶ms.protocol_version.as_str())
1737 {
1738 params.protocol_version
1739 } else {
1740 crate::protocol::LATEST_PROTOCOL_VERSION.to_string()
1741 };
1742
1743 self.session.mark_initializing();
1745
1746 Ok(McpResponse::Initialize(InitializeResult {
1747 protocol_version,
1748 capabilities: self.capabilities(),
1749 server_info: Implementation {
1750 name: self.inner.server_name.clone(),
1751 version: self.inner.server_version.clone(),
1752 title: self.inner.server_title.clone(),
1753 description: self.inner.server_description.clone(),
1754 icons: self.inner.server_icons.clone(),
1755 website_url: self.inner.server_website_url.clone(),
1756 meta: None,
1757 },
1758 instructions: if let Some(config) = &self.inner.auto_instructions {
1759 Some(self.inner.generate_instructions(config))
1760 } else {
1761 self.inner.instructions.clone()
1762 },
1763 meta: None,
1764 }))
1765 }
1766
1767 McpRequest::Discover(_) => {
1768 tracing::debug!("Stateless server/discover request");
1775 Ok(McpResponse::Discover(DiscoverResult {
1776 supported_versions: crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1777 .iter()
1778 .map(|v| (*v).to_string())
1779 .collect(),
1780 capabilities: self.capabilities(),
1781 server_info: Implementation {
1782 name: self.inner.server_name.clone(),
1783 version: self.inner.server_version.clone(),
1784 title: self.inner.server_title.clone(),
1785 description: self.inner.server_description.clone(),
1786 icons: self.inner.server_icons.clone(),
1787 website_url: self.inner.server_website_url.clone(),
1788 meta: None,
1789 },
1790 instructions: if let Some(config) = &self.inner.auto_instructions {
1791 Some(self.inner.generate_instructions(config))
1792 } else {
1793 self.inner.instructions.clone()
1794 },
1795 meta: None,
1796 }))
1797 }
1798
1799 McpRequest::ListTools(params) => {
1800 let filter = self.inner.tool_filter.as_ref();
1801 let disabled = self.inner.disabled_tools.read().unwrap().clone();
1802 let is_visible = |t: &Tool| {
1803 !disabled.contains(&t.name)
1804 && filter
1805 .map(|f| f.is_visible(&self.session, t))
1806 .unwrap_or(true)
1807 };
1808
1809 let mut tools: Vec<ToolDefinition> = self
1811 .inner
1812 .tools
1813 .values()
1814 .filter(|t| is_visible(t))
1815 .map(|t| t.definition())
1816 .collect();
1817
1818 #[cfg(feature = "dynamic-tools")]
1820 if let Some(ref dynamic) = self.inner.dynamic_tools {
1821 let static_names: HashSet<String> =
1822 tools.iter().map(|t| t.name.clone()).collect();
1823 for t in dynamic.list() {
1824 if !static_names.contains(&t.name) && is_visible(&t) {
1825 tools.push(t.definition());
1826 }
1827 }
1828 }
1829
1830 tools.sort_by(|a, b| a.name.cmp(&b.name));
1831
1832 let (tools, next_cursor) =
1833 paginate(tools, params.cursor.as_deref(), self.inner.page_size)?;
1834
1835 Ok(McpResponse::ListTools(ListToolsResult {
1836 tools,
1837 next_cursor,
1838 ttl_ms: self.inner.list_ttl_ms,
1839 cache_scope: None,
1840 meta: None,
1841 }))
1842 }
1843
1844 McpRequest::CallTool(params) => {
1845 if self
1847 .inner
1848 .disabled_tools
1849 .read()
1850 .unwrap()
1851 .contains(¶ms.name)
1852 {
1853 tracing::info!(
1854 target: "mcp::tools",
1855 tool = %params.name,
1856 status = "disabled",
1857 "tool call completed"
1858 );
1859 return Err(Error::JsonRpc(JsonRpcError::method_not_found(¶ms.name)));
1860 }
1861
1862 let tool = self.inner.tools.get(¶ms.name).cloned();
1864 #[cfg(feature = "dynamic-tools")]
1865 let tool = tool.or_else(|| {
1866 self.inner
1867 .dynamic_tools
1868 .as_ref()
1869 .and_then(|d| d.get(¶ms.name))
1870 });
1871
1872 let tool = match tool {
1873 Some(t) => t,
1874 None => {
1875 tracing::info!(
1876 target: "mcp::tools",
1877 tool = %params.name,
1878 status = "not_found",
1879 "tool call completed"
1880 );
1881 return Err(Error::JsonRpc(JsonRpcError::method_not_found(¶ms.name)));
1882 }
1883 };
1884
1885 if let Some(filter) = &self.inner.tool_filter
1887 && !filter.is_visible(&self.session, &tool)
1888 {
1889 tracing::info!(
1890 target: "mcp::tools",
1891 tool = %params.name,
1892 status = "denied",
1893 "tool call completed"
1894 );
1895 return Err(filter.denial_error(¶ms.name));
1896 }
1897
1898 if let Some(task_params) = params.task {
1899 if matches!(tool.task_support, TaskSupportMode::Forbidden) {
1901 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1902 "Tool '{}' does not support async tasks",
1903 params.name
1904 ))));
1905 }
1906
1907 let (task_id, cancellation_token) = self.inner.task_store.create_task(
1909 ¶ms.name,
1910 params.arguments.clone(),
1911 task_params.ttl,
1912 );
1913
1914 tracing::info!(task_id = %task_id, tool = %params.name, "Created async task");
1915
1916 let progress_token = params.meta.and_then(|m| m.progress_token);
1918 let ctx = self.create_context_with_extensions(
1919 request_id,
1920 progress_token,
1921 &extensions,
1922 );
1923
1924 let task_store = self.inner.task_store.clone();
1926 let tool = tool.clone();
1927 let arguments = params.arguments;
1928 let task_id_clone = task_id.clone();
1929
1930 let tool_name = params.name.clone();
1931 tokio::spawn(async move {
1932 if cancellation_token.is_cancelled() {
1934 tracing::debug!(task_id = %task_id_clone, "Task cancelled before execution");
1935 return;
1936 }
1937
1938 let start = std::time::Instant::now();
1940 let result = tool.call_with_context(ctx, arguments).await;
1941 let duration_ms = start.elapsed().as_secs_f64() * 1000.0;
1942
1943 if cancellation_token.is_cancelled() {
1944 tracing::debug!(task_id = %task_id_clone, "Task cancelled during execution");
1945 } else if result.is_error {
1946 let error_msg = result.first_text().unwrap_or("Tool execution failed");
1948 task_store.fail_task(&task_id_clone, error_msg);
1949 tracing::info!(
1950 target: "mcp::tools",
1951 tool = %tool_name,
1952 task_id = %task_id_clone,
1953 duration_ms,
1954 status = "error",
1955 error = %error_msg,
1956 "tool call completed"
1957 );
1958 } else {
1959 task_store.complete_task(&task_id_clone, result);
1960 tracing::info!(
1961 target: "mcp::tools",
1962 tool = %tool_name,
1963 task_id = %task_id_clone,
1964 duration_ms,
1965 status = "success",
1966 "tool call completed"
1967 );
1968 }
1969 });
1970
1971 let task = self.inner.task_store.get_task(&task_id).ok_or_else(|| {
1972 Error::JsonRpc(JsonRpcError::internal_error(
1973 "Failed to retrieve created task",
1974 ))
1975 })?;
1976
1977 Ok(McpResponse::CreateTask(CreateTaskResult::new(task)))
1978 } else {
1979 if matches!(tool.task_support, TaskSupportMode::Required) {
1981 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1982 "Tool '{}' requires async task execution (include 'task' in params)",
1983 params.name
1984 ))));
1985 }
1986
1987 let progress_token = params.meta.and_then(|m| m.progress_token);
1989 let ctx = self.create_context_with_extensions(
1990 request_id,
1991 progress_token,
1992 &extensions,
1993 );
1994
1995 let start = std::time::Instant::now();
1996 let result = tool.call_with_context(ctx, params.arguments).await;
1997 let duration_ms = start.elapsed().as_secs_f64() * 1000.0;
1998
1999 if result.is_error {
2000 tracing::info!(
2001 target: "mcp::tools",
2002 tool = %params.name,
2003 duration_ms,
2004 status = "error",
2005 "tool call completed"
2006 );
2007 } else {
2008 tracing::info!(
2009 target: "mcp::tools",
2010 tool = %params.name,
2011 duration_ms,
2012 status = "success",
2013 "tool call completed"
2014 );
2015 }
2016
2017 Ok(McpResponse::CallTool(result))
2018 }
2019 }
2020
2021 McpRequest::ListResources(params) => {
2022 let disabled = self.inner.disabled_resources.read().unwrap().clone();
2023 let is_visible = |r: &Resource| -> bool {
2024 !disabled.contains(&r.uri)
2025 && self
2026 .inner
2027 .resource_filter
2028 .as_ref()
2029 .map(|f| f.is_visible(&self.session, r))
2030 .unwrap_or(true)
2031 };
2032
2033 let mut resources: Vec<ResourceDefinition> = self
2034 .inner
2035 .resources
2036 .values()
2037 .filter(|r| is_visible(r))
2038 .map(|r| r.definition())
2039 .collect();
2040
2041 #[cfg(feature = "dynamic-tools")]
2043 if let Some(ref dynamic) = self.inner.dynamic_resources {
2044 let static_uris: HashSet<String> =
2045 resources.iter().map(|r| r.uri.clone()).collect();
2046 for r in dynamic.list() {
2047 if !static_uris.contains(&r.uri) && is_visible(&r) {
2048 resources.push(r.definition());
2049 }
2050 }
2051 }
2052
2053 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
2054
2055 let (resources, next_cursor) =
2056 paginate(resources, params.cursor.as_deref(), self.inner.page_size)?;
2057
2058 Ok(McpResponse::ListResources(ListResourcesResult {
2059 resources,
2060 next_cursor,
2061 ttl_ms: self.inner.list_ttl_ms,
2062 cache_scope: None,
2063 meta: None,
2064 }))
2065 }
2066
2067 McpRequest::ListResourceTemplates(params) => {
2068 let mut resource_templates: Vec<ResourceTemplateDefinition> = self
2069 .inner
2070 .resource_templates
2071 .iter()
2072 .map(|t| t.definition())
2073 .collect();
2074
2075 #[cfg(feature = "dynamic-tools")]
2077 if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
2078 let static_patterns: HashSet<String> = resource_templates
2079 .iter()
2080 .map(|t| t.uri_template.clone())
2081 .collect();
2082 for t in dynamic.list() {
2083 if !static_patterns.contains(&t.uri_template) {
2084 resource_templates.push(t.definition());
2085 }
2086 }
2087 }
2088
2089 resource_templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
2090
2091 let (resource_templates, next_cursor) = paginate(
2092 resource_templates,
2093 params.cursor.as_deref(),
2094 self.inner.page_size,
2095 )?;
2096
2097 Ok(McpResponse::ListResourceTemplates(
2098 ListResourceTemplatesResult {
2099 resource_templates,
2100 next_cursor,
2101 ttl_ms: self.inner.list_ttl_ms,
2102 cache_scope: None,
2103 meta: None,
2104 },
2105 ))
2106 }
2107
2108 McpRequest::ReadResource(params) => {
2109 if self
2111 .inner
2112 .disabled_resources
2113 .read()
2114 .unwrap()
2115 .contains(¶ms.uri)
2116 {
2117 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2118 ¶ms.uri,
2119 )));
2120 }
2121
2122 if let Some(resource) = self.inner.resources.get(¶ms.uri) {
2124 if let Some(filter) = &self.inner.resource_filter
2126 && !filter.is_visible(&self.session, resource)
2127 {
2128 return Err(filter.denial_error(¶ms.uri));
2129 }
2130
2131 tracing::debug!(uri = %params.uri, "Reading static resource");
2132 let ctx = self.create_context_with_extensions(request_id, None, &extensions);
2133 let result = resource.read_with_context(ctx).await;
2134 return Ok(McpResponse::ReadResource(result));
2135 }
2136
2137 #[cfg(feature = "dynamic-tools")]
2139 #[allow(clippy::collapsible_if)]
2140 if let Some(ref dynamic) = self.inner.dynamic_resources {
2141 if let Some(resource) = dynamic.get(¶ms.uri) {
2142 if let Some(filter) = &self.inner.resource_filter
2143 && !filter.is_visible(&self.session, &resource)
2144 {
2145 return Err(filter.denial_error(¶ms.uri));
2146 }
2147 tracing::debug!(uri = %params.uri, "Reading dynamic resource");
2148 let ctx =
2149 self.create_context_with_extensions(request_id, None, &extensions);
2150 let result = resource.read_with_context(ctx).await;
2151 return Ok(McpResponse::ReadResource(result));
2152 }
2153 }
2154
2155 for template in &self.inner.resource_templates {
2157 if let Some(variables) = template.match_uri(¶ms.uri) {
2158 tracing::debug!(
2159 uri = %params.uri,
2160 template = %template.uri_template,
2161 "Reading resource via template"
2162 );
2163 let result = template.read(¶ms.uri, variables).await?;
2164 return Ok(McpResponse::ReadResource(result));
2165 }
2166 }
2167
2168 #[cfg(feature = "dynamic-tools")]
2170 #[allow(clippy::collapsible_if)]
2171 if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
2172 if let Some((template, variables)) = dynamic.match_uri(¶ms.uri) {
2173 tracing::debug!(
2174 uri = %params.uri,
2175 template = %template.uri_template,
2176 "Reading resource via dynamic template"
2177 );
2178 let result = template.read(¶ms.uri, variables).await?;
2179 return Ok(McpResponse::ReadResource(result));
2180 }
2181 }
2182
2183 Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2185 ¶ms.uri,
2186 )))
2187 }
2188
2189 McpRequest::SubscribeResource(params) => {
2190 if !self.inner.resources.contains_key(¶ms.uri) {
2192 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2193 ¶ms.uri,
2194 )));
2195 }
2196
2197 tracing::debug!(uri = %params.uri, "Subscribing to resource");
2198 self.subscribe(¶ms.uri);
2199
2200 Ok(McpResponse::SubscribeResource(EmptyResult {}))
2201 }
2202
2203 McpRequest::UnsubscribeResource(params) => {
2204 if !self.inner.resources.contains_key(¶ms.uri) {
2206 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2207 ¶ms.uri,
2208 )));
2209 }
2210
2211 tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
2212 self.unsubscribe(¶ms.uri);
2213
2214 Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
2215 }
2216
2217 McpRequest::ListPrompts(params) => {
2218 let disabled = self.inner.disabled_prompts.read().unwrap().clone();
2219 let is_visible = |p: &Prompt| -> bool {
2220 !disabled.contains(&p.name)
2221 && self
2222 .inner
2223 .prompt_filter
2224 .as_ref()
2225 .map(|f| f.is_visible(&self.session, p))
2226 .unwrap_or(true)
2227 };
2228
2229 let mut prompts: Vec<PromptDefinition> = self
2230 .inner
2231 .prompts
2232 .values()
2233 .filter(|p| is_visible(p))
2234 .map(|p| p.definition())
2235 .collect();
2236
2237 #[cfg(feature = "dynamic-tools")]
2239 if let Some(ref dynamic) = self.inner.dynamic_prompts {
2240 let static_names: HashSet<String> =
2241 prompts.iter().map(|p| p.name.clone()).collect();
2242 for p in dynamic.list() {
2243 if !static_names.contains(&p.name) && is_visible(&p) {
2244 prompts.push(p.definition());
2245 }
2246 }
2247 }
2248
2249 prompts.sort_by(|a, b| a.name.cmp(&b.name));
2250
2251 let (prompts, next_cursor) =
2252 paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
2253
2254 Ok(McpResponse::ListPrompts(ListPromptsResult {
2255 prompts,
2256 next_cursor,
2257 ttl_ms: self.inner.list_ttl_ms,
2258 cache_scope: None,
2259 meta: None,
2260 }))
2261 }
2262
2263 McpRequest::GetPrompt(params) => {
2264 if self
2266 .inner
2267 .disabled_prompts
2268 .read()
2269 .unwrap()
2270 .contains(¶ms.name)
2271 {
2272 return Err(Error::JsonRpc(JsonRpcError::method_not_found(&format!(
2273 "Prompt not found: {}",
2274 params.name
2275 ))));
2276 }
2277
2278 let prompt = self.inner.prompts.get(¶ms.name).cloned();
2280 #[cfg(feature = "dynamic-tools")]
2281 let prompt = prompt.or_else(|| {
2282 self.inner
2283 .dynamic_prompts
2284 .as_ref()
2285 .and_then(|d| d.get(¶ms.name))
2286 });
2287 let prompt = prompt.ok_or_else(|| {
2288 Error::JsonRpc(JsonRpcError::method_not_found(&format!(
2289 "Prompt not found: {}",
2290 params.name
2291 )))
2292 })?;
2293
2294 if let Some(filter) = &self.inner.prompt_filter
2296 && !filter.is_visible(&self.session, &prompt)
2297 {
2298 return Err(filter.denial_error(¶ms.name));
2299 }
2300
2301 tracing::debug!(name = %params.name, "Getting prompt");
2302 let ctx = self.create_context_with_extensions(request_id, None, &extensions);
2303 let result = prompt.get_with_context(ctx, params.arguments).await?;
2304
2305 Ok(McpResponse::GetPrompt(result))
2306 }
2307
2308 McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
2309
2310 McpRequest::ListTasks(params) => {
2311 let tasks = self.inner.task_store.list_tasks(params.status);
2312
2313 let (tasks, next_cursor) =
2314 paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
2315
2316 Ok(McpResponse::ListTasks(ListTasksResult {
2317 tasks,
2318 next_cursor,
2319 }))
2320 }
2321
2322 McpRequest::GetTaskInfo(params) => {
2323 let task = self
2324 .inner
2325 .task_store
2326 .get_task(¶ms.task_id)
2327 .ok_or_else(|| {
2328 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2329 "Task not found: {}",
2330 params.task_id
2331 )))
2332 })?;
2333
2334 Ok(McpResponse::GetTaskInfo(task))
2335 }
2336
2337 McpRequest::GetTaskResult(params) => {
2338 let (task_obj, result, error) = self
2340 .inner
2341 .task_store
2342 .wait_for_completion(¶ms.task_id)
2343 .await
2344 .ok_or_else(|| {
2345 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2346 "Task not found: {}",
2347 params.task_id
2348 )))
2349 })?;
2350
2351 let meta = serde_json::json!({
2353 "io.modelcontextprotocol/related-task": task_obj
2354 });
2355
2356 match task_obj.status {
2357 TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
2358 format!("Task {} was cancelled", params.task_id),
2359 ))),
2360 TaskStatus::Failed => {
2361 let mut call_result = CallToolResult::error(
2362 error.unwrap_or_else(|| "Task failed".to_string()),
2363 );
2364 call_result.meta = Some(meta);
2365 Ok(McpResponse::GetTaskResult(call_result))
2366 }
2367 _ => {
2368 let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
2369 call_result.meta = Some(meta);
2370 Ok(McpResponse::GetTaskResult(call_result))
2371 }
2372 }
2373 }
2374
2375 McpRequest::UpdateTask(params) => {
2376 let _ = self
2384 .inner
2385 .task_store
2386 .get_task(¶ms.task_id)
2387 .ok_or_else(|| {
2388 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2389 "Task not found: {}",
2390 params.task_id
2391 )))
2392 })?;
2393 Ok(McpResponse::UpdateTask(EmptyResult {}))
2394 }
2395
2396 McpRequest::CancelTask(params) => {
2397 let current = self
2399 .inner
2400 .task_store
2401 .get_task(¶ms.task_id)
2402 .ok_or_else(|| {
2403 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2404 "Task not found: {}",
2405 params.task_id
2406 )))
2407 })?;
2408
2409 if current.status.is_terminal() {
2410 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
2411 "Task {} is already in terminal state: {}",
2412 params.task_id, current.status
2413 ))));
2414 }
2415
2416 let task_obj = self
2417 .inner
2418 .task_store
2419 .cancel_task(¶ms.task_id, params.reason.as_deref())
2420 .ok_or_else(|| {
2421 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2422 "Task not found: {}",
2423 params.task_id
2424 )))
2425 })?;
2426
2427 Ok(McpResponse::CancelTask(task_obj))
2428 }
2429
2430 McpRequest::SetLoggingLevel(params) => {
2431 tracing::debug!(level = ?params.level, "Client set logging level");
2432 if let Ok(mut level) = self.inner.min_log_level.write() {
2433 *level = params.level;
2434 }
2435 Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
2436 }
2437
2438 McpRequest::Complete(params) => {
2439 tracing::debug!(
2440 reference = ?params.reference,
2441 argument = %params.argument.name,
2442 "Completion request"
2443 );
2444
2445 if let Some(ref handler) = self.inner.completion_handler {
2447 let result = handler(params).await?;
2448 Ok(McpResponse::Complete(result))
2449 } else {
2450 Ok(McpResponse::Complete(CompleteResult::new(vec![])))
2452 }
2453 }
2454
2455 McpRequest::Unknown { method, .. } => {
2456 Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
2457 }
2458 _ => Err(Error::JsonRpc(JsonRpcError::method_not_found(
2459 "unknown method",
2460 ))),
2461 }
2462 }
2463
2464 pub fn handle_notification(&self, notification: McpNotification) {
2466 match notification {
2467 McpNotification::Initialized => {
2468 let phase_before = self.session.phase();
2469 if self.session.mark_initialized() {
2470 if phase_before == crate::session::SessionPhase::Uninitialized {
2471 tracing::info!(
2472 "Session initialized from uninitialized state (race resolved)"
2473 );
2474 } else {
2475 tracing::info!("Session initialized, entering operation phase");
2476 }
2477 } else {
2478 tracing::warn!(
2479 phase = ?self.session.phase(),
2480 "Received initialized notification in unexpected state"
2481 );
2482 }
2483 }
2484 McpNotification::Cancelled(params) => {
2485 if let Some(ref request_id) = params.request_id {
2486 if self.cancel_request(request_id) {
2487 tracing::info!(
2488 request_id = ?request_id,
2489 reason = ?params.reason,
2490 "Request cancelled"
2491 );
2492 } else {
2493 tracing::debug!(
2494 request_id = ?request_id,
2495 reason = ?params.reason,
2496 "Cancellation requested for unknown request"
2497 );
2498 }
2499 } else {
2500 tracing::debug!(
2501 reason = ?params.reason,
2502 "Cancellation notification received without request_id"
2503 );
2504 }
2505 }
2506 McpNotification::Progress(params) => {
2507 tracing::trace!(
2508 token = ?params.progress_token,
2509 progress = params.progress,
2510 total = ?params.total,
2511 "Progress notification"
2512 );
2513 }
2515 McpNotification::RootsListChanged => {
2516 tracing::info!("Client roots list changed");
2517 }
2520 McpNotification::Unknown { method, .. } => {
2521 tracing::debug!(method = %method, "Unknown notification received");
2522 }
2523 _ => {
2524 tracing::debug!("Unrecognized notification variant received");
2525 }
2526 }
2527 }
2528}
2529
2530impl Default for McpRouter {
2531 fn default() -> Self {
2532 Self::new()
2533 }
2534}
2535
2536pub use crate::context::Extensions;
2542
2543#[derive(Debug, Clone)]
2568pub struct ToolAnnotationsMap {
2569 map: Arc<HashMap<String, ToolAnnotations>>,
2570}
2571
2572impl ToolAnnotationsMap {
2573 pub fn get(&self, tool_name: &str) -> Option<&ToolAnnotations> {
2577 self.map.get(tool_name)
2578 }
2579
2580 pub fn is_read_only(&self, tool_name: &str) -> bool {
2585 self.map.get(tool_name).is_some_and(|a| a.read_only_hint)
2586 }
2587
2588 pub fn is_destructive(&self, tool_name: &str) -> bool {
2593 self.map.get(tool_name).is_none_or(|a| a.destructive_hint)
2594 }
2595
2596 pub fn is_idempotent(&self, tool_name: &str) -> bool {
2601 self.map.get(tool_name).is_some_and(|a| a.idempotent_hint)
2602 }
2603}
2604
2605#[derive(Debug, Clone)]
2627pub struct RouterRequest {
2628 pub id: RequestId,
2630 pub inner: McpRequest,
2632 pub extensions: Extensions,
2634}
2635
2636impl RouterRequest {
2637 pub fn new(id: RequestId, inner: McpRequest) -> Self {
2639 Self {
2640 id,
2641 inner,
2642 extensions: Extensions::new(),
2643 }
2644 }
2645
2646 pub fn with_inner(self, inner: McpRequest) -> Self {
2652 Self {
2653 id: self.id,
2654 inner,
2655 extensions: self.extensions,
2656 }
2657 }
2658
2659 pub fn with_id_and_inner(self, id: RequestId, inner: McpRequest) -> Self {
2665 Self {
2666 id,
2667 inner,
2668 extensions: self.extensions,
2669 }
2670 }
2671
2672 pub fn clone_with_inner(&self, inner: McpRequest) -> Self {
2680 Self {
2681 id: self.id.clone(),
2682 inner,
2683 extensions: self.extensions.clone(),
2684 }
2685 }
2686}
2687
2688#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
2690pub struct RouterResponse {
2691 pub id: RequestId,
2693 pub inner: std::result::Result<McpResponse, JsonRpcError>,
2695}
2696
2697impl RouterResponse {
2698 pub fn is_error(&self) -> bool {
2714 self.inner.is_err()
2715 }
2716
2717 pub fn into_jsonrpc(self) -> JsonRpcResponse {
2719 match self.inner {
2720 Ok(response) => match serde_json::to_value(response) {
2721 Ok(result) => JsonRpcResponse::result(self.id, result),
2722 Err(e) => {
2723 tracing::error!(error = %e, "Failed to serialize response");
2724 JsonRpcResponse::error(
2725 Some(self.id),
2726 JsonRpcError::internal_error(format!("Serialization error: {}", e)),
2727 )
2728 }
2729 },
2730 Err(error) => JsonRpcResponse::error(Some(self.id), error),
2731 }
2732 }
2733}
2734
2735impl Service<RouterRequest> for McpRouter {
2736 type Response = RouterResponse;
2737 type Error = std::convert::Infallible; type Future =
2739 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
2740
2741 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
2742 Poll::Ready(Ok(()))
2743 }
2744
2745 fn call(&mut self, req: RouterRequest) -> Self::Future {
2746 let router = self.clone();
2747 let request_id = req.id.clone();
2748 Box::pin(async move {
2749 let result = router.handle(req.id, req.inner, req.extensions).await;
2750 router.complete_request(&request_id);
2752 Ok(RouterResponse {
2753 id: request_id,
2754 inner: result.map_err(|e| match e {
2759 Error::JsonRpc(err) => err,
2760 Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
2761 e => JsonRpcError::internal_error(e.to_string()),
2762 }),
2763 })
2764 })
2765 }
2766}
2767
2768#[cfg(test)]
2769mod tests {
2770 use super::*;
2771 use crate::extract::{Context, Json};
2772 use crate::jsonrpc::JsonRpcService;
2773 use crate::tool::ToolBuilder;
2774 use schemars::JsonSchema;
2775 use serde::Deserialize;
2776 use tower::ServiceExt;
2777
2778 #[derive(Debug, Deserialize, JsonSchema)]
2779 struct AddInput {
2780 a: i64,
2781 b: i64,
2782 }
2783
2784 async fn init_router(router: &mut McpRouter) {
2786 let init_req = RouterRequest {
2788 id: RequestId::Number(0),
2789 inner: McpRequest::Initialize(InitializeParams {
2790 protocol_version: "2025-11-25".to_string(),
2791 capabilities: ClientCapabilities {
2792 roots: None,
2793 sampling: None,
2794 elicitation: None,
2795 tasks: None,
2796 experimental: None,
2797 extensions: None,
2798 },
2799 client_info: Implementation {
2800 name: "test".to_string(),
2801 version: "1.0".to_string(),
2802 ..Default::default()
2803 },
2804 meta: None,
2805 }),
2806 extensions: Extensions::new(),
2807 };
2808 let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
2809 router.handle_notification(McpNotification::Initialized);
2811 }
2812
2813 #[tokio::test]
2814 async fn test_router_list_tools() {
2815 let add_tool = ToolBuilder::new("add")
2816 .description("Add two numbers")
2817 .handler(|input: AddInput| async move {
2818 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2819 })
2820 .build();
2821
2822 let mut router = McpRouter::new().tool(add_tool);
2823
2824 init_router(&mut router).await;
2826
2827 let req = RouterRequest {
2828 id: RequestId::Number(1),
2829 inner: McpRequest::ListTools(ListToolsParams::default()),
2830 extensions: Extensions::new(),
2831 };
2832
2833 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2834
2835 match resp.inner {
2836 Ok(McpResponse::ListTools(result)) => {
2837 assert_eq!(result.tools.len(), 1);
2838 assert_eq!(result.tools[0].name, "add");
2839 }
2840 _ => panic!("Expected ListTools response"),
2841 }
2842 }
2843
2844 #[tokio::test]
2845 async fn test_router_call_tool() {
2846 let add_tool = ToolBuilder::new("add")
2847 .description("Add two numbers")
2848 .handler(|input: AddInput| async move {
2849 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2850 })
2851 .build();
2852
2853 let mut router = McpRouter::new().tool(add_tool);
2854
2855 init_router(&mut router).await;
2857
2858 let req = RouterRequest {
2859 id: RequestId::Number(1),
2860 inner: McpRequest::CallTool(CallToolParams {
2861 name: "add".to_string(),
2862 arguments: serde_json::json!({"a": 2, "b": 3}),
2863 meta: None,
2864 task: None,
2865 }),
2866 extensions: Extensions::new(),
2867 };
2868
2869 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2870
2871 match resp.inner {
2872 Ok(McpResponse::CallTool(result)) => {
2873 assert!(!result.is_error);
2874 match &result.content[0] {
2876 Content::Text { text, .. } => assert_eq!(text, "5"),
2877 _ => panic!("Expected text content"),
2878 }
2879 }
2880 _ => panic!("Expected CallTool response"),
2881 }
2882 }
2883
2884 async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2886 let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2887 "protocolVersion": "2025-11-25",
2888 "capabilities": {},
2889 "clientInfo": { "name": "test", "version": "1.0" }
2890 }));
2891 let _ = service.call_single(init_req).await.unwrap();
2892 router.handle_notification(McpNotification::Initialized);
2893 }
2894
2895 #[tokio::test]
2896 async fn test_jsonrpc_service() {
2897 let add_tool = ToolBuilder::new("add")
2898 .description("Add two numbers")
2899 .handler(|input: AddInput| async move {
2900 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2901 })
2902 .build();
2903
2904 let router = McpRouter::new().tool(add_tool);
2905 let mut service = JsonRpcService::new(router.clone());
2906
2907 init_jsonrpc_service(&mut service, &router).await;
2909
2910 let req = JsonRpcRequest::new(1, "tools/list");
2911
2912 let resp = service.call_single(req).await.unwrap();
2913
2914 match resp {
2915 JsonRpcResponse::Result(r) => {
2916 assert_eq!(r.id, RequestId::Number(1));
2917 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2918 assert_eq!(tools.len(), 1);
2919 }
2920 JsonRpcResponse::Error(_) => panic!("Expected success response"),
2921 _ => panic!("unexpected response variant"),
2922 }
2923 }
2924
2925 #[tokio::test]
2926 async fn test_batch_request() {
2927 let add_tool = ToolBuilder::new("add")
2928 .description("Add two numbers")
2929 .handler(|input: AddInput| async move {
2930 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2931 })
2932 .build();
2933
2934 let router = McpRouter::new().tool(add_tool);
2935 let mut service = JsonRpcService::new(router.clone());
2936
2937 init_jsonrpc_service(&mut service, &router).await;
2939
2940 let requests = vec![
2942 JsonRpcRequest::new(1, "tools/list"),
2943 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2944 "name": "add",
2945 "arguments": {"a": 10, "b": 20}
2946 })),
2947 JsonRpcRequest::new(3, "ping"),
2948 ];
2949
2950 let responses = service.call_batch(requests).await.unwrap();
2951
2952 assert_eq!(responses.len(), 3);
2953
2954 match &responses[0] {
2956 JsonRpcResponse::Result(r) => {
2957 assert_eq!(r.id, RequestId::Number(1));
2958 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2959 assert_eq!(tools.len(), 1);
2960 }
2961 JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2962 _ => panic!("unexpected response variant"),
2963 }
2964
2965 match &responses[1] {
2967 JsonRpcResponse::Result(r) => {
2968 assert_eq!(r.id, RequestId::Number(2));
2969 let content = r.result.get("content").unwrap().as_array().unwrap();
2970 let text = content[0].get("text").unwrap().as_str().unwrap();
2971 assert_eq!(text, "30");
2972 }
2973 JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2974 _ => panic!("unexpected response variant"),
2975 }
2976
2977 match &responses[2] {
2979 JsonRpcResponse::Result(r) => {
2980 assert_eq!(r.id, RequestId::Number(3));
2981 }
2982 JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2983 _ => panic!("unexpected response variant"),
2984 }
2985 }
2986
2987 #[tokio::test]
2988 async fn test_empty_batch_error() {
2989 let router = McpRouter::new();
2990 let mut service = JsonRpcService::new(router);
2991
2992 let result = service.call_batch(vec![]).await;
2993 assert!(result.is_err());
2994 }
2995
2996 #[tokio::test]
3001 async fn test_progress_token_extraction() {
3002 use crate::context::{ServerNotification, notification_channel};
3003 use crate::protocol::ProgressToken;
3004 use std::sync::Arc;
3005 use std::sync::atomic::{AtomicBool, Ordering};
3006
3007 let progress_reported = Arc::new(AtomicBool::new(false));
3009 let progress_ref = progress_reported.clone();
3010
3011 let tool = ToolBuilder::new("progress_tool")
3013 .description("Tool that reports progress")
3014 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
3015 let reported = progress_ref.clone();
3016 async move {
3017 ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
3019 .await;
3020 reported.store(true, Ordering::SeqCst);
3021 Ok(CallToolResult::text("done"))
3022 }
3023 })
3024 .build();
3025
3026 let (tx, mut rx) = notification_channel(10);
3028 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
3029 let mut service = JsonRpcService::new(router.clone());
3030
3031 init_jsonrpc_service(&mut service, &router).await;
3033
3034 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3036 "name": "progress_tool",
3037 "arguments": {"a": 1, "b": 2},
3038 "_meta": {
3039 "progressToken": "test-token-123"
3040 }
3041 }));
3042
3043 let resp = service.call_single(req).await.unwrap();
3044
3045 match resp {
3047 JsonRpcResponse::Result(_) => {}
3048 JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
3049 _ => panic!("unexpected response variant"),
3050 }
3051
3052 assert!(progress_reported.load(Ordering::SeqCst));
3054
3055 let notification = rx.try_recv().expect("Expected progress notification");
3057 match notification {
3058 ServerNotification::Progress(params) => {
3059 assert_eq!(
3060 params.progress_token,
3061 ProgressToken::String("test-token-123".to_string())
3062 );
3063 assert_eq!(params.progress, 50.0);
3064 assert_eq!(params.total, Some(100.0));
3065 assert_eq!(params.message.as_deref(), Some("Halfway"));
3066 }
3067 _ => panic!("Expected Progress notification"),
3068 }
3069 }
3070
3071 #[tokio::test]
3072 async fn test_tool_call_without_progress_token() {
3073 use crate::context::notification_channel;
3074 use std::sync::Arc;
3075 use std::sync::atomic::{AtomicBool, Ordering};
3076
3077 let progress_attempted = Arc::new(AtomicBool::new(false));
3078 let progress_ref = progress_attempted.clone();
3079
3080 let tool = ToolBuilder::new("no_token_tool")
3081 .description("Tool that tries to report progress without token")
3082 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
3083 let attempted = progress_ref.clone();
3084 async move {
3085 ctx.report_progress(50.0, Some(100.0), None).await;
3087 attempted.store(true, Ordering::SeqCst);
3088 Ok(CallToolResult::text("done"))
3089 }
3090 })
3091 .build();
3092
3093 let (tx, mut rx) = notification_channel(10);
3094 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
3095 let mut service = JsonRpcService::new(router.clone());
3096
3097 init_jsonrpc_service(&mut service, &router).await;
3098
3099 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3101 "name": "no_token_tool",
3102 "arguments": {"a": 1, "b": 2}
3103 }));
3104
3105 let resp = service.call_single(req).await.unwrap();
3106 assert!(matches!(resp, JsonRpcResponse::Result(_)));
3107
3108 assert!(progress_attempted.load(Ordering::SeqCst));
3110
3111 assert!(rx.try_recv().is_err());
3113 }
3114
3115 #[tokio::test]
3116 async fn test_batch_errors_returned_not_dropped() {
3117 let add_tool = ToolBuilder::new("add")
3118 .description("Add two numbers")
3119 .handler(|input: AddInput| async move {
3120 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3121 })
3122 .build();
3123
3124 let router = McpRouter::new().tool(add_tool);
3125 let mut service = JsonRpcService::new(router.clone());
3126
3127 init_jsonrpc_service(&mut service, &router).await;
3128
3129 let requests = vec![
3131 JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3133 "name": "add",
3134 "arguments": {"a": 10, "b": 20}
3135 })),
3136 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
3138 "name": "nonexistent_tool",
3139 "arguments": {}
3140 })),
3141 JsonRpcRequest::new(3, "ping"),
3143 ];
3144
3145 let responses = service.call_batch(requests).await.unwrap();
3146
3147 assert_eq!(responses.len(), 3);
3149
3150 match &responses[0] {
3152 JsonRpcResponse::Result(r) => {
3153 assert_eq!(r.id, RequestId::Number(1));
3154 }
3155 JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
3156 _ => panic!("unexpected response variant"),
3157 }
3158
3159 match &responses[1] {
3161 JsonRpcResponse::Error(e) => {
3162 assert_eq!(e.id, Some(RequestId::Number(2)));
3163 assert!(e.error.message.contains("not found") || e.error.code == -32601);
3165 }
3166 JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
3167 _ => panic!("unexpected response variant"),
3168 }
3169
3170 match &responses[2] {
3172 JsonRpcResponse::Result(r) => {
3173 assert_eq!(r.id, RequestId::Number(3));
3174 }
3175 JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
3176 _ => panic!("unexpected response variant"),
3177 }
3178 }
3179
3180 #[tokio::test]
3185 async fn test_list_resource_templates() {
3186 use crate::resource::ResourceTemplateBuilder;
3187 use std::collections::HashMap;
3188
3189 let template = ResourceTemplateBuilder::new("file:///{path}")
3190 .name("Project Files")
3191 .description("Access project files")
3192 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3193 Ok(ReadResourceResult {
3194 contents: vec![ResourceContent {
3195 uri,
3196 mime_type: None,
3197 text: None,
3198 blob: None,
3199 meta: None,
3200 }],
3201 meta: None,
3202 })
3203 });
3204
3205 let mut router = McpRouter::new().resource_template(template);
3206
3207 init_router(&mut router).await;
3209
3210 let req = RouterRequest {
3211 id: RequestId::Number(1),
3212 inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
3213 extensions: Extensions::new(),
3214 };
3215
3216 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3217
3218 match resp.inner {
3219 Ok(McpResponse::ListResourceTemplates(result)) => {
3220 assert_eq!(result.resource_templates.len(), 1);
3221 assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
3222 assert_eq!(result.resource_templates[0].name, "Project Files");
3223 }
3224 _ => panic!("Expected ListResourceTemplates response"),
3225 }
3226 }
3227
3228 #[tokio::test]
3229 async fn test_read_resource_via_template() {
3230 use crate::resource::ResourceTemplateBuilder;
3231 use std::collections::HashMap;
3232
3233 let template = ResourceTemplateBuilder::new("db://users/{id}")
3234 .name("User Records")
3235 .handler(|uri: String, vars: HashMap<String, String>| async move {
3236 let id = vars.get("id").unwrap().clone();
3237 Ok(ReadResourceResult {
3238 contents: vec![ResourceContent {
3239 uri,
3240 mime_type: Some("application/json".to_string()),
3241 text: Some(format!(r#"{{"id": "{}"}}"#, id)),
3242 blob: None,
3243 meta: None,
3244 }],
3245 meta: None,
3246 })
3247 });
3248
3249 let mut router = McpRouter::new().resource_template(template);
3250
3251 init_router(&mut router).await;
3253
3254 let req = RouterRequest {
3256 id: RequestId::Number(1),
3257 inner: McpRequest::ReadResource(ReadResourceParams {
3258 uri: "db://users/123".to_string(),
3259 meta: None,
3260 }),
3261 extensions: Extensions::new(),
3262 };
3263
3264 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3265
3266 match resp.inner {
3267 Ok(McpResponse::ReadResource(result)) => {
3268 assert_eq!(result.contents.len(), 1);
3269 assert_eq!(result.contents[0].uri, "db://users/123");
3270 assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
3271 }
3272 _ => panic!("Expected ReadResource response"),
3273 }
3274 }
3275
3276 #[tokio::test]
3277 async fn test_static_resource_takes_precedence_over_template() {
3278 use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
3279 use std::collections::HashMap;
3280
3281 let template = ResourceTemplateBuilder::new("file:///{path}")
3283 .name("Files Template")
3284 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3285 Ok(ReadResourceResult {
3286 contents: vec![ResourceContent {
3287 uri,
3288 mime_type: None,
3289 text: Some("from template".to_string()),
3290 blob: None,
3291 meta: None,
3292 }],
3293 meta: None,
3294 })
3295 });
3296
3297 let static_resource = ResourceBuilder::new("file:///README.md")
3299 .name("README")
3300 .text("from static resource");
3301
3302 let mut router = McpRouter::new()
3303 .resource_template(template)
3304 .resource(static_resource);
3305
3306 init_router(&mut router).await;
3308
3309 let req = RouterRequest {
3311 id: RequestId::Number(1),
3312 inner: McpRequest::ReadResource(ReadResourceParams {
3313 uri: "file:///README.md".to_string(),
3314 meta: None,
3315 }),
3316 extensions: Extensions::new(),
3317 };
3318
3319 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3320
3321 match resp.inner {
3322 Ok(McpResponse::ReadResource(result)) => {
3323 assert_eq!(
3325 result.contents[0].text.as_deref(),
3326 Some("from static resource")
3327 );
3328 }
3329 _ => panic!("Expected ReadResource response"),
3330 }
3331 }
3332
3333 #[tokio::test]
3334 async fn test_resource_not_found_when_no_match() {
3335 use crate::resource::ResourceTemplateBuilder;
3336 use std::collections::HashMap;
3337
3338 let template = ResourceTemplateBuilder::new("db://users/{id}")
3339 .name("Users")
3340 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3341 Ok(ReadResourceResult {
3342 contents: vec![ResourceContent {
3343 uri,
3344 mime_type: None,
3345 text: None,
3346 blob: None,
3347 meta: None,
3348 }],
3349 meta: None,
3350 })
3351 });
3352
3353 let mut router = McpRouter::new().resource_template(template);
3354
3355 init_router(&mut router).await;
3357
3358 let req = RouterRequest {
3360 id: RequestId::Number(1),
3361 inner: McpRequest::ReadResource(ReadResourceParams {
3362 uri: "db://posts/123".to_string(),
3363 meta: None,
3364 }),
3365 extensions: Extensions::new(),
3366 };
3367
3368 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3369
3370 match resp.inner {
3371 Err(err) => {
3372 assert!(err.message.contains("not found"));
3373 }
3374 Ok(_) => panic!("Expected error for non-matching URI"),
3375 }
3376 }
3377
3378 #[tokio::test]
3379 async fn test_capabilities_include_resources_with_only_templates() {
3380 use crate::resource::ResourceTemplateBuilder;
3381 use std::collections::HashMap;
3382
3383 let template = ResourceTemplateBuilder::new("file:///{path}")
3384 .name("Files")
3385 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3386 Ok(ReadResourceResult {
3387 contents: vec![ResourceContent {
3388 uri,
3389 mime_type: None,
3390 text: None,
3391 blob: None,
3392 meta: None,
3393 }],
3394 meta: None,
3395 })
3396 });
3397
3398 let mut router = McpRouter::new().resource_template(template);
3399
3400 let init_req = RouterRequest {
3402 id: RequestId::Number(0),
3403 inner: McpRequest::Initialize(InitializeParams {
3404 protocol_version: "2025-11-25".to_string(),
3405 capabilities: ClientCapabilities {
3406 roots: None,
3407 sampling: None,
3408 elicitation: None,
3409 tasks: None,
3410 experimental: None,
3411 extensions: None,
3412 },
3413 client_info: Implementation {
3414 name: "test".to_string(),
3415 version: "1.0".to_string(),
3416 ..Default::default()
3417 },
3418 meta: None,
3419 }),
3420 extensions: Extensions::new(),
3421 };
3422 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3423
3424 match resp.inner {
3425 Ok(McpResponse::Initialize(result)) => {
3426 assert!(result.capabilities.resources.is_some());
3428 }
3429 _ => panic!("Expected Initialize response"),
3430 }
3431 }
3432
3433 #[tokio::test]
3438 async fn test_log_sends_notification() {
3439 use crate::context::notification_channel;
3440
3441 let (tx, mut rx) = notification_channel(10);
3442 let router = McpRouter::new().with_notification_sender(tx);
3443
3444 let sent = router.log_info("Test message");
3446 assert!(sent);
3447
3448 let notification = rx.try_recv().unwrap();
3450 match notification {
3451 ServerNotification::LogMessage(params) => {
3452 assert_eq!(params.level, LogLevel::Info);
3453 let data = params.data;
3454 assert_eq!(
3455 data.get("message").unwrap().as_str().unwrap(),
3456 "Test message"
3457 );
3458 }
3459 _ => panic!("Expected LogMessage notification"),
3460 }
3461 }
3462
3463 #[tokio::test]
3464 async fn test_log_with_custom_params() {
3465 use crate::context::notification_channel;
3466
3467 let (tx, mut rx) = notification_channel(10);
3468 let router = McpRouter::new().with_notification_sender(tx);
3469
3470 let params = LoggingMessageParams::new(
3472 LogLevel::Error,
3473 serde_json::json!({
3474 "error": "Connection failed",
3475 "host": "localhost"
3476 }),
3477 )
3478 .with_logger("database");
3479
3480 let sent = router.log(params);
3481 assert!(sent);
3482
3483 let notification = rx.try_recv().unwrap();
3484 match notification {
3485 ServerNotification::LogMessage(params) => {
3486 assert_eq!(params.level, LogLevel::Error);
3487 assert_eq!(params.logger.as_deref(), Some("database"));
3488 let data = params.data;
3489 assert_eq!(
3490 data.get("error").unwrap().as_str().unwrap(),
3491 "Connection failed"
3492 );
3493 }
3494 _ => panic!("Expected LogMessage notification"),
3495 }
3496 }
3497
3498 #[tokio::test]
3499 async fn test_log_without_channel_returns_false() {
3500 let router = McpRouter::new();
3502
3503 assert!(!router.log_info("Test"));
3505 assert!(!router.log_warning("Test"));
3506 assert!(!router.log_error("Test"));
3507 assert!(!router.log_debug("Test"));
3508 }
3509
3510 #[tokio::test]
3511 async fn test_logging_capability_with_channel() {
3512 use crate::context::notification_channel;
3513
3514 let (tx, _rx) = notification_channel(10);
3515 let mut router = McpRouter::new().with_notification_sender(tx);
3516
3517 let init_req = RouterRequest {
3519 id: RequestId::Number(0),
3520 inner: McpRequest::Initialize(InitializeParams {
3521 protocol_version: "2025-11-25".to_string(),
3522 capabilities: ClientCapabilities {
3523 roots: None,
3524 sampling: None,
3525 elicitation: None,
3526 tasks: None,
3527 experimental: None,
3528 extensions: None,
3529 },
3530 client_info: Implementation {
3531 name: "test".to_string(),
3532 version: "1.0".to_string(),
3533 ..Default::default()
3534 },
3535 meta: None,
3536 }),
3537 extensions: Extensions::new(),
3538 };
3539 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3540
3541 match resp.inner {
3542 Ok(McpResponse::Initialize(result)) => {
3543 assert!(result.capabilities.logging.is_some());
3545 }
3546 _ => panic!("Expected Initialize response"),
3547 }
3548 }
3549
3550 #[tokio::test]
3551 async fn test_no_logging_capability_without_channel() {
3552 let mut router = McpRouter::new();
3553
3554 let init_req = RouterRequest {
3556 id: RequestId::Number(0),
3557 inner: McpRequest::Initialize(InitializeParams {
3558 protocol_version: "2025-11-25".to_string(),
3559 capabilities: ClientCapabilities {
3560 roots: None,
3561 sampling: None,
3562 elicitation: None,
3563 tasks: None,
3564 experimental: None,
3565 extensions: None,
3566 },
3567 client_info: Implementation {
3568 name: "test".to_string(),
3569 version: "1.0".to_string(),
3570 ..Default::default()
3571 },
3572 meta: None,
3573 }),
3574 extensions: Extensions::new(),
3575 };
3576 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3577
3578 match resp.inner {
3579 Ok(McpResponse::Initialize(result)) => {
3580 assert!(result.capabilities.logging.is_none());
3582 }
3583 _ => panic!("Expected Initialize response"),
3584 }
3585 }
3586
3587 #[tokio::test]
3592 async fn test_create_task_via_call_tool() {
3593 let add_tool = ToolBuilder::new("add")
3594 .description("Add two numbers")
3595 .task_support(TaskSupportMode::Optional)
3596 .handler(|input: AddInput| async move {
3597 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3598 })
3599 .build();
3600
3601 let mut router = McpRouter::new().tool(add_tool);
3602 init_router(&mut router).await;
3603
3604 let req = RouterRequest {
3605 id: RequestId::Number(1),
3606 inner: McpRequest::CallTool(CallToolParams {
3607 name: "add".to_string(),
3608 arguments: serde_json::json!({"a": 5, "b": 10}),
3609 meta: None,
3610 task: Some(TaskRequestParams { ttl: None }),
3611 }),
3612 extensions: Extensions::new(),
3613 };
3614
3615 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3616
3617 match resp.inner {
3618 Ok(McpResponse::CreateTask(result)) => {
3619 assert!(result.task.task_id.starts_with("task-"));
3620 assert_eq!(result.task.status, TaskStatus::Working);
3621 }
3622 _ => panic!("Expected CreateTask response"),
3623 }
3624 }
3625
3626 #[tokio::test]
3627 async fn test_list_tasks_empty() {
3628 let mut router = McpRouter::new();
3629 init_router(&mut router).await;
3630
3631 let req = RouterRequest {
3632 id: RequestId::Number(1),
3633 inner: McpRequest::ListTasks(ListTasksParams::default()),
3634 extensions: Extensions::new(),
3635 };
3636
3637 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3638
3639 match resp.inner {
3640 Ok(McpResponse::ListTasks(result)) => {
3641 assert!(result.tasks.is_empty());
3642 }
3643 _ => panic!("Expected ListTasks response"),
3644 }
3645 }
3646
3647 #[tokio::test]
3648 async fn test_task_lifecycle_complete() {
3649 let add_tool = ToolBuilder::new("add")
3650 .description("Add two numbers")
3651 .task_support(TaskSupportMode::Optional)
3652 .handler(|input: AddInput| async move {
3653 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3654 })
3655 .build();
3656
3657 let mut router = McpRouter::new().tool(add_tool);
3658 init_router(&mut router).await;
3659
3660 let req = RouterRequest {
3662 id: RequestId::Number(1),
3663 inner: McpRequest::CallTool(CallToolParams {
3664 name: "add".to_string(),
3665 arguments: serde_json::json!({"a": 7, "b": 8}),
3666 meta: None,
3667 task: Some(TaskRequestParams { ttl: None }),
3668 }),
3669 extensions: Extensions::new(),
3670 };
3671
3672 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3673 let task_id = match resp.inner {
3674 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3675 _ => panic!("Expected CreateTask response"),
3676 };
3677
3678 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3680
3681 let req = RouterRequest {
3683 id: RequestId::Number(2),
3684 inner: McpRequest::GetTaskResult(GetTaskResultParams {
3685 task_id: task_id.clone(),
3686 meta: None,
3687 }),
3688 extensions: Extensions::new(),
3689 };
3690
3691 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3692
3693 match resp.inner {
3694 Ok(McpResponse::GetTaskResult(result)) => {
3695 assert!(result.meta.is_some());
3697 match &result.content[0] {
3699 Content::Text { text, .. } => assert_eq!(text, "15"),
3700 _ => panic!("Expected text content"),
3701 }
3702 }
3703 _ => panic!("Expected GetTaskResult response"),
3704 }
3705 }
3706
3707 #[tokio::test]
3708 async fn test_task_cancellation() {
3709 let slow_tool = ToolBuilder::new("slow")
3711 .description("Slow tool")
3712 .task_support(TaskSupportMode::Optional)
3713 .handler(|_input: serde_json::Value| async move {
3714 tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
3715 Ok(CallToolResult::text("done"))
3716 })
3717 .build();
3718
3719 let mut router = McpRouter::new().tool(slow_tool);
3720 init_router(&mut router).await;
3721
3722 let req = RouterRequest {
3724 id: RequestId::Number(1),
3725 inner: McpRequest::CallTool(CallToolParams {
3726 name: "slow".to_string(),
3727 arguments: serde_json::json!({}),
3728 meta: None,
3729 task: Some(TaskRequestParams { ttl: None }),
3730 }),
3731 extensions: Extensions::new(),
3732 };
3733
3734 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3735 let task_id = match resp.inner {
3736 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3737 _ => panic!("Expected CreateTask response"),
3738 };
3739
3740 let req = RouterRequest {
3742 id: RequestId::Number(2),
3743 inner: McpRequest::CancelTask(CancelTaskParams {
3744 task_id: task_id.clone(),
3745 reason: Some("Test cancellation".to_string()),
3746 meta: None,
3747 }),
3748 extensions: Extensions::new(),
3749 };
3750
3751 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3752
3753 match resp.inner {
3754 Ok(McpResponse::CancelTask(task_obj)) => {
3755 assert_eq!(task_obj.status, TaskStatus::Cancelled);
3756 }
3757 _ => panic!("Expected CancelTask response"),
3758 }
3759 }
3760
3761 #[tokio::test]
3762 async fn test_get_task_info() {
3763 let add_tool = ToolBuilder::new("add")
3764 .description("Add two numbers")
3765 .task_support(TaskSupportMode::Optional)
3766 .handler(|input: AddInput| async move {
3767 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3768 })
3769 .build();
3770
3771 let mut router = McpRouter::new().tool(add_tool);
3772 init_router(&mut router).await;
3773
3774 let req = RouterRequest {
3776 id: RequestId::Number(1),
3777 inner: McpRequest::CallTool(CallToolParams {
3778 name: "add".to_string(),
3779 arguments: serde_json::json!({"a": 1, "b": 2}),
3780 meta: None,
3781 task: Some(TaskRequestParams { ttl: Some(600_000) }),
3782 }),
3783 extensions: Extensions::new(),
3784 };
3785
3786 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3787 let task_id = match resp.inner {
3788 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3789 _ => panic!("Expected CreateTask response"),
3790 };
3791
3792 let req = RouterRequest {
3794 id: RequestId::Number(2),
3795 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3796 task_id: task_id.clone(),
3797 meta: None,
3798 }),
3799 extensions: Extensions::new(),
3800 };
3801
3802 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3803
3804 match resp.inner {
3805 Ok(McpResponse::GetTaskInfo(info)) => {
3806 assert_eq!(info.task_id, task_id);
3807 assert!(info.created_at.contains('T')); assert_eq!(info.ttl, Some(600_000));
3809 }
3810 _ => panic!("Expected GetTaskInfo response"),
3811 }
3812 }
3813
3814 #[tokio::test]
3815 async fn test_task_forbidden_tool_rejects_task_params() {
3816 let tool = ToolBuilder::new("sync_only")
3817 .description("Sync only tool")
3818 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3819 .build();
3820
3821 let mut router = McpRouter::new().tool(tool);
3822 init_router(&mut router).await;
3823
3824 let req = RouterRequest {
3826 id: RequestId::Number(1),
3827 inner: McpRequest::CallTool(CallToolParams {
3828 name: "sync_only".to_string(),
3829 arguments: serde_json::json!({}),
3830 meta: None,
3831 task: Some(TaskRequestParams { ttl: None }),
3832 }),
3833 extensions: Extensions::new(),
3834 };
3835
3836 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3837
3838 match resp.inner {
3839 Err(e) => {
3840 assert!(e.message.contains("does not support async tasks"));
3841 }
3842 _ => panic!("Expected error response"),
3843 }
3844 }
3845
3846 #[tokio::test]
3847 async fn test_get_nonexistent_task() {
3848 let mut router = McpRouter::new();
3849 init_router(&mut router).await;
3850
3851 let req = RouterRequest {
3852 id: RequestId::Number(1),
3853 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3854 task_id: "task-999".to_string(),
3855 meta: None,
3856 }),
3857 extensions: Extensions::new(),
3858 };
3859
3860 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3861
3862 match resp.inner {
3863 Err(e) => {
3864 assert!(e.message.contains("not found"));
3865 }
3866 _ => panic!("Expected error response"),
3867 }
3868 }
3869
3870 #[tokio::test]
3875 async fn test_subscribe_to_resource() {
3876 use crate::resource::ResourceBuilder;
3877
3878 let resource = ResourceBuilder::new("file:///test.txt")
3879 .name("Test File")
3880 .text("Hello");
3881
3882 let mut router = McpRouter::new().resource(resource);
3883 init_router(&mut router).await;
3884
3885 let req = RouterRequest {
3887 id: RequestId::Number(1),
3888 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3889 uri: "file:///test.txt".to_string(),
3890 meta: None,
3891 }),
3892 extensions: Extensions::new(),
3893 };
3894
3895 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3896
3897 match resp.inner {
3898 Ok(McpResponse::SubscribeResource(_)) => {
3899 assert!(router.is_subscribed("file:///test.txt"));
3901 }
3902 _ => panic!("Expected SubscribeResource response"),
3903 }
3904 }
3905
3906 #[tokio::test]
3907 async fn test_unsubscribe_from_resource() {
3908 use crate::resource::ResourceBuilder;
3909
3910 let resource = ResourceBuilder::new("file:///test.txt")
3911 .name("Test File")
3912 .text("Hello");
3913
3914 let mut router = McpRouter::new().resource(resource);
3915 init_router(&mut router).await;
3916
3917 let req = RouterRequest {
3919 id: RequestId::Number(1),
3920 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3921 uri: "file:///test.txt".to_string(),
3922 meta: None,
3923 }),
3924 extensions: Extensions::new(),
3925 };
3926 let _ = router.ready().await.unwrap().call(req).await.unwrap();
3927 assert!(router.is_subscribed("file:///test.txt"));
3928
3929 let req = RouterRequest {
3931 id: RequestId::Number(2),
3932 inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3933 uri: "file:///test.txt".to_string(),
3934 meta: None,
3935 }),
3936 extensions: Extensions::new(),
3937 };
3938
3939 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3940
3941 match resp.inner {
3942 Ok(McpResponse::UnsubscribeResource(_)) => {
3943 assert!(!router.is_subscribed("file:///test.txt"));
3945 }
3946 _ => panic!("Expected UnsubscribeResource response"),
3947 }
3948 }
3949
3950 #[tokio::test]
3951 async fn test_subscribe_nonexistent_resource() {
3952 let mut router = McpRouter::new();
3953 init_router(&mut router).await;
3954
3955 let req = RouterRequest {
3956 id: RequestId::Number(1),
3957 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3958 uri: "file:///nonexistent.txt".to_string(),
3959 meta: None,
3960 }),
3961 extensions: Extensions::new(),
3962 };
3963
3964 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3965
3966 match resp.inner {
3967 Err(e) => {
3968 assert!(e.message.contains("not found"));
3969 }
3970 _ => panic!("Expected error response"),
3971 }
3972 }
3973
3974 #[tokio::test]
3975 async fn test_notify_resource_updated() {
3976 use crate::context::notification_channel;
3977 use crate::resource::ResourceBuilder;
3978
3979 let (tx, mut rx) = notification_channel(10);
3980
3981 let resource = ResourceBuilder::new("file:///test.txt")
3982 .name("Test File")
3983 .text("Hello");
3984
3985 let router = McpRouter::new()
3986 .resource(resource)
3987 .with_notification_sender(tx);
3988
3989 router.subscribe("file:///test.txt");
3991
3992 let sent = router.notify_resource_updated("file:///test.txt");
3994 assert!(sent);
3995
3996 let notification = rx.try_recv().unwrap();
3998 match notification {
3999 ServerNotification::ResourceUpdated { uri } => {
4000 assert_eq!(uri, "file:///test.txt");
4001 }
4002 _ => panic!("Expected ResourceUpdated notification"),
4003 }
4004 }
4005
4006 #[tokio::test]
4007 async fn test_notify_resource_updated_not_subscribed() {
4008 use crate::context::notification_channel;
4009 use crate::resource::ResourceBuilder;
4010
4011 let (tx, mut rx) = notification_channel(10);
4012
4013 let resource = ResourceBuilder::new("file:///test.txt")
4014 .name("Test File")
4015 .text("Hello");
4016
4017 let router = McpRouter::new()
4018 .resource(resource)
4019 .with_notification_sender(tx);
4020
4021 let sent = router.notify_resource_updated("file:///test.txt");
4023 assert!(!sent); assert!(rx.try_recv().is_err());
4027 }
4028
4029 #[tokio::test]
4030 async fn test_notify_resources_list_changed() {
4031 use crate::context::notification_channel;
4032
4033 let (tx, mut rx) = notification_channel(10);
4034 let router = McpRouter::new().with_notification_sender(tx);
4035
4036 let sent = router.notify_resources_list_changed();
4037 assert!(sent);
4038
4039 let notification = rx.try_recv().unwrap();
4040 match notification {
4041 ServerNotification::ResourcesListChanged => {}
4042 _ => panic!("Expected ResourcesListChanged notification"),
4043 }
4044 }
4045
4046 #[tokio::test]
4047 async fn test_subscribed_uris() {
4048 use crate::resource::ResourceBuilder;
4049
4050 let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
4051
4052 let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
4053
4054 let router = McpRouter::new().resource(resource1).resource(resource2);
4055
4056 router.subscribe("file:///a.txt");
4058 router.subscribe("file:///b.txt");
4059
4060 let uris = router.subscribed_uris();
4061 assert_eq!(uris.len(), 2);
4062 assert!(uris.contains(&"file:///a.txt".to_string()));
4063 assert!(uris.contains(&"file:///b.txt".to_string()));
4064 }
4065
4066 #[tokio::test]
4067 async fn test_subscription_capability_advertised() {
4068 use crate::resource::ResourceBuilder;
4069
4070 let resource = ResourceBuilder::new("file:///test.txt")
4071 .name("Test")
4072 .text("Hello");
4073
4074 let mut router = McpRouter::new().resource(resource);
4075
4076 let init_req = RouterRequest {
4078 id: RequestId::Number(0),
4079 inner: McpRequest::Initialize(InitializeParams {
4080 protocol_version: "2025-11-25".to_string(),
4081 capabilities: ClientCapabilities {
4082 roots: None,
4083 sampling: None,
4084 elicitation: None,
4085 tasks: None,
4086 experimental: None,
4087 extensions: None,
4088 },
4089 client_info: Implementation {
4090 name: "test".to_string(),
4091 version: "1.0".to_string(),
4092 ..Default::default()
4093 },
4094 meta: None,
4095 }),
4096 extensions: Extensions::new(),
4097 };
4098 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
4099
4100 match resp.inner {
4101 Ok(McpResponse::Initialize(result)) => {
4102 let resources_cap = result.capabilities.resources.unwrap();
4104 assert!(resources_cap.subscribe);
4105 }
4106 _ => panic!("Expected Initialize response"),
4107 }
4108 }
4109
4110 #[tokio::test]
4111 async fn test_completion_handler() {
4112 let router = McpRouter::new()
4113 .server_info("test", "1.0")
4114 .completion_handler(|params: CompleteParams| async move {
4115 let prefix = ¶ms.argument.value;
4117 let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
4118 .into_iter()
4119 .filter(|s| s.starts_with(prefix))
4120 .map(String::from)
4121 .collect();
4122 Ok(CompleteResult::new(suggestions))
4123 });
4124
4125 let init_req = RouterRequest {
4127 id: RequestId::Number(0),
4128 inner: McpRequest::Initialize(InitializeParams {
4129 protocol_version: "2025-11-25".to_string(),
4130 capabilities: ClientCapabilities::default(),
4131 client_info: Implementation {
4132 name: "test".to_string(),
4133 version: "1.0".to_string(),
4134 ..Default::default()
4135 },
4136 meta: None,
4137 }),
4138 extensions: Extensions::new(),
4139 };
4140 let resp = router
4141 .clone()
4142 .ready()
4143 .await
4144 .unwrap()
4145 .call(init_req)
4146 .await
4147 .unwrap();
4148
4149 match resp.inner {
4151 Ok(McpResponse::Initialize(result)) => {
4152 assert!(result.capabilities.completions.is_some());
4153 }
4154 _ => panic!("Expected Initialize response"),
4155 }
4156
4157 router.handle_notification(McpNotification::Initialized);
4159
4160 let complete_req = RouterRequest {
4162 id: RequestId::Number(1),
4163 inner: McpRequest::Complete(CompleteParams {
4164 reference: CompletionReference::prompt("test-prompt"),
4165 argument: CompletionArgument::new("query", "al"),
4166 context: None,
4167 meta: None,
4168 }),
4169 extensions: Extensions::new(),
4170 };
4171 let resp = router
4172 .clone()
4173 .ready()
4174 .await
4175 .unwrap()
4176 .call(complete_req)
4177 .await
4178 .unwrap();
4179
4180 match resp.inner {
4181 Ok(McpResponse::Complete(result)) => {
4182 assert_eq!(result.completion.values, vec!["alpha"]);
4183 }
4184 _ => panic!("Expected Complete response"),
4185 }
4186 }
4187
4188 #[tokio::test]
4189 async fn test_completion_without_handler_returns_empty() {
4190 let router = McpRouter::new().server_info("test", "1.0");
4191
4192 let init_req = RouterRequest {
4194 id: RequestId::Number(0),
4195 inner: McpRequest::Initialize(InitializeParams {
4196 protocol_version: "2025-11-25".to_string(),
4197 capabilities: ClientCapabilities::default(),
4198 client_info: Implementation {
4199 name: "test".to_string(),
4200 version: "1.0".to_string(),
4201 ..Default::default()
4202 },
4203 meta: None,
4204 }),
4205 extensions: Extensions::new(),
4206 };
4207 let resp = router
4208 .clone()
4209 .ready()
4210 .await
4211 .unwrap()
4212 .call(init_req)
4213 .await
4214 .unwrap();
4215
4216 match resp.inner {
4218 Ok(McpResponse::Initialize(result)) => {
4219 assert!(result.capabilities.completions.is_none());
4220 }
4221 _ => panic!("Expected Initialize response"),
4222 }
4223
4224 router.handle_notification(McpNotification::Initialized);
4226
4227 let complete_req = RouterRequest {
4229 id: RequestId::Number(1),
4230 inner: McpRequest::Complete(CompleteParams {
4231 reference: CompletionReference::prompt("test-prompt"),
4232 argument: CompletionArgument::new("query", "al"),
4233 context: None,
4234 meta: None,
4235 }),
4236 extensions: Extensions::new(),
4237 };
4238 let resp = router
4239 .clone()
4240 .ready()
4241 .await
4242 .unwrap()
4243 .call(complete_req)
4244 .await
4245 .unwrap();
4246
4247 match resp.inner {
4248 Ok(McpResponse::Complete(result)) => {
4249 assert!(result.completion.values.is_empty());
4250 }
4251 _ => panic!("Expected Complete response"),
4252 }
4253 }
4254
4255 #[tokio::test]
4256 async fn test_tool_filter_list() {
4257 use crate::filter::CapabilityFilter;
4258 use crate::tool::Tool;
4259
4260 let public_tool = ToolBuilder::new("public")
4261 .description("Public tool")
4262 .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
4263 .build();
4264
4265 let admin_tool = ToolBuilder::new("admin")
4266 .description("Admin tool")
4267 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4268 .build();
4269
4270 let mut router = McpRouter::new()
4271 .tool(public_tool)
4272 .tool(admin_tool)
4273 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
4274
4275 init_router(&mut router).await;
4277
4278 let req = RouterRequest {
4279 id: RequestId::Number(1),
4280 inner: McpRequest::ListTools(ListToolsParams::default()),
4281 extensions: Extensions::new(),
4282 };
4283
4284 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4285
4286 match resp.inner {
4287 Ok(McpResponse::ListTools(result)) => {
4288 assert_eq!(result.tools.len(), 1);
4290 assert_eq!(result.tools[0].name, "public");
4291 }
4292 _ => panic!("Expected ListTools response"),
4293 }
4294 }
4295
4296 #[tokio::test]
4297 async fn test_tool_filter_call_denied() {
4298 use crate::filter::CapabilityFilter;
4299 use crate::tool::Tool;
4300
4301 let admin_tool = ToolBuilder::new("admin")
4302 .description("Admin tool")
4303 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4304 .build();
4305
4306 let mut router = McpRouter::new()
4307 .tool(admin_tool)
4308 .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); init_router(&mut router).await;
4312
4313 let req = RouterRequest {
4314 id: RequestId::Number(1),
4315 inner: McpRequest::CallTool(CallToolParams {
4316 name: "admin".to_string(),
4317 arguments: serde_json::json!({"a": 1, "b": 2}),
4318 meta: None,
4319 task: None,
4320 }),
4321 extensions: Extensions::new(),
4322 };
4323
4324 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4325
4326 match resp.inner {
4328 Err(e) => {
4329 assert_eq!(e.code, -32601); }
4331 _ => panic!("Expected JsonRpc error"),
4332 }
4333 }
4334
4335 #[tokio::test]
4336 async fn test_tool_filter_call_allowed() {
4337 use crate::filter::CapabilityFilter;
4338 use crate::tool::Tool;
4339
4340 let public_tool = ToolBuilder::new("public")
4341 .description("Public tool")
4342 .handler(|input: AddInput| async move {
4343 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
4344 })
4345 .build();
4346
4347 let mut router = McpRouter::new()
4348 .tool(public_tool)
4349 .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); init_router(&mut router).await;
4353
4354 let req = RouterRequest {
4355 id: RequestId::Number(1),
4356 inner: McpRequest::CallTool(CallToolParams {
4357 name: "public".to_string(),
4358 arguments: serde_json::json!({"a": 1, "b": 2}),
4359 meta: None,
4360 task: None,
4361 }),
4362 extensions: Extensions::new(),
4363 };
4364
4365 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4366
4367 match resp.inner {
4368 Ok(McpResponse::CallTool(result)) => {
4369 assert!(!result.is_error);
4370 }
4371 _ => panic!("Expected CallTool response"),
4372 }
4373 }
4374
4375 #[tokio::test]
4376 async fn test_tool_filter_custom_denial() {
4377 use crate::filter::{CapabilityFilter, DenialBehavior};
4378 use crate::tool::Tool;
4379
4380 let admin_tool = ToolBuilder::new("admin")
4381 .description("Admin tool")
4382 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4383 .build();
4384
4385 let mut router = McpRouter::new().tool(admin_tool).tool_filter(
4386 CapabilityFilter::new(|_, _: &Tool| false)
4387 .denial_behavior(DenialBehavior::Unauthorized),
4388 );
4389
4390 init_router(&mut router).await;
4392
4393 let req = RouterRequest {
4394 id: RequestId::Number(1),
4395 inner: McpRequest::CallTool(CallToolParams {
4396 name: "admin".to_string(),
4397 arguments: serde_json::json!({"a": 1, "b": 2}),
4398 meta: None,
4399 task: None,
4400 }),
4401 extensions: Extensions::new(),
4402 };
4403
4404 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4405
4406 match resp.inner {
4408 Err(e) => {
4409 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4411 }
4412 _ => panic!("Expected JsonRpc error"),
4413 }
4414 }
4415
4416 #[tokio::test]
4417 async fn test_resource_filter_list() {
4418 use crate::filter::CapabilityFilter;
4419 use crate::resource::{Resource, ResourceBuilder};
4420
4421 let public_resource = ResourceBuilder::new("file:///public.txt")
4422 .name("Public File")
4423 .text("public content");
4424
4425 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4426 .name("Secret File")
4427 .text("secret content");
4428
4429 let mut router = McpRouter::new()
4430 .resource(public_resource)
4431 .resource(secret_resource)
4432 .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
4433 !r.name.contains("Secret")
4434 }));
4435
4436 init_router(&mut router).await;
4438
4439 let req = RouterRequest {
4440 id: RequestId::Number(1),
4441 inner: McpRequest::ListResources(ListResourcesParams::default()),
4442 extensions: Extensions::new(),
4443 };
4444
4445 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4446
4447 match resp.inner {
4448 Ok(McpResponse::ListResources(result)) => {
4449 assert_eq!(result.resources.len(), 1);
4451 assert_eq!(result.resources[0].name, "Public File");
4452 }
4453 _ => panic!("Expected ListResources response"),
4454 }
4455 }
4456
4457 #[tokio::test]
4458 async fn test_resource_filter_read_denied() {
4459 use crate::filter::CapabilityFilter;
4460 use crate::resource::{Resource, ResourceBuilder};
4461
4462 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4463 .name("Secret File")
4464 .text("secret content");
4465
4466 let mut router = McpRouter::new()
4467 .resource(secret_resource)
4468 .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); init_router(&mut router).await;
4472
4473 let req = RouterRequest {
4474 id: RequestId::Number(1),
4475 inner: McpRequest::ReadResource(ReadResourceParams {
4476 uri: "file:///secret.txt".to_string(),
4477 meta: None,
4478 }),
4479 extensions: Extensions::new(),
4480 };
4481
4482 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4483
4484 match resp.inner {
4486 Err(e) => {
4487 assert_eq!(e.code, -32601); }
4489 _ => panic!("Expected JsonRpc error"),
4490 }
4491 }
4492
4493 #[tokio::test]
4494 async fn test_resource_filter_read_allowed() {
4495 use crate::filter::CapabilityFilter;
4496 use crate::resource::{Resource, ResourceBuilder};
4497
4498 let public_resource = ResourceBuilder::new("file:///public.txt")
4499 .name("Public File")
4500 .text("public content");
4501
4502 let mut router = McpRouter::new()
4503 .resource(public_resource)
4504 .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); init_router(&mut router).await;
4508
4509 let req = RouterRequest {
4510 id: RequestId::Number(1),
4511 inner: McpRequest::ReadResource(ReadResourceParams {
4512 uri: "file:///public.txt".to_string(),
4513 meta: None,
4514 }),
4515 extensions: Extensions::new(),
4516 };
4517
4518 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4519
4520 match resp.inner {
4521 Ok(McpResponse::ReadResource(result)) => {
4522 assert_eq!(result.contents.len(), 1);
4523 assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
4524 }
4525 _ => panic!("Expected ReadResource response"),
4526 }
4527 }
4528
4529 #[tokio::test]
4530 async fn test_resource_filter_custom_denial() {
4531 use crate::filter::{CapabilityFilter, DenialBehavior};
4532 use crate::resource::{Resource, ResourceBuilder};
4533
4534 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4535 .name("Secret File")
4536 .text("secret content");
4537
4538 let mut router = McpRouter::new().resource(secret_resource).resource_filter(
4539 CapabilityFilter::new(|_, _: &Resource| false)
4540 .denial_behavior(DenialBehavior::Unauthorized),
4541 );
4542
4543 init_router(&mut router).await;
4545
4546 let req = RouterRequest {
4547 id: RequestId::Number(1),
4548 inner: McpRequest::ReadResource(ReadResourceParams {
4549 uri: "file:///secret.txt".to_string(),
4550 meta: None,
4551 }),
4552 extensions: Extensions::new(),
4553 };
4554
4555 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4556
4557 match resp.inner {
4559 Err(e) => {
4560 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4562 }
4563 _ => panic!("Expected JsonRpc error"),
4564 }
4565 }
4566
4567 #[tokio::test]
4568 async fn test_prompt_filter_list() {
4569 use crate::filter::CapabilityFilter;
4570 use crate::prompt::{Prompt, PromptBuilder};
4571
4572 let public_prompt = PromptBuilder::new("greeting")
4573 .description("A greeting")
4574 .user_message("Hello!");
4575
4576 let admin_prompt = PromptBuilder::new("system_debug")
4577 .description("Admin prompt")
4578 .user_message("Debug");
4579
4580 let mut router = McpRouter::new()
4581 .prompt(public_prompt)
4582 .prompt(admin_prompt)
4583 .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
4584 !p.name.contains("system")
4585 }));
4586
4587 init_router(&mut router).await;
4589
4590 let req = RouterRequest {
4591 id: RequestId::Number(1),
4592 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4593 extensions: Extensions::new(),
4594 };
4595
4596 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4597
4598 match resp.inner {
4599 Ok(McpResponse::ListPrompts(result)) => {
4600 assert_eq!(result.prompts.len(), 1);
4602 assert_eq!(result.prompts[0].name, "greeting");
4603 }
4604 _ => panic!("Expected ListPrompts response"),
4605 }
4606 }
4607
4608 #[tokio::test]
4609 async fn test_prompt_filter_get_denied() {
4610 use crate::filter::CapabilityFilter;
4611 use crate::prompt::{Prompt, PromptBuilder};
4612 use std::collections::HashMap;
4613
4614 let admin_prompt = PromptBuilder::new("system_debug")
4615 .description("Admin prompt")
4616 .user_message("Debug");
4617
4618 let mut router = McpRouter::new()
4619 .prompt(admin_prompt)
4620 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); init_router(&mut router).await;
4624
4625 let req = RouterRequest {
4626 id: RequestId::Number(1),
4627 inner: McpRequest::GetPrompt(GetPromptParams {
4628 name: "system_debug".to_string(),
4629 arguments: HashMap::new(),
4630 meta: None,
4631 }),
4632 extensions: Extensions::new(),
4633 };
4634
4635 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4636
4637 match resp.inner {
4639 Err(e) => {
4640 assert_eq!(e.code, -32601); }
4642 _ => panic!("Expected JsonRpc error"),
4643 }
4644 }
4645
4646 #[tokio::test]
4647 async fn test_prompt_filter_get_allowed() {
4648 use crate::filter::CapabilityFilter;
4649 use crate::prompt::{Prompt, PromptBuilder};
4650 use std::collections::HashMap;
4651
4652 let public_prompt = PromptBuilder::new("greeting")
4653 .description("A greeting")
4654 .user_message("Hello!");
4655
4656 let mut router = McpRouter::new()
4657 .prompt(public_prompt)
4658 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); init_router(&mut router).await;
4662
4663 let req = RouterRequest {
4664 id: RequestId::Number(1),
4665 inner: McpRequest::GetPrompt(GetPromptParams {
4666 name: "greeting".to_string(),
4667 arguments: HashMap::new(),
4668 meta: None,
4669 }),
4670 extensions: Extensions::new(),
4671 };
4672
4673 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4674
4675 match resp.inner {
4676 Ok(McpResponse::GetPrompt(result)) => {
4677 assert_eq!(result.messages.len(), 1);
4678 }
4679 _ => panic!("Expected GetPrompt response"),
4680 }
4681 }
4682
4683 #[tokio::test]
4684 async fn test_prompt_filter_custom_denial() {
4685 use crate::filter::{CapabilityFilter, DenialBehavior};
4686 use crate::prompt::{Prompt, PromptBuilder};
4687 use std::collections::HashMap;
4688
4689 let admin_prompt = PromptBuilder::new("system_debug")
4690 .description("Admin prompt")
4691 .user_message("Debug");
4692
4693 let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
4694 CapabilityFilter::new(|_, _: &Prompt| false)
4695 .denial_behavior(DenialBehavior::Unauthorized),
4696 );
4697
4698 init_router(&mut router).await;
4700
4701 let req = RouterRequest {
4702 id: RequestId::Number(1),
4703 inner: McpRequest::GetPrompt(GetPromptParams {
4704 name: "system_debug".to_string(),
4705 arguments: HashMap::new(),
4706 meta: None,
4707 }),
4708 extensions: Extensions::new(),
4709 };
4710
4711 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4712
4713 match resp.inner {
4715 Err(e) => {
4716 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4718 }
4719 _ => panic!("Expected JsonRpc error"),
4720 }
4721 }
4722
4723 #[derive(Debug, Deserialize, JsonSchema)]
4728 struct StringInput {
4729 value: String,
4730 }
4731
4732 #[tokio::test]
4733 async fn test_router_merge_tools() {
4734 let tool_a = ToolBuilder::new("tool_a")
4736 .description("Tool A")
4737 .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
4738 .build();
4739
4740 let router_a = McpRouter::new().tool(tool_a);
4741
4742 let tool_b = ToolBuilder::new("tool_b")
4744 .description("Tool B")
4745 .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
4746 .build();
4747 let tool_c = ToolBuilder::new("tool_c")
4748 .description("Tool C")
4749 .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
4750 .build();
4751
4752 let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
4753
4754 let mut merged = McpRouter::new()
4756 .server_info("merged", "1.0")
4757 .merge(router_a)
4758 .merge(router_b);
4759
4760 init_router(&mut merged).await;
4761
4762 let req = RouterRequest {
4764 id: RequestId::Number(1),
4765 inner: McpRequest::ListTools(ListToolsParams::default()),
4766 extensions: Extensions::new(),
4767 };
4768
4769 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4770
4771 match resp.inner {
4772 Ok(McpResponse::ListTools(result)) => {
4773 assert_eq!(result.tools.len(), 3);
4774 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4775 assert!(names.contains(&"tool_a"));
4776 assert!(names.contains(&"tool_b"));
4777 assert!(names.contains(&"tool_c"));
4778 }
4779 _ => panic!("Expected ListTools response"),
4780 }
4781 }
4782
4783 #[tokio::test]
4784 async fn test_router_merge_overwrites_duplicates() {
4785 let tool_v1 = ToolBuilder::new("shared")
4787 .description("Version 1")
4788 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
4789 .build();
4790
4791 let router_a = McpRouter::new().tool(tool_v1);
4792
4793 let tool_v2 = ToolBuilder::new("shared")
4795 .description("Version 2")
4796 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
4797 .build();
4798
4799 let router_b = McpRouter::new().tool(tool_v2);
4800
4801 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4803
4804 init_router(&mut merged).await;
4805
4806 let req = RouterRequest {
4807 id: RequestId::Number(1),
4808 inner: McpRequest::ListTools(ListToolsParams::default()),
4809 extensions: Extensions::new(),
4810 };
4811
4812 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4813
4814 match resp.inner {
4815 Ok(McpResponse::ListTools(result)) => {
4816 assert_eq!(result.tools.len(), 1);
4817 assert_eq!(result.tools[0].name, "shared");
4818 assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
4819 }
4820 _ => panic!("Expected ListTools response"),
4821 }
4822 }
4823
4824 #[tokio::test]
4825 async fn test_router_merge_resources() {
4826 use crate::resource::ResourceBuilder;
4827
4828 let router_a = McpRouter::new().resource(
4830 ResourceBuilder::new("file:///a.txt")
4831 .name("File A")
4832 .text("content a"),
4833 );
4834
4835 let router_b = McpRouter::new().resource(
4836 ResourceBuilder::new("file:///b.txt")
4837 .name("File B")
4838 .text("content b"),
4839 );
4840
4841 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4842
4843 init_router(&mut merged).await;
4844
4845 let req = RouterRequest {
4846 id: RequestId::Number(1),
4847 inner: McpRequest::ListResources(ListResourcesParams::default()),
4848 extensions: Extensions::new(),
4849 };
4850
4851 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4852
4853 match resp.inner {
4854 Ok(McpResponse::ListResources(result)) => {
4855 assert_eq!(result.resources.len(), 2);
4856 let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4857 assert!(uris.contains(&"file:///a.txt"));
4858 assert!(uris.contains(&"file:///b.txt"));
4859 }
4860 _ => panic!("Expected ListResources response"),
4861 }
4862 }
4863
4864 #[tokio::test]
4865 async fn test_router_merge_prompts() {
4866 use crate::prompt::PromptBuilder;
4867
4868 let router_a =
4869 McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4870
4871 let router_b =
4872 McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4873
4874 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4875
4876 init_router(&mut merged).await;
4877
4878 let req = RouterRequest {
4879 id: RequestId::Number(1),
4880 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4881 extensions: Extensions::new(),
4882 };
4883
4884 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4885
4886 match resp.inner {
4887 Ok(McpResponse::ListPrompts(result)) => {
4888 assert_eq!(result.prompts.len(), 2);
4889 let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4890 assert!(names.contains(&"prompt_a"));
4891 assert!(names.contains(&"prompt_b"));
4892 }
4893 _ => panic!("Expected ListPrompts response"),
4894 }
4895 }
4896
4897 #[tokio::test]
4898 async fn test_router_nest_prefixes_tools() {
4899 let tool_query = ToolBuilder::new("query")
4901 .description("Query the database")
4902 .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4903 .build();
4904 let tool_insert = ToolBuilder::new("insert")
4905 .description("Insert into database")
4906 .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4907 .build();
4908
4909 let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4910
4911 let mut router = McpRouter::new()
4913 .server_info("nested", "1.0")
4914 .nest("db", db_router);
4915
4916 init_router(&mut router).await;
4917
4918 let req = RouterRequest {
4919 id: RequestId::Number(1),
4920 inner: McpRequest::ListTools(ListToolsParams::default()),
4921 extensions: Extensions::new(),
4922 };
4923
4924 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4925
4926 match resp.inner {
4927 Ok(McpResponse::ListTools(result)) => {
4928 assert_eq!(result.tools.len(), 2);
4929 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4930 assert!(names.contains(&"db.query"));
4931 assert!(names.contains(&"db.insert"));
4932 }
4933 _ => panic!("Expected ListTools response"),
4934 }
4935 }
4936
4937 #[tokio::test]
4938 async fn test_router_nest_call_prefixed_tool() {
4939 let tool = ToolBuilder::new("echo")
4940 .description("Echo input")
4941 .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4942 .build();
4943
4944 let nested_router = McpRouter::new().tool(tool);
4945
4946 let mut router = McpRouter::new().nest("api", nested_router);
4947
4948 init_router(&mut router).await;
4949
4950 let req = RouterRequest {
4952 id: RequestId::Number(1),
4953 inner: McpRequest::CallTool(CallToolParams {
4954 name: "api.echo".to_string(),
4955 arguments: serde_json::json!({"value": "hello world"}),
4956 meta: None,
4957 task: None,
4958 }),
4959 extensions: Extensions::new(),
4960 };
4961
4962 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4963
4964 match resp.inner {
4965 Ok(McpResponse::CallTool(result)) => {
4966 assert!(!result.is_error);
4967 match &result.content[0] {
4968 Content::Text { text, .. } => assert_eq!(text, "hello world"),
4969 _ => panic!("Expected text content"),
4970 }
4971 }
4972 _ => panic!("Expected CallTool response"),
4973 }
4974 }
4975
4976 #[tokio::test]
4977 async fn test_router_multiple_nests() {
4978 let db_tool = ToolBuilder::new("query")
4979 .description("Database query")
4980 .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4981 .build();
4982
4983 let api_tool = ToolBuilder::new("fetch")
4984 .description("API fetch")
4985 .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4986 .build();
4987
4988 let db_router = McpRouter::new().tool(db_tool);
4989 let api_router = McpRouter::new().tool(api_tool);
4990
4991 let mut router = McpRouter::new()
4992 .nest("db", db_router)
4993 .nest("api", api_router);
4994
4995 init_router(&mut router).await;
4996
4997 let req = RouterRequest {
4998 id: RequestId::Number(1),
4999 inner: McpRequest::ListTools(ListToolsParams::default()),
5000 extensions: Extensions::new(),
5001 };
5002
5003 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5004
5005 match resp.inner {
5006 Ok(McpResponse::ListTools(result)) => {
5007 assert_eq!(result.tools.len(), 2);
5008 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5009 assert!(names.contains(&"db.query"));
5010 assert!(names.contains(&"api.fetch"));
5011 }
5012 _ => panic!("Expected ListTools response"),
5013 }
5014 }
5015
5016 #[tokio::test]
5017 async fn test_router_merge_and_nest_combined() {
5018 let tool_a = ToolBuilder::new("local")
5020 .description("Local tool")
5021 .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
5022 .build();
5023
5024 let nested_tool = ToolBuilder::new("remote")
5025 .description("Remote tool")
5026 .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
5027 .build();
5028
5029 let nested_router = McpRouter::new().tool(nested_tool);
5030
5031 let mut router = McpRouter::new()
5032 .tool(tool_a)
5033 .nest("external", nested_router);
5034
5035 init_router(&mut router).await;
5036
5037 let req = RouterRequest {
5038 id: RequestId::Number(1),
5039 inner: McpRequest::ListTools(ListToolsParams::default()),
5040 extensions: Extensions::new(),
5041 };
5042
5043 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5044
5045 match resp.inner {
5046 Ok(McpResponse::ListTools(result)) => {
5047 assert_eq!(result.tools.len(), 2);
5048 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5049 assert!(names.contains(&"local"));
5050 assert!(names.contains(&"external.remote"));
5051 }
5052 _ => panic!("Expected ListTools response"),
5053 }
5054 }
5055
5056 #[tokio::test]
5057 async fn test_router_merge_preserves_server_info() {
5058 let child_router = McpRouter::new()
5059 .server_info("child", "2.0")
5060 .instructions("Child instructions");
5061
5062 let mut router = McpRouter::new()
5063 .server_info("parent", "1.0")
5064 .instructions("Parent instructions")
5065 .merge(child_router);
5066
5067 init_router(&mut router).await;
5068
5069 let init_req = RouterRequest {
5071 id: RequestId::Number(99),
5072 inner: McpRequest::Initialize(InitializeParams {
5073 protocol_version: "2025-11-25".to_string(),
5074 capabilities: ClientCapabilities::default(),
5075 client_info: Implementation {
5076 name: "test".to_string(),
5077 version: "1.0".to_string(),
5078 ..Default::default()
5079 },
5080 meta: None,
5081 }),
5082 extensions: Extensions::new(),
5083 };
5084
5085 let child_router2 = McpRouter::new().server_info("child", "2.0");
5087 let mut fresh_router = McpRouter::new()
5088 .server_info("parent", "1.0")
5089 .merge(child_router2);
5090
5091 let resp = fresh_router
5092 .ready()
5093 .await
5094 .unwrap()
5095 .call(init_req)
5096 .await
5097 .unwrap();
5098
5099 match resp.inner {
5100 Ok(McpResponse::Initialize(result)) => {
5101 assert_eq!(result.server_info.name, "parent");
5102 assert_eq!(result.server_info.version, "1.0");
5103 }
5104 _ => panic!("Expected Initialize response"),
5105 }
5106 }
5107
5108 #[tokio::test]
5113 async fn test_auto_instructions_tools_only() {
5114 let tool_a = ToolBuilder::new("alpha")
5115 .description("Alpha tool")
5116 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5117 .build();
5118 let tool_b = ToolBuilder::new("beta")
5119 .description("Beta tool")
5120 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5121 .build();
5122
5123 let mut router = McpRouter::new()
5124 .auto_instructions()
5125 .tool(tool_a)
5126 .tool(tool_b);
5127
5128 let resp = send_initialize(&mut router).await;
5129 let instructions = resp.instructions.expect("should have instructions");
5130
5131 assert!(instructions.contains("## Tools"));
5132 assert!(instructions.contains("- **alpha**: Alpha tool"));
5133 assert!(instructions.contains("- **beta**: Beta tool"));
5134 assert!(!instructions.contains("## Resources"));
5136 assert!(!instructions.contains("## Prompts"));
5137 }
5138
5139 #[tokio::test]
5140 async fn test_auto_instructions_with_annotations() {
5141 let read_only_tool = ToolBuilder::new("query")
5142 .description("Run a query")
5143 .read_only()
5144 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5145 .build();
5146 let destructive_tool = ToolBuilder::new("delete")
5147 .description("Delete a record")
5148 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5149 .build();
5150 let idempotent_tool = ToolBuilder::new("upsert")
5151 .description("Upsert a record")
5152 .non_destructive()
5153 .idempotent()
5154 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5155 .build();
5156
5157 let mut router = McpRouter::new()
5158 .auto_instructions()
5159 .tool(read_only_tool)
5160 .tool(destructive_tool)
5161 .tool(idempotent_tool);
5162
5163 let resp = send_initialize(&mut router).await;
5164 let instructions = resp.instructions.unwrap();
5165
5166 assert!(instructions.contains("- **query**: Run a query [read-only]"));
5167 assert!(instructions.contains("- **delete**: Delete a record\n"));
5169 assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
5170 }
5171
5172 #[tokio::test]
5173 async fn test_auto_instructions_with_resources() {
5174 use crate::resource::ResourceBuilder;
5175
5176 let resource = ResourceBuilder::new("file:///schema.sql")
5177 .name("Schema")
5178 .description("Database schema")
5179 .text("CREATE TABLE ...");
5180
5181 let mut router = McpRouter::new().auto_instructions().resource(resource);
5182
5183 let resp = send_initialize(&mut router).await;
5184 let instructions = resp.instructions.unwrap();
5185
5186 assert!(instructions.contains("## Resources"));
5187 assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
5188 assert!(!instructions.contains("## Tools"));
5189 }
5190
5191 #[tokio::test]
5192 async fn test_auto_instructions_with_resource_templates() {
5193 use crate::resource::ResourceTemplateBuilder;
5194
5195 let template = ResourceTemplateBuilder::new("file:///{path}")
5196 .name("File")
5197 .description("Read a file by path")
5198 .handler(
5199 |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
5200 Ok(crate::ReadResourceResult::text("content", "text/plain"))
5201 },
5202 );
5203
5204 let mut router = McpRouter::new()
5205 .auto_instructions()
5206 .resource_template(template);
5207
5208 let resp = send_initialize(&mut router).await;
5209 let instructions = resp.instructions.unwrap();
5210
5211 assert!(instructions.contains("## Resources"));
5212 assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
5213 }
5214
5215 #[tokio::test]
5216 async fn test_auto_instructions_with_prompts() {
5217 use crate::prompt::PromptBuilder;
5218
5219 let prompt = PromptBuilder::new("write_query")
5220 .description("Help write a SQL query")
5221 .user_message("Write a query for: {task}");
5222
5223 let mut router = McpRouter::new().auto_instructions().prompt(prompt);
5224
5225 let resp = send_initialize(&mut router).await;
5226 let instructions = resp.instructions.unwrap();
5227
5228 assert!(instructions.contains("## Prompts"));
5229 assert!(instructions.contains("- **write_query**: Help write a SQL query"));
5230 assert!(!instructions.contains("## Tools"));
5231 }
5232
5233 #[tokio::test]
5234 async fn test_auto_instructions_all_sections() {
5235 use crate::prompt::PromptBuilder;
5236 use crate::resource::ResourceBuilder;
5237
5238 let tool = ToolBuilder::new("query")
5239 .description("Execute SQL")
5240 .read_only()
5241 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5242 .build();
5243 let resource = ResourceBuilder::new("db://schema")
5244 .name("Schema")
5245 .description("Full database schema")
5246 .text("schema");
5247 let prompt = PromptBuilder::new("write_query")
5248 .description("Help write a SQL query")
5249 .user_message("Write a query");
5250
5251 let mut router = McpRouter::new()
5252 .auto_instructions()
5253 .tool(tool)
5254 .resource(resource)
5255 .prompt(prompt);
5256
5257 let resp = send_initialize(&mut router).await;
5258 let instructions = resp.instructions.unwrap();
5259
5260 assert!(instructions.contains("## Tools"));
5262 assert!(instructions.contains("## Resources"));
5263 assert!(instructions.contains("## Prompts"));
5264
5265 let tools_pos = instructions.find("## Tools").unwrap();
5267 let resources_pos = instructions.find("## Resources").unwrap();
5268 let prompts_pos = instructions.find("## Prompts").unwrap();
5269 assert!(tools_pos < resources_pos);
5270 assert!(resources_pos < prompts_pos);
5271 }
5272
5273 #[tokio::test]
5274 async fn test_auto_instructions_with_prefix_and_suffix() {
5275 let tool = ToolBuilder::new("echo")
5276 .description("Echo input")
5277 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5278 .build();
5279
5280 let mut router = McpRouter::new()
5281 .auto_instructions_with(
5282 Some("This server provides echo capabilities."),
5283 Some("Contact admin@example.com for support."),
5284 )
5285 .tool(tool);
5286
5287 let resp = send_initialize(&mut router).await;
5288 let instructions = resp.instructions.unwrap();
5289
5290 assert!(instructions.starts_with("This server provides echo capabilities."));
5291 assert!(instructions.ends_with("Contact admin@example.com for support."));
5292 assert!(instructions.contains("## Tools"));
5293 assert!(instructions.contains("- **echo**: Echo input"));
5294 }
5295
5296 #[tokio::test]
5297 async fn test_auto_instructions_prefix_only() {
5298 let tool = ToolBuilder::new("echo")
5299 .description("Echo input")
5300 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5301 .build();
5302
5303 let mut router = McpRouter::new()
5304 .auto_instructions_with(Some("My server intro."), None::<String>)
5305 .tool(tool);
5306
5307 let resp = send_initialize(&mut router).await;
5308 let instructions = resp.instructions.unwrap();
5309
5310 assert!(instructions.starts_with("My server intro."));
5311 assert!(instructions.contains("- **echo**: Echo input"));
5312 }
5313
5314 #[tokio::test]
5315 async fn test_auto_instructions_empty_router() {
5316 let mut router = McpRouter::new().auto_instructions();
5317
5318 let resp = send_initialize(&mut router).await;
5319 let instructions = resp.instructions.expect("should have instructions");
5320
5321 assert!(!instructions.contains("## Tools"));
5323 assert!(!instructions.contains("## Resources"));
5324 assert!(!instructions.contains("## Prompts"));
5325 assert!(instructions.is_empty());
5326 }
5327
5328 #[tokio::test]
5329 async fn test_auto_instructions_overrides_manual() {
5330 let tool = ToolBuilder::new("echo")
5331 .description("Echo input")
5332 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5333 .build();
5334
5335 let mut router = McpRouter::new()
5336 .instructions("This will be overridden")
5337 .auto_instructions()
5338 .tool(tool);
5339
5340 let resp = send_initialize(&mut router).await;
5341 let instructions = resp.instructions.unwrap();
5342
5343 assert!(!instructions.contains("This will be overridden"));
5344 assert!(instructions.contains("- **echo**: Echo input"));
5345 }
5346
5347 #[tokio::test]
5348 async fn test_no_auto_instructions_returns_manual() {
5349 let tool = ToolBuilder::new("echo")
5350 .description("Echo input")
5351 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5352 .build();
5353
5354 let mut router = McpRouter::new()
5355 .instructions("Manual instructions here")
5356 .tool(tool);
5357
5358 let resp = send_initialize(&mut router).await;
5359 let instructions = resp.instructions.unwrap();
5360
5361 assert_eq!(instructions, "Manual instructions here");
5362 }
5363
5364 #[tokio::test]
5365 async fn test_auto_instructions_no_description_fallback() {
5366 let tool = ToolBuilder::new("mystery")
5367 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5368 .build();
5369
5370 let mut router = McpRouter::new().auto_instructions().tool(tool);
5371
5372 let resp = send_initialize(&mut router).await;
5373 let instructions = resp.instructions.unwrap();
5374
5375 assert!(instructions.contains("- **mystery**: No description"));
5376 }
5377
5378 #[tokio::test]
5379 async fn test_auto_instructions_sorted_alphabetically() {
5380 let tool_z = ToolBuilder::new("zebra")
5381 .description("Z tool")
5382 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5383 .build();
5384 let tool_a = ToolBuilder::new("alpha")
5385 .description("A tool")
5386 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5387 .build();
5388 let tool_m = ToolBuilder::new("middle")
5389 .description("M tool")
5390 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5391 .build();
5392
5393 let mut router = McpRouter::new()
5394 .auto_instructions()
5395 .tool(tool_z)
5396 .tool(tool_a)
5397 .tool(tool_m);
5398
5399 let resp = send_initialize(&mut router).await;
5400 let instructions = resp.instructions.unwrap();
5401
5402 let alpha_pos = instructions.find("**alpha**").unwrap();
5403 let middle_pos = instructions.find("**middle**").unwrap();
5404 let zebra_pos = instructions.find("**zebra**").unwrap();
5405 assert!(alpha_pos < middle_pos);
5406 assert!(middle_pos < zebra_pos);
5407 }
5408
5409 #[tokio::test]
5410 async fn test_auto_instructions_read_only_and_idempotent_tags() {
5411 let tool = ToolBuilder::new("safe_update")
5412 .description("Safe update operation")
5413 .idempotent()
5414 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5415 .build();
5416
5417 let mut router = McpRouter::new().auto_instructions().tool(tool);
5418
5419 let resp = send_initialize(&mut router).await;
5420 let instructions = resp.instructions.unwrap();
5421
5422 assert!(
5423 instructions.contains("[idempotent]"),
5424 "got: {}",
5425 instructions
5426 );
5427 }
5428
5429 #[tokio::test]
5430 async fn test_auto_instructions_lazy_generation() {
5431 let mut router = McpRouter::new().auto_instructions();
5434
5435 let tool = ToolBuilder::new("late_tool")
5436 .description("Added after auto_instructions")
5437 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5438 .build();
5439
5440 router = router.tool(tool);
5441
5442 let resp = send_initialize(&mut router).await;
5443 let instructions = resp.instructions.unwrap();
5444
5445 assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
5446 }
5447
5448 #[tokio::test]
5449 async fn test_auto_instructions_multiple_annotation_tags() {
5450 let tool = ToolBuilder::new("update")
5451 .description("Update a record")
5452 .annotations(ToolAnnotations {
5453 read_only_hint: true,
5454 idempotent_hint: true,
5455 ..Default::default()
5456 })
5457 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5458 .build();
5459
5460 let mut router = McpRouter::new().auto_instructions().tool(tool);
5461
5462 let resp = send_initialize(&mut router).await;
5463 let instructions = resp.instructions.unwrap();
5464
5465 assert!(
5466 instructions.contains("[read-only, idempotent]"),
5467 "got: {}",
5468 instructions
5469 );
5470 }
5471
5472 #[tokio::test]
5473 async fn test_auto_instructions_no_annotations_no_tags() {
5474 let tool = ToolBuilder::new("fetch")
5476 .description("Fetch data")
5477 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5478 .build();
5479
5480 let mut router = McpRouter::new().auto_instructions().tool(tool);
5481
5482 let resp = send_initialize(&mut router).await;
5483 let instructions = resp.instructions.unwrap();
5484
5485 assert!(
5487 !instructions.contains('['),
5488 "should have no tags, got: {}",
5489 instructions
5490 );
5491 assert!(instructions.contains("- **fetch**: Fetch data"));
5492 }
5493
5494 async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
5496 let init_req = RouterRequest {
5497 id: RequestId::Number(0),
5498 inner: McpRequest::Initialize(InitializeParams {
5499 protocol_version: "2025-11-25".to_string(),
5500 capabilities: ClientCapabilities {
5501 roots: None,
5502 sampling: None,
5503 elicitation: None,
5504 tasks: None,
5505 experimental: None,
5506 extensions: None,
5507 },
5508 client_info: Implementation {
5509 name: "test".to_string(),
5510 version: "1.0".to_string(),
5511 ..Default::default()
5512 },
5513 meta: None,
5514 }),
5515 extensions: Extensions::new(),
5516 };
5517 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5518 match resp.inner {
5519 Ok(McpResponse::Initialize(result)) => result,
5520 other => panic!("Expected Initialize response, got {:?}", other),
5521 }
5522 }
5523
5524 #[tokio::test]
5525 async fn test_notify_tools_list_changed() {
5526 let (tx, mut rx) = crate::context::notification_channel(16);
5527
5528 let router = McpRouter::new()
5529 .server_info("test", "1.0")
5530 .with_notification_sender(tx);
5531
5532 assert!(router.notify_tools_list_changed());
5533
5534 let notification = rx.recv().await.unwrap();
5535 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5536 }
5537
5538 #[tokio::test]
5539 async fn test_notify_prompts_list_changed() {
5540 let (tx, mut rx) = crate::context::notification_channel(16);
5541
5542 let router = McpRouter::new()
5543 .server_info("test", "1.0")
5544 .with_notification_sender(tx);
5545
5546 assert!(router.notify_prompts_list_changed());
5547
5548 let notification = rx.recv().await.unwrap();
5549 assert!(matches!(
5550 notification,
5551 ServerNotification::PromptsListChanged
5552 ));
5553 }
5554
5555 #[tokio::test]
5556 async fn test_notify_without_sender_returns_false() {
5557 let router = McpRouter::new().server_info("test", "1.0");
5558
5559 assert!(!router.notify_tools_list_changed());
5560 assert!(!router.notify_prompts_list_changed());
5561 assert!(!router.notify_resources_list_changed());
5562 }
5563
5564 #[tokio::test]
5565 async fn test_list_changed_capabilities_with_notification_sender() {
5566 let (tx, _rx) = crate::context::notification_channel(16);
5567 let tool = ToolBuilder::new("test")
5568 .description("test")
5569 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5570 .build();
5571
5572 let mut router = McpRouter::new()
5573 .server_info("test", "1.0")
5574 .tool(tool)
5575 .with_notification_sender(tx);
5576
5577 init_router(&mut router).await;
5578
5579 let caps = router.capabilities();
5580 let tools_cap = caps.tools.expect("tools capability should be present");
5581 assert!(
5582 tools_cap.list_changed,
5583 "tools.listChanged should be true when notification sender is configured"
5584 );
5585 }
5586
5587 #[tokio::test]
5588 async fn test_list_changed_capabilities_without_notification_sender() {
5589 let tool = ToolBuilder::new("test")
5590 .description("test")
5591 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5592 .build();
5593
5594 let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
5595
5596 init_router(&mut router).await;
5597
5598 let caps = router.capabilities();
5599 let tools_cap = caps.tools.expect("tools capability should be present");
5600 assert!(
5601 !tools_cap.list_changed,
5602 "tools.listChanged should be false without notification sender"
5603 );
5604 }
5605
5606 #[tokio::test]
5607 async fn test_set_logging_level_filters_messages() {
5608 let (tx, mut rx) = crate::context::notification_channel(16);
5609
5610 let mut router = McpRouter::new()
5611 .server_info("test", "1.0")
5612 .with_notification_sender(tx);
5613
5614 init_router(&mut router).await;
5615
5616 let set_level_req = RouterRequest {
5618 id: RequestId::Number(99),
5619 inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
5620 level: LogLevel::Warning,
5621 meta: None,
5622 }),
5623 extensions: crate::context::Extensions::new(),
5624 };
5625 let resp = router
5626 .ready()
5627 .await
5628 .unwrap()
5629 .call(set_level_req)
5630 .await
5631 .unwrap();
5632 assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
5633
5634 let ctx = router.create_context(RequestId::Number(100), None);
5636
5637 ctx.send_log(LoggingMessageParams::new(
5639 LogLevel::Error,
5640 serde_json::Value::Null,
5641 ));
5642 assert!(
5643 rx.try_recv().is_ok(),
5644 "Error should pass through Warning filter"
5645 );
5646
5647 ctx.send_log(LoggingMessageParams::new(
5649 LogLevel::Info,
5650 serde_json::Value::Null,
5651 ));
5652 assert!(
5653 rx.try_recv().is_err(),
5654 "Info should be filtered at Warning level"
5655 );
5656 }
5657
5658 #[test]
5659 fn test_paginate_no_page_size() {
5660 let items = vec![1, 2, 3, 4, 5];
5661 let (page, cursor) = paginate(items.clone(), None, None).unwrap();
5662 assert_eq!(page, items);
5663 assert!(cursor.is_none());
5664 }
5665
5666 #[test]
5667 fn test_paginate_first_page() {
5668 let items = vec![1, 2, 3, 4, 5];
5669 let (page, cursor) = paginate(items, None, Some(2)).unwrap();
5670 assert_eq!(page, vec![1, 2]);
5671 assert!(cursor.is_some());
5672 }
5673
5674 #[test]
5675 fn test_paginate_middle_page() {
5676 let items = vec![1, 2, 3, 4, 5];
5677 let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
5678 assert_eq!(page1, vec![1, 2]);
5679
5680 let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
5681 assert_eq!(page2, vec![3, 4]);
5682 assert!(cursor2.is_some());
5683 }
5684
5685 #[test]
5686 fn test_paginate_last_page() {
5687 let items = vec![1, 2, 3, 4, 5];
5688 let cursor = encode_cursor(4);
5690 let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
5691 assert_eq!(page, vec![5]);
5692 assert!(next.is_none());
5693 }
5694
5695 #[test]
5696 fn test_paginate_exact_boundary() {
5697 let items = vec![1, 2, 3, 4];
5698 let (page, cursor) = paginate(items, None, Some(4)).unwrap();
5699 assert_eq!(page, vec![1, 2, 3, 4]);
5700 assert!(cursor.is_none());
5701 }
5702
5703 #[test]
5704 fn test_paginate_invalid_cursor() {
5705 let items = vec![1, 2, 3];
5706 let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
5707 assert!(result.is_err());
5708 }
5709
5710 #[test]
5711 fn test_cursor_round_trip() {
5712 let offset = 42;
5713 let encoded = encode_cursor(offset);
5714 let decoded = decode_cursor(&encoded).unwrap();
5715 assert_eq!(decoded, offset);
5716 }
5717
5718 #[tokio::test]
5719 async fn test_list_tools_pagination() {
5720 let tool_a = ToolBuilder::new("alpha")
5721 .description("a")
5722 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5723 .build();
5724 let tool_b = ToolBuilder::new("beta")
5725 .description("b")
5726 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5727 .build();
5728 let tool_c = ToolBuilder::new("gamma")
5729 .description("c")
5730 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5731 .build();
5732
5733 let mut router = McpRouter::new()
5734 .server_info("test", "1.0")
5735 .page_size(2)
5736 .tool(tool_a)
5737 .tool(tool_b)
5738 .tool(tool_c);
5739
5740 init_router(&mut router).await;
5741
5742 let req = RouterRequest {
5744 id: RequestId::Number(1),
5745 inner: McpRequest::ListTools(ListToolsParams {
5746 cursor: None,
5747 meta: None,
5748 }),
5749 extensions: Extensions::new(),
5750 };
5751 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5752 let (tools, next_cursor) = match resp.inner {
5753 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5754 other => panic!("Expected ListTools, got {:?}", other),
5755 };
5756 assert_eq!(tools.len(), 2);
5757 assert_eq!(tools[0].name, "alpha");
5758 assert_eq!(tools[1].name, "beta");
5759 assert!(next_cursor.is_some());
5760
5761 let req = RouterRequest {
5763 id: RequestId::Number(2),
5764 inner: McpRequest::ListTools(ListToolsParams {
5765 cursor: next_cursor,
5766 meta: None,
5767 }),
5768 extensions: Extensions::new(),
5769 };
5770 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5771 let (tools, next_cursor) = match resp.inner {
5772 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5773 other => panic!("Expected ListTools, got {:?}", other),
5774 };
5775 assert_eq!(tools.len(), 1);
5776 assert_eq!(tools[0].name, "gamma");
5777 assert!(next_cursor.is_none());
5778 }
5779
5780 #[tokio::test]
5781 async fn test_list_tools_no_pagination_by_default() {
5782 let tool_a = ToolBuilder::new("alpha")
5783 .description("a")
5784 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5785 .build();
5786 let tool_b = ToolBuilder::new("beta")
5787 .description("b")
5788 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5789 .build();
5790
5791 let mut router = McpRouter::new()
5792 .server_info("test", "1.0")
5793 .tool(tool_a)
5794 .tool(tool_b);
5795
5796 init_router(&mut router).await;
5797
5798 let req = RouterRequest {
5799 id: RequestId::Number(1),
5800 inner: McpRequest::ListTools(ListToolsParams {
5801 cursor: None,
5802 meta: None,
5803 }),
5804 extensions: Extensions::new(),
5805 };
5806 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5807 match resp.inner {
5808 Ok(McpResponse::ListTools(result)) => {
5809 assert_eq!(result.tools.len(), 2);
5810 assert!(result.next_cursor.is_none());
5811 }
5812 other => panic!("Expected ListTools, got {:?}", other),
5813 }
5814 }
5815
5816 #[cfg(feature = "dynamic-tools")]
5821 mod dynamic_tools_tests {
5822 use super::*;
5823
5824 #[tokio::test]
5825 async fn test_dynamic_tools_register_and_list() {
5826 let (router, registry) = McpRouter::new()
5827 .server_info("test", "1.0")
5828 .with_dynamic_tools();
5829
5830 let tool = ToolBuilder::new("dynamic_echo")
5831 .description("Dynamic echo")
5832 .handler(|input: AddInput| async move {
5833 Ok(CallToolResult::text(format!("{}", input.a)))
5834 })
5835 .build();
5836
5837 registry.register(tool);
5838
5839 let mut router = router;
5840 init_router(&mut router).await;
5841
5842 let req = RouterRequest {
5843 id: RequestId::Number(1),
5844 inner: McpRequest::ListTools(ListToolsParams::default()),
5845 extensions: Extensions::new(),
5846 };
5847
5848 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5849 match resp.inner {
5850 Ok(McpResponse::ListTools(result)) => {
5851 assert_eq!(result.tools.len(), 1);
5852 assert_eq!(result.tools[0].name, "dynamic_echo");
5853 }
5854 _ => panic!("Expected ListTools response"),
5855 }
5856 }
5857
5858 #[tokio::test]
5859 async fn test_dynamic_tools_unregister() {
5860 let (router, registry) = McpRouter::new()
5861 .server_info("test", "1.0")
5862 .with_dynamic_tools();
5863
5864 let tool = ToolBuilder::new("temp")
5865 .description("Temporary")
5866 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5867 .build();
5868
5869 registry.register(tool);
5870 assert!(registry.contains("temp"));
5871
5872 let removed = registry.unregister("temp");
5873 assert!(removed);
5874 assert!(!registry.contains("temp"));
5875
5876 assert!(!registry.unregister("temp"));
5878
5879 let mut router = router;
5880 init_router(&mut router).await;
5881
5882 let req = RouterRequest {
5883 id: RequestId::Number(1),
5884 inner: McpRequest::ListTools(ListToolsParams::default()),
5885 extensions: Extensions::new(),
5886 };
5887
5888 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5889 match resp.inner {
5890 Ok(McpResponse::ListTools(result)) => {
5891 assert_eq!(result.tools.len(), 0);
5892 }
5893 _ => panic!("Expected ListTools response"),
5894 }
5895 }
5896
5897 #[tokio::test]
5898 async fn test_dynamic_tools_merged_with_static() {
5899 let static_tool = ToolBuilder::new("static_tool")
5900 .description("Static")
5901 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5902 .build();
5903
5904 let (router, registry) = McpRouter::new()
5905 .server_info("test", "1.0")
5906 .tool(static_tool)
5907 .with_dynamic_tools();
5908
5909 let dynamic_tool = ToolBuilder::new("dynamic_tool")
5910 .description("Dynamic")
5911 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5912 .build();
5913
5914 registry.register(dynamic_tool);
5915
5916 let mut router = router;
5917 init_router(&mut router).await;
5918
5919 let req = RouterRequest {
5920 id: RequestId::Number(1),
5921 inner: McpRequest::ListTools(ListToolsParams::default()),
5922 extensions: Extensions::new(),
5923 };
5924
5925 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5926 match resp.inner {
5927 Ok(McpResponse::ListTools(result)) => {
5928 assert_eq!(result.tools.len(), 2);
5929 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5930 assert!(names.contains(&"static_tool"));
5931 assert!(names.contains(&"dynamic_tool"));
5932 }
5933 _ => panic!("Expected ListTools response"),
5934 }
5935 }
5936
5937 #[tokio::test]
5938 async fn test_static_tools_shadow_dynamic() {
5939 let static_tool = ToolBuilder::new("shared")
5940 .description("Static version")
5941 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5942 .build();
5943
5944 let (router, registry) = McpRouter::new()
5945 .server_info("test", "1.0")
5946 .tool(static_tool)
5947 .with_dynamic_tools();
5948
5949 let dynamic_tool = ToolBuilder::new("shared")
5950 .description("Dynamic version")
5951 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5952 .build();
5953
5954 registry.register(dynamic_tool);
5955
5956 let mut router = router;
5957 init_router(&mut router).await;
5958
5959 let req = RouterRequest {
5961 id: RequestId::Number(1),
5962 inner: McpRequest::ListTools(ListToolsParams::default()),
5963 extensions: Extensions::new(),
5964 };
5965
5966 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5967 match resp.inner {
5968 Ok(McpResponse::ListTools(result)) => {
5969 assert_eq!(result.tools.len(), 1);
5970 assert_eq!(result.tools[0].name, "shared");
5971 assert_eq!(
5972 result.tools[0].description.as_deref(),
5973 Some("Static version")
5974 );
5975 }
5976 _ => panic!("Expected ListTools response"),
5977 }
5978
5979 let req = RouterRequest {
5981 id: RequestId::Number(2),
5982 inner: McpRequest::CallTool(CallToolParams {
5983 name: "shared".to_string(),
5984 arguments: serde_json::json!({"a": 1, "b": 2}),
5985 meta: None,
5986 task: None,
5987 }),
5988 extensions: Extensions::new(),
5989 };
5990
5991 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5992 match resp.inner {
5993 Ok(McpResponse::CallTool(result)) => {
5994 assert!(!result.is_error);
5995 match &result.content[0] {
5996 Content::Text { text, .. } => assert_eq!(text, "static"),
5997 _ => panic!("Expected text content"),
5998 }
5999 }
6000 _ => panic!("Expected CallTool response"),
6001 }
6002 }
6003
6004 #[tokio::test]
6005 async fn test_dynamic_tools_call() {
6006 let (router, registry) = McpRouter::new()
6007 .server_info("test", "1.0")
6008 .with_dynamic_tools();
6009
6010 let tool = ToolBuilder::new("add")
6011 .description("Add two numbers")
6012 .handler(|input: AddInput| async move {
6013 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
6014 })
6015 .build();
6016
6017 registry.register(tool);
6018
6019 let mut router = router;
6020 init_router(&mut router).await;
6021
6022 let req = RouterRequest {
6023 id: RequestId::Number(1),
6024 inner: McpRequest::CallTool(CallToolParams {
6025 name: "add".to_string(),
6026 arguments: serde_json::json!({"a": 3, "b": 4}),
6027 meta: None,
6028 task: None,
6029 }),
6030 extensions: Extensions::new(),
6031 };
6032
6033 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6034 match resp.inner {
6035 Ok(McpResponse::CallTool(result)) => {
6036 assert!(!result.is_error);
6037 match &result.content[0] {
6038 Content::Text { text, .. } => assert_eq!(text, "7"),
6039 _ => panic!("Expected text content"),
6040 }
6041 }
6042 _ => panic!("Expected CallTool response"),
6043 }
6044 }
6045
6046 #[tokio::test]
6047 async fn test_dynamic_tools_notification_on_register() {
6048 let (tx, mut rx) = crate::context::notification_channel(16);
6049 let (router, registry) = McpRouter::new()
6050 .server_info("test", "1.0")
6051 .with_dynamic_tools();
6052 let _router = router.with_notification_sender(tx);
6053
6054 let tool = ToolBuilder::new("notified")
6055 .description("Test")
6056 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6057 .build();
6058
6059 registry.register(tool);
6060
6061 let notification = rx.recv().await.unwrap();
6062 assert!(matches!(notification, ServerNotification::ToolsListChanged));
6063 }
6064
6065 #[tokio::test]
6066 async fn test_dynamic_tools_notification_on_unregister() {
6067 let (tx, mut rx) = crate::context::notification_channel(16);
6068 let (router, registry) = McpRouter::new()
6069 .server_info("test", "1.0")
6070 .with_dynamic_tools();
6071 let _router = router.with_notification_sender(tx);
6072
6073 let tool = ToolBuilder::new("notified")
6074 .description("Test")
6075 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6076 .build();
6077
6078 registry.register(tool);
6079 let _ = rx.recv().await.unwrap();
6081
6082 registry.unregister("notified");
6083 let notification = rx.recv().await.unwrap();
6084 assert!(matches!(notification, ServerNotification::ToolsListChanged));
6085 }
6086
6087 #[tokio::test]
6088 async fn test_dynamic_tools_no_notification_on_empty_unregister() {
6089 let (tx, mut rx) = crate::context::notification_channel(16);
6090 let (router, registry) = McpRouter::new()
6091 .server_info("test", "1.0")
6092 .with_dynamic_tools();
6093 let _router = router.with_notification_sender(tx);
6094
6095 assert!(!registry.unregister("nonexistent"));
6097
6098 assert!(rx.try_recv().is_err());
6100 }
6101
6102 #[tokio::test]
6103 async fn test_dynamic_tools_filter_applies() {
6104 use crate::filter::CapabilityFilter;
6105
6106 let (router, registry) = McpRouter::new()
6107 .server_info("test", "1.0")
6108 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
6109 tool.name != "hidden"
6110 }))
6111 .with_dynamic_tools();
6112
6113 let visible = ToolBuilder::new("visible")
6114 .description("Visible")
6115 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6116 .build();
6117
6118 let hidden = ToolBuilder::new("hidden")
6119 .description("Hidden")
6120 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6121 .build();
6122
6123 registry.register(visible);
6124 registry.register(hidden);
6125
6126 let mut router = router;
6127 init_router(&mut router).await;
6128
6129 let req = RouterRequest {
6131 id: RequestId::Number(1),
6132 inner: McpRequest::ListTools(ListToolsParams::default()),
6133 extensions: Extensions::new(),
6134 };
6135
6136 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6137 match resp.inner {
6138 Ok(McpResponse::ListTools(result)) => {
6139 assert_eq!(result.tools.len(), 1);
6140 assert_eq!(result.tools[0].name, "visible");
6141 }
6142 _ => panic!("Expected ListTools response"),
6143 }
6144
6145 let req = RouterRequest {
6147 id: RequestId::Number(2),
6148 inner: McpRequest::CallTool(CallToolParams {
6149 name: "hidden".to_string(),
6150 arguments: serde_json::json!({"a": 1, "b": 2}),
6151 meta: None,
6152 task: None,
6153 }),
6154 extensions: Extensions::new(),
6155 };
6156
6157 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6158 match resp.inner {
6159 Err(e) => {
6160 assert_eq!(e.code, -32601); }
6162 _ => panic!("Expected JsonRpc error"),
6163 }
6164 }
6165
6166 #[tokio::test]
6167 async fn test_dynamic_tools_capabilities_advertised() {
6168 let (mut router, _registry) = McpRouter::new()
6170 .server_info("test", "1.0")
6171 .with_dynamic_tools();
6172
6173 let init_req = RouterRequest {
6174 id: RequestId::Number(1),
6175 inner: McpRequest::Initialize(InitializeParams {
6176 protocol_version: "2025-11-25".to_string(),
6177 capabilities: ClientCapabilities::default(),
6178 client_info: Implementation {
6179 name: "test".to_string(),
6180 version: "1.0".to_string(),
6181 ..Default::default()
6182 },
6183 meta: None,
6184 }),
6185 extensions: Extensions::new(),
6186 };
6187
6188 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
6189 match resp.inner {
6190 Ok(McpResponse::Initialize(result)) => {
6191 assert!(result.capabilities.tools.is_some());
6192 }
6193 _ => panic!("Expected Initialize response"),
6194 }
6195 }
6196
6197 #[tokio::test]
6198 async fn test_dynamic_tools_multi_session_notification() {
6199 let (tx1, mut rx1) = crate::context::notification_channel(16);
6200 let (tx2, mut rx2) = crate::context::notification_channel(16);
6201
6202 let (router, registry) = McpRouter::new()
6203 .server_info("test", "1.0")
6204 .with_dynamic_tools();
6205
6206 let _session1 = router.clone().with_notification_sender(tx1);
6208 let _session2 = router.clone().with_notification_sender(tx2);
6209
6210 let tool = ToolBuilder::new("broadcast")
6211 .description("Test")
6212 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6213 .build();
6214
6215 registry.register(tool);
6216
6217 let n1 = rx1.recv().await.unwrap();
6219 let n2 = rx2.recv().await.unwrap();
6220 assert!(matches!(n1, ServerNotification::ToolsListChanged));
6221 assert!(matches!(n2, ServerNotification::ToolsListChanged));
6222 }
6223
6224 #[tokio::test]
6225 async fn test_dynamic_tools_call_not_found() {
6226 let (router, _registry) = McpRouter::new()
6227 .server_info("test", "1.0")
6228 .with_dynamic_tools();
6229
6230 let mut router = router;
6231 init_router(&mut router).await;
6232
6233 let req = RouterRequest {
6234 id: RequestId::Number(1),
6235 inner: McpRequest::CallTool(CallToolParams {
6236 name: "nonexistent".to_string(),
6237 arguments: serde_json::json!({}),
6238 meta: None,
6239 task: None,
6240 }),
6241 extensions: Extensions::new(),
6242 };
6243
6244 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6245 match resp.inner {
6246 Err(e) => {
6247 assert_eq!(e.code, -32601);
6248 }
6249 _ => panic!("Expected method not found error"),
6250 }
6251 }
6252
6253 #[tokio::test]
6254 async fn test_dynamic_tools_registry_list() {
6255 let (_, registry) = McpRouter::new()
6256 .server_info("test", "1.0")
6257 .with_dynamic_tools();
6258
6259 assert!(registry.list().is_empty());
6260
6261 let tool = ToolBuilder::new("tool_a")
6262 .description("A")
6263 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6264 .build();
6265 registry.register(tool);
6266
6267 let tool = ToolBuilder::new("tool_b")
6268 .description("B")
6269 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6270 .build();
6271 registry.register(tool);
6272
6273 let tools = registry.list();
6274 assert_eq!(tools.len(), 2);
6275 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
6276 assert!(names.contains(&"tool_a"));
6277 assert!(names.contains(&"tool_b"));
6278 }
6279 } #[tokio::test]
6282 async fn test_tool_if_true_registers() {
6283 let tool = ToolBuilder::new("conditional")
6284 .description("Conditional tool")
6285 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6286 .build();
6287
6288 let mut router = McpRouter::new().tool_if(true, tool);
6289 init_router(&mut router).await;
6290
6291 let req = RouterRequest {
6292 id: RequestId::Number(1),
6293 inner: McpRequest::ListTools(ListToolsParams::default()),
6294 extensions: Extensions::new(),
6295 };
6296 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6297 match resp.inner {
6298 Ok(McpResponse::ListTools(result)) => {
6299 assert_eq!(result.tools.len(), 1);
6300 assert_eq!(result.tools[0].name, "conditional");
6301 }
6302 _ => panic!("Expected ListTools response"),
6303 }
6304 }
6305
6306 #[tokio::test]
6307 async fn test_tool_if_false_skips() {
6308 let tool = ToolBuilder::new("conditional")
6309 .description("Conditional tool")
6310 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6311 .build();
6312
6313 let mut router = McpRouter::new().tool_if(false, tool);
6314 init_router(&mut router).await;
6315
6316 let req = RouterRequest {
6317 id: RequestId::Number(1),
6318 inner: McpRequest::ListTools(ListToolsParams::default()),
6319 extensions: Extensions::new(),
6320 };
6321 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6322 match resp.inner {
6323 Ok(McpResponse::ListTools(result)) => {
6324 assert_eq!(result.tools.len(), 0);
6325 }
6326 _ => panic!("Expected ListTools response"),
6327 }
6328 }
6329
6330 #[tokio::test]
6331 async fn test_tools_if_batch_conditional() {
6332 let tools = vec![
6333 ToolBuilder::new("a")
6334 .description("Tool A")
6335 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6336 .build(),
6337 ToolBuilder::new("b")
6338 .description("Tool B")
6339 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6340 .build(),
6341 ];
6342
6343 let mut router = McpRouter::new().tools_if(false, tools);
6344 init_router(&mut router).await;
6345
6346 let req = RouterRequest {
6347 id: RequestId::Number(1),
6348 inner: McpRequest::ListTools(ListToolsParams::default()),
6349 extensions: Extensions::new(),
6350 };
6351 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6352 match resp.inner {
6353 Ok(McpResponse::ListTools(result)) => {
6354 assert_eq!(result.tools.len(), 0);
6355 }
6356 _ => panic!("Expected ListTools response"),
6357 }
6358 }
6359
6360 #[test]
6361 fn test_resource_if_true_registers() {
6362 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6363 .name("test")
6364 .text("hello");
6365
6366 let router = McpRouter::new().resource_if(true, resource);
6367 assert_eq!(router.inner.resources.len(), 1);
6368 }
6369
6370 #[test]
6371 fn test_resource_if_false_skips() {
6372 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6373 .name("test")
6374 .text("hello");
6375
6376 let router = McpRouter::new().resource_if(false, resource);
6377 assert_eq!(router.inner.resources.len(), 0);
6378 }
6379
6380 #[test]
6381 fn test_prompt_if_true_registers() {
6382 let prompt = crate::prompt::PromptBuilder::new("greet")
6383 .description("Greeting")
6384 .user_message("Hello!");
6385
6386 let router = McpRouter::new().prompt_if(true, prompt);
6387 assert_eq!(router.inner.prompts.len(), 1);
6388 }
6389
6390 #[test]
6391 fn test_prompt_if_false_skips() {
6392 let prompt = crate::prompt::PromptBuilder::new("greet")
6393 .description("Greeting")
6394 .user_message("Hello!");
6395
6396 let router = McpRouter::new().prompt_if(false, prompt);
6397 assert_eq!(router.inner.prompts.len(), 0);
6398 }
6399
6400 #[tokio::test]
6401 async fn test_disable_tool_hides_from_list() {
6402 let safe = ToolBuilder::new("safe")
6403 .description("Safe tool")
6404 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6405 .build();
6406 let dangerous = ToolBuilder::new("dangerous")
6407 .description("Dangerous tool")
6408 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6409 .build();
6410 let mut router = McpRouter::new().tool(safe).tool(dangerous);
6411 init_router(&mut router).await;
6412
6413 router.disable_tool("dangerous");
6414 assert!(router.is_tool_enabled("safe"));
6415 assert!(!router.is_tool_enabled("dangerous"));
6416
6417 let req = RouterRequest {
6418 id: RequestId::Number(1),
6419 inner: McpRequest::ListTools(ListToolsParams::default()),
6420 extensions: Extensions::new(),
6421 };
6422 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6423 match resp.inner {
6424 Ok(McpResponse::ListTools(result)) => {
6425 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
6426 assert_eq!(names, vec!["safe"]);
6427 }
6428 _ => panic!("Expected ListTools response"),
6429 }
6430 }
6431
6432 #[tokio::test]
6433 async fn test_disable_tool_blocks_call() {
6434 let dangerous = ToolBuilder::new("dangerous")
6435 .description("Dangerous tool")
6436 .handler(|_: AddInput| async { Ok(CallToolResult::text("ran")) })
6437 .build();
6438 let mut router = McpRouter::new().tool(dangerous);
6439 init_router(&mut router).await;
6440
6441 router.disable_tool("dangerous");
6442
6443 let req = RouterRequest {
6444 id: RequestId::Number(2),
6445 inner: McpRequest::CallTool(CallToolParams {
6446 name: "dangerous".to_string(),
6447 arguments: serde_json::json!({"a": 1, "b": 2}),
6448 meta: None,
6449 task: None,
6450 }),
6451 extensions: Extensions::new(),
6452 };
6453 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6454 let err = resp.inner.expect_err("disabled tool should error");
6455 assert_eq!(err.code, crate::error::ErrorCode::MethodNotFound as i32);
6456 }
6457
6458 #[tokio::test]
6459 async fn test_enable_tool_restores_visibility() {
6460 let tool = ToolBuilder::new("flippy")
6461 .description("Toggleable tool")
6462 .handler(|_: AddInput| async { Ok(CallToolResult::text("ran")) })
6463 .build();
6464 let mut router = McpRouter::new().tool(tool);
6465 init_router(&mut router).await;
6466
6467 router.disable_tool("flippy");
6468 router.enable_tool("flippy");
6469 assert!(router.is_tool_enabled("flippy"));
6470
6471 let req = RouterRequest {
6472 id: RequestId::Number(3),
6473 inner: McpRequest::CallTool(CallToolParams {
6474 name: "flippy".to_string(),
6475 arguments: serde_json::json!({"a": 1, "b": 2}),
6476 meta: None,
6477 task: None,
6478 }),
6479 extensions: Extensions::new(),
6480 };
6481 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6482 match resp.inner {
6483 Ok(McpResponse::CallTool(result)) => {
6484 assert_eq!(result.first_text(), Some("ran"));
6485 }
6486 _ => panic!("Expected CallTool response"),
6487 }
6488 }
6489
6490 #[tokio::test]
6491 async fn test_disable_propagates_through_fresh_session() {
6492 let tool = ToolBuilder::new("shared")
6493 .description("Shared across sessions")
6494 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6495 .build();
6496 let router = McpRouter::new().tool(tool);
6497
6498 router.disable_tool("shared");
6500 let mut child = router.with_fresh_session();
6501 init_router(&mut child).await;
6502 assert!(!child.is_tool_enabled("shared"));
6503
6504 let req = RouterRequest {
6505 id: RequestId::Number(4),
6506 inner: McpRequest::ListTools(ListToolsParams::default()),
6507 extensions: Extensions::new(),
6508 };
6509 let resp = child.ready().await.unwrap().call(req).await.unwrap();
6510 match resp.inner {
6511 Ok(McpResponse::ListTools(result)) => {
6512 assert!(result.tools.is_empty());
6513 }
6514 _ => panic!("Expected ListTools response"),
6515 }
6516 }
6517
6518 #[tokio::test]
6519 async fn test_disable_resource_and_prompt() {
6520 let resource = crate::resource::ResourceBuilder::new("file:///hidden.txt")
6521 .name("hidden")
6522 .text("secret");
6523 let prompt = crate::prompt::PromptBuilder::new("hidden_prompt")
6524 .description("hidden")
6525 .user_message("hello");
6526
6527 let mut router = McpRouter::new().resource(resource).prompt(prompt);
6528 init_router(&mut router).await;
6529
6530 router.disable_resource("file:///hidden.txt");
6531 router.disable_prompt("hidden_prompt");
6532 assert!(!router.is_resource_enabled("file:///hidden.txt"));
6533 assert!(!router.is_prompt_enabled("hidden_prompt"));
6534
6535 let req = RouterRequest {
6537 id: RequestId::Number(5),
6538 inner: McpRequest::ListResources(ListResourcesParams::default()),
6539 extensions: Extensions::new(),
6540 };
6541 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6542 match resp.inner {
6543 Ok(McpResponse::ListResources(result)) => {
6544 assert!(result.resources.is_empty());
6545 }
6546 _ => panic!("Expected ListResources response"),
6547 }
6548
6549 let req = RouterRequest {
6551 id: RequestId::Number(6),
6552 inner: McpRequest::ReadResource(ReadResourceParams {
6553 uri: "file:///hidden.txt".to_string(),
6554 meta: None,
6555 }),
6556 extensions: Extensions::new(),
6557 };
6558 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6559 let err = resp.inner.expect_err("disabled resource should error");
6560 assert_eq!(err.code, -32602); let req = RouterRequest {
6564 id: RequestId::Number(7),
6565 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
6566 extensions: Extensions::new(),
6567 };
6568 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6569 match resp.inner {
6570 Ok(McpResponse::ListPrompts(result)) => {
6571 assert!(result.prompts.is_empty());
6572 }
6573 _ => panic!("Expected ListPrompts response"),
6574 }
6575
6576 let req = RouterRequest {
6578 id: RequestId::Number(8),
6579 inner: McpRequest::GetPrompt(GetPromptParams {
6580 name: "hidden_prompt".to_string(),
6581 arguments: Default::default(),
6582 meta: None,
6583 }),
6584 extensions: Extensions::new(),
6585 };
6586 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6587 let err = resp.inner.expect_err("disabled prompt should error");
6588 assert_eq!(err.code, crate::error::ErrorCode::MethodNotFound as i32);
6589 }
6590
6591 #[test]
6592 fn test_router_request_new() {
6593 let req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6594 assert_eq!(req.id, RequestId::Number(1));
6595 assert!(req.extensions.is_empty());
6596 }
6597
6598 #[test]
6599 fn test_with_inner_preserves_extensions() {
6600 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6601 req.extensions.insert(42u32);
6602
6603 let rewritten = req.with_inner(McpRequest::ListTools(Default::default()));
6604 assert!(matches!(rewritten.inner, McpRequest::ListTools(_)));
6605 assert_eq!(rewritten.id, RequestId::Number(1));
6606 assert_eq!(rewritten.extensions.get::<u32>(), Some(&42));
6607 }
6608
6609 #[test]
6610 fn test_with_id_and_inner_preserves_extensions() {
6611 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6612 req.extensions.insert(String::from("token-abc"));
6613
6614 let rewritten = req.with_id_and_inner(
6615 RequestId::Number(99),
6616 McpRequest::ListResources(Default::default()),
6617 );
6618 assert_eq!(rewritten.id, RequestId::Number(99));
6619 assert!(matches!(rewritten.inner, McpRequest::ListResources(_)));
6620 assert_eq!(
6621 rewritten.extensions.get::<String>(),
6622 Some(&String::from("token-abc"))
6623 );
6624 }
6625
6626 #[test]
6627 fn test_clone_with_inner_preserves_extensions() {
6628 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6629 req.extensions.insert(true);
6630
6631 let cloned = req.clone_with_inner(McpRequest::ListTools(Default::default()));
6632
6633 assert!(matches!(req.inner, McpRequest::Ping));
6635 assert_eq!(req.extensions.get::<bool>(), Some(&true));
6636
6637 assert!(matches!(cloned.inner, McpRequest::ListTools(_)));
6639 assert_eq!(cloned.extensions.get::<bool>(), Some(&true));
6640 }
6641
6642 #[test]
6643 fn test_router_response_is_error() {
6644 let ok_resp = RouterResponse {
6645 id: RequestId::Number(1),
6646 inner: Ok(McpResponse::Pong(Default::default())),
6647 };
6648 assert!(!ok_resp.is_error());
6649
6650 let err_resp = RouterResponse {
6651 id: RequestId::Number(2),
6652 inner: Err(JsonRpcError::internal_error("boom")),
6653 };
6654 assert!(err_resp.is_error());
6655 }
6656
6657 #[test]
6658 fn test_extensions_len_and_is_empty() {
6659 let mut ext = Extensions::new();
6660 assert!(ext.is_empty());
6661 assert_eq!(ext.len(), 0);
6662
6663 ext.insert(42u32);
6664 assert!(!ext.is_empty());
6665 assert_eq!(ext.len(), 1);
6666
6667 ext.insert(String::from("hello"));
6668 assert_eq!(ext.len(), 2);
6669 }
6670
6671 #[test]
6672 fn test_router_response_serde_roundtrip() {
6673 let response = RouterResponse {
6675 id: RequestId::Number(1),
6676 inner: Ok(McpResponse::Empty(EmptyResult {})),
6677 };
6678 let json = serde_json::to_string(&response).unwrap();
6679 let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6680 assert_eq!(deserialized.id, RequestId::Number(1));
6681 assert!(!deserialized.is_error());
6682
6683 let response = RouterResponse {
6685 id: RequestId::String("req-2".into()),
6686 inner: Err(JsonRpcError::method_not_found("unknown")),
6687 };
6688 let json = serde_json::to_string(&response).unwrap();
6689 let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6690 assert_eq!(deserialized.id, RequestId::String("req-2".into()));
6691 assert!(deserialized.is_error());
6692 }
6693
6694 #[tokio::test]
6701 async fn test_discover_dispatch_via_jsonrpc_service() {
6702 let router = McpRouter::new().server_info("unit-test-server", "4.2.0");
6705 let mut service = JsonRpcService::new(router);
6706
6707 let req = JsonRpcRequest::new(1, "server/discover");
6708 let resp = service.call_single(req).await.unwrap();
6709
6710 match resp {
6711 JsonRpcResponse::Result(r) => {
6712 let versions = r
6714 .result
6715 .get("supportedVersions")
6716 .and_then(|v| v.as_array())
6717 .expect("result.supportedVersions must be an array");
6718 assert!(!versions.is_empty(), "supportedVersions must not be empty");
6719
6720 assert_eq!(
6722 r.result["serverInfo"]["name"], "unit-test-server",
6723 "serverInfo.name must match configured value"
6724 );
6725 assert_eq!(
6726 r.result["serverInfo"]["version"], "4.2.0",
6727 "serverInfo.version must match configured value"
6728 );
6729
6730 assert!(
6733 r.result.get("protocolVersion").is_none(),
6734 "server/discover must NOT include protocolVersion: {:?}",
6735 r.result
6736 );
6737 }
6738 JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
6739 _ => panic!("unexpected response variant"),
6740 }
6741 }
6742
6743 #[tokio::test]
6744 async fn test_discover_does_not_require_initialization() {
6745 let router = McpRouter::new().server_info("fresh-router", "1.0.0");
6748 let mut service = JsonRpcService::new(router);
6749
6750 let req = JsonRpcRequest::new(2, "server/discover");
6751 let resp = service.call_single(req).await.unwrap();
6752
6753 assert!(
6755 !matches!(resp, JsonRpcResponse::Error(_)),
6756 "server/discover must not require initialization: {:?}",
6757 resp
6758 );
6759 }
6760}