feat: Add prompt support (#351)

* feat: add prompt support with typed argument handling

- Implement #[prompt], #[prompt_router], and #[prompt_handler] macros
- Add automatic JSON schema generation from Rust types for arguments
- Support flexible async handler signatures with automatic adaptation
- Create PromptRouter for efficient prompt dispatch
- Include comprehensive tests and example implementation

This enables MCP servers to provide reusable prompt templates that
LLMs can discover and invoke with strongly-typed parameters, similar
to the existing tool system but optimized for prompt use cases.

🤖 Generated with Claude Code
Co-Authored-By: Claude <noreply@anthropic.com>

* refactor: unify parameter handling between tools and prompts

- Replace Arguments<T> with Parameters<T> for consistent API
- Create shared common module for tool/prompt utilities
- Modernize async handling with futures::future::BoxFuture
- Move cached_schema_for_type to common module for reuse
- Update error types from rmcp::Error to rmcp::ErrorData
- Add comprehensive trait implementations for parameter extraction

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>

* refactor: extract Parameters wrapper to shared module and unify trait usage

- Move Parameters type from common.rs to dedicated wrapper/parameters.rs module
- Unify extractor traits by using FromContextPart for both tools and prompts
- Remove duplicate FromToolCallContextPart and FromPromptContextPart traits
- Add structured output support to Json wrapper with IntoCallToolResult impl
- Improve error messages with more descriptive panic for schema serialization
- Update all imports across codebase to use new module path
- Clean up trailing whitespace and formatting inconsistencies

This consolidates parameter extraction logic and reduces code duplication
between tool and prompt handlers while maintaining backward compatibility.

* chore: remove committed .egg-info directory and update gitignore

* docs: fix documentation formatting
This commit is contained in:
Michael Assaf 2025-08-20 23:22:19 -04:00 • committed by GitHub
commit 9349f5cb26
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
40 changed files with 3191 additions and 464 deletions

9
.gitignore vendored
View file

@ -13,9 +13,16 @@ Cargo.lock
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
.vscode/
# Python artifacts (for test directories)
*.egg-info/
__pycache__/
*.pyc
*.pyo
# RustRover
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
#.idea/

View file

@ -0,0 +1,65 @@
//! Common utilities shared between different macro implementations
use quote::quote;
use syn::{Attribute, Expr, FnArg, ImplItemFn, Signature, Type};
/// Parse a None expression
pub fn none_expr() -> syn::Result<Expr> {
syn::parse2::<Expr>(quote! { None })
}
/// Extract documentation from doc attributes
pub fn extract_doc_line(existing_docs: Option<String>, attr: &Attribute) -> Option<String> {
if !attr.path().is_ident("doc") {
return None;
}
let syn::Meta::NameValue(name_value) = &attr.meta else {
return None;
};
let syn::Expr::Lit(expr_lit) = &name_value.value else {
return None;
};
let syn::Lit::Str(lit_str) = &expr_lit.lit else {
return None;
};
let content = lit_str.value().trim().to_string();
match (existing_docs, content) {
(Some(mut existing_docs), content) if !content.is_empty() => {
existing_docs.push('\n');
existing_docs.push_str(&content);
Some(existing_docs)
}
(Some(existing_docs), _) => Some(existing_docs),
(None, content) if !content.is_empty() => Some(content),
_ => None,
}
}
/// Find Parameters<T> type in function signature
/// Returns the full Parameters<T> type if found
pub fn find_parameters_type_in_sig(sig: &Signature) -> Option<Box<Type>> {
sig.inputs.iter().find_map(|input| {
if let FnArg::Typed(pat_type) = input {
if let Type::Path(type_path) = &*pat_type.ty {
if type_path
.path
.segments
.last()
.is_some_and(|type_name| type_name.ident == "Parameters")
{
return Some(pat_type.ty.clone());
}
}
}
None
})
}
/// Find Parameters<T> type in ImplItemFn
pub fn find_parameters_type_impl(fn_item: &ImplItemFn) -> Option<Box<Type>> {
find_parameters_type_in_sig(&fn_item.sig)
}

View file

@ -1,6 +1,10 @@
#[allow(unused_imports)]
use proc_macro::TokenStream;
mod common;
mod prompt;
mod prompt_handler;
mod prompt_router;
mod tool;
mod tool_handler;
mod tool_router;
@ -160,3 +164,102 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
.unwrap_or_else(|err| err.to_compile_error())
.into()
}
/// # prompt
///
/// This macro is used to mark a function as a prompt handler.
///
/// This will generate a function that returns the attribute of this prompt, with type `rmcp::model::Prompt`.
///
/// ## Usage
///
/// | field | type | usage |
/// | :- | :- | :- |
/// | `name` | `String` | The name of the prompt. If not provided, it defaults to the function name. |
/// | `description` | `String` | A description of the prompt. The document of this function will be used if not provided. |
/// | `arguments` | `Expr` | An expression that evaluates to `Option<Vec<PromptArgument>>` defining the prompt's arguments. If not provided, it will automatically generate arguments from the `Parameters<T>` type found in the function signature. |
///
/// ## Example
///
/// ```rust,ignore
/// #[prompt(name = "code_review", description = "Reviews code for best practices")]
/// pub async fn code_review_prompt(&self, Parameters(args): Parameters<CodeReviewArgs>) -> Result<Vec<PromptMessage>> {
/// // Generate prompt messages based on arguments
/// }
/// ```
#[proc_macro_attribute]
pub fn prompt(attr: TokenStream, input: TokenStream) -> TokenStream {
prompt::prompt(attr.into(), input.into())
.unwrap_or_else(|err| err.to_compile_error())
.into()
}
/// # prompt_router
///
/// This macro generates a prompt router based on functions marked with `#[rmcp::prompt]` in an implementation block.
///
/// It creates a function that returns a `PromptRouter` instance.
///
/// ## Usage
///
/// | field | type | usage |
/// | :- | :- | :- |
/// | `router` | `Ident` | The name of the router function to be generated. Defaults to `prompt_router`. |
/// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
///
/// ## Example
///
/// ```rust,ignore
/// #[prompt_router]
/// impl MyPromptHandler {
/// #[prompt]
/// pub async fn greeting_prompt(&self, Parameters(args): Parameters<GreetingArgs>) -> Result<Vec<PromptMessage>, Error> {
/// // Generate greeting prompt using args
/// }
///
/// pub fn new() -> Self {
/// Self {
/// // the default name of prompt router will be `prompt_router`
/// prompt_router: Self::prompt_router(),
/// }
/// }
/// }
/// ```
#[proc_macro_attribute]
pub fn prompt_router(attr: TokenStream, input: TokenStream) -> TokenStream {
prompt_router::prompt_router(attr.into(), input.into())
.unwrap_or_else(|err| err.to_compile_error())
.into()
}
/// # prompt_handler
///
/// This macro generates handler methods for `get_prompt` and `list_prompts` in the implementation block, using an existing `PromptRouter` instance.
///
/// ## Usage
///
/// | field | type | usage |
/// | :- | :- | :- |
/// | `router` | `Expr` | The expression to access the `PromptRouter` instance. Defaults to `self.prompt_router`. |
///
/// ## Example
/// ```rust,ignore
/// #[prompt_handler]
/// impl ServerHandler for MyPromptHandler {
/// // ...implement other handler methods
/// }
/// ```
///
/// or using a custom router expression:
/// ```rust,ignore
/// #[prompt_handler(router = self.get_prompt_router())]
/// impl ServerHandler for MyPromptHandler {
/// // ...implement other handler methods
/// }
/// ```
#[proc_macro_attribute]
pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
prompt_handler::prompt_handler(attr.into(), input.into())
.unwrap_or_else(|err| err.to_compile_error())
.into()
}

View file

@ -0,0 +1,181 @@
use darling::{FromMeta, ast::NestedMeta};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{Expr, Ident, ImplItemFn, ReturnType};
use crate::common::{extract_doc_line, none_expr};
#[derive(FromMeta, Default, Debug)]
#[darling(default)]
pub struct PromptAttribute {
/// The name of the prompt
pub name: Option<String>,
/// Optional description of what the prompt does
pub description: Option<String>,
/// Arguments that can be passed to the prompt
pub arguments: Option<Expr>,
}
pub struct ResolvedPromptAttribute {
pub name: String,
pub description: Option<String>,
pub arguments: Expr,
}
impl ResolvedPromptAttribute {
pub fn into_fn(self, fn_ident: Ident) -> syn::Result<ImplItemFn> {
let Self {
name,
description,
arguments,
} = self;
let description = if let Some(description) = description {
quote! { Some(#description.into()) }
} else {
quote! { None }
};
let tokens = quote! {
pub fn #fn_ident() -> rmcp::model::Prompt {
rmcp::model::Prompt {
name: #name.into(),
description: #description,
arguments: #arguments,
}
}
};
syn::parse2::<ImplItemFn>(tokens)
}
}
pub fn prompt(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
let attribute = if attr.is_empty() {
Default::default()
} else {
let attr_args = NestedMeta::parse_meta_list(attr)?;
PromptAttribute::from_list(&attr_args)?
};
let mut fn_item = syn::parse2::<ImplItemFn>(input.clone())?;
let fn_ident = &fn_item.sig.ident;
let prompt_attr_fn_ident = format_ident!("{}_prompt_attr", fn_ident);
// Try to find prompt parameters from function parameters
let arguments_expr = if let Some(arguments) = attribute.arguments {
arguments
} else {
// Look for a type named Parameters in the function signature
let params_ty = crate::common::find_parameters_type_impl(&fn_item);
if let Some(params_ty) = params_ty {
// Generate arguments from the type's schema with caching
syn::parse2::<Expr>(quote! {
rmcp::handler::server::prompt::cached_arguments_from_schema::<#params_ty>()
})?
} else {
// No arguments
none_expr()?
}
};
let name = attribute.name.unwrap_or_else(|| fn_ident.to_string());
let description = attribute
.description
.or_else(|| fn_item.attrs.iter().fold(None, extract_doc_line));
let arguments = arguments_expr;
let resolved_prompt_attr = ResolvedPromptAttribute {
name: name.clone(),
description: description.clone(),
arguments: arguments.clone(),
};
let prompt_attr_fn = resolved_prompt_attr.into_fn(prompt_attr_fn_ident.clone())?;
// Modify the input function for async support (same as tool macro)
if fn_item.sig.asyncness.is_some() {
// 1. remove asyncness from sig
// 2. make return type: `futures::future::BoxFuture<'_, #ReturnType>`
// 3. make body: { Box::pin(async move { #body }) }
let new_output = syn::parse2::<ReturnType>({
let mut lt = quote! { 'static };
if let Some(receiver) = fn_item.sig.receiver() {
if let Some((_, receiver_lt)) = receiver.reference.as_ref() {
if let Some(receiver_lt) = receiver_lt {
lt = quote! { #receiver_lt };
} else {
lt = quote! { '_ };
}
}
}
match &fn_item.sig.output {
syn::ReturnType::Default => {
quote! { -> futures::future::BoxFuture<#lt, ()> }
}
syn::ReturnType::Type(_, ty) => {
quote! { -> futures::future::BoxFuture<#lt, #ty> }
}
}
})?;
let prev_block = &fn_item.block;
let new_block = syn::parse2::<syn::Block>(quote! {
{ Box::pin(async move #prev_block ) }
})?;
fn_item.sig.asyncness = None;
fn_item.sig.output = new_output;
fn_item.block = new_block;
}
Ok(quote! {
#prompt_attr_fn
#fn_item
})
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_prompt_macro() -> syn::Result<()> {
let attr = quote! {
name = "example-prompt",
description = "An example prompt"
};
let input = quote! {
async fn example_prompt(&self, Parameters(args): Parameters<ExampleArgs>) -> Result<String> {
Ok("Example prompt response".to_string())
}
};
let result = prompt(attr, input)?;
// Verify the output contains both the attribute function and the modified function
let result_str = result.to_string();
assert!(result_str.contains("example_prompt_prompt_attr"));
assert!(
result_str.contains("rmcp")
&& result_str.contains("model")
&& result_str.contains("Prompt")
);
Ok(())
}
#[test]
fn test_doc_comment_description() -> syn::Result<()> {
let attr = quote! {}; // No explicit description
let input = quote! {
/// This is a test prompt description
/// with multiple lines
fn test_prompt(&self) -> Result<String> {
Ok("Test".to_string())
}
};
let result = prompt(attr, input)?;
// The output should contain the description from doc comments
let result_str = result.to_string();
assert!(result_str.contains("This is a test prompt description"));
assert!(result_str.contains("with multiple lines"));
Ok(())
}
}

View file

@ -0,0 +1,141 @@
use darling::FromMeta;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{Expr, ImplItem, ItemImpl, parse_quote};
#[derive(FromMeta, Debug, Default)]
#[darling(default)]
pub struct PromptHandlerAttribute {
pub router: Option<Expr>,
}
pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
let attribute = if attr.is_empty() {
Default::default()
} else {
let attr_args = darling::ast::NestedMeta::parse_meta_list(attr)?;
PromptHandlerAttribute::from_list(&attr_args)?
};
let mut impl_block = syn::parse2::<ItemImpl>(input)?;
let router_expr = attribute
.router
.unwrap_or_else(|| syn::parse2(quote! { self.prompt_router }).unwrap());
// Add get_prompt implementation
let get_prompt_impl: ImplItem = parse_quote! {
async fn get_prompt(
&self,
request: GetPromptRequestParam,
context: RequestContext<RoleServer>,
) -> Result<GetPromptResult, rmcp::ErrorData> {
let prompt_context = rmcp::handler::server::prompt::PromptContext::new(
self,
request.name,
request.arguments,
context,
);
#router_expr.get_prompt(prompt_context).await
}
};
// Add list_prompts implementation
let list_prompts_impl: ImplItem = parse_quote! {
async fn list_prompts(
&self,
_request: Option<PaginatedRequestParam>,
_context: RequestContext<RoleServer>,
) -> Result<ListPromptsResult, rmcp::ErrorData> {
let prompts = #router_expr.list_all();
Ok(ListPromptsResult {
prompts,
next_cursor: None,
})
}
};
// Check if methods already exist and replace them if they do
let mut has_get_prompt = false;
let mut has_list_prompts = false;
for item in &mut impl_block.items {
if let ImplItem::Fn(fn_item) = item {
match fn_item.sig.ident.to_string().as_str() {
"get_prompt" => {
*item = get_prompt_impl.clone();
has_get_prompt = true;
}
"list_prompts" => {
*item = list_prompts_impl.clone();
has_list_prompts = true;
}
_ => {}
}
}
}
// Add methods if they don't exist
if !has_get_prompt {
impl_block.items.push(get_prompt_impl);
}
if !has_list_prompts {
impl_block.items.push(list_prompts_impl);
}
Ok(quote! {
#impl_block
})
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_prompt_handler_macro() -> syn::Result<()> {
let input = quote! {
impl ServerHandler for MyPromptHandler {
// Other handler methods...
}
};
let result = prompt_handler(TokenStream::new(), input)?;
let result_str = result.to_string();
// Check that the required methods were generated
assert!(result_str.contains("async fn get_prompt"));
assert!(result_str.contains("PromptContext") && result_str.contains("new"));
assert!(result_str.contains("async fn list_prompts"));
assert!(result_str.contains("ListPromptsResult"));
Ok(())
}
#[test]
fn test_prompt_handler_with_custom_router() -> syn::Result<()> {
let attr = quote! { router = self.get_prompt_router() };
let input = quote! {
impl ServerHandler for MyPromptHandler {
// Other handler methods...
}
};
let result = prompt_handler(attr, input)?;
let result_str = result.to_string();
// Check that the custom router expression is used
assert!(
result_str.contains("self")
&& result_str.contains("get_prompt_router")
&& result_str.contains("get_prompt")
);
assert!(
result_str.contains("self")
&& result_str.contains("get_prompt_router")
&& result_str.contains("list_all")
);
Ok(())
}
}

View file

@ -0,0 +1,111 @@
use darling::FromMeta;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{ImplItem, ItemImpl, Visibility, parse_quote};
#[derive(FromMeta, Debug, Default)]
#[darling(default)]
pub struct PromptRouterAttribute {
pub router: Option<String>,
pub vis: Option<Visibility>,
}
pub fn prompt_router(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
let attribute = if attr.is_empty() {
Default::default()
} else {
let attr_args = darling::ast::NestedMeta::parse_meta_list(attr)?;
PromptRouterAttribute::from_list(&attr_args)?
};
let mut impl_block = syn::parse2::<ItemImpl>(input)?;
let self_ty = &impl_block.self_ty;
let router_fn_ident = attribute
.router
.map(|s| format_ident!("{}", s))
.unwrap_or_else(|| format_ident!("prompt_router"));
let vis = attribute.vis.unwrap_or(Visibility::Inherited);
let mut prompt_route_fn_calls = Vec::new();
for item in &mut impl_block.items {
if let ImplItem::Fn(fn_item) = item {
let has_prompt_attr = fn_item.attrs.iter().any(|attr| {
attr.path()
.segments
.last()
.map(|seg| seg.ident == "prompt")
.unwrap_or(false)
});
if has_prompt_attr {
let fn_ident = &fn_item.sig.ident;
let attr_fn_ident = format_ident!("{}_prompt_attr", fn_ident);
// Check what parameters the function takes
let mut param_names = Vec::new();
let mut param_types = Vec::new();
for input in &fn_item.sig.inputs {
if let syn::FnArg::Typed(pat_type) = input {
// Extract parameter pattern and type
param_types.push(&*pat_type.ty);
param_names.push(&*pat_type.pat);
}
}
// Use the exact same pattern as tool_router
prompt_route_fn_calls.push(quote! {
.with_route((Self::#attr_fn_ident(), Self::#fn_ident))
});
}
}
}
let router_fn: ImplItem = parse_quote! {
#vis fn #router_fn_ident() -> rmcp::handler::server::router::prompt::PromptRouter<#self_ty> {
rmcp::handler::server::router::prompt::PromptRouter::new()
#(#prompt_route_fn_calls)*
}
};
impl_block.items.push(router_fn);
Ok(quote! {
#impl_block
})
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_prompt_router_macro() -> syn::Result<()> {
let input = quote! {
impl MyPromptHandler {
#[prompt]
async fn greeting_prompt(&self) -> Result<Vec<PromptMessage>, Error> {
Ok(vec![])
}
#[prompt]
async fn code_review_prompt(&self, Parameters(args): Parameters<CodeReviewArgs>) -> Result<Vec<PromptMessage>, Error> {
Ok(vec![])
}
}
};
let result = prompt_router(TokenStream::new(), input)?;
let result_str = result.to_string();
// Check that the prompt_router function was generated
assert!(result_str.contains("fn prompt_router"));
assert!(result_str.contains("PromptRouter :: new"));
assert!(result_str.contains("greeting_prompt_prompt_attr"));
assert!(result_str.contains("code_review_prompt_prompt_attr"));
Ok(())
}
}

View file

@ -2,6 +2,65 @@ use darling::{FromMeta, ast::NestedMeta};
use proc_macro2::TokenStream;
use quote::{ToTokens, format_ident, quote};
use syn::{Expr, Ident, ImplItemFn, ReturnType};
use crate::common::{extract_doc_line, none_expr};
/// Check if a type is Json<T> and extract the inner type T
fn extract_json_inner_type(ty: &syn::Type) -> Option<&syn::Type> {
if let syn::Type::Path(type_path) = ty {
if let Some(last_segment) = type_path.path.segments.last() {
if last_segment.ident == "Json" {
if let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments {
if let Some(syn::GenericArgument::Type(inner_type)) = args.args.first() {
return Some(inner_type);
}
}
}
}
}
None
}
/// Extract schema expression from a function's return type
/// Handles patterns like Json<T> and Result<Json<T>, E>
fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
// First, try direct Json<T>
if let Some(inner_type) = extract_json_inner_type(ret_type) {
return syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
})
.ok();
}
// Then, try Result<Json<T>, E>
let type_path = match ret_type {
syn::Type::Path(path) => path,
_ => return None,
};
let last_segment = type_path.path.segments.last()?;
if last_segment.ident != "Result" {
return None;
}
let args = match &last_segment.arguments {
syn::PathArguments::AngleBracketed(args) => args,
_ => return None,
};
let ok_type = match args.args.first()? {
syn::GenericArgument::Type(ty) => ty,
_ => return None,
};
let inner_type = extract_json_inner_type(ok_type)?;
syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
})
.ok()
}
#[derive(FromMeta, Default, Debug)]
#[darling(default)]
pub struct ToolAttribute {
@ -95,98 +154,6 @@ pub struct ToolAnnotationsAttribute {
pub open_world_hint: Option<bool>,
}
fn none_expr() -> Expr {
syn::parse2::<Expr>(quote! { None }).unwrap()
}
/// Check if a type is Json<T> and extract the inner type T
fn extract_json_inner_type(ty: &syn::Type) -> Option<&syn::Type> {
if let syn::Type::Path(type_path) = ty {
if let Some(last_segment) = type_path.path.segments.last() {
if last_segment.ident == "Json" {
if let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments {
if let Some(syn::GenericArgument::Type(inner_type)) = args.args.first() {
return Some(inner_type);
}
}
}
}
}
None
}
/// Extract schema expression from a function's return type
/// Handles patterns like Json<T> and Result<Json<T>, E>
fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
// First, try direct Json<T>
if let Some(inner_type) = extract_json_inner_type(ret_type) {
return syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
})
.ok();
}
// Then, try Result<Json<T>, E>
let type_path = match ret_type {
syn::Type::Path(path) => path,
_ => return None,
};
let last_segment = type_path.path.segments.last()?;
if last_segment.ident != "Result" {
return None;
}
let args = match &last_segment.arguments {
syn::PathArguments::AngleBracketed(args) => args,
_ => return None,
};
let ok_type = match args.args.first()? {
syn::GenericArgument::Type(ty) => ty,
_ => return None,
};
let inner_type = extract_json_inner_type(ok_type)?;
syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
})
.ok()
}
// extract doc line from attribute
fn extract_doc_line(existing_docs: Option<String>, attr: &syn::Attribute) -> Option<String> {
if !attr.path().is_ident("doc") {
return None;
}
let syn::Meta::NameValue(name_value) = &attr.meta else {
return None;
};
let syn::Expr::Lit(expr_lit) = &name_value.value else {
return None;
};
let syn::Lit::Str(lit_str) = &expr_lit.lit else {
return None;
};
let content = lit_str.value().trim().to_string();
match (existing_docs, content) {
(Some(mut existing_docs), content) if !content.is_empty() => {
existing_docs.push('\n');
existing_docs.push_str(&content);
Some(existing_docs)
}
(Some(existing_docs), _) => Some(existing_docs),
(None, content) if !content.is_empty() => Some(content),
_ => None,
}
}
pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
let attribute = if attr.is_empty() {
Default::default()
@ -202,30 +169,16 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
input_schema
} else {
// try to find some parameters wrapper in the function
let params_ty = fn_item.sig.inputs.iter().find_map(|input| {
if let syn::FnArg::Typed(pat_type) = input {
if let syn::Type::Path(type_path) = &*pat_type.ty {
if type_path
.path
.segments
.last()
.is_some_and(|type_name| type_name.ident == "Parameters")
{
return Some(pat_type.ty.clone());
}
}
}
None
});
let params_ty = crate::common::find_parameters_type_impl(&fn_item);
if let Some(params_ty) = params_ty {
// if found, use the Parameters schema
syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<#params_ty>()
rmcp::handler::server::common::cached_schema_for_type::<#params_ty>()
})?
} else {
// if not found, use the default EmptyObject schema
syn::parse2::<Expr>(quote! {
rmcp::handler::server::tool::cached_schema_for_type::<rmcp::model::EmptyObject>()
rmcp::handler::server::common::cached_schema_for_type::<rmcp::model::EmptyObject>()
})?
}
};
@ -257,7 +210,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
};
syn::parse2::<Expr>(token_stream)?
} else {
none_expr()
none_expr()?
};
// Handle output_schema - either explicit or generated from return type
let output_schema_expr = attribute.output_schema.or_else(|| {
@ -296,10 +249,10 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
}
match &fn_item.sig.output {
syn::ReturnType::Default => {
quote! { -> std::pin::Pin<Box<dyn Future<Output = ()> + Send + #lt>> }
quote! { -> futures::future::BoxFuture<#lt, ()> }
}
syn::ReturnType::Type(_, ty) => {
quote! { -> std::pin::Pin<Box<dyn Future<Output = #ty> + Send + #lt>> }
quote! { -> futures::future::BoxFuture<#lt, #ty> }
}
}
})?;

View file

@ -62,7 +62,7 @@ http-body = { version = "1", optional = true }
http-body-util = { version = "0.1", optional = true }
bytes = { version = "1", optional = true }
# macro
rmcp-macros = { version = "0.2.1", workspace = true, optional = true }
rmcp-macros = { workspace = true, optional = true }
[target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dependencies]
chrono = { version = "0.4.38", features = ["serde"] }

View file

@ -4,6 +4,8 @@ use crate::{
service::{NotificationContext, RequestContext, RoleServer, Service, ServiceRole},
};
pub mod common;
pub mod prompt;
mod resource;
pub mod router;
pub mod tool;

View file

@ -0,0 +1,146 @@
//! Common utilities shared between tool and prompt handlers
use std::{any::TypeId, collections::HashMap, sync::Arc};
use schemars::JsonSchema;
use crate::{
RoleServer, model::JsonObject, schemars::generate::SchemaSettings, service::RequestContext,
};
/// A shortcut for generating a JSON schema for a type.
pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
// explicitly to align json schema version to official specifications.
// https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.json
// TODO: update to 2020-12 waiting for the mcp spec update
let mut settings = SchemaSettings::draft07();
settings.transforms = vec![Box::new(schemars::transform::AddNullable::default())];
let generator = settings.into_generator();
let schema = generator.into_root_schema_for::<T>();
let object = serde_json::to_value(schema).expect("failed to serialize schema");
match object {
serde_json::Value::Object(object) => object,
_ => panic!(
"Schema serialization produced non-object value: expected JSON object but got {:?}",
object
),
}
}
/// Call [`schema_for_type`] with a cache
pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
thread_local! {
static CACHE_FOR_TYPE: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
};
CACHE_FOR_TYPE.with(|cache| {
if let Some(x) = cache
.read()
.expect("schema cache lock poisoned")
.get(&TypeId::of::<T>())
{
x.clone()
} else {
let schema = schema_for_type::<T>();
let schema = Arc::new(schema);
cache
.write()
.expect("schema cache lock poisoned")
.insert(TypeId::of::<T>(), schema.clone());
schema
}
})
}
/// Trait for extracting parts from a context, unifying tool and prompt extraction
pub trait FromContextPart<C>: Sized {
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData>;
}
/// Common extractors that can be used by both tool and prompt handlers
impl<C> FromContextPart<C> for RequestContext<RoleServer>
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().clone())
}
}
impl<C> FromContextPart<C> for tokio_util::sync::CancellationToken
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().ct.clone())
}
}
impl<C> FromContextPart<C> for crate::model::Extensions
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().extensions.clone())
}
}
pub struct Extension<T>(pub T);
impl<C, T> FromContextPart<C> for Extension<T>
where
C: AsRequestContext,
T: Send + Sync + 'static + Clone,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
let extension = context
.as_request_context()
.extensions
.get::<T>()
.cloned()
.ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("missing extension {}", std::any::type_name::<T>()),
None,
)
})?;
Ok(Extension(extension))
}
}
impl<C> FromContextPart<C> for crate::Peer<RoleServer>
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().peer.clone())
}
}
impl<C> FromContextPart<C> for crate::model::Meta
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
let request_context = context.as_request_context_mut();
let mut meta = crate::model::Meta::default();
std::mem::swap(&mut meta, &mut request_context.meta);
Ok(meta)
}
}
pub struct RequestId(pub crate::model::RequestId);
impl<C> FromContextPart<C> for RequestId
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(RequestId(context.as_request_context().id.clone()))
}
}
/// Trait for types that can provide access to RequestContext
pub trait AsRequestContext {
fn as_request_context(&self) -> &RequestContext<RoleServer>;
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer>;
}

View file

@ -0,0 +1,366 @@
//! Prompt handling infrastructure for MCP servers
//!
//! This module provides the core types and traits for implementing prompt handlers
//! in MCP servers. Prompts allow servers to provide reusable templates for LLM
//! interactions with customizable arguments.
use std::{future::Future, marker::PhantomData};
use futures::future::{BoxFuture, FutureExt};
use serde::de::DeserializeOwned;
use super::common::{AsRequestContext, FromContextPart};
pub use super::common::{Extension, RequestId};
use crate::{
RoleServer,
handler::server::wrapper::Parameters,
model::{GetPromptResult, PromptMessage},
service::RequestContext,
};
/// Context for prompt retrieval operations
pub struct PromptContext<'a, S> {
pub server: &'a S,
pub name: String,
pub arguments: Option<serde_json::Map<String, serde_json::Value>>,
pub context: RequestContext<RoleServer>,
}
impl<'a, S> PromptContext<'a, S> {
pub fn new(
server: &'a S,
name: String,
arguments: Option<serde_json::Map<String, serde_json::Value>>,
context: RequestContext<RoleServer>,
) -> Self {
Self {
server,
name,
arguments,
context,
}
}
}
impl<S> AsRequestContext for PromptContext<'_, S> {
fn as_request_context(&self) -> &RequestContext<RoleServer> {
&self.context
}
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
&mut self.context
}
}
/// Trait for handling prompt retrieval
pub trait GetPromptHandler<S, A> {
fn handle(
self,
context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>;
}
/// Type alias for dynamic prompt handlers
pub type DynGetPromptHandler<S> = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
+ Send
+ Sync;
/// Adapter type for async methods that return `Vec<PromptMessage>`
pub struct AsyncMethodAdapter<T>(PhantomData<T>);
/// Adapter type for async methods with parameters that return `Vec<PromptMessage>`
pub struct AsyncMethodWithArgsAdapter<T>(PhantomData<T>);
/// Adapter types for macro-generated implementations
#[allow(clippy::type_complexity)]
pub struct AsyncPromptAdapter<P, Fut, R>(PhantomData<fn(P) -> fn(Fut) -> R>);
pub struct SyncPromptAdapter<P, R>(PhantomData<fn(P) -> R>);
pub struct AsyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
pub struct SyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
/// Trait for types that can be converted into GetPromptResult
pub trait IntoGetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData>;
}
impl IntoGetPromptResult for GetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
Ok(self)
}
}
impl IntoGetPromptResult for Vec<PromptMessage> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
Ok(GetPromptResult {
description: None,
messages: self,
})
}
}
impl<T: IntoGetPromptResult> IntoGetPromptResult for Result<T, crate::ErrorData> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
self.and_then(|v| v.into_get_prompt_result())
}
}
// Future wrapper that automatically handles IntoGetPromptResult conversion
pin_project_lite::pin_project! {
#[project = IntoGetPromptResultFutProj]
pub enum IntoGetPromptResultFut<F, R> {
Pending {
#[pin]
fut: F,
_marker: PhantomData<R>,
},
Ready {
#[pin]
result: futures::future::Ready<Result<GetPromptResult, crate::ErrorData>>,
}
}
}
impl<F, R> Future for IntoGetPromptResultFut<F, R>
where
F: Future<Output = R>,
R: IntoGetPromptResult,
{
type Output = Result<GetPromptResult, crate::ErrorData>;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
match self.project() {
IntoGetPromptResultFutProj::Pending { fut, _marker } => fut
.poll(cx)
.map(IntoGetPromptResult::into_get_prompt_result),
IntoGetPromptResultFutProj::Ready { result } => result.poll(cx),
}
}
}
// Prompt-specific extractor for prompt name
pub struct PromptName(pub String);
impl<S> FromContextPart<PromptContext<'_, S>> for PromptName {
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
Ok(Self(context.name.clone()))
}
}
// Special implementation for Parameters that handles prompt arguments
impl<S, P> FromContextPart<PromptContext<'_, S>> for Parameters<P>
where
P: DeserializeOwned,
{
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
let params = if let Some(args_map) = context.arguments.take() {
let args_value = serde_json::Value::Object(args_map);
serde_json::from_value::<P>(args_value).map_err(|e| {
crate::ErrorData::invalid_params(format!("Failed to parse parameters: {}", e), None)
})?
} else {
// Try to deserialize from empty object for optional fields
serde_json::from_value::<P>(serde_json::json!({})).map_err(|e| {
crate::ErrorData::invalid_params(
format!("Missing required parameters: {}", e),
None,
)
})?
};
Ok(Parameters(params))
}
}
// Macro to generate GetPromptHandler implementations for various parameter combinations
macro_rules! impl_prompt_handler_for {
($($T: ident)*) => {
impl_prompt_handler_for!([] [$($T)*]);
};
// finished
([$($Tn: ident)*] []) => {
impl_prompt_handler_for!(@impl $($Tn)*);
};
([$($Tn: ident)*] [$Tn_1: ident $($Rest: ident)*]) => {
impl_prompt_handler_for!(@impl $($Tn)*);
impl_prompt_handler_for!([$($Tn)* $Tn_1] [$($Rest)*]);
};
(@impl $($Tn: ident)*) => {
// Implementation for async methods (transformed by #[prompt] macro)
impl<$($Tn,)* S, F, R> GetPromptHandler<S, ($($Tn,)*)> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
)*
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R> + Send,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync + 'static,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let service = context.server;
let fut = self(service, $($Tn,)*);
async move {
let result = fut.await;
result.into_get_prompt_result()
}.boxed()
}
}
// Implementation for sync methods
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
)*
F: FnOnce(&S, $($Tn,)*) -> R + Send,
R: IntoGetPromptResult + Send,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let service = context.server;
let result = self(service, $($Tn,)*);
std::future::ready(result.into_get_prompt_result()).boxed()
}
}
// AsyncPromptAdapter - for standalone functions returning GetPromptResult
impl<$($Tn,)* S, F, Fut, R> GetPromptHandler<S, AsyncPromptAdapter<($($Tn,)*), Fut, R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
)*
F: FnOnce($($Tn,)*) -> Fut + Send + 'static,
Fut: Future<Output = Result<R, crate::ErrorData>> + Send + 'static,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync + 'static,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
// Extract all parameters before moving into the async block
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
// Since we're dealing with standalone functions that don't take &S,
// we can return a 'static future
Box::pin(async move {
let result = self($($Tn,)*).await?;
result.into_get_prompt_result()
})
}
}
// SyncPromptAdapter - for standalone sync functions returning Result
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
)*
F: FnOnce($($Tn,)*) -> Result<R, crate::ErrorData> + Send + 'static,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let result = self($($Tn,)*);
std::future::ready(result.and_then(|r| r.into_get_prompt_result())).boxed()
}
}
};
}
// Invoke the macro to generate implementations for up to 16 parameters
impl_prompt_handler_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
/// Extract prompt arguments from a type's JSON schema
/// This function analyzes the schema of a type and extracts the properties
/// as PromptArgument entries with name, description, and required status
pub fn cached_arguments_from_schema<T: schemars::JsonSchema + std::any::Any>()
-> Option<Vec<crate::model::PromptArgument>> {
let schema = super::common::cached_schema_for_type::<T>();
let schema_value = serde_json::Value::Object((*schema).clone());
let properties = schema_value.get("properties").and_then(|p| p.as_object());
if let Some(props) = properties {
let required = schema_value
.get("required")
.and_then(|r| r.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str())
.collect::<std::collections::HashSet<_>>()
})
.unwrap_or_default();
let mut arguments = Vec::new();
for (name, prop_schema) in props {
let description = prop_schema
.get("description")
.and_then(|d| d.as_str())
.map(|s| s.to_string());
arguments.push(crate::model::PromptArgument {
name: name.clone(),
description,
required: Some(required.contains(name.as_str())),
});
}
if arguments.is_empty() {
None
} else {
Some(arguments)
}
} else {
None
}
}

View file

@ -1,18 +1,21 @@
use std::sync::Arc;
use prompt::{IntoPromptRoute, PromptRoute};
use tool::{IntoToolRoute, ToolRoute};
use super::ServerHandler;
use crate::{
RoleServer, Service,
model::{ClientRequest, ListToolsResult, ServerResult},
model::{ClientRequest, ListPromptsResult, ListToolsResult, ServerResult},
service::NotificationContext,
};
pub mod prompt;
pub mod tool;
pub struct Router<S> {
pub tool_router: tool::ToolRouter<S>,
pub prompt_router: prompt::PromptRouter<S>,
pub service: Arc<S>,
}
@ -23,6 +26,7 @@ where
pub fn new(service: S) -> Self {
Self {
tool_router: tool::ToolRouter::new(),
prompt_router: prompt::PromptRouter::new(),
service: Arc::new(service),
}
}
@ -41,6 +45,21 @@ where
}
self
}
pub fn with_prompt<R, A: 'static>(mut self, route: R) -> Self
where
R: IntoPromptRoute<S, A>,
{
self.prompt_router.add_route(route.into_prompt_route());
self
}
pub fn with_prompts(mut self, routes: impl IntoIterator<Item = PromptRoute<S>>) -> Self {
for route in routes {
self.prompt_router.add_route(route);
}
self
}
}
impl<S> Service<RoleServer> for Router<S>
@ -86,6 +105,29 @@ where
next_cursor: None,
}))
}
ClientRequest::GetPromptRequest(request) => {
if self.prompt_router.has_route(request.params.name.as_ref()) {
let prompt_context = crate::handler::server::prompt::PromptContext::new(
self.service.as_ref(),
request.params.name,
request.params.arguments,
context,
);
let result = self.prompt_router.get_prompt(prompt_context).await?;
Ok(ServerResult::GetPromptResult(result))
} else {
self.service
.handle_request(ClientRequest::GetPromptRequest(request), context)
.await
}
}
ClientRequest::ListPromptsRequest(_) => {
let prompts = self.prompt_router.list_all();
Ok(ServerResult::ListPromptsResult(ListPromptsResult {
prompts,
next_cursor: None,
}))
}
rest => self.service.handle_request(rest, context).await,
}
}

View file

@ -0,0 +1,213 @@
use std::{borrow::Cow, sync::Arc};
use futures::future::BoxFuture;
use crate::{
handler::server::prompt::{DynGetPromptHandler, GetPromptHandler, PromptContext},
model::{GetPromptResult, Prompt},
};
pub struct PromptRoute<S> {
#[allow(clippy::type_complexity)]
pub get: Arc<DynGetPromptHandler<S>>,
pub attr: crate::model::Prompt,
}
impl<S> std::fmt::Debug for PromptRoute<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PromptRoute")
.field("name", &self.attr.name)
.field("description", &self.attr.description)
.field("arguments", &self.attr.arguments)
.finish()
}
}
impl<S> Clone for PromptRoute<S> {
fn clone(&self) -> Self {
Self {
get: self.get.clone(),
attr: self.attr.clone(),
}
}
}
impl<S: Send + Sync + 'static> PromptRoute<S> {
pub fn new<H, A: 'static>(attr: impl Into<Prompt>, handler: H) -> Self
where
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
{
Self {
get: Arc::new(move |context: PromptContext<S>| {
let handler = handler.clone();
handler.handle(context)
}),
attr: attr.into(),
}
}
pub fn new_dyn<H>(attr: impl Into<Prompt>, handler: H) -> Self
where
H: for<'a> Fn(
PromptContext<'a, S>,
) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
+ Send
+ Sync
+ 'static,
{
Self {
get: Arc::new(handler),
attr: attr.into(),
}
}
pub fn name(&self) -> &str {
&self.attr.name
}
}
pub trait IntoPromptRoute<S, A> {
fn into_prompt_route(self) -> PromptRoute<S>;
}
impl<S, H, A, P> IntoPromptRoute<S, A> for (P, H)
where
S: Send + Sync + 'static,
A: 'static,
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
P: Into<Prompt>,
{
fn into_prompt_route(self) -> PromptRoute<S> {
PromptRoute::new(self.0.into(), self.1)
}
}
impl<S> IntoPromptRoute<S, ()> for PromptRoute<S>
where
S: Send + Sync + 'static,
{
fn into_prompt_route(self) -> PromptRoute<S> {
self
}
}
/// Adapter for functions generated by the #\[prompt\] macro
pub struct PromptAttrGenerateFunctionAdapter;
impl<S, F> IntoPromptRoute<S, PromptAttrGenerateFunctionAdapter> for F
where
S: Send + Sync + 'static,
F: Fn() -> PromptRoute<S>,
{
fn into_prompt_route(self) -> PromptRoute<S> {
(self)()
}
}
#[derive(Debug)]
pub struct PromptRouter<S> {
#[allow(clippy::type_complexity)]
pub map: std::collections::HashMap<Cow<'static, str>, PromptRoute<S>>,
}
impl<S> Default for PromptRouter<S> {
fn default() -> Self {
Self {
map: std::collections::HashMap::new(),
}
}
}
impl<S> Clone for PromptRouter<S> {
fn clone(&self) -> Self {
Self {
map: self.map.clone(),
}
}
}
impl<S> IntoIterator for PromptRouter<S> {
type Item = PromptRoute<S>;
type IntoIter = std::collections::hash_map::IntoValues<Cow<'static, str>, PromptRoute<S>>;
fn into_iter(self) -> Self::IntoIter {
self.map.into_values()
}
}
impl<S> PromptRouter<S>
where
S: Send + Sync + 'static,
{
pub fn new() -> Self {
Self {
map: std::collections::HashMap::new(),
}
}
pub fn with_route<R, A: 'static>(mut self, route: R) -> Self
where
R: IntoPromptRoute<S, A>,
{
self.add_route(route.into_prompt_route());
self
}
pub fn add_route(&mut self, item: PromptRoute<S>) {
self.map.insert(item.attr.name.clone().into(), item);
}
pub fn merge(&mut self, other: PromptRouter<S>) {
for item in other.map.into_values() {
self.add_route(item);
}
}
pub fn remove_route(&mut self, name: &str) {
self.map.remove(name);
}
pub fn has_route(&self, name: &str) -> bool {
self.map.contains_key(name)
}
pub async fn get_prompt(
&self,
context: PromptContext<'_, S>,
) -> Result<GetPromptResult, crate::ErrorData> {
let item = self.map.get(context.name.as_str()).ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("prompt '{}' not found", context.name),
Some(serde_json::json!({
"available_prompts": self.list_all().iter().map(|p| &p.name).collect::<Vec<_>>()
})),
)
})?;
(item.get)(context).await
}
pub fn list_all(&self) -> Vec<crate::model::Prompt> {
self.map.values().map(|item| item.attr.clone()).collect()
}
}
impl<S> std::ops::Add<PromptRouter<S>> for PromptRouter<S>
where
S: Send + Sync + 'static,
{
type Output = Self;
fn add(mut self, other: PromptRouter<S>) -> Self::Output {
self.merge(other);
self
}
}
impl<S> std::ops::AddAssign<PromptRouter<S>> for PromptRouter<S>
where
S: Send + Sync + 'static,
{
fn add_assign(&mut self, other: PromptRouter<S>) {
self.merge(other);
}
}

View file

@ -1,95 +1,23 @@
//! Tool handler traits and types for MCP servers.
//!
//! This module provides the infrastructure for implementing tools that can be called
//! by MCP clients. Tools can return either unstructured content (text, images) or
//! structured JSON data with schemas.
//!
//! # Structured Output
//!
//! Tools can return structured JSON data using the [`Json`] wrapper type.
//! When using `Json<T>`, the framework will:
//! - Automatically generate a JSON schema for the output type
//! - Validate the output against the schema
//! - Return the data in the `structured_content` field of [`CallToolResult`]
//!
//! # Example
//!
//! ```rust,ignore
//! use rmcp::{tool, Json};
//! use schemars::JsonSchema;
//! use serde::{Serialize, Deserialize};
//!
//! #[derive(Serialize, Deserialize, JsonSchema)]
//! struct AnalysisResult {
//! score: f64,
//! summary: String,
//! }
//!
//! #[tool(name = "analyze")]
//! async fn analyze(&self, text: String) -> Result<Json<AnalysisResult>, String> {
//! Ok(Json(AnalysisResult {
//! score: 0.95,
//! summary: "Positive sentiment".to_string(),
//! }))
//! }
//! ```
use std::{
any::TypeId, borrow::Cow, collections::HashMap, future::Ready, marker::PhantomData, sync::Arc,
borrow::Cow,
future::{Future, Ready},
marker::PhantomData,
};
use futures::future::{BoxFuture, FutureExt};
use schemars::{JsonSchema, transform::AddNullable};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use tokio_util::sync::CancellationToken;
use serde::de::DeserializeOwned;
pub use super::router::tool::{ToolRoute, ToolRouter};
use super::common::{AsRequestContext, FromContextPart};
pub use super::{
common::{Extension, RequestId, cached_schema_for_type, schema_for_type},
router::tool::{ToolRoute, ToolRouter},
};
use crate::{
RoleServer,
handler::server::wrapper::Json,
handler::server::wrapper::Parameters,
model::{CallToolRequestParam, CallToolResult, IntoContents, JsonObject},
schemars::generate::SchemaSettings,
service::RequestContext,
};
/// A shortcut for generating a JSON schema for a type.
pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
// explicitly to align json schema version to official specifications.
// https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.json
// TODO: update to 2020-12 waiting for the mcp spec update
let mut settings = SchemaSettings::draft07();
settings.transforms = vec![Box::new(AddNullable::default())];
let generator = settings.into_generator();
let schema = generator.into_root_schema_for::<T>();
let object = serde_json::to_value(schema).expect("failed to serialize schema");
match object {
serde_json::Value::Object(object) => object,
_ => panic!("unexpected schema value"),
}
}
/// Call [`schema_for_type`] with a cache
pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
thread_local! {
static CACHE_FOR_TYPE: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
};
CACHE_FOR_TYPE.with(|cache| {
if let Some(x) = cache
.read()
.expect("schema cache lock poisoned")
.get(&TypeId::of::<T>())
{
x.clone()
} else {
let schema = schema_for_type::<T>();
let schema = Arc::new(schema);
cache
.write()
.expect("schema cache lock poisoned")
.insert(TypeId::of::<T>(), schema.clone());
schema
}
})
}
/// Deserialize a JSON object into a type
pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, crate::ErrorData> {
@ -128,32 +56,18 @@ impl<'s, S> ToolCallContext<'s, S> {
}
}
pub trait FromToolCallContextPart<S>: Sized {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData>;
impl<S> AsRequestContext for ToolCallContext<'_, S> {
fn as_request_context(&self) -> &RequestContext<RoleServer> {
&self.request_context
}
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
&mut self.request_context
}
}
/// Trait for converting tool return values into [`CallToolResult`].
///
/// This trait is automatically implemented for:
/// - Types implementing [`IntoContents`] (returns unstructured content)
/// - `Result<T, E>` where both `T` and `E` implement [`IntoContents`]
/// - [`Json<T>`](crate::handler::server::wrapper::Json) where `T` implements [`Serialize`] (returns structured content)
/// - `Result<Json<T>, E>` for structured results with errors
///
/// The `#[tool]` macro uses this trait to convert tool function return values
/// into the appropriate [`CallToolResult`] format.
pub trait IntoCallToolResult {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData>;
/// Returns the output schema for this type, if any.
///
/// This is used by the macro to automatically generate output schemas
/// for tool functions that return structured data.
fn output_schema() -> Option<Arc<JsonObject>> {
None
}
}
impl<T: IntoContents> IntoCallToolResult for T {
@ -180,40 +94,6 @@ impl<T: IntoCallToolResult> IntoCallToolResult for Result<T, crate::ErrorData> {
}
}
// Implementation for Json<T> to create structured content
impl<T: Serialize + JsonSchema + 'static> IntoCallToolResult for Json<T> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
let value = serde_json::to_value(self.0).map_err(|e| {
crate::ErrorData::internal_error(
format!("Failed to serialize structured content: {}", e),
None,
)
})?;
Ok(CallToolResult::structured(value))
}
fn output_schema() -> Option<Arc<JsonObject>> {
Some(cached_schema_for_type::<T>())
}
}
// Implementation for Result<Json<T>, E>
impl<T: Serialize + JsonSchema + 'static, E: IntoContents> IntoCallToolResult
for Result<Json<T>, E>
{
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
match self {
Ok(value) => value.into_call_tool_result(),
Err(error) => Ok(CallToolResult::error(error.into_contents())),
}
}
fn output_schema() -> Option<Arc<JsonObject>> {
Json::<T>::output_schema()
}
}
pin_project_lite::pin_project! {
#[project = IntoCallToolResultFutProj]
pub enum IntoCallToolResultFut<F, R> {
@ -266,47 +146,21 @@ pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFu
+ Send
+ Sync;
/// Parameter Extractor
///
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Parameters<P>(pub P);
impl<P: JsonSchema> JsonSchema for Parameters<P> {
fn schema_name() -> Cow<'static, str> {
P::schema_name()
}
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
P::json_schema(generator)
}
}
impl<S> FromToolCallContextPart<S> for CancellationToken {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
Ok(context.request_context.ct.clone())
}
}
// Tool-specific extractor for tool name
pub struct ToolName(pub Cow<'static, str>);
impl<S> FromToolCallContextPart<S> for ToolName {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
impl<S> FromContextPart<ToolCallContext<'_, S>> for ToolName {
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
Ok(Self(context.name.clone()))
}
}
impl<S, P> FromToolCallContextPart<S> for Parameters<P>
// Special implementation for Parameters that handles tool arguments
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
where
P: DeserializeOwned,
{
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
let arguments = context.arguments.take().unwrap_or_default();
let value: P =
serde_json::from_value(serde_json::Value::Object(arguments)).map_err(|e| {
@ -319,84 +173,14 @@ where
}
}
impl<S> FromToolCallContextPart<S> for JsonObject {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
// Special implementation for JsonObject that takes tool arguments
impl<S> FromContextPart<ToolCallContext<'_, S>> for JsonObject {
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
let object = context.arguments.take().unwrap_or_default();
Ok(object)
}
}
impl<S> FromToolCallContextPart<S> for crate::model::Extensions {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
let extensions = context.request_context.extensions.clone();
Ok(extensions)
}
}
pub struct Extension<T>(pub T);
impl<S, T> FromToolCallContextPart<S> for Extension<T>
where
T: Send + Sync + 'static + Clone,
{
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
let extension = context
.request_context
.extensions
.get::<T>()
.cloned()
.ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("missing extension {}", std::any::type_name::<T>()),
None,
)
})?;
Ok(Extension(extension))
}
}
impl<S> FromToolCallContextPart<S> for crate::Peer<RoleServer> {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
let peer = context.request_context.peer.clone();
Ok(peer)
}
}
impl<S> FromToolCallContextPart<S> for crate::model::Meta {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
let mut meta = crate::model::Meta::default();
std::mem::swap(&mut meta, &mut context.request_context.meta);
Ok(meta)
}
}
pub struct RequestId(pub crate::model::RequestId);
impl<S> FromToolCallContextPart<S> for RequestId {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
Ok(RequestId(context.request_context.id.clone()))
}
}
impl<S> FromToolCallContextPart<S> for RequestContext<RoleServer> {
fn from_tool_call_context_part(
context: &mut ToolCallContext<S>,
) -> Result<Self, crate::ErrorData> {
Ok(context.request_context.clone())
}
}
impl<'s, S> ToolCallContext<'s, S> {
pub fn invoke<H, A>(self, h: H) -> BoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
where
@ -428,7 +212,7 @@ macro_rules! impl_for {
impl<$($Tn,)* S, F, R> CallToolHandler<S, AsyncMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: FromToolCallContextPart<S> ,
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
)*
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R>,
@ -443,7 +227,7 @@ macro_rules! impl_for {
mut context: ToolCallContext<'_, S>,
) -> BoxFuture<'_, Result<CallToolResult, crate::ErrorData>>{
$(
let result = $Tn::from_tool_call_context_part(&mut context);
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
@ -461,7 +245,7 @@ macro_rules! impl_for {
impl<$($Tn,)* S, F, Fut, R> CallToolHandler<S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
where
$(
$Tn: FromToolCallContextPart<S> ,
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
)*
F: FnOnce($($Tn,)*) -> Fut + Send + ,
Fut: Future<Output = R> + Send + 'static,
@ -474,7 +258,7 @@ macro_rules! impl_for {
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>>{
$(
let result = $Tn::from_tool_call_context_part(&mut context);
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
@ -491,7 +275,7 @@ macro_rules! impl_for {
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: FromToolCallContextPart<S> + ,
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
)*
F: FnOnce(&S, $($Tn,)*) -> R + Send + ,
R: IntoCallToolResult + Send + ,
@ -503,7 +287,7 @@ macro_rules! impl_for {
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
$(
let result = $Tn::from_tool_call_context_part(&mut context);
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
@ -516,7 +300,7 @@ macro_rules! impl_for {
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: FromToolCallContextPart<S> + ,
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
)*
F: FnOnce($($Tn,)*) -> R + Send + ,
R: IntoCallToolResult + Send + ,
@ -528,7 +312,7 @@ macro_rules! impl_for {
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
$(
let result = $Tn::from_tool_call_context_part(&mut context);
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),

View file

@ -1,2 +1,4 @@
mod json;
mod parameters;
pub use json::*;
pub use parameters::*;

View file

@ -1,6 +1,12 @@
use std::borrow::Cow;
use schemars::JsonSchema;
use serde::Serialize;
use crate::{
handler::server::tool::IntoCallToolResult,
model::{CallToolResult, IntoContents},
};
/// Json wrapper for structured output
///
@ -20,3 +26,29 @@ impl<T: JsonSchema> JsonSchema for Json<T> {
T::json_schema(generator)
}
}
// Implementation for Json<T> to create structured content
impl<T: Serialize + JsonSchema + 'static> IntoCallToolResult for Json<T> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
let value = serde_json::to_value(self.0).map_err(|e| {
crate::ErrorData::internal_error(
format!("Failed to serialize structured content: {}", e),
None,
)
})?;
Ok(CallToolResult::structured(value))
}
}
// Implementation for Result<Json<T>, E>
impl<T: Serialize + JsonSchema + 'static, E: IntoContents> IntoCallToolResult
for Result<Json<T>, E>
{
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
match self {
Ok(value) => value.into_call_tool_result(),
Err(error) => Ok(CallToolResult::error(error.into_contents())),
}
}
}

View file

@ -0,0 +1,55 @@
use schemars::JsonSchema;
/// Parameter extractor for tools and prompts
///
/// When used in tool and prompt handlers, this wrapper extracts and deserializes
/// parameters from the incoming request. The framework will automatically parse
/// the JSON arguments from tool calls or prompt arguments and deserialize them
/// into the specified type `P`.
///
/// The `#[serde(transparent)]` attribute ensures that the wrapper doesn't add
/// an extra layer in the JSON structure - it directly delegates serialization
/// and deserialization to the inner type `P`.
///
/// # Usage
///
/// Use `Parameters<T>` as a parameter in your tool or prompt handler functions:
///
/// ```rust
/// # use rmcp::handler::server::wrapper::Parameters;
/// # use schemars::JsonSchema;
/// # use serde::{Deserialize, Serialize};
/// #[derive(Deserialize, JsonSchema)]
/// struct CalculationRequest {
/// operation: String,
/// a: f64,
/// b: f64,
/// }
///
/// // In a tool handler
/// async fn calculate(params: Parameters<CalculationRequest>) -> Result<String, String> {
/// let request = params.0; // Extract the inner value
/// match request.operation.as_str() {
/// "add" => Ok((request.a + request.b).to_string()),
/// _ => Err("Unknown operation".to_string()),
/// }
/// }
/// ```
///
/// The framework handles the extraction automatically:
/// - For tools: Parses the `arguments` field from tool call requests
/// - For prompts: Parses the `arguments` field from prompt requests
/// - Returns appropriate error responses if deserialization fails
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(transparent)]
pub struct Parameters<P>(pub P);
impl<P: JsonSchema> JsonSchema for Parameters<P> {
fn schema_name() -> std::borrow::Cow<'static, str> {
P::schema_name()
}
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
P::json_schema(generator)
}
}

View file

@ -53,7 +53,7 @@
//! Tools can also return structured JSON data with schemas. Use the [`Json`] wrapper:
//!
//! ```rust
//! # use rmcp::{tool, tool_router, handler::server::tool::{ToolRouter, Parameters}, Json};
//! # use rmcp::{tool, tool_router, handler::server::{tool::ToolRouter, wrapper::Parameters}, Json};
//! # use schemars::JsonSchema;
//! # use serde::{Serialize, Deserialize};
//! #

View file

@ -162,12 +162,12 @@ pub trait DynService<R: ServiceRole>: Send + Sync {
&self,
request: R::PeerReq,
context: RequestContext<R>,
) -> BoxFuture<Result<R::Resp, McpError>>;
) -> BoxFuture<'_, Result<R::Resp, McpError>>;
fn handle_notification(
&self,
notification: R::PeerNot,
context: NotificationContext<R>,
) -> BoxFuture<Result<(), McpError>>;
) -> BoxFuture<'_, Result<(), McpError>>;
fn get_info(&self) -> R::Info;
}
@ -176,14 +176,14 @@ impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
&self,
request: R::PeerReq,
context: RequestContext<R>,
) -> BoxFuture<Result<R::Resp, McpError>> {
) -> BoxFuture<'_, Result<R::Resp, McpError>> {
Box::pin(self.handle_request(request, context))
}
fn handle_notification(
&self,
notification: R::PeerNot,
context: NotificationContext<R>,
) -> BoxFuture<Result<(), McpError>> {
) -> BoxFuture<'_, Result<(), McpError>> {
Box::pin(self.handle_notification(notification, context))
}
fn get_info(&self) -> R::Info {

View file

@ -1,7 +1,7 @@
#![allow(dead_code)]
use rmcp::{
ServerHandler,
handler::server::{router::tool::ToolRouter, tool::Parameters},
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
model::{ServerCapabilities, ServerInfo},
schemars, tool, tool_router,
};

Some files were not shown because too many files have changed in this diff Show more