refactor: refactor tool macros and router implementation (#261)

* refactor: refactor tool macros and router implementation

- Updated the `#[tool(tool_box)]` macro to `#[tool_router]` across various modules for consistency.
- Enhanced the `Calculator`, `Counter`, and `GenericService` structs to utilize `ToolRouter` for handling tool calls.
- Introduced `Parameters` struct for better parameter handling in tool functions.
- Added new methods for listing tools and calling tools in server handlers.
- Improved test cases to reflect changes in tool routing and parameter handling.
- Updated documentation and examples to align with the new router structure.

* fix: fix fmt and build error

* fix: fix test failure

* docs: documents for macros, fix ci

* fix: fix ci

* fix: fix wrongly replaced documents

* fix: remove useless file

* fix: change the parameter format for tool_router

* fix: update extract_doc_line to handle existing documentation and clean up unused code in server handler

* doc: update document for macro and examples

* doc: update readme and add contribute guide

* fix: fix type

crates/rmcp/src/handler/server/router.rs

@@ -0,0 +1,96 @@
+use std::sync::Arc;
+
+use tool::{IntoToolRoute, ToolRoute};
+
+use super::ServerHandler;
+use crate::{
+    RoleServer, Service,
+    model::{ClientRequest, ListToolsResult, ServerResult},
+    service::NotificationContext,
+};
+
+pub mod tool;
+
+pub struct Router<S> {
+    pub tool_router: tool::ToolRouter<S>,
+    pub service: Arc<S>,
+}
+
+impl<S> Router<S>
+where
+    S: ServerHandler,
+{
+    pub fn new(service: S) -> Self {
+        Self {
+            tool_router: tool::ToolRouter::new(),
+            service: Arc::new(service),
+        }
+    }
+
+    pub fn with_tool<R, A>(mut self, route: R) -> Self
+    where
+        R: IntoToolRoute<S, A>,
+    {
+        self.tool_router.add_route(route.into_tool_route());
+        self
+    }
+
+    pub fn with_tools(mut self, routes: impl IntoIterator<Item = ToolRoute<S>>) -> Self {
+        for route in routes {
+            self.tool_router.add_route(route);
+        }
+        self
+    }
+}
+
+impl<S> Service<RoleServer> for Router<S>
+where
+    S: ServerHandler,
+{
+    async fn handle_notification(
+        &self,
+        notification: <RoleServer as crate::service::ServiceRole>::PeerNot,
+        context: NotificationContext<RoleServer>,
+    ) -> Result<(), crate::Error> {
+        self.service
+            .handle_notification(notification, context)
+            .await
+    }
+    async fn handle_request(
+        &self,
+        request: <RoleServer as crate::service::ServiceRole>::PeerReq,
+        context: crate::service::RequestContext<RoleServer>,
+    ) -> Result<<RoleServer as crate::service::ServiceRole>::Resp, crate::Error> {
+        match request {
+            ClientRequest::CallToolRequest(request) => {
+                if self.tool_router.has_route(request.params.name.as_ref())
+                    || !self.tool_router.transparent_when_not_found
+                {
+                    let tool_call_context = crate::handler::server::tool::ToolCallContext::new(
+                        self.service.as_ref(),
+                        request.params,
+                        context,
+                    );
+                    let result = self.tool_router.call(tool_call_context).await?;
+                    Ok(ServerResult::CallToolResult(result))
+                } else {
+                    self.service
+                        .handle_request(ClientRequest::CallToolRequest(request), context)
+                        .await
+                }
+            }
+            ClientRequest::ListToolsRequest(_) => {
+                let tools = self.tool_router.list_all();
+                Ok(ServerResult::ListToolsResult(ListToolsResult {
+                    tools,
+                    next_cursor: None,
+                }))
+            }
+            rest => self.service.handle_request(rest, context).await,
+        }
+    }
+
+    fn get_info(&self) -> <RoleServer as crate::service::ServiceRole>::Info {
+        self.service.get_info()
+    }
+}

crates/rmcp/src/handler/server/router/tool.rs

@@ -0,0 +1,272 @@
+use std::{borrow::Cow, sync::Arc};
+
+use futures::{FutureExt, future::BoxFuture};
+use schemars::JsonSchema;
+
+use crate::{
+    handler::server::tool::{
+        CallToolHandler, DynCallToolHandler, ToolCallContext, schema_for_type,
+    },
+    model::{CallToolResult, Tool, ToolAnnotations},
+};
+
+pub struct ToolRoute<S> {
+    #[allow(clippy::type_complexity)]
+    pub call: Arc<DynCallToolHandler<S>>,
+    pub attr: crate::model::Tool,
+}
+
+impl<S> std::fmt::Debug for ToolRoute<S> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ToolRoute")
+            .field("name", &self.attr.name)
+            .field("description", &self.attr.description)
+            .field("input_schema", &self.attr.input_schema)
+            .finish()
+    }
+}
+
+impl<S> Clone for ToolRoute<S> {
+    fn clone(&self) -> Self {
+        Self {
+            call: self.call.clone(),
+            attr: self.attr.clone(),
+        }
+    }
+}
+
+impl<S: Send + Sync + 'static> ToolRoute<S> {
+    pub fn new<C, A>(attr: impl Into<Tool>, call: C) -> Self
+    where
+        C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+    {
+        Self {
+            call: Arc::new(move |context: ToolCallContext<S>| {
+                let call = call.clone();
+                context.invoke(call).boxed()
+            }),
+            attr: attr.into(),
+        }
+    }
+    pub fn new_dyn<C>(attr: impl Into<Tool>, call: C) -> Self
+    where
+        C: for<'a> Fn(
+                ToolCallContext<'a, S>,
+            ) -> BoxFuture<'a, Result<CallToolResult, crate::Error>>
+            + Send
+            + Sync
+            + 'static,
+    {
+        Self {
+            call: Arc::new(call),
+            attr: attr.into(),
+        }
+    }
+    pub fn name(&self) -> &str {
+        &self.attr.name
+    }
+}
+
+pub trait IntoToolRoute<S, A> {
+    fn into_tool_route(self) -> ToolRoute<S>;
+}
+
+impl<S, C, A, T> IntoToolRoute<S, A> for (T, C)
+where
+    S: Send + Sync + 'static,
+    C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+    T: Into<Tool>,
+{
+    fn into_tool_route(self) -> ToolRoute<S> {
+        ToolRoute::new(self.0.into(), self.1)
+    }
+}
+
+impl<S> IntoToolRoute<S, ()> for ToolRoute<S>
+where
+    S: Send + Sync + 'static,
+{
+    fn into_tool_route(self) -> ToolRoute<S> {
+        self
+    }
+}
+
+pub struct ToolAttrGenerateFunctionAdapter;
+impl<S, F> IntoToolRoute<S, ToolAttrGenerateFunctionAdapter> for F
+where
+    S: Send + Sync + 'static,
+    F: Fn() -> ToolRoute<S>,
+{
+    fn into_tool_route(self) -> ToolRoute<S> {
+        (self)()
+    }
+}
+
+pub trait CallToolHandlerExt<S, A>: Sized
+where
+    Self: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+{
+    fn name(self, name: impl Into<Cow<'static, str>>) -> WithToolAttr<Self, S, A>;
+}
+
+impl<C, S, A> CallToolHandlerExt<S, A> for C
+where
+    C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+{
+    fn name(self, name: impl Into<Cow<'static, str>>) -> WithToolAttr<Self, S, A> {
+        WithToolAttr {
+            attr: Tool::new(
+                name.into(),
+                "",
+                schema_for_type::<crate::model::JsonObject>(),
+            ),
+            call: self,
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+pub struct WithToolAttr<C, S, A>
+where
+    C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+{
+    pub attr: crate::model::Tool,
+    pub call: C,
+    pub _marker: std::marker::PhantomData<fn(S, A)>,
+}
+
+impl<C, S, A> IntoToolRoute<S, A> for WithToolAttr<C, S, A>
+where
+    C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+    S: Send + Sync + 'static,
+{
+    fn into_tool_route(self) -> ToolRoute<S> {
+        ToolRoute::new(self.attr, self.call)
+    }
+}
+
+impl<C, S, A> WithToolAttr<C, S, A>
+where
+    C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
+{
+    pub fn description(mut self, description: impl Into<Cow<'static, str>>) -> Self {
+        self.attr.description = Some(description.into());
+        self
+    }
+    pub fn parameters<T: JsonSchema>(mut self) -> Self {
+        self.attr.input_schema = schema_for_type::<T>().into();
+        self
+    }
+    pub fn parameters_value(mut self, schema: serde_json::Value) -> Self {
+        self.attr.input_schema = crate::model::object(schema).into();
+        self
+    }
+    pub fn annotation(mut self, annotation: impl Into<ToolAnnotations>) -> Self {
+        self.attr.annotations = Some(annotation.into());
+        self
+    }
+}
+#[derive(Debug)]
+pub struct ToolRouter<S> {
+    #[allow(clippy::type_complexity)]
+    pub map: std::collections::HashMap<Cow<'static, str>, ToolRoute<S>>,
+
+    pub transparent_when_not_found: bool,
+}
+
+impl<S> Default for ToolRouter<S> {
+    fn default() -> Self {
+        Self {
+            map: std::collections::HashMap::new(),
+            transparent_when_not_found: false,
+        }
+    }
+}
+impl<S> Clone for ToolRouter<S> {
+    fn clone(&self) -> Self {
+        Self {
+            map: self.map.clone(),
+            transparent_when_not_found: self.transparent_when_not_found,
+        }
+    }
+}
+
+impl<S> IntoIterator for ToolRouter<S> {
+    type Item = ToolRoute<S>;
+    type IntoIter = std::collections::hash_map::IntoValues<Cow<'static, str>, ToolRoute<S>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.map.into_values()
+    }
+}
+
+impl<S> ToolRouter<S>
+where
+    S: Send + Sync + 'static,
+{
+    pub fn new() -> Self {
+        Self {
+            map: std::collections::HashMap::new(),
+            transparent_when_not_found: false,
+        }
+    }
+    pub fn with_route<R, A>(mut self, route: R) -> Self
+    where
+        R: IntoToolRoute<S, A>,
+    {
+        self.add_route(route.into_tool_route());
+        self
+    }
+
+    pub fn add_route(&mut self, item: ToolRoute<S>) {
+        self.map.insert(item.attr.name.clone(), item);
+    }
+
+    pub fn merge(&mut self, other: ToolRouter<S>) {
+        for item in other.map.into_values() {
+            self.add_route(item);
+        }
+    }
+
+    pub fn remove_route<H, A>(&mut self, name: &str) {
+        self.map.remove(name);
+    }
+    pub fn has_route(&self, name: &str) -> bool {
+        self.map.contains_key(name)
+    }
+    pub async fn call(
+        &self,
+        context: ToolCallContext<'_, S>,
+    ) -> Result<CallToolResult, crate::Error> {
+        let item = self
+            .map
+            .get(context.name())
+            .ok_or_else(|| crate::Error::invalid_params("tool not found", None))?;
+        (item.call)(context).await
+    }
+
+    pub fn list_all(&self) -> Vec<crate::model::Tool> {
+        self.map.values().map(|item| item.attr.clone()).collect()
+    }
+}
+
+impl<S> std::ops::Add<ToolRouter<S>> for ToolRouter<S>
+where
+    S: Send + Sync + 'static,
+{
+    type Output = Self;
+
+    fn add(mut self, other: ToolRouter<S>) -> Self::Output {
+        self.merge(other);
+        self
+    }
+}
+
+impl<S> std::ops::AddAssign<ToolRouter<S>> for ToolRouter<S>
+where
+    S: Send + Sync + 'static,
+{
+    fn add_assign(&mut self, other: ToolRouter<S>) {
+        self.merge(other);
+    }
+}

crates/rmcp/src/handler/server/tool.rs

@@ -2,16 +2,18 @@ use std::{
     any::TypeId, borrow::Cow, collections::HashMap, future::Ready, marker::PhantomData, sync::Arc,
 };
 
-use futures::future::BoxFuture;
+use futures::future::{BoxFuture, FutureExt};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use tokio_util::sync::CancellationToken;
 
+pub use super::router::tool::{ToolRoute, ToolRouter};
 use crate::{
     RoleServer,
-    model::{CallToolRequestParam, CallToolResult, ConstString, IntoContents, JsonObject},
+    model::{CallToolRequestParam, CallToolResult, IntoContents, JsonObject},
     service::RequestContext,
 };
+
 /// A shortcut for generating a JSON schema for a type.
 pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
     // explicitly to align json schema version to official specifications.
@@ -62,16 +64,16 @@ pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, cr
         )
     })
 }
-pub struct ToolCallContext<'service, S> {
-    request_context: RequestContext<RoleServer>,
-    service: &'service S,
-    name: Cow<'static, str>,
-    arguments: Option<JsonObject>,
+pub struct ToolCallContext<'s, S> {
+    pub request_context: RequestContext<RoleServer>,
+    pub service: &'s S,
+    pub name: Cow<'static, str>,
+    pub arguments: Option<JsonObject>,
 }
 
-impl<'service, S> ToolCallContext<'service, S> {
+impl<'s, S> ToolCallContext<'s, S> {
     pub fn new(
-        service: &'service S,
+        service: &'s S,
         CallToolRequestParam { name, arguments }: CallToolRequestParam,
         request_context: RequestContext<RoleServer>,
     ) -> Self {
@@ -90,10 +92,8 @@ impl<'service, S> ToolCallContext<'service, S> {
     }
 }
 
-pub trait FromToolCallContextPart<'a, S>: Sized {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error>;
+pub trait FromToolCallContextPart<S>: Sized {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error>;
 }
 
 pub trait IntoCallToolResult {
@@ -161,16 +161,16 @@ impl IntoCallToolResult for Result<CallToolResult, crate::Error> {
     }
 }
 
-pub trait CallToolHandler<'a, S, A> {
-    type Fut: Future<Output = Result<CallToolResult, crate::Error>> + Send + 'a;
-    fn call(self, context: ToolCallContext<'a, S>) -> Self::Fut;
+pub trait CallToolHandler<S, A> {
+    fn call(
+        self,
+        context: ToolCallContext<'_, S>,
+    ) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>;
 }
 
-pub type DynCallToolHandler<S> = dyn Fn(ToolCallContext<'_, S>) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>
+pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result<CallToolResult, crate::Error>>
     + Send
     + Sync;
-/// Parameter Extractor
-pub struct Parameter<K: ConstString, V>(pub K, pub V);
 
 /// Parameter Extractor
 ///
@@ -188,83 +188,25 @@ impl<P: JsonSchema> JsonSchema for Parameters<P> {
     }
 }
 
-/// Callee Extractor
-pub struct Callee<'a, S>(pub &'a S);
-
-impl<'a, S> FromToolCallContextPart<'a, S> for CancellationToken {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((context.request_context.ct.clone(), context))
-    }
-}
-
-impl<'a, S> FromToolCallContextPart<'a, S> for Callee<'a, S> {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((Callee(context.service), context))
+impl<S> FromToolCallContextPart<S> for CancellationToken {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
+        Ok(context.request_context.ct.clone())
     }
 }
 
 pub struct ToolName(pub Cow<'static, str>);
 
-impl<'a, S> FromToolCallContextPart<'a, S> for ToolName {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((Self(context.name.clone()), context))
+impl<S> FromToolCallContextPart<S> for ToolName {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
+        Ok(Self(context.name.clone()))
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for &'a S {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((context.service, context))
-    }
-}
-
-impl<'a, S, K, V> FromToolCallContextPart<'a, S> for Parameter<K, V>
-where
-    K: ConstString,
-    V: DeserializeOwned,
-{
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        let arguments = context
-            .arguments
-            .as_ref()
-            .ok_or(crate::Error::invalid_params(
-                format!("missing parameter {field}", field = K::VALUE),
-                None,
-            ))?;
-        let value = arguments.get(K::VALUE).ok_or(crate::Error::invalid_params(
-            format!("missing parameter {field}", field = K::VALUE),
-            None,
-        ))?;
-        let value: V = serde_json::from_value(value.clone()).map_err(|e| {
-            crate::Error::invalid_params(
-                format!(
-                    "failed to deserialize parameter {field}: {error}",
-                    field = K::VALUE,
-                    error = e
-                ),
-                None,
-            )
-        })?;
-        Ok((Parameter(K::default(), value), context))
-    }
-}
-
-impl<'a, S, P> FromToolCallContextPart<'a, S> for Parameters<P>
+impl<S, P> FromToolCallContextPart<S> for Parameters<P>
 where
     P: DeserializeOwned,
 {
-    fn from_tool_call_context_part(
-        mut context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let arguments = context.arguments.take().unwrap_or_default();
         let value: P =
             serde_json::from_value(serde_json::Value::Object(arguments)).map_err(|e| {
@@ -273,37 +215,31 @@ where
                     None,
                 )
             })?;
-        Ok((Parameters(value), context))
+        Ok(Parameters(value))
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for JsonObject {
-    fn from_tool_call_context_part(
-        mut context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+impl<S> FromToolCallContextPart<S> for JsonObject {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let object = context.arguments.take().unwrap_or_default();
-        Ok((object, context))
+        Ok(object)
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for crate::model::Extensions {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+impl<S> FromToolCallContextPart<S> for crate::model::Extensions {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let extensions = context.request_context.extensions.clone();
-        Ok((extensions, context))
+        Ok(extensions)
     }
 }
 
 pub struct Extension<T>(pub T);
 
-impl<'a, S, T> FromToolCallContextPart<'a, S> for Extension<T>
+impl<S, T> FromToolCallContextPart<S> for Extension<T>
 where
     T: Send + Sync + 'static + Clone,
 {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let extension = context
             .request_context
             .extensions
@@ -315,58 +251,52 @@ where
                     None,
                 )
             })?;
-        Ok((Extension(extension), context))
+        Ok(Extension(extension))
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for crate::Peer<RoleServer> {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+impl<S> FromToolCallContextPart<S> for crate::Peer<RoleServer> {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let peer = context.request_context.peer.clone();
-        Ok((peer, context))
+        Ok(peer)
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for crate::model::Meta {
-    fn from_tool_call_context_part(
-        mut context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
+impl<S> FromToolCallContextPart<S> for crate::model::Meta {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
         let mut meta = crate::model::Meta::default();
         std::mem::swap(&mut meta, &mut context.request_context.meta);
-        Ok((meta, context))
+        Ok(meta)
     }
 }
 
 pub struct RequestId(pub crate::model::RequestId);
-impl<'a, S> FromToolCallContextPart<'a, S> for RequestId {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((RequestId(context.request_context.id.clone()), context))
+impl<S> FromToolCallContextPart<S> for RequestId {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
+        Ok(RequestId(context.request_context.id.clone()))
     }
 }
 
-impl<'a, S> FromToolCallContextPart<'a, S> for RequestContext<RoleServer> {
-    fn from_tool_call_context_part(
-        context: ToolCallContext<'a, S>,
-    ) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
-        Ok((context.request_context.clone(), context))
+impl<S> FromToolCallContextPart<S> for RequestContext<RoleServer> {
+    fn from_tool_call_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::Error> {
+        Ok(context.request_context.clone())
     }
 }
 
 impl<'s, S> ToolCallContext<'s, S> {
-    pub fn invoke<H, A>(self, h: H) -> H::Fut
+    pub fn invoke<H, A>(self, h: H) -> BoxFuture<'s, Result<CallToolResult, crate::Error>>
     where
-        H: CallToolHandler<'s, S, A>,
+        H: CallToolHandler<S, A>,
     {
         h.call(self)
     }
 }
-
 #[allow(clippy::type_complexity)]
-pub struct AsyncAdapter<P, Fut, R>(PhantomData<(fn(P) -> Fut, fn(Fut) -> R)>);
+pub struct AsyncAdapter<P, Fut, R>(PhantomData<fn(P) -> fn(Fut) -> R>);
 pub struct SyncAdapter<P, R>(PhantomData<fn(P) -> R>);
+// #[allow(clippy::type_complexity)]
+pub struct AsyncMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
+pub struct SyncMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
 
 macro_rules! impl_for {
     ($($T: ident)*) => {
@@ -381,174 +311,118 @@ macro_rules! impl_for {
         impl_for!([$($Tn)* $Tn_1] [$($Rest)*]);
     };
     (@impl $($Tn: ident)*) => {
-        impl<'s, $($Tn,)* S, F, Fut, R> CallToolHandler<'s, S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
+        impl<$($Tn,)* S, F,  R> CallToolHandler<S, AsyncMethodAdapter<($($Tn,)*), R>> for F
         where
             $(
-                $Tn: FromToolCallContextPart<'s, S> + 's,
+                $Tn: FromToolCallContextPart<S> ,
             )*
-            F: FnOnce($($Tn,)*) -> Fut + Send + 's,
-            Fut: Future<Output = R> + Send + 's,
-            R: IntoCallToolResult + Send + 's,
-            S: Send + Sync,
+            F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R>,
+
+            // Need RTN support here(I guess), https://github.com/rust-lang/rust/pull/138424
+            // Fut: Future<Output = R> + Send + 'a,
+            R: IntoCallToolResult + Send + 'static,
+            S: Send + Sync + 'static,
         {
-            type Fut = IntoCallToolResultFut<Fut, R>;
-            #[allow(unused_variables, non_snake_case)]
+            #[allow(unused_variables, non_snake_case, unused_mut)]
             fn call(
                 self,
-                context: ToolCallContext<'s, S>,
-            ) -> Self::Fut {
+                mut context: ToolCallContext<'_, S>,
+            ) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>{
                 $(
-                    let result = $Tn::from_tool_call_context_part(context);
-                    let ($Tn, context) = match result {
-                        Ok((value, context)) => (value, context),
-                        Err(e) => return IntoCallToolResultFut::Ready {
-                            result: std::future::ready(Err(e)),
-                        },
+                    let result = $Tn::from_tool_call_context_part(&mut context);
+                    let $Tn = match result {
+                        Ok(value) => value,
+                        Err(e) => return std::future::ready(Err(e)).boxed(),
                     };
                 )*
-                IntoCallToolResultFut::Pending {
-                    fut: self($($Tn,)*),
-                    _marker: PhantomData
-                }
+                let service = context.service;
+                let fut = self(service, $($Tn,)*);
+                async move {
+                    let result = fut.await;
+                    result.into_call_tool_result()
+                }.boxed()
             }
         }
 
-        impl<'s, $($Tn,)* S, F, R> CallToolHandler<'s, S, SyncAdapter<($($Tn,)*), R>> for F
+        impl<$($Tn,)* S, F, Fut, R> CallToolHandler<S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
         where
             $(
-                $Tn: FromToolCallContextPart<'s, S> + 's,
+                $Tn: FromToolCallContextPart<S> ,
             )*
-            F: FnOnce($($Tn,)*) -> R + Send + 's,
-            R: IntoCallToolResult + Send + 's,
+            F: FnOnce($($Tn,)*) -> Fut + Send + ,
+            Fut: Future<Output = R> + Send + 'static,
+            R: IntoCallToolResult + Send + 'static,
             S: Send + Sync,
         {
-            type Fut = Ready<Result<CallToolResult, crate::Error>>;
-            #[allow(unused_variables, non_snake_case)]
+            #[allow(unused_variables, non_snake_case, unused_mut)]
             fn call(
                 self,
-                context: ToolCallContext<'s, S>,
-            ) -> Self::Fut {
+                mut context: ToolCallContext<S>,
+            ) -> BoxFuture<'static, Result<CallToolResult, crate::Error>>{
                 $(
-                    let result = $Tn::from_tool_call_context_part(context);
-                    let ($Tn, context) = match result {
-                        Ok((value, context)) => (value, context),
-                        Err(e) => return std::future::ready(Err(e)),
+                    let result = $Tn::from_tool_call_context_part(&mut context);
+                    let $Tn = match result {
+                        Ok(value) => value,
+                        Err(e) => return std::future::ready(Err(e)).boxed(),
                     };
                 )*
-                std::future::ready(self($($Tn,)*).into_call_tool_result())
+                let fut = self($($Tn,)*);
+                async move {
+                    let result = fut.await;
+                    result.into_call_tool_result()
+                }.boxed()
+            }
+        }
+
+        impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncMethodAdapter<($($Tn,)*), R>> for F
+        where
+            $(
+                $Tn: FromToolCallContextPart<S> + ,
+            )*
+            F: FnOnce(&S, $($Tn,)*) -> R + Send + ,
+            R: IntoCallToolResult + Send + ,
+            S: Send + Sync,
+        {
+            #[allow(unused_variables, non_snake_case, unused_mut)]
+            fn call(
+                self,
+                mut context: ToolCallContext<S>,
+            ) -> BoxFuture<'static, Result<CallToolResult, crate::Error>> {
+                $(
+                    let result = $Tn::from_tool_call_context_part(&mut context);
+                    let $Tn = match result {
+                        Ok(value) => value,
+                        Err(e) => return std::future::ready(Err(e)).boxed(),
+                    };
+                )*
+                std::future::ready(self(context.service, $($Tn,)*).into_call_tool_result()).boxed()
+            }
+        }
+
+        impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncAdapter<($($Tn,)*), R>> for F
+        where
+            $(
+                $Tn: FromToolCallContextPart<S> + ,
+            )*
+            F: FnOnce($($Tn,)*) -> R + Send + ,
+            R: IntoCallToolResult + Send + ,
+            S: Send + Sync,
+        {
+            #[allow(unused_variables, non_snake_case, unused_mut)]
+            fn call(
+                self,
+                mut context: ToolCallContext<S>,
+            ) -> BoxFuture<'static, Result<CallToolResult, crate::Error>>  {
+                $(
+                    let result = $Tn::from_tool_call_context_part(&mut context);
+                    let $Tn = match result {
+                        Ok(value) => value,
+                        Err(e) => return std::future::ready(Err(e)).boxed(),
+                    };
+                )*
+                std::future::ready(self($($Tn,)*).into_call_tool_result()).boxed()
             }
         }
     };
 }
 impl_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
-pub struct ToolBoxItem<S> {
-    #[allow(clippy::type_complexity)]
-    pub call: Box<DynCallToolHandler<S>>,
-    pub attr: crate::model::Tool,
-}
-
-impl<S: Send + Sync + 'static + Clone> ToolBoxItem<S> {
-    pub fn new<C>(attr: crate::model::Tool, call: C) -> Self
-    where
-        C: Fn(ToolCallContext<'_, S>) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>
-            + Send
-            + Sync
-            + 'static,
-    {
-        Self {
-            call: Box::new(call),
-            attr,
-        }
-    }
-    pub fn name(&self) -> &str {
-        &self.attr.name
-    }
-}
-
-#[derive(Default)]
-pub struct ToolBox<S> {
-    #[allow(clippy::type_complexity)]
-    pub map: std::collections::HashMap<Cow<'static, str>, ToolBoxItem<S>>,
-}
-
-impl<S> ToolBox<S> {
-    pub fn new() -> Self {
-        Self {
-            map: std::collections::HashMap::new(),
-        }
-    }
-    pub fn add(&mut self, item: ToolBoxItem<S>) {
-        self.map.insert(item.attr.name.clone(), item);
-    }
-
-    pub fn remove<H, A>(&mut self, name: &str) {
-        self.map.remove(name);
-    }
-
-    pub async fn call(
-        &self,
-        context: ToolCallContext<'_, S>,
-    ) -> Result<CallToolResult, crate::Error> {
-        let item = self
-            .map
-            .get(context.name())
-            .ok_or_else(|| crate::Error::invalid_params("tool not found", None))?;
-        (item.call)(context).await
-    }
-
-    pub fn list(&self) -> Vec<crate::model::Tool> {
-        self.map.values().map(|item| item.attr.clone()).collect()
-    }
-}
-
-#[cfg(feature = "macros")]
-#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
-#[macro_export]
-macro_rules! tool_box {
-    (@pin_add $callee: ident, $attr: expr, $f: expr) => {
-        $callee.add(ToolBoxItem::new($attr, |context| Box::pin($f(context))));
-    };
-    ($server: ident { $($tool: ident),* $(,)?} ) => {
-        $crate::tool_box!($server { $($tool),* }  tool_box);
-    };
-    ($server: ident { $($tool: ident),* $(,)?} $tool_box: ident) => {
-        fn $tool_box() -> &'static $crate::handler::server::tool::ToolBox<$server> {
-            use $crate::handler::server::tool::{ToolBox, ToolBoxItem};
-            static TOOL_BOX: std::sync::OnceLock<ToolBox<$server>> = std::sync::OnceLock::new();
-            TOOL_BOX.get_or_init(|| {
-                let mut tool_box = ToolBox::new();
-                $crate::paste!{
-                    $(
-                        $crate::tool_box!(@pin_add tool_box, $server::[< $tool _tool_attr>](), $server::[<$tool _tool_call>]);
-                    )*
-                }
-                tool_box
-            })
-        }
-    };
-    (@derive) => {
-        $crate::tool_box!(@derive tool_box);
-    };
-
-    (@derive $tool_box:ident) => {
-        async fn list_tools(
-            &self,
-            _: Option<$crate::model::PaginatedRequestParam>,
-            _: $crate::service::RequestContext<$crate::service::RoleServer>,
-        ) -> Result<$crate::model::ListToolsResult, $crate::Error> {
-            Ok($crate::model::ListToolsResult {
-                next_cursor: None,
-                tools: Self::tool_box().list(),
-            })
-        }
-
-        async fn call_tool(
-            &self,
-            call_tool_request_param: $crate::model::CallToolRequestParam,
-            context: $crate::service::RequestContext<$crate::service::RoleServer>,
-        ) -> Result<$crate::model::CallToolResult, $crate::Error> {
-            let context = $crate::handler::server::tool::ToolCallContext::new(self, call_tool_request_param, context);
-            Self::$tool_box().call(context).await
-        }
-    }
-}