use std::collections::HashMap;
use std::convert::Infallible;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use pin_project_lite::pin_project;
use tower::util::BoxCloneService;
use tower_service::Service;
use crate::context::RequestContext;
use crate::error::{Error, Result};
use crate::protocol::{
ContentAnnotations, ReadResourceResult, ResourceContent, ResourceDefinition,
ResourceTemplateDefinition, ToolIcon,
};
#[derive(Debug, Clone)]
pub struct ResourceRequest {
pub ctx: RequestContext,
pub uri: String,
}
impl ResourceRequest {
pub fn new(ctx: RequestContext, uri: String) -> Self {
Self { ctx, uri }
}
}
pub type BoxResourceService = BoxCloneService<ResourceRequest, ReadResourceResult, Infallible>;
#[doc(hidden)]
pub struct ResourceCatchError<S> {
inner: S,
}
impl<S> ResourceCatchError<S> {
pub fn new(inner: S) -> Self {
Self { inner }
}
}
impl<S: Clone> Clone for ResourceCatchError<S> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
impl<S: fmt::Debug> fmt::Debug for ResourceCatchError<S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ResourceCatchError")
.field("inner", &self.inner)
.finish()
}
}
pin_project! {
#[doc(hidden)]
pub struct ResourceCatchErrorFuture<F> {
#[pin]
inner: F,
uri: Option<String>,
}
}
impl<F, E> Future for ResourceCatchErrorFuture<F>
where
F: Future<Output = std::result::Result<ReadResourceResult, E>>,
E: fmt::Display,
{
type Output = std::result::Result<ReadResourceResult, Infallible>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
match this.inner.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(result)) => Poll::Ready(Ok(result)),
Poll::Ready(Err(err)) => {
let uri = this.uri.take().unwrap_or_default();
Poll::Ready(Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: Some("text/plain".to_string()),
text: Some(format!("Error reading resource: {}", err)),
blob: None,
meta: None,
}],
meta: None,
}))
}
}
}
}
impl<S> Service<ResourceRequest> for ResourceCatchError<S>
where
S: Service<ResourceRequest, Response = ReadResourceResult> + Clone + Send + 'static,
S::Error: fmt::Display + Send,
S::Future: Send,
{
type Response = ReadResourceResult;
type Error = Infallible;
type Future = ResourceCatchErrorFuture<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
match self.inner.poll_ready(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(_)) => Poll::Ready(Ok(())),
Poll::Pending => Poll::Pending,
}
}
fn call(&mut self, req: ResourceRequest) -> Self::Future {
let uri = req.uri.clone();
let fut = self.inner.call(req);
ResourceCatchErrorFuture {
inner: fut,
uri: Some(uri),
}
}
}
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub trait ResourceHandler: Send + Sync {
fn read(&self) -> BoxFuture<'_, Result<ReadResourceResult>>;
fn read_with_context(&self, _ctx: RequestContext) -> BoxFuture<'_, Result<ReadResourceResult>> {
self.read()
}
fn uses_context(&self) -> bool {
false
}
}
struct ResourceHandlerService<H> {
handler: Arc<H>,
}
impl<H> ResourceHandlerService<H> {
fn new(handler: H) -> Self {
Self {
handler: Arc::new(handler),
}
}
}
impl<H> Clone for ResourceHandlerService<H> {
fn clone(&self) -> Self {
Self {
handler: self.handler.clone(),
}
}
}
impl<H> fmt::Debug for ResourceHandlerService<H> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ResourceHandlerService")
.finish_non_exhaustive()
}
}
impl<H> Service<ResourceRequest> for ResourceHandlerService<H>
where
H: ResourceHandler + 'static,
{
type Response = ReadResourceResult;
type Error = Error;
type Future =
Pin<Box<dyn Future<Output = std::result::Result<ReadResourceResult, Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: ResourceRequest) -> Self::Future {
let handler = self.handler.clone();
Box::pin(async move { handler.read_with_context(req.ctx).await })
}
}
pub struct Resource {
pub uri: String,
pub name: String,
pub title: Option<String>,
pub description: Option<String>,
pub mime_type: Option<String>,
pub icons: Option<Vec<ToolIcon>>,
pub size: Option<u64>,
pub annotations: Option<ContentAnnotations>,
service: BoxResourceService,
}
impl Clone for Resource {
fn clone(&self) -> Self {
Self {
uri: self.uri.clone(),
name: self.name.clone(),
title: self.title.clone(),
description: self.description.clone(),
mime_type: self.mime_type.clone(),
icons: self.icons.clone(),
size: self.size,
annotations: self.annotations.clone(),
service: self.service.clone(),
}
}
}
impl std::fmt::Debug for Resource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Resource")
.field("uri", &self.uri)
.field("name", &self.name)
.field("title", &self.title)
.field("description", &self.description)
.field("mime_type", &self.mime_type)
.field("icons", &self.icons)
.field("size", &self.size)
.field("annotations", &self.annotations)
.finish_non_exhaustive()
}
}
unsafe impl Send for Resource {}
unsafe impl Sync for Resource {}
impl Resource {
pub fn builder(uri: impl Into<String>) -> ResourceBuilder {
ResourceBuilder::new(uri)
}
pub fn definition(&self) -> ResourceDefinition {
ResourceDefinition {
uri: self.uri.clone(),
name: self.name.clone(),
title: self.title.clone(),
description: self.description.clone(),
mime_type: self.mime_type.clone(),
icons: self.icons.clone(),
size: self.size,
annotations: self.annotations.clone(),
meta: None,
}
}
pub fn read(&self) -> BoxFuture<'static, ReadResourceResult> {
let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
self.read_with_context(ctx)
}
pub fn read_with_context(&self, ctx: RequestContext) -> BoxFuture<'static, ReadResourceResult> {
use tower::ServiceExt;
let service = self.service.clone();
let uri = self.uri.clone();
Box::pin(async move {
service
.oneshot(ResourceRequest::new(ctx, uri))
.await
.unwrap()
})
}
#[allow(clippy::too_many_arguments)]
fn from_handler<H: ResourceHandler + 'static>(
uri: String,
name: String,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
handler: H,
) -> Self {
let handler_service = ResourceHandlerService::new(handler);
let catch_error = ResourceCatchError::new(handler_service);
let service = BoxCloneService::new(catch_error);
Self {
uri,
name,
title,
description,
mime_type,
icons,
size,
annotations,
service,
}
}
}
pub struct ResourceBuilder {
uri: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
}
impl ResourceBuilder {
pub fn new(uri: impl Into<String>) -> Self {
Self {
uri: uri.into(),
name: None,
title: None,
description: None,
mime_type: None,
icons: None,
size: None,
annotations: None,
}
}
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn title(mut self, title: impl Into<String>) -> Self {
self.title = Some(title.into());
self
}
pub fn description(mut self, description: impl Into<String>) -> Self {
self.description = Some(description.into());
self
}
pub fn mime_type(mut self, mime_type: impl Into<String>) -> Self {
self.mime_type = Some(mime_type.into());
self
}
pub fn icon(mut self, src: impl Into<String>) -> Self {
self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
src: src.into(),
mime_type: None,
sizes: None,
theme: None,
});
self
}
pub fn icon_with_meta(
mut self,
src: impl Into<String>,
mime_type: Option<String>,
sizes: Option<Vec<String>>,
) -> Self {
self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
src: src.into(),
mime_type,
sizes,
theme: None,
});
self
}
pub fn size(mut self, size: u64) -> Self {
self.size = Some(size);
self
}
pub fn annotations(mut self, annotations: ContentAnnotations) -> Self {
self.annotations = Some(annotations);
self
}
pub fn handler<F, Fut>(self, handler: F) -> ResourceBuilderWithHandler<F>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
ResourceBuilderWithHandler {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler,
}
}
pub fn handler_with_context<F, Fut>(self, handler: F) -> ResourceBuilderWithContextHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
ResourceBuilderWithContextHandler {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler,
}
}
pub fn text(self, content: impl Into<String>) -> Resource {
let uri = self.uri.clone();
let content = content.into();
let mime_type = self.mime_type.clone();
self.handler(move || {
let uri = uri.clone();
let content = content.clone();
let mime_type = mime_type.clone();
async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type,
text: Some(content),
blob: None,
meta: None,
}],
meta: None,
})
}
})
.build()
}
pub fn json(mut self, value: serde_json::Value) -> Resource {
let uri = self.uri.clone();
self.mime_type = Some("application/json".to_string());
let text = serde_json::to_string_pretty(&value).unwrap_or_else(|_| "{}".to_string());
self.handler(move || {
let uri = uri.clone();
let text = text.clone();
async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: Some("application/json".to_string()),
text: Some(text),
blob: None,
meta: None,
}],
meta: None,
})
}
})
.build()
}
}
#[doc(hidden)]
pub struct ResourceBuilderWithHandler<F> {
uri: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
handler: F,
}
impl<F, Fut> ResourceBuilderWithHandler<F>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
Resource::from_handler(
self.uri,
name,
self.title,
self.description,
self.mime_type,
self.icons,
self.size,
self.annotations,
FnHandler {
handler: self.handler,
},
)
}
pub fn layer<L>(self, layer: L) -> ResourceBuilderWithLayer<F, L> {
ResourceBuilderWithLayer {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler: self.handler,
layer,
}
}
}
#[doc(hidden)]
pub struct ResourceBuilderWithLayer<F, L> {
uri: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
handler: F,
layer: L,
}
#[allow(private_bounds)]
impl<F, Fut, L> ResourceBuilderWithLayer<F, L>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
L: tower::Layer<ResourceHandlerService<FnHandler<F>>> + Clone + Send + Sync + 'static,
L::Service: Service<ResourceRequest, Response = ReadResourceResult> + Clone + Send + 'static,
<L::Service as Service<ResourceRequest>>::Error: fmt::Display + Send,
<L::Service as Service<ResourceRequest>>::Future: Send,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
let handler_service = ResourceHandlerService::new(FnHandler {
handler: self.handler,
});
let layered = self.layer.layer(handler_service);
let catch_error = ResourceCatchError::new(layered);
let service = BoxCloneService::new(catch_error);
Resource {
uri: self.uri,
name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
service,
}
}
pub fn layer<L2>(
self,
layer: L2,
) -> ResourceBuilderWithLayer<F, tower::layer::util::Stack<L2, L>> {
ResourceBuilderWithLayer {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler: self.handler,
layer: tower::layer::util::Stack::new(layer, self.layer),
}
}
}
#[doc(hidden)]
pub struct ResourceBuilderWithContextHandler<F> {
uri: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
handler: F,
}
impl<F, Fut> ResourceBuilderWithContextHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
Resource::from_handler(
self.uri,
name,
self.title,
self.description,
self.mime_type,
self.icons,
self.size,
self.annotations,
ContextAwareHandler {
handler: self.handler,
},
)
}
pub fn layer<L>(self, layer: L) -> ResourceBuilderWithContextLayer<F, L> {
ResourceBuilderWithContextLayer {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler: self.handler,
layer,
}
}
}
#[doc(hidden)]
pub struct ResourceBuilderWithContextLayer<F, L> {
uri: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
size: Option<u64>,
annotations: Option<ContentAnnotations>,
handler: F,
layer: L,
}
#[allow(private_bounds)]
impl<F, Fut, L> ResourceBuilderWithContextLayer<F, L>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
L: tower::Layer<ResourceHandlerService<ContextAwareHandler<F>>> + Clone + Send + Sync + 'static,
L::Service: Service<ResourceRequest, Response = ReadResourceResult> + Clone + Send + 'static,
<L::Service as Service<ResourceRequest>>::Error: fmt::Display + Send,
<L::Service as Service<ResourceRequest>>::Future: Send,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
let handler_service = ResourceHandlerService::new(ContextAwareHandler {
handler: self.handler,
});
let layered = self.layer.layer(handler_service);
let catch_error = ResourceCatchError::new(layered);
let service = BoxCloneService::new(catch_error);
Resource {
uri: self.uri,
name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
service,
}
}
pub fn layer<L2>(
self,
layer: L2,
) -> ResourceBuilderWithContextLayer<F, tower::layer::util::Stack<L2, L>> {
ResourceBuilderWithContextLayer {
uri: self.uri,
name: self.name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
handler: self.handler,
layer: tower::layer::util::Stack::new(layer, self.layer),
}
}
}
struct FnHandler<F> {
handler: F,
}
impl<F, Fut> ResourceHandler for FnHandler<F>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
fn read(&self) -> BoxFuture<'_, Result<ReadResourceResult>> {
Box::pin((self.handler)())
}
}
struct ContextAwareHandler<F> {
handler: F,
}
impl<F, Fut> ResourceHandler for ContextAwareHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
fn read(&self) -> BoxFuture<'_, Result<ReadResourceResult>> {
let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
self.read_with_context(ctx)
}
fn read_with_context(&self, ctx: RequestContext) -> BoxFuture<'_, Result<ReadResourceResult>> {
Box::pin((self.handler)(ctx))
}
fn uses_context(&self) -> bool {
true
}
}
pub trait McpResource: Send + Sync + 'static {
const URI: &'static str;
const NAME: &'static str;
const DESCRIPTION: Option<&'static str> = None;
const MIME_TYPE: Option<&'static str> = None;
fn read(&self) -> impl Future<Output = Result<ReadResourceResult>> + Send;
fn into_resource(self) -> Resource
where
Self: Sized,
{
let resource = Arc::new(self);
Resource::from_handler(
Self::URI.to_string(),
Self::NAME.to_string(),
None,
Self::DESCRIPTION.map(|s| s.to_string()),
Self::MIME_TYPE.map(|s| s.to_string()),
None,
None,
None,
McpResourceHandler { resource },
)
}
}
struct McpResourceHandler<T: McpResource> {
resource: Arc<T>,
}
impl<T: McpResource> ResourceHandler for McpResourceHandler<T> {
fn read(&self) -> BoxFuture<'_, Result<ReadResourceResult>> {
let resource = self.resource.clone();
Box::pin(async move { resource.read().await })
}
}
pub trait ResourceTemplateHandler: Send + Sync {
fn read(
&self,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<ReadResourceResult>>;
}
pub struct ResourceTemplate {
pub uri_template: String,
pub name: String,
pub title: Option<String>,
pub description: Option<String>,
pub mime_type: Option<String>,
pub icons: Option<Vec<ToolIcon>>,
pub annotations: Option<ContentAnnotations>,
pattern: regex::Regex,
variables: Vec<String>,
handler: Arc<dyn ResourceTemplateHandler>,
}
impl Clone for ResourceTemplate {
fn clone(&self) -> Self {
Self {
uri_template: self.uri_template.clone(),
name: self.name.clone(),
title: self.title.clone(),
description: self.description.clone(),
mime_type: self.mime_type.clone(),
icons: self.icons.clone(),
annotations: self.annotations.clone(),
pattern: self.pattern.clone(),
variables: self.variables.clone(),
handler: self.handler.clone(),
}
}
}
impl std::fmt::Debug for ResourceTemplate {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ResourceTemplate")
.field("uri_template", &self.uri_template)
.field("name", &self.name)
.field("title", &self.title)
.field("description", &self.description)
.field("mime_type", &self.mime_type)
.field("icons", &self.icons)
.field("variables", &self.variables)
.finish_non_exhaustive()
}
}
impl ResourceTemplate {
pub fn builder(uri_template: impl Into<String>) -> ResourceTemplateBuilder {
ResourceTemplateBuilder::new(uri_template)
}
pub fn definition(&self) -> ResourceTemplateDefinition {
ResourceTemplateDefinition {
uri_template: self.uri_template.clone(),
name: self.name.clone(),
title: self.title.clone(),
description: self.description.clone(),
mime_type: self.mime_type.clone(),
icons: self.icons.clone(),
annotations: self.annotations.clone(),
arguments: Vec::new(),
meta: None,
}
}
pub fn match_uri(&self, uri: &str) -> Option<HashMap<String, String>> {
self.pattern.captures(uri).map(|caps| {
self.variables
.iter()
.enumerate()
.filter_map(|(i, name)| {
caps.get(i + 1)
.map(|m| (name.clone(), m.as_str().to_string()))
})
.collect()
})
}
pub fn read(
&self,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<ReadResourceResult>> {
self.handler.read(uri, variables)
}
}
pub struct ResourceTemplateBuilder {
uri_template: String,
name: Option<String>,
title: Option<String>,
description: Option<String>,
mime_type: Option<String>,
icons: Option<Vec<ToolIcon>>,
annotations: Option<ContentAnnotations>,
}
impl ResourceTemplateBuilder {
pub fn new(uri_template: impl Into<String>) -> Self {
Self {
uri_template: uri_template.into(),
name: None,
title: None,
description: None,
mime_type: None,
icons: None,
annotations: None,
}
}
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn title(mut self, title: impl Into<String>) -> Self {
self.title = Some(title.into());
self
}
pub fn description(mut self, description: impl Into<String>) -> Self {
self.description = Some(description.into());
self
}
pub fn mime_type(mut self, mime_type: impl Into<String>) -> Self {
self.mime_type = Some(mime_type.into());
self
}
pub fn icon(mut self, src: impl Into<String>) -> Self {
self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
src: src.into(),
mime_type: None,
sizes: None,
theme: None,
});
self
}
pub fn icon_with_meta(
mut self,
src: impl Into<String>,
mime_type: Option<String>,
sizes: Option<Vec<String>>,
) -> Self {
self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
src: src.into(),
mime_type,
sizes,
theme: None,
});
self
}
pub fn annotations(mut self, annotations: ContentAnnotations) -> Self {
self.annotations = Some(annotations);
self
}
pub fn handler<F, Fut>(self, handler: F) -> ResourceTemplate
where
F: Fn(String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
self.try_handler(handler).unwrap_or_else(|e| {
panic!("Invalid URI template: {e}");
})
}
pub fn try_handler<F, Fut>(self, handler: F) -> std::result::Result<ResourceTemplate, Error>
where
F: Fn(String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
let (pattern, variables) = compile_uri_template(&self.uri_template)?;
let name = self.name.unwrap_or_else(|| self.uri_template.clone());
Ok(ResourceTemplate {
uri_template: self.uri_template,
name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
annotations: self.annotations,
pattern,
variables,
handler: Arc::new(FnTemplateHandler { handler }),
})
}
}
struct FnTemplateHandler<F> {
handler: F,
}
impl<F, Fut> ResourceTemplateHandler for FnTemplateHandler<F>
where
F: Fn(String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ReadResourceResult>> + Send + 'static,
{
fn read(
&self,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<ReadResourceResult>> {
let uri = uri.to_string();
Box::pin((self.handler)(uri, variables))
}
}
fn compile_uri_template(template: &str) -> std::result::Result<(regex::Regex, Vec<String>), Error> {
let mut pattern = String::from("^");
let mut variables = Vec::new();
let mut chars = template.chars().peekable();
while let Some(c) = chars.next() {
if c == '{' {
let is_reserved = chars.peek() == Some(&'+');
if is_reserved {
chars.next();
}
let var_name: String = chars.by_ref().take_while(|&c| c != '}').collect();
variables.push(var_name);
if is_reserved {
pattern.push_str("(.+)");
} else {
pattern.push_str("([^/]+)");
}
} else {
match c {
'.' | '+' | '*' | '?' | '^' | '$' | '(' | ')' | '[' | ']' | '{' | '}' | '|'
| '\\' => {
pattern.push('\\');
pattern.push(c);
}
_ => pattern.push(c),
}
}
}
pattern.push('$');
let regex = regex::Regex::new(&pattern)
.map_err(|e| Error::Internal(format!("Invalid URI template '{}': {}", template, e)))?;
Ok((regex, variables))
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use tower::timeout::TimeoutLayer;
#[tokio::test]
async fn test_builder_resource() {
let resource = ResourceBuilder::new("file:///test.txt")
.name("Test File")
.description("A test file")
.text("Hello, World!");
assert_eq!(resource.uri, "file:///test.txt");
assert_eq!(resource.name, "Test File");
assert_eq!(resource.description.as_deref(), Some("A test file"));
let result = resource.read().await;
assert_eq!(result.contents.len(), 1);
assert_eq!(result.contents[0].text.as_deref(), Some("Hello, World!"));
}
#[tokio::test]
async fn test_json_resource() {
let resource = ResourceBuilder::new("file:///config.json")
.name("Config")
.json(serde_json::json!({"key": "value"}));
assert_eq!(resource.mime_type.as_deref(), Some("application/json"));
let result = resource.read().await;
assert!(result.contents[0].text.as_ref().unwrap().contains("key"));
}
#[tokio::test]
async fn test_handler_resource() {
let resource = ResourceBuilder::new("memory://counter")
.name("Counter")
.handler(|| async {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: "memory://counter".to_string(),
mime_type: Some("text/plain".to_string()),
text: Some("42".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
})
.build();
let result = resource.read().await;
assert_eq!(result.contents[0].text.as_deref(), Some("42"));
}
#[tokio::test]
async fn test_handler_resource_with_layer() {
let resource = ResourceBuilder::new("file:///with-timeout.txt")
.name("Resource with Timeout")
.handler(|| async {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: "file:///with-timeout.txt".to_string(),
mime_type: Some("text/plain".to_string()),
text: Some("content".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
})
.layer(TimeoutLayer::new(Duration::from_secs(30)))
.build();
let result = resource.read().await;
assert_eq!(result.contents[0].text.as_deref(), Some("content"));
}
#[tokio::test]
async fn test_handler_resource_with_timeout_error() {
let resource = ResourceBuilder::new("file:///slow.txt")
.name("Slow Resource")
.handler(|| async {
tokio::time::sleep(Duration::from_secs(1)).await;
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: "file:///slow.txt".to_string(),
mime_type: Some("text/plain".to_string()),
text: Some("content".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
})
.layer(TimeoutLayer::new(Duration::from_millis(50)))
.build();
let result = resource.read().await;
assert!(
result.contents[0]
.text
.as_ref()
.unwrap()
.contains("Error reading resource")
);
}
#[tokio::test]
async fn test_context_aware_handler() {
let resource = ResourceBuilder::new("file:///ctx.txt")
.name("Context Resource")
.handler_with_context(|_ctx: RequestContext| async {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: "file:///ctx.txt".to_string(),
mime_type: Some("text/plain".to_string()),
text: Some("context aware".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
})
.build();
let result = resource.read().await;
assert_eq!(result.contents[0].text.as_deref(), Some("context aware"));
}
#[tokio::test]
async fn test_context_aware_handler_with_layer() {
let resource = ResourceBuilder::new("file:///ctx-layer.txt")
.name("Context Resource with Layer")
.handler_with_context(|_ctx: RequestContext| async {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: "file:///ctx-layer.txt".to_string(),
mime_type: Some("text/plain".to_string()),
text: Some("context with layer".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
})
.layer(TimeoutLayer::new(Duration::from_secs(30)))
.build();
let result = resource.read().await;
assert_eq!(
result.contents[0].text.as_deref(),
Some("context with layer")
);
}
#[tokio::test]
async fn test_trait_resource() {
struct TestResource;
impl McpResource for TestResource {
const URI: &'static str = "test://resource";
const NAME: &'static str = "Test";
const DESCRIPTION: Option<&'static str> = Some("A test resource");
const MIME_TYPE: Option<&'static str> = Some("text/plain");
async fn read(&self) -> Result<ReadResourceResult> {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri: Self::URI.to_string(),
mime_type: Self::MIME_TYPE.map(|s| s.to_string()),
text: Some("test content".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
}
}
let resource = TestResource.into_resource();
assert_eq!(resource.uri, "test://resource");
assert_eq!(resource.name, "Test");
let result = resource.read().await;
assert_eq!(result.contents[0].text.as_deref(), Some("test content"));
}
#[test]
fn test_resource_definition() {
let resource = ResourceBuilder::new("file:///test.txt")
.name("Test")
.description("Description")
.mime_type("text/plain")
.text("content");
let def = resource.definition();
assert_eq!(def.uri, "file:///test.txt");
assert_eq!(def.name, "Test");
assert_eq!(def.description.as_deref(), Some("Description"));
assert_eq!(def.mime_type.as_deref(), Some("text/plain"));
}
#[test]
fn test_resource_request_new() {
let ctx = RequestContext::new(crate::protocol::RequestId::Number(1));
let req = ResourceRequest::new(ctx, "file:///test.txt".to_string());
assert_eq!(req.uri, "file:///test.txt");
}
#[test]
fn test_resource_catch_error_clone() {
let handler = FnHandler {
handler: || async {
Ok::<_, Error>(ReadResourceResult {
contents: vec![],
meta: None,
})
},
};
let service = ResourceHandlerService::new(handler);
let catch_error = ResourceCatchError::new(service);
let _clone = catch_error.clone();
}
#[test]
fn test_resource_catch_error_debug() {
let handler = FnHandler {
handler: || async {
Ok::<_, Error>(ReadResourceResult {
contents: vec![],
meta: None,
})
},
};
let service = ResourceHandlerService::new(handler);
let catch_error = ResourceCatchError::new(service);
let debug = format!("{:?}", catch_error);
assert!(debug.contains("ResourceCatchError"));
}
#[test]
fn test_compile_uri_template_simple() {
let (regex, vars) = compile_uri_template("file:///{path}").unwrap();
assert_eq!(vars, vec!["path"]);
assert!(regex.is_match("file:///README.md"));
assert!(!regex.is_match("file:///foo/bar")); }
#[test]
fn test_compile_uri_template_multiple_vars() {
let (regex, vars) = compile_uri_template("api://v1/{resource}/{id}").unwrap();
assert_eq!(vars, vec!["resource", "id"]);
assert!(regex.is_match("api://v1/users/123"));
assert!(regex.is_match("api://v1/posts/abc"));
assert!(!regex.is_match("api://v1/users")); }
#[test]
fn test_compile_uri_template_reserved_expansion() {
let (regex, vars) = compile_uri_template("file:///{+path}").unwrap();
assert_eq!(vars, vec!["path"]);
assert!(regex.is_match("file:///README.md"));
assert!(regex.is_match("file:///foo/bar/baz.txt")); }
#[test]
fn test_compile_uri_template_special_chars() {
let (regex, vars) = compile_uri_template("http://example.com/api?query={q}").unwrap();
assert_eq!(vars, vec!["q"]);
assert!(regex.is_match("http://example.com/api?query=hello"));
}
#[test]
fn test_resource_template_match_uri() {
let template = ResourceTemplateBuilder::new("db://users/{id}")
.name("User Records")
.handler(|uri: String, vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: Some(format!("User {}", vars.get("id").unwrap())),
blob: None,
meta: None,
}],
meta: None,
})
});
let vars = template.match_uri("db://users/123").unwrap();
assert_eq!(vars.get("id"), Some(&"123".to_string()));
assert!(template.match_uri("db://posts/123").is_none());
assert!(template.match_uri("db://users").is_none());
}
#[test]
fn test_resource_template_match_multiple_vars() {
let template = ResourceTemplateBuilder::new("api://{version}/{resource}/{id}")
.name("API Resources")
.handler(|uri: String, _vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: None,
blob: None,
meta: None,
}],
meta: None,
})
});
let vars = template.match_uri("api://v2/users/abc-123").unwrap();
assert_eq!(vars.get("version"), Some(&"v2".to_string()));
assert_eq!(vars.get("resource"), Some(&"users".to_string()));
assert_eq!(vars.get("id"), Some(&"abc-123".to_string()));
}
#[tokio::test]
async fn test_resource_template_read() {
let template = ResourceTemplateBuilder::new("file:///{path}")
.name("Files")
.mime_type("text/plain")
.handler(|uri: String, vars: HashMap<String, String>| async move {
let path = vars.get("path").unwrap().clone();
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: Some("text/plain".to_string()),
text: Some(format!("Contents of {}", path)),
blob: None,
meta: None,
}],
meta: None,
})
});
let vars = template.match_uri("file:///README.md").unwrap();
let result = template.read("file:///README.md", vars).await.unwrap();
assert_eq!(result.contents.len(), 1);
assert_eq!(result.contents[0].uri, "file:///README.md");
assert_eq!(
result.contents[0].text.as_deref(),
Some("Contents of README.md")
);
}
#[test]
fn test_resource_template_definition() {
let template = ResourceTemplateBuilder::new("db://records/{id}")
.name("Database Records")
.description("Access database records by ID")
.mime_type("application/json")
.handler(|uri: String, _vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: None,
blob: None,
meta: None,
}],
meta: None,
})
});
let def = template.definition();
assert_eq!(def.uri_template, "db://records/{id}");
assert_eq!(def.name, "Database Records");
assert_eq!(
def.description.as_deref(),
Some("Access database records by ID")
);
assert_eq!(def.mime_type.as_deref(), Some("application/json"));
}
#[test]
fn test_resource_template_reserved_path() {
let template = ResourceTemplateBuilder::new("file:///{+path}")
.name("Files with subpaths")
.handler(|uri: String, _vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: None,
blob: None,
meta: None,
}],
meta: None,
})
});
let vars = template.match_uri("file:///src/lib/utils.rs").unwrap();
assert_eq!(vars.get("path"), Some(&"src/lib/utils.rs".to_string()));
}
#[test]
fn test_resource_annotations() {
use crate::protocol::{ContentAnnotations, ContentRole};
let annotations = ContentAnnotations {
audience: Some(vec![ContentRole::User]),
priority: Some(0.8),
last_modified: None,
};
let resource = ResourceBuilder::new("file:///important.txt")
.name("Important File")
.annotations(annotations.clone())
.text("content");
let def = resource.definition();
assert!(def.annotations.is_some());
let ann = def.annotations.unwrap();
assert_eq!(ann.priority, Some(0.8));
assert_eq!(ann.audience.unwrap(), vec![ContentRole::User]);
}
#[test]
fn test_resource_template_annotations() {
use crate::protocol::{ContentAnnotations, ContentRole};
let annotations = ContentAnnotations {
audience: Some(vec![ContentRole::Assistant]),
priority: Some(0.5),
last_modified: None,
};
let template = ResourceTemplateBuilder::new("db://users/{id}")
.name("Users")
.annotations(annotations)
.handler(|uri: String, _vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: Some("data".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
});
let def = template.definition();
assert!(def.annotations.is_some());
let ann = def.annotations.unwrap();
assert_eq!(ann.priority, Some(0.5));
assert_eq!(ann.audience.unwrap(), vec![ContentRole::Assistant]);
}
#[test]
fn test_resource_no_annotations_by_default() {
let resource = ResourceBuilder::new("file:///test.txt")
.name("Test")
.text("content");
let def = resource.definition();
assert!(def.annotations.is_none());
}
#[test]
fn test_try_handler_success() {
let result = ResourceTemplateBuilder::new("db://users/{id}")
.name("Users")
.try_handler(|uri: String, _vars: HashMap<String, String>| async move {
Ok(ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: None,
text: Some("ok".to_string()),
blob: None,
meta: None,
}],
meta: None,
})
});
assert!(result.is_ok());
let template = result.unwrap();
assert_eq!(template.uri_template, "db://users/{id}");
}
#[test]
fn test_compile_uri_template_returns_result() {
assert!(compile_uri_template("file:///{path}").is_ok());
assert!(compile_uri_template("api://v1/{resource}/{id}").is_ok());
assert!(compile_uri_template("file:///{+path}").is_ok());
assert!(compile_uri_template("no-vars").is_ok());
assert!(compile_uri_template("").is_ok());
}
}