use std::cmp;
use std::collections::BTreeSet;
use std::sync::Arc;
use anyhow::Context as _;
use anyhow::Result;
use anyhow::anyhow;
use http::Uri;
use log::warn;
use tokio::sync::mpsc::Receiver;
use tokio_util::sync::CancellationToken;
use crate::mcp::list_resources_cursor::ListResourcesCursor;
use crate::mcp::resource::Resource;
use crate::mcp::resource_content_parts::ResourceContentParts;
use crate::mcp::resource_provider::ResourceProvider;
use crate::mcp::resource_provider_list_params::ResourceProviderListParams;
use crate::mcp::resource_provider_ordered::ResourceProviderOrdered;
use crate::mcp::resource_reference::ResourceReference;
use crate::mcp::resource_template::ResourceTemplate;
struct FoundProvider<'provider> {
pub provider: &'provider ResourceProviderOrdered,
pub resource_reference: ResourceReference,
}
pub struct ResourceListAggregate {
pub providers: BTreeSet<ResourceProviderOrdered>,
}
impl ResourceListAggregate {
pub fn total(&self) -> usize {
self.providers.iter().map(|p| p.0.total()).sum()
}
pub async fn list_resources(
&self,
ListResourcesCursor { offset, per_page }: ListResourcesCursor,
) -> Result<Vec<Resource>> {
let mut resources: Vec<Resource> = vec![];
let mut to_skip = offset;
let mut to_take = per_page;
for provider in &self.providers {
if to_take < 1 {
break;
}
let provider_total = provider.0.total();
if provider_total < to_skip {
to_skip -= provider_total;
} else {
let provider_available = provider_total - to_skip;
let provider_offset = to_skip;
let provider_to_take = cmp::min(provider_available, to_take);
let mut taken_resources = provider
.0
.list_resources(ResourceProviderListParams {
limit: provider_to_take,
offset: provider_offset,
})
.await?;
to_skip = 0;
to_take -= taken_resources.len();
resources.append(&mut taken_resources);
}
}
Ok(resources)
}
pub async fn read_resource_contents(&self, uri: &str) -> Result<Option<ResourceContentParts>> {
let FoundProvider {
provider,
resource_reference,
} = self.must_get_provider_for_uri(uri)?;
provider.0.read_resource_contents(resource_reference).await
}
pub async fn read_resources_templates_list(&self) -> Result<Vec<ResourceTemplate>> {
Ok(self
.providers
.iter()
.map(|provider| provider.0.resource_template())
.collect())
}
pub async fn subscribe(
&self,
cancellation_token: CancellationToken,
uri: &str,
) -> Result<Option<Receiver<ResourceContentParts>>> {
let FoundProvider {
provider,
resource_reference,
} = self.must_get_provider_for_uri(uri)?;
provider
.0
.clone()
.subscribe(cancellation_token, resource_reference)
.await
}
fn must_get_provider_for_uri<'provider>(
&'provider self,
uri: &str,
) -> Result<FoundProvider<'provider>> {
let parsed_uri: Uri = uri
.try_into()
.map_err(|err| anyhow!("{err:#?}"))
.context(format!("Unable to parse resource URI string: {uri}"))?;
let resource_reference: ResourceReference = parsed_uri.try_into()?;
for provider in &self.providers {
if provider.0.can_handle(&resource_reference) {
return Ok(FoundProvider {
provider,
resource_reference,
});
}
}
let message = anyhow!("There is no provider that can handle resource: {uri}");
warn!("{message}");
Err(message)
}
}
impl From<Vec<Arc<dyn ResourceProvider>>> for ResourceListAggregate {
fn from(providers: Vec<Arc<dyn ResourceProvider>>) -> Self {
Self {
providers: providers.into_iter().map(ResourceProviderOrdered).collect(),
}
}
}
#[cfg(test)]
mod tests {
use async_trait::async_trait;
use tokio::sync::Notify;
use super::*;
use crate::mcp::resource_provider::ResourceProvider;
use crate::mcp::resource_template_provider::ResourceTemplateProvider;
struct TestResourceProvider {
pub class: String,
pub total: usize,
}
impl ResourceTemplateProvider for TestResourceProvider {
fn mime_type(&self) -> String {
"x+foo/bar".to_string()
}
fn resource_class(&self) -> String {
self.class.clone()
}
fn resource_scheme(&self) -> String {
"foo".to_string()
}
}
#[async_trait]
impl ResourceProvider for TestResourceProvider {
async fn list_resources(
&self,
params: ResourceProviderListParams,
) -> Result<Vec<Resource>> {
let resource_class = self.resource_class();
let mut resources: Vec<Resource> = Vec::new();
for i in params.range() {
resources.push(Resource {
description: format!("description_p{resource_class}_r{i}"),
name: format!("name_p{resource_class}_r{i}"),
title: format!("title_p{resource_class}_r{i}"),
uri: format!("uri_p{resource_class}_r{i}"),
});
}
Ok(resources)
}
async fn read_resource_contents(
&self,
_: ResourceReference,
) -> Result<Option<ResourceContentParts>> {
Ok(None)
}
async fn resource_update_notifier(
self: Arc<Self>,
_: CancellationToken,
_: ResourceReference,
) -> Result<Option<Arc<Notify>>> {
unimplemented!();
}
fn total(&self) -> usize {
self.total
}
}
#[tokio::test]
async fn test_pagination_first_page() -> Result<()> {
let total = 4;
let per_page = 2;
let class = "1";
let aggregate: ResourceListAggregate = vec![Arc::new(TestResourceProvider {
class: class.to_string(),
total,
}) as Arc<dyn ResourceProvider>]
.into();
let first_page = aggregate
.list_resources(ListResourcesCursor {
offset: 0,
per_page,
})
.await?;
assert_eq!(first_page.len(), per_page);
assert_eq!(
first_page.first().unwrap().name,
format!("name_p{class}_r0")
);
assert_eq!(first_page.get(1).unwrap().name, format!("name_p{class}_r1"));
Ok(())
}
#[tokio::test]
async fn test_pagination_second_page() -> Result<()> {
let total = 4;
let per_page = 2;
let class = "1";
let aggregate: ResourceListAggregate = vec![Arc::new(TestResourceProvider {
class: class.to_string(),
total,
}) as Arc<dyn ResourceProvider>]
.into();
let second_page = aggregate
.list_resources(ListResourcesCursor {
offset: per_page,
per_page,
})
.await?;
assert_eq!(second_page.len(), per_page);
assert_eq!(
second_page.first().unwrap().name,
format!("name_p{class}_r{per_page}")
);
assert_eq!(
second_page.get(1).unwrap().name,
format!("name_p{class}_r{}", per_page + 1)
);
Ok(())
}
#[tokio::test]
async fn test_pagination_beyond_last_page() -> Result<()> {
let total = 4;
let per_page = 2;
let aggregate: ResourceListAggregate = vec![Arc::new(TestResourceProvider {
class: "1".to_string(),
total,
}) as Arc<dyn ResourceProvider>]
.into();
let beyond = aggregate
.list_resources(ListResourcesCursor {
offset: total,
per_page,
})
.await?;
assert_eq!(beyond.len(), 0);
Ok(())
}
#[tokio::test]
async fn test_pagination_zero_per_page_returns_empty() -> Result<()> {
let total = 4;
let aggregate: ResourceListAggregate = vec![Arc::new(TestResourceProvider {
class: "1".to_string(),
total,
}) as Arc<dyn ResourceProvider>]
.into();
let result = aggregate
.list_resources(ListResourcesCursor {
offset: 0,
per_page: 0,
})
.await?;
assert_eq!(result.len(), 0);
Ok(())
}
#[tokio::test]
async fn test_pagination_offset_overflow_returns_empty() -> Result<()> {
let total = 4;
let aggregate: ResourceListAggregate = vec![Arc::new(TestResourceProvider {
class: "1".to_string(),
total,
}) as Arc<dyn ResourceProvider>]
.into();
let result = aggregate
.list_resources(ListResourcesCursor {
offset: usize::MAX,
per_page: 2,
})
.await?;
assert_eq!(result.len(), 0);
Ok(())
}
#[tokio::test]
async fn test_list_resources() -> Result<()> {
let resource_list_aggregate: ResourceListAggregate = vec![
Arc::new(TestResourceProvider {
class: "1".to_string(),
total: 3,
}) as Arc<dyn ResourceProvider>,
Arc::new(TestResourceProvider {
class: "2".to_string(),
total: 2,
}) as Arc<dyn ResourceProvider>,
]
.try_into()?;
let resources_batch_1 = resource_list_aggregate
.list_resources(ListResourcesCursor {
offset: 0,
per_page: 2,
})
.await?;
assert_eq!(resources_batch_1.len(), 2);
assert_eq!(
resources_batch_1.first().unwrap().name,
"name_p1_r0".to_string()
);
assert_eq!(
resources_batch_1.get(1).unwrap().name,
"name_p1_r1".to_string()
);
let resources_batch_2 = resource_list_aggregate
.list_resources(ListResourcesCursor {
offset: 2,
per_page: 5,
})
.await?;
assert_eq!(resources_batch_2.len(), 3);
assert_eq!(
resources_batch_2.first().unwrap().name,
"name_p1_r2".to_string()
);
assert_eq!(
resources_batch_2.get(1).unwrap().name,
"name_p2_r0".to_string()
);
assert_eq!(
resources_batch_2.get(2).unwrap().name,
"name_p2_r1".to_string()
);
Ok(())
}
fn aggregate_with_classes(classes: &[&str]) -> ResourceListAggregate {
classes
.iter()
.map(|class| {
Arc::new(TestResourceProvider {
class: class.to_string(),
total: 1,
}) as Arc<dyn ResourceProvider>
})
.collect::<Vec<_>>()
.into()
}
#[test]
fn total_sums_every_provider_total() {
let aggregate: ResourceListAggregate = vec![
Arc::new(TestResourceProvider {
class: "1".to_string(),
total: 3,
}) as Arc<dyn ResourceProvider>,
Arc::new(TestResourceProvider {
class: "2".to_string(),
total: 2,
}) as Arc<dyn ResourceProvider>,
]
.into();
assert_eq!(aggregate.total(), 5);
}
#[tokio::test]
async fn templates_list_has_one_entry_per_provider() -> Result<()> {
let templates = aggregate_with_classes(&["1", "2"])
.read_resources_templates_list()
.await?;
assert_eq!(templates.len(), 2);
Ok(())
}
#[tokio::test]
async fn read_resource_contents_routes_to_matching_provider() -> Result<()> {
let contents = aggregate_with_classes(&["1"])
.read_resource_contents("foo://1/example")
.await?;
assert!(contents.is_none());
Ok(())
}
#[tokio::test]
async fn read_resource_contents_errors_for_malformed_uri() {
assert!(
aggregate_with_classes(&["1"])
.read_resource_contents("not a uri")
.await
.is_err()
);
}
#[tokio::test]
async fn read_resource_contents_errors_when_no_provider_handles_scheme() {
assert!(
aggregate_with_classes(&["1"])
.read_resource_contents("bar://1/example")
.await
.is_err()
);
}
}