mcp_host/registry/
resources.rs

1//! Resource registry for MCP servers
2//!
3//! Provides registration and reading of MCP resources with URI templates and retry logic
4
5use std::collections::HashMap;
6use std::sync::{Arc, RwLock};
7
8use async_trait::async_trait;
9use chrono_machines::{BackoffStrategy, ExponentialBackoff};
10use dashmap::DashMap;
11use rand::SeedableRng;
12use rand::rngs::SmallRng;
13use serde::{Deserialize, Serialize};
14use thiserror::Error;
15use tokio::sync::mpsc;
16
17use crate::content::resource::ResourceContent;
18use crate::server::session::Session;
19use crate::server::visibility::{ExecutionContext, VisibilityContext};
20use crate::transport::traits::JsonRpcNotification;
21
22/// Resource execution errors
23#[derive(Debug, Error)]
24pub enum ResourceError {
25    /// Resource not found
26    #[error("Resource not found: {0}")]
27    NotFound(String),
28
29    /// Invalid URI
30    #[error("Invalid URI: {0}")]
31    InvalidUri(String),
32
33    /// Read error (retryable)
34    #[error("Read error: {0}")]
35    Read(String),
36
37    /// Internal error
38    #[error("Internal error: {0}")]
39    Internal(String),
40
41    /// Retry exhausted
42    #[error("Retry exhausted after {attempts} attempts: {message}")]
43    RetryExhausted { attempts: u8, message: String },
44}
45
46/// Type alias for resource lookup results to avoid type complexity warnings
47pub type ResourceLookupResult = Result<(Arc<dyn Resource>, HashMap<String, String>), ResourceError>;
48
49impl ResourceError {
50    /// Check if this error is retryable
51    pub fn is_retryable(&self) -> bool {
52        matches!(self, Self::Read(_) | Self::Internal(_))
53    }
54}
55
56/// Retry configuration for resources
57#[derive(Debug, Clone)]
58pub struct ResourceRetryConfig {
59    /// Maximum number of retry attempts
60    pub max_attempts: u8,
61    /// Base delay in milliseconds
62    pub base_delay_ms: u64,
63    /// Exponential multiplier
64    pub multiplier: f64,
65    /// Maximum delay cap in milliseconds
66    pub max_delay_ms: u64,
67    /// Jitter factor (0.0 = no jitter, 1.0 = full jitter)
68    pub jitter_factor: f64,
69}
70
71impl Default for ResourceRetryConfig {
72    fn default() -> Self {
73        Self {
74            max_attempts: 3,
75            base_delay_ms: 100,
76            multiplier: 2.0,
77            max_delay_ms: 10_000,
78            jitter_factor: 1.0, // Full jitter by default
79        }
80    }
81}
82
83/// Resource metadata for listing
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct ResourceInfo {
86    /// Resource URI or URI template
87    pub uri: String,
88
89    /// Resource name
90    pub name: String,
91
92    /// Resource description
93    #[serde(skip_serializing_if = "Option::is_none")]
94    pub description: Option<String>,
95
96    /// MIME type
97    #[serde(skip_serializing_if = "Option::is_none", rename = "mimeType")]
98    pub mime_type: Option<String>,
99}
100
101/// Resource template metadata for listing
102#[derive(Debug, Clone, Serialize, Deserialize)]
103#[serde(rename_all = "camelCase")]
104pub struct ResourceTemplateInfo {
105    /// URI template with placeholders (e.g., "file:///{path}")
106    pub uri_template: String,
107
108    /// Template name
109    pub name: String,
110
111    /// Optional display title
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub title: Option<String>,
114
115    /// Template description
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub description: Option<String>,
118
119    /// MIME type for resources matching this template
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub mime_type: Option<String>,
122}
123
124/// Resource trait for implementing MCP resources
125#[async_trait]
126pub trait Resource: Send + Sync {
127    /// Get resource URI or URI template
128    fn uri(&self) -> &str;
129
130    /// Get resource name
131    fn name(&self) -> &str;
132
133    /// Get resource description
134    fn description(&self) -> Option<&str> {
135        None
136    }
137
138    /// Get MIME type
139    fn mime_type(&self) -> Option<&str> {
140        None
141    }
142
143    /// Check if this resource should be visible in the given context
144    ///
145    /// Override this to implement contextual visibility. The default implementation
146    /// always returns true (always visible).
147    fn is_visible(&self, _ctx: &VisibilityContext) -> bool {
148        true
149    }
150
151    /// Read the resource with execution context
152    ///
153    /// The execution context provides access to:
154    /// - `ctx.uri_params`: Parameters extracted from URI template
155    /// - `ctx.session`: Current session (roles, state, client info)
156    /// - `ctx.environment`: Optional environment state
157    ///
158    /// Returns ResourceContent which automatically includes uri and mimeType.
159    ///
160    /// # Example
161    ///
162    /// ```rust,ignore
163    /// async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
164    ///     let id = ctx.get_uri_param("id").unwrap_or("unknown");
165    ///     let content = format!("Resource {}", id);
166    ///
167    ///     Ok(vec![self.text_content(content)])
168    /// }
169    /// ```
170    async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError>;
171
172    /// Helper: Create a text ResourceContent with this resource's URI and mime_type
173    ///
174    /// This is a convenience method that automatically includes the resource's
175    /// URI and MIME type in the response.
176    ///
177    /// # Example
178    ///
179    /// ```rust,ignore
180    /// async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
181    ///     Ok(vec![self.text_content("Resource content here")])
182    /// }
183    /// ```
184    fn text_content(&self, text: &str) -> ResourceContent {
185        ResourceContent {
186            uri: self.uri().to_string(),
187            mime_type: self.mime_type().map(|s| s.to_string()),
188            text: Some(text.to_string()),
189            blob: None,
190        }
191    }
192
193    /// Helper: Create a blob ResourceContent with this resource's URI and mime_type
194    ///
195    /// # Example
196    ///
197    /// ```rust,ignore
198    /// async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
199    ///     Ok(vec![self.blob_content(base64_data)])
200    /// }
201    /// ```
202    fn blob_content(&self, data: &str) -> ResourceContent {
203        ResourceContent {
204            uri: self.uri().to_string(),
205            mime_type: self.mime_type().map(|s| s.to_string()),
206            text: None,
207            blob: Some(data.to_string()),
208        }
209    }
210}
211
212/// Resource template trait for implementing URI templates
213///
214/// Resource templates allow dynamic resources with URI patterns like `file:///{path}`.
215/// The template defines the pattern, and resources matching it provide the actual data.
216#[async_trait]
217pub trait ResourceTemplate: Send + Sync {
218    /// Get URI template pattern (e.g., "file:///{path}")
219    fn uri_template(&self) -> &str;
220
221    /// Get template name
222    fn name(&self) -> &str;
223
224    /// Get optional display title
225    fn title(&self) -> Option<&str> {
226        None
227    }
228
229    /// Get template description
230    fn description(&self) -> Option<&str> {
231        None
232    }
233
234    /// Get MIME type for resources matching this template
235    fn mime_type(&self) -> Option<&str> {
236        None
237    }
238
239    /// Check if this template should be visible in the given context
240    fn is_visible(&self, _ctx: &VisibilityContext) -> bool {
241        true
242    }
243
244    /// Read a resource matching this template
245    ///
246    /// The URI parameters are extracted from the concrete URI and passed via `ctx.uri_params`.
247    ///
248    /// # Example
249    ///
250    /// ```rust,ignore
251    /// async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
252    ///     let path = ctx.get_uri_param("path").ok_or_else(||
253    ///         ResourceError::InvalidUri("Missing path parameter".into())
254    ///     )?;
255    ///
256    ///     let content = std::fs::read_to_string(path)
257    ///         .map_err(|e| ResourceError::Read(e.to_string()))?;
258    ///
259    ///     Ok(vec![ResourceContent::text(
260    ///         format!("file:///{}", path),
261    ///         content
262    ///     )])
263    /// }
264    /// ```
265    async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError>;
266}
267
268/// Dynamic resource that wraps a template for on-demand resource creation
269struct DynamicTemplateResource {
270    template: Arc<dyn ResourceTemplate>,
271    uri: String,
272}
273
274#[async_trait]
275impl Resource for DynamicTemplateResource {
276    fn uri(&self) -> &str {
277        &self.uri
278    }
279
280    fn name(&self) -> &str {
281        self.template.name()
282    }
283
284    fn description(&self) -> Option<&str> {
285        self.template.description()
286    }
287
288    fn mime_type(&self) -> Option<&str> {
289        self.template.mime_type()
290    }
291
292    fn is_visible(&self, ctx: &VisibilityContext) -> bool {
293        self.template.is_visible(ctx)
294    }
295
296    async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
297        self.template.read(ctx).await
298    }
299}
300
301/// Resource manager for managing available resources with retry logic
302#[derive(Clone)]
303pub struct ResourceManager {
304    resources: Arc<DashMap<String, Arc<dyn Resource>>>,
305    templates: Arc<DashMap<String, Arc<dyn ResourceTemplate>>>,
306    retry_config: Arc<RwLock<ResourceRetryConfig>>,
307    notification_tx: Option<mpsc::UnboundedSender<JsonRpcNotification>>,
308}
309
310impl ResourceManager {
311    /// Create new resource manager
312    pub fn new() -> Self {
313        Self {
314            resources: Arc::new(DashMap::new()),
315            templates: Arc::new(DashMap::new()),
316            retry_config: Arc::new(RwLock::new(ResourceRetryConfig::default())),
317            notification_tx: None,
318        }
319    }
320
321    /// Create resource manager with notification channel
322    pub fn with_notifications(notification_tx: mpsc::UnboundedSender<JsonRpcNotification>) -> Self {
323        Self {
324            resources: Arc::new(DashMap::new()),
325            templates: Arc::new(DashMap::new()),
326            retry_config: Arc::new(RwLock::new(ResourceRetryConfig::default())),
327            notification_tx: Some(notification_tx),
328        }
329    }
330
331    /// Set notification channel
332    pub fn set_notification_tx(&mut self, tx: mpsc::UnboundedSender<JsonRpcNotification>) {
333        self.notification_tx = Some(tx);
334    }
335
336    /// Configure retry behavior
337    pub fn set_retry_config(&self, config: ResourceRetryConfig) {
338        if let Ok(mut cfg) = self.retry_config.write() {
339            *cfg = config;
340        }
341    }
342
343    /// Send notification
344    fn send_notification(&self, method: &str, params: Option<serde_json::Value>) {
345        if let Some(tx) = &self.notification_tx {
346            let notification = JsonRpcNotification::new(method, params);
347            let _ = tx.send(notification);
348        }
349    }
350
351    /// Notify resources list changed
352    fn notify_resources_changed(&self) {
353        self.send_notification("notifications/resources/list_changed", None);
354    }
355
356    /// Send logging message notification (visible to LLM)
357    fn notify_message(&self, level: &str, logger: &str, message: &str) {
358        self.send_notification(
359            "notifications/message",
360            Some(serde_json::json!({
361                "level": level,
362                "logger": logger,
363                "data": message
364            })),
365        );
366    }
367
368    /// Register a resource
369    pub fn register<R: Resource + 'static>(&self, resource: R) {
370        let uri = resource.uri().to_string();
371        self.resources.insert(uri, Arc::new(resource));
372    }
373
374    /// Register a boxed resource
375    pub fn register_boxed(&self, resource: Arc<dyn Resource>) {
376        let uri = resource.uri().to_string();
377        self.resources.insert(uri, resource);
378    }
379
380    /// Register a resource template
381    pub fn register_template<T: ResourceTemplate + 'static>(&self, template: T) {
382        let name = template.name().to_string();
383        self.templates.insert(name, Arc::new(template));
384    }
385
386    /// Register a boxed resource template
387    pub fn register_template_boxed(&self, template: Arc<dyn ResourceTemplate>) {
388        let name = template.name().to_string();
389        self.templates.insert(name, template);
390    }
391
392    /// Get a template by name
393    pub fn get_template(&self, name: &str) -> Option<Arc<dyn ResourceTemplate>> {
394        self.templates.get(name).map(|t| Arc::clone(&t))
395    }
396
397    /// Get a resource by URI
398    pub fn get(&self, uri: &str) -> Option<Arc<dyn Resource>> {
399        self.resources.get(uri).map(|r| Arc::clone(&r))
400    }
401
402    /// List all registered resources
403    pub fn list(&self) -> Vec<ResourceInfo> {
404        self.resources
405            .iter()
406            .map(|entry| {
407                let resource = entry.value();
408                ResourceInfo {
409                    uri: resource.uri().to_string(),
410                    name: resource.name().to_string(),
411                    description: resource.description().map(|s| s.to_string()),
412                    mime_type: resource.mime_type().map(|s| s.to_string()),
413                }
414            })
415            .collect()
416    }
417
418    /// List all resource templates
419    pub fn list_templates(&self) -> Vec<ResourceTemplateInfo> {
420        self.templates
421            .iter()
422            .map(|entry| {
423                let template = entry.value();
424                ResourceTemplateInfo {
425                    uri_template: template.uri_template().to_string(),
426                    name: template.name().to_string(),
427                    title: template.title().map(|s| s.to_string()),
428                    description: template.description().map(|s| s.to_string()),
429                    mime_type: template.mime_type().map(|s| s.to_string()),
430                }
431            })
432            .collect()
433    }
434
435    /// List templates visible in the given session/context
436    pub fn list_templates_for_session(
437        &self,
438        _session: &Session,
439        ctx: &VisibilityContext<'_>,
440    ) -> Vec<ResourceTemplateInfo> {
441        self.templates
442            .iter()
443            .filter(|entry| entry.value().is_visible(ctx))
444            .map(|entry| {
445                let template = entry.value();
446                ResourceTemplateInfo {
447                    uri_template: template.uri_template().to_string(),
448                    name: template.name().to_string(),
449                    title: template.title().map(|s| s.to_string()),
450                    description: template.description().map(|s| s.to_string()),
451                    mime_type: template.mime_type().map(|s| s.to_string()),
452                }
453            })
454            .collect()
455    }
456
457    /// Read a resource by URI with parameters and retry logic
458    pub async fn read(
459        &self,
460        uri: &str,
461        params: HashMap<String, String>,
462        session: &Session,
463        logger: &crate::logging::McpLogger,
464    ) -> Result<Vec<ResourceContent>, ResourceError> {
465        // Find the resource first
466        let (resource, combined_params) = self.find_resource(uri, params)?;
467
468        // Get retry config
469        let retry_config = self
470            .retry_config
471            .read()
472            .map(|c| c.clone())
473            .unwrap_or_default();
474
475        // Build backoff strategy using chrono-machines
476        let backoff = ExponentialBackoff::default()
477            .max_attempts(retry_config.max_attempts)
478            .base_delay_ms(retry_config.base_delay_ms)
479            .multiplier(retry_config.multiplier)
480            .max_delay_ms(retry_config.max_delay_ms)
481            .jitter_factor(retry_config.jitter_factor);
482
483        let mut rng = SmallRng::from_os_rng();
484        let mut attempt: u8 = 1;
485
486        loop {
487            // Create execution context for resource
488            let ctx = ExecutionContext::for_resource(combined_params.clone(), session, logger);
489
490            match resource.read(ctx).await {
491                Ok(content) => {
492                    // If we recovered after failures, notify
493                    if attempt > 1 {
494                        self.notify_message(
495                            "info",
496                            "chrono-machines",
497                            &format!("Resource '{}' succeeded after {} attempts", uri, attempt),
498                        );
499                    }
500                    return Ok(content);
501                }
502                Err(e) => {
503                    // Check if error is retryable
504                    if !e.is_retryable() {
505                        return Err(e);
506                    }
507
508                    // Check if we should retry
509                    if !backoff.should_retry(attempt) {
510                        // Exhausted retries, send notifications
511                        self.notify_resources_changed();
512                        self.notify_message(
513                            "warning",
514                            "chrono-machines",
515                            &format!(
516                                "Resource '{}' failed after {} attempts: {}",
517                                uri, attempt, e
518                            ),
519                        );
520
521                        return Err(ResourceError::RetryExhausted {
522                            attempts: attempt,
523                            message: e.to_string(),
524                        });
525                    }
526
527                    // Calculate delay and sleep
528                    if let Some(delay_ms) = backoff.delay(attempt, &mut rng) {
529                        tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
530                    }
531
532                    attempt = attempt.saturating_add(1);
533                }
534            }
535        }
536    }
537
538    /// Find a resource by URI (exact match or template match)
539    fn find_resource(&self, uri: &str, params: HashMap<String, String>) -> ResourceLookupResult {
540        // Try exact match first
541        if let Some(resource) = self.get(uri) {
542            return Ok((resource, params));
543        }
544
545        // Try template matching in static resources
546        for entry in self.resources.iter() {
547            let template_uri = entry.key();
548            if let Some(extracted_params) = self.match_template(template_uri, uri) {
549                let resource = Arc::clone(entry.value());
550                let mut combined_params = params.clone();
551                combined_params.extend(extracted_params);
552                return Ok((resource, combined_params));
553            }
554        }
555
556        // Try matching against resource templates
557        for entry in self.templates.iter() {
558            let template = entry.value();
559            let template_uri = template.uri_template();
560            if let Some(extracted_params) = self.match_template(template_uri, uri) {
561                // Create a dynamic resource from the template
562                let dynamic_resource = DynamicTemplateResource {
563                    template: Arc::clone(template),
564                    uri: uri.to_string(),
565                };
566                let mut combined_params = params.clone();
567                combined_params.extend(extracted_params);
568                return Ok((Arc::new(dynamic_resource), combined_params));
569            }
570        }
571
572        Err(ResourceError::NotFound(uri.to_string()))
573    }
574
575    /// Match URI against template and extract parameters
576    /// Template format: "scheme://path/{param}/more/{param2}"
577    /// If the last part is a parameter, it greedily matches the rest of the URI
578    fn match_template(&self, template: &str, uri: &str) -> Option<HashMap<String, String>> {
579        let template_parts: Vec<&str> = template.split('/').collect();
580        let uri_parts: Vec<&str> = uri.split('/').collect();
581
582        // Check if template has fewer parts (only valid if last part is a parameter)
583        if template_parts.len() > uri_parts.len() {
584            return None;
585        }
586
587        let mut params = HashMap::new();
588        let last_idx = template_parts.len() - 1;
589
590        for (i, template_part) in template_parts.iter().enumerate() {
591            if template_part.starts_with('{') && template_part.ends_with('}') {
592                let param_name = &template_part[1..template_part.len() - 1];
593
594                if i == last_idx {
595                    // Last parameter: greedily match remaining URI parts
596                    let remaining = uri_parts[i..].join("/");
597                    params.insert(param_name.to_string(), remaining);
598                    break;
599                } else {
600                    // Middle parameter: match single part
601                    if i >= uri_parts.len() {
602                        return None;
603                    }
604                    params.insert(param_name.to_string(), uri_parts[i].to_string());
605                }
606            } else if i >= uri_parts.len() || template_part != &uri_parts[i] {
607                // Static part doesn't match or URI too short
608                return None;
609            }
610        }
611
612        Some(params)
613    }
614
615    // ==================== Session-Aware Methods ====================
616
617    /// List resources visible to a specific session
618    ///
619    /// Resolution order:
620    /// 1. Session overrides (replace global resources with session-specific implementations)
621    /// 2. Session extras (additional resources added to session)
622    /// 3. Global resources (filtered by hidden list and visibility predicate)
623    pub fn list_for_session(
624        &self,
625        session: &Session,
626        ctx: &VisibilityContext<'_>,
627    ) -> Vec<ResourceInfo> {
628        let mut resources = std::collections::HashMap::new();
629
630        // 1. Add global resources (filtered by hidden and visibility)
631        for entry in self.resources.iter() {
632            let uri = entry.key().clone();
633            if !session.is_resource_hidden(&uri) {
634                let resource = entry.value();
635                if resource.is_visible(ctx) {
636                    resources.insert(
637                        uri,
638                        ResourceInfo {
639                            uri: resource.uri().to_string(),
640                            name: resource.name().to_string(),
641                            description: resource.description().map(|s| s.to_string()),
642                            mime_type: resource.mime_type().map(|s| s.to_string()),
643                        },
644                    );
645                }
646            }
647        }
648
649        // 2. Add session extras
650        for entry in session.resource_extras().iter() {
651            let uri = entry.key().clone();
652            let resource = entry.value();
653            if resource.is_visible(ctx) {
654                resources.insert(
655                    uri,
656                    ResourceInfo {
657                        uri: resource.uri().to_string(),
658                        name: resource.name().to_string(),
659                        description: resource.description().map(|s| s.to_string()),
660                        mime_type: resource.mime_type().map(|s| s.to_string()),
661                    },
662                );
663            }
664        }
665
666        // 3. Apply session overrides
667        for entry in session.resource_overrides().iter() {
668            let uri = entry.key().clone();
669            let resource = entry.value();
670            if resource.is_visible(ctx) {
671                resources.insert(
672                    uri,
673                    ResourceInfo {
674                        uri: resource.uri().to_string(),
675                        name: resource.name().to_string(),
676                        description: resource.description().map(|s| s.to_string()),
677                        mime_type: resource.mime_type().map(|s| s.to_string()),
678                    },
679                );
680            }
681        }
682
683        resources.into_values().collect()
684    }
685
686    /// Read a resource with session context
687    ///
688    /// Resolution order:
689    /// 1. Check session overrides
690    /// 2. Check session extras
691    /// 3. Check session hidden
692    /// 4. Check visibility predicate
693    /// 5. Fall back to global registry with retry logic
694    pub async fn read_for_session(
695        &self,
696        uri: &str,
697        params: HashMap<String, String>,
698        session: &Session,
699        logger: &crate::logging::McpLogger,
700        visibility_ctx: &VisibilityContext<'_>,
701    ) -> Result<Vec<ResourceContent>, ResourceError> {
702        // Create execution context (reuse environment from visibility context)
703        let exec_ctx = match visibility_ctx.environment {
704            Some(env) => ExecutionContext::for_resource_with_environment(
705                params.clone(),
706                session,
707                logger,
708                env,
709            ),
710            None => ExecutionContext::for_resource(params.clone(), session, logger),
711        };
712
713        // 1. Check session override first
714        if let Some(resource) = session.get_resource_override(uri) {
715            if !resource.is_visible(visibility_ctx) {
716                return Err(ResourceError::NotFound(uri.to_string()));
717            }
718            return resource.read(exec_ctx).await;
719        }
720
721        // 2. Check session extras
722        if let Some(resource) = session.get_resource_extra(uri) {
723            if !resource.is_visible(visibility_ctx) {
724                return Err(ResourceError::NotFound(uri.to_string()));
725            }
726            return resource.read(exec_ctx).await;
727        }
728
729        // 3. Check if hidden in session
730        if session.is_resource_hidden(uri) {
731            return Err(ResourceError::NotFound(uri.to_string()));
732        }
733
734        // 4. Check global registry with visibility check
735        if let Some(resource) = self.get(uri)
736            && !resource.is_visible(visibility_ctx)
737        {
738            return Err(ResourceError::NotFound(uri.to_string()));
739        }
740
741        // 5. Fall back to global registry with retry logic
742        self.read(uri, params, session, logger).await
743    }
744
745    /// Get number of registered resources
746    pub fn len(&self) -> usize {
747        self.resources.len()
748    }
749
750    /// Check if manager is empty
751    pub fn is_empty(&self) -> bool {
752        self.resources.is_empty()
753    }
754}
755
756impl Default for ResourceManager {
757    fn default() -> Self {
758        Self::new()
759    }
760}
761
762#[cfg(test)]
763mod tests {
764    use super::*;
765
766    // Example static resource for testing
767    struct HelloResource;
768
769    #[async_trait]
770    impl Resource for HelloResource {
771        fn uri(&self) -> &str {
772            "test://hello"
773        }
774
775        fn name(&self) -> &str {
776            "hello"
777        }
778
779        fn description(&self) -> Option<&str> {
780            Some("Returns a greeting")
781        }
782
783        fn mime_type(&self) -> Option<&str> {
784            Some("text/plain")
785        }
786
787        async fn read(
788            &self,
789            _ctx: ExecutionContext<'_>,
790        ) -> Result<Vec<ResourceContent>, ResourceError> {
791            Ok(vec![self.text_content("Hello, World!")])
792        }
793    }
794
795    // Example template resource for testing
796    struct UserResource;
797
798    #[async_trait]
799    impl Resource for UserResource {
800        fn uri(&self) -> &str {
801            "test://users/{id}"
802        }
803
804        fn name(&self) -> &str {
805            "user"
806        }
807
808        fn description(&self) -> Option<&str> {
809            Some("Returns user information")
810        }
811
812        async fn read(
813            &self,
814            ctx: ExecutionContext<'_>,
815        ) -> Result<Vec<ResourceContent>, ResourceError> {
816            let id = ctx
817                .get_uri_param("id")
818                .ok_or_else(|| ResourceError::InvalidUri("Missing 'id' parameter".to_string()))?;
819
820            Ok(vec![self.text_content(&format!("User ID: {}", id))])
821        }
822    }
823
824    #[test]
825    fn test_manager_creation() {
826        let manager = ResourceManager::new();
827        assert!(manager.is_empty());
828    }
829
830    #[test]
831    fn test_resource_registration() {
832        let manager = ResourceManager::new();
833        manager.register(HelloResource);
834
835        assert_eq!(manager.len(), 1);
836        assert!(!manager.is_empty());
837    }
838
839    #[test]
840    fn test_get_resource() {
841        let manager = ResourceManager::new();
842        manager.register(HelloResource);
843
844        let resource = manager.get("test://hello");
845        assert!(resource.is_some());
846        assert_eq!(resource.unwrap().name(), "hello");
847
848        let missing = manager.get("test://nonexistent");
849        assert!(missing.is_none());
850    }
851
852    #[test]
853    fn test_list_resources() {
854        let manager = ResourceManager::new();
855        manager.register(HelloResource);
856
857        let resources = manager.list();
858        assert_eq!(resources.len(), 1);
859        assert_eq!(resources[0].uri, "test://hello");
860        assert_eq!(resources[0].name, "hello");
861        assert_eq!(
862            resources[0].description,
863            Some("Returns a greeting".to_string())
864        );
865        assert_eq!(resources[0].mime_type, Some("text/plain".to_string()));
866    }
867
868    #[tokio::test]
869    async fn test_read_static_resource() {
870        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
871        let logger = crate::logging::McpLogger::new(_tx, "test");
872        let manager = ResourceManager::new();
873        manager.register(HelloResource);
874        let session = Session::new();
875
876        let result = manager
877            .read("test://hello", HashMap::new(), &session, &logger)
878            .await
879            .unwrap();
880        assert_eq!(result.len(), 1);
881    }
882
883    #[tokio::test]
884    async fn test_read_missing_resource() {
885        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
886        let logger = crate::logging::McpLogger::new(_tx, "test");
887        let manager = ResourceManager::new();
888        let session = Session::new();
889
890        let result = manager
891            .read("test://nonexistent", HashMap::new(), &session, &logger)
892            .await;
893        assert!(matches!(result, Err(ResourceError::NotFound(_))));
894    }
895
896    #[tokio::test]
897    async fn test_template_matching() {
898        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
899        let logger = crate::logging::McpLogger::new(_tx, "test");
900        let manager = ResourceManager::new();
901        manager.register(UserResource);
902        let session = Session::new();
903
904        // Test template parameter extraction
905        let result = manager
906            .read("test://users/123", HashMap::new(), &session, &logger)
907            .await
908            .unwrap();
909        assert_eq!(result.len(), 1);
910    }
911
912    #[test]
913    fn test_template_matching_internal() {
914        let manager = ResourceManager::new();
915
916        // Test exact template matching
917        let params = manager.match_template("test://users/{id}", "test://users/123");
918        assert!(params.is_some());
919        let params = params.unwrap();
920        assert_eq!(params.get("id"), Some(&"123".to_string()));
921
922        // Test multiple parameters
923        let params = manager.match_template(
924            "test://org/{org}/repo/{repo}",
925            "test://org/myorg/repo/myrepo",
926        );
927        assert!(params.is_some());
928        let params = params.unwrap();
929        assert_eq!(params.get("org"), Some(&"myorg".to_string()));
930        assert_eq!(params.get("repo"), Some(&"myrepo".to_string()));
931
932        // Test non-matching
933        let params = manager.match_template("test://users/{id}", "test://posts/123");
934        assert!(params.is_none());
935
936        // Test greedy matching (last parameter matches remaining segments)
937        let params = manager.match_template("test://users/{id}", "test://users/123/extra");
938        assert!(params.is_some());
939        let params = params.unwrap();
940        assert_eq!(params.get("id"), Some(&"123/extra".to_string()));
941    }
942
943    #[tokio::test]
944    async fn test_resource_non_matching_template() {
945        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
946        let logger = crate::logging::McpLogger::new(_tx, "test");
947        let manager = ResourceManager::new();
948        manager.register(UserResource);
949        let session = Session::new();
950
951        // This won't match because "posts" != "users"
952        let result = manager
953            .read("test://posts/123", HashMap::new(), &session, &logger)
954            .await;
955        assert!(matches!(result, Err(ResourceError::NotFound(_))));
956    }
957
958    #[tokio::test]
959    async fn test_resource_greedy_matching() {
960        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
961        let logger = crate::logging::McpLogger::new(_tx, "test");
962        let manager = ResourceManager::new();
963        manager.register(UserResource);
964        let session = Session::new();
965
966        // The {id} parameter greedily matches the remaining segments
967        let result = manager
968            .read("test://users/123/extra", HashMap::new(), &session, &logger)
969            .await;
970        assert!(result.is_ok());
971        let content = result.unwrap();
972        assert_eq!(content.len(), 1);
973        // Verify the id parameter captured the greedy match
974        assert!(content[0].text.as_ref().unwrap().contains("123/extra"));
975    }
976
977    // ==================== Resource Template Tests ====================
978
979    // Example template for testing
980    struct FileTemplate;
981
982    #[async_trait]
983    impl ResourceTemplate for FileTemplate {
984        fn uri_template(&self) -> &str {
985            "file:///{path}"
986        }
987
988        fn name(&self) -> &str {
989            "project_files"
990        }
991
992        fn title(&self) -> Option<&str> {
993            Some("Project Files")
994        }
995
996        fn description(&self) -> Option<&str> {
997            Some("Access files in the project directory")
998        }
999
1000        fn mime_type(&self) -> Option<&str> {
1001            Some("application/octet-stream")
1002        }
1003
1004        async fn read(
1005            &self,
1006            ctx: ExecutionContext<'_>,
1007        ) -> Result<Vec<ResourceContent>, ResourceError> {
1008            let path = ctx
1009                .get_uri_param("path")
1010                .ok_or_else(|| ResourceError::InvalidUri("Missing 'path' parameter".to_string()))?;
1011
1012            Ok(vec![ResourceContent::text(
1013                format!("file:///{}", path),
1014                format!("Mock content for file: {}", path),
1015            )])
1016        }
1017    }
1018
1019    // Template with visibility predicate
1020    struct AdminOnlyTemplate;
1021
1022    #[async_trait]
1023    impl ResourceTemplate for AdminOnlyTemplate {
1024        fn uri_template(&self) -> &str {
1025            "admin:///{resource}"
1026        }
1027
1028        fn name(&self) -> &str {
1029            "admin_resources"
1030        }
1031
1032        fn is_visible(&self, ctx: &VisibilityContext) -> bool {
1033            ctx.has_role("admin")
1034        }
1035
1036        async fn read(
1037            &self,
1038            _ctx: ExecutionContext<'_>,
1039        ) -> Result<Vec<ResourceContent>, ResourceError> {
1040            Ok(vec![ResourceContent::text(
1041                "admin:///test",
1042                "Admin content",
1043            )])
1044        }
1045    }
1046
1047    #[test]
1048    fn test_template_registration() {
1049        let manager = ResourceManager::new();
1050        manager.register_template(FileTemplate);
1051
1052        let template = manager.get_template("project_files");
1053        assert!(template.is_some());
1054        assert_eq!(template.unwrap().name(), "project_files");
1055    }
1056
1057    #[test]
1058    fn test_get_template() {
1059        let manager = ResourceManager::new();
1060        manager.register_template(FileTemplate);
1061
1062        let template = manager.get_template("project_files");
1063        assert!(template.is_some());
1064        assert_eq!(template.unwrap().uri_template(), "file:///{path}");
1065
1066        let missing = manager.get_template("nonexistent");
1067        assert!(missing.is_none());
1068    }
1069
1070    #[test]
1071    fn test_list_templates() {
1072        let manager = ResourceManager::new();
1073        manager.register_template(FileTemplate);
1074
1075        let templates = manager.list_templates();
1076        assert_eq!(templates.len(), 1);
1077        assert_eq!(templates[0].uri_template, "file:///{path}");
1078        assert_eq!(templates[0].name, "project_files");
1079        assert_eq!(templates[0].title, Some("Project Files".to_string()));
1080        assert_eq!(
1081            templates[0].description,
1082            Some("Access files in the project directory".to_string())
1083        );
1084        assert_eq!(
1085            templates[0].mime_type,
1086            Some("application/octet-stream".to_string())
1087        );
1088    }
1089
1090    #[test]
1091    fn test_list_templates_for_session_visibility() {
1092        let manager = ResourceManager::new();
1093        manager.register_template(FileTemplate);
1094        manager.register_template(AdminOnlyTemplate);
1095
1096        // Regular session without admin role
1097        let session = Session::new();
1098        let ctx = VisibilityContext::new(&session);
1099        let templates = manager.list_templates_for_session(&session, &ctx);
1100
1101        // Should only see FileTemplate, not AdminOnlyTemplate
1102        assert_eq!(templates.len(), 1);
1103        assert_eq!(templates[0].name, "project_files");
1104
1105        // Admin session
1106        let admin_session = Session::new();
1107        admin_session.set_state("roles", serde_json::json!(["admin"]));
1108        let admin_ctx = VisibilityContext::new(&admin_session);
1109        let admin_templates = manager.list_templates_for_session(&admin_session, &admin_ctx);
1110
1111        // Should see both templates
1112        assert_eq!(admin_templates.len(), 2);
1113        let names: Vec<_> = admin_templates.iter().map(|t| t.name.as_str()).collect();
1114        assert!(names.contains(&"project_files"));
1115        assert!(names.contains(&"admin_resources"));
1116    }
1117
1118    #[test]
1119    fn test_template_info_serialization() {
1120        let manager = ResourceManager::new();
1121        manager.register_template(FileTemplate);
1122
1123        let templates = manager.list_templates();
1124        let serialized = serde_json::to_value(&templates[0]).unwrap();
1125
1126        // Verify MCP-compliant field names (camelCase)
1127        assert!(serialized.get("uriTemplate").is_some());
1128        assert_eq!(serialized["uriTemplate"], "file:///{path}");
1129        assert_eq!(serialized["name"], "project_files");
1130        assert_eq!(serialized["mimeType"], "application/octet-stream");
1131
1132        // Verify snake_case fields are NOT present
1133        assert!(serialized.get("uri_template").is_none());
1134        assert!(serialized.get("mime_type").is_none());
1135    }
1136
1137    #[tokio::test]
1138    async fn test_template_read_with_uri_params() {
1139        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
1140        let logger = crate::logging::McpLogger::new(_tx, "test");
1141        let session = Session::new();
1142        let ctx = ExecutionContext::for_resource(
1143            vec![("path".to_string(), "src/main.rs".to_string())]
1144                .into_iter()
1145                .collect(),
1146            &session,
1147            &logger,
1148        );
1149
1150        let template = FileTemplate;
1151        let result = template.read(ctx).await.unwrap();
1152
1153        assert_eq!(result.len(), 1);
1154        assert_eq!(result[0].uri, "file:///src/main.rs");
1155        assert!(result[0].text.as_ref().unwrap().contains("src/main.rs"));
1156    }
1157
1158    #[tokio::test]
1159    async fn test_template_read_missing_param() {
1160        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
1161        let logger = crate::logging::McpLogger::new(_tx, "test");
1162        let session = Session::new();
1163        let ctx = ExecutionContext::for_resource(HashMap::new(), &session, &logger);
1164
1165        let template = FileTemplate;
1166        let result = template.read(ctx).await;
1167
1168        // Should fail with InvalidUri when path parameter is missing
1169        assert!(matches!(result, Err(ResourceError::InvalidUri(_))));
1170    }
1171}
mcp_host/registry/resources.rs

mcp_host/registry/
resources.rs
