332 lines
11 KiB
Rust
332 lines
11 KiB
Rust
#![doc = include_str!("../README.md")]
|
|
|
|
#[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;
|
|
/// # tool
|
|
///
|
|
/// This macro is used to mark a function as a tool handler.
|
|
///
|
|
/// This will generate a function that return the attribute of this tool, with type `rmcp::model::Tool`.
|
|
///
|
|
/// ## Usage
|
|
///
|
|
/// | field | type | usage |
|
|
/// | :- | :- | :- |
|
|
/// | `name` | `String` | The name of the tool. If not provided, it defaults to the function name. |
|
|
/// | `description` | `String` | A description of the tool. The document of this function will be used. |
|
|
/// | `input_schema` | `Expr` | A JSON Schema object defining the expected parameters for the tool. If not provide, if will use the json schema of its argument with type `Parameters<T>` |
|
|
/// | `annotations` | `ToolAnnotationsAttribute` | Additional tool information. Defaults to `None`. |
|
|
///
|
|
/// ## Example
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool(name = "my_tool", description = "This is my tool", annotations(title = "我的工具", read_only_hint = true))]
|
|
/// pub async fn my_tool(param: Parameters<MyToolParam>) {
|
|
/// // handling tool request
|
|
/// }
|
|
/// ```
|
|
#[proc_macro_attribute]
|
|
pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream {
|
|
tool::tool(attr.into(), input.into())
|
|
.unwrap_or_else(|err| err.to_compile_error())
|
|
.into()
|
|
}
|
|
|
|
/// # tool_router
|
|
///
|
|
/// This macro is used to generate a tool router based on functions marked with `#[rmcp::tool]` in an implementation block.
|
|
///
|
|
/// It creates a function that returns a `ToolRouter` instance.
|
|
///
|
|
/// The generated function is used by `#[tool_handler]` by default (via `Self::tool_router()`),
|
|
/// so in most cases you do not need to store the router in a field.
|
|
///
|
|
/// ## Usage
|
|
///
|
|
/// | field | type | usage |
|
|
/// | :- | :- | :- |
|
|
/// | `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
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool_router]
|
|
/// impl MyToolHandler {
|
|
/// #[tool]
|
|
/// pub fn my_tool() {
|
|
///
|
|
/// }
|
|
/// }
|
|
///
|
|
/// // #[tool_handler] calls Self::tool_router() automatically
|
|
/// #[tool_handler]
|
|
/// impl ServerHandler for MyToolHandler {}
|
|
/// ```
|
|
///
|
|
/// ### Eliding `#[tool_handler]`
|
|
///
|
|
/// For a tools-only server, pass `server_handler` so the `impl ServerHandler` block is not written by hand:
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool_router(server_handler)]
|
|
/// impl MyToolHandler {
|
|
/// #[tool]
|
|
/// fn my_tool() {}
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// This expands in two steps: first `#[tool_router]` emits the inherent impl plus
|
|
/// `#[::rmcp::tool_handler] impl ServerHandler for MyToolHandler {}`, then `#[tool_handler]`
|
|
/// fills in `call_tool`, `list_tools`, `get_info`, and related methods. If you combine tools with
|
|
/// prompts or tasks on the **same** `impl ServerHandler` block (stacked `#[tool_handler]` /
|
|
/// `#[prompt_handler]` attributes), keep using an explicit `#[tool_handler]` impl instead of `server_handler`.
|
|
///
|
|
/// Or specify the visibility and router name, which would be helpful when you want to combine multiple routers into one:
|
|
///
|
|
/// ```rust,ignore
|
|
/// mod a {
|
|
/// #[tool_router(router = tool_router_a, vis = "pub")]
|
|
/// impl MyToolHandler {
|
|
/// #[tool]
|
|
/// fn my_tool_a() {
|
|
///
|
|
/// }
|
|
/// }
|
|
/// }
|
|
///
|
|
/// mod b {
|
|
/// #[tool_router(router = tool_router_b, vis = "pub")]
|
|
/// impl MyToolHandler {
|
|
/// #[tool]
|
|
/// fn my_tool_b() {
|
|
///
|
|
/// }
|
|
/// }
|
|
/// }
|
|
///
|
|
/// impl MyToolHandler {
|
|
/// fn new() -> Self {
|
|
/// Self {
|
|
/// tool_router: self::tool_router_a() + self::tool_router_b(),
|
|
/// }
|
|
/// }
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// ### 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())
|
|
.unwrap_or_else(|err| err.to_compile_error())
|
|
.into()
|
|
}
|
|
|
|
/// # tool_handler
|
|
///
|
|
/// This macro generates the `call_tool`, `list_tools`, `get_tool`, and (optionally)
|
|
/// `get_info` methods for a `ServerHandler` implementation, using a `ToolRouter`.
|
|
///
|
|
/// ## Usage
|
|
///
|
|
/// | field | type | usage |
|
|
/// | :- | :- | :- |
|
|
/// | `router` | `Expr` | The expression to access the `ToolRouter` instance. Defaults to `Self::tool_router()`. |
|
|
/// | `meta` | `Expr` | Optional metadata for `ListToolsResult`. |
|
|
/// | `name` | `String` | Custom server name. Defaults to `CARGO_CRATE_NAME`. |
|
|
/// | `version` | `String` | Custom server version. Defaults to `CARGO_PKG_VERSION`. |
|
|
/// | `instructions` | `String` | Optional human-readable instructions about using this server. |
|
|
///
|
|
/// ## Minimal example (no boilerplate)
|
|
///
|
|
/// The macro automatically generates `get_info()` with tools capability enabled
|
|
/// and reads the server name/version from `Cargo.toml`:
|
|
///
|
|
/// ```rust,ignore
|
|
/// struct TimeServer;
|
|
///
|
|
/// #[tool_router]
|
|
/// impl TimeServer {
|
|
/// #[tool(description = "Get current time")]
|
|
/// async fn get_time(&self) -> String { "12:00".into() }
|
|
/// }
|
|
///
|
|
/// #[tool_handler]
|
|
/// impl ServerHandler for TimeServer {}
|
|
/// ```
|
|
///
|
|
/// ## Custom server info
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool_handler(name = "my-server", version = "1.0.0", instructions = "A helpful server")]
|
|
/// impl ServerHandler for MyToolHandler {}
|
|
/// ```
|
|
///
|
|
/// ## Custom router expression
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool_handler(router = self.tool_router)]
|
|
/// impl ServerHandler for MyToolHandler {
|
|
/// // ...implement other handler
|
|
/// }
|
|
/// ```
|
|
///
|
|
/// ## Manual `get_info()`
|
|
///
|
|
/// If you provide your own `get_info()`, the macro will not generate one:
|
|
///
|
|
/// ```rust,ignore
|
|
/// #[tool_handler]
|
|
/// impl ServerHandler for MyToolHandler {
|
|
/// fn get_info(&self) -> ServerConfig {
|
|
/// ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
|
|
/// }
|
|
/// }
|
|
/// ```
|
|
#[proc_macro_attribute]
|
|
pub fn tool_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
|
|
tool_handler::tool_handler(attr.into(), input.into())
|
|
.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 a `PromptRouter`. It also auto-generates `get_info()`
|
|
/// with prompts capability enabled if not already provided.
|
|
///
|
|
/// ## Usage
|
|
///
|
|
/// | field | type | usage |
|
|
/// | :- | :- | :- |
|
|
/// | `router` | `Expr` | The expression to access the `PromptRouter` instance. Defaults to `Self::prompt_router()`. |
|
|
/// | `meta` | `Expr` | Optional metadata for `ListPromptsResult`. |
|
|
///
|
|
/// ## 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.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()
|
|
}
|