tower_mcp/
registry.rs

1//! Dynamic tool registry for runtime tool (de)registration.
2//!
3//! The [`DynamicToolRegistry`] provides a thread-safe, cloneable handle for
4//! adding and removing tools at runtime. When tools change, all connected
5//! sessions are notified via `notifications/tools/list_changed`.
6//!
7//! # Example
8//!
9//! ```rust
10//! use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
11//! use schemars::JsonSchema;
12//! use serde::Deserialize;
13//!
14//! #[derive(Debug, Deserialize, JsonSchema)]
15//! struct Input { value: String }
16//!
17//! let (router, registry) = McpRouter::new()
18//!     .server_info("my-server", "1.0.0")
19//!     .with_dynamic_tools();
20//!
21//! // Register a tool at runtime
22//! let tool = ToolBuilder::new("echo")
23//!     .description("Echo input")
24//!     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
25//!     .build();
26//!
27//! registry.register(tool);
28//! ```
29
30use std::collections::HashMap;
31use std::sync::{Arc, RwLock};
32
33use crate::context::{NotificationSender, ServerNotification};
34use crate::tool::Tool;
35
36/// Inner state shared between the registry handle and the router.
37pub(crate) struct DynamicToolsInner {
38    tools: RwLock<HashMap<String, Arc<Tool>>>,
39    notification_senders: RwLock<Vec<NotificationSender>>,
40}
41
42impl DynamicToolsInner {
43    pub(crate) fn new() -> Self {
44        Self {
45            tools: RwLock::new(HashMap::new()),
46            notification_senders: RwLock::new(Vec::new()),
47        }
48    }
49
50    /// Register a notification sender for a new session.
51    pub(crate) fn add_notification_sender(&self, sender: NotificationSender) {
52        let mut senders = self.notification_senders.write().unwrap();
53        senders.push(sender);
54    }
55
56    /// Broadcast `ToolsListChanged` to all sessions, lazily cleaning up closed channels.
57    fn broadcast_tools_changed(&self) {
58        let mut senders = self.notification_senders.write().unwrap();
59        senders.retain(|tx| !tx.is_closed());
60        for tx in senders.iter() {
61            let _ = tx.try_send(ServerNotification::ToolsListChanged);
62        }
63    }
64
65    /// Get a snapshot of all dynamic tools.
66    pub(crate) fn list(&self) -> Vec<Arc<Tool>> {
67        let tools = self.tools.read().unwrap();
68        tools.values().cloned().collect()
69    }
70
71    /// Look up a dynamic tool by name.
72    pub(crate) fn get(&self, name: &str) -> Option<Arc<Tool>> {
73        let tools = self.tools.read().unwrap();
74        tools.get(name).cloned()
75    }
76
77    /// Check if a dynamic tool exists.
78    pub(crate) fn contains(&self, name: &str) -> bool {
79        let tools = self.tools.read().unwrap();
80        tools.contains_key(name)
81    }
82}
83
84/// A thread-safe, cloneable handle for runtime tool management.
85///
86/// Obtained from [`McpRouter::with_dynamic_tools()`](crate::McpRouter::with_dynamic_tools).
87/// Tools registered here are merged with the router's static tools when
88/// handling `tools/list` and `tools/call` requests.
89///
90/// When a tool is registered or unregistered, all connected sessions receive a
91/// `notifications/tools/list_changed` notification.
92#[derive(Clone)]
93pub struct DynamicToolRegistry {
94    inner: Arc<DynamicToolsInner>,
95}
96
97impl DynamicToolRegistry {
98    pub(crate) fn new(inner: Arc<DynamicToolsInner>) -> Self {
99        Self { inner }
100    }
101
102    /// Register a tool, replacing any existing tool with the same name.
103    ///
104    /// Broadcasts `ToolsListChanged` to all connected sessions.
105    pub fn register(&self, tool: Tool) {
106        {
107            let mut tools = self.inner.tools.write().unwrap();
108            tools.insert(tool.name.clone(), Arc::new(tool));
109        }
110        self.inner.broadcast_tools_changed();
111    }
112
113    /// Unregister a tool by name.
114    ///
115    /// Returns `true` if the tool existed and was removed.
116    /// Broadcasts `ToolsListChanged` only if the tool was actually removed.
117    pub fn unregister(&self, name: &str) -> bool {
118        let removed = {
119            let mut tools = self.inner.tools.write().unwrap();
120            tools.remove(name).is_some()
121        };
122        if removed {
123            self.inner.broadcast_tools_changed();
124        }
125        removed
126    }
127
128    /// List all currently registered dynamic tools.
129    pub fn list(&self) -> Vec<Arc<Tool>> {
130        self.inner.list()
131    }
132
133    /// Check if a tool with the given name is registered.
134    pub fn contains(&self, name: &str) -> bool {
135        self.inner.contains(name)
136    }
137}
138
139// =============================================================================
140// Dynamic Prompt Registry
141// =============================================================================
142
143/// Inner state shared between the prompt registry handle and the router.
144pub(crate) struct DynamicPromptsInner {
145    prompts: RwLock<HashMap<String, Arc<crate::prompt::Prompt>>>,
146    notification_senders: RwLock<Vec<NotificationSender>>,
147}
148
149impl DynamicPromptsInner {
150    pub(crate) fn new() -> Self {
151        Self {
152            prompts: RwLock::new(HashMap::new()),
153            notification_senders: RwLock::new(Vec::new()),
154        }
155    }
156
157    /// Register a notification sender for a new session.
158    pub(crate) fn add_notification_sender(&self, sender: NotificationSender) {
159        let mut senders = self.notification_senders.write().unwrap();
160        senders.push(sender);
161    }
162
163    /// Broadcast `PromptsListChanged` to all sessions, lazily cleaning up closed channels.
164    fn broadcast_prompts_changed(&self) {
165        let mut senders = self.notification_senders.write().unwrap();
166        senders.retain(|tx| !tx.is_closed());
167        for tx in senders.iter() {
168            let _ = tx.try_send(ServerNotification::PromptsListChanged);
169        }
170    }
171
172    /// Get a snapshot of all dynamic prompts.
173    pub(crate) fn list(&self) -> Vec<Arc<crate::prompt::Prompt>> {
174        let prompts = self.prompts.read().unwrap();
175        prompts.values().cloned().collect()
176    }
177
178    /// Look up a dynamic prompt by name.
179    pub(crate) fn get(&self, name: &str) -> Option<Arc<crate::prompt::Prompt>> {
180        let prompts = self.prompts.read().unwrap();
181        prompts.get(name).cloned()
182    }
183
184    /// Check if a dynamic prompt exists.
185    pub(crate) fn contains(&self, name: &str) -> bool {
186        let prompts = self.prompts.read().unwrap();
187        prompts.contains_key(name)
188    }
189}
190
191/// A thread-safe, cloneable handle for runtime prompt management.
192///
193/// Obtained from [`McpRouter::with_dynamic_prompts()`](crate::McpRouter::with_dynamic_prompts).
194/// Prompts registered here are merged with the router's static prompts when
195/// handling `prompts/list` and `prompts/get` requests.
196///
197/// When a prompt is registered or unregistered, all connected sessions receive a
198/// `notifications/prompts/list_changed` notification.
199///
200/// # Example
201///
202/// ```rust
203/// use tower_mcp::{McpRouter, PromptBuilder};
204///
205/// let (router, registry) = McpRouter::new()
206///     .server_info("my-server", "1.0.0")
207///     .with_dynamic_prompts();
208///
209/// // Register a prompt at runtime
210/// let prompt = PromptBuilder::new("greet")
211///     .description("Greet someone")
212///     .user_message("Hello!");
213///
214/// registry.register(prompt);
215/// ```
216#[derive(Clone)]
217pub struct DynamicPromptRegistry {
218    inner: Arc<DynamicPromptsInner>,
219}
220
221impl DynamicPromptRegistry {
222    pub(crate) fn new(inner: Arc<DynamicPromptsInner>) -> Self {
223        Self { inner }
224    }
225
226    /// Register a prompt, replacing any existing prompt with the same name.
227    ///
228    /// Broadcasts `PromptsListChanged` to all connected sessions.
229    pub fn register(&self, prompt: crate::prompt::Prompt) {
230        {
231            let mut prompts = self.inner.prompts.write().unwrap();
232            prompts.insert(prompt.name.clone(), Arc::new(prompt));
233        }
234        self.inner.broadcast_prompts_changed();
235    }
236
237    /// Unregister a prompt by name.
238    ///
239    /// Returns `true` if the prompt existed and was removed.
240    /// Broadcasts `PromptsListChanged` only if the prompt was actually removed.
241    pub fn unregister(&self, name: &str) -> bool {
242        let removed = {
243            let mut prompts = self.inner.prompts.write().unwrap();
244            prompts.remove(name).is_some()
245        };
246        if removed {
247            self.inner.broadcast_prompts_changed();
248        }
249        removed
250    }
251
252    /// List all currently registered dynamic prompts.
253    pub fn list(&self) -> Vec<Arc<crate::prompt::Prompt>> {
254        self.inner.list()
255    }
256
257    /// Check if a prompt with the given name is registered.
258    pub fn contains(&self, name: &str) -> bool {
259        self.inner.contains(name)
260    }
261}
262
263// =============================================================================
264// Dynamic Resource Registry
265// =============================================================================
266
267/// Inner state shared between the resource registry handle and the router.
268pub(crate) struct DynamicResourcesInner {
269    resources: RwLock<HashMap<String, Arc<crate::resource::Resource>>>,
270    notification_senders: RwLock<Vec<NotificationSender>>,
271}
272
273impl DynamicResourcesInner {
274    pub(crate) fn new() -> Self {
275        Self {
276            resources: RwLock::new(HashMap::new()),
277            notification_senders: RwLock::new(Vec::new()),
278        }
279    }
280
281    pub(crate) fn add_notification_sender(&self, sender: NotificationSender) {
282        let mut senders = self.notification_senders.write().unwrap();
283        senders.push(sender);
284    }
285
286    fn broadcast_resources_changed(&self) {
287        let mut senders = self.notification_senders.write().unwrap();
288        senders.retain(|tx| !tx.is_closed());
289        for tx in senders.iter() {
290            let _ = tx.try_send(ServerNotification::ResourcesListChanged);
291        }
292    }
293
294    pub(crate) fn list(&self) -> Vec<Arc<crate::resource::Resource>> {
295        let resources = self.resources.read().unwrap();
296        resources.values().cloned().collect()
297    }
298
299    pub(crate) fn get(&self, uri: &str) -> Option<Arc<crate::resource::Resource>> {
300        let resources = self.resources.read().unwrap();
301        resources.get(uri).cloned()
302    }
303}
304
305/// A thread-safe, cloneable handle for runtime resource management.
306///
307/// Obtained from [`McpRouter::with_dynamic_resources()`](crate::McpRouter::with_dynamic_resources).
308/// Resources registered here are merged with the router's static resources
309/// when handling `resources/list` and `resources/read` requests. Static
310/// resources take precedence over dynamic resources when URIs collide.
311///
312/// When a resource is registered or unregistered, all connected sessions
313/// receive a `notifications/resources/list_changed` notification.
314///
315/// # Example
316///
317/// ```rust
318/// use tower_mcp::{McpRouter, ResourceBuilder};
319///
320/// let (router, registry) = McpRouter::new()
321///     .server_info("my-server", "1.0.0")
322///     .with_dynamic_resources();
323///
324/// let resource = ResourceBuilder::new("file:///data.json")
325///     .name("Data")
326///     .text(r#"{"key": "value"}"#);
327///
328/// registry.register(resource);
329/// ```
330#[derive(Clone)]
331pub struct DynamicResourceRegistry {
332    inner: Arc<DynamicResourcesInner>,
333}
334
335impl DynamicResourceRegistry {
336    pub(crate) fn new(inner: Arc<DynamicResourcesInner>) -> Self {
337        Self { inner }
338    }
339
340    /// Register a resource, replacing any existing resource with the same URI.
341    ///
342    /// Broadcasts `ResourcesListChanged` to all connected sessions.
343    pub fn register(&self, resource: crate::resource::Resource) {
344        {
345            let mut resources = self.inner.resources.write().unwrap();
346            resources.insert(resource.uri.clone(), Arc::new(resource));
347        }
348        self.inner.broadcast_resources_changed();
349    }
350
351    /// Unregister a resource by URI.
352    ///
353    /// Returns `true` if the resource existed and was removed.
354    /// Broadcasts `ResourcesListChanged` only if the resource was actually removed.
355    pub fn unregister(&self, uri: &str) -> bool {
356        let removed = {
357            let mut resources = self.inner.resources.write().unwrap();
358            resources.remove(uri).is_some()
359        };
360        if removed {
361            self.inner.broadcast_resources_changed();
362        }
363        removed
364    }
365
366    /// List all currently registered dynamic resources.
367    pub fn list(&self) -> Vec<Arc<crate::resource::Resource>> {
368        self.inner.list()
369    }
370
371    /// Check if a resource with the given URI is registered.
372    pub fn contains(&self, uri: &str) -> bool {
373        let resources = self.inner.resources.read().unwrap();
374        resources.contains_key(uri)
375    }
376}
377
378// =============================================================================
379// Dynamic Resource Template Registry
380// =============================================================================
381
382/// Inner state shared between the resource template registry handle and the router.
383pub(crate) struct DynamicResourceTemplatesInner {
384    templates: RwLock<Vec<Arc<crate::resource::ResourceTemplate>>>,
385    notification_senders: RwLock<Vec<NotificationSender>>,
386}
387
388impl DynamicResourceTemplatesInner {
389    pub(crate) fn new() -> Self {
390        Self {
391            templates: RwLock::new(Vec::new()),
392            notification_senders: RwLock::new(Vec::new()),
393        }
394    }
395
396    pub(crate) fn add_notification_sender(&self, sender: NotificationSender) {
397        let mut senders = self.notification_senders.write().unwrap();
398        senders.push(sender);
399    }
400
401    fn broadcast_resources_changed(&self) {
402        let mut senders = self.notification_senders.write().unwrap();
403        senders.retain(|tx| !tx.is_closed());
404        for tx in senders.iter() {
405            let _ = tx.try_send(ServerNotification::ResourcesListChanged);
406        }
407    }
408
409    pub(crate) fn list(&self) -> Vec<Arc<crate::resource::ResourceTemplate>> {
410        let templates = self.templates.read().unwrap();
411        templates.clone()
412    }
413
414    pub(crate) fn match_uri(
415        &self,
416        uri: &str,
417    ) -> Option<(
418        Arc<crate::resource::ResourceTemplate>,
419        std::collections::HashMap<String, String>,
420    )> {
421        let templates = self.templates.read().unwrap();
422        for template in templates.iter() {
423            if let Some(variables) = template.match_uri(uri) {
424                return Some((Arc::clone(template), variables));
425            }
426        }
427        None
428    }
429}
430
431/// A thread-safe, cloneable handle for runtime resource template management.
432///
433/// Obtained from [`McpRouter::with_dynamic_resource_templates()`](crate::McpRouter::with_dynamic_resource_templates).
434/// Templates registered here are merged with the router's static templates
435/// when handling `resources/templates/list` and `resources/read` requests.
436/// Static templates are checked before dynamic ones.
437///
438/// When a template is registered or unregistered, all connected sessions
439/// receive a `notifications/resources/list_changed` notification.
440///
441/// # Example
442///
443/// ```rust,ignore
444/// use tower_mcp::{McpRouter, ResourceTemplateBuilder};
445///
446/// let (router, registry) = McpRouter::new()
447///     .server_info("my-server", "1.0.0")
448///     .with_dynamic_resource_templates();
449///
450/// let template = ResourceTemplateBuilder::new("db://tables/{table}")
451///     .name("Database Table")
452///     .handler(|uri, vars| async move { /* ... */ });
453///
454/// registry.register(template);
455/// ```
456#[derive(Clone)]
457pub struct DynamicResourceTemplateRegistry {
458    inner: Arc<DynamicResourceTemplatesInner>,
459}
460
461impl DynamicResourceTemplateRegistry {
462    pub(crate) fn new(inner: Arc<DynamicResourceTemplatesInner>) -> Self {
463        Self { inner }
464    }
465
466    /// Register a resource template.
467    ///
468    /// Broadcasts `ResourcesListChanged` to all connected sessions.
469    pub fn register(&self, template: crate::resource::ResourceTemplate) {
470        {
471            let mut templates = self.inner.templates.write().unwrap();
472            // Remove any existing template with the same URI pattern
473            templates.retain(|t| t.uri_template != template.uri_template);
474            templates.push(Arc::new(template));
475        }
476        self.inner.broadcast_resources_changed();
477    }
478
479    /// Unregister a resource template by URI pattern.
480    ///
481    /// Returns `true` if the template existed and was removed.
482    /// Broadcasts `ResourcesListChanged` only if the template was actually removed.
483    pub fn unregister(&self, uri_template: &str) -> bool {
484        let removed = {
485            let mut templates = self.inner.templates.write().unwrap();
486            let before = templates.len();
487            templates.retain(|t| t.uri_template != uri_template);
488            templates.len() < before
489        };
490        if removed {
491            self.inner.broadcast_resources_changed();
492        }
493        removed
494    }
495
496    /// List all currently registered dynamic resource templates.
497    pub fn list(&self) -> Vec<Arc<crate::resource::ResourceTemplate>> {
498        self.inner.list()
499    }
500
501    /// Check if a template with the given URI pattern is registered.
502    pub fn contains(&self, uri_template: &str) -> bool {
503        let templates = self.inner.templates.read().unwrap();
504        templates.iter().any(|t| t.uri_template == uri_template)
505    }
506}
507
508#[cfg(test)]
509mod tests {
510    use super::*;
511    use crate::CallToolResult;
512    use crate::tool::ToolBuilder;
513    use tokio::sync::mpsc;
514
515    fn make_tool(name: &str) -> Tool {
516        ToolBuilder::new(name)
517            .description(format!("Test tool: {name}"))
518            .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
519            .build()
520    }
521
522    fn make_registry() -> (DynamicToolRegistry, Arc<DynamicToolsInner>) {
523        let inner = Arc::new(DynamicToolsInner::new());
524        let registry = DynamicToolRegistry::new(inner.clone());
525        (registry, inner)
526    }
527
528    #[test]
529    fn test_register_and_list() {
530        let (registry, _) = make_registry();
531
532        assert!(registry.list().is_empty());
533
534        registry.register(make_tool("tool_a"));
535        assert_eq!(registry.list().len(), 1);
536        assert!(registry.contains("tool_a"));
537
538        registry.register(make_tool("tool_b"));
539        assert_eq!(registry.list().len(), 2);
540        assert!(registry.contains("tool_b"));
541    }
542
543    #[test]
544    fn test_unregister() {
545        let (registry, _) = make_registry();
546
547        registry.register(make_tool("tool_a"));
548        registry.register(make_tool("tool_b"));
549        assert_eq!(registry.list().len(), 2);
550
551        assert!(registry.unregister("tool_a"));
552        assert_eq!(registry.list().len(), 1);
553        assert!(!registry.contains("tool_a"));
554        assert!(registry.contains("tool_b"));
555    }
556
557    #[test]
558    fn test_unregister_nonexistent() {
559        let (registry, _) = make_registry();
560        assert!(!registry.unregister("no_such_tool"));
561    }
562
563    #[test]
564    fn test_register_replaces_existing() {
565        let (registry, _) = make_registry();
566
567        registry.register(make_tool("tool_a"));
568        registry.register(make_tool("tool_a"));
569        assert_eq!(registry.list().len(), 1);
570    }
571
572    #[test]
573    fn test_contains() {
574        let (registry, _) = make_registry();
575
576        assert!(!registry.contains("tool_a"));
577        registry.register(make_tool("tool_a"));
578        assert!(registry.contains("tool_a"));
579        registry.unregister("tool_a");
580        assert!(!registry.contains("tool_a"));
581    }
582
583    #[test]
584    fn test_inner_get() {
585        let (registry, inner) = make_registry();
586
587        assert!(inner.get("tool_a").is_none());
588        registry.register(make_tool("tool_a"));
589        let tool = inner.get("tool_a").unwrap();
590        assert_eq!(tool.name, "tool_a");
591    }
592
593    #[tokio::test]
594    async fn test_broadcast_on_register() {
595        let (registry, inner) = make_registry();
596
597        let (tx, mut rx) = mpsc::channel(16);
598        inner.add_notification_sender(tx);
599
600        registry.register(make_tool("tool_a"));
601
602        let notification = rx.try_recv().unwrap();
603        assert!(matches!(notification, ServerNotification::ToolsListChanged));
604    }
605
606    #[tokio::test]
607    async fn test_broadcast_on_unregister() {
608        let (registry, inner) = make_registry();
609
610        registry.register(make_tool("tool_a"));
611
612        let (tx, mut rx) = mpsc::channel(16);
613        inner.add_notification_sender(tx);
614
615        registry.unregister("tool_a");
616
617        let notification = rx.try_recv().unwrap();
618        assert!(matches!(notification, ServerNotification::ToolsListChanged));
619    }
620
621    #[tokio::test]
622    async fn test_no_broadcast_on_unregister_nonexistent() {
623        let (registry, inner) = make_registry();
624
625        let (tx, mut rx) = mpsc::channel(16);
626        inner.add_notification_sender(tx);
627
628        registry.unregister("no_such_tool");
629
630        assert!(rx.try_recv().is_err());
631    }
632
633    #[tokio::test]
634    async fn test_closed_senders_are_cleaned_up() {
635        let (registry, inner) = make_registry();
636
637        let (tx, rx) = mpsc::channel(16);
638        inner.add_notification_sender(tx);
639        // Drop the receiver to close the channel
640        drop(rx);
641
642        // This should not panic, and should clean up the closed sender
643        registry.register(make_tool("tool_a"));
644
645        // Add a new sender and verify it still works
646        let (tx2, mut rx2) = mpsc::channel(16);
647        inner.add_notification_sender(tx2);
648
649        registry.register(make_tool("tool_b"));
650        let notification = rx2.try_recv().unwrap();
651        assert!(matches!(notification, ServerNotification::ToolsListChanged));
652    }
653
654    // =========================================================================
655    // DynamicPromptRegistry tests
656    // =========================================================================
657
658    fn make_prompt(name: &str) -> crate::prompt::Prompt {
659        crate::prompt::PromptBuilder::new(name)
660            .description(format!("Test prompt: {name}"))
661            .user_message("ok")
662    }
663
664    fn make_prompt_registry() -> (DynamicPromptRegistry, Arc<DynamicPromptsInner>) {
665        let inner = Arc::new(DynamicPromptsInner::new());
666        let registry = DynamicPromptRegistry::new(inner.clone());
667        (registry, inner)
668    }
669
670    #[test]
671    fn test_prompt_register_and_list() {
672        let (registry, _) = make_prompt_registry();
673
674        assert!(registry.list().is_empty());
675
676        registry.register(make_prompt("prompt_a"));
677        assert_eq!(registry.list().len(), 1);
678        assert!(registry.contains("prompt_a"));
679
680        registry.register(make_prompt("prompt_b"));
681        assert_eq!(registry.list().len(), 2);
682        assert!(registry.contains("prompt_b"));
683    }
684
685    #[test]
686    fn test_prompt_unregister() {
687        let (registry, _) = make_prompt_registry();
688
689        registry.register(make_prompt("prompt_a"));
690        registry.register(make_prompt("prompt_b"));
691
692        assert!(registry.unregister("prompt_a"));
693        assert_eq!(registry.list().len(), 1);
694        assert!(!registry.contains("prompt_a"));
695        assert!(registry.contains("prompt_b"));
696    }
697
698    #[test]
699    fn test_prompt_unregister_nonexistent() {
700        let (registry, _) = make_prompt_registry();
701        assert!(!registry.unregister("no_such_prompt"));
702    }
703
704    #[tokio::test]
705    async fn test_prompt_broadcast_on_register() {
706        let (registry, inner) = make_prompt_registry();
707
708        let (tx, mut rx) = mpsc::channel(16);
709        inner.add_notification_sender(tx);
710
711        registry.register(make_prompt("prompt_a"));
712
713        let notification = rx.try_recv().unwrap();
714        assert!(matches!(
715            notification,
716            ServerNotification::PromptsListChanged
717        ));
718    }
719
720    #[tokio::test]
721    async fn test_prompt_broadcast_on_unregister() {
722        let (registry, inner) = make_prompt_registry();
723
724        registry.register(make_prompt("prompt_a"));
725
726        let (tx, mut rx) = mpsc::channel(16);
727        inner.add_notification_sender(tx);
728
729        registry.unregister("prompt_a");
730
731        let notification = rx.try_recv().unwrap();
732        assert!(matches!(
733            notification,
734            ServerNotification::PromptsListChanged
735        ));
736    }
737
738    // =========================================================================
739    // DynamicResourceRegistry tests
740    // =========================================================================
741
742    fn make_resource(uri: &str) -> crate::resource::Resource {
743        crate::resource::ResourceBuilder::new(uri)
744            .name(uri)
745            .text("content")
746    }
747
748    fn make_resource_registry() -> (DynamicResourceRegistry, Arc<DynamicResourcesInner>) {
749        let inner = Arc::new(DynamicResourcesInner::new());
750        let registry = DynamicResourceRegistry::new(inner.clone());
751        (registry, inner)
752    }
753
754    #[test]
755    fn test_resource_register_and_list() {
756        let (registry, _) = make_resource_registry();
757
758        assert!(registry.list().is_empty());
759
760        registry.register(make_resource("file:///a.txt"));
761        assert_eq!(registry.list().len(), 1);
762        assert!(registry.contains("file:///a.txt"));
763
764        registry.register(make_resource("file:///b.txt"));
765        assert_eq!(registry.list().len(), 2);
766    }
767
768    #[test]
769    fn test_resource_unregister() {
770        let (registry, _) = make_resource_registry();
771
772        registry.register(make_resource("file:///a.txt"));
773        registry.register(make_resource("file:///b.txt"));
774
775        assert!(registry.unregister("file:///a.txt"));
776        assert_eq!(registry.list().len(), 1);
777        assert!(!registry.contains("file:///a.txt"));
778        assert!(registry.contains("file:///b.txt"));
779    }
780
781    #[test]
782    fn test_resource_unregister_nonexistent() {
783        let (registry, _) = make_resource_registry();
784        assert!(!registry.unregister("file:///nope"));
785    }
786
787    #[tokio::test]
788    async fn test_resource_broadcast_on_register() {
789        let (registry, inner) = make_resource_registry();
790
791        let (tx, mut rx) = mpsc::channel(16);
792        inner.add_notification_sender(tx);
793
794        registry.register(make_resource("file:///a.txt"));
795
796        let notification = rx.try_recv().unwrap();
797        assert!(matches!(
798            notification,
799            ServerNotification::ResourcesListChanged
800        ));
801    }
802
803    // =========================================================================
804    // DynamicResourceTemplateRegistry tests
805    // =========================================================================
806
807    fn make_template_registry() -> (
808        DynamicResourceTemplateRegistry,
809        Arc<DynamicResourceTemplatesInner>,
810    ) {
811        let inner = Arc::new(DynamicResourceTemplatesInner::new());
812        let registry = DynamicResourceTemplateRegistry::new(inner.clone());
813        (registry, inner)
814    }
815
816    #[test]
817    fn test_template_register_and_list() {
818        use crate::resource::ResourceTemplateBuilder;
819
820        let (registry, _) = make_template_registry();
821        assert!(registry.list().is_empty());
822
823        let template = ResourceTemplateBuilder::new("db://tables/{table}")
824            .name("Tables")
825            .handler(
826                |uri: String, _vars: std::collections::HashMap<String, String>| async move {
827                    Ok(crate::protocol::ReadResourceResult {
828                        contents: vec![crate::protocol::ResourceContent {
829                            uri,
830                            mime_type: None,
831                            text: Some("data".to_string()),
832                            blob: None,
833                            meta: None,
834                        }],
835                        meta: None,
836                        ..Default::default()
837                    })
838                },
839            );
840
841        registry.register(template);
842        assert_eq!(registry.list().len(), 1);
843        assert!(registry.contains("db://tables/{table}"));
844    }
845
846    #[test]
847    fn test_template_unregister() {
848        use crate::resource::ResourceTemplateBuilder;
849
850        let (registry, _) = make_template_registry();
851
852        let template = ResourceTemplateBuilder::new("db://tables/{table}")
853            .name("Tables")
854            .handler(
855                |uri: String, _vars: std::collections::HashMap<String, String>| async move {
856                    Ok(crate::protocol::ReadResourceResult {
857                        contents: vec![crate::protocol::ResourceContent {
858                            uri,
859                            mime_type: None,
860                            text: Some("data".to_string()),
861                            blob: None,
862                            meta: None,
863                        }],
864                        meta: None,
865                        ..Default::default()
866                    })
867                },
868            );
869
870        registry.register(template);
871        assert!(registry.unregister("db://tables/{table}"));
872        assert!(registry.list().is_empty());
873        assert!(!registry.unregister("db://tables/{table}"));
874    }
875
876    #[tokio::test]
877    async fn test_template_broadcast_on_register() {
878        use crate::resource::ResourceTemplateBuilder;
879
880        let (registry, inner) = make_template_registry();
881
882        let (tx, mut rx) = mpsc::channel(16);
883        inner.add_notification_sender(tx);
884
885        let template = ResourceTemplateBuilder::new("db://tables/{table}")
886            .name("Tables")
887            .handler(
888                |uri: String, _vars: std::collections::HashMap<String, String>| async move {
889                    Ok(crate::protocol::ReadResourceResult {
890                        contents: vec![crate::protocol::ResourceContent {
891                            uri,
892                            mime_type: None,
893                            text: Some("data".to_string()),
894                            blob: None,
895                            meta: None,
896                        }],
897                        meta: None,
898                        ..Default::default()
899                    })
900                },
901            );
902
903        registry.register(template);
904
905        let notification = rx.try_recv().unwrap();
906        assert!(matches!(
907            notification,
908            ServerNotification::ResourcesListChanged
909        ));
910    }
911
912    #[tokio::test]
913    async fn test_template_match_uri() {
914        use crate::resource::ResourceTemplateBuilder;
915
916        let (_, inner) = make_template_registry();
917
918        let template = ResourceTemplateBuilder::new("db://tables/{table}")
919            .name("Tables")
920            .handler(
921                |uri: String, _vars: std::collections::HashMap<String, String>| async move {
922                    Ok(crate::protocol::ReadResourceResult {
923                        contents: vec![crate::protocol::ResourceContent {
924                            uri,
925                            mime_type: None,
926                            text: Some("data".to_string()),
927                            blob: None,
928                            meta: None,
929                        }],
930                        meta: None,
931                        ..Default::default()
932                    })
933                },
934            );
935
936        {
937            let mut templates = inner.templates.write().unwrap();
938            templates.push(Arc::new(template));
939        }
940
941        let result = inner.match_uri("db://tables/users");
942        assert!(result.is_some());
943        let (_, vars) = result.unwrap();
944        assert_eq!(vars.get("table").unwrap(), "users");
945
946        assert!(inner.match_uri("db://other/path").is_none());
947    }
948}
tower_mcp/registry.rs

tower_mcp/
registry.rs
