1use dashmap::DashMap;
59use parking_lot::RwLock;
60use std::collections::HashMap;
61use std::sync::Arc;
62use turbomcp_protocol::types::{Prompt, Resource, Root, Tool};
63
64use crate::handlers::{
65 HandlerMetadata, LoggingHandler, PromptHandler, ResourceHandler, SamplingHandler, ToolHandler,
66};
67use crate::{ServerError, ServerResult};
68
69pub struct HandlerRegistry {
71 pub tools: DashMap<String, Arc<dyn ToolHandler + 'static>>,
73 pub prompts: DashMap<String, Arc<dyn PromptHandler + 'static>>,
75 pub resources: DashMap<String, Arc<dyn ResourceHandler + 'static>>,
77 pub sampling: DashMap<String, Arc<dyn SamplingHandler + 'static>>,
79 pub logging: DashMap<String, Arc<dyn LoggingHandler + 'static>>,
81 pub roots: Arc<RwLock<Vec<Root>>>,
83 metadata: DashMap<String, HandlerMetadata>,
85 config: Arc<RwLock<RegistryConfig>>,
87}
88
89impl std::fmt::Debug for HandlerRegistry {
90 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91 f.debug_struct("HandlerRegistry")
92 .field("tools_count", &self.tools.len())
93 .field("prompts_count", &self.prompts.len())
94 .field("resources_count", &self.resources.len())
95 .field("sampling_count", &self.sampling.len())
96 .field("logging_count", &self.logging.len())
97 .field("roots_count", &self.roots.read().len())
98 .finish()
99 }
100}
101
102#[derive(Debug, Clone)]
104pub struct RegistryConfig {
105 pub max_handlers_per_type: usize,
107 pub enable_metrics: bool,
109 pub enable_validation: bool,
111 pub handler_timeout_ms: u64,
113 pub enable_hot_reload: bool,
115 pub event_listeners: Vec<String>,
117}
118
119impl Default for RegistryConfig {
120 fn default() -> Self {
121 Self {
122 max_handlers_per_type: 1000,
123 enable_metrics: true,
124 enable_validation: true,
125 handler_timeout_ms: 30_000,
126 enable_hot_reload: false,
127 event_listeners: Vec::new(),
128 }
129 }
130}
131
132#[derive(Debug, Clone)]
134pub enum RegistryEvent {
135 HandlerRegistered {
137 handler_type: String,
139 name: String,
141 timestamp: chrono::DateTime<chrono::Utc>,
143 },
144 HandlerUnregistered {
146 handler_type: String,
148 name: String,
150 timestamp: chrono::DateTime<chrono::Utc>,
152 },
153 HandlerUpdated {
155 handler_type: String,
157 name: String,
159 timestamp: chrono::DateTime<chrono::Utc>,
161 },
162 RegistryCleared {
164 timestamp: chrono::DateTime<chrono::Utc>,
166 },
167}
168
169impl HandlerRegistry {
170 #[must_use]
187 pub fn new() -> Self {
188 Self {
189 tools: DashMap::new(),
190 prompts: DashMap::new(),
191 resources: DashMap::new(),
192 sampling: DashMap::new(),
193 logging: DashMap::new(),
194 roots: Arc::new(RwLock::new(Vec::new())),
195 metadata: DashMap::new(),
196 config: Arc::new(RwLock::new(RegistryConfig::default())),
197 }
198 }
199
200 #[must_use]
202 pub fn with_config(config: RegistryConfig) -> Self {
203 Self {
204 tools: DashMap::new(),
205 prompts: DashMap::new(),
206 resources: DashMap::new(),
207 sampling: DashMap::new(),
208 logging: DashMap::new(),
209 roots: Arc::new(RwLock::new(Vec::new())),
210 metadata: DashMap::new(),
211 config: Arc::new(RwLock::new(config)),
212 }
213 }
214
215 pub fn register_tool<T>(&self, name: impl Into<String>, handler: T) -> ServerResult<()>
224 where
225 T: ToolHandler + 'static,
226 {
227 let name = name.into();
228
229 crate::handler_validation::validate_handler_name(&name)?;
231
232 if self.tools.len() >= self.config.read().max_handlers_per_type {
234 return Err(ServerError::handler(format!(
235 "Maximum number of tool handlers ({}) exceeded",
236 self.config.read().max_handlers_per_type
237 )));
238 }
239
240 if self.config.read().enable_validation {
242 self.validate_tool_handler(&handler)?;
243 }
244
245 self.tools.insert(name.clone(), Arc::new(handler));
247
248 let metadata = HandlerMetadata {
250 name: name.clone(),
251 version: "1.0.0".to_string(),
252 description: None,
253 tags: vec!["tool".to_string()],
254 created_at: chrono::Utc::now(),
255 config: HashMap::new(),
256 metrics_enabled: self.config.read().enable_metrics,
257 rate_limit: None,
258 allowed_roles: None,
259 };
260 self.metadata.insert(format!("tool:{name}"), metadata);
261
262 tracing::info!("Registered tool handler: {}", name);
263 Ok(())
264 }
265
266 pub fn register_prompt<P>(&self, name: impl Into<String>, handler: P) -> ServerResult<()>
275 where
276 P: PromptHandler + 'static,
277 {
278 let name = name.into();
279
280 crate::handler_validation::validate_handler_name(&name)?;
282
283 if self.prompts.len() >= self.config.read().max_handlers_per_type {
285 return Err(ServerError::handler(format!(
286 "Maximum number of prompt handlers ({}) exceeded",
287 self.config.read().max_handlers_per_type
288 )));
289 }
290
291 if self.config.read().enable_validation {
293 self.validate_prompt_handler(&handler)?;
294 }
295
296 self.prompts.insert(name.clone(), Arc::new(handler));
298
299 let metadata = HandlerMetadata {
301 name: name.clone(),
302 version: "1.0.0".to_string(),
303 description: None,
304 tags: vec!["prompt".to_string()],
305 created_at: chrono::Utc::now(),
306 config: HashMap::new(),
307 metrics_enabled: self.config.read().enable_metrics,
308 rate_limit: None,
309 allowed_roles: None,
310 };
311 self.metadata.insert(format!("prompt:{name}"), metadata);
312
313 tracing::info!("Registered prompt handler: {}", name);
314 Ok(())
315 }
316
317 pub fn register_resource<R>(&self, name: impl Into<String>, handler: R) -> ServerResult<()>
326 where
327 R: ResourceHandler + 'static,
328 {
329 let name = name.into();
330
331 if self.resources.len() >= self.config.read().max_handlers_per_type {
337 return Err(ServerError::handler(format!(
338 "Maximum number of resource handlers ({}) exceeded",
339 self.config.read().max_handlers_per_type
340 )));
341 }
342
343 if self.config.read().enable_validation {
345 self.validate_resource_handler(&handler)?;
346 }
347
348 self.resources.insert(name.clone(), Arc::new(handler));
350
351 let metadata = HandlerMetadata {
353 name: name.clone(),
354 version: "1.0.0".to_string(),
355 description: None,
356 tags: vec!["resource".to_string()],
357 created_at: chrono::Utc::now(),
358 config: HashMap::new(),
359 metrics_enabled: self.config.read().enable_metrics,
360 rate_limit: None,
361 allowed_roles: None,
362 };
363 self.metadata.insert(format!("resource:{name}"), metadata);
364
365 tracing::info!("Registered resource handler: {}", name);
366 Ok(())
367 }
368
369 pub fn register_sampling<S>(&self, name: impl Into<String>, handler: S) -> ServerResult<()>
375 where
376 S: SamplingHandler + 'static,
377 {
378 let name = name.into();
379
380 crate::handler_validation::validate_handler_name(&name)?;
382
383 if self.sampling.len() >= self.config.read().max_handlers_per_type {
385 return Err(ServerError::handler(format!(
386 "Maximum number of sampling handlers ({}) exceeded",
387 self.config.read().max_handlers_per_type
388 )));
389 }
390
391 self.sampling.insert(name.clone(), Arc::new(handler));
392
393 let metadata = HandlerMetadata {
395 name: name.clone(),
396 version: "1.0.0".to_string(),
397 description: None,
398 tags: vec!["sampling".to_string()],
399 created_at: chrono::Utc::now(),
400 config: HashMap::new(),
401 metrics_enabled: self.config.read().enable_metrics,
402 rate_limit: None,
403 allowed_roles: None,
404 };
405 self.metadata.insert(format!("sampling:{name}"), metadata);
406
407 tracing::info!("Registered sampling handler: {}", name);
408 Ok(())
409 }
410
411 pub fn register_logging<L>(&self, name: impl Into<String>, handler: L) -> ServerResult<()>
417 where
418 L: LoggingHandler + 'static,
419 {
420 let name = name.into();
421
422 crate::handler_validation::validate_handler_name(&name)?;
424
425 if self.logging.len() >= self.config.read().max_handlers_per_type {
427 return Err(ServerError::handler(format!(
428 "Maximum number of logging handlers ({}) exceeded",
429 self.config.read().max_handlers_per_type
430 )));
431 }
432
433 self.logging.insert(name.clone(), Arc::new(handler));
434
435 let metadata = HandlerMetadata {
437 name: name.clone(),
438 version: "1.0.0".to_string(),
439 description: None,
440 tags: vec!["logging".to_string()],
441 created_at: chrono::Utc::now(),
442 config: HashMap::new(),
443 metrics_enabled: self.config.read().enable_metrics,
444 rate_limit: None,
445 allowed_roles: None,
446 };
447 self.metadata.insert(format!("logging:{name}"), metadata);
448
449 tracing::info!("Registered logging handler: {}", name);
450 Ok(())
451 }
452
453 #[must_use]
455 pub fn get_tool(&self, name: &str) -> Option<Arc<dyn ToolHandler>> {
456 self.tools.get(name).map(|entry| Arc::clone(entry.value()))
457 }
458
459 #[must_use]
461 pub fn get_prompt(&self, name: &str) -> Option<Arc<dyn PromptHandler>> {
462 self.prompts
463 .get(name)
464 .map(|entry| Arc::clone(entry.value()))
465 }
466
467 #[must_use]
469 pub fn get_resource(&self, name: &str) -> Option<Arc<dyn ResourceHandler>> {
470 self.resources
471 .get(name)
472 .map(|entry| Arc::clone(entry.value()))
473 }
474
475 #[must_use]
477 pub fn get_sampling(&self, name: &str) -> Option<Arc<dyn SamplingHandler>> {
478 self.sampling
479 .get(name)
480 .map(|entry| Arc::clone(entry.value()))
481 }
482
483 #[must_use]
485 pub fn get_logging(&self, name: &str) -> Option<Arc<dyn LoggingHandler>> {
486 self.logging
487 .get(name)
488 .map(|entry| Arc::clone(entry.value()))
489 }
490
491 #[must_use]
493 pub fn list_tools(&self) -> Vec<String> {
494 self.tools.iter().map(|entry| entry.key().clone()).collect()
495 }
496
497 #[must_use]
499 pub fn list_prompts(&self) -> Vec<String> {
500 self.prompts
501 .iter()
502 .map(|entry| entry.key().clone())
503 .collect()
504 }
505
506 #[must_use]
508 pub fn list_resources(&self) -> Vec<String> {
509 self.resources
510 .iter()
511 .map(|entry| entry.key().clone())
512 .collect()
513 }
514
515 #[must_use]
517 pub fn list_sampling(&self) -> Vec<String> {
518 self.sampling
519 .iter()
520 .map(|entry| entry.key().clone())
521 .collect()
522 }
523
524 #[must_use]
526 pub fn list_logging(&self) -> Vec<String> {
527 self.logging
528 .iter()
529 .map(|entry| entry.key().clone())
530 .collect()
531 }
532
533 #[must_use]
535 pub fn get_tool_definitions(&self) -> Vec<Tool> {
536 self.tools
537 .iter()
538 .map(|entry| entry.value().tool_definition())
539 .collect()
540 }
541
542 #[must_use]
544 pub fn get_prompt_definitions(&self) -> Vec<Prompt> {
545 self.prompts
546 .iter()
547 .map(|entry| entry.value().prompt_definition())
548 .collect()
549 }
550
551 #[must_use]
553 pub fn get_resource_definitions(&self) -> Vec<Resource> {
554 self.resources
555 .iter()
556 .map(|entry| entry.value().resource_definition())
557 .collect()
558 }
559
560 pub fn add_root(&self, root: Root) {
562 self.roots.write().push(root);
563 }
564
565 pub fn set_roots(&self, roots: Vec<Root>) {
567 *self.roots.write() = roots;
568 }
569
570 #[must_use]
572 pub fn get_roots(&self) -> Vec<Root> {
573 self.roots.read().clone()
574 }
575
576 pub fn clear_roots(&self) {
578 self.roots.write().clear();
579 }
580
581 pub fn unregister_tool(&self, name: &str) -> bool {
583 let removed = self.tools.remove(name).is_some();
584 if removed {
585 self.metadata.remove(&format!("tool:{name}"));
586 tracing::info!("Unregistered tool handler: {}", name);
587 }
588 removed
589 }
590
591 pub fn unregister_prompt(&self, name: &str) -> bool {
593 let removed = self.prompts.remove(name).is_some();
594 if removed {
595 self.metadata.remove(&format!("prompt:{name}"));
596 tracing::info!("Unregistered prompt handler: {}", name);
597 }
598 removed
599 }
600
601 pub fn unregister_resource(&self, name: &str) -> bool {
603 let removed = self.resources.remove(name).is_some();
604 if removed {
605 self.metadata.remove(&format!("resource:{name}"));
606 tracing::info!("Unregistered resource handler: {}", name);
607 }
608 removed
609 }
610
611 pub fn clear(&self) {
613 self.tools.clear();
614 self.prompts.clear();
615 self.resources.clear();
616 self.sampling.clear();
617 self.logging.clear();
618 self.metadata.clear();
619 tracing::info!("Cleared all handlers from registry");
620 }
621
622 #[must_use]
624 pub fn stats(&self) -> RegistryStats {
625 RegistryStats {
626 tool_count: self.tools.len(),
627 prompt_count: self.prompts.len(),
628 resource_count: self.resources.len(),
629 sampling_count: self.sampling.len(),
630 logging_count: self.logging.len(),
631 total_count: self.tools.len()
632 + self.prompts.len()
633 + self.resources.len()
634 + self.sampling.len()
635 + self.logging.len(),
636 }
637 }
638
639 #[must_use]
641 pub fn get_metadata(&self, key: &str) -> Option<HandlerMetadata> {
642 self.metadata.get(key).map(|entry| entry.value().clone())
643 }
644
645 pub fn update_config<F>(&self, f: F)
647 where
648 F: FnOnce(&mut RegistryConfig),
649 {
650 let mut config = self.config.write();
651 f(&mut config);
652 }
653
654 fn validate_tool_handler(&self, handler: &dyn ToolHandler) -> ServerResult<()> {
657 let tool_def = handler.tool_definition();
658
659 if tool_def.name.is_empty() {
660 return Err(ServerError::handler("Tool name cannot be empty"));
661 }
662
663 if tool_def.name.len() > 100 {
664 return Err(ServerError::handler(
665 "Tool name too long (max 100 characters)",
666 ));
667 }
668
669 if self.tools.contains_key(&tool_def.name) {
671 return Err(ServerError::handler(format!(
672 "Tool with name '{}' already exists",
673 tool_def.name
674 )));
675 }
676
677 Ok(())
678 }
679
680 fn validate_prompt_handler(&self, handler: &dyn PromptHandler) -> ServerResult<()> {
681 let prompt_def = handler.prompt_definition();
682
683 if prompt_def.name.is_empty() {
684 return Err(ServerError::handler("Prompt name cannot be empty"));
685 }
686
687 if prompt_def.name.len() > 100 {
688 return Err(ServerError::handler(
689 "Prompt name too long (max 100 characters)",
690 ));
691 }
692
693 if self.prompts.contains_key(&prompt_def.name) {
695 return Err(ServerError::handler(format!(
696 "Prompt with name '{}' already exists",
697 prompt_def.name
698 )));
699 }
700
701 Ok(())
702 }
703
704 fn validate_resource_handler(&self, handler: &dyn ResourceHandler) -> ServerResult<()> {
705 let resource_def = handler.resource_definition();
706
707 if resource_def.uri.is_empty() {
708 return Err(ServerError::handler("Resource URI cannot be empty"));
709 }
710
711 if resource_def.name.is_empty() {
712 return Err(ServerError::handler("Resource name cannot be empty"));
713 }
714
715 for entry in &self.resources {
717 if entry.value().resource_definition().uri == resource_def.uri {
718 return Err(ServerError::handler(format!(
719 "Resource with URI '{}' already exists",
720 resource_def.uri
721 )));
722 }
723 }
724
725 Ok(())
726 }
727}
728
729impl Default for HandlerRegistry {
730 fn default() -> Self {
731 Self::new()
732 }
733}
734
735#[derive(Debug, Clone)]
737pub struct RegistryStats {
738 pub tool_count: usize,
740 pub prompt_count: usize,
742 pub resource_count: usize,
744 pub sampling_count: usize,
746 pub logging_count: usize,
748 pub total_count: usize,
750}
751
752#[derive(Debug)]
754pub struct RegistryBuilder {
755 config: RegistryConfig,
756}
757
758impl RegistryBuilder {
759 #[must_use]
761 pub fn new() -> Self {
762 Self {
763 config: RegistryConfig::default(),
764 }
765 }
766
767 #[must_use]
769 pub const fn max_handlers_per_type(mut self, max: usize) -> Self {
770 self.config.max_handlers_per_type = max;
771 self
772 }
773
774 #[must_use]
776 pub const fn enable_metrics(mut self, enable: bool) -> Self {
777 self.config.enable_metrics = enable;
778 self
779 }
780
781 #[must_use]
783 pub const fn enable_validation(mut self, enable: bool) -> Self {
784 self.config.enable_validation = enable;
785 self
786 }
787
788 #[must_use]
790 pub const fn handler_timeout_ms(mut self, timeout: u64) -> Self {
791 self.config.handler_timeout_ms = timeout;
792 self
793 }
794
795 #[must_use]
797 pub const fn enable_hot_reload(mut self, enable: bool) -> Self {
798 self.config.enable_hot_reload = enable;
799 self
800 }
801
802 #[must_use]
804 pub fn build(self) -> HandlerRegistry {
805 HandlerRegistry::with_config(self.config)
806 }
807}
808
809impl Default for RegistryBuilder {
810 fn default() -> Self {
811 Self::new()
812 }
813}
814
815pub type Registry = HandlerRegistry;
817
818#[cfg(test)]
820mod tests;