From 12b4310e00628c5aaf8492b8d80b711fad95c06e Mon Sep 17 00:00:00 2001 From: Abhay Agarwal Date: Sat, 17 May 2025 16:44:01 -0700 Subject: [PATCH] feat(server): Add annotation to tool macro skip serializing the tool annotations if none, fixes mcp inspector --- crates/rmcp-macros/Cargo.toml | 1 + crates/rmcp-macros/src/tool.rs | 62 ++++++++++++++++++- crates/rmcp/src/model/tool.rs | 5 ++ .../rmcp/tests/test_tool_macro_annotations.rs | 59 ++++++++++++++++++ 4 files changed, 124 insertions(+), 3 deletions(-) create mode 100644 crates/rmcp/tests/test_tool_macro_annotations.rs diff --git a/crates/rmcp-macros/Cargo.toml b/crates/rmcp-macros/Cargo.toml index 3f46bdf14..9afa8b695 100644 --- a/crates/rmcp-macros/Cargo.toml +++ b/crates/rmcp-macros/Cargo.toml @@ -18,6 +18,7 @@ proc-macro = true syn = {version = "2", features = ["full"]} quote = "1" proc-macro2 = "1" +serde_json = "1.0" [features] diff --git a/crates/rmcp-macros/src/tool.rs b/crates/rmcp-macros/src/tool.rs index 5425412ce..c17ef2f80 100644 --- a/crates/rmcp-macros/src/tool.rs +++ b/crates/rmcp-macros/src/tool.rs @@ -2,11 +2,45 @@ use std::collections::HashSet; use proc_macro2::TokenStream; use quote::{ToTokens, quote}; +use serde_json::json; use syn::{ - Expr, FnArg, Ident, ItemFn, ItemImpl, MetaList, PatType, Token, Type, Visibility, parse::Parse, - parse_quote, spanned::Spanned, + Expr, FnArg, Ident, ItemFn, ItemImpl, Lit, MetaList, PatType, Token, Type, Visibility, + parse::Parse, parse_quote, spanned::Spanned, }; +/// Stores tool annotation attributes +#[derive(Default, Clone)] +struct ToolAnnotationAttrs(pub serde_json::Map); + +impl Parse for ToolAnnotationAttrs { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let mut attrs = serde_json::Map::new(); + + while !input.is_empty() { + let key: Ident = input.parse()?; + input.parse::()?; + let value: Lit = input.parse()?; + let value = match value { + Lit::Str(s) => json!(s.value()), + Lit::Bool(b) => json!(b.value), + _ => { + return Err(syn::Error::new( + key.span(), + "annotations must be string or boolean literals", + )); + } + }; + attrs.insert(key.to_string(), value); + if input.is_empty() { + break; + } + input.parse::()?; + } + + Ok(ToolAnnotationAttrs(attrs)) + } +} + #[derive(Default)] struct ToolImplItemAttrs { tool_box: Option>, @@ -45,6 +79,7 @@ struct ToolFnItemAttrs { name: Option, description: Option, vis: Option, + annotations: Option, } impl Parse for ToolFnItemAttrs { @@ -52,6 +87,8 @@ impl Parse for ToolFnItemAttrs { let mut name = None; let mut description = None; let mut vis = None; + let mut annotations = None; + while !input.is_empty() { let key: Ident = input.parse()?; input.parse::()?; @@ -68,6 +105,13 @@ impl Parse for ToolFnItemAttrs { let value: Visibility = input.parse()?; vis = Some(value); } + "annotations" => { + // Parse the annotations as a nested structure + let content; + syn::braced!(content in input); + let value = content.parse()?; + annotations = Some(value); + } _ => { return Err(syn::Error::new(key.span(), "unknown attribute")); } @@ -82,6 +126,7 @@ impl Parse for ToolFnItemAttrs { name, description, vis, + annotations, }) } } @@ -470,6 +515,17 @@ pub(crate) fn tool_fn_item(attr: TokenStream, mut input_fn: ItemFn) -> syn::Resu }; let input_fn_attrs = &input_fn.attrs; let input_fn_vis = &input_fn.vis; + + let annotations_code = if let Some(annotations) = &tool_macro_attrs.fn_item.annotations { + let annotations = + serde_json::to_string(&annotations.0).expect("failed to serialize annotations"); + quote! { + Some(serde_json::from_str::(&#annotations).expect("Could not parse tool annotations")) + } + } else { + quote! { None } + }; + quote! { #(#input_fn_attrs)* #input_fn_vis fn #tool_attr_fn_ident() -> rmcp::model::Tool { @@ -477,7 +533,7 @@ pub(crate) fn tool_fn_item(attr: TokenStream, mut input_fn: ItemFn) -> syn::Resu name: #name.into(), description: Some(#description.into()), input_schema: #schema.into(), - annotations: None + annotations: #annotations_code, } } } diff --git a/crates/rmcp/src/model/tool.rs b/crates/rmcp/src/model/tool.rs index d392f9b85..e24cf44c1 100644 --- a/crates/rmcp/src/model/tool.rs +++ b/crates/rmcp/src/model/tool.rs @@ -37,11 +37,13 @@ pub struct Tool { #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] pub struct ToolAnnotations { /// A human-readable title for the tool. + #[serde(skip_serializing_if = "Option::is_none")] pub title: Option, /// If true, the tool does not modify its environment. /// /// Default: false + #[serde(skip_serializing_if = "Option::is_none")] pub read_only_hint: Option, /// If true, the tool may perform destructive updates to its environment. @@ -51,6 +53,7 @@ pub struct ToolAnnotations { /// /// Default: true /// A human-readable description of the tool's purpose. + #[serde(skip_serializing_if = "Option::is_none")] pub destructive_hint: Option, /// If true, calling the tool repeatedly with the same arguments @@ -59,6 +62,7 @@ pub struct ToolAnnotations { /// (This property is meaningful only when `readOnlyHint == false`) /// /// Default: false. + #[serde(skip_serializing_if = "Option::is_none")] pub idempotent_hint: Option, /// If true, this tool may interact with an "open world" of external @@ -67,6 +71,7 @@ pub struct ToolAnnotations { /// of a memory tool is not. /// /// Default: true + #[serde(skip_serializing_if = "Option::is_none")] pub open_world_hint: Option, } diff --git a/crates/rmcp/tests/test_tool_macro_annotations.rs b/crates/rmcp/tests/test_tool_macro_annotations.rs new file mode 100644 index 000000000..0b57dffd6 --- /dev/null +++ b/crates/rmcp/tests/test_tool_macro_annotations.rs @@ -0,0 +1,59 @@ +#[cfg(test)] +mod tests { + use rmcp::{ServerHandler, tool}; + + #[derive(Debug, Clone, Default)] + pub struct AnnotatedServer {} + + impl AnnotatedServer { + // Tool with inline comments for documentation + /// Direct annotation test tool + /// This is used to test tool annotations + #[tool( + name = "direct-annotated-tool", + annotations = { + title: "Annotated Tool", + readOnlyHint: true + } + )] + pub async fn direct_annotated_tool(&self, #[tool(param)] input: String) -> String { + format!("Direct: {}", input) + } + } + + impl ServerHandler for AnnotatedServer { + async fn call_tool( + &self, + request: rmcp::model::CallToolRequestParam, + context: rmcp::service::RequestContext, + ) -> Result { + let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context); + match tcc.name() { + "direct-annotated-tool" => Self::direct_annotated_tool_tool_call(tcc).await, + _ => Err(rmcp::Error::invalid_params("method not found", None)), + } + } + } + + #[test] + fn test_direct_tool_attributes() { + // Get the tool definition + let tool = AnnotatedServer::direct_annotated_tool_tool_attr(); + + // Verify basic properties + assert_eq!(tool.name, "direct-annotated-tool"); + + // Verify description is extracted from doc comments + assert!(tool.description.is_some()); + assert!( + tool.description + .as_ref() + .unwrap() + .contains("Direct annotation test tool") + ); + + let annotations = tool.annotations.unwrap(); + assert_eq!(annotations.title.as_ref().unwrap(), "Annotated Tool"); + assert_eq!(annotations.read_only_hint, Some(true)); + } +}