Enhancement: Refactoring tool macro #195

jokemanfire · 2025-05-19T01:50:08Z

Due to some historical reasons, there are some different logical operations in the macro of the tool for template and non template operations. The purpose of this refactoring is to unify the operations of these two different types of macros and make the tool macro support multiple tool macro declarations for structures. For the implementation of service macros, it may be necessary to split them.

4t145 · 2025-05-19T05:23:23Z

The design of macros do requires thorough discussion.

jokemanfire · 2025-05-19T06:21:54Z

I will provide a specific plan and demo in the next few days, and we will review them together later.

jokemanfire · 2025-05-19T07:38:21Z

I think the tool macro doesn't need to be too complex or even use the static variable TOOL_BOX. I think the entire tool macro can be implemented like codegen, which can easily cover template and non template types, and the code will be very simple and obvious. What do you think @4t145

4t145 · 2025-05-19T08:56:06Z

There's a gap: the tool router implementaion is in trait Service, while the tool defination is in the impl block.

I came up with two solutions:

One solution is to generate some router functions, so user can call it from handler implementation.

The another one, is to create a static type with shared state just like build a service in axum. And then we can discard the macro.

4t145 · 2025-05-19T09:30:45Z

#62 #159

4t145 · 2025-05-19T09:46:39Z

For current CallToolHandler<'a, S, A>, we can remove the lifetime a. Since we always call on Arc<S>
We can have a router type Router<S = ()>

pub struct Router<S> {
    tool_router: ToolRouter<S>,
    resource_router: ResourceRouter<S>,
    // ... and so on.
    service: S,
}

pub fn build_router() {
    Router::new(Calculator::<u64>::default())
        .tool(
            Calculator::<u64>::sum,
            Tool::builder()
                .name("sum")
                .discription("sum two numbers")
                .parameters_type::<SumRequest>()
                .build(),
        )
        .build()
}

and then implement Handler for Router, router will filter the tool related request, can call it by tool handlers in tool router, which is the toolbox now. So we don't need static toolbox anymore, therefore the generic type would not be an issue.

Pros: this would be much easier to implement, and no magic for user.
Cons: user need to write the tool defination and implementation in different place.

4t145 · 2025-05-19T10:04:15Z

And also @yIllusionSky. Maybe you could have a discussion here before writing any code.

jokemanfire · 2025-05-19T10:55:44Z

I think macros should still exist. and I think we can refer to the implementation of routing macros in actix web, which can generate adapted code without using any static variables.

4t145 · 2025-05-19T11:36:05Z

We should not use static router in a lib anyway.

4t145 · 2025-05-20T08:46:21Z

@jokemanfire
https://github.com/modelcontextprotocol/rust-sdk/pull/201/files#diff-ce4632077fe5bb8c0c5c6acf054464f194cd1fcffe38b875d67d047971f19d58

here's draft, a preview look like below:

    let router = Router::new(test_handler)
        // build with function
        .with_tool(
            TestHandler::sync_method
                .name("sync_method")
                .description("a sync method tool")
                .parameters::<Request>(),
        )
        // build with closure
        .with_tool(
            (|Parameters(Sum { a, b }): Parameters<Sum>| (a + b).to_string())
                .name("add")
                .parameters::<Sum>(),
        )
        // build with generator function, which can be generated by macros
        .with_tool(attr_generator_fn)
        .with_tools(tool_router);

And so does the promt

jokemanfire · 2025-05-21T01:17:27Z

Actually, from my perspective as a developer who calls RUST-SDK, I prefer to use macros, so I don't really want to use this method until I confirm that macros are completely unsuitable.

4t145 · 2025-05-21T02:51:25Z

@jokemanfire

We can have both.

There is a confliction by using current handler, if you generate the tool attibute outside the impl trait block, then you have to inject it into impl trait block. Probably by using serveral generated function to take over the request call. We can move the logic of ToolRouter::call and ToolRouter::list_all to a generated function, and then call it in impl trait block.

#[rmcp(tool)]
impl Handler<T> {
    #[rmcp(tool)]
    pub fn tool_a() {}

    ...
    fn __generated_tool_call(...) {}
}

#[rmcp(use_tool)]
impl<T> ServerHandler for Handler<T> {

}

But for users, they won't know what happend, and may meet some problems when they are trying to use those macro separately。

mod feature_a {
    #[rmcp(tool)]
    impl Handler<T> {
        #[rmcp(tool)]
        pub fn tool_a() {}

        // !!! cannot get tools from other impl block
        fn __generated_tool_call(...) {}
    }
}

mod feature_b {
    #[rmcp(tool)]
    impl Handler<T> {
        #[rmcp(tool)]
        pub fn tool_b() {}

        // !!! This can cause a compile error
        fn __generated_tool_call(...) {}
    }
}

// and this looks like a magic. 
#[rmcp(use_tool)]
impl<T> ServerHandler for Handler<T> {

}

And you mentioned actix-web sytle macro, I do discussed it at here in the draft, (perhaps you missed it). Even in actix-web, you have to register it into router manually as well as what was mentioned in the draft.

jokemanfire · 2025-05-21T03:11:52Z

// and this looks like a magic.
#[rmcp(use_tool)]
impl ServerHandler for Handler {

}

This is exactly what I want to consider now. I want to make this part more developer friendly, and indeed, I need to take over the macro again here.

We can have both.

I agree with this.

Sorry, there have been too many things lately and I haven't read the draft yet. I just looked at this issue.Maybe I'll have time to watch today.

jokemanfire · 2025-05-21T12:45:29Z

actully , we can create a ToolRegistry.Registering traits at runtime avoids generic issues this should user to do.
To sum up, two macros are provided: one is the tool macro, and the other is the tool_function macro. The tool macro only provides the registration function declaration, and registers the entire tool structure entity. The tool_function registers the entire function entity, and then the call time is actually when the function calls register.

the daemon like this

This is a registry

pub struct ToolRegistry {
    tools: HashMap<String, Arc<dyn Any + Send + Sync>>,
    functions: HashMap<String, ToolFunction>,
}

tool will gen a registry func,

#[tool]
impl Counter {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            counter: Arc::new(Mutex::new(0)),
        }
    }

    #[tool_function(description = "Decrement the counter by 1")]
    async fn decrement(&self) -> Result<CallToolResult, McpError> {
        let mut counter = self.counter.lock().await;
        *counter -= 1;
        Ok(CallToolResult::success(vec![Content::text(
            counter.to_string(),
        )]))
    }

    #[tool_function(description = "Get the current counter value")]
    async fn get_value(&self) -> Result<CallToolResult, McpError> {
        let counter = self.counter.lock().await;
        Ok(CallToolResult::success(vec![Content::text(
            counter.to_string(),
        )]))
    }
}

In main func ,user should init a registry

#[tokio::main]
async fn main() {
    // Create tool registry
    let mut registry = ToolRegistry::new();
    
    // Register Counter tool , it gen by tool macro
    Counter::register(&mut registry);

    // Call tool function
    match registry.call_function("Counter.get_value", vec![]).await {
        Ok(result) => {
            if let Content::Text(value) = &result.contents[0] {
                println!("Counter value: {}", value);
            }
        }
        Err(err) => {
            println!("Error: {:?}", err);
        }
    }
    
    // Call decrement function
    match registry.call_function("Counter.decrement", vec![]).await {
        Ok(result) => {
            if let Content::Text(value) = &result.contents[0] {
                println!("Counter after decrement: {}", value);
            }
        }
        Err(err) => {
            println!("Error: {:?}", err);
        }
    }
    
    println!("Tool execution completed!");
}

the marco impl like this

/// Tool attribute macro, used to mark struct implementations
#[proc_macro_attribute]
pub fn tool(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemImpl);
    let ty = &input.self_ty;
    
    // Find all methods using #[tool_function]
    let mut register_calls = Vec::new();
    
    for item in &input.items {
        if let syn::ImplItem::Fn(method) = item {
            let is_tool_function = method.attrs.iter().any(|attr| {
                attr.path().is_ident("tool_function")
            });
            
            if is_tool_function {
                let method_name = &method.sig.ident;
                let register_name = format!("register_{}", method_name);
                let register_ident = Ident::new(&register_name, Span::call_site());
                
                register_calls.push(quote! {
                    // Get the short form of the struct name
                    let struct_name = type_name.split("::").last().unwrap_or(type_name);
                    // Register tool function using "struct_name.method_name" format
                    let name = format!("{}.{}", struct_name, stringify!(#method_name));
                    registry.register_function(&name, tool_ref.#register_ident());
                });
            }
        }
    }
    
    let expanded = quote! {
        #input

        impl Tool for #ty {
            fn register(registry: &mut ToolRegistry) {
                let tool = Arc::new(Self::new());
                let type_name = std::any::type_name::<Self>();
                let tool_ref = &tool; // Create reference to avoid moving
                
                // Get the short form of the struct name
                let struct_name = type_name.split("::").last().unwrap_or(type_name);
                
                // Register all tool functions first
                #(#register_calls)*
                
                // Finally register the tool itself
                registry.register_tool(struct_name, tool);
            }
        }
    };

    TokenStream::from(expanded)
}

/// Tool function attribute macro, used to mark methods in tool implementations
#[proc_macro_attribute]
pub fn tool_function(args: TokenStream, input: TokenStream) -> TokenStream {
    let attr_args = parse_macro_input!(args as AttributeArgs);
    let input_fn = parse_macro_input!(input as ItemFn);
    
    let fn_name = &input_fn.sig.ident;
    let fn_args = &input_fn.sig.inputs;
    let fn_body = &input_fn.block;
    let fn_async = input_fn.sig.asyncness.is_some();
    
    // Create new identifier for registration function
    let register_fn_name_str = format!("register_{}", fn_name);
    let register_fn_name = Ident::new(&register_fn_name_str, Span::call_site());

    // Extract description parameter
    let mut description = String::new();
    for meta in attr_args.args.iter() {
        if let Meta::NameValue(name_value) = meta {
            if name_value.path.is_ident("description") {
                if let Expr::Lit(ExprLit { lit: Lit::Str(lit_str), .. }) = &name_value.value {
                    description = lit_str.value();
                }
            }
        }
    }

    let expanded = if fn_async {
        quote! {
            #[allow(dead_code)]
            #input_fn

            fn #register_fn_name(&self) -> ToolFunction {
                ToolFunction {
                    name: stringify!(#fn_name).to_string(),
                    description: #description.to_string(),
                    handler: Box::new(|args, tool_arc| {
                        Box::pin(async move {
                            let tool = tool_arc.downcast_ref::<Self>().unwrap();
                            tool.#fn_name().await
                        })
                    }),
                }
            }
        }
    } else {
        quote! {
            #[allow(dead_code)]
            #input_fn

            fn #register_fn_name(&self) -> ToolFunction {
                ToolFunction {
                    name: stringify!(#fn_name).to_string(),
                    description: #description.to_string(),
                    handler: Box::new(|args, tool_arc| {
                        Box::pin(async move {
                            let tool = tool_arc.downcast_ref::<Self>().unwrap();
                            tool.#fn_name()
                        })
                    }),
                }
            }
        }
    };

    TokenStream::from(expanded)
}

jokemanfire · 2025-05-21T12:49:27Z

tool_test.zip
The daemon code here.

4t145 · 2025-05-21T16:45:35Z

I don't understant why to use a totally dynamic version instead of there's already a typed one, where's the advantage comparing this tool router?

And how do you pass non-json-value argument like tool call context into it?

And what does it do with generic problem, the generic issue is introduced by static registry. You can't declare a static variable with generic type parameter. And it won't be a problem if we don't use static registry.

Have you checked the draft I posted yesterday? The tool router there is your tool registry, they have the same function but with type.

jokemanfire · 2025-05-22T01:05:20Z

I don't understant why to use a totally dynamic version instead of there's already a typed one, where's the advantage comparing this tool router?

I think using a fully dynamic implementation can reduce a lot of code and make implementation and maintenance simpler.Actually, it's just a complete transformation of the toolbox from static to dynamic.

And how do you pass non-json-value argument like tool call context into it?

Now I have only provided the general framework of the design, without considering the details. But for external access to MCP, the parameters passed in should be able to be serialized by JSON, so it should be acceptable for only JSON serialization parameters to be supported internally

And what does it do with generic problem, the generic issue is introduced by static registry. You can't declare a static variable with generic type parameter. And it won't be a problem if we don't use static registry.

And this should be considered a common design, delaying the registration of the toolbox after macro expansion.So generic expansion shouldn't be a problem.

Have you checked the draft I posted yesterday? The tool router there is your tool registry, they have the same function but with type.

I have reviewed the draft and in fact, it did not make significant changes to the macros. Personally, I do not like the Axum style approach, as it makes me feel like I am developing a web backend（Sorry, forgive my personal opinion ）. Therefore, I am more inclined to use the function dynamic registration form. (I referred to some tool routing in the draft for tool registration)

4t145 · 2025-05-22T02:37:43Z

Now I have only provided the general framework of the design, without considering the details. But for external access to MCP, the parameters passed in should be able to be serialized by JSON, so it should be acceptable for only JSON serialization parameters to be supported internally

Currently, we already support pass things like progress token, cancellation token and http header into tool call (#61, #153).

And this should be considered a common design, delaying the registration of the toolbox after macro expansion.So generic expansion shouldn't be a problem.

我的意思是，是否动态类型和泛型参数是没有关系的。之前不能用泛型参数是因为有一个静态仓库，只要不用静态仓库就没有泛型问题。所以我不是想问，你做了什么，而是想说，有什么关系。这里应该是What does it to do with，我少打了一个to，可能翻译没有传达到位。

I have reviewed the draft and in fact, it did not make significant changes to the macros. Personally, I do not like the Axum style approach, as it makes me feel like I am developing a web backend（Sorry, forgive my personal opinion ）. Therefore, I am more inclined to use the function dynamic registration form. (I referred to some tool routing in the draft for tool registration)

You can build upon it in different level. It only looks like axum when you use Router, while macros still could be built upon ToolRouter.

it makes me feel like I am developing a web backend

They are all share the request-handler model, so it's not suprised if they looks similar.

jokemanfire · 2025-05-22T03:20:26Z

Currently, we already support pass things like progress token, cancellation token and http header into tool call

这里我没理解到哈，从我这边看的话这几个功能都是动态的并非静态注册进去，那么从架构上来看应该不存在冲突啊？

jokemanfire · 2025-05-22T03:21:49Z

所以现在的问题应该在
是否能使用全动态框架是吧？

jokemanfire · 2025-05-22T03:32:08Z

从整个设计上我更倾向于注册方式实际上抽象出来整个宏实现就是工具的注册，并且拆分tool宏，现在的tool宏需要做的事情太多，我认为从代码设计角度看他的实现过于臃肿，且职责过泛，使用动态注册方式可以拆分为两种宏
1.使用tool宏生成注册函数，并注册结构体
2.tool_function宏注册函数

最后在实现时调用注册函数以注册结构体和函数
好处宏职责更加清晰，调用者代码也更加清晰，不需要过于复杂的静态宏实现

调用者更好理解代码逻辑：

    // Create tool registry
    let mut registry = ToolRegistry::new();
    
    // Register Counter tool , it gen by tool macro
    Counter::register(&mut registry);

    // Call tool function
    registry.call_function("Counter.get_value", vec![])

4t145 · 2025-05-22T03:42:53Z

从整个设计上我更倾向于注册方式实际上抽象出来整个宏实现就是工具的注册

你说的注册register和路由的添加add本质上没有任何区别啊？路由不就是handler的仓库吗？所以这里不存在倾向的说法吧，这两个其实就是换了个名字而已

所以现在的问题应该在是否能使用全动态框架是吧？

这里我没理解到哈，从我这边看的话这几个功能都是动态的并非静态注册进去，那么从架构上来看应该不存在冲突啊？

现在是已经有了一套Extractor + Handler的系统了，这一套设计模式广泛应用于rust的各种框架中，比如web框架基本都采用这一套，比如游戏引擎bevy的ECS，等等。我不知道你是否了解这一套模式，我觉得这一套可以在强类型的情况下实现你动态的功能。

我是想说，要做宏可以在保留类型的ToolRouter基础上继续做宏。可以从不同级别渐进式的满足需求：可以是生成一个函数返回ToolRouteItem（actix-web like），也可以全包裹的，在impl块生成两个处理list和call的方法，在impl trait块调用这两个方法 (原来的方法)

registry.call_function("Counter.get_value", vec![])

对于调用者而言，是传入函数指针更清晰还是传入函数名字符串更清晰？如果我是调用者我肯定希望传入函数指针。

我认为从代码设计角度看他的实现过于臃肿，且职责过泛

的确原来的写的代码结构不是很好，但是tool臃肿是必然的，在这一个语义上承担了太多功能。所以可以把它的语义分割一下，成为#[rmcp(tool)], #[rmcp(generate(tool_call, tool_list))], #[rmcp(use(tool_call, tool_list))]这三个，对于调用者来说他们也会很清楚他们在做什么

4t145 · 2025-05-22T03:55:11Z

我觉得这种纯动态的方式，应该是在一些需要隔离Caller和Callee的情况下，也就是说，依赖注入，听起来有点java。我原来也用过这种结构，是在处理网关插件的场景下，这种可能是用户自己编写的插件，甚至是动态链接进来的插件 https://github.com/ideal-world/spacegate/blob/master/crates/plugin/src/lib.rs

jokemanfire · 2025-05-22T12:52:24Z

现在是已经有了一套Extractor + Handler的系统了，这一套设计模式广泛应用于rust的各种框架中

我其实认为这个需求从我这里来看并不需要过分复杂的实现，我想把这块的代码越简化越好。完全使用动态注册在我看来，实现简单代码量少，当然使用静态注册的方式也可以的。

registry.call_function("Counter.get_value", vec![])

这种更像是python的eval函数

但是可能还有个问题我需要指出
现在的方法不能多次对结构体实现函数，你应该是清楚地例如
impl test{
fn test1(){}
}

impl test{
fn test2(){}
}

这样是不行的，所以我同样想解决这个问题，特别是我在写MCP上层应用时，更喜欢拆成多个impl实现块。这样让我的MCP服务更加清晰。

4t145 · 2025-05-22T19:49:04Z

对于用户来说，从强类型简化到动态很容易，从动态fallback到强类型却很难。

对于多impl块，可以考虑为每个块生成一个函数，然后最后调用时再把这些路由合并起来，然后在合并后路由上调用。

#[rmcp(tool)], #[rmcp(tool_router(mod_1))], #[rmcp(use_tool_router(mod_1, mod_2))]

但是我觉得，对于复杂的项目，我不会在意是否多写几行代码来手写路由，如果是我我就不会使用宏了

Pluto-Y · 2025-05-28T03:13:44Z

这里重构版本是否有考虑过在#[tool]定义的函数内拿到peer和context?

    #[tool(description = "Demonstrates a long running operation with progress updates")]
    pub async fn long_running_operation(
        &self,
        #[tool(param)]
        #[schemars(description = "Duration of the operation in seconds")]
        duration: u64,
        #[tool(param)]
        #[schemars(description = "Number of steps in the operation")]
        steps: u64,
        // _context: RequestContext<RoleServer>,
    ) -> String {
        let total_steps = steps;
        for i in 0..total_steps {
            tokio::time::sleep(tokio::time::Duration::from_secs(duration / total_steps)).await;
            let progress = (i + 1) * 10;
            // TODO: support progress notification
            // self.peer
            //     .as_ref()
            //     .unwrap()
            //     .notify_progress(ProgressNotificationParam { progress })
            //     .await
            //     .unwrap();
            println!("Progress: {}%", progress);
        }
        format!(
            "Long running operation completed. Duration: {} seconds, Steps: {}.",
            duration, total_steps
        )
    }

比如上面的一个长耗时行为通过#tool声明的宏内就无法拿到peer和progressToken。老的ServerHandler内是可以通过get_peer获得，新的接口只有在context内拿到peer信息

4t145 · 2025-05-28T03:29:37Z

@Pluto-Y 是有考虑的，其实之前的handler设想中推荐的做法是通过extractor来获取，也就是直接把peer: Peer放在参数中，提取器会自动从context中提取出peer并传入。

Pluto-Y · 2025-05-28T03:32:10Z

有考虑就好，因为我们现在再做对应功能开发的时候没办法通过#[tool]宏拿到上下文，只能那个退回到 ServerHandler 里面处理，写起来有点别扭。

期待新版本

Pluto-Y · 2025-05-28T03:35:22Z

另外问个问题，这个重构会在0.2.0的里程碑里吗？

4t145 · 2025-05-28T03:47:05Z

@Pluto-Y

因为我们现在再做对应功能开发的时候没办法通过#[tool]宏拿到上下文

tool宏可能会有较大的改动。然后你拿不到上下文是我们这里有遗漏，没有给Peer实现FromToolCallContextPart这个trait，其实小修补一下就可以从参数中拿到了。

另外问个问题，这个重构会在0.2.0的里程碑里吗？

目前是放在里面的，但是这个设计的方案我们还在犹豫。

jokemanfire added this to the Pre release milestones milestone May 19, 2025

This was referenced May 20, 2025

Types with generics do not work with tool macro #91

Open

draft: handler router, tool router, with axum style. #201

Draft

jokemanfire added the enhancement New feature or request label May 21, 2025

4t145 mentioned this issue May 25, 2025

Question: Dynamic tool generation #221

Open

Enhancement: Refactoring tool macro #195

Enhancement: Refactoring tool macro #195

Comments

jokemanfire commented May 19, 2025

4t145 commented May 19, 2025

Uh oh!

jokemanfire commented May 19, 2025

Uh oh!

jokemanfire commented May 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

4t145 commented May 19, 2025

Uh oh!

4t145 commented May 19, 2025

Uh oh!

4t145 commented May 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

4t145 commented May 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jokemanfire commented May 19, 2025

Uh oh!

4t145 commented May 19, 2025

Uh oh!

4t145 commented May 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jokemanfire commented May 21, 2025

Uh oh!

4t145 commented May 21, 2025

Uh oh!

jokemanfire commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jokemanfire commented May 21, 2025 • edited by 4t145 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jokemanfire commented May 21, 2025

Uh oh!

4t145 commented May 21, 2025

Uh oh!

jokemanfire commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

4t145 commented May 22, 2025

Uh oh!

jokemanfire commented May 22, 2025

Uh oh!

jokemanfire commented May 22, 2025

Uh oh!

jokemanfire commented May 22, 2025

Uh oh!

4t145 commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

4t145 commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jokemanfire commented May 22, 2025

Uh oh!

4t145 commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Pluto-Y commented May 28, 2025

Uh oh!

4t145 commented May 28, 2025

Uh oh!

Pluto-Y commented May 28, 2025

Uh oh!

Pluto-Y commented May 28, 2025

Uh oh!

4t145 commented May 28, 2025

Uh oh!

jokemanfire commented May 19, 2025 •

edited

Loading

4t145 commented May 19, 2025 •

edited

Loading

4t145 commented May 19, 2025 •

edited

Loading

4t145 commented May 20, 2025 •

edited

Loading

jokemanfire commented May 21, 2025 •

edited

Loading

jokemanfire commented May 21, 2025 •

edited by 4t145

Loading

jokemanfire commented May 22, 2025 •

edited

Loading

4t145 commented May 22, 2025 •

edited

Loading

4t145 commented May 22, 2025 •

edited

Loading

4t145 commented May 22, 2025 •

edited

Loading