Skip to content

Commit 9349f5c

Browse files
authored
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 <[email protected]> * 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 <[email protected]> * 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
1 parent c0452f9 commit 9349f5c

40 files changed

+3191
-464
lines changed

.gitignore

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,16 @@ Cargo.lock
1313
# MSVC Windows builds of rustc generate these, which store debugging information
1414
*.pdb
1515
.vscode/
16+
17+
# Python artifacts (for test directories)
18+
*.egg-info/
19+
__pycache__/
20+
*.pyc
21+
*.pyo
22+
1623
# RustRover
1724
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
1825
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
1926
# and can be added to the global gitignore or merged into this file. For a more nuclear
2027
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
21-
#.idea/
28+
#.idea/

crates/rmcp-macros/src/common.rs

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
//! Common utilities shared between different macro implementations
2+
3+
use quote::quote;
4+
use syn::{Attribute, Expr, FnArg, ImplItemFn, Signature, Type};
5+
6+
/// Parse a None expression
7+
pub fn none_expr() -> syn::Result<Expr> {
8+
syn::parse2::<Expr>(quote! { None })
9+
}
10+
11+
/// Extract documentation from doc attributes
12+
pub fn extract_doc_line(existing_docs: Option<String>, attr: &Attribute) -> Option<String> {
13+
if !attr.path().is_ident("doc") {
14+
return None;
15+
}
16+
17+
let syn::Meta::NameValue(name_value) = &attr.meta else {
18+
return None;
19+
};
20+
21+
let syn::Expr::Lit(expr_lit) = &name_value.value else {
22+
return None;
23+
};
24+
25+
let syn::Lit::Str(lit_str) = &expr_lit.lit else {
26+
return None;
27+
};
28+
29+
let content = lit_str.value().trim().to_string();
30+
match (existing_docs, content) {
31+
(Some(mut existing_docs), content) if !content.is_empty() => {
32+
existing_docs.push('\n');
33+
existing_docs.push_str(&content);
34+
Some(existing_docs)
35+
}
36+
(Some(existing_docs), _) => Some(existing_docs),
37+
(None, content) if !content.is_empty() => Some(content),
38+
_ => None,
39+
}
40+
}
41+
42+
/// Find Parameters<T> type in function signature
43+
/// Returns the full Parameters<T> type if found
44+
pub fn find_parameters_type_in_sig(sig: &Signature) -> Option<Box<Type>> {
45+
sig.inputs.iter().find_map(|input| {
46+
if let FnArg::Typed(pat_type) = input {
47+
if let Type::Path(type_path) = &*pat_type.ty {
48+
if type_path
49+
.path
50+
.segments
51+
.last()
52+
.is_some_and(|type_name| type_name.ident == "Parameters")
53+
{
54+
return Some(pat_type.ty.clone());
55+
}
56+
}
57+
}
58+
None
59+
})
60+
}
61+
62+
/// Find Parameters<T> type in ImplItemFn
63+
pub fn find_parameters_type_impl(fn_item: &ImplItemFn) -> Option<Box<Type>> {
64+
find_parameters_type_in_sig(&fn_item.sig)
65+
}

crates/rmcp-macros/src/lib.rs

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
#[allow(unused_imports)]
22
use proc_macro::TokenStream;
33

4+
mod common;
5+
mod prompt;
6+
mod prompt_handler;
7+
mod prompt_router;
48
mod tool;
59
mod tool_handler;
610
mod tool_router;
@@ -160,3 +164,102 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
160164
.unwrap_or_else(|err| err.to_compile_error())
161165
.into()
162166
}
167+
168+
/// # prompt
169+
///
170+
/// This macro is used to mark a function as a prompt handler.
171+
///
172+
/// This will generate a function that returns the attribute of this prompt, with type `rmcp::model::Prompt`.
173+
///
174+
/// ## Usage
175+
///
176+
/// | field | type | usage |
177+
/// | :- | :- | :- |
178+
/// | `name` | `String` | The name of the prompt. If not provided, it defaults to the function name. |
179+
/// | `description` | `String` | A description of the prompt. The document of this function will be used if not provided. |
180+
/// | `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. |
181+
///
182+
/// ## Example
183+
///
184+
/// ```rust,ignore
185+
/// #[prompt(name = "code_review", description = "Reviews code for best practices")]
186+
/// pub async fn code_review_prompt(&self, Parameters(args): Parameters<CodeReviewArgs>) -> Result<Vec<PromptMessage>> {
187+
/// // Generate prompt messages based on arguments
188+
/// }
189+
/// ```
190+
#[proc_macro_attribute]
191+
pub fn prompt(attr: TokenStream, input: TokenStream) -> TokenStream {
192+
prompt::prompt(attr.into(), input.into())
193+
.unwrap_or_else(|err| err.to_compile_error())
194+
.into()
195+
}
196+
197+
/// # prompt_router
198+
///
199+
/// This macro generates a prompt router based on functions marked with `#[rmcp::prompt]` in an implementation block.
200+
///
201+
/// It creates a function that returns a `PromptRouter` instance.
202+
///
203+
/// ## Usage
204+
///
205+
/// | field | type | usage |
206+
/// | :- | :- | :- |
207+
/// | `router` | `Ident` | The name of the router function to be generated. Defaults to `prompt_router`. |
208+
/// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
209+
///
210+
/// ## Example
211+
///
212+
/// ```rust,ignore
213+
/// #[prompt_router]
214+
/// impl MyPromptHandler {
215+
/// #[prompt]
216+
/// pub async fn greeting_prompt(&self, Parameters(args): Parameters<GreetingArgs>) -> Result<Vec<PromptMessage>, Error> {
217+
/// // Generate greeting prompt using args
218+
/// }
219+
///
220+
/// pub fn new() -> Self {
221+
/// Self {
222+
/// // the default name of prompt router will be `prompt_router`
223+
/// prompt_router: Self::prompt_router(),
224+
/// }
225+
/// }
226+
/// }
227+
/// ```
228+
#[proc_macro_attribute]
229+
pub fn prompt_router(attr: TokenStream, input: TokenStream) -> TokenStream {
230+
prompt_router::prompt_router(attr.into(), input.into())
231+
.unwrap_or_else(|err| err.to_compile_error())
232+
.into()
233+
}
234+
235+
/// # prompt_handler
236+
///
237+
/// This macro generates handler methods for `get_prompt` and `list_prompts` in the implementation block, using an existing `PromptRouter` instance.
238+
///
239+
/// ## Usage
240+
///
241+
/// | field | type | usage |
242+
/// | :- | :- | :- |
243+
/// | `router` | `Expr` | The expression to access the `PromptRouter` instance. Defaults to `self.prompt_router`. |
244+
///
245+
/// ## Example
246+
/// ```rust,ignore
247+
/// #[prompt_handler]
248+
/// impl ServerHandler for MyPromptHandler {
249+
/// // ...implement other handler methods
250+
/// }
251+
/// ```
252+
///
253+
/// or using a custom router expression:
254+
/// ```rust,ignore
255+
/// #[prompt_handler(router = self.get_prompt_router())]
256+
/// impl ServerHandler for MyPromptHandler {
257+
/// // ...implement other handler methods
258+
/// }
259+
/// ```
260+
#[proc_macro_attribute]
261+
pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
262+
prompt_handler::prompt_handler(attr.into(), input.into())
263+
.unwrap_or_else(|err| err.to_compile_error())
264+
.into()
265+
}

crates/rmcp-macros/src/prompt.rs

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
use darling::{FromMeta, ast::NestedMeta};
2+
use proc_macro2::TokenStream;
3+
use quote::{format_ident, quote};
4+
use syn::{Expr, Ident, ImplItemFn, ReturnType};
5+
6+
use crate::common::{extract_doc_line, none_expr};
7+
8+
#[derive(FromMeta, Default, Debug)]
9+
#[darling(default)]
10+
pub struct PromptAttribute {
11+
/// The name of the prompt
12+
pub name: Option<String>,
13+
/// Optional description of what the prompt does
14+
pub description: Option<String>,
15+
/// Arguments that can be passed to the prompt
16+
pub arguments: Option<Expr>,
17+
}
18+
19+
pub struct ResolvedPromptAttribute {
20+
pub name: String,
21+
pub description: Option<String>,
22+
pub arguments: Expr,
23+
}
24+
25+
impl ResolvedPromptAttribute {
26+
pub fn into_fn(self, fn_ident: Ident) -> syn::Result<ImplItemFn> {
27+
let Self {
28+
name,
29+
description,
30+
arguments,
31+
} = self;
32+
let description = if let Some(description) = description {
33+
quote! { Some(#description.into()) }
34+
} else {
35+
quote! { None }
36+
};
37+
let tokens = quote! {
38+
pub fn #fn_ident() -> rmcp::model::Prompt {
39+
rmcp::model::Prompt {
40+
name: #name.into(),
41+
description: #description,
42+
arguments: #arguments,
43+
}
44+
}
45+
};
46+
syn::parse2::<ImplItemFn>(tokens)
47+
}
48+
}
49+
50+
pub fn prompt(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
51+
let attribute = if attr.is_empty() {
52+
Default::default()
53+
} else {
54+
let attr_args = NestedMeta::parse_meta_list(attr)?;
55+
PromptAttribute::from_list(&attr_args)?
56+
};
57+
let mut fn_item = syn::parse2::<ImplItemFn>(input.clone())?;
58+
let fn_ident = &fn_item.sig.ident;
59+
60+
let prompt_attr_fn_ident = format_ident!("{}_prompt_attr", fn_ident);
61+
62+
// Try to find prompt parameters from function parameters
63+
let arguments_expr = if let Some(arguments) = attribute.arguments {
64+
arguments
65+
} else {
66+
// Look for a type named Parameters in the function signature
67+
let params_ty = crate::common::find_parameters_type_impl(&fn_item);
68+
69+
if let Some(params_ty) = params_ty {
70+
// Generate arguments from the type's schema with caching
71+
syn::parse2::<Expr>(quote! {
72+
rmcp::handler::server::prompt::cached_arguments_from_schema::<#params_ty>()
73+
})?
74+
} else {
75+
// No arguments
76+
none_expr()?
77+
}
78+
};
79+
80+
let name = attribute.name.unwrap_or_else(|| fn_ident.to_string());
81+
let description = attribute
82+
.description
83+
.or_else(|| fn_item.attrs.iter().fold(None, extract_doc_line));
84+
let arguments = arguments_expr;
85+
86+
let resolved_prompt_attr = ResolvedPromptAttribute {
87+
name: name.clone(),
88+
description: description.clone(),
89+
arguments: arguments.clone(),
90+
};
91+
let prompt_attr_fn = resolved_prompt_attr.into_fn(prompt_attr_fn_ident.clone())?;
92+
93+
// Modify the input function for async support (same as tool macro)
94+
if fn_item.sig.asyncness.is_some() {
95+
// 1. remove asyncness from sig
96+
// 2. make return type: `futures::future::BoxFuture<'_, #ReturnType>`
97+
// 3. make body: { Box::pin(async move { #body }) }
98+
let new_output = syn::parse2::<ReturnType>({
99+
let mut lt = quote! { 'static };
100+
if let Some(receiver) = fn_item.sig.receiver() {
101+
if let Some((_, receiver_lt)) = receiver.reference.as_ref() {
102+
if let Some(receiver_lt) = receiver_lt {
103+
lt = quote! { #receiver_lt };
104+
} else {
105+
lt = quote! { '_ };
106+
}
107+
}
108+
}
109+
match &fn_item.sig.output {
110+
syn::ReturnType::Default => {
111+
quote! { -> futures::future::BoxFuture<#lt, ()> }
112+
}
113+
syn::ReturnType::Type(_, ty) => {
114+
quote! { -> futures::future::BoxFuture<#lt, #ty> }
115+
}
116+
}
117+
})?;
118+
let prev_block = &fn_item.block;
119+
let new_block = syn::parse2::<syn::Block>(quote! {
120+
{ Box::pin(async move #prev_block ) }
121+
})?;
122+
fn_item.sig.asyncness = None;
123+
fn_item.sig.output = new_output;
124+
fn_item.block = new_block;
125+
}
126+
127+
Ok(quote! {
128+
#prompt_attr_fn
129+
#fn_item
130+
})
131+
}
132+
133+
#[cfg(test)]
134+
mod test {
135+
use super::*;
136+
137+
#[test]
138+
fn test_prompt_macro() -> syn::Result<()> {
139+
let attr = quote! {
140+
name = "example-prompt",
141+
description = "An example prompt"
142+
};
143+
let input = quote! {
144+
async fn example_prompt(&self, Parameters(args): Parameters<ExampleArgs>) -> Result<String> {
145+
Ok("Example prompt response".to_string())
146+
}
147+
};
148+
let result = prompt(attr, input)?;
149+
150+
// Verify the output contains both the attribute function and the modified function
151+
let result_str = result.to_string();
152+
assert!(result_str.contains("example_prompt_prompt_attr"));
153+
assert!(
154+
result_str.contains("rmcp")
155+
&& result_str.contains("model")
156+
&& result_str.contains("Prompt")
157+
);
158+
159+
Ok(())
160+
}
161+
162+
#[test]
163+
fn test_doc_comment_description() -> syn::Result<()> {
164+
let attr = quote! {}; // No explicit description
165+
let input = quote! {
166+
/// This is a test prompt description
167+
/// with multiple lines
168+
fn test_prompt(&self) -> Result<String> {
169+
Ok("Test".to_string())
170+
}
171+
};
172+
let result = prompt(attr, input)?;
173+
174+
// The output should contain the description from doc comments
175+
let result_str = result.to_string();
176+
assert!(result_str.contains("This is a test prompt description"));
177+
assert!(result_str.contains("with multiple lines"));
178+
179+
Ok(())
180+
}
181+
}

0 commit comments

Comments
 (0)