fix(router): simplify disable tool api

Author: lutz-grex

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

@@ -83,8 +83,10 @@ where
     ) -> Result<<RoleServer as crate::service::ServiceRole>::Resp, crate::ErrorData> {
         match request {
             ClientRequest::CallToolRequest(request) => {
-                if self.tool_router.has_route(request.params.name.as_ref())
-                    || self.tool_router.is_disabled(request.params.name.as_ref())
+                if self
+                    .tool_router
+                    .map
+                    .contains_key(request.params.name.as_ref())
                     || !self.tool_router.transparent_when_not_found
                 {
                     let tool_call_context = crate::handler::server::tool::ToolCallContext::new(

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

@@ -409,11 +409,19 @@ where
         }
     }
 
+    /// Remove a tool route from the router.
+    ///
+    /// The disabled state is **preserved**: if the name was in the disabled
+    /// set, it stays there so that a future [`add_route`](Self::add_route)
+    /// or [`merge`](Self::merge) with the same name will inherit the
+    /// disabled state. To also clear the disabled marker, call
+    /// [`enable_route`](Self::enable_route) afterwards.
     pub fn remove_route(&mut self, name: &str) {
         self.map.remove(name);
-        self.disabled.remove(name);
     }
 
+    /// Returns `true` if the tool is registered **and** not currently
+    /// disabled.
     pub fn has_route(&self, name: &str) -> bool {
         self.map.contains_key(name) && !self.disabled.contains(name)
     }
@@ -422,20 +430,30 @@ where
     /// rejected by `call`. The tool remains in the router and can be
     /// re-enabled later with [`enable_route`](Self::enable_route).
     ///
+    /// Returns `true` if the name was newly added to the disabled set.
     /// The name is recorded even if no matching route exists yet, so routes
     /// added later (via [`add_route`](Self::add_route) or
     /// [`merge`](Self::merge)) will inherit the disabled state.
-    pub fn disable_route(&mut self, name: &str) {
-        self.disabled.insert(Cow::Owned(name.to_owned()));
+    ///
+    /// Callers should send `Peer::notify_tool_list_changed` when the
+    /// visible tool list changes. Accepts `&'static str` or `String`;
+    /// for a non-static `&str`, call `.to_owned()` first.
+    pub fn disable_route(&mut self, name: impl Into<Cow<'static, str>>) -> bool {
+        self.disabled.insert(name.into())
     }
 
-    /// Re-enable a previously disabled tool.
-    pub fn enable_route(&mut self, name: &str) {
-        self.disabled.remove(name);
+    /// Re-enable a previously disabled tool. Returns `true` if the name
+    /// was present in the disabled set and was removed.
+    ///
+    /// Callers should send `Peer::notify_tool_list_changed` when the
+    /// visible tool list changes.
+    pub fn enable_route(&mut self, name: &str) -> bool {
+        self.disabled.remove(name)
     }
 
-    /// Returns `true` if the tool exists in the router but is currently
-    /// disabled.
+    /// Returns `true` if the tool exists in the router **and** is currently
+    /// disabled. Returns `false` if the tool does not exist or if the name
+    /// was pre-disabled without a matching route.
     pub fn is_disabled(&self, name: &str) -> bool {
         self.map.contains_key(name) && self.disabled.contains(name)
     }
@@ -511,3 +529,49 @@ where
         self.merge(other);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::*;
+    use crate::{
+        RoleServer,
+        model::{CallToolRequestParams, ErrorCode, NumberOrString},
+        service::{AtomicU32RequestIdProvider, Peer, RequestContext},
+    };
+
+    struct DummyService;
+    impl crate::handler::server::ServerHandler for DummyService {}
+
+    #[tokio::test]
+    async fn test_call_disabled_tool_returns_error() {
+        let service = DummyService;
+        let mut router = ToolRouter::new().with_route(ToolRoute::new_dyn(
+            crate::model::Tool::new("test_tool", "a test tool", Arc::new(Default::default())),
+            |_ctx| Box::pin(async { Ok(CallToolResult::default()) }),
+        ));
+        router.disable_route("test_tool");
+
+        let id_provider: Arc<dyn crate::service::RequestIdProvider> =
+            Arc::new(AtomicU32RequestIdProvider::default());
+        let (peer, _rx) = Peer::<RoleServer>::new(id_provider, None);
+        let ctx = crate::handler::server::tool::ToolCallContext::new(
+            &service,
+            CallToolRequestParams {
+                meta: None,
+                name: Cow::Borrowed("test_tool"),
+                arguments: None,
+                task: None,
+            },
+            RequestContext::new(NumberOrString::Number(1), peer),
+        );
+
+        let err = router
+            .call(ctx)
+            .await
+            .expect_err("disabled tool should reject");
+        assert_eq!(err.code, ErrorCode::INVALID_PARAMS);
+        assert_eq!(err.message, "tool not found");
+    }
+}

crates/rmcp/tests/test_tool_routers.rs

@@ -100,7 +100,7 @@ fn test_disable_route() {
     assert!(router.has_route("async_function"));
     assert!(router.get("async_function").is_some());
 
-    router.disable_route("async_function");
+    assert!(router.disable_route("async_function"));
 
     assert_eq!(router.list_all().len(), 3);
     assert!(!router.has_route("async_function"));
@@ -116,10 +116,10 @@ fn test_disable_route() {
 #[test]
 fn test_enable_route() {
     let mut router = build_router();
-    router.disable_route("async_function");
+    assert!(router.disable_route("async_function"));
     assert!(!router.has_route("async_function"));
 
-    router.enable_route("async_function");
+    assert!(router.enable_route("async_function"));
     assert!(router.has_route("async_function"));
     assert!(router.get("async_function").is_some());
     assert!(!router.is_disabled("async_function"));
@@ -143,7 +143,7 @@ fn test_with_disabled_builder() {
 fn test_disabled_tools_survive_merge() {
     let mut router_a = ToolRouter::<TestHandler<()>>::new()
         .with_route((async_function_tool_attr(), async_function));
-    router_a.disable_route("async_function");
+    assert!(router_a.disable_route("async_function"));
 
     let router_b = ToolRouter::<TestHandler<()>>::new()
         .with_route((async_function2_tool_attr(), async_function2));
@@ -158,22 +158,42 @@ fn test_disabled_tools_survive_merge() {
 #[test]
 fn test_disable_nonexistent_tool() {
     let mut router = build_router();
-    // should not panic
-    router.disable_route("does_not_exist");
+    // should not panic; returns true because the name is newly added to disabled set
+    assert!(router.disable_route("does_not_exist"));
     assert_eq!(router.list_all().len(), 4);
     // is_disabled returns false for tools not in the map
     assert!(!router.is_disabled("does_not_exist"));
 }
 
 #[test]
-fn test_remove_route_clears_disabled_state() {
+fn test_remove_route_preserves_disabled_state() {
     let mut router = build_router();
-    router.disable_route("async_function");
+    assert!(router.disable_route("async_function"));
     assert!(router.is_disabled("async_function"));
 
     router.remove_route("async_function");
+    assert!(!router.has_route("async_function"));
+    // Disabled marker is preserved — is_disabled returns false (no route in map)
+    // but re-adding will inherit the disabled state (tested separately)
     assert!(!router.is_disabled("async_function"));
+}
+
+#[test]
+fn test_remove_route_then_readd_stays_disabled() {
+    let mut router = build_router();
+    assert!(router.disable_route("async_function"));
+
+    router.remove_route("async_function");
+    assert!(!router.has_route("async_function"));
+
+    // Re-add the route — it should inherit the disabled state
+    let other = ToolRouter::<TestHandler<()>>::new()
+        .with_route((async_function_tool_attr(), async_function));
+    router.merge(other);
+
     assert!(!router.has_route("async_function"));
+    assert!(router.is_disabled("async_function"));
+    assert!(router.get("async_function").is_none());
 }
 
 #[test]
@@ -211,6 +231,42 @@ fn test_disabled_tool_invisible_across_all_queries() {
     assert!(router.get("async_function").is_none());
     // Not routable
     assert!(!router.has_route("async_function"));
-    // But still known as disabled
+    // But known as disabled
     assert!(router.is_disabled("async_function"));
 }
+
+#[test]
+fn test_disable_route_then_add_route_blocks_tool() {
+    // Full pre-disable lifecycle via runtime mutation (not builder)
+    let mut router = ToolRouter::<TestHandler<()>>::new();
+    router.disable_route("async_function");
+
+    // Add route after disabling — tool should be blocked
+    let other = ToolRouter::<TestHandler<()>>::new()
+        .with_route((async_function_tool_attr(), async_function));
+    router.merge(other);
+
+    assert!(router.is_disabled("async_function"));
+    assert!(!router.has_route("async_function"));
+    assert!(router.get("async_function").is_none());
+    assert_eq!(router.list_all().len(), 0);
+}
+
+#[test]
+fn test_disable_enable_return_false_cases() {
+    let mut router = build_router();
+
+    // Repeated disable returns false
+    assert!(router.disable_route("async_function"));
+    assert!(!router.disable_route("async_function"));
+
+    // Enable returns true, then false on repeat
+    assert!(router.enable_route("async_function"));
+    assert!(!router.enable_route("async_function"));
+
+    // Enable on name never disabled returns false
+    assert!(!router.enable_route("async_function2"));
+
+    // Enable on unknown name returns false
+    assert!(!router.enable_route("unknown"));
+}