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::{DynamicToolRegistry, DynamicToolsInner};
27use crate::resource::{Resource, ResourceTemplate};
28use crate::session::SessionState;
29use crate::tool::Tool;
30
31pub type CompletionHandler = Arc<
33 dyn Fn(CompleteParams) -> Pin<Box<dyn Future<Output = Result<CompleteResult>> + Send>>
34 + Send
35 + Sync,
36>;
37
38fn decode_cursor(cursor: &str) -> Result<usize> {
42 let bytes = BASE64
43 .decode(cursor)
44 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
45 let s = String::from_utf8(bytes)
46 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
47 s.parse::<usize>()
48 .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))
49}
50
51fn encode_cursor(offset: usize) -> String {
53 BASE64.encode(offset.to_string())
54}
55
56fn paginate<T>(
60 items: Vec<T>,
61 cursor: Option<&str>,
62 page_size: Option<usize>,
63) -> Result<(Vec<T>, Option<String>)> {
64 let Some(page_size) = page_size else {
65 return Ok((items, None));
66 };
67
68 let offset = match cursor {
69 Some(c) => decode_cursor(c)?,
70 None => 0,
71 };
72
73 if offset >= items.len() {
74 return Ok((Vec::new(), None));
75 }
76
77 let end = (offset + page_size).min(items.len());
78 let next_cursor = if end < items.len() {
79 Some(encode_cursor(end))
80 } else {
81 None
82 };
83
84 let mut items = items;
85 let page = items.drain(offset..end).collect();
86 Ok((page, next_cursor))
87}
88
89#[derive(Clone)]
113pub struct McpRouter {
114 inner: Arc<McpRouterInner>,
115 session: SessionState,
116}
117
118impl std::fmt::Debug for McpRouter {
119 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120 f.debug_struct("McpRouter")
121 .field("server_name", &self.inner.server_name)
122 .field("server_version", &self.inner.server_version)
123 .field("tools_count", &self.inner.tools.len())
124 .field("resources_count", &self.inner.resources.len())
125 .field("prompts_count", &self.inner.prompts.len())
126 .field("session_phase", &self.session.phase())
127 .finish()
128 }
129}
130
131#[derive(Clone, Debug)]
133struct AutoInstructionsConfig {
134 prefix: Option<String>,
135 suffix: Option<String>,
136}
137
138#[derive(Clone)]
140struct McpRouterInner {
141 server_name: String,
142 server_version: String,
143 server_title: Option<String>,
145 server_description: Option<String>,
147 server_icons: Option<Vec<ToolIcon>>,
149 server_website_url: Option<String>,
151 instructions: Option<String>,
152 auto_instructions: Option<AutoInstructionsConfig>,
153 tools: HashMap<String, Arc<Tool>>,
154 resources: HashMap<String, Arc<Resource>>,
155 resource_templates: Vec<Arc<ResourceTemplate>>,
157 prompts: HashMap<String, Arc<Prompt>>,
158 in_flight: Arc<RwLock<HashMap<RequestId, CancellationToken>>>,
160 notification_tx: Option<NotificationSender>,
162 client_requester: Option<ClientRequesterHandle>,
164 task_store: TaskStore,
166 subscriptions: Arc<RwLock<HashSet<String>>>,
168 completion_handler: Option<CompletionHandler>,
170 tool_filter: Option<ToolFilter>,
172 resource_filter: Option<ResourceFilter>,
174 prompt_filter: Option<PromptFilter>,
176 extensions: Arc<crate::context::Extensions>,
178 min_log_level: Arc<RwLock<LogLevel>>,
180 page_size: Option<usize>,
182 #[cfg(feature = "dynamic-tools")]
184 dynamic_tools: Option<Arc<DynamicToolsInner>>,
185}
186
187impl McpRouterInner {
188 fn generate_instructions(&self, config: &AutoInstructionsConfig) -> String {
190 let mut parts = Vec::new();
191
192 if let Some(prefix) = &config.prefix {
193 parts.push(prefix.clone());
194 }
195
196 if !self.tools.is_empty() {
198 let mut lines = vec!["## Tools".to_string(), String::new()];
199 let mut tools: Vec<_> = self.tools.values().collect();
200 tools.sort_by(|a, b| a.name.cmp(&b.name));
201 for tool in tools {
202 let desc = tool.description.as_deref().unwrap_or("No description");
203 let tags = annotation_tags(tool.annotations.as_ref());
204 if tags.is_empty() {
205 lines.push(format!("- **{}**: {}", tool.name, desc));
206 } else {
207 lines.push(format!("- **{}**: {} [{}]", tool.name, desc, tags));
208 }
209 }
210 parts.push(lines.join("\n"));
211 }
212
213 if !self.resources.is_empty() || !self.resource_templates.is_empty() {
215 let mut lines = vec!["## Resources".to_string(), String::new()];
216 let mut resources: Vec<_> = self.resources.values().collect();
217 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
218 for resource in resources {
219 let desc = resource.description.as_deref().unwrap_or("No description");
220 lines.push(format!("- **{}**: {}", resource.uri, desc));
221 }
222 let mut templates: Vec<_> = self.resource_templates.iter().collect();
223 templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
224 for template in templates {
225 let desc = template.description.as_deref().unwrap_or("No description");
226 lines.push(format!("- **{}**: {}", template.uri_template, desc));
227 }
228 parts.push(lines.join("\n"));
229 }
230
231 if !self.prompts.is_empty() {
233 let mut lines = vec!["## Prompts".to_string(), String::new()];
234 let mut prompts: Vec<_> = self.prompts.values().collect();
235 prompts.sort_by(|a, b| a.name.cmp(&b.name));
236 for prompt in prompts {
237 let desc = prompt.description.as_deref().unwrap_or("No description");
238 lines.push(format!("- **{}**: {}", prompt.name, desc));
239 }
240 parts.push(lines.join("\n"));
241 }
242
243 if let Some(suffix) = &config.suffix {
244 parts.push(suffix.clone());
245 }
246
247 parts.join("\n\n")
248 }
249}
250
251fn annotation_tags(annotations: Option<&crate::protocol::ToolAnnotations>) -> String {
257 let Some(ann) = annotations else {
258 return String::new();
259 };
260 let mut tags = Vec::new();
261 if ann.is_read_only() {
262 tags.push("read-only");
263 }
264 if ann.is_idempotent() {
265 tags.push("idempotent");
266 }
267 tags.join(", ")
268}
269
270impl McpRouter {
271 pub fn new() -> Self {
273 Self {
274 inner: Arc::new(McpRouterInner {
275 server_name: "tower-mcp".to_string(),
276 server_version: env!("CARGO_PKG_VERSION").to_string(),
277 server_title: None,
278 server_description: None,
279 server_icons: None,
280 server_website_url: None,
281 instructions: None,
282 auto_instructions: None,
283 tools: HashMap::new(),
284 resources: HashMap::new(),
285 resource_templates: Vec::new(),
286 prompts: HashMap::new(),
287 in_flight: Arc::new(RwLock::new(HashMap::new())),
288 notification_tx: None,
289 client_requester: None,
290 task_store: TaskStore::new(),
291 subscriptions: Arc::new(RwLock::new(HashSet::new())),
292 extensions: Arc::new(crate::context::Extensions::new()),
293 completion_handler: None,
294 tool_filter: None,
295 resource_filter: None,
296 prompt_filter: None,
297 min_log_level: Arc::new(RwLock::new(LogLevel::Debug)),
298 page_size: None,
299 #[cfg(feature = "dynamic-tools")]
300 dynamic_tools: None,
301 }),
302 session: SessionState::new(),
303 }
304 }
305
306 pub fn with_fresh_session(&self) -> Self {
314 Self {
315 inner: self.inner.clone(),
316 session: SessionState::new(),
317 }
318 }
319
320 pub fn task_store(&self) -> &TaskStore {
322 &self.inner.task_store
323 }
324
325 #[cfg(feature = "dynamic-tools")]
355 pub fn with_dynamic_tools(mut self) -> (Self, DynamicToolRegistry) {
356 let inner_dyn = Arc::new(DynamicToolsInner::new());
357 Arc::make_mut(&mut self.inner).dynamic_tools = Some(inner_dyn.clone());
358 (self, DynamicToolRegistry::new(inner_dyn))
359 }
360
361 pub fn with_notification_sender(mut self, tx: NotificationSender) -> Self {
365 let inner = Arc::make_mut(&mut self.inner);
366 #[cfg(feature = "dynamic-tools")]
369 if let Some(ref dynamic_tools) = inner.dynamic_tools {
370 dynamic_tools.add_notification_sender(tx.clone());
371 }
372 inner.notification_tx = Some(tx);
373 self
374 }
375
376 pub fn notification_sender(&self) -> Option<&NotificationSender> {
378 self.inner.notification_tx.as_ref()
379 }
380
381 pub fn with_client_requester(mut self, requester: ClientRequesterHandle) -> Self {
386 Arc::make_mut(&mut self.inner).client_requester = Some(requester);
387 self
388 }
389
390 pub fn client_requester(&self) -> Option<&ClientRequesterHandle> {
392 self.inner.client_requester.as_ref()
393 }
394
395 pub fn with_state<T: Clone + Send + Sync + 'static>(mut self, state: T) -> Self {
438 let inner = Arc::make_mut(&mut self.inner);
439 Arc::make_mut(&mut inner.extensions).insert(state);
440 self
441 }
442
443 pub fn with_extension<T: Clone + Send + Sync + 'static>(self, value: T) -> Self {
448 self.with_state(value)
449 }
450
451 pub fn extensions(&self) -> &crate::context::Extensions {
453 &self.inner.extensions
454 }
455
456 pub fn create_context(
461 &self,
462 request_id: RequestId,
463 progress_token: Option<ProgressToken>,
464 ) -> RequestContext {
465 let ctx = RequestContext::new(request_id.clone());
466
467 let ctx = if let Some(token) = progress_token {
469 ctx.with_progress_token(token)
470 } else {
471 ctx
472 };
473
474 let ctx = if let Some(tx) = &self.inner.notification_tx {
476 ctx.with_notification_sender(tx.clone())
477 } else {
478 ctx
479 };
480
481 let ctx = if let Some(requester) = &self.inner.client_requester {
483 ctx.with_client_requester(requester.clone())
484 } else {
485 ctx
486 };
487
488 let ctx = ctx.with_extensions(self.inner.extensions.clone());
490
491 let ctx = ctx.with_min_log_level(self.inner.min_log_level.clone());
493
494 let token = ctx.cancellation_token();
496 if let Ok(mut in_flight) = self.inner.in_flight.write() {
497 in_flight.insert(request_id, token);
498 }
499
500 ctx
501 }
502
503 pub fn complete_request(&self, request_id: &RequestId) {
505 if let Ok(mut in_flight) = self.inner.in_flight.write() {
506 in_flight.remove(request_id);
507 }
508 }
509
510 fn cancel_request(&self, request_id: &RequestId) -> bool {
512 let Ok(in_flight) = self.inner.in_flight.read() else {
513 return false;
514 };
515 let Some(token) = in_flight.get(request_id) else {
516 return false;
517 };
518 token.cancel();
519 true
520 }
521
522 pub fn server_info(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
524 let inner = Arc::make_mut(&mut self.inner);
525 inner.server_name = name.into();
526 inner.server_version = version.into();
527 self
528 }
529
530 pub fn page_size(mut self, size: usize) -> Self {
537 Arc::make_mut(&mut self.inner).page_size = Some(size);
538 self
539 }
540
541 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
543 Arc::make_mut(&mut self.inner).instructions = Some(instructions.into());
544 self
545 }
546
547 pub fn auto_instructions(mut self) -> Self {
579 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
580 prefix: None,
581 suffix: None,
582 });
583 self
584 }
585
586 pub fn auto_instructions_with(
603 mut self,
604 prefix: Option<impl Into<String>>,
605 suffix: Option<impl Into<String>>,
606 ) -> Self {
607 Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
608 prefix: prefix.map(Into::into),
609 suffix: suffix.map(Into::into),
610 });
611 self
612 }
613
614 pub fn server_title(mut self, title: impl Into<String>) -> Self {
616 Arc::make_mut(&mut self.inner).server_title = Some(title.into());
617 self
618 }
619
620 pub fn server_description(mut self, description: impl Into<String>) -> Self {
622 Arc::make_mut(&mut self.inner).server_description = Some(description.into());
623 self
624 }
625
626 pub fn server_icons(mut self, icons: Vec<ToolIcon>) -> Self {
628 Arc::make_mut(&mut self.inner).server_icons = Some(icons);
629 self
630 }
631
632 pub fn server_website_url(mut self, url: impl Into<String>) -> Self {
634 Arc::make_mut(&mut self.inner).server_website_url = Some(url.into());
635 self
636 }
637
638 pub fn tool(mut self, tool: Tool) -> Self {
640 Arc::make_mut(&mut self.inner)
641 .tools
642 .insert(tool.name.clone(), Arc::new(tool));
643 self
644 }
645
646 pub fn resource(mut self, resource: Resource) -> Self {
648 Arc::make_mut(&mut self.inner)
649 .resources
650 .insert(resource.uri.clone(), Arc::new(resource));
651 self
652 }
653
654 pub fn resource_template(mut self, template: ResourceTemplate) -> Self {
687 Arc::make_mut(&mut self.inner)
688 .resource_templates
689 .push(Arc::new(template));
690 self
691 }
692
693 pub fn prompt(mut self, prompt: Prompt) -> Self {
695 Arc::make_mut(&mut self.inner)
696 .prompts
697 .insert(prompt.name.clone(), Arc::new(prompt));
698 self
699 }
700
701 pub fn tools(self, tools: impl IntoIterator<Item = Tool>) -> Self {
727 tools
728 .into_iter()
729 .fold(self, |router, tool| router.tool(tool))
730 }
731
732 pub fn resources(self, resources: impl IntoIterator<Item = Resource>) -> Self {
751 resources
752 .into_iter()
753 .fold(self, |router, resource| router.resource(resource))
754 }
755
756 pub fn prompts(self, prompts: impl IntoIterator<Item = Prompt>) -> Self {
775 prompts
776 .into_iter()
777 .fold(self, |router, prompt| router.prompt(prompt))
778 }
779
780 pub fn merge(mut self, other: McpRouter) -> Self {
825 let inner = Arc::make_mut(&mut self.inner);
826 let other_inner = other.inner;
827
828 for (name, tool) in &other_inner.tools {
830 inner.tools.insert(name.clone(), tool.clone());
831 }
832
833 for (uri, resource) in &other_inner.resources {
835 inner.resources.insert(uri.clone(), resource.clone());
836 }
837
838 for template in &other_inner.resource_templates {
841 inner.resource_templates.push(template.clone());
842 }
843
844 for (name, prompt) in &other_inner.prompts {
846 inner.prompts.insert(name.clone(), prompt.clone());
847 }
848
849 self
850 }
851
852 pub fn nest(mut self, prefix: impl Into<String>, other: McpRouter) -> Self {
892 let prefix = prefix.into();
893 let inner = Arc::make_mut(&mut self.inner);
894 let other_inner = other.inner;
895
896 for tool in other_inner.tools.values() {
898 let prefixed_tool = tool.with_name_prefix(&prefix);
899 inner
900 .tools
901 .insert(prefixed_tool.name.clone(), Arc::new(prefixed_tool));
902 }
903
904 for (uri, resource) in &other_inner.resources {
906 inner.resources.insert(uri.clone(), resource.clone());
907 }
908
909 for template in &other_inner.resource_templates {
911 inner.resource_templates.push(template.clone());
912 }
913
914 for (name, prompt) in &other_inner.prompts {
916 inner.prompts.insert(name.clone(), prompt.clone());
917 }
918
919 self
920 }
921
922 pub fn completion_handler<F, Fut>(mut self, handler: F) -> Self
949 where
950 F: Fn(CompleteParams) -> Fut + Send + Sync + 'static,
951 Fut: Future<Output = Result<CompleteResult>> + Send + 'static,
952 {
953 Arc::make_mut(&mut self.inner).completion_handler =
954 Some(Arc::new(move |params| Box::pin(handler(params))));
955 self
956 }
957
958 pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
993 Arc::make_mut(&mut self.inner).tool_filter = Some(filter);
994 self
995 }
996
997 pub fn resource_filter(mut self, filter: ResourceFilter) -> Self {
1028 Arc::make_mut(&mut self.inner).resource_filter = Some(filter);
1029 self
1030 }
1031
1032 pub fn prompt_filter(mut self, filter: PromptFilter) -> Self {
1061 Arc::make_mut(&mut self.inner).prompt_filter = Some(filter);
1062 self
1063 }
1064
1065 pub fn session(&self) -> &SessionState {
1067 &self.session
1068 }
1069
1070 pub fn log(&self, params: LoggingMessageParams) -> bool {
1092 let Some(tx) = &self.inner.notification_tx else {
1093 return false;
1094 };
1095 tx.try_send(ServerNotification::LogMessage(params)).is_ok()
1096 }
1097
1098 pub fn log_info(&self, message: &str) -> bool {
1102 self.log(LoggingMessageParams::new(
1103 LogLevel::Info,
1104 serde_json::json!({ "message": message }),
1105 ))
1106 }
1107
1108 pub fn log_warning(&self, message: &str) -> bool {
1110 self.log(LoggingMessageParams::new(
1111 LogLevel::Warning,
1112 serde_json::json!({ "message": message }),
1113 ))
1114 }
1115
1116 pub fn log_error(&self, message: &str) -> bool {
1118 self.log(LoggingMessageParams::new(
1119 LogLevel::Error,
1120 serde_json::json!({ "message": message }),
1121 ))
1122 }
1123
1124 pub fn log_debug(&self, message: &str) -> bool {
1126 self.log(LoggingMessageParams::new(
1127 LogLevel::Debug,
1128 serde_json::json!({ "message": message }),
1129 ))
1130 }
1131
1132 pub fn is_subscribed(&self, uri: &str) -> bool {
1134 if let Ok(subs) = self.inner.subscriptions.read() {
1135 return subs.contains(uri);
1136 }
1137 false
1138 }
1139
1140 pub fn subscribed_uris(&self) -> Vec<String> {
1142 if let Ok(subs) = self.inner.subscriptions.read() {
1143 return subs.iter().cloned().collect();
1144 }
1145 Vec::new()
1146 }
1147
1148 fn subscribe(&self, uri: &str) -> bool {
1150 if let Ok(mut subs) = self.inner.subscriptions.write() {
1151 return subs.insert(uri.to_string());
1152 }
1153 false
1154 }
1155
1156 fn unsubscribe(&self, uri: &str) -> bool {
1158 if let Ok(mut subs) = self.inner.subscriptions.write() {
1159 return subs.remove(uri);
1160 }
1161 false
1162 }
1163
1164 pub fn notify_resource_updated(&self, uri: &str) -> bool {
1169 if !self.is_subscribed(uri) {
1171 return false;
1172 }
1173
1174 let Some(tx) = &self.inner.notification_tx else {
1175 return false;
1176 };
1177 tx.try_send(ServerNotification::ResourceUpdated {
1178 uri: uri.to_string(),
1179 })
1180 .is_ok()
1181 }
1182
1183 pub fn notify_resources_list_changed(&self) -> bool {
1187 let Some(tx) = &self.inner.notification_tx else {
1188 return false;
1189 };
1190 tx.try_send(ServerNotification::ResourcesListChanged)
1191 .is_ok()
1192 }
1193
1194 pub fn notify_tools_list_changed(&self) -> bool {
1198 let Some(tx) = &self.inner.notification_tx else {
1199 return false;
1200 };
1201 tx.try_send(ServerNotification::ToolsListChanged).is_ok()
1202 }
1203
1204 pub fn notify_prompts_list_changed(&self) -> bool {
1208 let Some(tx) = &self.inner.notification_tx else {
1209 return false;
1210 };
1211 tx.try_send(ServerNotification::PromptsListChanged).is_ok()
1212 }
1213
1214 fn capabilities(&self) -> ServerCapabilities {
1216 let has_resources =
1217 !self.inner.resources.is_empty() || !self.inner.resource_templates.is_empty();
1218 let has_notifications = self.inner.notification_tx.is_some();
1219
1220 #[cfg(feature = "dynamic-tools")]
1221 let has_dynamic_tools = self.inner.dynamic_tools.is_some();
1222 #[cfg(not(feature = "dynamic-tools"))]
1223 let has_dynamic_tools = false;
1224
1225 ServerCapabilities {
1226 tools: if self.inner.tools.is_empty() && !has_dynamic_tools {
1227 None
1228 } else {
1229 Some(ToolsCapability {
1230 list_changed: has_notifications,
1231 })
1232 },
1233 resources: if has_resources {
1234 Some(ResourcesCapability {
1235 subscribe: true,
1236 list_changed: has_notifications,
1237 })
1238 } else {
1239 None
1240 },
1241 prompts: if self.inner.prompts.is_empty() {
1242 None
1243 } else {
1244 Some(PromptsCapability {
1245 list_changed: has_notifications,
1246 })
1247 },
1248 logging: if self.inner.notification_tx.is_some() {
1250 Some(LoggingCapability::default())
1251 } else {
1252 None
1253 },
1254 tasks: {
1256 let has_task_support = self
1257 .inner
1258 .tools
1259 .values()
1260 .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1261 if has_task_support {
1262 Some(TasksCapability {
1263 list: Some(TasksListCapability {}),
1264 cancel: Some(TasksCancelCapability {}),
1265 requests: Some(TasksRequestsCapability {
1266 tools: Some(TasksToolsRequestsCapability {
1267 call: Some(TasksToolsCallCapability {}),
1268 }),
1269 }),
1270 })
1271 } else {
1272 None
1273 }
1274 },
1275 completions: if self.inner.completion_handler.is_some() {
1277 Some(CompletionsCapability::default())
1278 } else {
1279 None
1280 },
1281 experimental: None,
1282 extensions: None,
1283 }
1284 }
1285
1286 async fn handle(&self, request_id: RequestId, request: McpRequest) -> Result<McpResponse> {
1288 let method = request.method_name();
1290 if !self.session.is_request_allowed(method) {
1291 tracing::warn!(
1292 method = %method,
1293 phase = ?self.session.phase(),
1294 "Request rejected: session not initialized"
1295 );
1296 return Err(Error::JsonRpc(JsonRpcError::invalid_request(format!(
1297 "Session not initialized. Only 'initialize' and 'ping' are allowed before initialization. Got: {}",
1298 method
1299 ))));
1300 }
1301
1302 match request {
1303 McpRequest::Initialize(params) => {
1304 tracing::info!(
1305 client = %params.client_info.name,
1306 version = %params.client_info.version,
1307 "Client initializing"
1308 );
1309
1310 let protocol_version = if crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1313 .contains(¶ms.protocol_version.as_str())
1314 {
1315 params.protocol_version
1316 } else {
1317 crate::protocol::LATEST_PROTOCOL_VERSION.to_string()
1318 };
1319
1320 self.session.mark_initializing();
1322
1323 Ok(McpResponse::Initialize(InitializeResult {
1324 protocol_version,
1325 capabilities: self.capabilities(),
1326 server_info: Implementation {
1327 name: self.inner.server_name.clone(),
1328 version: self.inner.server_version.clone(),
1329 title: self.inner.server_title.clone(),
1330 description: self.inner.server_description.clone(),
1331 icons: self.inner.server_icons.clone(),
1332 website_url: self.inner.server_website_url.clone(),
1333 meta: None,
1334 },
1335 instructions: if let Some(config) = &self.inner.auto_instructions {
1336 Some(self.inner.generate_instructions(config))
1337 } else {
1338 self.inner.instructions.clone()
1339 },
1340 meta: None,
1341 }))
1342 }
1343
1344 McpRequest::ListTools(params) => {
1345 let filter = self.inner.tool_filter.as_ref();
1346 let is_visible = |t: &Tool| {
1347 filter
1348 .map(|f| f.is_visible(&self.session, t))
1349 .unwrap_or(true)
1350 };
1351
1352 let mut tools: Vec<ToolDefinition> = self
1354 .inner
1355 .tools
1356 .values()
1357 .filter(|t| is_visible(t))
1358 .map(|t| t.definition())
1359 .collect();
1360
1361 #[cfg(feature = "dynamic-tools")]
1363 if let Some(ref dynamic) = self.inner.dynamic_tools {
1364 let static_names: HashSet<String> =
1365 tools.iter().map(|t| t.name.clone()).collect();
1366 for t in dynamic.list() {
1367 if !static_names.contains(&t.name) && is_visible(&t) {
1368 tools.push(t.definition());
1369 }
1370 }
1371 }
1372
1373 tools.sort_by(|a, b| a.name.cmp(&b.name));
1374
1375 let (tools, next_cursor) =
1376 paginate(tools, params.cursor.as_deref(), self.inner.page_size)?;
1377
1378 Ok(McpResponse::ListTools(ListToolsResult {
1379 tools,
1380 next_cursor,
1381 meta: None,
1382 }))
1383 }
1384
1385 McpRequest::CallTool(params) => {
1386 let tool = self.inner.tools.get(¶ms.name).cloned();
1388 #[cfg(feature = "dynamic-tools")]
1389 let tool = tool.or_else(|| {
1390 self.inner
1391 .dynamic_tools
1392 .as_ref()
1393 .and_then(|d| d.get(¶ms.name))
1394 });
1395
1396 let tool = tool
1397 .ok_or_else(|| Error::JsonRpc(JsonRpcError::method_not_found(¶ms.name)))?;
1398
1399 if let Some(filter) = &self.inner.tool_filter
1401 && !filter.is_visible(&self.session, &tool)
1402 {
1403 return Err(filter.denial_error(¶ms.name));
1404 }
1405
1406 if let Some(task_params) = params.task {
1407 if matches!(tool.task_support, TaskSupportMode::Forbidden) {
1409 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1410 "Tool '{}' does not support async tasks",
1411 params.name
1412 ))));
1413 }
1414
1415 let (task_id, cancellation_token) = self.inner.task_store.create_task(
1417 ¶ms.name,
1418 params.arguments.clone(),
1419 task_params.ttl,
1420 );
1421
1422 tracing::info!(task_id = %task_id, tool = %params.name, "Created async task");
1423
1424 let progress_token = params.meta.and_then(|m| m.progress_token);
1426 let ctx = self.create_context(request_id, progress_token);
1427
1428 let task_store = self.inner.task_store.clone();
1430 let tool = tool.clone();
1431 let arguments = params.arguments;
1432 let task_id_clone = task_id.clone();
1433
1434 tokio::spawn(async move {
1435 if cancellation_token.is_cancelled() {
1437 tracing::debug!(task_id = %task_id_clone, "Task cancelled before execution");
1438 return;
1439 }
1440
1441 let result = tool.call_with_context(ctx, arguments).await;
1443
1444 if cancellation_token.is_cancelled() {
1445 tracing::debug!(task_id = %task_id_clone, "Task cancelled during execution");
1446 } else if result.is_error {
1447 let error_msg = result.first_text().unwrap_or("Tool execution failed");
1449 task_store.fail_task(&task_id_clone, error_msg);
1450 tracing::warn!(task_id = %task_id_clone, error = %error_msg, "Task failed");
1451 } else {
1452 task_store.complete_task(&task_id_clone, result);
1453 tracing::debug!(task_id = %task_id_clone, "Task completed successfully");
1454 }
1455 });
1456
1457 let task = self.inner.task_store.get_task(&task_id).ok_or_else(|| {
1458 Error::JsonRpc(JsonRpcError::internal_error(
1459 "Failed to retrieve created task",
1460 ))
1461 })?;
1462
1463 Ok(McpResponse::CreateTask(CreateTaskResult {
1464 task,
1465 meta: None,
1466 }))
1467 } else {
1468 if matches!(tool.task_support, TaskSupportMode::Required) {
1470 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1471 "Tool '{}' requires async task execution (include 'task' in params)",
1472 params.name
1473 ))));
1474 }
1475
1476 let progress_token = params.meta.and_then(|m| m.progress_token);
1478 let ctx = self.create_context(request_id, progress_token);
1479
1480 tracing::debug!(tool = %params.name, "Calling tool");
1481 let result = tool.call_with_context(ctx, params.arguments).await;
1482
1483 Ok(McpResponse::CallTool(result))
1484 }
1485 }
1486
1487 McpRequest::ListResources(params) => {
1488 let mut resources: Vec<ResourceDefinition> = self
1489 .inner
1490 .resources
1491 .values()
1492 .filter(|r| {
1493 self.inner
1495 .resource_filter
1496 .as_ref()
1497 .map(|f| f.is_visible(&self.session, r))
1498 .unwrap_or(true)
1499 })
1500 .map(|r| r.definition())
1501 .collect();
1502 resources.sort_by(|a, b| a.uri.cmp(&b.uri));
1503
1504 let (resources, next_cursor) =
1505 paginate(resources, params.cursor.as_deref(), self.inner.page_size)?;
1506
1507 Ok(McpResponse::ListResources(ListResourcesResult {
1508 resources,
1509 next_cursor,
1510 meta: None,
1511 }))
1512 }
1513
1514 McpRequest::ListResourceTemplates(params) => {
1515 let mut resource_templates: Vec<ResourceTemplateDefinition> = self
1516 .inner
1517 .resource_templates
1518 .iter()
1519 .map(|t| t.definition())
1520 .collect();
1521 resource_templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
1522
1523 let (resource_templates, next_cursor) = paginate(
1524 resource_templates,
1525 params.cursor.as_deref(),
1526 self.inner.page_size,
1527 )?;
1528
1529 Ok(McpResponse::ListResourceTemplates(
1530 ListResourceTemplatesResult {
1531 resource_templates,
1532 next_cursor,
1533 meta: None,
1534 },
1535 ))
1536 }
1537
1538 McpRequest::ReadResource(params) => {
1539 if let Some(resource) = self.inner.resources.get(¶ms.uri) {
1541 if let Some(filter) = &self.inner.resource_filter
1543 && !filter.is_visible(&self.session, resource)
1544 {
1545 return Err(filter.denial_error(¶ms.uri));
1546 }
1547
1548 tracing::debug!(uri = %params.uri, "Reading static resource");
1549 let result = resource.read().await;
1550 return Ok(McpResponse::ReadResource(result));
1551 }
1552
1553 for template in &self.inner.resource_templates {
1555 if let Some(variables) = template.match_uri(¶ms.uri) {
1556 tracing::debug!(
1557 uri = %params.uri,
1558 template = %template.uri_template,
1559 "Reading resource via template"
1560 );
1561 let result = template.read(¶ms.uri, variables).await?;
1562 return Ok(McpResponse::ReadResource(result));
1563 }
1564 }
1565
1566 Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1568 ¶ms.uri,
1569 )))
1570 }
1571
1572 McpRequest::SubscribeResource(params) => {
1573 if !self.inner.resources.contains_key(¶ms.uri) {
1575 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1576 ¶ms.uri,
1577 )));
1578 }
1579
1580 tracing::debug!(uri = %params.uri, "Subscribing to resource");
1581 self.subscribe(¶ms.uri);
1582
1583 Ok(McpResponse::SubscribeResource(EmptyResult {}))
1584 }
1585
1586 McpRequest::UnsubscribeResource(params) => {
1587 if !self.inner.resources.contains_key(¶ms.uri) {
1589 return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1590 ¶ms.uri,
1591 )));
1592 }
1593
1594 tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
1595 self.unsubscribe(¶ms.uri);
1596
1597 Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
1598 }
1599
1600 McpRequest::ListPrompts(params) => {
1601 let mut prompts: Vec<PromptDefinition> = self
1602 .inner
1603 .prompts
1604 .values()
1605 .filter(|p| {
1606 self.inner
1608 .prompt_filter
1609 .as_ref()
1610 .map(|f| f.is_visible(&self.session, p))
1611 .unwrap_or(true)
1612 })
1613 .map(|p| p.definition())
1614 .collect();
1615 prompts.sort_by(|a, b| a.name.cmp(&b.name));
1616
1617 let (prompts, next_cursor) =
1618 paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
1619
1620 Ok(McpResponse::ListPrompts(ListPromptsResult {
1621 prompts,
1622 next_cursor,
1623 meta: None,
1624 }))
1625 }
1626
1627 McpRequest::GetPrompt(params) => {
1628 let prompt = self.inner.prompts.get(¶ms.name).ok_or_else(|| {
1629 Error::JsonRpc(JsonRpcError::method_not_found(&format!(
1630 "Prompt not found: {}",
1631 params.name
1632 )))
1633 })?;
1634
1635 if let Some(filter) = &self.inner.prompt_filter
1637 && !filter.is_visible(&self.session, prompt)
1638 {
1639 return Err(filter.denial_error(¶ms.name));
1640 }
1641
1642 tracing::debug!(name = %params.name, "Getting prompt");
1643 let result = prompt.get(params.arguments).await?;
1644
1645 Ok(McpResponse::GetPrompt(result))
1646 }
1647
1648 McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
1649
1650 McpRequest::ListTasks(params) => {
1651 let tasks = self.inner.task_store.list_tasks(params.status);
1652
1653 let (tasks, next_cursor) =
1654 paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
1655
1656 Ok(McpResponse::ListTasks(ListTasksResult {
1657 tasks,
1658 next_cursor,
1659 }))
1660 }
1661
1662 McpRequest::GetTaskInfo(params) => {
1663 let task = self
1664 .inner
1665 .task_store
1666 .get_task(¶ms.task_id)
1667 .ok_or_else(|| {
1668 Error::JsonRpc(JsonRpcError::invalid_params(format!(
1669 "Task not found: {}",
1670 params.task_id
1671 )))
1672 })?;
1673
1674 Ok(McpResponse::GetTaskInfo(task))
1675 }
1676
1677 McpRequest::GetTaskResult(params) => {
1678 let (task_obj, result, error) = self
1680 .inner
1681 .task_store
1682 .wait_for_completion(¶ms.task_id)
1683 .await
1684 .ok_or_else(|| {
1685 Error::JsonRpc(JsonRpcError::invalid_params(format!(
1686 "Task not found: {}",
1687 params.task_id
1688 )))
1689 })?;
1690
1691 let meta = serde_json::json!({
1693 "io.modelcontextprotocol/related-task": task_obj
1694 });
1695
1696 match task_obj.status {
1697 TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
1698 format!("Task {} was cancelled", params.task_id),
1699 ))),
1700 TaskStatus::Failed => {
1701 let mut call_result = CallToolResult::error(
1702 error.unwrap_or_else(|| "Task failed".to_string()),
1703 );
1704 call_result.meta = Some(meta);
1705 Ok(McpResponse::GetTaskResult(call_result))
1706 }
1707 _ => {
1708 let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
1709 call_result.meta = Some(meta);
1710 Ok(McpResponse::GetTaskResult(call_result))
1711 }
1712 }
1713 }
1714
1715 McpRequest::CancelTask(params) => {
1716 let current = self
1718 .inner
1719 .task_store
1720 .get_task(¶ms.task_id)
1721 .ok_or_else(|| {
1722 Error::JsonRpc(JsonRpcError::invalid_params(format!(
1723 "Task not found: {}",
1724 params.task_id
1725 )))
1726 })?;
1727
1728 if current.status.is_terminal() {
1729 return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1730 "Task {} is already in terminal state: {}",
1731 params.task_id, current.status
1732 ))));
1733 }
1734
1735 let task_obj = self
1736 .inner
1737 .task_store
1738 .cancel_task(¶ms.task_id, params.reason.as_deref())
1739 .ok_or_else(|| {
1740 Error::JsonRpc(JsonRpcError::invalid_params(format!(
1741 "Task not found: {}",
1742 params.task_id
1743 )))
1744 })?;
1745
1746 Ok(McpResponse::CancelTask(task_obj))
1747 }
1748
1749 McpRequest::SetLoggingLevel(params) => {
1750 tracing::debug!(level = ?params.level, "Client set logging level");
1751 if let Ok(mut level) = self.inner.min_log_level.write() {
1752 *level = params.level;
1753 }
1754 Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
1755 }
1756
1757 McpRequest::Complete(params) => {
1758 tracing::debug!(
1759 reference = ?params.reference,
1760 argument = %params.argument.name,
1761 "Completion request"
1762 );
1763
1764 if let Some(ref handler) = self.inner.completion_handler {
1766 let result = handler(params).await?;
1767 Ok(McpResponse::Complete(result))
1768 } else {
1769 Ok(McpResponse::Complete(CompleteResult::new(vec![])))
1771 }
1772 }
1773
1774 McpRequest::Unknown { method, .. } => {
1775 Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
1776 }
1777 }
1778 }
1779
1780 pub fn handle_notification(&self, notification: McpNotification) {
1782 match notification {
1783 McpNotification::Initialized => {
1784 let phase_before = self.session.phase();
1785 if self.session.mark_initialized() {
1786 if phase_before == crate::session::SessionPhase::Uninitialized {
1787 tracing::info!(
1788 "Session initialized from uninitialized state (race resolved)"
1789 );
1790 } else {
1791 tracing::info!("Session initialized, entering operation phase");
1792 }
1793 } else {
1794 tracing::warn!(
1795 phase = ?self.session.phase(),
1796 "Received initialized notification in unexpected state"
1797 );
1798 }
1799 }
1800 McpNotification::Cancelled(params) => {
1801 if let Some(ref request_id) = params.request_id {
1802 if self.cancel_request(request_id) {
1803 tracing::info!(
1804 request_id = ?request_id,
1805 reason = ?params.reason,
1806 "Request cancelled"
1807 );
1808 } else {
1809 tracing::debug!(
1810 request_id = ?request_id,
1811 reason = ?params.reason,
1812 "Cancellation requested for unknown request"
1813 );
1814 }
1815 } else {
1816 tracing::debug!(
1817 reason = ?params.reason,
1818 "Cancellation notification received without request_id"
1819 );
1820 }
1821 }
1822 McpNotification::Progress(params) => {
1823 tracing::trace!(
1824 token = ?params.progress_token,
1825 progress = params.progress,
1826 total = ?params.total,
1827 "Progress notification"
1828 );
1829 }
1831 McpNotification::RootsListChanged => {
1832 tracing::info!("Client roots list changed");
1833 }
1836 McpNotification::Unknown { method, .. } => {
1837 tracing::debug!(method = %method, "Unknown notification received");
1838 }
1839 }
1840 }
1841}
1842
1843impl Default for McpRouter {
1844 fn default() -> Self {
1845 Self::new()
1846 }
1847}
1848
1849pub use crate::context::Extensions;
1855
1856#[derive(Debug, Clone)]
1858pub struct RouterRequest {
1859 pub id: RequestId,
1860 pub inner: McpRequest,
1861 pub extensions: Extensions,
1863}
1864
1865#[derive(Debug, Clone)]
1867pub struct RouterResponse {
1868 pub id: RequestId,
1869 pub inner: std::result::Result<McpResponse, JsonRpcError>,
1870}
1871
1872impl RouterResponse {
1873 pub fn into_jsonrpc(self) -> JsonRpcResponse {
1875 match self.inner {
1876 Ok(response) => match serde_json::to_value(response) {
1877 Ok(result) => JsonRpcResponse::result(self.id, result),
1878 Err(e) => {
1879 tracing::error!(error = %e, "Failed to serialize response");
1880 JsonRpcResponse::error(
1881 Some(self.id),
1882 JsonRpcError::internal_error(format!("Serialization error: {}", e)),
1883 )
1884 }
1885 },
1886 Err(error) => JsonRpcResponse::error(Some(self.id), error),
1887 }
1888 }
1889}
1890
1891impl Service<RouterRequest> for McpRouter {
1892 type Response = RouterResponse;
1893 type Error = std::convert::Infallible; type Future =
1895 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
1896
1897 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
1898 Poll::Ready(Ok(()))
1899 }
1900
1901 fn call(&mut self, req: RouterRequest) -> Self::Future {
1902 let router = self.clone();
1903 let request_id = req.id.clone();
1904 Box::pin(async move {
1905 let result = router.handle(req.id, req.inner).await;
1906 router.complete_request(&request_id);
1908 Ok(RouterResponse {
1909 id: request_id,
1910 inner: result.map_err(|e| match e {
1915 Error::JsonRpc(err) => err,
1916 Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
1917 e => JsonRpcError::internal_error(e.to_string()),
1918 }),
1919 })
1920 })
1921 }
1922}
1923
1924#[cfg(test)]
1925mod tests {
1926 use super::*;
1927 use crate::extract::{Context, Json};
1928 use crate::jsonrpc::JsonRpcService;
1929 use crate::tool::ToolBuilder;
1930 use schemars::JsonSchema;
1931 use serde::Deserialize;
1932 use tower::ServiceExt;
1933
1934 #[derive(Debug, Deserialize, JsonSchema)]
1935 struct AddInput {
1936 a: i64,
1937 b: i64,
1938 }
1939
1940 async fn init_router(router: &mut McpRouter) {
1942 let init_req = RouterRequest {
1944 id: RequestId::Number(0),
1945 inner: McpRequest::Initialize(InitializeParams {
1946 protocol_version: "2025-11-25".to_string(),
1947 capabilities: ClientCapabilities {
1948 roots: None,
1949 sampling: None,
1950 elicitation: None,
1951 tasks: None,
1952 experimental: None,
1953 extensions: None,
1954 },
1955 client_info: Implementation {
1956 name: "test".to_string(),
1957 version: "1.0".to_string(),
1958 ..Default::default()
1959 },
1960 meta: None,
1961 }),
1962 extensions: Extensions::new(),
1963 };
1964 let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
1965 router.handle_notification(McpNotification::Initialized);
1967 }
1968
1969 #[tokio::test]
1970 async fn test_router_list_tools() {
1971 let add_tool = ToolBuilder::new("add")
1972 .description("Add two numbers")
1973 .handler(|input: AddInput| async move {
1974 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
1975 })
1976 .build();
1977
1978 let mut router = McpRouter::new().tool(add_tool);
1979
1980 init_router(&mut router).await;
1982
1983 let req = RouterRequest {
1984 id: RequestId::Number(1),
1985 inner: McpRequest::ListTools(ListToolsParams::default()),
1986 extensions: Extensions::new(),
1987 };
1988
1989 let resp = router.ready().await.unwrap().call(req).await.unwrap();
1990
1991 match resp.inner {
1992 Ok(McpResponse::ListTools(result)) => {
1993 assert_eq!(result.tools.len(), 1);
1994 assert_eq!(result.tools[0].name, "add");
1995 }
1996 _ => panic!("Expected ListTools response"),
1997 }
1998 }
1999
2000 #[tokio::test]
2001 async fn test_router_call_tool() {
2002 let add_tool = ToolBuilder::new("add")
2003 .description("Add two numbers")
2004 .handler(|input: AddInput| async move {
2005 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2006 })
2007 .build();
2008
2009 let mut router = McpRouter::new().tool(add_tool);
2010
2011 init_router(&mut router).await;
2013
2014 let req = RouterRequest {
2015 id: RequestId::Number(1),
2016 inner: McpRequest::CallTool(CallToolParams {
2017 name: "add".to_string(),
2018 arguments: serde_json::json!({"a": 2, "b": 3}),
2019 meta: None,
2020 task: None,
2021 }),
2022 extensions: Extensions::new(),
2023 };
2024
2025 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2026
2027 match resp.inner {
2028 Ok(McpResponse::CallTool(result)) => {
2029 assert!(!result.is_error);
2030 match &result.content[0] {
2032 Content::Text { text, .. } => assert_eq!(text, "5"),
2033 _ => panic!("Expected text content"),
2034 }
2035 }
2036 _ => panic!("Expected CallTool response"),
2037 }
2038 }
2039
2040 async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2042 let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2043 "protocolVersion": "2025-11-25",
2044 "capabilities": {},
2045 "clientInfo": { "name": "test", "version": "1.0" }
2046 }));
2047 let _ = service.call_single(init_req).await.unwrap();
2048 router.handle_notification(McpNotification::Initialized);
2049 }
2050
2051 #[tokio::test]
2052 async fn test_jsonrpc_service() {
2053 let add_tool = ToolBuilder::new("add")
2054 .description("Add two numbers")
2055 .handler(|input: AddInput| async move {
2056 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2057 })
2058 .build();
2059
2060 let router = McpRouter::new().tool(add_tool);
2061 let mut service = JsonRpcService::new(router.clone());
2062
2063 init_jsonrpc_service(&mut service, &router).await;
2065
2066 let req = JsonRpcRequest::new(1, "tools/list");
2067
2068 let resp = service.call_single(req).await.unwrap();
2069
2070 match resp {
2071 JsonRpcResponse::Result(r) => {
2072 assert_eq!(r.id, RequestId::Number(1));
2073 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2074 assert_eq!(tools.len(), 1);
2075 }
2076 JsonRpcResponse::Error(_) => panic!("Expected success response"),
2077 }
2078 }
2079
2080 #[tokio::test]
2081 async fn test_batch_request() {
2082 let add_tool = ToolBuilder::new("add")
2083 .description("Add two numbers")
2084 .handler(|input: AddInput| async move {
2085 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2086 })
2087 .build();
2088
2089 let router = McpRouter::new().tool(add_tool);
2090 let mut service = JsonRpcService::new(router.clone());
2091
2092 init_jsonrpc_service(&mut service, &router).await;
2094
2095 let requests = vec![
2097 JsonRpcRequest::new(1, "tools/list"),
2098 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2099 "name": "add",
2100 "arguments": {"a": 10, "b": 20}
2101 })),
2102 JsonRpcRequest::new(3, "ping"),
2103 ];
2104
2105 let responses = service.call_batch(requests).await.unwrap();
2106
2107 assert_eq!(responses.len(), 3);
2108
2109 match &responses[0] {
2111 JsonRpcResponse::Result(r) => {
2112 assert_eq!(r.id, RequestId::Number(1));
2113 let tools = r.result.get("tools").unwrap().as_array().unwrap();
2114 assert_eq!(tools.len(), 1);
2115 }
2116 JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2117 }
2118
2119 match &responses[1] {
2121 JsonRpcResponse::Result(r) => {
2122 assert_eq!(r.id, RequestId::Number(2));
2123 let content = r.result.get("content").unwrap().as_array().unwrap();
2124 let text = content[0].get("text").unwrap().as_str().unwrap();
2125 assert_eq!(text, "30");
2126 }
2127 JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2128 }
2129
2130 match &responses[2] {
2132 JsonRpcResponse::Result(r) => {
2133 assert_eq!(r.id, RequestId::Number(3));
2134 }
2135 JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2136 }
2137 }
2138
2139 #[tokio::test]
2140 async fn test_empty_batch_error() {
2141 let router = McpRouter::new();
2142 let mut service = JsonRpcService::new(router);
2143
2144 let result = service.call_batch(vec![]).await;
2145 assert!(result.is_err());
2146 }
2147
2148 #[tokio::test]
2153 async fn test_progress_token_extraction() {
2154 use crate::context::{ServerNotification, notification_channel};
2155 use crate::protocol::ProgressToken;
2156 use std::sync::Arc;
2157 use std::sync::atomic::{AtomicBool, Ordering};
2158
2159 let progress_reported = Arc::new(AtomicBool::new(false));
2161 let progress_ref = progress_reported.clone();
2162
2163 let tool = ToolBuilder::new("progress_tool")
2165 .description("Tool that reports progress")
2166 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2167 let reported = progress_ref.clone();
2168 async move {
2169 ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
2171 .await;
2172 reported.store(true, Ordering::SeqCst);
2173 Ok(CallToolResult::text("done"))
2174 }
2175 })
2176 .build();
2177
2178 let (tx, mut rx) = notification_channel(10);
2180 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2181 let mut service = JsonRpcService::new(router.clone());
2182
2183 init_jsonrpc_service(&mut service, &router).await;
2185
2186 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2188 "name": "progress_tool",
2189 "arguments": {"a": 1, "b": 2},
2190 "_meta": {
2191 "progressToken": "test-token-123"
2192 }
2193 }));
2194
2195 let resp = service.call_single(req).await.unwrap();
2196
2197 match resp {
2199 JsonRpcResponse::Result(_) => {}
2200 JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
2201 }
2202
2203 assert!(progress_reported.load(Ordering::SeqCst));
2205
2206 let notification = rx.try_recv().expect("Expected progress notification");
2208 match notification {
2209 ServerNotification::Progress(params) => {
2210 assert_eq!(
2211 params.progress_token,
2212 ProgressToken::String("test-token-123".to_string())
2213 );
2214 assert_eq!(params.progress, 50.0);
2215 assert_eq!(params.total, Some(100.0));
2216 assert_eq!(params.message.as_deref(), Some("Halfway"));
2217 }
2218 _ => panic!("Expected Progress notification"),
2219 }
2220 }
2221
2222 #[tokio::test]
2223 async fn test_tool_call_without_progress_token() {
2224 use crate::context::notification_channel;
2225 use std::sync::Arc;
2226 use std::sync::atomic::{AtomicBool, Ordering};
2227
2228 let progress_attempted = Arc::new(AtomicBool::new(false));
2229 let progress_ref = progress_attempted.clone();
2230
2231 let tool = ToolBuilder::new("no_token_tool")
2232 .description("Tool that tries to report progress without token")
2233 .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2234 let attempted = progress_ref.clone();
2235 async move {
2236 ctx.report_progress(50.0, Some(100.0), None).await;
2238 attempted.store(true, Ordering::SeqCst);
2239 Ok(CallToolResult::text("done"))
2240 }
2241 })
2242 .build();
2243
2244 let (tx, mut rx) = notification_channel(10);
2245 let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2246 let mut service = JsonRpcService::new(router.clone());
2247
2248 init_jsonrpc_service(&mut service, &router).await;
2249
2250 let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2252 "name": "no_token_tool",
2253 "arguments": {"a": 1, "b": 2}
2254 }));
2255
2256 let resp = service.call_single(req).await.unwrap();
2257 assert!(matches!(resp, JsonRpcResponse::Result(_)));
2258
2259 assert!(progress_attempted.load(Ordering::SeqCst));
2261
2262 assert!(rx.try_recv().is_err());
2264 }
2265
2266 #[tokio::test]
2267 async fn test_batch_errors_returned_not_dropped() {
2268 let add_tool = ToolBuilder::new("add")
2269 .description("Add two numbers")
2270 .handler(|input: AddInput| async move {
2271 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2272 })
2273 .build();
2274
2275 let router = McpRouter::new().tool(add_tool);
2276 let mut service = JsonRpcService::new(router.clone());
2277
2278 init_jsonrpc_service(&mut service, &router).await;
2279
2280 let requests = vec![
2282 JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2284 "name": "add",
2285 "arguments": {"a": 10, "b": 20}
2286 })),
2287 JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2289 "name": "nonexistent_tool",
2290 "arguments": {}
2291 })),
2292 JsonRpcRequest::new(3, "ping"),
2294 ];
2295
2296 let responses = service.call_batch(requests).await.unwrap();
2297
2298 assert_eq!(responses.len(), 3);
2300
2301 match &responses[0] {
2303 JsonRpcResponse::Result(r) => {
2304 assert_eq!(r.id, RequestId::Number(1));
2305 }
2306 JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
2307 }
2308
2309 match &responses[1] {
2311 JsonRpcResponse::Error(e) => {
2312 assert_eq!(e.id, Some(RequestId::Number(2)));
2313 assert!(e.error.message.contains("not found") || e.error.code == -32601);
2315 }
2316 JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
2317 }
2318
2319 match &responses[2] {
2321 JsonRpcResponse::Result(r) => {
2322 assert_eq!(r.id, RequestId::Number(3));
2323 }
2324 JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
2325 }
2326 }
2327
2328 #[tokio::test]
2333 async fn test_list_resource_templates() {
2334 use crate::resource::ResourceTemplateBuilder;
2335 use std::collections::HashMap;
2336
2337 let template = ResourceTemplateBuilder::new("file:///{path}")
2338 .name("Project Files")
2339 .description("Access project files")
2340 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2341 Ok(ReadResourceResult {
2342 contents: vec![ResourceContent {
2343 uri,
2344 mime_type: None,
2345 text: None,
2346 blob: None,
2347 meta: None,
2348 }],
2349 meta: None,
2350 })
2351 });
2352
2353 let mut router = McpRouter::new().resource_template(template);
2354
2355 init_router(&mut router).await;
2357
2358 let req = RouterRequest {
2359 id: RequestId::Number(1),
2360 inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
2361 extensions: Extensions::new(),
2362 };
2363
2364 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2365
2366 match resp.inner {
2367 Ok(McpResponse::ListResourceTemplates(result)) => {
2368 assert_eq!(result.resource_templates.len(), 1);
2369 assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
2370 assert_eq!(result.resource_templates[0].name, "Project Files");
2371 }
2372 _ => panic!("Expected ListResourceTemplates response"),
2373 }
2374 }
2375
2376 #[tokio::test]
2377 async fn test_read_resource_via_template() {
2378 use crate::resource::ResourceTemplateBuilder;
2379 use std::collections::HashMap;
2380
2381 let template = ResourceTemplateBuilder::new("db://users/{id}")
2382 .name("User Records")
2383 .handler(|uri: String, vars: HashMap<String, String>| async move {
2384 let id = vars.get("id").unwrap().clone();
2385 Ok(ReadResourceResult {
2386 contents: vec![ResourceContent {
2387 uri,
2388 mime_type: Some("application/json".to_string()),
2389 text: Some(format!(r#"{{"id": "{}"}}"#, id)),
2390 blob: None,
2391 meta: None,
2392 }],
2393 meta: None,
2394 })
2395 });
2396
2397 let mut router = McpRouter::new().resource_template(template);
2398
2399 init_router(&mut router).await;
2401
2402 let req = RouterRequest {
2404 id: RequestId::Number(1),
2405 inner: McpRequest::ReadResource(ReadResourceParams {
2406 uri: "db://users/123".to_string(),
2407 meta: None,
2408 }),
2409 extensions: Extensions::new(),
2410 };
2411
2412 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2413
2414 match resp.inner {
2415 Ok(McpResponse::ReadResource(result)) => {
2416 assert_eq!(result.contents.len(), 1);
2417 assert_eq!(result.contents[0].uri, "db://users/123");
2418 assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
2419 }
2420 _ => panic!("Expected ReadResource response"),
2421 }
2422 }
2423
2424 #[tokio::test]
2425 async fn test_static_resource_takes_precedence_over_template() {
2426 use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
2427 use std::collections::HashMap;
2428
2429 let template = ResourceTemplateBuilder::new("file:///{path}")
2431 .name("Files Template")
2432 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2433 Ok(ReadResourceResult {
2434 contents: vec![ResourceContent {
2435 uri,
2436 mime_type: None,
2437 text: Some("from template".to_string()),
2438 blob: None,
2439 meta: None,
2440 }],
2441 meta: None,
2442 })
2443 });
2444
2445 let static_resource = ResourceBuilder::new("file:///README.md")
2447 .name("README")
2448 .text("from static resource");
2449
2450 let mut router = McpRouter::new()
2451 .resource_template(template)
2452 .resource(static_resource);
2453
2454 init_router(&mut router).await;
2456
2457 let req = RouterRequest {
2459 id: RequestId::Number(1),
2460 inner: McpRequest::ReadResource(ReadResourceParams {
2461 uri: "file:///README.md".to_string(),
2462 meta: None,
2463 }),
2464 extensions: Extensions::new(),
2465 };
2466
2467 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2468
2469 match resp.inner {
2470 Ok(McpResponse::ReadResource(result)) => {
2471 assert_eq!(
2473 result.contents[0].text.as_deref(),
2474 Some("from static resource")
2475 );
2476 }
2477 _ => panic!("Expected ReadResource response"),
2478 }
2479 }
2480
2481 #[tokio::test]
2482 async fn test_resource_not_found_when_no_match() {
2483 use crate::resource::ResourceTemplateBuilder;
2484 use std::collections::HashMap;
2485
2486 let template = ResourceTemplateBuilder::new("db://users/{id}")
2487 .name("Users")
2488 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2489 Ok(ReadResourceResult {
2490 contents: vec![ResourceContent {
2491 uri,
2492 mime_type: None,
2493 text: None,
2494 blob: None,
2495 meta: None,
2496 }],
2497 meta: None,
2498 })
2499 });
2500
2501 let mut router = McpRouter::new().resource_template(template);
2502
2503 init_router(&mut router).await;
2505
2506 let req = RouterRequest {
2508 id: RequestId::Number(1),
2509 inner: McpRequest::ReadResource(ReadResourceParams {
2510 uri: "db://posts/123".to_string(),
2511 meta: None,
2512 }),
2513 extensions: Extensions::new(),
2514 };
2515
2516 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2517
2518 match resp.inner {
2519 Err(err) => {
2520 assert!(err.message.contains("not found"));
2521 }
2522 Ok(_) => panic!("Expected error for non-matching URI"),
2523 }
2524 }
2525
2526 #[tokio::test]
2527 async fn test_capabilities_include_resources_with_only_templates() {
2528 use crate::resource::ResourceTemplateBuilder;
2529 use std::collections::HashMap;
2530
2531 let template = ResourceTemplateBuilder::new("file:///{path}")
2532 .name("Files")
2533 .handler(|uri: String, _vars: HashMap<String, String>| async move {
2534 Ok(ReadResourceResult {
2535 contents: vec![ResourceContent {
2536 uri,
2537 mime_type: None,
2538 text: None,
2539 blob: None,
2540 meta: None,
2541 }],
2542 meta: None,
2543 })
2544 });
2545
2546 let mut router = McpRouter::new().resource_template(template);
2547
2548 let init_req = RouterRequest {
2550 id: RequestId::Number(0),
2551 inner: McpRequest::Initialize(InitializeParams {
2552 protocol_version: "2025-11-25".to_string(),
2553 capabilities: ClientCapabilities {
2554 roots: None,
2555 sampling: None,
2556 elicitation: None,
2557 tasks: None,
2558 experimental: None,
2559 extensions: None,
2560 },
2561 client_info: Implementation {
2562 name: "test".to_string(),
2563 version: "1.0".to_string(),
2564 ..Default::default()
2565 },
2566 meta: None,
2567 }),
2568 extensions: Extensions::new(),
2569 };
2570 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
2571
2572 match resp.inner {
2573 Ok(McpResponse::Initialize(result)) => {
2574 assert!(result.capabilities.resources.is_some());
2576 }
2577 _ => panic!("Expected Initialize response"),
2578 }
2579 }
2580
2581 #[tokio::test]
2586 async fn test_log_sends_notification() {
2587 use crate::context::notification_channel;
2588
2589 let (tx, mut rx) = notification_channel(10);
2590 let router = McpRouter::new().with_notification_sender(tx);
2591
2592 let sent = router.log_info("Test message");
2594 assert!(sent);
2595
2596 let notification = rx.try_recv().unwrap();
2598 match notification {
2599 ServerNotification::LogMessage(params) => {
2600 assert_eq!(params.level, LogLevel::Info);
2601 let data = params.data;
2602 assert_eq!(
2603 data.get("message").unwrap().as_str().unwrap(),
2604 "Test message"
2605 );
2606 }
2607 _ => panic!("Expected LogMessage notification"),
2608 }
2609 }
2610
2611 #[tokio::test]
2612 async fn test_log_with_custom_params() {
2613 use crate::context::notification_channel;
2614
2615 let (tx, mut rx) = notification_channel(10);
2616 let router = McpRouter::new().with_notification_sender(tx);
2617
2618 let params = LoggingMessageParams::new(
2620 LogLevel::Error,
2621 serde_json::json!({
2622 "error": "Connection failed",
2623 "host": "localhost"
2624 }),
2625 )
2626 .with_logger("database");
2627
2628 let sent = router.log(params);
2629 assert!(sent);
2630
2631 let notification = rx.try_recv().unwrap();
2632 match notification {
2633 ServerNotification::LogMessage(params) => {
2634 assert_eq!(params.level, LogLevel::Error);
2635 assert_eq!(params.logger.as_deref(), Some("database"));
2636 let data = params.data;
2637 assert_eq!(
2638 data.get("error").unwrap().as_str().unwrap(),
2639 "Connection failed"
2640 );
2641 }
2642 _ => panic!("Expected LogMessage notification"),
2643 }
2644 }
2645
2646 #[tokio::test]
2647 async fn test_log_without_channel_returns_false() {
2648 let router = McpRouter::new();
2650
2651 assert!(!router.log_info("Test"));
2653 assert!(!router.log_warning("Test"));
2654 assert!(!router.log_error("Test"));
2655 assert!(!router.log_debug("Test"));
2656 }
2657
2658 #[tokio::test]
2659 async fn test_logging_capability_with_channel() {
2660 use crate::context::notification_channel;
2661
2662 let (tx, _rx) = notification_channel(10);
2663 let mut router = McpRouter::new().with_notification_sender(tx);
2664
2665 let init_req = RouterRequest {
2667 id: RequestId::Number(0),
2668 inner: McpRequest::Initialize(InitializeParams {
2669 protocol_version: "2025-11-25".to_string(),
2670 capabilities: ClientCapabilities {
2671 roots: None,
2672 sampling: None,
2673 elicitation: None,
2674 tasks: None,
2675 experimental: None,
2676 extensions: None,
2677 },
2678 client_info: Implementation {
2679 name: "test".to_string(),
2680 version: "1.0".to_string(),
2681 ..Default::default()
2682 },
2683 meta: None,
2684 }),
2685 extensions: Extensions::new(),
2686 };
2687 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
2688
2689 match resp.inner {
2690 Ok(McpResponse::Initialize(result)) => {
2691 assert!(result.capabilities.logging.is_some());
2693 }
2694 _ => panic!("Expected Initialize response"),
2695 }
2696 }
2697
2698 #[tokio::test]
2699 async fn test_no_logging_capability_without_channel() {
2700 let mut router = McpRouter::new();
2701
2702 let init_req = RouterRequest {
2704 id: RequestId::Number(0),
2705 inner: McpRequest::Initialize(InitializeParams {
2706 protocol_version: "2025-11-25".to_string(),
2707 capabilities: ClientCapabilities {
2708 roots: None,
2709 sampling: None,
2710 elicitation: None,
2711 tasks: None,
2712 experimental: None,
2713 extensions: None,
2714 },
2715 client_info: Implementation {
2716 name: "test".to_string(),
2717 version: "1.0".to_string(),
2718 ..Default::default()
2719 },
2720 meta: None,
2721 }),
2722 extensions: Extensions::new(),
2723 };
2724 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
2725
2726 match resp.inner {
2727 Ok(McpResponse::Initialize(result)) => {
2728 assert!(result.capabilities.logging.is_none());
2730 }
2731 _ => panic!("Expected Initialize response"),
2732 }
2733 }
2734
2735 #[tokio::test]
2740 async fn test_create_task_via_call_tool() {
2741 let add_tool = ToolBuilder::new("add")
2742 .description("Add two numbers")
2743 .task_support(TaskSupportMode::Optional)
2744 .handler(|input: AddInput| async move {
2745 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2746 })
2747 .build();
2748
2749 let mut router = McpRouter::new().tool(add_tool);
2750 init_router(&mut router).await;
2751
2752 let req = RouterRequest {
2753 id: RequestId::Number(1),
2754 inner: McpRequest::CallTool(CallToolParams {
2755 name: "add".to_string(),
2756 arguments: serde_json::json!({"a": 5, "b": 10}),
2757 meta: None,
2758 task: Some(TaskRequestParams { ttl: None }),
2759 }),
2760 extensions: Extensions::new(),
2761 };
2762
2763 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2764
2765 match resp.inner {
2766 Ok(McpResponse::CreateTask(result)) => {
2767 assert!(result.task.task_id.starts_with("task-"));
2768 assert_eq!(result.task.status, TaskStatus::Working);
2769 }
2770 _ => panic!("Expected CreateTask response"),
2771 }
2772 }
2773
2774 #[tokio::test]
2775 async fn test_list_tasks_empty() {
2776 let mut router = McpRouter::new();
2777 init_router(&mut router).await;
2778
2779 let req = RouterRequest {
2780 id: RequestId::Number(1),
2781 inner: McpRequest::ListTasks(ListTasksParams::default()),
2782 extensions: Extensions::new(),
2783 };
2784
2785 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2786
2787 match resp.inner {
2788 Ok(McpResponse::ListTasks(result)) => {
2789 assert!(result.tasks.is_empty());
2790 }
2791 _ => panic!("Expected ListTasks response"),
2792 }
2793 }
2794
2795 #[tokio::test]
2796 async fn test_task_lifecycle_complete() {
2797 let add_tool = ToolBuilder::new("add")
2798 .description("Add two numbers")
2799 .task_support(TaskSupportMode::Optional)
2800 .handler(|input: AddInput| async move {
2801 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2802 })
2803 .build();
2804
2805 let mut router = McpRouter::new().tool(add_tool);
2806 init_router(&mut router).await;
2807
2808 let req = RouterRequest {
2810 id: RequestId::Number(1),
2811 inner: McpRequest::CallTool(CallToolParams {
2812 name: "add".to_string(),
2813 arguments: serde_json::json!({"a": 7, "b": 8}),
2814 meta: None,
2815 task: Some(TaskRequestParams { ttl: None }),
2816 }),
2817 extensions: Extensions::new(),
2818 };
2819
2820 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2821 let task_id = match resp.inner {
2822 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
2823 _ => panic!("Expected CreateTask response"),
2824 };
2825
2826 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
2828
2829 let req = RouterRequest {
2831 id: RequestId::Number(2),
2832 inner: McpRequest::GetTaskResult(GetTaskResultParams {
2833 task_id: task_id.clone(),
2834 meta: None,
2835 }),
2836 extensions: Extensions::new(),
2837 };
2838
2839 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2840
2841 match resp.inner {
2842 Ok(McpResponse::GetTaskResult(result)) => {
2843 assert!(result.meta.is_some());
2845 match &result.content[0] {
2847 Content::Text { text, .. } => assert_eq!(text, "15"),
2848 _ => panic!("Expected text content"),
2849 }
2850 }
2851 _ => panic!("Expected GetTaskResult response"),
2852 }
2853 }
2854
2855 #[tokio::test]
2856 async fn test_task_cancellation() {
2857 let slow_tool = ToolBuilder::new("slow")
2859 .description("Slow tool")
2860 .task_support(TaskSupportMode::Optional)
2861 .handler(|_input: serde_json::Value| async move {
2862 tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
2863 Ok(CallToolResult::text("done"))
2864 })
2865 .build();
2866
2867 let mut router = McpRouter::new().tool(slow_tool);
2868 init_router(&mut router).await;
2869
2870 let req = RouterRequest {
2872 id: RequestId::Number(1),
2873 inner: McpRequest::CallTool(CallToolParams {
2874 name: "slow".to_string(),
2875 arguments: serde_json::json!({}),
2876 meta: None,
2877 task: Some(TaskRequestParams { ttl: None }),
2878 }),
2879 extensions: Extensions::new(),
2880 };
2881
2882 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2883 let task_id = match resp.inner {
2884 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
2885 _ => panic!("Expected CreateTask response"),
2886 };
2887
2888 let req = RouterRequest {
2890 id: RequestId::Number(2),
2891 inner: McpRequest::CancelTask(CancelTaskParams {
2892 task_id: task_id.clone(),
2893 reason: Some("Test cancellation".to_string()),
2894 meta: None,
2895 }),
2896 extensions: Extensions::new(),
2897 };
2898
2899 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2900
2901 match resp.inner {
2902 Ok(McpResponse::CancelTask(task_obj)) => {
2903 assert_eq!(task_obj.status, TaskStatus::Cancelled);
2904 }
2905 _ => panic!("Expected CancelTask response"),
2906 }
2907 }
2908
2909 #[tokio::test]
2910 async fn test_get_task_info() {
2911 let add_tool = ToolBuilder::new("add")
2912 .description("Add two numbers")
2913 .task_support(TaskSupportMode::Optional)
2914 .handler(|input: AddInput| async move {
2915 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2916 })
2917 .build();
2918
2919 let mut router = McpRouter::new().tool(add_tool);
2920 init_router(&mut router).await;
2921
2922 let req = RouterRequest {
2924 id: RequestId::Number(1),
2925 inner: McpRequest::CallTool(CallToolParams {
2926 name: "add".to_string(),
2927 arguments: serde_json::json!({"a": 1, "b": 2}),
2928 meta: None,
2929 task: Some(TaskRequestParams { ttl: Some(600_000) }),
2930 }),
2931 extensions: Extensions::new(),
2932 };
2933
2934 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2935 let task_id = match resp.inner {
2936 Ok(McpResponse::CreateTask(result)) => result.task.task_id,
2937 _ => panic!("Expected CreateTask response"),
2938 };
2939
2940 let req = RouterRequest {
2942 id: RequestId::Number(2),
2943 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
2944 task_id: task_id.clone(),
2945 meta: None,
2946 }),
2947 extensions: Extensions::new(),
2948 };
2949
2950 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2951
2952 match resp.inner {
2953 Ok(McpResponse::GetTaskInfo(info)) => {
2954 assert_eq!(info.task_id, task_id);
2955 assert!(info.created_at.contains('T')); assert_eq!(info.ttl, Some(600_000));
2957 }
2958 _ => panic!("Expected GetTaskInfo response"),
2959 }
2960 }
2961
2962 #[tokio::test]
2963 async fn test_task_forbidden_tool_rejects_task_params() {
2964 let tool = ToolBuilder::new("sync_only")
2965 .description("Sync only tool")
2966 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
2967 .build();
2968
2969 let mut router = McpRouter::new().tool(tool);
2970 init_router(&mut router).await;
2971
2972 let req = RouterRequest {
2974 id: RequestId::Number(1),
2975 inner: McpRequest::CallTool(CallToolParams {
2976 name: "sync_only".to_string(),
2977 arguments: serde_json::json!({}),
2978 meta: None,
2979 task: Some(TaskRequestParams { ttl: None }),
2980 }),
2981 extensions: Extensions::new(),
2982 };
2983
2984 let resp = router.ready().await.unwrap().call(req).await.unwrap();
2985
2986 match resp.inner {
2987 Err(e) => {
2988 assert!(e.message.contains("does not support async tasks"));
2989 }
2990 _ => panic!("Expected error response"),
2991 }
2992 }
2993
2994 #[tokio::test]
2995 async fn test_get_nonexistent_task() {
2996 let mut router = McpRouter::new();
2997 init_router(&mut router).await;
2998
2999 let req = RouterRequest {
3000 id: RequestId::Number(1),
3001 inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3002 task_id: "task-999".to_string(),
3003 meta: None,
3004 }),
3005 extensions: Extensions::new(),
3006 };
3007
3008 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3009
3010 match resp.inner {
3011 Err(e) => {
3012 assert!(e.message.contains("not found"));
3013 }
3014 _ => panic!("Expected error response"),
3015 }
3016 }
3017
3018 #[tokio::test]
3023 async fn test_subscribe_to_resource() {
3024 use crate::resource::ResourceBuilder;
3025
3026 let resource = ResourceBuilder::new("file:///test.txt")
3027 .name("Test File")
3028 .text("Hello");
3029
3030 let mut router = McpRouter::new().resource(resource);
3031 init_router(&mut router).await;
3032
3033 let req = RouterRequest {
3035 id: RequestId::Number(1),
3036 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3037 uri: "file:///test.txt".to_string(),
3038 meta: None,
3039 }),
3040 extensions: Extensions::new(),
3041 };
3042
3043 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3044
3045 match resp.inner {
3046 Ok(McpResponse::SubscribeResource(_)) => {
3047 assert!(router.is_subscribed("file:///test.txt"));
3049 }
3050 _ => panic!("Expected SubscribeResource response"),
3051 }
3052 }
3053
3054 #[tokio::test]
3055 async fn test_unsubscribe_from_resource() {
3056 use crate::resource::ResourceBuilder;
3057
3058 let resource = ResourceBuilder::new("file:///test.txt")
3059 .name("Test File")
3060 .text("Hello");
3061
3062 let mut router = McpRouter::new().resource(resource);
3063 init_router(&mut router).await;
3064
3065 let req = RouterRequest {
3067 id: RequestId::Number(1),
3068 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3069 uri: "file:///test.txt".to_string(),
3070 meta: None,
3071 }),
3072 extensions: Extensions::new(),
3073 };
3074 let _ = router.ready().await.unwrap().call(req).await.unwrap();
3075 assert!(router.is_subscribed("file:///test.txt"));
3076
3077 let req = RouterRequest {
3079 id: RequestId::Number(2),
3080 inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3081 uri: "file:///test.txt".to_string(),
3082 meta: None,
3083 }),
3084 extensions: Extensions::new(),
3085 };
3086
3087 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3088
3089 match resp.inner {
3090 Ok(McpResponse::UnsubscribeResource(_)) => {
3091 assert!(!router.is_subscribed("file:///test.txt"));
3093 }
3094 _ => panic!("Expected UnsubscribeResource response"),
3095 }
3096 }
3097
3098 #[tokio::test]
3099 async fn test_subscribe_nonexistent_resource() {
3100 let mut router = McpRouter::new();
3101 init_router(&mut router).await;
3102
3103 let req = RouterRequest {
3104 id: RequestId::Number(1),
3105 inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3106 uri: "file:///nonexistent.txt".to_string(),
3107 meta: None,
3108 }),
3109 extensions: Extensions::new(),
3110 };
3111
3112 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3113
3114 match resp.inner {
3115 Err(e) => {
3116 assert!(e.message.contains("not found"));
3117 }
3118 _ => panic!("Expected error response"),
3119 }
3120 }
3121
3122 #[tokio::test]
3123 async fn test_notify_resource_updated() {
3124 use crate::context::notification_channel;
3125 use crate::resource::ResourceBuilder;
3126
3127 let (tx, mut rx) = notification_channel(10);
3128
3129 let resource = ResourceBuilder::new("file:///test.txt")
3130 .name("Test File")
3131 .text("Hello");
3132
3133 let router = McpRouter::new()
3134 .resource(resource)
3135 .with_notification_sender(tx);
3136
3137 router.subscribe("file:///test.txt");
3139
3140 let sent = router.notify_resource_updated("file:///test.txt");
3142 assert!(sent);
3143
3144 let notification = rx.try_recv().unwrap();
3146 match notification {
3147 ServerNotification::ResourceUpdated { uri } => {
3148 assert_eq!(uri, "file:///test.txt");
3149 }
3150 _ => panic!("Expected ResourceUpdated notification"),
3151 }
3152 }
3153
3154 #[tokio::test]
3155 async fn test_notify_resource_updated_not_subscribed() {
3156 use crate::context::notification_channel;
3157 use crate::resource::ResourceBuilder;
3158
3159 let (tx, mut rx) = notification_channel(10);
3160
3161 let resource = ResourceBuilder::new("file:///test.txt")
3162 .name("Test File")
3163 .text("Hello");
3164
3165 let router = McpRouter::new()
3166 .resource(resource)
3167 .with_notification_sender(tx);
3168
3169 let sent = router.notify_resource_updated("file:///test.txt");
3171 assert!(!sent); assert!(rx.try_recv().is_err());
3175 }
3176
3177 #[tokio::test]
3178 async fn test_notify_resources_list_changed() {
3179 use crate::context::notification_channel;
3180
3181 let (tx, mut rx) = notification_channel(10);
3182 let router = McpRouter::new().with_notification_sender(tx);
3183
3184 let sent = router.notify_resources_list_changed();
3185 assert!(sent);
3186
3187 let notification = rx.try_recv().unwrap();
3188 match notification {
3189 ServerNotification::ResourcesListChanged => {}
3190 _ => panic!("Expected ResourcesListChanged notification"),
3191 }
3192 }
3193
3194 #[tokio::test]
3195 async fn test_subscribed_uris() {
3196 use crate::resource::ResourceBuilder;
3197
3198 let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
3199
3200 let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
3201
3202 let router = McpRouter::new().resource(resource1).resource(resource2);
3203
3204 router.subscribe("file:///a.txt");
3206 router.subscribe("file:///b.txt");
3207
3208 let uris = router.subscribed_uris();
3209 assert_eq!(uris.len(), 2);
3210 assert!(uris.contains(&"file:///a.txt".to_string()));
3211 assert!(uris.contains(&"file:///b.txt".to_string()));
3212 }
3213
3214 #[tokio::test]
3215 async fn test_subscription_capability_advertised() {
3216 use crate::resource::ResourceBuilder;
3217
3218 let resource = ResourceBuilder::new("file:///test.txt")
3219 .name("Test")
3220 .text("Hello");
3221
3222 let mut router = McpRouter::new().resource(resource);
3223
3224 let init_req = RouterRequest {
3226 id: RequestId::Number(0),
3227 inner: McpRequest::Initialize(InitializeParams {
3228 protocol_version: "2025-11-25".to_string(),
3229 capabilities: ClientCapabilities {
3230 roots: None,
3231 sampling: None,
3232 elicitation: None,
3233 tasks: None,
3234 experimental: None,
3235 extensions: None,
3236 },
3237 client_info: Implementation {
3238 name: "test".to_string(),
3239 version: "1.0".to_string(),
3240 ..Default::default()
3241 },
3242 meta: None,
3243 }),
3244 extensions: Extensions::new(),
3245 };
3246 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3247
3248 match resp.inner {
3249 Ok(McpResponse::Initialize(result)) => {
3250 let resources_cap = result.capabilities.resources.unwrap();
3252 assert!(resources_cap.subscribe);
3253 }
3254 _ => panic!("Expected Initialize response"),
3255 }
3256 }
3257
3258 #[tokio::test]
3259 async fn test_completion_handler() {
3260 let router = McpRouter::new()
3261 .server_info("test", "1.0")
3262 .completion_handler(|params: CompleteParams| async move {
3263 let prefix = ¶ms.argument.value;
3265 let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
3266 .into_iter()
3267 .filter(|s| s.starts_with(prefix))
3268 .map(String::from)
3269 .collect();
3270 Ok(CompleteResult::new(suggestions))
3271 });
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::default(),
3279 client_info: Implementation {
3280 name: "test".to_string(),
3281 version: "1.0".to_string(),
3282 ..Default::default()
3283 },
3284 meta: None,
3285 }),
3286 extensions: Extensions::new(),
3287 };
3288 let resp = router
3289 .clone()
3290 .ready()
3291 .await
3292 .unwrap()
3293 .call(init_req)
3294 .await
3295 .unwrap();
3296
3297 match resp.inner {
3299 Ok(McpResponse::Initialize(result)) => {
3300 assert!(result.capabilities.completions.is_some());
3301 }
3302 _ => panic!("Expected Initialize response"),
3303 }
3304
3305 router.handle_notification(McpNotification::Initialized);
3307
3308 let complete_req = RouterRequest {
3310 id: RequestId::Number(1),
3311 inner: McpRequest::Complete(CompleteParams {
3312 reference: CompletionReference::prompt("test-prompt"),
3313 argument: CompletionArgument::new("query", "al"),
3314 context: None,
3315 meta: None,
3316 }),
3317 extensions: Extensions::new(),
3318 };
3319 let resp = router
3320 .clone()
3321 .ready()
3322 .await
3323 .unwrap()
3324 .call(complete_req)
3325 .await
3326 .unwrap();
3327
3328 match resp.inner {
3329 Ok(McpResponse::Complete(result)) => {
3330 assert_eq!(result.completion.values, vec!["alpha"]);
3331 }
3332 _ => panic!("Expected Complete response"),
3333 }
3334 }
3335
3336 #[tokio::test]
3337 async fn test_completion_without_handler_returns_empty() {
3338 let router = McpRouter::new().server_info("test", "1.0");
3339
3340 let init_req = RouterRequest {
3342 id: RequestId::Number(0),
3343 inner: McpRequest::Initialize(InitializeParams {
3344 protocol_version: "2025-11-25".to_string(),
3345 capabilities: ClientCapabilities::default(),
3346 client_info: Implementation {
3347 name: "test".to_string(),
3348 version: "1.0".to_string(),
3349 ..Default::default()
3350 },
3351 meta: None,
3352 }),
3353 extensions: Extensions::new(),
3354 };
3355 let resp = router
3356 .clone()
3357 .ready()
3358 .await
3359 .unwrap()
3360 .call(init_req)
3361 .await
3362 .unwrap();
3363
3364 match resp.inner {
3366 Ok(McpResponse::Initialize(result)) => {
3367 assert!(result.capabilities.completions.is_none());
3368 }
3369 _ => panic!("Expected Initialize response"),
3370 }
3371
3372 router.handle_notification(McpNotification::Initialized);
3374
3375 let complete_req = RouterRequest {
3377 id: RequestId::Number(1),
3378 inner: McpRequest::Complete(CompleteParams {
3379 reference: CompletionReference::prompt("test-prompt"),
3380 argument: CompletionArgument::new("query", "al"),
3381 context: None,
3382 meta: None,
3383 }),
3384 extensions: Extensions::new(),
3385 };
3386 let resp = router
3387 .clone()
3388 .ready()
3389 .await
3390 .unwrap()
3391 .call(complete_req)
3392 .await
3393 .unwrap();
3394
3395 match resp.inner {
3396 Ok(McpResponse::Complete(result)) => {
3397 assert!(result.completion.values.is_empty());
3398 }
3399 _ => panic!("Expected Complete response"),
3400 }
3401 }
3402
3403 #[tokio::test]
3404 async fn test_tool_filter_list() {
3405 use crate::filter::CapabilityFilter;
3406 use crate::tool::Tool;
3407
3408 let public_tool = ToolBuilder::new("public")
3409 .description("Public tool")
3410 .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
3411 .build();
3412
3413 let admin_tool = ToolBuilder::new("admin")
3414 .description("Admin tool")
3415 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3416 .build();
3417
3418 let mut router = McpRouter::new()
3419 .tool(public_tool)
3420 .tool(admin_tool)
3421 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
3422
3423 init_router(&mut router).await;
3425
3426 let req = RouterRequest {
3427 id: RequestId::Number(1),
3428 inner: McpRequest::ListTools(ListToolsParams::default()),
3429 extensions: Extensions::new(),
3430 };
3431
3432 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3433
3434 match resp.inner {
3435 Ok(McpResponse::ListTools(result)) => {
3436 assert_eq!(result.tools.len(), 1);
3438 assert_eq!(result.tools[0].name, "public");
3439 }
3440 _ => panic!("Expected ListTools response"),
3441 }
3442 }
3443
3444 #[tokio::test]
3445 async fn test_tool_filter_call_denied() {
3446 use crate::filter::CapabilityFilter;
3447 use crate::tool::Tool;
3448
3449 let admin_tool = ToolBuilder::new("admin")
3450 .description("Admin tool")
3451 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3452 .build();
3453
3454 let mut router = McpRouter::new()
3455 .tool(admin_tool)
3456 .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); init_router(&mut router).await;
3460
3461 let req = RouterRequest {
3462 id: RequestId::Number(1),
3463 inner: McpRequest::CallTool(CallToolParams {
3464 name: "admin".to_string(),
3465 arguments: serde_json::json!({"a": 1, "b": 2}),
3466 meta: None,
3467 task: None,
3468 }),
3469 extensions: Extensions::new(),
3470 };
3471
3472 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3473
3474 match resp.inner {
3476 Err(e) => {
3477 assert_eq!(e.code, -32601); }
3479 _ => panic!("Expected JsonRpc error"),
3480 }
3481 }
3482
3483 #[tokio::test]
3484 async fn test_tool_filter_call_allowed() {
3485 use crate::filter::CapabilityFilter;
3486 use crate::tool::Tool;
3487
3488 let public_tool = ToolBuilder::new("public")
3489 .description("Public tool")
3490 .handler(|input: AddInput| async move {
3491 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3492 })
3493 .build();
3494
3495 let mut router = McpRouter::new()
3496 .tool(public_tool)
3497 .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); init_router(&mut router).await;
3501
3502 let req = RouterRequest {
3503 id: RequestId::Number(1),
3504 inner: McpRequest::CallTool(CallToolParams {
3505 name: "public".to_string(),
3506 arguments: serde_json::json!({"a": 1, "b": 2}),
3507 meta: None,
3508 task: None,
3509 }),
3510 extensions: Extensions::new(),
3511 };
3512
3513 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3514
3515 match resp.inner {
3516 Ok(McpResponse::CallTool(result)) => {
3517 assert!(!result.is_error);
3518 }
3519 _ => panic!("Expected CallTool response"),
3520 }
3521 }
3522
3523 #[tokio::test]
3524 async fn test_tool_filter_custom_denial() {
3525 use crate::filter::{CapabilityFilter, DenialBehavior};
3526 use crate::tool::Tool;
3527
3528 let admin_tool = ToolBuilder::new("admin")
3529 .description("Admin tool")
3530 .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3531 .build();
3532
3533 let mut router = McpRouter::new().tool(admin_tool).tool_filter(
3534 CapabilityFilter::new(|_, _: &Tool| false)
3535 .denial_behavior(DenialBehavior::Unauthorized),
3536 );
3537
3538 init_router(&mut router).await;
3540
3541 let req = RouterRequest {
3542 id: RequestId::Number(1),
3543 inner: McpRequest::CallTool(CallToolParams {
3544 name: "admin".to_string(),
3545 arguments: serde_json::json!({"a": 1, "b": 2}),
3546 meta: None,
3547 task: None,
3548 }),
3549 extensions: Extensions::new(),
3550 };
3551
3552 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3553
3554 match resp.inner {
3556 Err(e) => {
3557 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
3559 }
3560 _ => panic!("Expected JsonRpc error"),
3561 }
3562 }
3563
3564 #[tokio::test]
3565 async fn test_resource_filter_list() {
3566 use crate::filter::CapabilityFilter;
3567 use crate::resource::{Resource, ResourceBuilder};
3568
3569 let public_resource = ResourceBuilder::new("file:///public.txt")
3570 .name("Public File")
3571 .text("public content");
3572
3573 let secret_resource = ResourceBuilder::new("file:///secret.txt")
3574 .name("Secret File")
3575 .text("secret content");
3576
3577 let mut router = McpRouter::new()
3578 .resource(public_resource)
3579 .resource(secret_resource)
3580 .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
3581 !r.name.contains("Secret")
3582 }));
3583
3584 init_router(&mut router).await;
3586
3587 let req = RouterRequest {
3588 id: RequestId::Number(1),
3589 inner: McpRequest::ListResources(ListResourcesParams::default()),
3590 extensions: Extensions::new(),
3591 };
3592
3593 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3594
3595 match resp.inner {
3596 Ok(McpResponse::ListResources(result)) => {
3597 assert_eq!(result.resources.len(), 1);
3599 assert_eq!(result.resources[0].name, "Public File");
3600 }
3601 _ => panic!("Expected ListResources response"),
3602 }
3603 }
3604
3605 #[tokio::test]
3606 async fn test_resource_filter_read_denied() {
3607 use crate::filter::CapabilityFilter;
3608 use crate::resource::{Resource, ResourceBuilder};
3609
3610 let secret_resource = ResourceBuilder::new("file:///secret.txt")
3611 .name("Secret File")
3612 .text("secret content");
3613
3614 let mut router = McpRouter::new()
3615 .resource(secret_resource)
3616 .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); init_router(&mut router).await;
3620
3621 let req = RouterRequest {
3622 id: RequestId::Number(1),
3623 inner: McpRequest::ReadResource(ReadResourceParams {
3624 uri: "file:///secret.txt".to_string(),
3625 meta: None,
3626 }),
3627 extensions: Extensions::new(),
3628 };
3629
3630 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3631
3632 match resp.inner {
3634 Err(e) => {
3635 assert_eq!(e.code, -32601); }
3637 _ => panic!("Expected JsonRpc error"),
3638 }
3639 }
3640
3641 #[tokio::test]
3642 async fn test_resource_filter_read_allowed() {
3643 use crate::filter::CapabilityFilter;
3644 use crate::resource::{Resource, ResourceBuilder};
3645
3646 let public_resource = ResourceBuilder::new("file:///public.txt")
3647 .name("Public File")
3648 .text("public content");
3649
3650 let mut router = McpRouter::new()
3651 .resource(public_resource)
3652 .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); init_router(&mut router).await;
3656
3657 let req = RouterRequest {
3658 id: RequestId::Number(1),
3659 inner: McpRequest::ReadResource(ReadResourceParams {
3660 uri: "file:///public.txt".to_string(),
3661 meta: None,
3662 }),
3663 extensions: Extensions::new(),
3664 };
3665
3666 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3667
3668 match resp.inner {
3669 Ok(McpResponse::ReadResource(result)) => {
3670 assert_eq!(result.contents.len(), 1);
3671 assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
3672 }
3673 _ => panic!("Expected ReadResource response"),
3674 }
3675 }
3676
3677 #[tokio::test]
3678 async fn test_resource_filter_custom_denial() {
3679 use crate::filter::{CapabilityFilter, DenialBehavior};
3680 use crate::resource::{Resource, ResourceBuilder};
3681
3682 let secret_resource = ResourceBuilder::new("file:///secret.txt")
3683 .name("Secret File")
3684 .text("secret content");
3685
3686 let mut router = McpRouter::new().resource(secret_resource).resource_filter(
3687 CapabilityFilter::new(|_, _: &Resource| false)
3688 .denial_behavior(DenialBehavior::Unauthorized),
3689 );
3690
3691 init_router(&mut router).await;
3693
3694 let req = RouterRequest {
3695 id: RequestId::Number(1),
3696 inner: McpRequest::ReadResource(ReadResourceParams {
3697 uri: "file:///secret.txt".to_string(),
3698 meta: None,
3699 }),
3700 extensions: Extensions::new(),
3701 };
3702
3703 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3704
3705 match resp.inner {
3707 Err(e) => {
3708 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
3710 }
3711 _ => panic!("Expected JsonRpc error"),
3712 }
3713 }
3714
3715 #[tokio::test]
3716 async fn test_prompt_filter_list() {
3717 use crate::filter::CapabilityFilter;
3718 use crate::prompt::{Prompt, PromptBuilder};
3719
3720 let public_prompt = PromptBuilder::new("greeting")
3721 .description("A greeting")
3722 .user_message("Hello!");
3723
3724 let admin_prompt = PromptBuilder::new("system_debug")
3725 .description("Admin prompt")
3726 .user_message("Debug");
3727
3728 let mut router = McpRouter::new()
3729 .prompt(public_prompt)
3730 .prompt(admin_prompt)
3731 .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
3732 !p.name.contains("system")
3733 }));
3734
3735 init_router(&mut router).await;
3737
3738 let req = RouterRequest {
3739 id: RequestId::Number(1),
3740 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
3741 extensions: Extensions::new(),
3742 };
3743
3744 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3745
3746 match resp.inner {
3747 Ok(McpResponse::ListPrompts(result)) => {
3748 assert_eq!(result.prompts.len(), 1);
3750 assert_eq!(result.prompts[0].name, "greeting");
3751 }
3752 _ => panic!("Expected ListPrompts response"),
3753 }
3754 }
3755
3756 #[tokio::test]
3757 async fn test_prompt_filter_get_denied() {
3758 use crate::filter::CapabilityFilter;
3759 use crate::prompt::{Prompt, PromptBuilder};
3760 use std::collections::HashMap;
3761
3762 let admin_prompt = PromptBuilder::new("system_debug")
3763 .description("Admin prompt")
3764 .user_message("Debug");
3765
3766 let mut router = McpRouter::new()
3767 .prompt(admin_prompt)
3768 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); init_router(&mut router).await;
3772
3773 let req = RouterRequest {
3774 id: RequestId::Number(1),
3775 inner: McpRequest::GetPrompt(GetPromptParams {
3776 name: "system_debug".to_string(),
3777 arguments: HashMap::new(),
3778 meta: None,
3779 }),
3780 extensions: Extensions::new(),
3781 };
3782
3783 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3784
3785 match resp.inner {
3787 Err(e) => {
3788 assert_eq!(e.code, -32601); }
3790 _ => panic!("Expected JsonRpc error"),
3791 }
3792 }
3793
3794 #[tokio::test]
3795 async fn test_prompt_filter_get_allowed() {
3796 use crate::filter::CapabilityFilter;
3797 use crate::prompt::{Prompt, PromptBuilder};
3798 use std::collections::HashMap;
3799
3800 let public_prompt = PromptBuilder::new("greeting")
3801 .description("A greeting")
3802 .user_message("Hello!");
3803
3804 let mut router = McpRouter::new()
3805 .prompt(public_prompt)
3806 .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); init_router(&mut router).await;
3810
3811 let req = RouterRequest {
3812 id: RequestId::Number(1),
3813 inner: McpRequest::GetPrompt(GetPromptParams {
3814 name: "greeting".to_string(),
3815 arguments: HashMap::new(),
3816 meta: None,
3817 }),
3818 extensions: Extensions::new(),
3819 };
3820
3821 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3822
3823 match resp.inner {
3824 Ok(McpResponse::GetPrompt(result)) => {
3825 assert_eq!(result.messages.len(), 1);
3826 }
3827 _ => panic!("Expected GetPrompt response"),
3828 }
3829 }
3830
3831 #[tokio::test]
3832 async fn test_prompt_filter_custom_denial() {
3833 use crate::filter::{CapabilityFilter, DenialBehavior};
3834 use crate::prompt::{Prompt, PromptBuilder};
3835 use std::collections::HashMap;
3836
3837 let admin_prompt = PromptBuilder::new("system_debug")
3838 .description("Admin prompt")
3839 .user_message("Debug");
3840
3841 let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
3842 CapabilityFilter::new(|_, _: &Prompt| false)
3843 .denial_behavior(DenialBehavior::Unauthorized),
3844 );
3845
3846 init_router(&mut router).await;
3848
3849 let req = RouterRequest {
3850 id: RequestId::Number(1),
3851 inner: McpRequest::GetPrompt(GetPromptParams {
3852 name: "system_debug".to_string(),
3853 arguments: HashMap::new(),
3854 meta: None,
3855 }),
3856 extensions: Extensions::new(),
3857 };
3858
3859 let resp = router.ready().await.unwrap().call(req).await.unwrap();
3860
3861 match resp.inner {
3863 Err(e) => {
3864 assert_eq!(e.code, -32007); assert!(e.message.contains("Unauthorized"));
3866 }
3867 _ => panic!("Expected JsonRpc error"),
3868 }
3869 }
3870
3871 #[derive(Debug, Deserialize, JsonSchema)]
3876 struct StringInput {
3877 value: String,
3878 }
3879
3880 #[tokio::test]
3881 async fn test_router_merge_tools() {
3882 let tool_a = ToolBuilder::new("tool_a")
3884 .description("Tool A")
3885 .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
3886 .build();
3887
3888 let router_a = McpRouter::new().tool(tool_a);
3889
3890 let tool_b = ToolBuilder::new("tool_b")
3892 .description("Tool B")
3893 .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
3894 .build();
3895 let tool_c = ToolBuilder::new("tool_c")
3896 .description("Tool C")
3897 .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
3898 .build();
3899
3900 let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
3901
3902 let mut merged = McpRouter::new()
3904 .server_info("merged", "1.0")
3905 .merge(router_a)
3906 .merge(router_b);
3907
3908 init_router(&mut merged).await;
3909
3910 let req = RouterRequest {
3912 id: RequestId::Number(1),
3913 inner: McpRequest::ListTools(ListToolsParams::default()),
3914 extensions: Extensions::new(),
3915 };
3916
3917 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
3918
3919 match resp.inner {
3920 Ok(McpResponse::ListTools(result)) => {
3921 assert_eq!(result.tools.len(), 3);
3922 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
3923 assert!(names.contains(&"tool_a"));
3924 assert!(names.contains(&"tool_b"));
3925 assert!(names.contains(&"tool_c"));
3926 }
3927 _ => panic!("Expected ListTools response"),
3928 }
3929 }
3930
3931 #[tokio::test]
3932 async fn test_router_merge_overwrites_duplicates() {
3933 let tool_v1 = ToolBuilder::new("shared")
3935 .description("Version 1")
3936 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
3937 .build();
3938
3939 let router_a = McpRouter::new().tool(tool_v1);
3940
3941 let tool_v2 = ToolBuilder::new("shared")
3943 .description("Version 2")
3944 .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
3945 .build();
3946
3947 let router_b = McpRouter::new().tool(tool_v2);
3948
3949 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
3951
3952 init_router(&mut merged).await;
3953
3954 let req = RouterRequest {
3955 id: RequestId::Number(1),
3956 inner: McpRequest::ListTools(ListToolsParams::default()),
3957 extensions: Extensions::new(),
3958 };
3959
3960 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
3961
3962 match resp.inner {
3963 Ok(McpResponse::ListTools(result)) => {
3964 assert_eq!(result.tools.len(), 1);
3965 assert_eq!(result.tools[0].name, "shared");
3966 assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
3967 }
3968 _ => panic!("Expected ListTools response"),
3969 }
3970 }
3971
3972 #[tokio::test]
3973 async fn test_router_merge_resources() {
3974 use crate::resource::ResourceBuilder;
3975
3976 let router_a = McpRouter::new().resource(
3978 ResourceBuilder::new("file:///a.txt")
3979 .name("File A")
3980 .text("content a"),
3981 );
3982
3983 let router_b = McpRouter::new().resource(
3984 ResourceBuilder::new("file:///b.txt")
3985 .name("File B")
3986 .text("content b"),
3987 );
3988
3989 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
3990
3991 init_router(&mut merged).await;
3992
3993 let req = RouterRequest {
3994 id: RequestId::Number(1),
3995 inner: McpRequest::ListResources(ListResourcesParams::default()),
3996 extensions: Extensions::new(),
3997 };
3998
3999 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4000
4001 match resp.inner {
4002 Ok(McpResponse::ListResources(result)) => {
4003 assert_eq!(result.resources.len(), 2);
4004 let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4005 assert!(uris.contains(&"file:///a.txt"));
4006 assert!(uris.contains(&"file:///b.txt"));
4007 }
4008 _ => panic!("Expected ListResources response"),
4009 }
4010 }
4011
4012 #[tokio::test]
4013 async fn test_router_merge_prompts() {
4014 use crate::prompt::PromptBuilder;
4015
4016 let router_a =
4017 McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4018
4019 let router_b =
4020 McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4021
4022 let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4023
4024 init_router(&mut merged).await;
4025
4026 let req = RouterRequest {
4027 id: RequestId::Number(1),
4028 inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4029 extensions: Extensions::new(),
4030 };
4031
4032 let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4033
4034 match resp.inner {
4035 Ok(McpResponse::ListPrompts(result)) => {
4036 assert_eq!(result.prompts.len(), 2);
4037 let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4038 assert!(names.contains(&"prompt_a"));
4039 assert!(names.contains(&"prompt_b"));
4040 }
4041 _ => panic!("Expected ListPrompts response"),
4042 }
4043 }
4044
4045 #[tokio::test]
4046 async fn test_router_nest_prefixes_tools() {
4047 let tool_query = ToolBuilder::new("query")
4049 .description("Query the database")
4050 .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4051 .build();
4052 let tool_insert = ToolBuilder::new("insert")
4053 .description("Insert into database")
4054 .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4055 .build();
4056
4057 let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4058
4059 let mut router = McpRouter::new()
4061 .server_info("nested", "1.0")
4062 .nest("db", db_router);
4063
4064 init_router(&mut router).await;
4065
4066 let req = RouterRequest {
4067 id: RequestId::Number(1),
4068 inner: McpRequest::ListTools(ListToolsParams::default()),
4069 extensions: Extensions::new(),
4070 };
4071
4072 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4073
4074 match resp.inner {
4075 Ok(McpResponse::ListTools(result)) => {
4076 assert_eq!(result.tools.len(), 2);
4077 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4078 assert!(names.contains(&"db.query"));
4079 assert!(names.contains(&"db.insert"));
4080 }
4081 _ => panic!("Expected ListTools response"),
4082 }
4083 }
4084
4085 #[tokio::test]
4086 async fn test_router_nest_call_prefixed_tool() {
4087 let tool = ToolBuilder::new("echo")
4088 .description("Echo input")
4089 .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4090 .build();
4091
4092 let nested_router = McpRouter::new().tool(tool);
4093
4094 let mut router = McpRouter::new().nest("api", nested_router);
4095
4096 init_router(&mut router).await;
4097
4098 let req = RouterRequest {
4100 id: RequestId::Number(1),
4101 inner: McpRequest::CallTool(CallToolParams {
4102 name: "api.echo".to_string(),
4103 arguments: serde_json::json!({"value": "hello world"}),
4104 meta: None,
4105 task: None,
4106 }),
4107 extensions: Extensions::new(),
4108 };
4109
4110 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4111
4112 match resp.inner {
4113 Ok(McpResponse::CallTool(result)) => {
4114 assert!(!result.is_error);
4115 match &result.content[0] {
4116 Content::Text { text, .. } => assert_eq!(text, "hello world"),
4117 _ => panic!("Expected text content"),
4118 }
4119 }
4120 _ => panic!("Expected CallTool response"),
4121 }
4122 }
4123
4124 #[tokio::test]
4125 async fn test_router_multiple_nests() {
4126 let db_tool = ToolBuilder::new("query")
4127 .description("Database query")
4128 .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4129 .build();
4130
4131 let api_tool = ToolBuilder::new("fetch")
4132 .description("API fetch")
4133 .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4134 .build();
4135
4136 let db_router = McpRouter::new().tool(db_tool);
4137 let api_router = McpRouter::new().tool(api_tool);
4138
4139 let mut router = McpRouter::new()
4140 .nest("db", db_router)
4141 .nest("api", api_router);
4142
4143 init_router(&mut router).await;
4144
4145 let req = RouterRequest {
4146 id: RequestId::Number(1),
4147 inner: McpRequest::ListTools(ListToolsParams::default()),
4148 extensions: Extensions::new(),
4149 };
4150
4151 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4152
4153 match resp.inner {
4154 Ok(McpResponse::ListTools(result)) => {
4155 assert_eq!(result.tools.len(), 2);
4156 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4157 assert!(names.contains(&"db.query"));
4158 assert!(names.contains(&"api.fetch"));
4159 }
4160 _ => panic!("Expected ListTools response"),
4161 }
4162 }
4163
4164 #[tokio::test]
4165 async fn test_router_merge_and_nest_combined() {
4166 let tool_a = ToolBuilder::new("local")
4168 .description("Local tool")
4169 .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
4170 .build();
4171
4172 let nested_tool = ToolBuilder::new("remote")
4173 .description("Remote tool")
4174 .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
4175 .build();
4176
4177 let nested_router = McpRouter::new().tool(nested_tool);
4178
4179 let mut router = McpRouter::new()
4180 .tool(tool_a)
4181 .nest("external", nested_router);
4182
4183 init_router(&mut router).await;
4184
4185 let req = RouterRequest {
4186 id: RequestId::Number(1),
4187 inner: McpRequest::ListTools(ListToolsParams::default()),
4188 extensions: Extensions::new(),
4189 };
4190
4191 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4192
4193 match resp.inner {
4194 Ok(McpResponse::ListTools(result)) => {
4195 assert_eq!(result.tools.len(), 2);
4196 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4197 assert!(names.contains(&"local"));
4198 assert!(names.contains(&"external.remote"));
4199 }
4200 _ => panic!("Expected ListTools response"),
4201 }
4202 }
4203
4204 #[tokio::test]
4205 async fn test_router_merge_preserves_server_info() {
4206 let child_router = McpRouter::new()
4207 .server_info("child", "2.0")
4208 .instructions("Child instructions");
4209
4210 let mut router = McpRouter::new()
4211 .server_info("parent", "1.0")
4212 .instructions("Parent instructions")
4213 .merge(child_router);
4214
4215 init_router(&mut router).await;
4216
4217 let init_req = RouterRequest {
4219 id: RequestId::Number(99),
4220 inner: McpRequest::Initialize(InitializeParams {
4221 protocol_version: "2025-11-25".to_string(),
4222 capabilities: ClientCapabilities::default(),
4223 client_info: Implementation {
4224 name: "test".to_string(),
4225 version: "1.0".to_string(),
4226 ..Default::default()
4227 },
4228 meta: None,
4229 }),
4230 extensions: Extensions::new(),
4231 };
4232
4233 let child_router2 = McpRouter::new().server_info("child", "2.0");
4235 let mut fresh_router = McpRouter::new()
4236 .server_info("parent", "1.0")
4237 .merge(child_router2);
4238
4239 let resp = fresh_router
4240 .ready()
4241 .await
4242 .unwrap()
4243 .call(init_req)
4244 .await
4245 .unwrap();
4246
4247 match resp.inner {
4248 Ok(McpResponse::Initialize(result)) => {
4249 assert_eq!(result.server_info.name, "parent");
4250 assert_eq!(result.server_info.version, "1.0");
4251 }
4252 _ => panic!("Expected Initialize response"),
4253 }
4254 }
4255
4256 #[tokio::test]
4261 async fn test_auto_instructions_tools_only() {
4262 let tool_a = ToolBuilder::new("alpha")
4263 .description("Alpha tool")
4264 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4265 .build();
4266 let tool_b = ToolBuilder::new("beta")
4267 .description("Beta tool")
4268 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4269 .build();
4270
4271 let mut router = McpRouter::new()
4272 .auto_instructions()
4273 .tool(tool_a)
4274 .tool(tool_b);
4275
4276 let resp = send_initialize(&mut router).await;
4277 let instructions = resp.instructions.expect("should have instructions");
4278
4279 assert!(instructions.contains("## Tools"));
4280 assert!(instructions.contains("- **alpha**: Alpha tool"));
4281 assert!(instructions.contains("- **beta**: Beta tool"));
4282 assert!(!instructions.contains("## Resources"));
4284 assert!(!instructions.contains("## Prompts"));
4285 }
4286
4287 #[tokio::test]
4288 async fn test_auto_instructions_with_annotations() {
4289 let read_only_tool = ToolBuilder::new("query")
4290 .description("Run a query")
4291 .read_only()
4292 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4293 .build();
4294 let destructive_tool = ToolBuilder::new("delete")
4295 .description("Delete a record")
4296 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4297 .build();
4298 let idempotent_tool = ToolBuilder::new("upsert")
4299 .description("Upsert a record")
4300 .non_destructive()
4301 .idempotent()
4302 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4303 .build();
4304
4305 let mut router = McpRouter::new()
4306 .auto_instructions()
4307 .tool(read_only_tool)
4308 .tool(destructive_tool)
4309 .tool(idempotent_tool);
4310
4311 let resp = send_initialize(&mut router).await;
4312 let instructions = resp.instructions.unwrap();
4313
4314 assert!(instructions.contains("- **query**: Run a query [read-only]"));
4315 assert!(instructions.contains("- **delete**: Delete a record\n"));
4317 assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
4318 }
4319
4320 #[tokio::test]
4321 async fn test_auto_instructions_with_resources() {
4322 use crate::resource::ResourceBuilder;
4323
4324 let resource = ResourceBuilder::new("file:///schema.sql")
4325 .name("Schema")
4326 .description("Database schema")
4327 .text("CREATE TABLE ...");
4328
4329 let mut router = McpRouter::new().auto_instructions().resource(resource);
4330
4331 let resp = send_initialize(&mut router).await;
4332 let instructions = resp.instructions.unwrap();
4333
4334 assert!(instructions.contains("## Resources"));
4335 assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
4336 assert!(!instructions.contains("## Tools"));
4337 }
4338
4339 #[tokio::test]
4340 async fn test_auto_instructions_with_resource_templates() {
4341 use crate::resource::ResourceTemplateBuilder;
4342
4343 let template = ResourceTemplateBuilder::new("file:///{path}")
4344 .name("File")
4345 .description("Read a file by path")
4346 .handler(
4347 |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
4348 Ok(crate::ReadResourceResult::text("content", "text/plain"))
4349 },
4350 );
4351
4352 let mut router = McpRouter::new()
4353 .auto_instructions()
4354 .resource_template(template);
4355
4356 let resp = send_initialize(&mut router).await;
4357 let instructions = resp.instructions.unwrap();
4358
4359 assert!(instructions.contains("## Resources"));
4360 assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
4361 }
4362
4363 #[tokio::test]
4364 async fn test_auto_instructions_with_prompts() {
4365 use crate::prompt::PromptBuilder;
4366
4367 let prompt = PromptBuilder::new("write_query")
4368 .description("Help write a SQL query")
4369 .user_message("Write a query for: {task}");
4370
4371 let mut router = McpRouter::new().auto_instructions().prompt(prompt);
4372
4373 let resp = send_initialize(&mut router).await;
4374 let instructions = resp.instructions.unwrap();
4375
4376 assert!(instructions.contains("## Prompts"));
4377 assert!(instructions.contains("- **write_query**: Help write a SQL query"));
4378 assert!(!instructions.contains("## Tools"));
4379 }
4380
4381 #[tokio::test]
4382 async fn test_auto_instructions_all_sections() {
4383 use crate::prompt::PromptBuilder;
4384 use crate::resource::ResourceBuilder;
4385
4386 let tool = ToolBuilder::new("query")
4387 .description("Execute SQL")
4388 .read_only()
4389 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4390 .build();
4391 let resource = ResourceBuilder::new("db://schema")
4392 .name("Schema")
4393 .description("Full database schema")
4394 .text("schema");
4395 let prompt = PromptBuilder::new("write_query")
4396 .description("Help write a SQL query")
4397 .user_message("Write a query");
4398
4399 let mut router = McpRouter::new()
4400 .auto_instructions()
4401 .tool(tool)
4402 .resource(resource)
4403 .prompt(prompt);
4404
4405 let resp = send_initialize(&mut router).await;
4406 let instructions = resp.instructions.unwrap();
4407
4408 assert!(instructions.contains("## Tools"));
4410 assert!(instructions.contains("## Resources"));
4411 assert!(instructions.contains("## Prompts"));
4412
4413 let tools_pos = instructions.find("## Tools").unwrap();
4415 let resources_pos = instructions.find("## Resources").unwrap();
4416 let prompts_pos = instructions.find("## Prompts").unwrap();
4417 assert!(tools_pos < resources_pos);
4418 assert!(resources_pos < prompts_pos);
4419 }
4420
4421 #[tokio::test]
4422 async fn test_auto_instructions_with_prefix_and_suffix() {
4423 let tool = ToolBuilder::new("echo")
4424 .description("Echo input")
4425 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4426 .build();
4427
4428 let mut router = McpRouter::new()
4429 .auto_instructions_with(
4430 Some("This server provides echo capabilities."),
4431 Some("Contact admin@example.com for support."),
4432 )
4433 .tool(tool);
4434
4435 let resp = send_initialize(&mut router).await;
4436 let instructions = resp.instructions.unwrap();
4437
4438 assert!(instructions.starts_with("This server provides echo capabilities."));
4439 assert!(instructions.ends_with("Contact admin@example.com for support."));
4440 assert!(instructions.contains("## Tools"));
4441 assert!(instructions.contains("- **echo**: Echo input"));
4442 }
4443
4444 #[tokio::test]
4445 async fn test_auto_instructions_prefix_only() {
4446 let tool = ToolBuilder::new("echo")
4447 .description("Echo input")
4448 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4449 .build();
4450
4451 let mut router = McpRouter::new()
4452 .auto_instructions_with(Some("My server intro."), None::<String>)
4453 .tool(tool);
4454
4455 let resp = send_initialize(&mut router).await;
4456 let instructions = resp.instructions.unwrap();
4457
4458 assert!(instructions.starts_with("My server intro."));
4459 assert!(instructions.contains("- **echo**: Echo input"));
4460 }
4461
4462 #[tokio::test]
4463 async fn test_auto_instructions_empty_router() {
4464 let mut router = McpRouter::new().auto_instructions();
4465
4466 let resp = send_initialize(&mut router).await;
4467 let instructions = resp.instructions.expect("should have instructions");
4468
4469 assert!(!instructions.contains("## Tools"));
4471 assert!(!instructions.contains("## Resources"));
4472 assert!(!instructions.contains("## Prompts"));
4473 assert!(instructions.is_empty());
4474 }
4475
4476 #[tokio::test]
4477 async fn test_auto_instructions_overrides_manual() {
4478 let tool = ToolBuilder::new("echo")
4479 .description("Echo input")
4480 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4481 .build();
4482
4483 let mut router = McpRouter::new()
4484 .instructions("This will be overridden")
4485 .auto_instructions()
4486 .tool(tool);
4487
4488 let resp = send_initialize(&mut router).await;
4489 let instructions = resp.instructions.unwrap();
4490
4491 assert!(!instructions.contains("This will be overridden"));
4492 assert!(instructions.contains("- **echo**: Echo input"));
4493 }
4494
4495 #[tokio::test]
4496 async fn test_no_auto_instructions_returns_manual() {
4497 let tool = ToolBuilder::new("echo")
4498 .description("Echo input")
4499 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4500 .build();
4501
4502 let mut router = McpRouter::new()
4503 .instructions("Manual instructions here")
4504 .tool(tool);
4505
4506 let resp = send_initialize(&mut router).await;
4507 let instructions = resp.instructions.unwrap();
4508
4509 assert_eq!(instructions, "Manual instructions here");
4510 }
4511
4512 #[tokio::test]
4513 async fn test_auto_instructions_no_description_fallback() {
4514 let tool = ToolBuilder::new("mystery")
4515 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4516 .build();
4517
4518 let mut router = McpRouter::new().auto_instructions().tool(tool);
4519
4520 let resp = send_initialize(&mut router).await;
4521 let instructions = resp.instructions.unwrap();
4522
4523 assert!(instructions.contains("- **mystery**: No description"));
4524 }
4525
4526 #[tokio::test]
4527 async fn test_auto_instructions_sorted_alphabetically() {
4528 let tool_z = ToolBuilder::new("zebra")
4529 .description("Z tool")
4530 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4531 .build();
4532 let tool_a = ToolBuilder::new("alpha")
4533 .description("A tool")
4534 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4535 .build();
4536 let tool_m = ToolBuilder::new("middle")
4537 .description("M tool")
4538 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4539 .build();
4540
4541 let mut router = McpRouter::new()
4542 .auto_instructions()
4543 .tool(tool_z)
4544 .tool(tool_a)
4545 .tool(tool_m);
4546
4547 let resp = send_initialize(&mut router).await;
4548 let instructions = resp.instructions.unwrap();
4549
4550 let alpha_pos = instructions.find("**alpha**").unwrap();
4551 let middle_pos = instructions.find("**middle**").unwrap();
4552 let zebra_pos = instructions.find("**zebra**").unwrap();
4553 assert!(alpha_pos < middle_pos);
4554 assert!(middle_pos < zebra_pos);
4555 }
4556
4557 #[tokio::test]
4558 async fn test_auto_instructions_read_only_and_idempotent_tags() {
4559 let tool = ToolBuilder::new("safe_update")
4560 .description("Safe update operation")
4561 .idempotent()
4562 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4563 .build();
4564
4565 let mut router = McpRouter::new().auto_instructions().tool(tool);
4566
4567 let resp = send_initialize(&mut router).await;
4568 let instructions = resp.instructions.unwrap();
4569
4570 assert!(
4571 instructions.contains("[idempotent]"),
4572 "got: {}",
4573 instructions
4574 );
4575 }
4576
4577 #[tokio::test]
4578 async fn test_auto_instructions_lazy_generation() {
4579 let mut router = McpRouter::new().auto_instructions();
4582
4583 let tool = ToolBuilder::new("late_tool")
4584 .description("Added after auto_instructions")
4585 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4586 .build();
4587
4588 router = router.tool(tool);
4589
4590 let resp = send_initialize(&mut router).await;
4591 let instructions = resp.instructions.unwrap();
4592
4593 assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
4594 }
4595
4596 #[tokio::test]
4597 async fn test_auto_instructions_multiple_annotation_tags() {
4598 let tool = ToolBuilder::new("update")
4599 .description("Update a record")
4600 .annotations(ToolAnnotations {
4601 read_only_hint: true,
4602 idempotent_hint: true,
4603 ..Default::default()
4604 })
4605 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4606 .build();
4607
4608 let mut router = McpRouter::new().auto_instructions().tool(tool);
4609
4610 let resp = send_initialize(&mut router).await;
4611 let instructions = resp.instructions.unwrap();
4612
4613 assert!(
4614 instructions.contains("[read-only, idempotent]"),
4615 "got: {}",
4616 instructions
4617 );
4618 }
4619
4620 #[tokio::test]
4621 async fn test_auto_instructions_no_annotations_no_tags() {
4622 let tool = ToolBuilder::new("fetch")
4624 .description("Fetch data")
4625 .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4626 .build();
4627
4628 let mut router = McpRouter::new().auto_instructions().tool(tool);
4629
4630 let resp = send_initialize(&mut router).await;
4631 let instructions = resp.instructions.unwrap();
4632
4633 assert!(
4635 !instructions.contains('['),
4636 "should have no tags, got: {}",
4637 instructions
4638 );
4639 assert!(instructions.contains("- **fetch**: Fetch data"));
4640 }
4641
4642 async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
4644 let init_req = RouterRequest {
4645 id: RequestId::Number(0),
4646 inner: McpRequest::Initialize(InitializeParams {
4647 protocol_version: "2025-11-25".to_string(),
4648 capabilities: ClientCapabilities {
4649 roots: None,
4650 sampling: None,
4651 elicitation: None,
4652 tasks: None,
4653 experimental: None,
4654 extensions: None,
4655 },
4656 client_info: Implementation {
4657 name: "test".to_string(),
4658 version: "1.0".to_string(),
4659 ..Default::default()
4660 },
4661 meta: None,
4662 }),
4663 extensions: Extensions::new(),
4664 };
4665 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
4666 match resp.inner {
4667 Ok(McpResponse::Initialize(result)) => result,
4668 other => panic!("Expected Initialize response, got {:?}", other),
4669 }
4670 }
4671
4672 #[tokio::test]
4673 async fn test_notify_tools_list_changed() {
4674 let (tx, mut rx) = crate::context::notification_channel(16);
4675
4676 let router = McpRouter::new()
4677 .server_info("test", "1.0")
4678 .with_notification_sender(tx);
4679
4680 assert!(router.notify_tools_list_changed());
4681
4682 let notification = rx.recv().await.unwrap();
4683 assert!(matches!(notification, ServerNotification::ToolsListChanged));
4684 }
4685
4686 #[tokio::test]
4687 async fn test_notify_prompts_list_changed() {
4688 let (tx, mut rx) = crate::context::notification_channel(16);
4689
4690 let router = McpRouter::new()
4691 .server_info("test", "1.0")
4692 .with_notification_sender(tx);
4693
4694 assert!(router.notify_prompts_list_changed());
4695
4696 let notification = rx.recv().await.unwrap();
4697 assert!(matches!(
4698 notification,
4699 ServerNotification::PromptsListChanged
4700 ));
4701 }
4702
4703 #[tokio::test]
4704 async fn test_notify_without_sender_returns_false() {
4705 let router = McpRouter::new().server_info("test", "1.0");
4706
4707 assert!(!router.notify_tools_list_changed());
4708 assert!(!router.notify_prompts_list_changed());
4709 assert!(!router.notify_resources_list_changed());
4710 }
4711
4712 #[tokio::test]
4713 async fn test_list_changed_capabilities_with_notification_sender() {
4714 let (tx, _rx) = crate::context::notification_channel(16);
4715 let tool = ToolBuilder::new("test")
4716 .description("test")
4717 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4718 .build();
4719
4720 let mut router = McpRouter::new()
4721 .server_info("test", "1.0")
4722 .tool(tool)
4723 .with_notification_sender(tx);
4724
4725 init_router(&mut router).await;
4726
4727 let caps = router.capabilities();
4728 let tools_cap = caps.tools.expect("tools capability should be present");
4729 assert!(
4730 tools_cap.list_changed,
4731 "tools.listChanged should be true when notification sender is configured"
4732 );
4733 }
4734
4735 #[tokio::test]
4736 async fn test_list_changed_capabilities_without_notification_sender() {
4737 let tool = ToolBuilder::new("test")
4738 .description("test")
4739 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4740 .build();
4741
4742 let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
4743
4744 init_router(&mut router).await;
4745
4746 let caps = router.capabilities();
4747 let tools_cap = caps.tools.expect("tools capability should be present");
4748 assert!(
4749 !tools_cap.list_changed,
4750 "tools.listChanged should be false without notification sender"
4751 );
4752 }
4753
4754 #[tokio::test]
4755 async fn test_set_logging_level_filters_messages() {
4756 let (tx, mut rx) = crate::context::notification_channel(16);
4757
4758 let mut router = McpRouter::new()
4759 .server_info("test", "1.0")
4760 .with_notification_sender(tx);
4761
4762 init_router(&mut router).await;
4763
4764 let set_level_req = RouterRequest {
4766 id: RequestId::Number(99),
4767 inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
4768 level: LogLevel::Warning,
4769 meta: None,
4770 }),
4771 extensions: crate::context::Extensions::new(),
4772 };
4773 let resp = router
4774 .ready()
4775 .await
4776 .unwrap()
4777 .call(set_level_req)
4778 .await
4779 .unwrap();
4780 assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
4781
4782 let ctx = router.create_context(RequestId::Number(100), None);
4784
4785 ctx.send_log(LoggingMessageParams::new(
4787 LogLevel::Error,
4788 serde_json::Value::Null,
4789 ));
4790 assert!(
4791 rx.try_recv().is_ok(),
4792 "Error should pass through Warning filter"
4793 );
4794
4795 ctx.send_log(LoggingMessageParams::new(
4797 LogLevel::Info,
4798 serde_json::Value::Null,
4799 ));
4800 assert!(
4801 rx.try_recv().is_err(),
4802 "Info should be filtered at Warning level"
4803 );
4804 }
4805
4806 #[test]
4807 fn test_paginate_no_page_size() {
4808 let items = vec![1, 2, 3, 4, 5];
4809 let (page, cursor) = paginate(items.clone(), None, None).unwrap();
4810 assert_eq!(page, items);
4811 assert!(cursor.is_none());
4812 }
4813
4814 #[test]
4815 fn test_paginate_first_page() {
4816 let items = vec![1, 2, 3, 4, 5];
4817 let (page, cursor) = paginate(items, None, Some(2)).unwrap();
4818 assert_eq!(page, vec![1, 2]);
4819 assert!(cursor.is_some());
4820 }
4821
4822 #[test]
4823 fn test_paginate_middle_page() {
4824 let items = vec![1, 2, 3, 4, 5];
4825 let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
4826 assert_eq!(page1, vec![1, 2]);
4827
4828 let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
4829 assert_eq!(page2, vec![3, 4]);
4830 assert!(cursor2.is_some());
4831 }
4832
4833 #[test]
4834 fn test_paginate_last_page() {
4835 let items = vec![1, 2, 3, 4, 5];
4836 let cursor = encode_cursor(4);
4838 let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
4839 assert_eq!(page, vec![5]);
4840 assert!(next.is_none());
4841 }
4842
4843 #[test]
4844 fn test_paginate_exact_boundary() {
4845 let items = vec![1, 2, 3, 4];
4846 let (page, cursor) = paginate(items, None, Some(4)).unwrap();
4847 assert_eq!(page, vec![1, 2, 3, 4]);
4848 assert!(cursor.is_none());
4849 }
4850
4851 #[test]
4852 fn test_paginate_invalid_cursor() {
4853 let items = vec![1, 2, 3];
4854 let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
4855 assert!(result.is_err());
4856 }
4857
4858 #[test]
4859 fn test_cursor_round_trip() {
4860 let offset = 42;
4861 let encoded = encode_cursor(offset);
4862 let decoded = decode_cursor(&encoded).unwrap();
4863 assert_eq!(decoded, offset);
4864 }
4865
4866 #[tokio::test]
4867 async fn test_list_tools_pagination() {
4868 let tool_a = ToolBuilder::new("alpha")
4869 .description("a")
4870 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4871 .build();
4872 let tool_b = ToolBuilder::new("beta")
4873 .description("b")
4874 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4875 .build();
4876 let tool_c = ToolBuilder::new("gamma")
4877 .description("c")
4878 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4879 .build();
4880
4881 let mut router = McpRouter::new()
4882 .server_info("test", "1.0")
4883 .page_size(2)
4884 .tool(tool_a)
4885 .tool(tool_b)
4886 .tool(tool_c);
4887
4888 init_router(&mut router).await;
4889
4890 let req = RouterRequest {
4892 id: RequestId::Number(1),
4893 inner: McpRequest::ListTools(ListToolsParams {
4894 cursor: None,
4895 meta: None,
4896 }),
4897 extensions: Extensions::new(),
4898 };
4899 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4900 let (tools, next_cursor) = match resp.inner {
4901 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
4902 other => panic!("Expected ListTools, got {:?}", other),
4903 };
4904 assert_eq!(tools.len(), 2);
4905 assert_eq!(tools[0].name, "alpha");
4906 assert_eq!(tools[1].name, "beta");
4907 assert!(next_cursor.is_some());
4908
4909 let req = RouterRequest {
4911 id: RequestId::Number(2),
4912 inner: McpRequest::ListTools(ListToolsParams {
4913 cursor: next_cursor,
4914 meta: None,
4915 }),
4916 extensions: Extensions::new(),
4917 };
4918 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4919 let (tools, next_cursor) = match resp.inner {
4920 Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
4921 other => panic!("Expected ListTools, got {:?}", other),
4922 };
4923 assert_eq!(tools.len(), 1);
4924 assert_eq!(tools[0].name, "gamma");
4925 assert!(next_cursor.is_none());
4926 }
4927
4928 #[tokio::test]
4929 async fn test_list_tools_no_pagination_by_default() {
4930 let tool_a = ToolBuilder::new("alpha")
4931 .description("a")
4932 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4933 .build();
4934 let tool_b = ToolBuilder::new("beta")
4935 .description("b")
4936 .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
4937 .build();
4938
4939 let mut router = McpRouter::new()
4940 .server_info("test", "1.0")
4941 .tool(tool_a)
4942 .tool(tool_b);
4943
4944 init_router(&mut router).await;
4945
4946 let req = RouterRequest {
4947 id: RequestId::Number(1),
4948 inner: McpRequest::ListTools(ListToolsParams {
4949 cursor: None,
4950 meta: None,
4951 }),
4952 extensions: Extensions::new(),
4953 };
4954 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4955 match resp.inner {
4956 Ok(McpResponse::ListTools(result)) => {
4957 assert_eq!(result.tools.len(), 2);
4958 assert!(result.next_cursor.is_none());
4959 }
4960 other => panic!("Expected ListTools, got {:?}", other),
4961 }
4962 }
4963
4964 #[cfg(feature = "dynamic-tools")]
4969 mod dynamic_tools_tests {
4970 use super::*;
4971
4972 #[tokio::test]
4973 async fn test_dynamic_tools_register_and_list() {
4974 let (router, registry) = McpRouter::new()
4975 .server_info("test", "1.0")
4976 .with_dynamic_tools();
4977
4978 let tool = ToolBuilder::new("dynamic_echo")
4979 .description("Dynamic echo")
4980 .handler(|input: AddInput| async move {
4981 Ok(CallToolResult::text(format!("{}", input.a)))
4982 })
4983 .build();
4984
4985 registry.register(tool);
4986
4987 let mut router = router;
4988 init_router(&mut router).await;
4989
4990 let req = RouterRequest {
4991 id: RequestId::Number(1),
4992 inner: McpRequest::ListTools(ListToolsParams::default()),
4993 extensions: Extensions::new(),
4994 };
4995
4996 let resp = router.ready().await.unwrap().call(req).await.unwrap();
4997 match resp.inner {
4998 Ok(McpResponse::ListTools(result)) => {
4999 assert_eq!(result.tools.len(), 1);
5000 assert_eq!(result.tools[0].name, "dynamic_echo");
5001 }
5002 _ => panic!("Expected ListTools response"),
5003 }
5004 }
5005
5006 #[tokio::test]
5007 async fn test_dynamic_tools_unregister() {
5008 let (router, registry) = McpRouter::new()
5009 .server_info("test", "1.0")
5010 .with_dynamic_tools();
5011
5012 let tool = ToolBuilder::new("temp")
5013 .description("Temporary")
5014 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5015 .build();
5016
5017 registry.register(tool);
5018 assert!(registry.contains("temp"));
5019
5020 let removed = registry.unregister("temp");
5021 assert!(removed);
5022 assert!(!registry.contains("temp"));
5023
5024 assert!(!registry.unregister("temp"));
5026
5027 let mut router = router;
5028 init_router(&mut router).await;
5029
5030 let req = RouterRequest {
5031 id: RequestId::Number(1),
5032 inner: McpRequest::ListTools(ListToolsParams::default()),
5033 extensions: Extensions::new(),
5034 };
5035
5036 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5037 match resp.inner {
5038 Ok(McpResponse::ListTools(result)) => {
5039 assert_eq!(result.tools.len(), 0);
5040 }
5041 _ => panic!("Expected ListTools response"),
5042 }
5043 }
5044
5045 #[tokio::test]
5046 async fn test_dynamic_tools_merged_with_static() {
5047 let static_tool = ToolBuilder::new("static_tool")
5048 .description("Static")
5049 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5050 .build();
5051
5052 let (router, registry) = McpRouter::new()
5053 .server_info("test", "1.0")
5054 .tool(static_tool)
5055 .with_dynamic_tools();
5056
5057 let dynamic_tool = ToolBuilder::new("dynamic_tool")
5058 .description("Dynamic")
5059 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5060 .build();
5061
5062 registry.register(dynamic_tool);
5063
5064 let mut router = router;
5065 init_router(&mut router).await;
5066
5067 let req = RouterRequest {
5068 id: RequestId::Number(1),
5069 inner: McpRequest::ListTools(ListToolsParams::default()),
5070 extensions: Extensions::new(),
5071 };
5072
5073 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5074 match resp.inner {
5075 Ok(McpResponse::ListTools(result)) => {
5076 assert_eq!(result.tools.len(), 2);
5077 let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5078 assert!(names.contains(&"static_tool"));
5079 assert!(names.contains(&"dynamic_tool"));
5080 }
5081 _ => panic!("Expected ListTools response"),
5082 }
5083 }
5084
5085 #[tokio::test]
5086 async fn test_static_tools_shadow_dynamic() {
5087 let static_tool = ToolBuilder::new("shared")
5088 .description("Static version")
5089 .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5090 .build();
5091
5092 let (router, registry) = McpRouter::new()
5093 .server_info("test", "1.0")
5094 .tool(static_tool)
5095 .with_dynamic_tools();
5096
5097 let dynamic_tool = ToolBuilder::new("shared")
5098 .description("Dynamic version")
5099 .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5100 .build();
5101
5102 registry.register(dynamic_tool);
5103
5104 let mut router = router;
5105 init_router(&mut router).await;
5106
5107 let req = RouterRequest {
5109 id: RequestId::Number(1),
5110 inner: McpRequest::ListTools(ListToolsParams::default()),
5111 extensions: Extensions::new(),
5112 };
5113
5114 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5115 match resp.inner {
5116 Ok(McpResponse::ListTools(result)) => {
5117 assert_eq!(result.tools.len(), 1);
5118 assert_eq!(result.tools[0].name, "shared");
5119 assert_eq!(
5120 result.tools[0].description.as_deref(),
5121 Some("Static version")
5122 );
5123 }
5124 _ => panic!("Expected ListTools response"),
5125 }
5126
5127 let req = RouterRequest {
5129 id: RequestId::Number(2),
5130 inner: McpRequest::CallTool(CallToolParams {
5131 name: "shared".to_string(),
5132 arguments: serde_json::json!({"a": 1, "b": 2}),
5133 meta: None,
5134 task: None,
5135 }),
5136 extensions: Extensions::new(),
5137 };
5138
5139 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5140 match resp.inner {
5141 Ok(McpResponse::CallTool(result)) => {
5142 assert!(!result.is_error);
5143 match &result.content[0] {
5144 Content::Text { text, .. } => assert_eq!(text, "static"),
5145 _ => panic!("Expected text content"),
5146 }
5147 }
5148 _ => panic!("Expected CallTool response"),
5149 }
5150 }
5151
5152 #[tokio::test]
5153 async fn test_dynamic_tools_call() {
5154 let (router, registry) = McpRouter::new()
5155 .server_info("test", "1.0")
5156 .with_dynamic_tools();
5157
5158 let tool = ToolBuilder::new("add")
5159 .description("Add two numbers")
5160 .handler(|input: AddInput| async move {
5161 Ok(CallToolResult::text(format!("{}", input.a + input.b)))
5162 })
5163 .build();
5164
5165 registry.register(tool);
5166
5167 let mut router = router;
5168 init_router(&mut router).await;
5169
5170 let req = RouterRequest {
5171 id: RequestId::Number(1),
5172 inner: McpRequest::CallTool(CallToolParams {
5173 name: "add".to_string(),
5174 arguments: serde_json::json!({"a": 3, "b": 4}),
5175 meta: None,
5176 task: None,
5177 }),
5178 extensions: Extensions::new(),
5179 };
5180
5181 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5182 match resp.inner {
5183 Ok(McpResponse::CallTool(result)) => {
5184 assert!(!result.is_error);
5185 match &result.content[0] {
5186 Content::Text { text, .. } => assert_eq!(text, "7"),
5187 _ => panic!("Expected text content"),
5188 }
5189 }
5190 _ => panic!("Expected CallTool response"),
5191 }
5192 }
5193
5194 #[tokio::test]
5195 async fn test_dynamic_tools_notification_on_register() {
5196 let (tx, mut rx) = crate::context::notification_channel(16);
5197 let (router, registry) = McpRouter::new()
5198 .server_info("test", "1.0")
5199 .with_dynamic_tools();
5200 let _router = router.with_notification_sender(tx);
5201
5202 let tool = ToolBuilder::new("notified")
5203 .description("Test")
5204 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5205 .build();
5206
5207 registry.register(tool);
5208
5209 let notification = rx.recv().await.unwrap();
5210 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5211 }
5212
5213 #[tokio::test]
5214 async fn test_dynamic_tools_notification_on_unregister() {
5215 let (tx, mut rx) = crate::context::notification_channel(16);
5216 let (router, registry) = McpRouter::new()
5217 .server_info("test", "1.0")
5218 .with_dynamic_tools();
5219 let _router = router.with_notification_sender(tx);
5220
5221 let tool = ToolBuilder::new("notified")
5222 .description("Test")
5223 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5224 .build();
5225
5226 registry.register(tool);
5227 let _ = rx.recv().await.unwrap();
5229
5230 registry.unregister("notified");
5231 let notification = rx.recv().await.unwrap();
5232 assert!(matches!(notification, ServerNotification::ToolsListChanged));
5233 }
5234
5235 #[tokio::test]
5236 async fn test_dynamic_tools_no_notification_on_empty_unregister() {
5237 let (tx, mut rx) = crate::context::notification_channel(16);
5238 let (router, registry) = McpRouter::new()
5239 .server_info("test", "1.0")
5240 .with_dynamic_tools();
5241 let _router = router.with_notification_sender(tx);
5242
5243 assert!(!registry.unregister("nonexistent"));
5245
5246 assert!(rx.try_recv().is_err());
5248 }
5249
5250 #[tokio::test]
5251 async fn test_dynamic_tools_filter_applies() {
5252 use crate::filter::CapabilityFilter;
5253
5254 let (router, registry) = McpRouter::new()
5255 .server_info("test", "1.0")
5256 .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
5257 tool.name != "hidden"
5258 }))
5259 .with_dynamic_tools();
5260
5261 let visible = ToolBuilder::new("visible")
5262 .description("Visible")
5263 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5264 .build();
5265
5266 let hidden = ToolBuilder::new("hidden")
5267 .description("Hidden")
5268 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5269 .build();
5270
5271 registry.register(visible);
5272 registry.register(hidden);
5273
5274 let mut router = router;
5275 init_router(&mut router).await;
5276
5277 let req = RouterRequest {
5279 id: RequestId::Number(1),
5280 inner: McpRequest::ListTools(ListToolsParams::default()),
5281 extensions: Extensions::new(),
5282 };
5283
5284 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5285 match resp.inner {
5286 Ok(McpResponse::ListTools(result)) => {
5287 assert_eq!(result.tools.len(), 1);
5288 assert_eq!(result.tools[0].name, "visible");
5289 }
5290 _ => panic!("Expected ListTools response"),
5291 }
5292
5293 let req = RouterRequest {
5295 id: RequestId::Number(2),
5296 inner: McpRequest::CallTool(CallToolParams {
5297 name: "hidden".to_string(),
5298 arguments: serde_json::json!({"a": 1, "b": 2}),
5299 meta: None,
5300 task: None,
5301 }),
5302 extensions: Extensions::new(),
5303 };
5304
5305 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5306 match resp.inner {
5307 Err(e) => {
5308 assert_eq!(e.code, -32601); }
5310 _ => panic!("Expected JsonRpc error"),
5311 }
5312 }
5313
5314 #[tokio::test]
5315 async fn test_dynamic_tools_capabilities_advertised() {
5316 let (mut router, _registry) = McpRouter::new()
5318 .server_info("test", "1.0")
5319 .with_dynamic_tools();
5320
5321 let init_req = RouterRequest {
5322 id: RequestId::Number(1),
5323 inner: McpRequest::Initialize(InitializeParams {
5324 protocol_version: "2025-11-25".to_string(),
5325 capabilities: ClientCapabilities::default(),
5326 client_info: Implementation {
5327 name: "test".to_string(),
5328 version: "1.0".to_string(),
5329 ..Default::default()
5330 },
5331 meta: None,
5332 }),
5333 extensions: Extensions::new(),
5334 };
5335
5336 let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5337 match resp.inner {
5338 Ok(McpResponse::Initialize(result)) => {
5339 assert!(result.capabilities.tools.is_some());
5340 }
5341 _ => panic!("Expected Initialize response"),
5342 }
5343 }
5344
5345 #[tokio::test]
5346 async fn test_dynamic_tools_multi_session_notification() {
5347 let (tx1, mut rx1) = crate::context::notification_channel(16);
5348 let (tx2, mut rx2) = crate::context::notification_channel(16);
5349
5350 let (router, registry) = McpRouter::new()
5351 .server_info("test", "1.0")
5352 .with_dynamic_tools();
5353
5354 let _session1 = router.clone().with_notification_sender(tx1);
5356 let _session2 = router.clone().with_notification_sender(tx2);
5357
5358 let tool = ToolBuilder::new("broadcast")
5359 .description("Test")
5360 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5361 .build();
5362
5363 registry.register(tool);
5364
5365 let n1 = rx1.recv().await.unwrap();
5367 let n2 = rx2.recv().await.unwrap();
5368 assert!(matches!(n1, ServerNotification::ToolsListChanged));
5369 assert!(matches!(n2, ServerNotification::ToolsListChanged));
5370 }
5371
5372 #[tokio::test]
5373 async fn test_dynamic_tools_call_not_found() {
5374 let (router, _registry) = McpRouter::new()
5375 .server_info("test", "1.0")
5376 .with_dynamic_tools();
5377
5378 let mut router = router;
5379 init_router(&mut router).await;
5380
5381 let req = RouterRequest {
5382 id: RequestId::Number(1),
5383 inner: McpRequest::CallTool(CallToolParams {
5384 name: "nonexistent".to_string(),
5385 arguments: serde_json::json!({}),
5386 meta: None,
5387 task: None,
5388 }),
5389 extensions: Extensions::new(),
5390 };
5391
5392 let resp = router.ready().await.unwrap().call(req).await.unwrap();
5393 match resp.inner {
5394 Err(e) => {
5395 assert_eq!(e.code, -32601);
5396 }
5397 _ => panic!("Expected method not found error"),
5398 }
5399 }
5400
5401 #[tokio::test]
5402 async fn test_dynamic_tools_registry_list() {
5403 let (_, registry) = McpRouter::new()
5404 .server_info("test", "1.0")
5405 .with_dynamic_tools();
5406
5407 assert!(registry.list().is_empty());
5408
5409 let tool = ToolBuilder::new("tool_a")
5410 .description("A")
5411 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5412 .build();
5413 registry.register(tool);
5414
5415 let tool = ToolBuilder::new("tool_b")
5416 .description("B")
5417 .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5418 .build();
5419 registry.register(tool);
5420
5421 let tools = registry.list();
5422 assert_eq!(tools.len(), 2);
5423 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
5424 assert!(names.contains(&"tool_a"));
5425 assert!(names.contains(&"tool_b"));
5426 }
5427 } }