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