use crate::server::cancellation::RequestHandlerExtra;
use crate::shared::uri_template::UriTemplate;
use crate::types::{ReadResourceResult, ResourceTemplate};
use crate::Result;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct UriParams {
variables: HashMap<String, String>,
}
impl UriParams {
pub fn new(variables: HashMap<String, String>) -> Self {
Self { variables }
}
pub fn get(&self, name: &str) -> Option<&String> {
self.variables.get(name)
}
pub fn get_or(&self, name: &str, default: &str) -> String {
self.variables
.get(name)
.cloned()
.unwrap_or_else(|| default.to_string())
}
pub fn contains(&self, name: &str) -> bool {
self.variables.contains_key(name)
}
pub fn keys(&self) -> Vec<&String> {
self.variables.keys().collect()
}
pub fn values(&self) -> Vec<&String> {
self.variables.values().collect()
}
pub fn len(&self) -> usize {
self.variables.len()
}
pub fn is_empty(&self) -> bool {
self.variables.is_empty()
}
pub fn into_inner(self) -> HashMap<String, String> {
self.variables
}
}
#[derive(Debug, Clone)]
pub struct RequestContext {
pub session_id: Option<String>,
pub metadata: HashMap<String, String>,
pub extra: RequestHandlerExtra,
}
impl RequestContext {
pub fn new(extra: RequestHandlerExtra) -> Self {
Self {
session_id: None,
metadata: HashMap::new(),
extra,
}
}
pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
pub fn get_metadata(&self, key: &str) -> Option<&String> {
self.metadata.get(key)
}
}
#[async_trait]
pub trait DynamicResourceProvider: Send + Sync {
fn templates(&self) -> Vec<ResourceTemplate>;
async fn fetch(
&self,
uri: &str,
params: UriParams,
context: RequestContext,
) -> Result<ReadResourceResult>;
fn priority(&self) -> i32 {
50
}
}
#[derive(Clone)]
pub(crate) struct MatchedProvider {
pub(crate) provider: Arc<dyn DynamicResourceProvider>,
pub(crate) params: UriParams,
#[allow(dead_code)] pub(crate) template: ResourceTemplate,
}
struct ParsedTemplate {
template: ResourceTemplate,
parsed: UriTemplate,
}
pub(crate) struct ResourceRouter {
providers: Vec<Arc<dyn DynamicResourceProvider>>,
parsed_templates: Vec<(Arc<dyn DynamicResourceProvider>, ParsedTemplate)>,
}
impl ResourceRouter {
pub(crate) fn new() -> Self {
Self {
providers: Vec::new(),
parsed_templates: Vec::new(),
}
}
pub(crate) fn add_provider(&mut self, provider: Arc<dyn DynamicResourceProvider>) {
for template in provider.templates() {
if let Ok(parsed) = UriTemplate::new(&template.uri_template) {
self.parsed_templates.push((
Arc::clone(&provider),
ParsedTemplate {
template: template.clone(),
parsed,
},
));
} else {
tracing::warn!("Failed to parse URI template: {}", template.uri_template);
}
}
self.providers.push(provider);
self.providers.sort_by_key(|p| p.priority());
self.parsed_templates
.sort_by_key(|(provider, _)| provider.priority());
}
pub(crate) fn match_uri(&self, uri: &str) -> Option<MatchedProvider> {
for (provider, parsed_template) in &self.parsed_templates {
if let Some(variables) = parsed_template.parsed.match_uri(uri) {
return Some(MatchedProvider {
provider: Arc::clone(provider),
params: UriParams::new(variables),
template: parsed_template.template.clone(),
});
}
}
None
}
pub(crate) fn all_templates(&self) -> Vec<ResourceTemplate> {
self.providers.iter().flat_map(|p| p.templates()).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::Content;
#[test]
fn test_uri_params() {
let mut vars = HashMap::new();
vars.insert("id".to_string(), "123".to_string());
vars.insert("type".to_string(), "schema".to_string());
let params = UriParams::new(vars);
assert_eq!(params.get("id"), Some(&"123".to_string()));
assert_eq!(params.get("type"), Some(&"schema".to_string()));
assert_eq!(params.get("missing"), None);
assert_eq!(params.get_or("missing", "default"), "default");
assert!(params.contains("id"));
assert!(!params.contains("missing"));
assert_eq!(params.len(), 2);
assert!(!params.is_empty());
}
#[test]
fn test_request_context() {
use tokio_util::sync::CancellationToken;
let extra = RequestHandlerExtra::new("test-request".to_string(), CancellationToken::new());
let context = RequestContext::new(extra)
.with_session_id("session-123")
.with_metadata("user", "alice");
assert_eq!(context.session_id, Some("session-123".to_string()));
assert_eq!(context.get_metadata("user"), Some(&"alice".to_string()));
assert_eq!(context.get_metadata("missing"), None);
}
struct TestProvider {
priority: i32,
}
#[async_trait]
impl DynamicResourceProvider for TestProvider {
fn templates(&self) -> Vec<ResourceTemplate> {
vec![ResourceTemplate::new("test://{id}", "Test Resource")
.with_description("A test resource")
.with_mime_type("text/plain")]
}
fn priority(&self) -> i32 {
self.priority
}
async fn fetch(
&self,
_uri: &str,
params: UriParams,
_context: RequestContext,
) -> Result<ReadResourceResult> {
let id = params.get("id").unwrap();
Ok(ReadResourceResult::new(vec![Content::text(format!(
"Resource {}",
id
))]))
}
}
#[test]
fn test_resource_router_matching() {
let mut router = ResourceRouter::new();
router.add_provider(Arc::new(TestProvider { priority: 50 }));
let matched = router.match_uri("test://123");
assert!(matched.is_some());
let matched = matched.unwrap();
assert_eq!(matched.params.get("id"), Some(&"123".to_string()));
assert_eq!(matched.template.name, "Test Resource");
let no_match = router.match_uri("other://123");
assert!(no_match.is_none());
}
#[test]
fn test_resource_router_priority() {
let mut router = ResourceRouter::new();
router.add_provider(Arc::new(TestProvider { priority: 100 }));
router.add_provider(Arc::new(TestProvider { priority: 10 }));
router.add_provider(Arc::new(TestProvider { priority: 50 }));
assert_eq!(router.providers[0].priority(), 10);
assert_eq!(router.providers[1].priority(), 50);
assert_eq!(router.providers[2].priority(), 100);
}
#[tokio::test]
async fn test_provider_fetch() {
use tokio_util::sync::CancellationToken;
let provider = TestProvider { priority: 50 };
let mut vars = HashMap::new();
vars.insert("id".to_string(), "456".to_string());
let params = UriParams::new(vars);
let extra = RequestHandlerExtra::new("test-request".to_string(), CancellationToken::new());
let context = RequestContext::new(extra);
let result = provider.fetch("test://456", params, context).await.unwrap();
match &result.contents[0] {
Content::Text { text } => assert_eq!(text, "Resource 456"),
_ => panic!("Expected text content"),
}
}
}