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 serde_json::Value;
#[cfg(feature = "stateless")]
use tower::ServiceExt;
use tower::util::BoxCloneService;
use tower_service::Service;
#[cfg(feature = "stateless")]
use tokio::sync::Mutex;
use crate::context::RequestContext;
use crate::error::{Error, JsonRpcError, Result};
use crate::protocol::{
ContentAnnotations, PromptArgument, ReadResourceResult, RequestOutcome, 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,
std::result::Result<ReadResourceResult, JsonRpcError>,
Infallible,
>;
#[cfg(feature = "stateless")]
type BoxMrtrResourceService = BoxCloneService<
ResourceRequest,
std::result::Result<RequestOutcome<ReadResourceResult>, JsonRpcError>,
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,
}
}
impl<F, E> Future for ResourceCatchErrorFuture<F>
where
F: Future<
Output = std::result::Result<std::result::Result<ReadResourceResult, JsonRpcError>, E>,
>,
E: fmt::Display,
{
type Output =
std::result::Result<std::result::Result<ReadResourceResult, JsonRpcError>, 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(inner)) => Poll::Ready(Ok(inner)),
Poll::Ready(Err(err)) => {
Poll::Ready(Ok(Err(JsonRpcError::internal_error(err.to_string()))))
}
}
}
}
impl<S> Service<ResourceRequest> for ResourceCatchError<S>
where
S: Service<ResourceRequest, Response = std::result::Result<ReadResourceResult, JsonRpcError>>
+ Clone
+ Send
+ 'static,
S::Error: fmt::Display + Send,
S::Future: Send,
{
type Response = std::result::Result<ReadResourceResult, JsonRpcError>;
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 fut = self.inner.call(req);
ResourceCatchErrorFuture { inner: fut }
}
}
#[cfg(feature = "stateless")]
#[derive(Clone)]
struct MrtrResourceCatchError<S> {
inner: S,
}
#[cfg(feature = "stateless")]
impl<S> MrtrResourceCatchError<S> {
fn new(inner: S) -> Self {
Self { inner }
}
}
#[cfg(feature = "stateless")]
impl<S> Service<ResourceRequest> for MrtrResourceCatchError<S>
where
S: Service<
ResourceRequest,
Response = std::result::Result<RequestOutcome<ReadResourceResult>, JsonRpcError>,
> + Clone
+ Send
+ 'static,
S::Error: fmt::Display + Send + 'static,
S::Future: Send + 'static,
{
type Response = std::result::Result<RequestOutcome<ReadResourceResult>, JsonRpcError>;
type Error = Infallible;
type Future =
Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
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(Err(_)) => Poll::Ready(Ok(())),
Poll::Pending => Poll::Pending,
}
}
fn call(&mut self, req: ResourceRequest) -> Self::Future {
let future = self.inner.call(req);
Box::pin(async move {
Ok(match future.await {
Ok(inner) => inner,
Err(error) => Err(JsonRpcError::internal_error(error.to_string())),
})
})
}
}
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub(crate) trait ResourceHandler: Send + Sync {
fn read(&self) -> BoxFuture<'_, Result<ReadResourceResult>>;
fn read_with_context(&self, _ctx: RequestContext) -> BoxFuture<'_, Result<ReadResourceResult>> {
self.read()
}
}
#[cfg(feature = "stateless")]
pub(crate) trait MrtrResourceHandler: Send + Sync {
fn read(
&self,
ctx: RequestContext,
) -> BoxFuture<'_, Result<RequestOutcome<ReadResourceResult>>>;
}
#[cfg(feature = "stateless")]
struct MrtrResourceHandlerService<H> {
handler: Arc<H>,
}
#[cfg(feature = "stateless")]
impl<H> MrtrResourceHandlerService<H> {
fn new(handler: H) -> Self {
Self {
handler: Arc::new(handler),
}
}
}
#[cfg(feature = "stateless")]
impl<H> Clone for MrtrResourceHandlerService<H> {
fn clone(&self) -> Self {
Self {
handler: self.handler.clone(),
}
}
}
#[cfg(feature = "stateless")]
impl<H> Service<ResourceRequest> for MrtrResourceHandlerService<H>
where
H: MrtrResourceHandler + 'static,
{
type Response = std::result::Result<RequestOutcome<ReadResourceResult>, JsonRpcError>;
type Error = Infallible;
type Future =
Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Infallible>> + 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 {
Ok(handler
.read(req.ctx)
.await
.map_err(Error::into_json_rpc_error))
})
}
}
#[cfg(feature = "stateless")]
struct ServiceMrtrResourceHandler {
service: Mutex<BoxMrtrResourceService>,
uri: String,
}
#[cfg(feature = "stateless")]
impl MrtrResourceHandler for ServiceMrtrResourceHandler {
fn read(
&self,
ctx: RequestContext,
) -> BoxFuture<'_, Result<RequestOutcome<ReadResourceResult>>> {
Box::pin(async move {
let request = ResourceRequest::new(ctx, self.uri.clone());
let mut service = self.service.lock().await.clone();
let outcome = service
.ready()
.await
.expect("MRTR resource service is infallible")
.call(request)
.await
.expect("MRTR resource service is infallible");
outcome.map_err(Into::into)
})
}
}
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 = std::result::Result<ReadResourceResult, JsonRpcError>;
type Error = Infallible;
type Future =
Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Infallible>> + 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 {
Ok(handler
.read_with_context(req.ctx)
.await
.map_err(Error::into_json_rpc_error))
})
}
}
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>,
pub meta: Option<Value>,
service: Option<BoxResourceService>,
#[cfg(feature = "stateless")]
mrtr_handler: Option<Arc<dyn MrtrResourceHandler>>,
}
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(),
meta: self.meta.clone(),
service: self.service.clone(),
#[cfg(feature = "stateless")]
mrtr_handler: self.mrtr_handler.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)
.field("meta", &self.meta)
.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: self.meta.clone(),
}
}
pub fn with_meta(
mut self,
meta: Value,
) -> std::result::Result<Self, crate::protocol::MetaValidationError> {
crate::protocol::validate_meta_object(&meta)?;
self.meta = Some(meta);
Ok(self)
}
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> {
let resource = self.clone();
let uri = self.uri.clone();
Box::pin(async move {
match resource.read_outcome_with_context(ctx).await {
Ok(RequestOutcome::Complete(result)) => result,
Ok(RequestOutcome::InputRequired(_)) => ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: Some("text/plain".into()),
text: Some(
"resource requires additional client input; use read_outcome_with_context"
.into(),
),
blob: None,
meta: None,
}],
..ReadResourceResult::default()
},
Err(error) => ReadResourceResult {
contents: vec![ResourceContent {
uri,
mime_type: Some("text/plain".into()),
text: Some(error.to_string()),
blob: None,
meta: None,
}],
..ReadResourceResult::default()
},
}
})
}
pub fn read_outcome_with_context(
&self,
ctx: RequestContext,
) -> BoxFuture<'static, Result<RequestOutcome<ReadResourceResult>>> {
use tower::ServiceExt;
#[cfg(feature = "stateless")]
if let Some(handler) = self.mrtr_handler.clone() {
return Box::pin(async move { handler.read(ctx).await });
}
let service = self
.service
.clone()
.expect("resource must have a complete or MRTR handler");
let uri = self.uri.clone();
Box::pin(async move {
let result = service
.oneshot(ResourceRequest::new(ctx, uri))
.await
.unwrap();
match result {
Ok(read_result) => Ok(RequestOutcome::Complete(read_result)),
Err(json_rpc_err) => Err(json_rpc_err.into()),
}
})
}
#[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,
meta: None,
service: Some(service),
#[cfg(feature = "stateless")]
mrtr_handler: None,
}
}
#[cfg(feature = "stateless")]
#[allow(clippy::too_many_arguments)]
fn from_mrtr_handler<H: MrtrResourceHandler + '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 {
Self {
uri,
name,
title,
description,
mime_type,
icons,
size,
annotations,
meta: None,
service: None,
mrtr_handler: Some(Arc::new(handler)),
}
}
}
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,
}
}
#[cfg(feature = "stateless")]
pub fn mrtr_handler<F, Fut>(self, handler: F) -> ResourceBuilderWithMrtrHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
ResourceBuilderWithMrtrHandler {
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,
..Default::default()
})
}
})
.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,
..Default::default()
})
}
})
.build()
}
}
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,
}
#[cfg(feature = "stateless")]
pub struct ResourceBuilderWithMrtrHandler<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,
}
#[cfg(feature = "stateless")]
pub struct ResourceBuilderWithMrtrLayer<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,
}
#[cfg(feature = "stateless")]
impl<F, Fut> ResourceBuilderWithMrtrHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
Resource::from_mrtr_handler(
self.uri,
name,
self.title,
self.description,
self.mime_type,
self.icons,
self.size,
self.annotations,
MrtrContextHandler {
handler: self.handler,
},
)
}
pub fn layer<L>(self, layer: L) -> ResourceBuilderWithMrtrLayer<F, L> {
ResourceBuilderWithMrtrLayer {
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,
}
}
}
#[cfg(feature = "stateless")]
#[allow(private_bounds)]
impl<F, Fut, L> ResourceBuilderWithMrtrLayer<F, L>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
L: tower::Layer<MrtrResourceHandlerService<MrtrContextHandler<F>>>
+ Clone
+ Send
+ Sync
+ 'static,
L::Service: Service<
ResourceRequest,
Response = std::result::Result<RequestOutcome<ReadResourceResult>, JsonRpcError>,
> + Clone
+ Send
+ 'static,
<L::Service as Service<ResourceRequest>>::Error: fmt::Display + Send + 'static,
<L::Service as Service<ResourceRequest>>::Future: Send + 'static,
{
pub fn build(self) -> Resource {
let name = self.name.unwrap_or_else(|| self.uri.clone());
let handler = MrtrContextHandler {
handler: self.handler,
};
let service = MrtrResourceHandlerService::new(handler);
let service = self.layer.layer(service);
let service = BoxCloneService::new(MrtrResourceCatchError::new(service));
Resource {
uri: self.uri.clone(),
name,
title: self.title,
description: self.description,
mime_type: self.mime_type,
icons: self.icons,
size: self.size,
annotations: self.annotations,
meta: None,
service: None,
mrtr_handler: Some(Arc::new(ServiceMrtrResourceHandler {
service: Mutex::new(service),
uri: self.uri,
})),
}
}
pub fn layer<L2>(
self,
layer: L2,
) -> ResourceBuilderWithMrtrLayer<F, tower::layer::util::Stack<L2, L>> {
ResourceBuilderWithMrtrLayer {
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),
}
}
}
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,
}
}
}
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 = std::result::Result<ReadResourceResult, JsonRpcError>>
+ 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,
meta: None,
service: Some(service),
#[cfg(feature = "stateless")]
mrtr_handler: None,
}
}
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),
}
}
}
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,
}
}
}
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 = std::result::Result<ReadResourceResult, JsonRpcError>>
+ 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,
meta: None,
service: Some(service),
#[cfg(feature = "stateless")]
mrtr_handler: None,
}
}
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))
}
}
#[cfg(feature = "stateless")]
struct MrtrContextHandler<F> {
handler: F,
}
#[cfg(feature = "stateless")]
impl<F, Fut> MrtrResourceHandler for MrtrContextHandler<F>
where
F: Fn(RequestContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
fn read(
&self,
ctx: RequestContext,
) -> BoxFuture<'_, Result<RequestOutcome<ReadResourceResult>>> {
Box::pin((self.handler)(ctx))
}
}
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(crate) trait ResourceTemplateHandler: Send + Sync {
fn read(
&self,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<ReadResourceResult>>;
}
#[cfg(feature = "stateless")]
pub(crate) trait MrtrResourceTemplateHandler: Send + Sync {
fn read(
&self,
ctx: RequestContext,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<RequestOutcome<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>,
pub arguments: Vec<PromptArgument>,
meta: Option<Value>,
pattern: regex::Regex,
query_variables: Vec<String>,
variables: Vec<String>,
handler: Option<Arc<dyn ResourceTemplateHandler>>,
#[cfg(feature = "stateless")]
mrtr_handler: Option<Arc<dyn MrtrResourceTemplateHandler>>,
}
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(),
arguments: self.arguments.clone(),
meta: self.meta.clone(),
pattern: self.pattern.clone(),
query_variables: self.query_variables.clone(),
variables: self.variables.clone(),
handler: self.handler.clone(),
#[cfg(feature = "stateless")]
mrtr_handler: self.mrtr_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: self.arguments.clone(),
meta: self.meta.clone(),
}
}
pub fn with_meta(
mut self,
meta: Value,
) -> std::result::Result<Self, crate::protocol::MetaValidationError> {
crate::protocol::validate_meta_object(&meta)?;
self.meta = Some(meta);
Ok(self)
}
pub fn match_uri(&self, uri: &str) -> Option<HashMap<String, String>> {
let (path, query) = if self.query_variables.is_empty() {
(uri, None)
} else {
match uri.split_once('?') {
Some((path, query)) => (path, Some(query)),
None => (uri, None),
}
};
let mut matched: HashMap<String, String> = self.pattern.captures(path).map(|caps| {
self.variables
.iter()
.enumerate()
.filter_map(|(i, name)| {
caps.get(i + 1)
.map(|m| (name.clone(), m.as_str().to_string()))
})
.collect()
})?;
if let Some(query) = query {
matched.extend(extract_query_variables(query, &self.query_variables));
}
Some(matched)
}
pub fn read(
&self,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<ReadResourceResult>> {
match &self.handler {
Some(handler) => handler.read(uri, variables),
None => Box::pin(async {
Err(Error::invalid_params(
"MRTR resource template requires read_outcome_with_context",
))
}),
}
}
pub fn read_outcome_with_context(
&self,
ctx: RequestContext,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<RequestOutcome<ReadResourceResult>>> {
let _ = &ctx;
#[cfg(feature = "stateless")]
if let Some(handler) = &self.mrtr_handler {
return handler.read(ctx, uri, variables);
}
match &self.handler {
Some(handler) => {
let handler = handler.clone();
let uri = uri.to_string();
Box::pin(async move {
handler
.read(&uri, variables)
.await
.map(RequestOutcome::Complete)
})
}
None => Box::pin(async {
Err(Error::invalid_params(
"resource template has neither a complete nor MRTR handler",
))
}),
}
}
}
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>,
arguments: Vec<PromptArgument>,
}
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,
arguments: Vec::new(),
}
}
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 argument(
mut self,
name: impl Into<String>,
description: Option<impl Into<String>>,
required: bool,
) -> Self {
self.arguments.push(PromptArgument {
name: name.into(),
description: description.map(Into::into),
required,
});
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 CompiledTemplate {
pattern,
variables,
query_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,
arguments: self.arguments,
meta: None,
pattern,
query_variables,
variables,
handler: Some(Arc::new(FnTemplateHandler { handler })),
#[cfg(feature = "stateless")]
mrtr_handler: None,
})
}
#[cfg(feature = "stateless")]
pub fn mrtr_handler<F, Fut>(self, handler: F) -> ResourceTemplate
where
F: Fn(RequestContext, String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
self.try_mrtr_handler(handler)
.unwrap_or_else(|error| panic!("Invalid URI template: {error}"))
}
#[cfg(feature = "stateless")]
pub fn try_mrtr_handler<F, Fut>(
self,
handler: F,
) -> std::result::Result<ResourceTemplate, Error>
where
F: Fn(RequestContext, String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
let CompiledTemplate {
pattern,
variables,
query_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,
arguments: self.arguments,
meta: None,
pattern,
query_variables,
variables,
handler: None,
mrtr_handler: Some(Arc::new(MrtrFnTemplateHandler { 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))
}
}
#[cfg(feature = "stateless")]
struct MrtrFnTemplateHandler<F> {
handler: F,
}
#[cfg(feature = "stateless")]
impl<F, Fut> MrtrResourceTemplateHandler for MrtrFnTemplateHandler<F>
where
F: Fn(RequestContext, String, HashMap<String, String>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<RequestOutcome<ReadResourceResult>>> + Send + 'static,
{
fn read(
&self,
ctx: RequestContext,
uri: &str,
variables: HashMap<String, String>,
) -> BoxFuture<'_, Result<RequestOutcome<ReadResourceResult>>> {
Box::pin((self.handler)(ctx, uri.to_string(), variables))
}
}
struct CompiledTemplate {
pattern: regex::Regex,
variables: Vec<String>,
query_variables: Vec<String>,
}
fn compile_uri_template(template: &str) -> std::result::Result<CompiledTemplate, Error> {
let mut pattern = String::from("^");
let mut variables = Vec::new();
let mut query_variables: Vec<String> = Vec::new();
let mut chars = template.chars().peekable();
while let Some(c) = chars.next() {
if c == '{' {
if matches!(chars.peek(), Some('?') | Some('&')) {
chars.next();
let body: String = chars.by_ref().take_while(|&c| c != '}').collect();
if !query_variables.is_empty() {
return Err(Error::Internal(format!(
"URI template '{template}' declares more than one query expression"
)));
}
for name in body.split(',') {
let name = name.trim();
if name.is_empty() {
return Err(Error::Internal(format!(
"URI template '{template}' has an empty query variable name"
)));
}
query_variables.push(name.to_string());
}
if chars.peek().is_some() {
return Err(Error::Internal(format!(
"URI template '{template}' has text after its query expression, \
which describes the end of the URI"
)));
}
continue;
}
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(CompiledTemplate {
pattern: regex,
variables,
query_variables,
})
}
fn decode_query_component(value: &str) -> String {
let bytes = value.as_bytes();
let mut out: Vec<u8> = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
match bytes[i] {
b'+' => {
out.push(b' ');
i += 1;
}
b'%' if i + 2 < bytes.len() => {
let high = (bytes[i + 1] as char).to_digit(16);
let low = (bytes[i + 2] as char).to_digit(16);
match (high, low) {
(Some(high), Some(low)) => {
out.push((high * 16 + low) as u8);
i += 3;
}
_ => {
out.push(bytes[i]);
i += 1;
}
}
}
byte => {
out.push(byte);
i += 1;
}
}
}
String::from_utf8(out).unwrap_or_else(|_| value.to_string())
}
fn extract_query_variables(query: &str, declared: &[String]) -> HashMap<String, String> {
let mut found = HashMap::new();
for pair in query.split('&').filter(|p| !p.is_empty()) {
let (raw_key, raw_value) = match pair.split_once('=') {
Some((key, value)) => (key, value),
None => (pair, ""),
};
let key = decode_query_component(raw_key);
if !declared.contains(&key) {
continue;
}
found
.entry(key)
.or_insert_with(|| decode_query_component(raw_value));
}
found
}
#[cfg(test)]
mod query_expansion_tests;
#[cfg(test)]
mod tests;