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 if self.tools.len() >= self.config.read().max_handlers_per_type {
231 return Err(ServerError::handler(format!(
232 "Maximum number of tool handlers ({}) exceeded",
233 self.config.read().max_handlers_per_type
234 )));
235 }
236
237 if self.config.read().enable_validation {
239 self.validate_tool_handler(&handler)?;
240 }
241
242 self.tools.insert(name.clone(), Arc::new(handler));
244
245 let metadata = HandlerMetadata {
247 name: name.clone(),
248 version: "1.0.0".to_string(),
249 description: None,
250 tags: vec!["tool".to_string()],
251 created_at: chrono::Utc::now(),
252 config: HashMap::new(),
253 metrics_enabled: self.config.read().enable_metrics,
254 rate_limit: None,
255 allowed_roles: None,
256 };
257 self.metadata.insert(format!("tool:{name}"), metadata);
258
259 tracing::info!("Registered tool handler: {}", name);
260 Ok(())
261 }
262
263 pub fn register_prompt<P>(&self, name: impl Into<String>, handler: P) -> ServerResult<()>
272 where
273 P: PromptHandler + 'static,
274 {
275 let name = name.into();
276
277 if self.prompts.len() >= self.config.read().max_handlers_per_type {
279 return Err(ServerError::handler(format!(
280 "Maximum number of prompt handlers ({}) exceeded",
281 self.config.read().max_handlers_per_type
282 )));
283 }
284
285 if self.config.read().enable_validation {
287 self.validate_prompt_handler(&handler)?;
288 }
289
290 self.prompts.insert(name.clone(), Arc::new(handler));
292
293 let metadata = HandlerMetadata {
295 name: name.clone(),
296 version: "1.0.0".to_string(),
297 description: None,
298 tags: vec!["prompt".to_string()],
299 created_at: chrono::Utc::now(),
300 config: HashMap::new(),
301 metrics_enabled: self.config.read().enable_metrics,
302 rate_limit: None,
303 allowed_roles: None,
304 };
305 self.metadata.insert(format!("prompt:{name}"), metadata);
306
307 tracing::info!("Registered prompt handler: {}", name);
308 Ok(())
309 }
310
311 pub fn register_resource<R>(&self, name: impl Into<String>, handler: R) -> ServerResult<()>
320 where
321 R: ResourceHandler + 'static,
322 {
323 let name = name.into();
324
325 if self.resources.len() >= self.config.read().max_handlers_per_type {
327 return Err(ServerError::handler(format!(
328 "Maximum number of resource handlers ({}) exceeded",
329 self.config.read().max_handlers_per_type
330 )));
331 }
332
333 if self.config.read().enable_validation {
335 self.validate_resource_handler(&handler)?;
336 }
337
338 self.resources.insert(name.clone(), Arc::new(handler));
340
341 let metadata = HandlerMetadata {
343 name: name.clone(),
344 version: "1.0.0".to_string(),
345 description: None,
346 tags: vec!["resource".to_string()],
347 created_at: chrono::Utc::now(),
348 config: HashMap::new(),
349 metrics_enabled: self.config.read().enable_metrics,
350 rate_limit: None,
351 allowed_roles: None,
352 };
353 self.metadata.insert(format!("resource:{name}"), metadata);
354
355 tracing::info!("Registered resource handler: {}", name);
356 Ok(())
357 }
358
359 pub fn register_sampling<S>(&self, name: impl Into<String>, handler: S) -> ServerResult<()>
365 where
366 S: SamplingHandler + 'static,
367 {
368 let name = name.into();
369
370 if self.sampling.len() >= self.config.read().max_handlers_per_type {
372 return Err(ServerError::handler(format!(
373 "Maximum number of sampling handlers ({}) exceeded",
374 self.config.read().max_handlers_per_type
375 )));
376 }
377
378 self.sampling.insert(name.clone(), Arc::new(handler));
379
380 let metadata = HandlerMetadata {
382 name: name.clone(),
383 version: "1.0.0".to_string(),
384 description: None,
385 tags: vec!["sampling".to_string()],
386 created_at: chrono::Utc::now(),
387 config: HashMap::new(),
388 metrics_enabled: self.config.read().enable_metrics,
389 rate_limit: None,
390 allowed_roles: None,
391 };
392 self.metadata.insert(format!("sampling:{name}"), metadata);
393
394 tracing::info!("Registered sampling handler: {}", name);
395 Ok(())
396 }
397
398 pub fn register_logging<L>(&self, name: impl Into<String>, handler: L) -> ServerResult<()>
404 where
405 L: LoggingHandler + 'static,
406 {
407 let name = name.into();
408
409 if self.logging.len() >= self.config.read().max_handlers_per_type {
411 return Err(ServerError::handler(format!(
412 "Maximum number of logging handlers ({}) exceeded",
413 self.config.read().max_handlers_per_type
414 )));
415 }
416
417 self.logging.insert(name.clone(), Arc::new(handler));
418
419 let metadata = HandlerMetadata {
421 name: name.clone(),
422 version: "1.0.0".to_string(),
423 description: None,
424 tags: vec!["logging".to_string()],
425 created_at: chrono::Utc::now(),
426 config: HashMap::new(),
427 metrics_enabled: self.config.read().enable_metrics,
428 rate_limit: None,
429 allowed_roles: None,
430 };
431 self.metadata.insert(format!("logging:{name}"), metadata);
432
433 tracing::info!("Registered logging handler: {}", name);
434 Ok(())
435 }
436
437 #[must_use]
439 pub fn get_tool(&self, name: &str) -> Option<Arc<dyn ToolHandler>> {
440 self.tools.get(name).map(|entry| Arc::clone(entry.value()))
441 }
442
443 #[must_use]
445 pub fn get_prompt(&self, name: &str) -> Option<Arc<dyn PromptHandler>> {
446 self.prompts
447 .get(name)
448 .map(|entry| Arc::clone(entry.value()))
449 }
450
451 #[must_use]
453 pub fn get_resource(&self, name: &str) -> Option<Arc<dyn ResourceHandler>> {
454 self.resources
455 .get(name)
456 .map(|entry| Arc::clone(entry.value()))
457 }
458
459 #[must_use]
461 pub fn get_sampling(&self, name: &str) -> Option<Arc<dyn SamplingHandler>> {
462 self.sampling
463 .get(name)
464 .map(|entry| Arc::clone(entry.value()))
465 }
466
467 #[must_use]
469 pub fn get_logging(&self, name: &str) -> Option<Arc<dyn LoggingHandler>> {
470 self.logging
471 .get(name)
472 .map(|entry| Arc::clone(entry.value()))
473 }
474
475 #[must_use]
477 pub fn list_tools(&self) -> Vec<String> {
478 self.tools.iter().map(|entry| entry.key().clone()).collect()
479 }
480
481 #[must_use]
483 pub fn list_prompts(&self) -> Vec<String> {
484 self.prompts
485 .iter()
486 .map(|entry| entry.key().clone())
487 .collect()
488 }
489
490 #[must_use]
492 pub fn list_resources(&self) -> Vec<String> {
493 self.resources
494 .iter()
495 .map(|entry| entry.key().clone())
496 .collect()
497 }
498
499 #[must_use]
501 pub fn list_sampling(&self) -> Vec<String> {
502 self.sampling
503 .iter()
504 .map(|entry| entry.key().clone())
505 .collect()
506 }
507
508 #[must_use]
510 pub fn list_logging(&self) -> Vec<String> {
511 self.logging
512 .iter()
513 .map(|entry| entry.key().clone())
514 .collect()
515 }
516
517 #[must_use]
519 pub fn get_tool_definitions(&self) -> Vec<Tool> {
520 self.tools
521 .iter()
522 .map(|entry| entry.value().tool_definition())
523 .collect()
524 }
525
526 #[must_use]
528 pub fn get_prompt_definitions(&self) -> Vec<Prompt> {
529 self.prompts
530 .iter()
531 .map(|entry| entry.value().prompt_definition())
532 .collect()
533 }
534
535 #[must_use]
537 pub fn get_resource_definitions(&self) -> Vec<Resource> {
538 self.resources
539 .iter()
540 .map(|entry| entry.value().resource_definition())
541 .collect()
542 }
543
544 pub fn add_root(&self, root: Root) {
546 self.roots.write().push(root);
547 }
548
549 pub fn set_roots(&self, roots: Vec<Root>) {
551 *self.roots.write() = roots;
552 }
553
554 #[must_use]
556 pub fn get_roots(&self) -> Vec<Root> {
557 self.roots.read().clone()
558 }
559
560 pub fn clear_roots(&self) {
562 self.roots.write().clear();
563 }
564
565 pub fn unregister_tool(&self, name: &str) -> bool {
567 let removed = self.tools.remove(name).is_some();
568 if removed {
569 self.metadata.remove(&format!("tool:{name}"));
570 tracing::info!("Unregistered tool handler: {}", name);
571 }
572 removed
573 }
574
575 pub fn unregister_prompt(&self, name: &str) -> bool {
577 let removed = self.prompts.remove(name).is_some();
578 if removed {
579 self.metadata.remove(&format!("prompt:{name}"));
580 tracing::info!("Unregistered prompt handler: {}", name);
581 }
582 removed
583 }
584
585 pub fn unregister_resource(&self, name: &str) -> bool {
587 let removed = self.resources.remove(name).is_some();
588 if removed {
589 self.metadata.remove(&format!("resource:{name}"));
590 tracing::info!("Unregistered resource handler: {}", name);
591 }
592 removed
593 }
594
595 pub fn clear(&self) {
597 self.tools.clear();
598 self.prompts.clear();
599 self.resources.clear();
600 self.sampling.clear();
601 self.logging.clear();
602 self.metadata.clear();
603 tracing::info!("Cleared all handlers from registry");
604 }
605
606 #[must_use]
608 pub fn stats(&self) -> RegistryStats {
609 RegistryStats {
610 tool_count: self.tools.len(),
611 prompt_count: self.prompts.len(),
612 resource_count: self.resources.len(),
613 sampling_count: self.sampling.len(),
614 logging_count: self.logging.len(),
615 total_count: self.tools.len()
616 + self.prompts.len()
617 + self.resources.len()
618 + self.sampling.len()
619 + self.logging.len(),
620 }
621 }
622
623 #[must_use]
625 pub fn get_metadata(&self, key: &str) -> Option<HandlerMetadata> {
626 self.metadata.get(key).map(|entry| entry.value().clone())
627 }
628
629 pub fn update_config<F>(&self, f: F)
631 where
632 F: FnOnce(&mut RegistryConfig),
633 {
634 let mut config = self.config.write();
635 f(&mut config);
636 }
637
638 fn validate_tool_handler(&self, handler: &dyn ToolHandler) -> ServerResult<()> {
641 let tool_def = handler.tool_definition();
642
643 if tool_def.name.is_empty() {
644 return Err(ServerError::handler("Tool name cannot be empty"));
645 }
646
647 if tool_def.name.len() > 100 {
648 return Err(ServerError::handler(
649 "Tool name too long (max 100 characters)",
650 ));
651 }
652
653 if self.tools.contains_key(&tool_def.name) {
655 return Err(ServerError::handler(format!(
656 "Tool with name '{}' already exists",
657 tool_def.name
658 )));
659 }
660
661 Ok(())
662 }
663
664 fn validate_prompt_handler(&self, handler: &dyn PromptHandler) -> ServerResult<()> {
665 let prompt_def = handler.prompt_definition();
666
667 if prompt_def.name.is_empty() {
668 return Err(ServerError::handler("Prompt name cannot be empty"));
669 }
670
671 if prompt_def.name.len() > 100 {
672 return Err(ServerError::handler(
673 "Prompt name too long (max 100 characters)",
674 ));
675 }
676
677 if self.prompts.contains_key(&prompt_def.name) {
679 return Err(ServerError::handler(format!(
680 "Prompt with name '{}' already exists",
681 prompt_def.name
682 )));
683 }
684
685 Ok(())
686 }
687
688 fn validate_resource_handler(&self, handler: &dyn ResourceHandler) -> ServerResult<()> {
689 let resource_def = handler.resource_definition();
690
691 if resource_def.uri.is_empty() {
692 return Err(ServerError::handler("Resource URI cannot be empty"));
693 }
694
695 if resource_def.name.is_empty() {
696 return Err(ServerError::handler("Resource name cannot be empty"));
697 }
698
699 for entry in &self.resources {
701 if entry.value().resource_definition().uri == resource_def.uri {
702 return Err(ServerError::handler(format!(
703 "Resource with URI '{}' already exists",
704 resource_def.uri
705 )));
706 }
707 }
708
709 Ok(())
710 }
711}
712
713impl Default for HandlerRegistry {
714 fn default() -> Self {
715 Self::new()
716 }
717}
718
719#[derive(Debug, Clone)]
721pub struct RegistryStats {
722 pub tool_count: usize,
724 pub prompt_count: usize,
726 pub resource_count: usize,
728 pub sampling_count: usize,
730 pub logging_count: usize,
732 pub total_count: usize,
734}
735
736#[derive(Debug)]
738pub struct RegistryBuilder {
739 config: RegistryConfig,
740}
741
742impl RegistryBuilder {
743 #[must_use]
745 pub fn new() -> Self {
746 Self {
747 config: RegistryConfig::default(),
748 }
749 }
750
751 #[must_use]
753 pub const fn max_handlers_per_type(mut self, max: usize) -> Self {
754 self.config.max_handlers_per_type = max;
755 self
756 }
757
758 #[must_use]
760 pub const fn enable_metrics(mut self, enable: bool) -> Self {
761 self.config.enable_metrics = enable;
762 self
763 }
764
765 #[must_use]
767 pub const fn enable_validation(mut self, enable: bool) -> Self {
768 self.config.enable_validation = enable;
769 self
770 }
771
772 #[must_use]
774 pub const fn handler_timeout_ms(mut self, timeout: u64) -> Self {
775 self.config.handler_timeout_ms = timeout;
776 self
777 }
778
779 #[must_use]
781 pub const fn enable_hot_reload(mut self, enable: bool) -> Self {
782 self.config.enable_hot_reload = enable;
783 self
784 }
785
786 #[must_use]
788 pub fn build(self) -> HandlerRegistry {
789 HandlerRegistry::with_config(self.config)
790 }
791}
792
793impl Default for RegistryBuilder {
794 fn default() -> Self {
795 Self::new()
796 }
797}
798
799pub type Registry = HandlerRegistry;
801
802#[cfg(test)]
804mod tests;