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::rngs::SmallRng;
12use rand::SeedableRng;
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    ) -> Result<Vec<ResourceContent>, ResourceError> {
464        // Find the resource first
465        let (resource, combined_params) = self.find_resource(uri, params)?;
466
467        // Get retry config
468        let retry_config = self.retry_config.read()
469            .map(|c| c.clone())
470            .unwrap_or_default();
471
472        // Build backoff strategy using chrono-machines
473        let backoff = ExponentialBackoff::default()
474            .max_attempts(retry_config.max_attempts)
475            .base_delay_ms(retry_config.base_delay_ms)
476            .multiplier(retry_config.multiplier)
477            .max_delay_ms(retry_config.max_delay_ms)
478            .jitter_factor(retry_config.jitter_factor);
479
480        let mut rng = SmallRng::from_os_rng();
481        let mut attempt: u8 = 1;
482
483        loop {
484            // Create execution context for resource
485            let ctx = ExecutionContext::for_resource(combined_params.clone(), session);
486
487            match resource.read(ctx).await {
488                Ok(content) => {
489                    // If we recovered after failures, notify
490                    if attempt > 1 {
491                        self.notify_message(
492                            "info",
493                            "chrono-machines",
494                            &format!("Resource '{}' succeeded after {} attempts", uri, attempt),
495                        );
496                    }
497                    return Ok(content);
498                }
499                Err(e) => {
500                    // Check if error is retryable
501                    if !e.is_retryable() {
502                        return Err(e);
503                    }
504
505                    // Check if we should retry
506                    if !backoff.should_retry(attempt) {
507                        // Exhausted retries, send notifications
508                        self.notify_resources_changed();
509                        self.notify_message(
510                            "warning",
511                            "chrono-machines",
512                            &format!(
513                                "Resource '{}' failed after {} attempts: {}",
514                                uri, attempt, e
515                            ),
516                        );
517
518                        return Err(ResourceError::RetryExhausted {
519                            attempts: attempt,
520                            message: e.to_string(),
521                        });
522                    }
523
524                    // Calculate delay and sleep
525                    if let Some(delay_ms) = backoff.delay(attempt, &mut rng) {
526                        tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
527                    }
528
529                    attempt = attempt.saturating_add(1);
530                }
531            }
532        }
533    }
534
535    /// Find a resource by URI (exact match or template match)
536    fn find_resource(
537        &self,
538        uri: &str,
539        params: HashMap<String, String>,
540    ) -> ResourceLookupResult {
541        // Try exact match first
542        if let Some(resource) = self.get(uri) {
543            return Ok((resource, params));
544        }
545
546        // Try template matching in static resources
547        for entry in self.resources.iter() {
548            let template_uri = entry.key();
549            if let Some(extracted_params) = self.match_template(template_uri, uri) {
550                let resource = Arc::clone(entry.value());
551                let mut combined_params = params.clone();
552                combined_params.extend(extracted_params);
553                return Ok((resource, combined_params));
554            }
555        }
556
557        // Try matching against resource templates
558        for entry in self.templates.iter() {
559            let template = entry.value();
560            let template_uri = template.uri_template();
561            if let Some(extracted_params) = self.match_template(template_uri, uri) {
562                // Create a dynamic resource from the template
563                let dynamic_resource = DynamicTemplateResource {
564                    template: Arc::clone(template),
565                    uri: uri.to_string(),
566                };
567                let mut combined_params = params.clone();
568                combined_params.extend(extracted_params);
569                return Ok((Arc::new(dynamic_resource), combined_params));
570            }
571        }
572
573        Err(ResourceError::NotFound(uri.to_string()))
574    }
575
576    /// Match URI against template and extract parameters
577    /// Template format: "scheme://path/{param}/more/{param2}"
578    /// If the last part is a parameter, it greedily matches the rest of the URI
579    fn match_template(&self, template: &str, uri: &str) -> Option<HashMap<String, String>> {
580        let template_parts: Vec<&str> = template.split('/').collect();
581        let uri_parts: Vec<&str> = uri.split('/').collect();
582
583        // Check if template has fewer parts (only valid if last part is a parameter)
584        if template_parts.len() > uri_parts.len() {
585            return None;
586        }
587
588        let mut params = HashMap::new();
589        let last_idx = template_parts.len() - 1;
590
591        for (i, template_part) in template_parts.iter().enumerate() {
592            if template_part.starts_with('{') && template_part.ends_with('}') {
593                let param_name = &template_part[1..template_part.len() - 1];
594
595                if i == last_idx {
596                    // Last parameter: greedily match remaining URI parts
597                    let remaining = uri_parts[i..].join("/");
598                    params.insert(param_name.to_string(), remaining);
599                    break;
600                } else {
601                    // Middle parameter: match single part
602                    if i >= uri_parts.len() {
603                        return None;
604                    }
605                    params.insert(param_name.to_string(), uri_parts[i].to_string());
606                }
607            } else if i >= uri_parts.len() || template_part != &uri_parts[i] {
608                // Static part doesn't match or URI too short
609                return None;
610            }
611        }
612
613        Some(params)
614    }
615
616    // ==================== Session-Aware Methods ====================
617
618    /// List resources visible to a specific session
619    ///
620    /// Resolution order:
621    /// 1. Session overrides (replace global resources with session-specific implementations)
622    /// 2. Session extras (additional resources added to session)
623    /// 3. Global resources (filtered by hidden list and visibility predicate)
624    pub fn list_for_session(&self, session: &Session, ctx: &VisibilityContext<'_>) -> Vec<ResourceInfo> {
625        let mut resources = std::collections::HashMap::new();
626
627        // 1. Add global resources (filtered by hidden and visibility)
628        for entry in self.resources.iter() {
629            let uri = entry.key().clone();
630            if !session.is_resource_hidden(&uri) {
631                let resource = entry.value();
632                if resource.is_visible(ctx) {
633                    resources.insert(
634                        uri,
635                        ResourceInfo {
636                            uri: resource.uri().to_string(),
637                            name: resource.name().to_string(),
638                            description: resource.description().map(|s| s.to_string()),
639                            mime_type: resource.mime_type().map(|s| s.to_string()),
640                        },
641                    );
642                }
643            }
644        }
645
646        // 2. Add session extras
647        for entry in session.resource_extras().iter() {
648            let uri = entry.key().clone();
649            let resource = entry.value();
650            if resource.is_visible(ctx) {
651                resources.insert(
652                    uri,
653                    ResourceInfo {
654                        uri: resource.uri().to_string(),
655                        name: resource.name().to_string(),
656                        description: resource.description().map(|s| s.to_string()),
657                        mime_type: resource.mime_type().map(|s| s.to_string()),
658                    },
659                );
660            }
661        }
662
663        // 3. Apply session overrides
664        for entry in session.resource_overrides().iter() {
665            let uri = entry.key().clone();
666            let resource = entry.value();
667            if resource.is_visible(ctx) {
668                resources.insert(
669                    uri,
670                    ResourceInfo {
671                        uri: resource.uri().to_string(),
672                        name: resource.name().to_string(),
673                        description: resource.description().map(|s| s.to_string()),
674                        mime_type: resource.mime_type().map(|s| s.to_string()),
675                    },
676                );
677            }
678        }
679
680        resources.into_values().collect()
681    }
682
683    /// Read a resource with session context
684    ///
685    /// Resolution order:
686    /// 1. Check session overrides
687    /// 2. Check session extras
688    /// 3. Check session hidden
689    /// 4. Check visibility predicate
690    /// 5. Fall back to global registry with retry logic
691    pub async fn read_for_session(
692        &self,
693        uri: &str,
694        params: HashMap<String, String>,
695        session: &Session,
696        visibility_ctx: &VisibilityContext<'_>,
697    ) -> Result<Vec<ResourceContent>, ResourceError> {
698        // Create execution context (reuse environment from visibility context)
699        let exec_ctx = match visibility_ctx.environment {
700            Some(env) => ExecutionContext::for_resource_with_environment(params.clone(), session, env),
701            None => ExecutionContext::for_resource(params.clone(), session),
702        };
703
704        // 1. Check session override first
705        if let Some(resource) = session.get_resource_override(uri) {
706            if !resource.is_visible(visibility_ctx) {
707                return Err(ResourceError::NotFound(uri.to_string()));
708            }
709            return resource.read(exec_ctx).await;
710        }
711
712        // 2. Check session extras
713        if let Some(resource) = session.get_resource_extra(uri) {
714            if !resource.is_visible(visibility_ctx) {
715                return Err(ResourceError::NotFound(uri.to_string()));
716            }
717            return resource.read(exec_ctx).await;
718        }
719
720        // 3. Check if hidden in session
721        if session.is_resource_hidden(uri) {
722            return Err(ResourceError::NotFound(uri.to_string()));
723        }
724
725        // 4. Check global registry with visibility check
726        if let Some(resource) = self.get(uri)
727            && !resource.is_visible(visibility_ctx) {
728                return Err(ResourceError::NotFound(uri.to_string()));
729            }
730
731        // 5. Fall back to global registry with retry logic
732        self.read(uri, params, session).await
733    }
734
735    /// Get number of registered resources
736    pub fn len(&self) -> usize {
737        self.resources.len()
738    }
739
740    /// Check if manager is empty
741    pub fn is_empty(&self) -> bool {
742        self.resources.is_empty()
743    }
744}
745
746impl Default for ResourceManager {
747    fn default() -> Self {
748        Self::new()
749    }
750}
751
752#[cfg(test)]
753mod tests {
754    use super::*;
755
756    // Example static resource for testing
757    struct HelloResource;
758
759    #[async_trait]
760    impl Resource for HelloResource {
761        fn uri(&self) -> &str {
762            "test://hello"
763        }
764
765        fn name(&self) -> &str {
766            "hello"
767        }
768
769        fn description(&self) -> Option<&str> {
770            Some("Returns a greeting")
771        }
772
773        fn mime_type(&self) -> Option<&str> {
774            Some("text/plain")
775        }
776
777        async fn read(&self, _ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
778            Ok(vec![self.text_content("Hello, World!")])
779        }
780    }
781
782    // Example template resource for testing
783    struct UserResource;
784
785    #[async_trait]
786    impl Resource for UserResource {
787        fn uri(&self) -> &str {
788            "test://users/{id}"
789        }
790
791        fn name(&self) -> &str {
792            "user"
793        }
794
795        fn description(&self) -> Option<&str> {
796            Some("Returns user information")
797        }
798
799        async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
800            let id = ctx
801                .get_uri_param("id")
802                .ok_or_else(|| ResourceError::InvalidUri("Missing 'id' parameter".to_string()))?;
803
804            Ok(vec![self.text_content(&format!("User ID: {}", id))])
805        }
806    }
807
808    #[test]
809    fn test_manager_creation() {
810        let manager = ResourceManager::new();
811        assert!(manager.is_empty());
812    }
813
814    #[test]
815    fn test_resource_registration() {
816        let manager = ResourceManager::new();
817        manager.register(HelloResource);
818
819        assert_eq!(manager.len(), 1);
820        assert!(!manager.is_empty());
821    }
822
823    #[test]
824    fn test_get_resource() {
825        let manager = ResourceManager::new();
826        manager.register(HelloResource);
827
828        let resource = manager.get("test://hello");
829        assert!(resource.is_some());
830        assert_eq!(resource.unwrap().name(), "hello");
831
832        let missing = manager.get("test://nonexistent");
833        assert!(missing.is_none());
834    }
835
836    #[test]
837    fn test_list_resources() {
838        let manager = ResourceManager::new();
839        manager.register(HelloResource);
840
841        let resources = manager.list();
842        assert_eq!(resources.len(), 1);
843        assert_eq!(resources[0].uri, "test://hello");
844        assert_eq!(resources[0].name, "hello");
845        assert_eq!(resources[0].description, Some("Returns a greeting".to_string()));
846        assert_eq!(resources[0].mime_type, Some("text/plain".to_string()));
847    }
848
849    #[tokio::test]
850    async fn test_read_static_resource() {
851        let manager = ResourceManager::new();
852        manager.register(HelloResource);
853        let session = Session::new();
854
855        let result = manager.read("test://hello", HashMap::new(), &session).await.unwrap();
856        assert_eq!(result.len(), 1);
857    }
858
859    #[tokio::test]
860    async fn test_read_missing_resource() {
861        let manager = ResourceManager::new();
862        let session = Session::new();
863
864        let result = manager.read("test://nonexistent", HashMap::new(), &session).await;
865        assert!(matches!(result, Err(ResourceError::NotFound(_))));
866    }
867
868    #[tokio::test]
869    async fn test_template_matching() {
870        let manager = ResourceManager::new();
871        manager.register(UserResource);
872        let session = Session::new();
873
874        // Test template parameter extraction
875        let result = manager.read("test://users/123", HashMap::new(), &session).await.unwrap();
876        assert_eq!(result.len(), 1);
877    }
878
879    #[test]
880    fn test_template_matching_internal() {
881        let manager = ResourceManager::new();
882
883        // Test exact template matching
884        let params = manager.match_template("test://users/{id}", "test://users/123");
885        assert!(params.is_some());
886        let params = params.unwrap();
887        assert_eq!(params.get("id"), Some(&"123".to_string()));
888
889        // Test multiple parameters
890        let params = manager.match_template(
891            "test://org/{org}/repo/{repo}",
892            "test://org/myorg/repo/myrepo",
893        );
894        assert!(params.is_some());
895        let params = params.unwrap();
896        assert_eq!(params.get("org"), Some(&"myorg".to_string()));
897        assert_eq!(params.get("repo"), Some(&"myrepo".to_string()));
898
899        // Test non-matching
900        let params = manager.match_template("test://users/{id}", "test://posts/123");
901        assert!(params.is_none());
902
903        // Test greedy matching (last parameter matches remaining segments)
904        let params = manager.match_template("test://users/{id}", "test://users/123/extra");
905        assert!(params.is_some());
906        let params = params.unwrap();
907        assert_eq!(params.get("id"), Some(&"123/extra".to_string()));
908    }
909
910    #[tokio::test]
911    async fn test_resource_non_matching_template() {
912        let manager = ResourceManager::new();
913        manager.register(UserResource);
914        let session = Session::new();
915
916        // This won't match because "posts" != "users"
917        let result = manager.read("test://posts/123", HashMap::new(), &session).await;
918        assert!(matches!(result, Err(ResourceError::NotFound(_))));
919    }
920
921    #[tokio::test]
922    async fn test_resource_greedy_matching() {
923        let manager = ResourceManager::new();
924        manager.register(UserResource);
925        let session = Session::new();
926
927        // The {id} parameter greedily matches the remaining segments
928        let result = manager.read("test://users/123/extra", HashMap::new(), &session).await;
929        assert!(result.is_ok());
930        let content = result.unwrap();
931        assert_eq!(content.len(), 1);
932        // Verify the id parameter captured the greedy match
933        assert!(content[0].text.as_ref().unwrap().contains("123/extra"));
934    }
935
936    // ==================== Resource Template Tests ====================
937
938    // Example template for testing
939    struct FileTemplate;
940
941    #[async_trait]
942    impl ResourceTemplate for FileTemplate {
943        fn uri_template(&self) -> &str {
944            "file:///{path}"
945        }
946
947        fn name(&self) -> &str {
948            "project_files"
949        }
950
951        fn title(&self) -> Option<&str> {
952            Some("Project Files")
953        }
954
955        fn description(&self) -> Option<&str> {
956            Some("Access files in the project directory")
957        }
958
959        fn mime_type(&self) -> Option<&str> {
960            Some("application/octet-stream")
961        }
962
963        async fn read(&self, ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
964            let path = ctx
965                .get_uri_param("path")
966                .ok_or_else(|| ResourceError::InvalidUri("Missing 'path' parameter".to_string()))?;
967
968            Ok(vec![ResourceContent::text(
969                format!("file:///{}", path),
970                format!("Mock content for file: {}", path),
971            )])
972        }
973    }
974
975    // Template with visibility predicate
976    struct AdminOnlyTemplate;
977
978    #[async_trait]
979    impl ResourceTemplate for AdminOnlyTemplate {
980        fn uri_template(&self) -> &str {
981            "admin:///{resource}"
982        }
983
984        fn name(&self) -> &str {
985            "admin_resources"
986        }
987
988        fn is_visible(&self, ctx: &VisibilityContext) -> bool {
989            ctx.has_role("admin")
990        }
991
992        async fn read(&self, _ctx: ExecutionContext<'_>) -> Result<Vec<ResourceContent>, ResourceError> {
993            Ok(vec![ResourceContent::text("admin:///test", "Admin content")])
994        }
995    }
996
997    #[test]
998    fn test_template_registration() {
999        let manager = ResourceManager::new();
1000        manager.register_template(FileTemplate);
1001
1002        let template = manager.get_template("project_files");
1003        assert!(template.is_some());
1004        assert_eq!(template.unwrap().name(), "project_files");
1005    }
1006
1007    #[test]
1008    fn test_get_template() {
1009        let manager = ResourceManager::new();
1010        manager.register_template(FileTemplate);
1011
1012        let template = manager.get_template("project_files");
1013        assert!(template.is_some());
1014        assert_eq!(template.unwrap().uri_template(), "file:///{path}");
1015
1016        let missing = manager.get_template("nonexistent");
1017        assert!(missing.is_none());
1018    }
1019
1020    #[test]
1021    fn test_list_templates() {
1022        let manager = ResourceManager::new();
1023        manager.register_template(FileTemplate);
1024
1025        let templates = manager.list_templates();
1026        assert_eq!(templates.len(), 1);
1027        assert_eq!(templates[0].uri_template, "file:///{path}");
1028        assert_eq!(templates[0].name, "project_files");
1029        assert_eq!(templates[0].title, Some("Project Files".to_string()));
1030        assert_eq!(templates[0].description, Some("Access files in the project directory".to_string()));
1031        assert_eq!(templates[0].mime_type, Some("application/octet-stream".to_string()));
1032    }
1033
1034    #[test]
1035    fn test_list_templates_for_session_visibility() {
1036        let manager = ResourceManager::new();
1037        manager.register_template(FileTemplate);
1038        manager.register_template(AdminOnlyTemplate);
1039
1040        // Regular session without admin role
1041        let session = Session::new();
1042        let ctx = VisibilityContext::new(&session);
1043        let templates = manager.list_templates_for_session(&session, &ctx);
1044
1045        // Should only see FileTemplate, not AdminOnlyTemplate
1046        assert_eq!(templates.len(), 1);
1047        assert_eq!(templates[0].name, "project_files");
1048
1049        // Admin session
1050        let admin_session = Session::new();
1051        admin_session.set_state("roles", serde_json::json!(["admin"]));
1052        let admin_ctx = VisibilityContext::new(&admin_session);
1053        let admin_templates = manager.list_templates_for_session(&admin_session, &admin_ctx);
1054
1055        // Should see both templates
1056        assert_eq!(admin_templates.len(), 2);
1057        let names: Vec<_> = admin_templates.iter().map(|t| t.name.as_str()).collect();
1058        assert!(names.contains(&"project_files"));
1059        assert!(names.contains(&"admin_resources"));
1060    }
1061
1062    #[test]
1063    fn test_template_info_serialization() {
1064        let manager = ResourceManager::new();
1065        manager.register_template(FileTemplate);
1066
1067        let templates = manager.list_templates();
1068        let serialized = serde_json::to_value(&templates[0]).unwrap();
1069
1070        // Verify MCP-compliant field names (camelCase)
1071        assert!(serialized.get("uriTemplate").is_some());
1072        assert_eq!(serialized["uriTemplate"], "file:///{path}");
1073        assert_eq!(serialized["name"], "project_files");
1074        assert_eq!(serialized["mimeType"], "application/octet-stream");
1075
1076        // Verify snake_case fields are NOT present
1077        assert!(serialized.get("uri_template").is_none());
1078        assert!(serialized.get("mime_type").is_none());
1079    }
1080
1081    #[tokio::test]
1082    async fn test_template_read_with_uri_params() {
1083        let session = Session::new();
1084        let ctx = ExecutionContext::for_resource(
1085            vec![("path".to_string(), "src/main.rs".to_string())].into_iter().collect(),
1086            &session,
1087        );
1088
1089        let template = FileTemplate;
1090        let result = template.read(ctx).await.unwrap();
1091
1092        assert_eq!(result.len(), 1);
1093        assert_eq!(result[0].uri, "file:///src/main.rs");
1094        assert!(result[0].text.as_ref().unwrap().contains("src/main.rs"));
1095    }
1096
1097    #[tokio::test]
1098    async fn test_template_read_missing_param() {
1099        let session = Session::new();
1100        let ctx = ExecutionContext::for_resource(HashMap::new(), &session);
1101
1102        let template = FileTemplate;
1103        let result = template.read(ctx).await;
1104
1105        // Should fail with InvalidUri when path parameter is missing
1106        assert!(matches!(result, Err(ResourceError::InvalidUri(_))));
1107    }
1108}
mcp_host/registry/resources.rs

mcp_host/registry/
resources.rs
