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 ctx = self.create_context(request_id, None);
1901 let result = resource.read_with_context(ctx).await;
1902 return Ok(McpResponse::ReadResource(result));
1903 }
1904
1905 #[cfg(feature = "dynamic-tools")]
1907 #[allow(clippy::collapsible_if)]
1908 if let Some(ref dynamic) = self.inner.dynamic_resources {
1909 if let Some(resource) = dynamic.get(¶ms.uri) {
1910 if let Some(filter) = &self.inner.resource_filter
1911 && !filter.is_visible(&self.session, &resource)
1912 {
1913 return Err(filter.denial_error(¶ms.uri));
1914 }
1915 tracing::debug!(uri = %params.uri, "Reading dynamic resource");
1916 let ctx = self.create_context(request_id, None);
1917 let result = resource.read_with_context(ctx).await;
1918 return Ok(McpResponse::ReadResource(result));
1919 }
1920 }
1921
1922 for template in &self.inner.resource_templates {
1924 if let Some(variables) = template.match_uri(¶ms.uri) {
1925 tracing::debug!(
1926 uri = %params.uri,
1927 template = %template.uri_template,
1928 "Reading resource via template"
1929 );
1930 let result = template.read(¶ms.uri, variables).await?;
1931 return Ok(McpResponse::ReadResource(result));
1932 }
1933 }
1934
1935 #[cfg(feature = "dynamic-tools")]
1937 #[allow(clippy::collapsible_if)]
1938 if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
1939 if let Some((template, variables)) = dynamic.match_uri(¶ms.uri) {
1940 tracing::debug!(
1941 uri = %params.uri,
1942 template = %template.uri_template,
1943 "Reading resource via dynamic template"
1944 );
1945 let result = template.read(¶ms.uri, variables).await?;
1946 return Ok(McpResponse::ReadResource(result));
1947 }
1948 }
1949
1950 Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1952 ¶ms.uri,
1953 )))
1954 }
1955
1956 McpRequest::SubscribeResource(params) => {
1957 if !self.inner.resources.contains_key(¶ms.uri) {
1959 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1960 ¶ms.uri,
1961 )));
1962 }
1963
1964 tracing::debug!(uri = %params.uri, "Subscribing to resource");
1965 self.subscribe(¶ms.uri);
1966
1967 Ok(McpResponse::SubscribeResource(EmptyResult {}))
1968 }
1969
1970 McpRequest::UnsubscribeResource(params) => {
1971 if !self.inner.resources.contains_key(¶ms.uri) {
1973 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1974 ¶ms.uri,
1975 )));
1976 }
1977
1978 tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
1979 self.unsubscribe(¶ms.uri);
1980
1981 Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
1982 }
1983
1984 McpRequest::ListPrompts(params) => {
1985 let is_visible = |p: &Prompt| -> bool {
1986 self.inner
1987 .prompt_filter
1988 .as_ref()
1989 .map(|f| f.is_visible(&self.session, p))
1990 .unwrap_or(true)
1991 };
1992
1993 let mut prompts: Vec<PromptDefinition> = self
1994 .inner
1995 .prompts
1996 .values()
1997 .filter(|p| is_visible(p))
1998 .map(|p| p.definition())
1999 .collect();
2000
2001 #[cfg(feature = "dynamic-tools")]
2003 if let Some(ref dynamic) = self.inner.dynamic_prompts {
2004 let static_names: HashSet<String> =
2005 prompts.iter().map(|p| p.name.clone()).collect();
2006 for p in dynamic.list() {
2007 if !static_names.contains(&p.name) && is_visible(&p) {
2008 prompts.push(p.definition());
2009 }
2010 }
2011 }
2012
2013 prompts.sort_by(|a, b| a.name.cmp(&b.name));
2014
2015 let (prompts, next_cursor) =
2016 paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
2017
2018 Ok(McpResponse::ListPrompts(ListPromptsResult {
2019 prompts,
2020 next_cursor,
2021 meta: None,
2022 }))
2023 }
2024
2025 McpRequest::GetPrompt(params) => {
2026 let prompt = self.inner.prompts.get(¶ms.name).cloned();
2028 #[cfg(feature = "dynamic-tools")]
2029 let prompt = prompt.or_else(|| {
2030 self.inner
2031 .dynamic_prompts
2032 .as_ref()
2033 .and_then(|d| d.get(¶ms.name))
2034 });
2035 let prompt = prompt.ok_or_else(|| {
2036 Error::JsonRpc(JsonRpcError::method_not_found(&format!(
2037 "Prompt not found: {}",
2038 params.name
2039 )))
2040 })?;
2041
2042 if let Some(filter) = &self.inner.prompt_filter
2044 && !filter.is_visible(&self.session, &prompt)
2045 {
2046 return Err(filter.denial_error(¶ms.name));
2047 }
2048
2049 tracing::debug!(name = %params.name, "Getting prompt");
2050 let ctx = self.create_context(request_id, None);
2051 let result = prompt.get_with_context(ctx, params.arguments).await?;
2052
2053 Ok(McpResponse::GetPrompt(result))
2054 }
2055
2056 McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
2057
2058 McpRequest::ListTasks(params) => {
2059 let tasks = self.inner.task_store.list_tasks(params.status);
2060
2061 let (tasks, next_cursor) =
2062 paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
2063
2064 Ok(McpResponse::ListTasks(ListTasksResult {
2065 tasks,
2066 next_cursor,
2067 }))
2068 }
2069
2070 McpRequest::GetTaskInfo(params) => {
2071 let task = self
2072 .inner
2073 .task_store
2074 .get_task(¶ms.task_id)
2075 .ok_or_else(|| {
2076 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2077 "Task not found: {}",
2078 params.task_id
2079 )))
2080 })?;
2081
2082 Ok(McpResponse::GetTaskInfo(task))
2083 }
2084
2085 McpRequest::GetTaskResult(params) => {
2086 let (task_obj, result, error) = self
2088 .inner
2089 .task_store
2090 .wait_for_completion(¶ms.task_id)
2091 .await
2092 .ok_or_else(|| {
2093 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2094 "Task not found: {}",
2095 params.task_id
2096 )))
2097 })?;
2098
2099 let meta = serde_json::json!({
2101 "io.modelcontextprotocol/related-task": task_obj
2102 });
2103
2104 match task_obj.status {
2105 TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
2106 format!("Task {} was cancelled", params.task_id),
2107 ))),
2108 TaskStatus::Failed => {
2109 let mut call_result = CallToolResult::error(
2110 error.unwrap_or_else(|| "Task failed".to_string()),
2111 );
2112 call_result.meta = Some(meta);
2113 Ok(McpResponse::GetTaskResult(call_result))
2114 }
2115 _ => {
2116 let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
2117 call_result.meta = Some(meta);
2118 Ok(McpResponse::GetTaskResult(call_result))
2119 }
2120 }
2121 }
2122
2123 McpRequest::CancelTask(params) => {
2124 let current = self
2126 .inner
2127 .task_store
2128 .get_task(¶ms.task_id)
2129 .ok_or_else(|| {
2130 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2131 "Task not found: {}",
2132 params.task_id
2133 )))
2134 })?;
2135
2136 if current.status.is_terminal() {
2137 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
2138 "Task {} is already in terminal state: {}",
2139 params.task_id, current.status
2140 ))));
2141 }
2142
2143 let task_obj = self
2144 .inner
2145 .task_store
2146 .cancel_task(¶ms.task_id, params.reason.as_deref())
2147 .ok_or_else(|| {
2148 Error::JsonRpc(JsonRpcError::invalid_params(format!(
2149 "Task not found: {}",
2150 params.task_id
2151 )))
2152 })?;
2153
2154 Ok(McpResponse::CancelTask(task_obj))
2155 }
2156
2157 McpRequest::SetLoggingLevel(params) => {
2158 tracing::debug!(level = ?params.level, "Client set logging level");
2159 if let Ok(mut level) = self.inner.min_log_level.write() {
2160 *level = params.level;
2161 }
2162 Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
2163 }
2164
2165 McpRequest::Complete(params) => {
2166 tracing::debug!(
2167 reference = ?params.reference,
2168 argument = %params.argument.name,
2169 "Completion request"
2170 );
2171
2172 if let Some(ref handler) = self.inner.completion_handler {
2174 let result = handler(params).await?;
2175 Ok(McpResponse::Complete(result))
2176 } else {
2177 Ok(McpResponse::Complete(CompleteResult::new(vec![])))
2179 }
2180 }
2181
2182 McpRequest::Unknown { ref method, .. } if method == "server/discover" => {
2183 #[cfg(feature = "stateless")]
2184 {
2185 use crate::protocol::SUPPORTED_PROTOCOL_VERSIONS;
2186 let result = crate::stateless::DiscoverResult {
2187 supported_versions: SUPPORTED_PROTOCOL_VERSIONS
2188 .iter()
2189 .map(|v| v.to_string())
2190 .collect(),
2191 capabilities: self.capabilities(),
2192 server_info: Implementation {
2193 name: self.inner.server_name.clone(),
2194 version: self.inner.server_version.clone(),
2195 title: self.inner.server_title.clone(),
2196 description: self.inner.server_description.clone(),
2197 icons: self.inner.server_icons.clone(),
2198 website_url: None,
2199 meta: None,
2200 },
2201 instructions: self.inner.instructions.clone(),
2202 };
2203 Ok(McpResponse::Raw(serde_json::to_value(result).unwrap()))
2204 }
2205 #[cfg(not(feature = "stateless"))]
2206 {
2207 Err(Error::JsonRpc(JsonRpcError::method_not_found(method)))
2208 }
2209 }
2210
2211 McpRequest::Unknown { method, .. } => {
2212 Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
2213 }
2214 _ => Err(Error::JsonRpc(JsonRpcError::method_not_found(
2215 "unknown method",
2216 ))),
2217 }
2218 }
2219
2220 pub fn handle_notification(&self, notification: McpNotification) {
2222 match notification {
2223 McpNotification::Initialized => {
2224 let phase_before = self.session.phase();
2225 if self.session.mark_initialized() {
2226 if phase_before == crate::session::SessionPhase::Uninitialized {
2227 tracing::info!(
2228 "Session initialized from uninitialized state (race resolved)"
2229 );
2230 } else {
2231 tracing::info!("Session initialized, entering operation phase");
2232 }
2233 } else {
2234 tracing::warn!(
2235 phase = ?self.session.phase(),
2236 "Received initialized notification in unexpected state"
2237 );
2238 }
2239 }
2240 McpNotification::Cancelled(params) => {
2241 if let Some(ref request_id) = params.request_id {
2242 if self.cancel_request(request_id) {
2243 tracing::info!(
2244 request_id = ?request_id,
2245 reason = ?params.reason,
2246 "Request cancelled"
2247 );
2248 } else {
2249 tracing::debug!(
2250 request_id = ?request_id,
2251 reason = ?params.reason,
2252 "Cancellation requested for unknown request"
2253 );
2254 }
2255 } else {
2256 tracing::debug!(
2257 reason = ?params.reason,
2258 "Cancellation notification received without request_id"
2259 );
2260 }
2261 }
2262 McpNotification::Progress(params) => {
2263 tracing::trace!(
2264 token = ?params.progress_token,
2265 progress = params.progress,
2266 total = ?params.total,
2267 "Progress notification"
2268 );
2269 }
2271 McpNotification::RootsListChanged => {
2272 tracing::info!("Client roots list changed");
2273 }
2276 McpNotification::Unknown { method, .. } => {
2277 tracing::debug!(method = %method, "Unknown notification received");
2278 }
2279 _ => {
2280 tracing::debug!("Unrecognized notification variant received");
2281 }
2282 }
2283 }
2284}
2285
2286impl Default for McpRouter {
2287 fn default() -> Self {
2288 Self::new()
2289 }
2290}
2291
2292pub use crate::context::Extensions;
2298
2299#[derive(Debug, Clone)]
2324pub struct ToolAnnotationsMap {
2325 map: Arc<HashMap<String, ToolAnnotations>>,
2326}
2327
2328impl ToolAnnotationsMap {
2329 pub fn get(&self, tool_name: &str) -> Option<&ToolAnnotations> {
2333 self.map.get(tool_name)
2334 }
2335
2336 pub fn is_read_only(&self, tool_name: &str) -> bool {
2341 self.map.get(tool_name).is_some_and(|a| a.read_only_hint)
2342 }
2343
2344 pub fn is_destructive(&self, tool_name: &str) -> bool {
2349 self.map.get(tool_name).is_none_or(|a| a.destructive_hint)
2350 }
2351
2352 pub fn is_idempotent(&self, tool_name: &str) -> bool {
2357 self.map.get(tool_name).is_some_and(|a| a.idempotent_hint)
2358 }
2359}
2360
2361#[derive(Debug, Clone)]
2383pub struct RouterRequest {
2384 pub id: RequestId,
2386 pub inner: McpRequest,
2388 pub extensions: Extensions,
2390}
2391
2392impl RouterRequest {
2393 pub fn new(id: RequestId, inner: McpRequest) -> Self {
2395 Self {
2396 id,
2397 inner,
2398 extensions: Extensions::new(),
2399 }
2400 }
2401
2402 pub fn with_inner(self, inner: McpRequest) -> Self {
2408 Self {
2409 id: self.id,
2410 inner,
2411 extensions: self.extensions,
2412 }
2413 }
2414
2415 pub fn with_id_and_inner(self, id: RequestId, inner: McpRequest) -> Self {
2421 Self {
2422 id,
2423 inner,
2424 extensions: self.extensions,
2425 }
2426 }
2427
2428 pub fn clone_with_inner(&self, inner: McpRequest) -> Self {
2436 Self {
2437 id: self.id.clone(),
2438 inner,
2439 extensions: self.extensions.clone(),
2440 }
2441 }
2442}
2443
2444#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
2446pub struct RouterResponse {
2447 pub id: RequestId,
2449 pub inner: std::result::Result<McpResponse, JsonRpcError>,
2451}
2452
2453impl RouterResponse {
2454 pub fn is_error(&self) -> bool {
2470 self.inner.is_err()
2471 }
2472
2473 pub fn into_jsonrpc(self) -> JsonRpcResponse {
2475 match self.inner {
2476 Ok(response) => match serde_json::to_value(response) {
2477 Ok(result) => JsonRpcResponse::result(self.id, result),
2478 Err(e) => {
2479 tracing::error!(error = %e, "Failed to serialize response");
2480 JsonRpcResponse::error(
2481 Some(self.id),
2482 JsonRpcError::internal_error(format!("Serialization error: {}", e)),
2483 )
2484 }
2485 },
2486 Err(error) => JsonRpcResponse::error(Some(self.id), error),
2487 }
2488 }
2489}
2490
2491impl Service<RouterRequest> for McpRouter {
2492 type Response = RouterResponse;
2493 type Error = std::convert::Infallible; type Future =
2495 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
2496
2497 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
2498 Poll::Ready(Ok(()))
2499 }
2500
2501 fn call(&mut self, req: RouterRequest) -> Self::Future {
2502 let router = self.clone();
2503 let request_id = req.id.clone();
2504 Box::pin(async move {
2505 let result = router.handle(req.id, req.inner).await;
2506 router.complete_request(&request_id);
2508 Ok(RouterResponse {
2509 id: request_id,
2510 inner: result.map_err(|e| match e {
2515 Error::JsonRpc(err) => err,
2516 Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
2517 e => JsonRpcError::internal_error(e.to_string()),
2518 }),
2519 })
2520 })
2521 }
2522}
2523
2524#[cfg(test)]
2525mod tests {
2526 use super::*;
2527 use crate::extract::{Context, Json};
2528 use crate::jsonrpc::JsonRpcService;
2529 use crate::tool::ToolBuilder;
2530 use schemars::JsonSchema;
2531 use serde::Deserialize;
2532 use tower::ServiceExt;
2533
2534 #[derive(Debug, Deserialize, JsonSchema)]
2535 struct AddInput {
2536 a: i64,
2537 b: i64,
2538 }
2539
2540 async fn init_router(router: &mut McpRouter) {
2542 let init_req = RouterRequest {
2544 id: RequestId::Number(0),
2545 inner: McpRequest::Initialize(InitializeParams {
2546 protocol_version: "2025-11-25".to_string(),
2547 capabilities: ClientCapabilities {
2548 roots: None,
2549 sampling: None,
2550 elicitation: None,
2551 tasks: None,
2552 experimental: None,
2553 extensions: None,
2554 },
2555 client_info: Implementation {
2556 name: "test".to_string(),
2557 version: "1.0".to_string(),
2558 ..Default::default()
2559 },
2560 meta: None,
2561 }),
2562 extensions: Extensions::new(),
2563 };
2564 let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
2565 router.handle_notification(McpNotification::Initialized);
2567 }
2568
2569 #[tokio::test]
2570 async fn test_router_list_tools() {
2571 let add_tool = ToolBuilder::new("add")
2572 .description("Add two numbers")
2573 .handler(|input: AddInput| async move {
2574 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2575 })
2576 .build();
2577
2578 let mut router = McpRouter::new().tool(add_tool);
2579
2580 init_router(&mut router).await;
2582
2583 let req = RouterRequest {
2584 id: RequestId::Number(1),
2585 inner: McpRequest::ListTools(ListToolsParams::default()),
2586 extensions: Extensions::new(),
2587 };
2588
2589 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2590
2591 match resp.inner {
2592 Ok(McpResponse::ListTools(result)) => {
2593 assert_eq!(result.tools.len(), 1);
2594 assert_eq!(result.tools[0].name, "add");
2595 }
2596 _ => panic!("Expected ListTools response"),
2597 }
2598 }
2599
2600 #[tokio::test]
2601 async fn test_router_call_tool() {
2602 let add_tool = ToolBuilder::new("add")
2603 .description("Add two numbers")
2604 .handler(|input: AddInput| async move {
2605 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2606 })
2607 .build();
2608
2609 let mut router = McpRouter::new().tool(add_tool);
2610
2611 init_router(&mut router).await;
2613
2614 let req = RouterRequest {
2615 id: RequestId::Number(1),
2616 inner: McpRequest::CallTool(CallToolParams {
2617 name: "add".to_string(),
2618 arguments: serde_json::json!({"a": 2, "b": 3}),
2619 meta: None,
2620 task: None,
2621 }),
2622 extensions: Extensions::new(),
2623 };
2624
2625 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2626
2627 match resp.inner {
2628 Ok(McpResponse::CallTool(result)) => {
2629 assert!(!result.is_error);
2630 match &result.content[0] {
2632 Content::Text { text, .. } => assert_eq!(text, "5"),
2633 _ => panic!("Expected text content"),
2634 }
2635 }
2636 _ => panic!("Expected CallTool response"),
2637 }
2638 }
2639
2640 async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2642 let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2643 "protocolVersion": "2025-11-25",
2644 "capabilities": {},
2645 "clientInfo": { "name": "test", "version": "1.0" }
2646 }));
2647 let _ = service.call_single(init_req).await.unwrap();
2648 router.handle_notification(McpNotification::Initialized);
2649 }
2650
2651 #[tokio::test]
2652 async fn test_jsonrpc_service() {
2653 let add_tool = ToolBuilder::new("add")
2654 .description("Add two numbers")
2655 .handler(|input: AddInput| async move {
2656 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2657 })
2658 .build();
2659
2660 let router = McpRouter::new().tool(add_tool);
2661 let mut service = JsonRpcService::new(router.clone());
2662
2663 init_jsonrpc_service(&mut service, &router).await;
2665
2666 let req = JsonRpcRequest::new(1, "tools/list");
2667
2668 let resp = service.call_single(req).await.unwrap();
2669
2670 match resp {
2671 JsonRpcResponse::Result(r) => {
2672 assert_eq!(r.id, RequestId::Number(1));
2673 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2674 assert_eq!(tools.len(), 1);
2675 }
2676 JsonRpcResponse::Error(_) => panic!("Expected success response"),
2677 _ => panic!("unexpected response variant"),
2678 }
2679 }
2680
2681 #[tokio::test]
2682 async fn test_batch_request() {
2683 let add_tool = ToolBuilder::new("add")
2684 .description("Add two numbers")
2685 .handler(|input: AddInput| async move {
2686 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2687 })
2688 .build();
2689
2690 let router = McpRouter::new().tool(add_tool);
2691 let mut service = JsonRpcService::new(router.clone());
2692
2693 init_jsonrpc_service(&mut service, &router).await;
2695
2696 let requests = vec![
2698 JsonRpcRequest::new(1, "tools/list"),
2699 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2700 "name": "add",
2701 "arguments": {"a": 10, "b": 20}
2702 })),
2703 JsonRpcRequest::new(3, "ping"),
2704 ];
2705
2706 let responses = service.call_batch(requests).await.unwrap();
2707
2708 assert_eq!(responses.len(), 3);
2709
2710 match &responses[0] {
2712 JsonRpcResponse::Result(r) => {
2713 assert_eq!(r.id, RequestId::Number(1));
2714 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2715 assert_eq!(tools.len(), 1);
2716 }
2717 JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2718 _ => panic!("unexpected response variant"),
2719 }
2720
2721 match &responses[1] {
2723 JsonRpcResponse::Result(r) => {
2724 assert_eq!(r.id, RequestId::Number(2));
2725 let content = r.result.get("content").unwrap().as_array().unwrap();
2726 let text = content[0].get("text").unwrap().as_str().unwrap();
2727 assert_eq!(text, "30");
2728 }
2729 JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2730 _ => panic!("unexpected response variant"),
2731 }
2732
2733 match &responses[2] {
2735 JsonRpcResponse::Result(r) => {
2736 assert_eq!(r.id, RequestId::Number(3));
2737 }
2738 JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2739 _ => panic!("unexpected response variant"),
2740 }
2741 }
2742
2743 #[tokio::test]
2744 async fn test_empty_batch_error() {
2745 let router = McpRouter::new();
2746 let mut service = JsonRpcService::new(router);
2747
2748 let result = service.call_batch(vec![]).await;
2749 assert!(result.is_err());
2750 }
2751
2752 #[tokio::test]
2757 async fn test_progress_token_extraction() {
2758 use crate::context::{ServerNotification, notification_channel};
2759 use crate::protocol::ProgressToken;
2760 use std::sync::Arc;
2761 use std::sync::atomic::{AtomicBool, Ordering};
2762
2763 let progress_reported = Arc::new(AtomicBool::new(false));
2765 let progress_ref = progress_reported.clone();
2766
2767 let tool = ToolBuilder::new("progress_tool")
2769 .description("Tool that reports progress")
2770 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2771 let reported = progress_ref.clone();
2772 async move {
2773 ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
2775 .await;
2776 reported.store(true, Ordering::SeqCst);
2777 Ok(CallToolResult::text("done"))
2778 }
2779 })
2780 .build();
2781
2782 let (tx, mut rx) = notification_channel(10);
2784 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2785 let mut service = JsonRpcService::new(router.clone());
2786
2787 init_jsonrpc_service(&mut service, &router).await;
2789
2790 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2792 "name": "progress_tool",
2793 "arguments": {"a": 1, "b": 2},
2794 "_meta": {
2795 "progressToken": "test-token-123"
2796 }
2797 }));
2798
2799 let resp = service.call_single(req).await.unwrap();
2800
2801 match resp {
2803 JsonRpcResponse::Result(_) => {}
2804 JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
2805 _ => panic!("unexpected response variant"),
2806 }
2807
2808 assert!(progress_reported.load(Ordering::SeqCst));
2810
2811 let notification = rx.try_recv().expect("Expected progress notification");
2813 match notification {
2814 ServerNotification::Progress(params) => {
2815 assert_eq!(
2816 params.progress_token,
2817 ProgressToken::String("test-token-123".to_string())
2818 );
2819 assert_eq!(params.progress, 50.0);
2820 assert_eq!(params.total, Some(100.0));
2821 assert_eq!(params.message.as_deref(), Some("Halfway"));
2822 }
2823 _ => panic!("Expected Progress notification"),
2824 }
2825 }
2826
2827 #[tokio::test]
2828 async fn test_tool_call_without_progress_token() {
2829 use crate::context::notification_channel;
2830 use std::sync::Arc;
2831 use std::sync::atomic::{AtomicBool, Ordering};
2832
2833 let progress_attempted = Arc::new(AtomicBool::new(false));
2834 let progress_ref = progress_attempted.clone();
2835
2836 let tool = ToolBuilder::new("no_token_tool")
2837 .description("Tool that tries to report progress without token")
2838 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2839 let attempted = progress_ref.clone();
2840 async move {
2841 ctx.report_progress(50.0, Some(100.0), None).await;
2843 attempted.store(true, Ordering::SeqCst);
2844 Ok(CallToolResult::text("done"))
2845 }
2846 })
2847 .build();
2848
2849 let (tx, mut rx) = notification_channel(10);
2850 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2851 let mut service = JsonRpcService::new(router.clone());
2852
2853 init_jsonrpc_service(&mut service, &router).await;
2854
2855 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2857 "name": "no_token_tool",
2858 "arguments": {"a": 1, "b": 2}
2859 }));
2860
2861 let resp = service.call_single(req).await.unwrap();
2862 assert!(matches!(resp, JsonRpcResponse::Result(_)));
2863
2864 assert!(progress_attempted.load(Ordering::SeqCst));
2866
2867 assert!(rx.try_recv().is_err());
2869 }
2870
2871 #[tokio::test]
2872 async fn test_batch_errors_returned_not_dropped() {
2873 let add_tool = ToolBuilder::new("add")
2874 .description("Add two numbers")
2875 .handler(|input: AddInput| async move {
2876 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2877 })
2878 .build();
2879
2880 let router = McpRouter::new().tool(add_tool);
2881 let mut service = JsonRpcService::new(router.clone());
2882
2883 init_jsonrpc_service(&mut service, &router).await;
2884
2885 let requests = vec![
2887 JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2889 "name": "add",
2890 "arguments": {"a": 10, "b": 20}
2891 })),
2892 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2894 "name": "nonexistent_tool",
2895 "arguments": {}
2896 })),
2897 JsonRpcRequest::new(3, "ping"),
2899 ];
2900
2901 let responses = service.call_batch(requests).await.unwrap();
2902
2903 assert_eq!(responses.len(), 3);
2905
2906 match &responses[0] {
2908 JsonRpcResponse::Result(r) => {
2909 assert_eq!(r.id, RequestId::Number(1));
2910 }
2911 JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
2912 _ => panic!("unexpected response variant"),
2913 }
2914
2915 match &responses[1] {
2917 JsonRpcResponse::Error(e) => {
2918 assert_eq!(e.id, Some(RequestId::Number(2)));
2919 assert!(e.error.message.contains("not found") || e.error.code == -32601);
2921 }
2922 JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
2923 _ => panic!("unexpected response variant"),
2924 }
2925
2926 match &responses[2] {
2928 JsonRpcResponse::Result(r) => {
2929 assert_eq!(r.id, RequestId::Number(3));
2930 }
2931 JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
2932 _ => panic!("unexpected response variant"),
2933 }
2934 }
2935
2936 #[tokio::test]
2941 async fn test_list_resource_templates() {
2942 use crate::resource::ResourceTemplateBuilder;
2943 use std::collections::HashMap;
2944
2945 let template = ResourceTemplateBuilder::new("file:///{path}")
2946 .name("Project Files")
2947 .description("Access project files")
2948 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2949 Ok(ReadResourceResult {
2950 contents: vec![ResourceContent {
2951 uri,
2952 mime_type: None,
2953 text: None,
2954 blob: None,
2955 meta: None,
2956 }],
2957 meta: None,
2958 })
2959 });
2960
2961 let mut router = McpRouter::new().resource_template(template);
2962
2963 init_router(&mut router).await;
2965
2966 let req = RouterRequest {
2967 id: RequestId::Number(1),
2968 inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
2969 extensions: Extensions::new(),
2970 };
2971
2972 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2973
2974 match resp.inner {
2975 Ok(McpResponse::ListResourceTemplates(result)) => {
2976 assert_eq!(result.resource_templates.len(), 1);
2977 assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
2978 assert_eq!(result.resource_templates[0].name, "Project Files");
2979 }
2980 _ => panic!("Expected ListResourceTemplates response"),
2981 }
2982 }
2983
2984 #[tokio::test]
2985 async fn test_read_resource_via_template() {
2986 use crate::resource::ResourceTemplateBuilder;
2987 use std::collections::HashMap;
2988
2989 let template = ResourceTemplateBuilder::new("db://users/{id}")
2990 .name("User Records")
2991 .handler(|uri: String, vars: HashMap<String, String>| async move {
2992 let id = vars.get("id").unwrap().clone();
2993 Ok(ReadResourceResult {
2994 contents: vec![ResourceContent {
2995 uri,
2996 mime_type: Some("application/json".to_string()),
2997 text: Some(format!(r#"{{"id": "{}"}}"#, id)),
2998 blob: None,
2999 meta: None,
3000 }],
3001 meta: None,
3002 })
3003 });
3004
3005 let mut router = McpRouter::new().resource_template(template);
3006
3007 init_router(&mut router).await;
3009
3010 let req = RouterRequest {
3012 id: RequestId::Number(1),
3013 inner: McpRequest::ReadResource(ReadResourceParams {
3014 uri: "db://users/123".to_string(),
3015 meta: None,
3016 }),
3017 extensions: Extensions::new(),
3018 };
3019
3020 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3021
3022 match resp.inner {
3023 Ok(McpResponse::ReadResource(result)) => {
3024 assert_eq!(result.contents.len(), 1);
3025 assert_eq!(result.contents[0].uri, "db://users/123");
3026 assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
3027 }
3028 _ => panic!("Expected ReadResource response"),
3029 }
3030 }
3031
3032 #[tokio::test]
3033 async fn test_static_resource_takes_precedence_over_template() {
3034 use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
3035 use std::collections::HashMap;
3036
3037 let template = ResourceTemplateBuilder::new("file:///{path}")
3039 .name("Files Template")
3040 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3041 Ok(ReadResourceResult {
3042 contents: vec![ResourceContent {
3043 uri,
3044 mime_type: None,
3045 text: Some("from template".to_string()),
3046 blob: None,
3047 meta: None,
3048 }],
3049 meta: None,
3050 })
3051 });
3052
3053 let static_resource = ResourceBuilder::new("file:///README.md")
3055 .name("README")
3056 .text("from static resource");
3057
3058 let mut router = McpRouter::new()
3059 .resource_template(template)
3060 .resource(static_resource);
3061
3062 init_router(&mut router).await;
3064
3065 let req = RouterRequest {
3067 id: RequestId::Number(1),
3068 inner: McpRequest::ReadResource(ReadResourceParams {
3069 uri: "file:///README.md".to_string(),
3070 meta: None,
3071 }),
3072 extensions: Extensions::new(),
3073 };
3074
3075 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3076
3077 match resp.inner {
3078 Ok(McpResponse::ReadResource(result)) => {
3079 assert_eq!(
3081 result.contents[0].text.as_deref(),
3082 Some("from static resource")
3083 );
3084 }
3085 _ => panic!("Expected ReadResource response"),
3086 }
3087 }
3088
3089 #[tokio::test]
3090 async fn test_resource_not_found_when_no_match() {
3091 use crate::resource::ResourceTemplateBuilder;
3092 use std::collections::HashMap;
3093
3094 let template = ResourceTemplateBuilder::new("db://users/{id}")
3095 .name("Users")
3096 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3097 Ok(ReadResourceResult {
3098 contents: vec![ResourceContent {
3099 uri,
3100 mime_type: None,
3101 text: None,
3102 blob: None,
3103 meta: None,
3104 }],
3105 meta: None,
3106 })
3107 });
3108
3109 let mut router = McpRouter::new().resource_template(template);
3110
3111 init_router(&mut router).await;
3113
3114 let req = RouterRequest {
3116 id: RequestId::Number(1),
3117 inner: McpRequest::ReadResource(ReadResourceParams {
3118 uri: "db://posts/123".to_string(),
3119 meta: None,
3120 }),
3121 extensions: Extensions::new(),
3122 };
3123
3124 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3125
3126 match resp.inner {
3127 Err(err) => {
3128 assert!(err.message.contains("not found"));
3129 }
3130 Ok(_) => panic!("Expected error for non-matching URI"),
3131 }
3132 }
3133
3134 #[tokio::test]
3135 async fn test_capabilities_include_resources_with_only_templates() {
3136 use crate::resource::ResourceTemplateBuilder;
3137 use std::collections::HashMap;
3138
3139 let template = ResourceTemplateBuilder::new("file:///{path}")
3140 .name("Files")
3141 .handler(|uri: String, _vars: HashMap<String, String>| async move {
3142 Ok(ReadResourceResult {
3143 contents: vec![ResourceContent {
3144 uri,
3145 mime_type: None,
3146 text: None,
3147 blob: None,
3148 meta: None,
3149 }],
3150 meta: None,
3151 })
3152 });
3153
3154 let mut router = McpRouter::new().resource_template(template);
3155
3156 let init_req = RouterRequest {
3158 id: RequestId::Number(0),
3159 inner: McpRequest::Initialize(InitializeParams {
3160 protocol_version: "2025-11-25".to_string(),
3161 capabilities: ClientCapabilities {
3162 roots: None,
3163 sampling: None,
3164 elicitation: None,
3165 tasks: None,
3166 experimental: None,
3167 extensions: None,
3168 },
3169 client_info: Implementation {
3170 name: "test".to_string(),
3171 version: "1.0".to_string(),
3172 ..Default::default()
3173 },
3174 meta: None,
3175 }),
3176 extensions: Extensions::new(),
3177 };
3178 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3179
3180 match resp.inner {
3181 Ok(McpResponse::Initialize(result)) => {
3182 assert!(result.capabilities.resources.is_some());
3184 }
3185 _ => panic!("Expected Initialize response"),
3186 }
3187 }
3188
3189 #[tokio::test]
3194 async fn test_log_sends_notification() {
3195 use crate::context::notification_channel;
3196
3197 let (tx, mut rx) = notification_channel(10);
3198 let router = McpRouter::new().with_notification_sender(tx);
3199
3200 let sent = router.log_info("Test message");
3202 assert!(sent);
3203
3204 let notification = rx.try_recv().unwrap();
3206 match notification {
3207 ServerNotification::LogMessage(params) => {
3208 assert_eq!(params.level, LogLevel::Info);
3209 let data = params.data;
3210 assert_eq!(
3211 data.get("message").unwrap().as_str().unwrap(),
3212 "Test message"
3213 );
3214 }
3215 _ => panic!("Expected LogMessage notification"),
3216 }
3217 }
3218
3219 #[tokio::test]
3220 async fn test_log_with_custom_params() {
3221 use crate::context::notification_channel;
3222
3223 let (tx, mut rx) = notification_channel(10);
3224 let router = McpRouter::new().with_notification_sender(tx);
3225
3226 let params = LoggingMessageParams::new(
3228 LogLevel::Error,
3229 serde_json::json!({
3230 "error": "Connection failed",
3231 "host": "localhost"
3232 }),
3233 )
3234 .with_logger("database");
3235
3236 let sent = router.log(params);
3237 assert!(sent);
3238
3239 let notification = rx.try_recv().unwrap();
3240 match notification {
3241 ServerNotification::LogMessage(params) => {
3242 assert_eq!(params.level, LogLevel::Error);
3243 assert_eq!(params.logger.as_deref(), Some("database"));
3244 let data = params.data;
3245 assert_eq!(
3246 data.get("error").unwrap().as_str().unwrap(),
3247 "Connection failed"
3248 );
3249 }
3250 _ => panic!("Expected LogMessage notification"),
3251 }
3252 }
3253
3254 #[tokio::test]
3255 async fn test_log_without_channel_returns_false() {
3256 let router = McpRouter::new();
3258
3259 assert!(!router.log_info("Test"));
3261 assert!(!router.log_warning("Test"));
3262 assert!(!router.log_error("Test"));
3263 assert!(!router.log_debug("Test"));
3264 }
3265
3266 #[tokio::test]
3267 async fn test_logging_capability_with_channel() {
3268 use crate::context::notification_channel;
3269
3270 let (tx, _rx) = notification_channel(10);
3271 let mut router = McpRouter::new().with_notification_sender(tx);
3272
3273 let init_req = RouterRequest {
3275 id: RequestId::Number(0),
3276 inner: McpRequest::Initialize(InitializeParams {
3277 protocol_version: "2025-11-25".to_string(),
3278 capabilities: ClientCapabilities {
3279 roots: None,
3280 sampling: None,
3281 elicitation: None,
3282 tasks: None,
3283 experimental: None,
3284 extensions: None,
3285 },
3286 client_info: Implementation {
3287 name: "test".to_string(),
3288 version: "1.0".to_string(),
3289 ..Default::default()
3290 },
3291 meta: None,
3292 }),
3293 extensions: Extensions::new(),
3294 };
3295 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3296
3297 match resp.inner {
3298 Ok(McpResponse::Initialize(result)) => {
3299 assert!(result.capabilities.logging.is_some());
3301 }
3302 _ => panic!("Expected Initialize response"),
3303 }
3304 }
3305
3306 #[tokio::test]
3307 async fn test_no_logging_capability_without_channel() {
3308 let mut router = McpRouter::new();
3309
3310 let init_req = RouterRequest {
3312 id: RequestId::Number(0),
3313 inner: McpRequest::Initialize(InitializeParams {
3314 protocol_version: "2025-11-25".to_string(),
3315 capabilities: ClientCapabilities {
3316 roots: None,
3317 sampling: None,
3318 elicitation: None,
3319 tasks: None,
3320 experimental: None,
3321 extensions: None,
3322 },
3323 client_info: Implementation {
3324 name: "test".to_string(),
3325 version: "1.0".to_string(),
3326 ..Default::default()
3327 },
3328 meta: None,
3329 }),
3330 extensions: Extensions::new(),
3331 };
3332 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3333
3334 match resp.inner {
3335 Ok(McpResponse::Initialize(result)) => {
3336 assert!(result.capabilities.logging.is_none());
3338 }
3339 _ => panic!("Expected Initialize response"),
3340 }
3341 }
3342
3343 #[tokio::test]
3348 async fn test_create_task_via_call_tool() {
3349 let add_tool = ToolBuilder::new("add")
3350 .description("Add two numbers")
3351 .task_support(TaskSupportMode::Optional)
3352 .handler(|input: AddInput| async move {
3353 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3354 })
3355 .build();
3356
3357 let mut router = McpRouter::new().tool(add_tool);
3358 init_router(&mut router).await;
3359
3360 let req = RouterRequest {
3361 id: RequestId::Number(1),
3362 inner: McpRequest::CallTool(CallToolParams {
3363 name: "add".to_string(),
3364 arguments: serde_json::json!({"a": 5, "b": 10}),
3365 meta: None,
3366 task: Some(TaskRequestParams { ttl: None }),
3367 }),
3368 extensions: Extensions::new(),
3369 };
3370
3371 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3372
3373 match resp.inner {
3374 Ok(McpResponse::CreateTask(result)) => {
3375 assert!(result.task.task_id.starts_with("task-"));
3376 assert_eq!(result.task.status, TaskStatus::Working);
3377 }
3378 _ => panic!("Expected CreateTask response"),
3379 }
3380 }
3381
3382 #[tokio::test]
3383 async fn test_list_tasks_empty() {
3384 let mut router = McpRouter::new();
3385 init_router(&mut router).await;
3386
3387 let req = RouterRequest {
3388 id: RequestId::Number(1),
3389 inner: McpRequest::ListTasks(ListTasksParams::default()),
3390 extensions: Extensions::new(),
3391 };
3392
3393 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3394
3395 match resp.inner {
3396 Ok(McpResponse::ListTasks(result)) => {
3397 assert!(result.tasks.is_empty());
3398 }
3399 _ => panic!("Expected ListTasks response"),
3400 }
3401 }
3402
3403 #[tokio::test]
3404 async fn test_task_lifecycle_complete() {
3405 let add_tool = ToolBuilder::new("add")
3406 .description("Add two numbers")
3407 .task_support(TaskSupportMode::Optional)
3408 .handler(|input: AddInput| async move {
3409 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3410 })
3411 .build();
3412
3413 let mut router = McpRouter::new().tool(add_tool);
3414 init_router(&mut router).await;
3415
3416 let req = RouterRequest {
3418 id: RequestId::Number(1),
3419 inner: McpRequest::CallTool(CallToolParams {
3420 name: "add".to_string(),
3421 arguments: serde_json::json!({"a": 7, "b": 8}),
3422 meta: None,
3423 task: Some(TaskRequestParams { ttl: None }),
3424 }),
3425 extensions: Extensions::new(),
3426 };
3427
3428 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3429 let task_id = match resp.inner {
3430 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3431 _ => panic!("Expected CreateTask response"),
3432 };
3433
3434 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3436
3437 let req = RouterRequest {
3439 id: RequestId::Number(2),
3440 inner: McpRequest::GetTaskResult(GetTaskResultParams {
3441 task_id: task_id.clone(),
3442 meta: None,
3443 }),
3444 extensions: Extensions::new(),
3445 };
3446
3447 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3448
3449 match resp.inner {
3450 Ok(McpResponse::GetTaskResult(result)) => {
3451 assert!(result.meta.is_some());
3453 match &result.content[0] {
3455 Content::Text { text, .. } => assert_eq!(text, "15"),
3456 _ => panic!("Expected text content"),
3457 }
3458 }
3459 _ => panic!("Expected GetTaskResult response"),
3460 }
3461 }
3462
3463 #[tokio::test]
3464 async fn test_task_cancellation() {
3465 let slow_tool = ToolBuilder::new("slow")
3467 .description("Slow tool")
3468 .task_support(TaskSupportMode::Optional)
3469 .handler(|_input: serde_json::Value| async move {
3470 tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
3471 Ok(CallToolResult::text("done"))
3472 })
3473 .build();
3474
3475 let mut router = McpRouter::new().tool(slow_tool);
3476 init_router(&mut router).await;
3477
3478 let req = RouterRequest {
3480 id: RequestId::Number(1),
3481 inner: McpRequest::CallTool(CallToolParams {
3482 name: "slow".to_string(),
3483 arguments: serde_json::json!({}),
3484 meta: None,
3485 task: Some(TaskRequestParams { ttl: None }),
3486 }),
3487 extensions: Extensions::new(),
3488 };
3489
3490 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3491 let task_id = match resp.inner {
3492 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3493 _ => panic!("Expected CreateTask response"),
3494 };
3495
3496 let req = RouterRequest {
3498 id: RequestId::Number(2),
3499 inner: McpRequest::CancelTask(CancelTaskParams {
3500 task_id: task_id.clone(),
3501 reason: Some("Test cancellation".to_string()),
3502 meta: None,
3503 }),
3504 extensions: Extensions::new(),
3505 };
3506
3507 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3508
3509 match resp.inner {
3510 Ok(McpResponse::CancelTask(task_obj)) => {
3511 assert_eq!(task_obj.status, TaskStatus::Cancelled);
3512 }
3513 _ => panic!("Expected CancelTask response"),
3514 }
3515 }
3516
3517 #[tokio::test]
3518 async fn test_get_task_info() {
3519 let add_tool = ToolBuilder::new("add")
3520 .description("Add two numbers")
3521 .task_support(TaskSupportMode::Optional)
3522 .handler(|input: AddInput| async move {
3523 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3524 })
3525 .build();
3526
3527 let mut router = McpRouter::new().tool(add_tool);
3528 init_router(&mut router).await;
3529
3530 let req = RouterRequest {
3532 id: RequestId::Number(1),
3533 inner: McpRequest::CallTool(CallToolParams {
3534 name: "add".to_string(),
3535 arguments: serde_json::json!({"a": 1, "b": 2}),
3536 meta: None,
3537 task: Some(TaskRequestParams { ttl: Some(600_000) }),
3538 }),
3539 extensions: Extensions::new(),
3540 };
3541
3542 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3543 let task_id = match resp.inner {
3544 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3545 _ => panic!("Expected CreateTask response"),
3546 };
3547
3548 let req = RouterRequest {
3550 id: RequestId::Number(2),
3551 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3552 task_id: task_id.clone(),
3553 meta: None,
3554 }),
3555 extensions: Extensions::new(),
3556 };
3557
3558 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3559
3560 match resp.inner {
3561 Ok(McpResponse::GetTaskInfo(info)) => {
3562 assert_eq!(info.task_id, task_id);
3563 assert!(info.created_at.contains('T')); assert_eq!(info.ttl, Some(600_000));
3565 }
3566 _ => panic!("Expected GetTaskInfo response"),
3567 }
3568 }
3569
3570 #[tokio::test]
3571 async fn test_task_forbidden_tool_rejects_task_params() {
3572 let tool = ToolBuilder::new("sync_only")
3573 .description("Sync only tool")
3574 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3575 .build();
3576
3577 let mut router = McpRouter::new().tool(tool);
3578 init_router(&mut router).await;
3579
3580 let req = RouterRequest {
3582 id: RequestId::Number(1),
3583 inner: McpRequest::CallTool(CallToolParams {
3584 name: "sync_only".to_string(),
3585 arguments: serde_json::json!({}),
3586 meta: None,
3587 task: Some(TaskRequestParams { ttl: None }),
3588 }),
3589 extensions: Extensions::new(),
3590 };
3591
3592 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3593
3594 match resp.inner {
3595 Err(e) => {
3596 assert!(e.message.contains("does not support async tasks"));
3597 }
3598 _ => panic!("Expected error response"),
3599 }
3600 }
3601
3602 #[tokio::test]
3603 async fn test_get_nonexistent_task() {
3604 let mut router = McpRouter::new();
3605 init_router(&mut router).await;
3606
3607 let req = RouterRequest {
3608 id: RequestId::Number(1),
3609 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3610 task_id: "task-999".to_string(),
3611 meta: None,
3612 }),
3613 extensions: Extensions::new(),
3614 };
3615
3616 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3617
3618 match resp.inner {
3619 Err(e) => {
3620 assert!(e.message.contains("not found"));
3621 }
3622 _ => panic!("Expected error response"),
3623 }
3624 }
3625
3626 #[tokio::test]
3631 async fn test_subscribe_to_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
3651 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3652
3653 match resp.inner {
3654 Ok(McpResponse::SubscribeResource(_)) => {
3655 assert!(router.is_subscribed("file:///test.txt"));
3657 }
3658 _ => panic!("Expected SubscribeResource response"),
3659 }
3660 }
3661
3662 #[tokio::test]
3663 async fn test_unsubscribe_from_resource() {
3664 use crate::resource::ResourceBuilder;
3665
3666 let resource = ResourceBuilder::new("file:///test.txt")
3667 .name("Test File")
3668 .text("Hello");
3669
3670 let mut router = McpRouter::new().resource(resource);
3671 init_router(&mut router).await;
3672
3673 let req = RouterRequest {
3675 id: RequestId::Number(1),
3676 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3677 uri: "file:///test.txt".to_string(),
3678 meta: None,
3679 }),
3680 extensions: Extensions::new(),
3681 };
3682 let _ = router.ready().await.unwrap().call(req).await.unwrap();
3683 assert!(router.is_subscribed("file:///test.txt"));
3684
3685 let req = RouterRequest {
3687 id: RequestId::Number(2),
3688 inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3689 uri: "file:///test.txt".to_string(),
3690 meta: None,
3691 }),
3692 extensions: Extensions::new(),
3693 };
3694
3695 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3696
3697 match resp.inner {
3698 Ok(McpResponse::UnsubscribeResource(_)) => {
3699 assert!(!router.is_subscribed("file:///test.txt"));
3701 }
3702 _ => panic!("Expected UnsubscribeResource response"),
3703 }
3704 }
3705
3706 #[tokio::test]
3707 async fn test_subscribe_nonexistent_resource() {
3708 let mut router = McpRouter::new();
3709 init_router(&mut router).await;
3710
3711 let req = RouterRequest {
3712 id: RequestId::Number(1),
3713 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3714 uri: "file:///nonexistent.txt".to_string(),
3715 meta: None,
3716 }),
3717 extensions: Extensions::new(),
3718 };
3719
3720 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3721
3722 match resp.inner {
3723 Err(e) => {
3724 assert!(e.message.contains("not found"));
3725 }
3726 _ => panic!("Expected error response"),
3727 }
3728 }
3729
3730 #[tokio::test]
3731 async fn test_notify_resource_updated() {
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 router.subscribe("file:///test.txt");
3747
3748 let sent = router.notify_resource_updated("file:///test.txt");
3750 assert!(sent);
3751
3752 let notification = rx.try_recv().unwrap();
3754 match notification {
3755 ServerNotification::ResourceUpdated { uri } => {
3756 assert_eq!(uri, "file:///test.txt");
3757 }
3758 _ => panic!("Expected ResourceUpdated notification"),
3759 }
3760 }
3761
3762 #[tokio::test]
3763 async fn test_notify_resource_updated_not_subscribed() {
3764 use crate::context::notification_channel;
3765 use crate::resource::ResourceBuilder;
3766
3767 let (tx, mut rx) = notification_channel(10);
3768
3769 let resource = ResourceBuilder::new("file:///test.txt")
3770 .name("Test File")
3771 .text("Hello");
3772
3773 let router = McpRouter::new()
3774 .resource(resource)
3775 .with_notification_sender(tx);
3776
3777 let sent = router.notify_resource_updated("file:///test.txt");
3779 assert!(!sent); assert!(rx.try_recv().is_err());
3783 }
3784
3785 #[tokio::test]
3786 async fn test_notify_resources_list_changed() {
3787 use crate::context::notification_channel;
3788
3789 let (tx, mut rx) = notification_channel(10);
3790 let router = McpRouter::new().with_notification_sender(tx);
3791
3792 let sent = router.notify_resources_list_changed();
3793 assert!(sent);
3794
3795 let notification = rx.try_recv().unwrap();
3796 match notification {
3797 ServerNotification::ResourcesListChanged => {}
3798 _ => panic!("Expected ResourcesListChanged notification"),
3799 }
3800 }
3801
3802 #[tokio::test]
3803 async fn test_subscribed_uris() {
3804 use crate::resource::ResourceBuilder;
3805
3806 let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
3807
3808 let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
3809
3810 let router = McpRouter::new().resource(resource1).resource(resource2);
3811
3812 router.subscribe("file:///a.txt");
3814 router.subscribe("file:///b.txt");
3815
3816 let uris = router.subscribed_uris();
3817 assert_eq!(uris.len(), 2);
3818 assert!(uris.contains(&"file:///a.txt".to_string()));
3819 assert!(uris.contains(&"file:///b.txt".to_string()));
3820 }
3821
3822 #[tokio::test]
3823 async fn test_subscription_capability_advertised() {
3824 use crate::resource::ResourceBuilder;
3825
3826 let resource = ResourceBuilder::new("file:///test.txt")
3827 .name("Test")
3828 .text("Hello");
3829
3830 let mut router = McpRouter::new().resource(resource);
3831
3832 let init_req = RouterRequest {
3834 id: RequestId::Number(0),
3835 inner: McpRequest::Initialize(InitializeParams {
3836 protocol_version: "2025-11-25".to_string(),
3837 capabilities: ClientCapabilities {
3838 roots: None,
3839 sampling: None,
3840 elicitation: None,
3841 tasks: None,
3842 experimental: None,
3843 extensions: None,
3844 },
3845 client_info: Implementation {
3846 name: "test".to_string(),
3847 version: "1.0".to_string(),
3848 ..Default::default()
3849 },
3850 meta: None,
3851 }),
3852 extensions: Extensions::new(),
3853 };
3854 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3855
3856 match resp.inner {
3857 Ok(McpResponse::Initialize(result)) => {
3858 let resources_cap = result.capabilities.resources.unwrap();
3860 assert!(resources_cap.subscribe);
3861 }
3862 _ => panic!("Expected Initialize response"),
3863 }
3864 }
3865
3866 #[tokio::test]
3867 async fn test_completion_handler() {
3868 let router = McpRouter::new()
3869 .server_info("test", "1.0")
3870 .completion_handler(|params: CompleteParams| async move {
3871 let prefix = ¶ms.argument.value;
3873 let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
3874 .into_iter()
3875 .filter(|s| s.starts_with(prefix))
3876 .map(String::from)
3877 .collect();
3878 Ok(CompleteResult::new(suggestions))
3879 });
3880
3881 let init_req = RouterRequest {
3883 id: RequestId::Number(0),
3884 inner: McpRequest::Initialize(InitializeParams {
3885 protocol_version: "2025-11-25".to_string(),
3886 capabilities: ClientCapabilities::default(),
3887 client_info: Implementation {
3888 name: "test".to_string(),
3889 version: "1.0".to_string(),
3890 ..Default::default()
3891 },
3892 meta: None,
3893 }),
3894 extensions: Extensions::new(),
3895 };
3896 let resp = router
3897 .clone()
3898 .ready()
3899 .await
3900 .unwrap()
3901 .call(init_req)
3902 .await
3903 .unwrap();
3904
3905 match resp.inner {
3907 Ok(McpResponse::Initialize(result)) => {
3908 assert!(result.capabilities.completions.is_some());
3909 }
3910 _ => panic!("Expected Initialize response"),
3911 }
3912
3913 router.handle_notification(McpNotification::Initialized);
3915
3916 let complete_req = RouterRequest {
3918 id: RequestId::Number(1),
3919 inner: McpRequest::Complete(CompleteParams {
3920 reference: CompletionReference::prompt("test-prompt"),
3921 argument: CompletionArgument::new("query", "al"),
3922 context: None,
3923 meta: None,
3924 }),
3925 extensions: Extensions::new(),
3926 };
3927 let resp = router
3928 .clone()
3929 .ready()
3930 .await
3931 .unwrap()
3932 .call(complete_req)
3933 .await
3934 .unwrap();
3935
3936 match resp.inner {
3937 Ok(McpResponse::Complete(result)) => {
3938 assert_eq!(result.completion.values, vec!["alpha"]);
3939 }
3940 _ => panic!("Expected Complete response"),
3941 }
3942 }
3943
3944 #[tokio::test]
3945 async fn test_completion_without_handler_returns_empty() {
3946 let router = McpRouter::new().server_info("test", "1.0");
3947
3948 let init_req = RouterRequest {
3950 id: RequestId::Number(0),
3951 inner: McpRequest::Initialize(InitializeParams {
3952 protocol_version: "2025-11-25".to_string(),
3953 capabilities: ClientCapabilities::default(),
3954 client_info: Implementation {
3955 name: "test".to_string(),
3956 version: "1.0".to_string(),
3957 ..Default::default()
3958 },
3959 meta: None,
3960 }),
3961 extensions: Extensions::new(),
3962 };
3963 let resp = router
3964 .clone()
3965 .ready()
3966 .await
3967 .unwrap()
3968 .call(init_req)
3969 .await
3970 .unwrap();
3971
3972 match resp.inner {
3974 Ok(McpResponse::Initialize(result)) => {
3975 assert!(result.capabilities.completions.is_none());
3976 }
3977 _ => panic!("Expected Initialize response"),
3978 }
3979
3980 router.handle_notification(McpNotification::Initialized);
3982
3983 let complete_req = RouterRequest {
3985 id: RequestId::Number(1),
3986 inner: McpRequest::Complete(CompleteParams {
3987 reference: CompletionReference::prompt("test-prompt"),
3988 argument: CompletionArgument::new("query", "al"),
3989 context: None,
3990 meta: None,
3991 }),
3992 extensions: Extensions::new(),
3993 };
3994 let resp = router
3995 .clone()
3996 .ready()
3997 .await
3998 .unwrap()
3999 .call(complete_req)
4000 .await
4001 .unwrap();
4002
4003 match resp.inner {
4004 Ok(McpResponse::Complete(result)) => {
4005 assert!(result.completion.values.is_empty());
4006 }
4007 _ => panic!("Expected Complete response"),
4008 }
4009 }
4010
4011 #[tokio::test]
4012 async fn test_tool_filter_list() {
4013 use crate::filter::CapabilityFilter;
4014 use crate::tool::Tool;
4015
4016 let public_tool = ToolBuilder::new("public")
4017 .description("Public tool")
4018 .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
4019 .build();
4020
4021 let admin_tool = ToolBuilder::new("admin")
4022 .description("Admin tool")
4023 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4024 .build();
4025
4026 let mut router = McpRouter::new()
4027 .tool(public_tool)
4028 .tool(admin_tool)
4029 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
4030
4031 init_router(&mut router).await;
4033
4034 let req = RouterRequest {
4035 id: RequestId::Number(1),
4036 inner: McpRequest::ListTools(ListToolsParams::default()),
4037 extensions: Extensions::new(),
4038 };
4039
4040 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4041
4042 match resp.inner {
4043 Ok(McpResponse::ListTools(result)) => {
4044 assert_eq!(result.tools.len(), 1);
4046 assert_eq!(result.tools[0].name, "public");
4047 }
4048 _ => panic!("Expected ListTools response"),
4049 }
4050 }
4051
4052 #[tokio::test]
4053 async fn test_tool_filter_call_denied() {
4054 use crate::filter::CapabilityFilter;
4055 use crate::tool::Tool;
4056
4057 let admin_tool = ToolBuilder::new("admin")
4058 .description("Admin tool")
4059 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4060 .build();
4061
4062 let mut router = McpRouter::new()
4063 .tool(admin_tool)
4064 .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); init_router(&mut router).await;
4068
4069 let req = RouterRequest {
4070 id: RequestId::Number(1),
4071 inner: McpRequest::CallTool(CallToolParams {
4072 name: "admin".to_string(),
4073 arguments: serde_json::json!({"a": 1, "b": 2}),
4074 meta: None,
4075 task: None,
4076 }),
4077 extensions: Extensions::new(),
4078 };
4079
4080 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4081
4082 match resp.inner {
4084 Err(e) => {
4085 assert_eq!(e.code, -32601); }
4087 _ => panic!("Expected JsonRpc error"),
4088 }
4089 }
4090
4091 #[tokio::test]
4092 async fn test_tool_filter_call_allowed() {
4093 use crate::filter::CapabilityFilter;
4094 use crate::tool::Tool;
4095
4096 let public_tool = ToolBuilder::new("public")
4097 .description("Public tool")
4098 .handler(|input: AddInput| async move {
4099 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
4100 })
4101 .build();
4102
4103 let mut router = McpRouter::new()
4104 .tool(public_tool)
4105 .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); init_router(&mut router).await;
4109
4110 let req = RouterRequest {
4111 id: RequestId::Number(1),
4112 inner: McpRequest::CallTool(CallToolParams {
4113 name: "public".to_string(),
4114 arguments: serde_json::json!({"a": 1, "b": 2}),
4115 meta: None,
4116 task: None,
4117 }),
4118 extensions: Extensions::new(),
4119 };
4120
4121 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4122
4123 match resp.inner {
4124 Ok(McpResponse::CallTool(result)) => {
4125 assert!(!result.is_error);
4126 }
4127 _ => panic!("Expected CallTool response"),
4128 }
4129 }
4130
4131 #[tokio::test]
4132 async fn test_tool_filter_custom_denial() {
4133 use crate::filter::{CapabilityFilter, DenialBehavior};
4134 use crate::tool::Tool;
4135
4136 let admin_tool = ToolBuilder::new("admin")
4137 .description("Admin tool")
4138 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4139 .build();
4140
4141 let mut router = McpRouter::new().tool(admin_tool).tool_filter(
4142 CapabilityFilter::new(|_, _: &Tool| false)
4143 .denial_behavior(DenialBehavior::Unauthorized),
4144 );
4145
4146 init_router(&mut router).await;
4148
4149 let req = RouterRequest {
4150 id: RequestId::Number(1),
4151 inner: McpRequest::CallTool(CallToolParams {
4152 name: "admin".to_string(),
4153 arguments: serde_json::json!({"a": 1, "b": 2}),
4154 meta: None,
4155 task: None,
4156 }),
4157 extensions: Extensions::new(),
4158 };
4159
4160 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4161
4162 match resp.inner {
4164 Err(e) => {
4165 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4167 }
4168 _ => panic!("Expected JsonRpc error"),
4169 }
4170 }
4171
4172 #[tokio::test]
4173 async fn test_resource_filter_list() {
4174 use crate::filter::CapabilityFilter;
4175 use crate::resource::{Resource, ResourceBuilder};
4176
4177 let public_resource = ResourceBuilder::new("file:///public.txt")
4178 .name("Public File")
4179 .text("public content");
4180
4181 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4182 .name("Secret File")
4183 .text("secret content");
4184
4185 let mut router = McpRouter::new()
4186 .resource(public_resource)
4187 .resource(secret_resource)
4188 .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
4189 !r.name.contains("Secret")
4190 }));
4191
4192 init_router(&mut router).await;
4194
4195 let req = RouterRequest {
4196 id: RequestId::Number(1),
4197 inner: McpRequest::ListResources(ListResourcesParams::default()),
4198 extensions: Extensions::new(),
4199 };
4200
4201 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4202
4203 match resp.inner {
4204 Ok(McpResponse::ListResources(result)) => {
4205 assert_eq!(result.resources.len(), 1);
4207 assert_eq!(result.resources[0].name, "Public File");
4208 }
4209 _ => panic!("Expected ListResources response"),
4210 }
4211 }
4212
4213 #[tokio::test]
4214 async fn test_resource_filter_read_denied() {
4215 use crate::filter::CapabilityFilter;
4216 use crate::resource::{Resource, ResourceBuilder};
4217
4218 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4219 .name("Secret File")
4220 .text("secret content");
4221
4222 let mut router = McpRouter::new()
4223 .resource(secret_resource)
4224 .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); init_router(&mut router).await;
4228
4229 let req = RouterRequest {
4230 id: RequestId::Number(1),
4231 inner: McpRequest::ReadResource(ReadResourceParams {
4232 uri: "file:///secret.txt".to_string(),
4233 meta: None,
4234 }),
4235 extensions: Extensions::new(),
4236 };
4237
4238 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4239
4240 match resp.inner {
4242 Err(e) => {
4243 assert_eq!(e.code, -32601); }
4245 _ => panic!("Expected JsonRpc error"),
4246 }
4247 }
4248
4249 #[tokio::test]
4250 async fn test_resource_filter_read_allowed() {
4251 use crate::filter::CapabilityFilter;
4252 use crate::resource::{Resource, ResourceBuilder};
4253
4254 let public_resource = ResourceBuilder::new("file:///public.txt")
4255 .name("Public File")
4256 .text("public content");
4257
4258 let mut router = McpRouter::new()
4259 .resource(public_resource)
4260 .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); init_router(&mut router).await;
4264
4265 let req = RouterRequest {
4266 id: RequestId::Number(1),
4267 inner: McpRequest::ReadResource(ReadResourceParams {
4268 uri: "file:///public.txt".to_string(),
4269 meta: None,
4270 }),
4271 extensions: Extensions::new(),
4272 };
4273
4274 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4275
4276 match resp.inner {
4277 Ok(McpResponse::ReadResource(result)) => {
4278 assert_eq!(result.contents.len(), 1);
4279 assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
4280 }
4281 _ => panic!("Expected ReadResource response"),
4282 }
4283 }
4284
4285 #[tokio::test]
4286 async fn test_resource_filter_custom_denial() {
4287 use crate::filter::{CapabilityFilter, DenialBehavior};
4288 use crate::resource::{Resource, ResourceBuilder};
4289
4290 let secret_resource = ResourceBuilder::new("file:///secret.txt")
4291 .name("Secret File")
4292 .text("secret content");
4293
4294 let mut router = McpRouter::new().resource(secret_resource).resource_filter(
4295 CapabilityFilter::new(|_, _: &Resource| false)
4296 .denial_behavior(DenialBehavior::Unauthorized),
4297 );
4298
4299 init_router(&mut router).await;
4301
4302 let req = RouterRequest {
4303 id: RequestId::Number(1),
4304 inner: McpRequest::ReadResource(ReadResourceParams {
4305 uri: "file:///secret.txt".to_string(),
4306 meta: None,
4307 }),
4308 extensions: Extensions::new(),
4309 };
4310
4311 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4312
4313 match resp.inner {
4315 Err(e) => {
4316 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4318 }
4319 _ => panic!("Expected JsonRpc error"),
4320 }
4321 }
4322
4323 #[tokio::test]
4324 async fn test_prompt_filter_list() {
4325 use crate::filter::CapabilityFilter;
4326 use crate::prompt::{Prompt, PromptBuilder};
4327
4328 let public_prompt = PromptBuilder::new("greeting")
4329 .description("A greeting")
4330 .user_message("Hello!");
4331
4332 let admin_prompt = PromptBuilder::new("system_debug")
4333 .description("Admin prompt")
4334 .user_message("Debug");
4335
4336 let mut router = McpRouter::new()
4337 .prompt(public_prompt)
4338 .prompt(admin_prompt)
4339 .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
4340 !p.name.contains("system")
4341 }));
4342
4343 init_router(&mut router).await;
4345
4346 let req = RouterRequest {
4347 id: RequestId::Number(1),
4348 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4349 extensions: Extensions::new(),
4350 };
4351
4352 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4353
4354 match resp.inner {
4355 Ok(McpResponse::ListPrompts(result)) => {
4356 assert_eq!(result.prompts.len(), 1);
4358 assert_eq!(result.prompts[0].name, "greeting");
4359 }
4360 _ => panic!("Expected ListPrompts response"),
4361 }
4362 }
4363
4364 #[tokio::test]
4365 async fn test_prompt_filter_get_denied() {
4366 use crate::filter::CapabilityFilter;
4367 use crate::prompt::{Prompt, PromptBuilder};
4368 use std::collections::HashMap;
4369
4370 let admin_prompt = PromptBuilder::new("system_debug")
4371 .description("Admin prompt")
4372 .user_message("Debug");
4373
4374 let mut router = McpRouter::new()
4375 .prompt(admin_prompt)
4376 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); init_router(&mut router).await;
4380
4381 let req = RouterRequest {
4382 id: RequestId::Number(1),
4383 inner: McpRequest::GetPrompt(GetPromptParams {
4384 name: "system_debug".to_string(),
4385 arguments: HashMap::new(),
4386 meta: None,
4387 }),
4388 extensions: Extensions::new(),
4389 };
4390
4391 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4392
4393 match resp.inner {
4395 Err(e) => {
4396 assert_eq!(e.code, -32601); }
4398 _ => panic!("Expected JsonRpc error"),
4399 }
4400 }
4401
4402 #[tokio::test]
4403 async fn test_prompt_filter_get_allowed() {
4404 use crate::filter::CapabilityFilter;
4405 use crate::prompt::{Prompt, PromptBuilder};
4406 use std::collections::HashMap;
4407
4408 let public_prompt = PromptBuilder::new("greeting")
4409 .description("A greeting")
4410 .user_message("Hello!");
4411
4412 let mut router = McpRouter::new()
4413 .prompt(public_prompt)
4414 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); init_router(&mut router).await;
4418
4419 let req = RouterRequest {
4420 id: RequestId::Number(1),
4421 inner: McpRequest::GetPrompt(GetPromptParams {
4422 name: "greeting".to_string(),
4423 arguments: HashMap::new(),
4424 meta: None,
4425 }),
4426 extensions: Extensions::new(),
4427 };
4428
4429 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4430
4431 match resp.inner {
4432 Ok(McpResponse::GetPrompt(result)) => {
4433 assert_eq!(result.messages.len(), 1);
4434 }
4435 _ => panic!("Expected GetPrompt response"),
4436 }
4437 }
4438
4439 #[tokio::test]
4440 async fn test_prompt_filter_custom_denial() {
4441 use crate::filter::{CapabilityFilter, DenialBehavior};
4442 use crate::prompt::{Prompt, PromptBuilder};
4443 use std::collections::HashMap;
4444
4445 let admin_prompt = PromptBuilder::new("system_debug")
4446 .description("Admin prompt")
4447 .user_message("Debug");
4448
4449 let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
4450 CapabilityFilter::new(|_, _: &Prompt| false)
4451 .denial_behavior(DenialBehavior::Unauthorized),
4452 );
4453
4454 init_router(&mut router).await;
4456
4457 let req = RouterRequest {
4458 id: RequestId::Number(1),
4459 inner: McpRequest::GetPrompt(GetPromptParams {
4460 name: "system_debug".to_string(),
4461 arguments: HashMap::new(),
4462 meta: None,
4463 }),
4464 extensions: Extensions::new(),
4465 };
4466
4467 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4468
4469 match resp.inner {
4471 Err(e) => {
4472 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
4474 }
4475 _ => panic!("Expected JsonRpc error"),
4476 }
4477 }
4478
4479 #[derive(Debug, Deserialize, JsonSchema)]
4484 struct StringInput {
4485 value: String,
4486 }
4487
4488 #[tokio::test]
4489 async fn test_router_merge_tools() {
4490 let tool_a = ToolBuilder::new("tool_a")
4492 .description("Tool A")
4493 .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
4494 .build();
4495
4496 let router_a = McpRouter::new().tool(tool_a);
4497
4498 let tool_b = ToolBuilder::new("tool_b")
4500 .description("Tool B")
4501 .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
4502 .build();
4503 let tool_c = ToolBuilder::new("tool_c")
4504 .description("Tool C")
4505 .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
4506 .build();
4507
4508 let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
4509
4510 let mut merged = McpRouter::new()
4512 .server_info("merged", "1.0")
4513 .merge(router_a)
4514 .merge(router_b);
4515
4516 init_router(&mut merged).await;
4517
4518 let req = RouterRequest {
4520 id: RequestId::Number(1),
4521 inner: McpRequest::ListTools(ListToolsParams::default()),
4522 extensions: Extensions::new(),
4523 };
4524
4525 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4526
4527 match resp.inner {
4528 Ok(McpResponse::ListTools(result)) => {
4529 assert_eq!(result.tools.len(), 3);
4530 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4531 assert!(names.contains(&"tool_a"));
4532 assert!(names.contains(&"tool_b"));
4533 assert!(names.contains(&"tool_c"));
4534 }
4535 _ => panic!("Expected ListTools response"),
4536 }
4537 }
4538
4539 #[tokio::test]
4540 async fn test_router_merge_overwrites_duplicates() {
4541 let tool_v1 = ToolBuilder::new("shared")
4543 .description("Version 1")
4544 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
4545 .build();
4546
4547 let router_a = McpRouter::new().tool(tool_v1);
4548
4549 let tool_v2 = ToolBuilder::new("shared")
4551 .description("Version 2")
4552 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
4553 .build();
4554
4555 let router_b = McpRouter::new().tool(tool_v2);
4556
4557 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4559
4560 init_router(&mut merged).await;
4561
4562 let req = RouterRequest {
4563 id: RequestId::Number(1),
4564 inner: McpRequest::ListTools(ListToolsParams::default()),
4565 extensions: Extensions::new(),
4566 };
4567
4568 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4569
4570 match resp.inner {
4571 Ok(McpResponse::ListTools(result)) => {
4572 assert_eq!(result.tools.len(), 1);
4573 assert_eq!(result.tools[0].name, "shared");
4574 assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
4575 }
4576 _ => panic!("Expected ListTools response"),
4577 }
4578 }
4579
4580 #[tokio::test]
4581 async fn test_router_merge_resources() {
4582 use crate::resource::ResourceBuilder;
4583
4584 let router_a = McpRouter::new().resource(
4586 ResourceBuilder::new("file:///a.txt")
4587 .name("File A")
4588 .text("content a"),
4589 );
4590
4591 let router_b = McpRouter::new().resource(
4592 ResourceBuilder::new("file:///b.txt")
4593 .name("File B")
4594 .text("content b"),
4595 );
4596
4597 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4598
4599 init_router(&mut merged).await;
4600
4601 let req = RouterRequest {
4602 id: RequestId::Number(1),
4603 inner: McpRequest::ListResources(ListResourcesParams::default()),
4604 extensions: Extensions::new(),
4605 };
4606
4607 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4608
4609 match resp.inner {
4610 Ok(McpResponse::ListResources(result)) => {
4611 assert_eq!(result.resources.len(), 2);
4612 let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4613 assert!(uris.contains(&"file:///a.txt"));
4614 assert!(uris.contains(&"file:///b.txt"));
4615 }
4616 _ => panic!("Expected ListResources response"),
4617 }
4618 }
4619
4620 #[tokio::test]
4621 async fn test_router_merge_prompts() {
4622 use crate::prompt::PromptBuilder;
4623
4624 let router_a =
4625 McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4626
4627 let router_b =
4628 McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4629
4630 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4631
4632 init_router(&mut merged).await;
4633
4634 let req = RouterRequest {
4635 id: RequestId::Number(1),
4636 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4637 extensions: Extensions::new(),
4638 };
4639
4640 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4641
4642 match resp.inner {
4643 Ok(McpResponse::ListPrompts(result)) => {
4644 assert_eq!(result.prompts.len(), 2);
4645 let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4646 assert!(names.contains(&"prompt_a"));
4647 assert!(names.contains(&"prompt_b"));
4648 }
4649 _ => panic!("Expected ListPrompts response"),
4650 }
4651 }
4652
4653 #[tokio::test]
4654 async fn test_router_nest_prefixes_tools() {
4655 let tool_query = ToolBuilder::new("query")
4657 .description("Query the database")
4658 .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4659 .build();
4660 let tool_insert = ToolBuilder::new("insert")
4661 .description("Insert into database")
4662 .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4663 .build();
4664
4665 let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4666
4667 let mut router = McpRouter::new()
4669 .server_info("nested", "1.0")
4670 .nest("db", db_router);
4671
4672 init_router(&mut router).await;
4673
4674 let req = RouterRequest {
4675 id: RequestId::Number(1),
4676 inner: McpRequest::ListTools(ListToolsParams::default()),
4677 extensions: Extensions::new(),
4678 };
4679
4680 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4681
4682 match resp.inner {
4683 Ok(McpResponse::ListTools(result)) => {
4684 assert_eq!(result.tools.len(), 2);
4685 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4686 assert!(names.contains(&"db.query"));
4687 assert!(names.contains(&"db.insert"));
4688 }
4689 _ => panic!("Expected ListTools response"),
4690 }
4691 }
4692
4693 #[tokio::test]
4694 async fn test_router_nest_call_prefixed_tool() {
4695 let tool = ToolBuilder::new("echo")
4696 .description("Echo input")
4697 .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4698 .build();
4699
4700 let nested_router = McpRouter::new().tool(tool);
4701
4702 let mut router = McpRouter::new().nest("api", nested_router);
4703
4704 init_router(&mut router).await;
4705
4706 let req = RouterRequest {
4708 id: RequestId::Number(1),
4709 inner: McpRequest::CallTool(CallToolParams {
4710 name: "api.echo".to_string(),
4711 arguments: serde_json::json!({"value": "hello world"}),
4712 meta: None,
4713 task: None,
4714 }),
4715 extensions: Extensions::new(),
4716 };
4717
4718 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4719
4720 match resp.inner {
4721 Ok(McpResponse::CallTool(result)) => {
4722 assert!(!result.is_error);
4723 match &result.content[0] {
4724 Content::Text { text, .. } => assert_eq!(text, "hello world"),
4725 _ => panic!("Expected text content"),
4726 }
4727 }
4728 _ => panic!("Expected CallTool response"),
4729 }
4730 }
4731
4732 #[tokio::test]
4733 async fn test_router_multiple_nests() {
4734 let db_tool = ToolBuilder::new("query")
4735 .description("Database query")
4736 .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4737 .build();
4738
4739 let api_tool = ToolBuilder::new("fetch")
4740 .description("API fetch")
4741 .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4742 .build();
4743
4744 let db_router = McpRouter::new().tool(db_tool);
4745 let api_router = McpRouter::new().tool(api_tool);
4746
4747 let mut router = McpRouter::new()
4748 .nest("db", db_router)
4749 .nest("api", api_router);
4750
4751 init_router(&mut router).await;
4752
4753 let req = RouterRequest {
4754 id: RequestId::Number(1),
4755 inner: McpRequest::ListTools(ListToolsParams::default()),
4756 extensions: Extensions::new(),
4757 };
4758
4759 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4760
4761 match resp.inner {
4762 Ok(McpResponse::ListTools(result)) => {
4763 assert_eq!(result.tools.len(), 2);
4764 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4765 assert!(names.contains(&"db.query"));
4766 assert!(names.contains(&"api.fetch"));
4767 }
4768 _ => panic!("Expected ListTools response"),
4769 }
4770 }
4771
4772 #[tokio::test]
4773 async fn test_router_merge_and_nest_combined() {
4774 let tool_a = ToolBuilder::new("local")
4776 .description("Local tool")
4777 .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
4778 .build();
4779
4780 let nested_tool = ToolBuilder::new("remote")
4781 .description("Remote tool")
4782 .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
4783 .build();
4784
4785 let nested_router = McpRouter::new().tool(nested_tool);
4786
4787 let mut router = McpRouter::new()
4788 .tool(tool_a)
4789 .nest("external", nested_router);
4790
4791 init_router(&mut router).await;
4792
4793 let req = RouterRequest {
4794 id: RequestId::Number(1),
4795 inner: McpRequest::ListTools(ListToolsParams::default()),
4796 extensions: Extensions::new(),
4797 };
4798
4799 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4800
4801 match resp.inner {
4802 Ok(McpResponse::ListTools(result)) => {
4803 assert_eq!(result.tools.len(), 2);
4804 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4805 assert!(names.contains(&"local"));
4806 assert!(names.contains(&"external.remote"));
4807 }
4808 _ => panic!("Expected ListTools response"),
4809 }
4810 }
4811
4812 #[tokio::test]
4813 async fn test_router_merge_preserves_server_info() {
4814 let child_router = McpRouter::new()
4815 .server_info("child", "2.0")
4816 .instructions("Child instructions");
4817
4818 let mut router = McpRouter::new()
4819 .server_info("parent", "1.0")
4820 .instructions("Parent instructions")
4821 .merge(child_router);
4822
4823 init_router(&mut router).await;
4824
4825 let init_req = RouterRequest {
4827 id: RequestId::Number(99),
4828 inner: McpRequest::Initialize(InitializeParams {
4829 protocol_version: "2025-11-25".to_string(),
4830 capabilities: ClientCapabilities::default(),
4831 client_info: Implementation {
4832 name: "test".to_string(),
4833 version: "1.0".to_string(),
4834 ..Default::default()
4835 },
4836 meta: None,
4837 }),
4838 extensions: Extensions::new(),
4839 };
4840
4841 let child_router2 = McpRouter::new().server_info("child", "2.0");
4843 let mut fresh_router = McpRouter::new()
4844 .server_info("parent", "1.0")
4845 .merge(child_router2);
4846
4847 let resp = fresh_router
4848 .ready()
4849 .await
4850 .unwrap()
4851 .call(init_req)
4852 .await
4853 .unwrap();
4854
4855 match resp.inner {
4856 Ok(McpResponse::Initialize(result)) => {
4857 assert_eq!(result.server_info.name, "parent");
4858 assert_eq!(result.server_info.version, "1.0");
4859 }
4860 _ => panic!("Expected Initialize response"),
4861 }
4862 }
4863
4864 #[tokio::test]
4869 async fn test_auto_instructions_tools_only() {
4870 let tool_a = ToolBuilder::new("alpha")
4871 .description("Alpha tool")
4872 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4873 .build();
4874 let tool_b = ToolBuilder::new("beta")
4875 .description("Beta tool")
4876 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4877 .build();
4878
4879 let mut router = McpRouter::new()
4880 .auto_instructions()
4881 .tool(tool_a)
4882 .tool(tool_b);
4883
4884 let resp = send_initialize(&mut router).await;
4885 let instructions = resp.instructions.expect("should have instructions");
4886
4887 assert!(instructions.contains("## Tools"));
4888 assert!(instructions.contains("- **alpha**: Alpha tool"));
4889 assert!(instructions.contains("- **beta**: Beta tool"));
4890 assert!(!instructions.contains("## Resources"));
4892 assert!(!instructions.contains("## Prompts"));
4893 }
4894
4895 #[tokio::test]
4896 async fn test_auto_instructions_with_annotations() {
4897 let read_only_tool = ToolBuilder::new("query")
4898 .description("Run a query")
4899 .read_only()
4900 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4901 .build();
4902 let destructive_tool = ToolBuilder::new("delete")
4903 .description("Delete a record")
4904 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4905 .build();
4906 let idempotent_tool = ToolBuilder::new("upsert")
4907 .description("Upsert a record")
4908 .non_destructive()
4909 .idempotent()
4910 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4911 .build();
4912
4913 let mut router = McpRouter::new()
4914 .auto_instructions()
4915 .tool(read_only_tool)
4916 .tool(destructive_tool)
4917 .tool(idempotent_tool);
4918
4919 let resp = send_initialize(&mut router).await;
4920 let instructions = resp.instructions.unwrap();
4921
4922 assert!(instructions.contains("- **query**: Run a query [read-only]"));
4923 assert!(instructions.contains("- **delete**: Delete a record\n"));
4925 assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
4926 }
4927
4928 #[tokio::test]
4929 async fn test_auto_instructions_with_resources() {
4930 use crate::resource::ResourceBuilder;
4931
4932 let resource = ResourceBuilder::new("file:///schema.sql")
4933 .name("Schema")
4934 .description("Database schema")
4935 .text("CREATE TABLE ...");
4936
4937 let mut router = McpRouter::new().auto_instructions().resource(resource);
4938
4939 let resp = send_initialize(&mut router).await;
4940 let instructions = resp.instructions.unwrap();
4941
4942 assert!(instructions.contains("## Resources"));
4943 assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
4944 assert!(!instructions.contains("## Tools"));
4945 }
4946
4947 #[tokio::test]
4948 async fn test_auto_instructions_with_resource_templates() {
4949 use crate::resource::ResourceTemplateBuilder;
4950
4951 let template = ResourceTemplateBuilder::new("file:///{path}")
4952 .name("File")
4953 .description("Read a file by path")
4954 .handler(
4955 |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
4956 Ok(crate::ReadResourceResult::text("content", "text/plain"))
4957 },
4958 );
4959
4960 let mut router = McpRouter::new()
4961 .auto_instructions()
4962 .resource_template(template);
4963
4964 let resp = send_initialize(&mut router).await;
4965 let instructions = resp.instructions.unwrap();
4966
4967 assert!(instructions.contains("## Resources"));
4968 assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
4969 }
4970
4971 #[tokio::test]
4972 async fn test_auto_instructions_with_prompts() {
4973 use crate::prompt::PromptBuilder;
4974
4975 let prompt = PromptBuilder::new("write_query")
4976 .description("Help write a SQL query")
4977 .user_message("Write a query for: {task}");
4978
4979 let mut router = McpRouter::new().auto_instructions().prompt(prompt);
4980
4981 let resp = send_initialize(&mut router).await;
4982 let instructions = resp.instructions.unwrap();
4983
4984 assert!(instructions.contains("## Prompts"));
4985 assert!(instructions.contains("- **write_query**: Help write a SQL query"));
4986 assert!(!instructions.contains("## Tools"));
4987 }
4988
4989 #[tokio::test]
4990 async fn test_auto_instructions_all_sections() {
4991 use crate::prompt::PromptBuilder;
4992 use crate::resource::ResourceBuilder;
4993
4994 let tool = ToolBuilder::new("query")
4995 .description("Execute SQL")
4996 .read_only()
4997 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4998 .build();
4999 let resource = ResourceBuilder::new("db://schema")
5000 .name("Schema")
5001 .description("Full database schema")
5002 .text("schema");
5003 let prompt = PromptBuilder::new("write_query")
5004 .description("Help write a SQL query")
5005 .user_message("Write a query");
5006
5007 let mut router = McpRouter::new()
5008 .auto_instructions()
5009 .tool(tool)
5010 .resource(resource)
5011 .prompt(prompt);
5012
5013 let resp = send_initialize(&mut router).await;
5014 let instructions = resp.instructions.unwrap();
5015
5016 assert!(instructions.contains("## Tools"));
5018 assert!(instructions.contains("## Resources"));
5019 assert!(instructions.contains("## Prompts"));
5020
5021 let tools_pos = instructions.find("## Tools").unwrap();
5023 let resources_pos = instructions.find("## Resources").unwrap();
5024 let prompts_pos = instructions.find("## Prompts").unwrap();
5025 assert!(tools_pos < resources_pos);
5026 assert!(resources_pos < prompts_pos);
5027 }
5028
5029 #[tokio::test]
5030 async fn test_auto_instructions_with_prefix_and_suffix() {
5031 let tool = ToolBuilder::new("echo")
5032 .description("Echo input")
5033 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5034 .build();
5035
5036 let mut router = McpRouter::new()
5037 .auto_instructions_with(
5038 Some("This server provides echo capabilities."),
5039 Some("Contact admin@example.com for support."),
5040 )
5041 .tool(tool);
5042
5043 let resp = send_initialize(&mut router).await;
5044 let instructions = resp.instructions.unwrap();
5045
5046 assert!(instructions.starts_with("This server provides echo capabilities."));
5047 assert!(instructions.ends_with("Contact admin@example.com for support."));
5048 assert!(instructions.contains("## Tools"));
5049 assert!(instructions.contains("- **echo**: Echo input"));
5050 }
5051
5052 #[tokio::test]
5053 async fn test_auto_instructions_prefix_only() {
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 .auto_instructions_with(Some("My server intro."), None::<String>)
5061 .tool(tool);
5062
5063 let resp = send_initialize(&mut router).await;
5064 let instructions = resp.instructions.unwrap();
5065
5066 assert!(instructions.starts_with("My server intro."));
5067 assert!(instructions.contains("- **echo**: Echo input"));
5068 }
5069
5070 #[tokio::test]
5071 async fn test_auto_instructions_empty_router() {
5072 let mut router = McpRouter::new().auto_instructions();
5073
5074 let resp = send_initialize(&mut router).await;
5075 let instructions = resp.instructions.expect("should have instructions");
5076
5077 assert!(!instructions.contains("## Tools"));
5079 assert!(!instructions.contains("## Resources"));
5080 assert!(!instructions.contains("## Prompts"));
5081 assert!(instructions.is_empty());
5082 }
5083
5084 #[tokio::test]
5085 async fn test_auto_instructions_overrides_manual() {
5086 let tool = ToolBuilder::new("echo")
5087 .description("Echo input")
5088 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5089 .build();
5090
5091 let mut router = McpRouter::new()
5092 .instructions("This will be overridden")
5093 .auto_instructions()
5094 .tool(tool);
5095
5096 let resp = send_initialize(&mut router).await;
5097 let instructions = resp.instructions.unwrap();
5098
5099 assert!(!instructions.contains("This will be overridden"));
5100 assert!(instructions.contains("- **echo**: Echo input"));
5101 }
5102
5103 #[tokio::test]
5104 async fn test_no_auto_instructions_returns_manual() {
5105 let tool = ToolBuilder::new("echo")
5106 .description("Echo input")
5107 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5108 .build();
5109
5110 let mut router = McpRouter::new()
5111 .instructions("Manual instructions here")
5112 .tool(tool);
5113
5114 let resp = send_initialize(&mut router).await;
5115 let instructions = resp.instructions.unwrap();
5116
5117 assert_eq!(instructions, "Manual instructions here");
5118 }
5119
5120 #[tokio::test]
5121 async fn test_auto_instructions_no_description_fallback() {
5122 let tool = ToolBuilder::new("mystery")
5123 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5124 .build();
5125
5126 let mut router = McpRouter::new().auto_instructions().tool(tool);
5127
5128 let resp = send_initialize(&mut router).await;
5129 let instructions = resp.instructions.unwrap();
5130
5131 assert!(instructions.contains("- **mystery**: No description"));
5132 }
5133
5134 #[tokio::test]
5135 async fn test_auto_instructions_sorted_alphabetically() {
5136 let tool_z = ToolBuilder::new("zebra")
5137 .description("Z tool")
5138 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5139 .build();
5140 let tool_a = ToolBuilder::new("alpha")
5141 .description("A tool")
5142 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5143 .build();
5144 let tool_m = ToolBuilder::new("middle")
5145 .description("M tool")
5146 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5147 .build();
5148
5149 let mut router = McpRouter::new()
5150 .auto_instructions()
5151 .tool(tool_z)
5152 .tool(tool_a)
5153 .tool(tool_m);
5154
5155 let resp = send_initialize(&mut router).await;
5156 let instructions = resp.instructions.unwrap();
5157
5158 let alpha_pos = instructions.find("**alpha**").unwrap();
5159 let middle_pos = instructions.find("**middle**").unwrap();
5160 let zebra_pos = instructions.find("**zebra**").unwrap();
5161 assert!(alpha_pos < middle_pos);
5162 assert!(middle_pos < zebra_pos);
5163 }
5164
5165 #[tokio::test]
5166 async fn test_auto_instructions_read_only_and_idempotent_tags() {
5167 let tool = ToolBuilder::new("safe_update")
5168 .description("Safe update operation")
5169 .idempotent()
5170 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5171 .build();
5172
5173 let mut router = McpRouter::new().auto_instructions().tool(tool);
5174
5175 let resp = send_initialize(&mut router).await;
5176 let instructions = resp.instructions.unwrap();
5177
5178 assert!(
5179 instructions.contains("[idempotent]"),
5180 "got: {}",
5181 instructions
5182 );
5183 }
5184
5185 #[tokio::test]
5186 async fn test_auto_instructions_lazy_generation() {
5187 let mut router = McpRouter::new().auto_instructions();
5190
5191 let tool = ToolBuilder::new("late_tool")
5192 .description("Added after auto_instructions")
5193 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5194 .build();
5195
5196 router = router.tool(tool);
5197
5198 let resp = send_initialize(&mut router).await;
5199 let instructions = resp.instructions.unwrap();
5200
5201 assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
5202 }
5203
5204 #[tokio::test]
5205 async fn test_auto_instructions_multiple_annotation_tags() {
5206 let tool = ToolBuilder::new("update")
5207 .description("Update a record")
5208 .annotations(ToolAnnotations {
5209 read_only_hint: true,
5210 idempotent_hint: true,
5211 ..Default::default()
5212 })
5213 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5214 .build();
5215
5216 let mut router = McpRouter::new().auto_instructions().tool(tool);
5217
5218 let resp = send_initialize(&mut router).await;
5219 let instructions = resp.instructions.unwrap();
5220
5221 assert!(
5222 instructions.contains("[read-only, idempotent]"),
5223 "got: {}",
5224 instructions
5225 );
5226 }
5227
5228 #[tokio::test]
5229 async fn test_auto_instructions_no_annotations_no_tags() {
5230 let tool = ToolBuilder::new("fetch")
5232 .description("Fetch data")
5233 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5234 .build();
5235
5236 let mut router = McpRouter::new().auto_instructions().tool(tool);
5237
5238 let resp = send_initialize(&mut router).await;
5239 let instructions = resp.instructions.unwrap();
5240
5241 assert!(
5243 !instructions.contains('['),
5244 "should have no tags, got: {}",
5245 instructions
5246 );
5247 assert!(instructions.contains("- **fetch**: Fetch data"));
5248 }
5249
5250 async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
5252 let init_req = RouterRequest {
5253 id: RequestId::Number(0),
5254 inner: McpRequest::Initialize(InitializeParams {
5255 protocol_version: "2025-11-25".to_string(),
5256 capabilities: ClientCapabilities {
5257 roots: None,
5258 sampling: None,
5259 elicitation: None,
5260 tasks: None,
5261 experimental: None,
5262 extensions: None,
5263 },
5264 client_info: Implementation {
5265 name: "test".to_string(),
5266 version: "1.0".to_string(),
5267 ..Default::default()
5268 },
5269 meta: None,
5270 }),
5271 extensions: Extensions::new(),
5272 };
5273 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5274 match resp.inner {
5275 Ok(McpResponse::Initialize(result)) => result,
5276 other => panic!("Expected Initialize response, got {:?}", other),
5277 }
5278 }
5279
5280 #[tokio::test]
5281 async fn test_notify_tools_list_changed() {
5282 let (tx, mut rx) = crate::context::notification_channel(16);
5283
5284 let router = McpRouter::new()
5285 .server_info("test", "1.0")
5286 .with_notification_sender(tx);
5287
5288 assert!(router.notify_tools_list_changed());
5289
5290 let notification = rx.recv().await.unwrap();
5291 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5292 }
5293
5294 #[tokio::test]
5295 async fn test_notify_prompts_list_changed() {
5296 let (tx, mut rx) = crate::context::notification_channel(16);
5297
5298 let router = McpRouter::new()
5299 .server_info("test", "1.0")
5300 .with_notification_sender(tx);
5301
5302 assert!(router.notify_prompts_list_changed());
5303
5304 let notification = rx.recv().await.unwrap();
5305 assert!(matches!(
5306 notification,
5307 ServerNotification::PromptsListChanged
5308 ));
5309 }
5310
5311 #[tokio::test]
5312 async fn test_notify_without_sender_returns_false() {
5313 let router = McpRouter::new().server_info("test", "1.0");
5314
5315 assert!(!router.notify_tools_list_changed());
5316 assert!(!router.notify_prompts_list_changed());
5317 assert!(!router.notify_resources_list_changed());
5318 }
5319
5320 #[tokio::test]
5321 async fn test_list_changed_capabilities_with_notification_sender() {
5322 let (tx, _rx) = crate::context::notification_channel(16);
5323 let tool = ToolBuilder::new("test")
5324 .description("test")
5325 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5326 .build();
5327
5328 let mut router = McpRouter::new()
5329 .server_info("test", "1.0")
5330 .tool(tool)
5331 .with_notification_sender(tx);
5332
5333 init_router(&mut router).await;
5334
5335 let caps = router.capabilities();
5336 let tools_cap = caps.tools.expect("tools capability should be present");
5337 assert!(
5338 tools_cap.list_changed,
5339 "tools.listChanged should be true when notification sender is configured"
5340 );
5341 }
5342
5343 #[tokio::test]
5344 async fn test_list_changed_capabilities_without_notification_sender() {
5345 let tool = ToolBuilder::new("test")
5346 .description("test")
5347 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5348 .build();
5349
5350 let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
5351
5352 init_router(&mut router).await;
5353
5354 let caps = router.capabilities();
5355 let tools_cap = caps.tools.expect("tools capability should be present");
5356 assert!(
5357 !tools_cap.list_changed,
5358 "tools.listChanged should be false without notification sender"
5359 );
5360 }
5361
5362 #[tokio::test]
5363 async fn test_set_logging_level_filters_messages() {
5364 let (tx, mut rx) = crate::context::notification_channel(16);
5365
5366 let mut router = McpRouter::new()
5367 .server_info("test", "1.0")
5368 .with_notification_sender(tx);
5369
5370 init_router(&mut router).await;
5371
5372 let set_level_req = RouterRequest {
5374 id: RequestId::Number(99),
5375 inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
5376 level: LogLevel::Warning,
5377 meta: None,
5378 }),
5379 extensions: crate::context::Extensions::new(),
5380 };
5381 let resp = router
5382 .ready()
5383 .await
5384 .unwrap()
5385 .call(set_level_req)
5386 .await
5387 .unwrap();
5388 assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
5389
5390 let ctx = router.create_context(RequestId::Number(100), None);
5392
5393 ctx.send_log(LoggingMessageParams::new(
5395 LogLevel::Error,
5396 serde_json::Value::Null,
5397 ));
5398 assert!(
5399 rx.try_recv().is_ok(),
5400 "Error should pass through Warning filter"
5401 );
5402
5403 ctx.send_log(LoggingMessageParams::new(
5405 LogLevel::Info,
5406 serde_json::Value::Null,
5407 ));
5408 assert!(
5409 rx.try_recv().is_err(),
5410 "Info should be filtered at Warning level"
5411 );
5412 }
5413
5414 #[test]
5415 fn test_paginate_no_page_size() {
5416 let items = vec![1, 2, 3, 4, 5];
5417 let (page, cursor) = paginate(items.clone(), None, None).unwrap();
5418 assert_eq!(page, items);
5419 assert!(cursor.is_none());
5420 }
5421
5422 #[test]
5423 fn test_paginate_first_page() {
5424 let items = vec![1, 2, 3, 4, 5];
5425 let (page, cursor) = paginate(items, None, Some(2)).unwrap();
5426 assert_eq!(page, vec![1, 2]);
5427 assert!(cursor.is_some());
5428 }
5429
5430 #[test]
5431 fn test_paginate_middle_page() {
5432 let items = vec![1, 2, 3, 4, 5];
5433 let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
5434 assert_eq!(page1, vec![1, 2]);
5435
5436 let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
5437 assert_eq!(page2, vec![3, 4]);
5438 assert!(cursor2.is_some());
5439 }
5440
5441 #[test]
5442 fn test_paginate_last_page() {
5443 let items = vec![1, 2, 3, 4, 5];
5444 let cursor = encode_cursor(4);
5446 let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
5447 assert_eq!(page, vec![5]);
5448 assert!(next.is_none());
5449 }
5450
5451 #[test]
5452 fn test_paginate_exact_boundary() {
5453 let items = vec![1, 2, 3, 4];
5454 let (page, cursor) = paginate(items, None, Some(4)).unwrap();
5455 assert_eq!(page, vec![1, 2, 3, 4]);
5456 assert!(cursor.is_none());
5457 }
5458
5459 #[test]
5460 fn test_paginate_invalid_cursor() {
5461 let items = vec![1, 2, 3];
5462 let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
5463 assert!(result.is_err());
5464 }
5465
5466 #[test]
5467 fn test_cursor_round_trip() {
5468 let offset = 42;
5469 let encoded = encode_cursor(offset);
5470 let decoded = decode_cursor(&encoded).unwrap();
5471 assert_eq!(decoded, offset);
5472 }
5473
5474 #[tokio::test]
5475 async fn test_list_tools_pagination() {
5476 let tool_a = ToolBuilder::new("alpha")
5477 .description("a")
5478 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5479 .build();
5480 let tool_b = ToolBuilder::new("beta")
5481 .description("b")
5482 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5483 .build();
5484 let tool_c = ToolBuilder::new("gamma")
5485 .description("c")
5486 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5487 .build();
5488
5489 let mut router = McpRouter::new()
5490 .server_info("test", "1.0")
5491 .page_size(2)
5492 .tool(tool_a)
5493 .tool(tool_b)
5494 .tool(tool_c);
5495
5496 init_router(&mut router).await;
5497
5498 let req = RouterRequest {
5500 id: RequestId::Number(1),
5501 inner: McpRequest::ListTools(ListToolsParams {
5502 cursor: None,
5503 meta: None,
5504 }),
5505 extensions: Extensions::new(),
5506 };
5507 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5508 let (tools, next_cursor) = match resp.inner {
5509 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5510 other => panic!("Expected ListTools, got {:?}", other),
5511 };
5512 assert_eq!(tools.len(), 2);
5513 assert_eq!(tools[0].name, "alpha");
5514 assert_eq!(tools[1].name, "beta");
5515 assert!(next_cursor.is_some());
5516
5517 let req = RouterRequest {
5519 id: RequestId::Number(2),
5520 inner: McpRequest::ListTools(ListToolsParams {
5521 cursor: next_cursor,
5522 meta: None,
5523 }),
5524 extensions: Extensions::new(),
5525 };
5526 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5527 let (tools, next_cursor) = match resp.inner {
5528 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5529 other => panic!("Expected ListTools, got {:?}", other),
5530 };
5531 assert_eq!(tools.len(), 1);
5532 assert_eq!(tools[0].name, "gamma");
5533 assert!(next_cursor.is_none());
5534 }
5535
5536 #[tokio::test]
5537 async fn test_list_tools_no_pagination_by_default() {
5538 let tool_a = ToolBuilder::new("alpha")
5539 .description("a")
5540 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5541 .build();
5542 let tool_b = ToolBuilder::new("beta")
5543 .description("b")
5544 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5545 .build();
5546
5547 let mut router = McpRouter::new()
5548 .server_info("test", "1.0")
5549 .tool(tool_a)
5550 .tool(tool_b);
5551
5552 init_router(&mut router).await;
5553
5554 let req = RouterRequest {
5555 id: RequestId::Number(1),
5556 inner: McpRequest::ListTools(ListToolsParams {
5557 cursor: None,
5558 meta: None,
5559 }),
5560 extensions: Extensions::new(),
5561 };
5562 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5563 match resp.inner {
5564 Ok(McpResponse::ListTools(result)) => {
5565 assert_eq!(result.tools.len(), 2);
5566 assert!(result.next_cursor.is_none());
5567 }
5568 other => panic!("Expected ListTools, got {:?}", other),
5569 }
5570 }
5571
5572 #[cfg(feature = "dynamic-tools")]
5577 mod dynamic_tools_tests {
5578 use super::*;
5579
5580 #[tokio::test]
5581 async fn test_dynamic_tools_register_and_list() {
5582 let (router, registry) = McpRouter::new()
5583 .server_info("test", "1.0")
5584 .with_dynamic_tools();
5585
5586 let tool = ToolBuilder::new("dynamic_echo")
5587 .description("Dynamic echo")
5588 .handler(|input: AddInput| async move {
5589 Ok(CallToolResult::text(format!("{}", input.a)))
5590 })
5591 .build();
5592
5593 registry.register(tool);
5594
5595 let mut router = router;
5596 init_router(&mut router).await;
5597
5598 let req = RouterRequest {
5599 id: RequestId::Number(1),
5600 inner: McpRequest::ListTools(ListToolsParams::default()),
5601 extensions: Extensions::new(),
5602 };
5603
5604 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5605 match resp.inner {
5606 Ok(McpResponse::ListTools(result)) => {
5607 assert_eq!(result.tools.len(), 1);
5608 assert_eq!(result.tools[0].name, "dynamic_echo");
5609 }
5610 _ => panic!("Expected ListTools response"),
5611 }
5612 }
5613
5614 #[tokio::test]
5615 async fn test_dynamic_tools_unregister() {
5616 let (router, registry) = McpRouter::new()
5617 .server_info("test", "1.0")
5618 .with_dynamic_tools();
5619
5620 let tool = ToolBuilder::new("temp")
5621 .description("Temporary")
5622 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5623 .build();
5624
5625 registry.register(tool);
5626 assert!(registry.contains("temp"));
5627
5628 let removed = registry.unregister("temp");
5629 assert!(removed);
5630 assert!(!registry.contains("temp"));
5631
5632 assert!(!registry.unregister("temp"));
5634
5635 let mut router = router;
5636 init_router(&mut router).await;
5637
5638 let req = RouterRequest {
5639 id: RequestId::Number(1),
5640 inner: McpRequest::ListTools(ListToolsParams::default()),
5641 extensions: Extensions::new(),
5642 };
5643
5644 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5645 match resp.inner {
5646 Ok(McpResponse::ListTools(result)) => {
5647 assert_eq!(result.tools.len(), 0);
5648 }
5649 _ => panic!("Expected ListTools response"),
5650 }
5651 }
5652
5653 #[tokio::test]
5654 async fn test_dynamic_tools_merged_with_static() {
5655 let static_tool = ToolBuilder::new("static_tool")
5656 .description("Static")
5657 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5658 .build();
5659
5660 let (router, registry) = McpRouter::new()
5661 .server_info("test", "1.0")
5662 .tool(static_tool)
5663 .with_dynamic_tools();
5664
5665 let dynamic_tool = ToolBuilder::new("dynamic_tool")
5666 .description("Dynamic")
5667 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5668 .build();
5669
5670 registry.register(dynamic_tool);
5671
5672 let mut router = router;
5673 init_router(&mut router).await;
5674
5675 let req = RouterRequest {
5676 id: RequestId::Number(1),
5677 inner: McpRequest::ListTools(ListToolsParams::default()),
5678 extensions: Extensions::new(),
5679 };
5680
5681 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5682 match resp.inner {
5683 Ok(McpResponse::ListTools(result)) => {
5684 assert_eq!(result.tools.len(), 2);
5685 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5686 assert!(names.contains(&"static_tool"));
5687 assert!(names.contains(&"dynamic_tool"));
5688 }
5689 _ => panic!("Expected ListTools response"),
5690 }
5691 }
5692
5693 #[tokio::test]
5694 async fn test_static_tools_shadow_dynamic() {
5695 let static_tool = ToolBuilder::new("shared")
5696 .description("Static version")
5697 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5698 .build();
5699
5700 let (router, registry) = McpRouter::new()
5701 .server_info("test", "1.0")
5702 .tool(static_tool)
5703 .with_dynamic_tools();
5704
5705 let dynamic_tool = ToolBuilder::new("shared")
5706 .description("Dynamic version")
5707 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5708 .build();
5709
5710 registry.register(dynamic_tool);
5711
5712 let mut router = router;
5713 init_router(&mut router).await;
5714
5715 let req = RouterRequest {
5717 id: RequestId::Number(1),
5718 inner: McpRequest::ListTools(ListToolsParams::default()),
5719 extensions: Extensions::new(),
5720 };
5721
5722 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5723 match resp.inner {
5724 Ok(McpResponse::ListTools(result)) => {
5725 assert_eq!(result.tools.len(), 1);
5726 assert_eq!(result.tools[0].name, "shared");
5727 assert_eq!(
5728 result.tools[0].description.as_deref(),
5729 Some("Static version")
5730 );
5731 }
5732 _ => panic!("Expected ListTools response"),
5733 }
5734
5735 let req = RouterRequest {
5737 id: RequestId::Number(2),
5738 inner: McpRequest::CallTool(CallToolParams {
5739 name: "shared".to_string(),
5740 arguments: serde_json::json!({"a": 1, "b": 2}),
5741 meta: None,
5742 task: None,
5743 }),
5744 extensions: Extensions::new(),
5745 };
5746
5747 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5748 match resp.inner {
5749 Ok(McpResponse::CallTool(result)) => {
5750 assert!(!result.is_error);
5751 match &result.content[0] {
5752 Content::Text { text, .. } => assert_eq!(text, "static"),
5753 _ => panic!("Expected text content"),
5754 }
5755 }
5756 _ => panic!("Expected CallTool response"),
5757 }
5758 }
5759
5760 #[tokio::test]
5761 async fn test_dynamic_tools_call() {
5762 let (router, registry) = McpRouter::new()
5763 .server_info("test", "1.0")
5764 .with_dynamic_tools();
5765
5766 let tool = ToolBuilder::new("add")
5767 .description("Add two numbers")
5768 .handler(|input: AddInput| async move {
5769 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
5770 })
5771 .build();
5772
5773 registry.register(tool);
5774
5775 let mut router = router;
5776 init_router(&mut router).await;
5777
5778 let req = RouterRequest {
5779 id: RequestId::Number(1),
5780 inner: McpRequest::CallTool(CallToolParams {
5781 name: "add".to_string(),
5782 arguments: serde_json::json!({"a": 3, "b": 4}),
5783 meta: None,
5784 task: None,
5785 }),
5786 extensions: Extensions::new(),
5787 };
5788
5789 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5790 match resp.inner {
5791 Ok(McpResponse::CallTool(result)) => {
5792 assert!(!result.is_error);
5793 match &result.content[0] {
5794 Content::Text { text, .. } => assert_eq!(text, "7"),
5795 _ => panic!("Expected text content"),
5796 }
5797 }
5798 _ => panic!("Expected CallTool response"),
5799 }
5800 }
5801
5802 #[tokio::test]
5803 async fn test_dynamic_tools_notification_on_register() {
5804 let (tx, mut rx) = crate::context::notification_channel(16);
5805 let (router, registry) = McpRouter::new()
5806 .server_info("test", "1.0")
5807 .with_dynamic_tools();
5808 let _router = router.with_notification_sender(tx);
5809
5810 let tool = ToolBuilder::new("notified")
5811 .description("Test")
5812 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5813 .build();
5814
5815 registry.register(tool);
5816
5817 let notification = rx.recv().await.unwrap();
5818 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5819 }
5820
5821 #[tokio::test]
5822 async fn test_dynamic_tools_notification_on_unregister() {
5823 let (tx, mut rx) = crate::context::notification_channel(16);
5824 let (router, registry) = McpRouter::new()
5825 .server_info("test", "1.0")
5826 .with_dynamic_tools();
5827 let _router = router.with_notification_sender(tx);
5828
5829 let tool = ToolBuilder::new("notified")
5830 .description("Test")
5831 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5832 .build();
5833
5834 registry.register(tool);
5835 let _ = rx.recv().await.unwrap();
5837
5838 registry.unregister("notified");
5839 let notification = rx.recv().await.unwrap();
5840 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5841 }
5842
5843 #[tokio::test]
5844 async fn test_dynamic_tools_no_notification_on_empty_unregister() {
5845 let (tx, mut rx) = crate::context::notification_channel(16);
5846 let (router, registry) = McpRouter::new()
5847 .server_info("test", "1.0")
5848 .with_dynamic_tools();
5849 let _router = router.with_notification_sender(tx);
5850
5851 assert!(!registry.unregister("nonexistent"));
5853
5854 assert!(rx.try_recv().is_err());
5856 }
5857
5858 #[tokio::test]
5859 async fn test_dynamic_tools_filter_applies() {
5860 use crate::filter::CapabilityFilter;
5861
5862 let (router, registry) = McpRouter::new()
5863 .server_info("test", "1.0")
5864 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
5865 tool.name != "hidden"
5866 }))
5867 .with_dynamic_tools();
5868
5869 let visible = ToolBuilder::new("visible")
5870 .description("Visible")
5871 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5872 .build();
5873
5874 let hidden = ToolBuilder::new("hidden")
5875 .description("Hidden")
5876 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5877 .build();
5878
5879 registry.register(visible);
5880 registry.register(hidden);
5881
5882 let mut router = router;
5883 init_router(&mut router).await;
5884
5885 let req = RouterRequest {
5887 id: RequestId::Number(1),
5888 inner: McpRequest::ListTools(ListToolsParams::default()),
5889 extensions: Extensions::new(),
5890 };
5891
5892 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5893 match resp.inner {
5894 Ok(McpResponse::ListTools(result)) => {
5895 assert_eq!(result.tools.len(), 1);
5896 assert_eq!(result.tools[0].name, "visible");
5897 }
5898 _ => panic!("Expected ListTools response"),
5899 }
5900
5901 let req = RouterRequest {
5903 id: RequestId::Number(2),
5904 inner: McpRequest::CallTool(CallToolParams {
5905 name: "hidden".to_string(),
5906 arguments: serde_json::json!({"a": 1, "b": 2}),
5907 meta: None,
5908 task: None,
5909 }),
5910 extensions: Extensions::new(),
5911 };
5912
5913 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5914 match resp.inner {
5915 Err(e) => {
5916 assert_eq!(e.code, -32601); }
5918 _ => panic!("Expected JsonRpc error"),
5919 }
5920 }
5921
5922 #[tokio::test]
5923 async fn test_dynamic_tools_capabilities_advertised() {
5924 let (mut router, _registry) = McpRouter::new()
5926 .server_info("test", "1.0")
5927 .with_dynamic_tools();
5928
5929 let init_req = RouterRequest {
5930 id: RequestId::Number(1),
5931 inner: McpRequest::Initialize(InitializeParams {
5932 protocol_version: "2025-11-25".to_string(),
5933 capabilities: ClientCapabilities::default(),
5934 client_info: Implementation {
5935 name: "test".to_string(),
5936 version: "1.0".to_string(),
5937 ..Default::default()
5938 },
5939 meta: None,
5940 }),
5941 extensions: Extensions::new(),
5942 };
5943
5944 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5945 match resp.inner {
5946 Ok(McpResponse::Initialize(result)) => {
5947 assert!(result.capabilities.tools.is_some());
5948 }
5949 _ => panic!("Expected Initialize response"),
5950 }
5951 }
5952
5953 #[tokio::test]
5954 async fn test_dynamic_tools_multi_session_notification() {
5955 let (tx1, mut rx1) = crate::context::notification_channel(16);
5956 let (tx2, mut rx2) = crate::context::notification_channel(16);
5957
5958 let (router, registry) = McpRouter::new()
5959 .server_info("test", "1.0")
5960 .with_dynamic_tools();
5961
5962 let _session1 = router.clone().with_notification_sender(tx1);
5964 let _session2 = router.clone().with_notification_sender(tx2);
5965
5966 let tool = ToolBuilder::new("broadcast")
5967 .description("Test")
5968 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5969 .build();
5970
5971 registry.register(tool);
5972
5973 let n1 = rx1.recv().await.unwrap();
5975 let n2 = rx2.recv().await.unwrap();
5976 assert!(matches!(n1, ServerNotification::ToolsListChanged));
5977 assert!(matches!(n2, ServerNotification::ToolsListChanged));
5978 }
5979
5980 #[tokio::test]
5981 async fn test_dynamic_tools_call_not_found() {
5982 let (router, _registry) = McpRouter::new()
5983 .server_info("test", "1.0")
5984 .with_dynamic_tools();
5985
5986 let mut router = router;
5987 init_router(&mut router).await;
5988
5989 let req = RouterRequest {
5990 id: RequestId::Number(1),
5991 inner: McpRequest::CallTool(CallToolParams {
5992 name: "nonexistent".to_string(),
5993 arguments: serde_json::json!({}),
5994 meta: None,
5995 task: None,
5996 }),
5997 extensions: Extensions::new(),
5998 };
5999
6000 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6001 match resp.inner {
6002 Err(e) => {
6003 assert_eq!(e.code, -32601);
6004 }
6005 _ => panic!("Expected method not found error"),
6006 }
6007 }
6008
6009 #[tokio::test]
6010 async fn test_dynamic_tools_registry_list() {
6011 let (_, registry) = McpRouter::new()
6012 .server_info("test", "1.0")
6013 .with_dynamic_tools();
6014
6015 assert!(registry.list().is_empty());
6016
6017 let tool = ToolBuilder::new("tool_a")
6018 .description("A")
6019 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6020 .build();
6021 registry.register(tool);
6022
6023 let tool = ToolBuilder::new("tool_b")
6024 .description("B")
6025 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6026 .build();
6027 registry.register(tool);
6028
6029 let tools = registry.list();
6030 assert_eq!(tools.len(), 2);
6031 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
6032 assert!(names.contains(&"tool_a"));
6033 assert!(names.contains(&"tool_b"));
6034 }
6035 } #[tokio::test]
6038 async fn test_tool_if_true_registers() {
6039 let tool = ToolBuilder::new("conditional")
6040 .description("Conditional tool")
6041 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6042 .build();
6043
6044 let mut router = McpRouter::new().tool_if(true, tool);
6045 init_router(&mut router).await;
6046
6047 let req = RouterRequest {
6048 id: RequestId::Number(1),
6049 inner: McpRequest::ListTools(ListToolsParams::default()),
6050 extensions: Extensions::new(),
6051 };
6052 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6053 match resp.inner {
6054 Ok(McpResponse::ListTools(result)) => {
6055 assert_eq!(result.tools.len(), 1);
6056 assert_eq!(result.tools[0].name, "conditional");
6057 }
6058 _ => panic!("Expected ListTools response"),
6059 }
6060 }
6061
6062 #[tokio::test]
6063 async fn test_tool_if_false_skips() {
6064 let tool = ToolBuilder::new("conditional")
6065 .description("Conditional tool")
6066 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6067 .build();
6068
6069 let mut router = McpRouter::new().tool_if(false, tool);
6070 init_router(&mut router).await;
6071
6072 let req = RouterRequest {
6073 id: RequestId::Number(1),
6074 inner: McpRequest::ListTools(ListToolsParams::default()),
6075 extensions: Extensions::new(),
6076 };
6077 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6078 match resp.inner {
6079 Ok(McpResponse::ListTools(result)) => {
6080 assert_eq!(result.tools.len(), 0);
6081 }
6082 _ => panic!("Expected ListTools response"),
6083 }
6084 }
6085
6086 #[tokio::test]
6087 async fn test_tools_if_batch_conditional() {
6088 let tools = vec![
6089 ToolBuilder::new("a")
6090 .description("Tool A")
6091 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6092 .build(),
6093 ToolBuilder::new("b")
6094 .description("Tool B")
6095 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6096 .build(),
6097 ];
6098
6099 let mut router = McpRouter::new().tools_if(false, tools);
6100 init_router(&mut router).await;
6101
6102 let req = RouterRequest {
6103 id: RequestId::Number(1),
6104 inner: McpRequest::ListTools(ListToolsParams::default()),
6105 extensions: Extensions::new(),
6106 };
6107 let resp = router.ready().await.unwrap().call(req).await.unwrap();
6108 match resp.inner {
6109 Ok(McpResponse::ListTools(result)) => {
6110 assert_eq!(result.tools.len(), 0);
6111 }
6112 _ => panic!("Expected ListTools response"),
6113 }
6114 }
6115
6116 #[test]
6117 fn test_resource_if_true_registers() {
6118 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6119 .name("test")
6120 .text("hello");
6121
6122 let router = McpRouter::new().resource_if(true, resource);
6123 assert_eq!(router.inner.resources.len(), 1);
6124 }
6125
6126 #[test]
6127 fn test_resource_if_false_skips() {
6128 let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6129 .name("test")
6130 .text("hello");
6131
6132 let router = McpRouter::new().resource_if(false, resource);
6133 assert_eq!(router.inner.resources.len(), 0);
6134 }
6135
6136 #[test]
6137 fn test_prompt_if_true_registers() {
6138 let prompt = crate::prompt::PromptBuilder::new("greet")
6139 .description("Greeting")
6140 .user_message("Hello!");
6141
6142 let router = McpRouter::new().prompt_if(true, prompt);
6143 assert_eq!(router.inner.prompts.len(), 1);
6144 }
6145
6146 #[test]
6147 fn test_prompt_if_false_skips() {
6148 let prompt = crate::prompt::PromptBuilder::new("greet")
6149 .description("Greeting")
6150 .user_message("Hello!");
6151
6152 let router = McpRouter::new().prompt_if(false, prompt);
6153 assert_eq!(router.inner.prompts.len(), 0);
6154 }
6155
6156 #[test]
6157 fn test_router_request_new() {
6158 let req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6159 assert_eq!(req.id, RequestId::Number(1));
6160 assert!(req.extensions.is_empty());
6161 }
6162
6163 #[test]
6164 fn test_with_inner_preserves_extensions() {
6165 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6166 req.extensions.insert(42u32);
6167
6168 let rewritten = req.with_inner(McpRequest::ListTools(Default::default()));
6169 assert!(matches!(rewritten.inner, McpRequest::ListTools(_)));
6170 assert_eq!(rewritten.id, RequestId::Number(1));
6171 assert_eq!(rewritten.extensions.get::<u32>(), Some(&42));
6172 }
6173
6174 #[test]
6175 fn test_with_id_and_inner_preserves_extensions() {
6176 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6177 req.extensions.insert(String::from("token-abc"));
6178
6179 let rewritten = req.with_id_and_inner(
6180 RequestId::Number(99),
6181 McpRequest::ListResources(Default::default()),
6182 );
6183 assert_eq!(rewritten.id, RequestId::Number(99));
6184 assert!(matches!(rewritten.inner, McpRequest::ListResources(_)));
6185 assert_eq!(
6186 rewritten.extensions.get::<String>(),
6187 Some(&String::from("token-abc"))
6188 );
6189 }
6190
6191 #[test]
6192 fn test_clone_with_inner_preserves_extensions() {
6193 let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6194 req.extensions.insert(true);
6195
6196 let cloned = req.clone_with_inner(McpRequest::ListTools(Default::default()));
6197
6198 assert!(matches!(req.inner, McpRequest::Ping));
6200 assert_eq!(req.extensions.get::<bool>(), Some(&true));
6201
6202 assert!(matches!(cloned.inner, McpRequest::ListTools(_)));
6204 assert_eq!(cloned.extensions.get::<bool>(), Some(&true));
6205 }
6206
6207 #[test]
6208 fn test_router_response_is_error() {
6209 let ok_resp = RouterResponse {
6210 id: RequestId::Number(1),
6211 inner: Ok(McpResponse::Pong(Default::default())),
6212 };
6213 assert!(!ok_resp.is_error());
6214
6215 let err_resp = RouterResponse {
6216 id: RequestId::Number(2),
6217 inner: Err(JsonRpcError::internal_error("boom")),
6218 };
6219 assert!(err_resp.is_error());
6220 }
6221
6222 #[test]
6223 fn test_extensions_len_and_is_empty() {
6224 let mut ext = Extensions::new();
6225 assert!(ext.is_empty());
6226 assert_eq!(ext.len(), 0);
6227
6228 ext.insert(42u32);
6229 assert!(!ext.is_empty());
6230 assert_eq!(ext.len(), 1);
6231
6232 ext.insert(String::from("hello"));
6233 assert_eq!(ext.len(), 2);
6234 }
6235
6236 #[test]
6237 fn test_router_response_serde_roundtrip() {
6238 let response = RouterResponse {
6240 id: RequestId::Number(1),
6241 inner: Ok(McpResponse::Empty(EmptyResult {})),
6242 };
6243 let json = serde_json::to_string(&response).unwrap();
6244 let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6245 assert_eq!(deserialized.id, RequestId::Number(1));
6246 assert!(!deserialized.is_error());
6247
6248 let response = RouterResponse {
6250 id: RequestId::String("req-2".into()),
6251 inner: Err(JsonRpcError::method_not_found("unknown")),
6252 };
6253 let json = serde_json::to_string(&response).unwrap();
6254 let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6255 assert_eq!(deserialized.id, RequestId::String("req-2".into()));
6256 assert!(deserialized.is_error());
6257 }
6258}