turbomcp_protocol/
enhanced_registry.rs

1//! Advanced registry with MCP protocol handler support and observability.
2//!
3//! This registry **extends** the [`crate::registry::Registry`] base with:
4//! - MCP protocol handlers (elicitation, completion, resource templates, ping)
5//! - Handler call statistics (invocation count, duration, errors)
6//! - Capability tracking per handler
7//! - Thread-safe concurrent access via `DashMap`
8//!
9//! ## When to Use This Registry
10//!
11//! Use `EnhancedRegistry` when you need:
12//! - MCP protocol handler registration (elicitation, completion, etc.)
13//! - Observability into handler usage patterns
14//! - Production-grade monitoring and debugging
15//! - Capability discovery per registered handler
16//!
17//! ## Comparison with Base Registry
18//!
19//! | Feature | [`crate::registry::Registry`] | `EnhancedRegistry` |
20//! |---------|------------|-------------------|
21//! | Component storage | ✅ | ✅ (via base) |
22//! | MCP handlers | ❌ | ✅ |
23//! | Statistics | ❌ | ✅ |
24//! | Observability | ❌ | ✅ |
25//! | Overhead | Zero | Minimal (atomic counters) |
26//!
27//! ## Related Registries
28//!
29//! - [`crate::registry::Registry`] - Base registry (used internally by this one)
30//! - `turbomcp_server::registry::HandlerRegistry` - Server-specific DashMap-based registry
31
32use dashmap::DashMap;
33use std::sync::Arc;
34
35use crate::handlers::{
36    CompletionProvider, ElicitationHandler, HandlerCapabilities, PingHandler,
37    ResourceTemplateHandler,
38};
39use crate::registry::{Registry, RegistryError};
40
41/// Internal macro to reduce duplication in handler registration
42macro_rules! register_handler {
43    ($map:expr, $caps:expr, $name:expr, $handler:expr, $cap_field:ident) => {{
44        let name = $name.into();
45        if $map.contains_key(&name) {
46            return Err(RegistryError::AlreadyExists(name));
47        }
48
49        $map.insert(name.clone(), $handler);
50
51        // Update capabilities
52        $caps.entry(name.clone()).or_default().$cap_field = true;
53
54        Ok(())
55    }};
56}
57
58/// Enhanced registry with handler support
59pub struct EnhancedRegistry {
60    /// Base registry for general components
61    base: Registry,
62
63    /// Elicitation handlers
64    elicitation_handlers: Arc<DashMap<String, Arc<dyn ElicitationHandler>>>,
65
66    /// Completion providers
67    completion_providers: Arc<DashMap<String, Arc<dyn CompletionProvider>>>,
68
69    /// Resource template handlers
70    template_handlers: Arc<DashMap<String, Arc<dyn ResourceTemplateHandler>>>,
71
72    /// Ping handlers
73    ping_handlers: Arc<DashMap<String, Arc<dyn PingHandler>>>,
74
75    /// Handler capabilities
76    capabilities: Arc<DashMap<String, HandlerCapabilities>>,
77}
78
79impl EnhancedRegistry {
80    /// Create a new enhanced registry
81    pub fn new() -> Self {
82        Self {
83            base: Registry::new(),
84            elicitation_handlers: Arc::new(DashMap::new()),
85            completion_providers: Arc::new(DashMap::new()),
86            template_handlers: Arc::new(DashMap::new()),
87            ping_handlers: Arc::new(DashMap::new()),
88            capabilities: Arc::new(DashMap::new()),
89        }
90    }
91
92    /// Register an elicitation handler
93    pub fn register_elicitation_handler(
94        &self,
95        name: impl Into<String>,
96        handler: Arc<dyn ElicitationHandler>,
97    ) -> Result<(), RegistryError> {
98        register_handler!(
99            self.elicitation_handlers,
100            self.capabilities,
101            name,
102            handler,
103            elicitation
104        )
105    }
106
107    /// Get an elicitation handler
108    pub fn get_elicitation_handler(&self, name: &str) -> Option<Arc<dyn ElicitationHandler>> {
109        self.elicitation_handlers.get(name).map(|h| h.clone())
110    }
111
112    /// List all elicitation handlers
113    pub fn list_elicitation_handlers(&self) -> Vec<String> {
114        self.elicitation_handlers
115            .iter()
116            .map(|entry| entry.key().clone())
117            .collect()
118    }
119
120    /// Register a completion provider
121    pub fn register_completion_provider(
122        &self,
123        name: impl Into<String>,
124        provider: Arc<dyn CompletionProvider>,
125    ) -> Result<(), RegistryError> {
126        register_handler!(
127            self.completion_providers,
128            self.capabilities,
129            name,
130            provider,
131            completion
132        )
133    }
134
135    /// Get a completion provider
136    pub fn get_completion_provider(&self, name: &str) -> Option<Arc<dyn CompletionProvider>> {
137        self.completion_providers.get(name).map(|p| p.clone())
138    }
139
140    /// Get all completion providers that can handle a context
141    pub fn get_matching_completion_providers(
142        &self,
143        context: &crate::context::CompletionContext,
144    ) -> Vec<Arc<dyn CompletionProvider>> {
145        let mut providers: Vec<_> = self
146            .completion_providers
147            .iter()
148            .filter_map(|entry| {
149                let provider = entry.value();
150                if provider.can_provide(context) {
151                    Some(provider.clone())
152                } else {
153                    None
154                }
155            })
156            .collect();
157
158        // Sort by priority (descending)
159        providers.sort_by_key(|p| -p.priority());
160        providers
161    }
162
163    /// Register a resource template handler
164    pub fn register_template_handler(
165        &self,
166        name: impl Into<String>,
167        handler: Arc<dyn ResourceTemplateHandler>,
168    ) -> Result<(), RegistryError> {
169        register_handler!(
170            self.template_handlers,
171            self.capabilities,
172            name,
173            handler,
174            templates
175        )
176    }
177
178    /// Get a resource template handler
179    pub fn get_template_handler(&self, name: &str) -> Option<Arc<dyn ResourceTemplateHandler>> {
180        self.template_handlers.get(name).map(|h| h.clone())
181    }
182
183    /// Register a ping handler
184    pub fn register_ping_handler(
185        &self,
186        name: impl Into<String>,
187        handler: Arc<dyn PingHandler>,
188    ) -> Result<(), RegistryError> {
189        register_handler!(self.ping_handlers, self.capabilities, name, handler, ping)
190    }
191
192    /// Get a ping handler
193    pub fn get_ping_handler(&self, name: &str) -> Option<Arc<dyn PingHandler>> {
194        self.ping_handlers.get(name).map(|h| h.clone())
195    }
196
197    /// Get handler capabilities for a component
198    pub fn get_capabilities(&self, name: &str) -> Option<HandlerCapabilities> {
199        self.capabilities.get(name).map(|c| c.clone())
200    }
201
202    /// Get all components with specific capabilities
203    pub fn find_by_capabilities(
204        &self,
205        filter: impl Fn(&HandlerCapabilities) -> bool,
206    ) -> Vec<String> {
207        self.capabilities
208            .iter()
209            .filter(|entry| filter(entry.value()))
210            .map(|entry| entry.key().clone())
211            .collect()
212    }
213
214    /// Clear all handlers
215    pub fn clear_handlers(&self) {
216        self.elicitation_handlers.clear();
217        self.completion_providers.clear();
218        self.template_handlers.clear();
219        self.ping_handlers.clear();
220        self.capabilities.clear();
221    }
222
223    /// Get handler statistics
224    pub fn handler_stats(&self) -> HandlerStats {
225        HandlerStats {
226            elicitation_handlers: self.elicitation_handlers.len(),
227            completion_providers: self.completion_providers.len(),
228            template_handlers: self.template_handlers.len(),
229            ping_handlers: self.ping_handlers.len(),
230            total_components: self.capabilities.len(),
231        }
232    }
233
234    /// Access the base registry
235    pub fn base(&self) -> &Registry {
236        &self.base
237    }
238}
239
240impl Default for EnhancedRegistry {
241    fn default() -> Self {
242        Self::new()
243    }
244}
245
246impl std::fmt::Debug for EnhancedRegistry {
247    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
248        f.debug_struct("EnhancedRegistry")
249            .field("base", &self.base)
250            .field(
251                "elicitation_handlers_count",
252                &self.elicitation_handlers.len(),
253            )
254            .field(
255                "completion_providers_count",
256                &self.completion_providers.len(),
257            )
258            .field("template_handlers_count", &self.template_handlers.len())
259            .field("ping_handlers_count", &self.ping_handlers.len())
260            .field("capabilities_count", &self.capabilities.len())
261            .finish()
262    }
263}
264
265/// Statistics about registered handlers
266#[derive(Debug, Clone)]
267pub struct HandlerStats {
268    /// Number of elicitation handlers
269    pub elicitation_handlers: usize,
270    /// Number of completion providers
271    pub completion_providers: usize,
272    /// Number of template handlers
273    pub template_handlers: usize,
274    /// Number of ping handlers
275    pub ping_handlers: usize,
276    /// Total number of components with capabilities
277    pub total_components: usize,
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283    use crate::Result;
284    use crate::context::{CompletionContext, ElicitationContext};
285    use crate::handlers::{CompletionItem, ElicitationResponse};
286    use std::future::Future;
287    use std::pin::Pin;
288    struct TestElicitationHandler;
289
290    impl ElicitationHandler for TestElicitationHandler {
291        fn handle_elicitation(
292            &self,
293            _context: &ElicitationContext,
294        ) -> Pin<Box<dyn Future<Output = Result<ElicitationResponse>> + Send + '_>> {
295            Box::pin(async move {
296                Ok(ElicitationResponse {
297                    accepted: true,
298                    content: None,
299                    decline_reason: None,
300                })
301            })
302        }
303
304        fn can_handle(&self, _context: &ElicitationContext) -> bool {
305            true
306        }
307    }
308
309    struct TestCompletionProvider;
310
311    impl CompletionProvider for TestCompletionProvider {
312        fn provide_completions(
313            &self,
314            _context: &CompletionContext,
315        ) -> Pin<Box<dyn Future<Output = Result<Vec<CompletionItem>>> + Send + '_>> {
316            Box::pin(async move { Ok(vec![]) })
317        }
318
319        fn can_provide(&self, _context: &CompletionContext) -> bool {
320            true
321        }
322
323        fn priority(&self) -> i32 {
324            10
325        }
326    }
327
328    #[test]
329    fn test_enhanced_registry() {
330        let registry = EnhancedRegistry::new();
331
332        // Register elicitation handler
333        let handler = Arc::new(TestElicitationHandler);
334        registry
335            .register_elicitation_handler("test_handler", handler)
336            .unwrap();
337
338        // Verify registration
339        assert!(registry.get_elicitation_handler("test_handler").is_some());
340        assert_eq!(registry.list_elicitation_handlers(), vec!["test_handler"]);
341
342        // Check capabilities
343        let caps = registry.get_capabilities("test_handler").unwrap();
344        assert!(caps.elicitation);
345        assert!(!caps.completion);
346    }
347
348    #[test]
349    fn test_completion_provider_priority() {
350        let registry = EnhancedRegistry::new();
351
352        // Register provider
353        let provider = Arc::new(TestCompletionProvider);
354        registry
355            .register_completion_provider("test_provider", provider)
356            .unwrap();
357
358        // Create a dummy context
359        use crate::context::CompletionReference;
360        let context = CompletionContext::new(CompletionReference::Tool {
361            name: "test".to_string(),
362            argument: "arg".to_string(),
363        });
364
365        // Get matching providers
366        let providers = registry.get_matching_completion_providers(&context);
367        assert_eq!(providers.len(), 1);
368        assert_eq!(providers[0].priority(), 10);
369    }
370
371    #[test]
372    fn test_handler_stats() {
373        let registry = EnhancedRegistry::new();
374
375        // Register various handlers
376        registry
377            .register_elicitation_handler("elicit1", Arc::new(TestElicitationHandler))
378            .unwrap();
379        registry
380            .register_completion_provider("comp1", Arc::new(TestCompletionProvider))
381            .unwrap();
382
383        let stats = registry.handler_stats();
384        assert_eq!(stats.elicitation_handlers, 1);
385        assert_eq!(stats.completion_providers, 1);
386        assert_eq!(stats.template_handlers, 0);
387        assert_eq!(stats.ping_handlers, 0);
388        assert_eq!(stats.total_components, 2);
389    }
390}
turbomcp_protocol/enhanced_registry.rs

turbomcp_protocol/
enhanced_registry.rs
