Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions crates/rmcp-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream {
/// | `router` | `Ident` | The name of the router function to be generated. Defaults to `tool_router`. |
/// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
/// | `server_handler` | `flag` | When set, also emits `#[::rmcp::tool_handler]` on `impl ServerHandler for Self` so you can omit a separate `#[tool_handler]` block. |
/// | `allow_empty` | `flag` | When set, accepts an impl block with no `#[tool]` fn. Without it, an empty router is a compile error. |
///
/// ## Example
///
Expand Down Expand Up @@ -122,6 +123,37 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream {
/// }
/// }
/// ```
///
/// ### Empty routers
///
/// Collecting tools is this attribute's whole purpose, so an impl block with no `#[tool]` fn is a
/// compile error rather than a router that silently serves nothing. Pass `allow_empty` when that
/// is what you want:
///
/// ```rust,ignore
/// #[tool_router(allow_empty)]
/// impl MyToolHandler {}
/// ```
///
/// The usual way to hit this by accident is a `macro_rules!` helper *inside* the impl block. An
/// attribute macro receives the unexpanded item, so `#[tool]` fns produced by such a helper are
/// invisible to `#[tool_router]`. Let the `macro_rules!` emit the whole annotated impl instead:
///
/// ```rust,ignore
/// macro_rules! define_tools {
/// ($($name:ident => $description:literal),* $(,)?) => {
/// #[tool_router]
/// impl MyToolHandler {
/// $(
/// #[tool(description = $description)]
/// async fn $name(&self) -> String { stringify!($name).to_owned() }
/// )*
/// }
/// };
/// }
///
/// define_tools!(my_tool => "what my tool does");
/// ```
#[proc_macro_attribute]
pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream {
tool_router::tool_router(attr.into(), input.into())
Expand Down
72 changes: 72 additions & 0 deletions crates/rmcp-macros/src/tool_router.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ pub struct ToolRouterAttribute {
/// When set, also emit `#[::rmcp::tool_handler]` on `impl ServerHandler for Self` so callers
/// can skip a separate `#[tool_handler]` block (expanded in a later macro pass).
pub server_handler: bool,
/// When set, accept an impl block with no `#[tool]` fn instead of reporting an error.
pub allow_empty: bool,
}

impl Default for ToolRouterAttribute {
Expand All @@ -25,6 +27,7 @@ impl Default for ToolRouterAttribute {
router: format_ident!("tool_router"),
vis: None,
server_handler: false,
allow_empty: false,
}
}
}
Expand All @@ -35,6 +38,7 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result<TokenSt
router,
vis,
server_handler,
allow_empty,
} = ToolRouterAttribute::from_list(&attr_args)?;
let mut item_impl = syn::parse2::<ItemImpl>(input)?;
// find all function marked with `#[rmcp::tool]`
Expand All @@ -58,6 +62,19 @@ pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result<TokenSt
}
})
.collect();
if tool_attr_fns.is_empty() && !allow_empty {
return Err(syn::Error::new_spanned(
&item_impl.self_ty,
format!(
"`#[tool_router]` found no `#[tool]` fn in this impl block, so `Self::{router}()` \
would serve no tools\n\
note: a `macro_rules!` invocation inside the impl block is not expanded before \
this attribute runs, so any `#[tool]` fn it generates is invisible here; let the \
`macro_rules!` emit the whole `#[tool_router] impl` instead\n\
note: use `#[tool_router(allow_empty)]` if an empty router is intended"
),
));
}
let mut routers = Vec::with_capacity(tool_attr_fns.len());
for handler in tool_attr_fns {
let tool_attr_fn_ident = format_ident!("{handler}_tool_attr");
Expand Down Expand Up @@ -111,13 +128,18 @@ mod test {
router,
vis,
server_handler,
allow_empty,
} = ToolRouterAttribute::from_list(&attr_args)?;
assert_eq!(router.to_string(), "test_router");
assert!(vis.is_some(), "vis = \"pub(crate)\" should parse");
assert!(
!server_handler,
"server_handler should default to false when omitted"
);
assert!(
!allow_empty,
"allow_empty should default to false when omitted"
);
Ok(())
}

Expand All @@ -137,4 +159,54 @@ mod test {
assert!(server_handler);
Ok(())
}

#[test]
fn tool_router_rejects_an_impl_without_directly_visible_tool_fns() {
let input = quote! {
impl Probe {
a_capability!(broadcast = "what a capability would own");
}
};
let message = tool_router(TokenStream::new(), input)
.expect_err("an impl with no `#[tool]` fn should not compile")
.to_string();
assert!(
message.contains("`Self::tool_router()` would serve no tools"),
"{message}"
);
assert!(message.contains("macro_rules!"), "{message}");
assert!(message.contains("allow_empty"), "{message}");
}

#[test]
fn tool_router_names_the_custom_router_fn_when_it_is_empty() {
let message = tool_router(quote! { router = custom_router }, quote! { impl Probe {} })
.expect_err("an impl with no `#[tool]` fn should not compile")
.to_string();
assert!(
message.contains("`Self::custom_router()` would serve no tools"),
"{message}"
);
}

#[test]
fn tool_router_accepts_an_impl_with_a_tool_fn() -> syn::Result<()> {
let input = quote! {
impl Probe {
#[tool(description = "probe")]
async fn probe(&self) -> String { "probed".to_owned() }
}
};
let generated = tool_router(TokenStream::new(), input)?.to_string();
assert!(generated.contains("with_route"), "{generated}");
Ok(())
}

#[test]
fn tool_router_allow_empty_generates_a_router_without_routes() -> syn::Result<()> {
let generated = tool_router(quote! { allow_empty }, quote! { impl Probe {} })?.to_string();
assert!(generated.contains("fn tool_router"), "{generated}");
assert!(!generated.contains("with_route"), "{generated}");
Ok(())
}
}
44 changes: 44 additions & 0 deletions crates/rmcp/tests/test_tool_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -573,3 +573,47 @@ fn test_manual_get_info_not_overridden() {
"manual resources should be preserved"
);
}

/// Server whose tools come from a `macro_rules!` helper wrapping the whole annotated impl.
#[derive(Debug, Clone)]
struct MacroGeneratedServer;

macro_rules! define_tools {
($($name:ident => $description:literal),* $(,)?) => {
#[tool_router]
impl MacroGeneratedServer {
$(
#[tool(description = $description)]
async fn $name(&self) -> String {
stringify!($name).to_owned()
}
)*
}
};
}

define_tools!(probe => "what a capability would own");

#[test]
fn test_macro_rules_around_the_impl_registers_tools() {
let tools = MacroGeneratedServer::tool_router().list_all();

assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "probe");
assert_eq!(
tools[0].description.as_deref(),
Some("what a capability would own")
);
}

/// Server that opts in to a router with no tools.
#[derive(Debug, Clone)]
struct EmptyRouterServer;

#[tool_router(allow_empty)]
impl EmptyRouterServer {}

#[test]
fn test_allow_empty_builds_a_router_without_tools() {
assert!(EmptyRouterServer::tool_router().list_all().is_empty());
}