feat(router): support runtime disabling of tools (#809)

* feat(router): support runtime disabling of tools

Add methods to disable/enable tools at runtime.
Disabled tools are hidden from listing, lookup,
and execution, including in composed routers.

Closes #477

* fix(router): simplify disable tool api

* feat(router): auto-send tools/list_changed on disable/enable

* refactor(router): simplify disable_route and notifier call
This commit is contained in:
lutz-grex 2026-04-22 14:10:29 +02:00 • committed by GitHub
commit 63583b164f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 750 additions and 15 deletions

View file

@ -0,0 +1,172 @@
//! Integration tests for tool list change notifications.
#![cfg(all(feature = "client", not(feature = "local")))]
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use rmcp::{
ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceExt,
handler::server::{router::tool::ToolRoute, tool::ToolCallContext},
model::{CallToolResult, ServerCapabilities, ServerInfo, Tool},
service::{MaybeSendFuture, NotificationContext},
};
use tokio::sync::{Notify, RwLock};
#[derive(Clone)]
struct TestToolServer {
router: Arc<RwLock<rmcp::handler::server::router::tool::ToolRouter<Self>>>,
trigger_disable: Arc<Notify>,
trigger_enable: Arc<Notify>,
}
impl TestToolServer {
fn new() -> Self {
let mut tool_router = rmcp::handler::server::router::tool::ToolRouter::<Self>::new();
tool_router.add_route(ToolRoute::new_dyn(
Tool::new("tool_a", "Tool A", Arc::new(Default::default())),
|_ctx| Box::pin(async { Ok(CallToolResult::default()) }),
));
tool_router.add_route(ToolRoute::new_dyn(
Tool::new("tool_b", "Tool B", Arc::new(Default::default())),
|_ctx| Box::pin(async { Ok(CallToolResult::default()) }),
));
Self {
router: Arc::new(RwLock::new(tool_router)),
trigger_disable: Arc::new(Notify::new()),
trigger_enable: Arc::new(Notify::new()),
}
}
}
impl ServerHandler for TestToolServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
}
fn call_tool(
&self,
request: rmcp::model::CallToolRequestParams,
context: rmcp::service::RequestContext<RoleServer>,
) -> impl std::future::Future<Output = Result<CallToolResult, rmcp::ErrorData>> + MaybeSendFuture + '_
{
async move {
let router = self.router.read().await;
let tcc = ToolCallContext::new(self, request, context);
router.call(tcc).await
}
}
fn list_tools(
&self,
_request: Option<rmcp::model::PaginatedRequestParams>,
_context: rmcp::service::RequestContext<RoleServer>,
) -> impl std::future::Future<Output = Result<rmcp::model::ListToolsResult, rmcp::ErrorData>>
+ MaybeSendFuture
+ '_ {
async move {
let router = self.router.read().await;
Ok(rmcp::model::ListToolsResult {
tools: router.list_all(),
..Default::default()
})
}
}
fn on_initialized(
&self,
context: NotificationContext<RoleServer>,
) -> impl std::future::Future<Output = ()> + MaybeSendFuture + '_ {
let router = self.router.clone();
let trigger_disable = self.trigger_disable.clone();
let trigger_enable = self.trigger_enable.clone();
let peer = context.peer.clone();
async move {
router.write().await.bind_peer_notifier(&peer);
let router = router.clone();
tokio::spawn(async move {
trigger_disable.notified().await;
{
let mut r = router.write().await;
r.disable_route("tool_a");
}
trigger_enable.notified().await;
{
let mut r = router.write().await;
r.enable_route("tool_a");
}
});
}
}
}
#[derive(Clone)]
struct TestToolClient {
notification_count: Arc<AtomicUsize>,
notify: Arc<Notify>,
}
impl TestToolClient {
fn new() -> Self {
Self {
notification_count: Arc::new(AtomicUsize::new(0)),
notify: Arc::new(Notify::new()),
}
}
}
impl ClientHandler for TestToolClient {
fn on_tool_list_changed(
&self,
_context: NotificationContext<RoleClient>,
) -> impl std::future::Future<Output = ()> + MaybeSendFuture + '_ {
self.notification_count.fetch_add(1, Ordering::SeqCst);
self.notify.notify_one();
std::future::ready(())
}
}
#[tokio::test]
async fn test_disable_enable_sends_tool_list_changed() {
let server = TestToolServer::new();
let trigger_disable = server.trigger_disable.clone();
let trigger_enable = server.trigger_enable.clone();
let client = TestToolClient::new();
let notification_count = client.notification_count.clone();
let client_notify = client.notify.clone();
let (server_transport, client_transport) = tokio::io::duplex(4096);
let server_handle = tokio::spawn(async move { server.serve(server_transport).await });
let client_service = client.serve(client_transport).await.unwrap();
let tools = client_service.peer().list_tools(None).await.unwrap();
assert_eq!(tools.tools.len(), 2);
trigger_disable.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(5), client_notify.notified())
.await
.expect("timed out waiting for tool_list_changed");
assert_eq!(notification_count.load(Ordering::SeqCst), 1);
let tools = client_service.peer().list_tools(None).await.unwrap();
assert_eq!(tools.tools.len(), 1);
assert_eq!(tools.tools[0].name, "tool_b");
trigger_enable.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(5), client_notify.notified())
.await
.expect("timed out waiting for tool_list_changed");
assert_eq!(notification_count.load(Ordering::SeqCst), 2);
let tools = client_service.peer().list_tools(None).await.unwrap();
assert_eq!(tools.tools.len(), 2);
client_service.cancel().await.unwrap();
server_handle.abort();
}

View file

@ -1,5 +1,8 @@
#![cfg(not(feature = "local"))]
use std::collections::HashMap;
use std::{
collections::HashMap,
sync::atomic::{AtomicUsize, Ordering},
};
use futures::future::BoxFuture;
use rmcp::{
@ -84,3 +87,286 @@ fn test_tool_router_list_all_is_sorted() {
"list_all() should return tools sorted alphabetically by name"
);
}
fn build_router() -> ToolRouter<TestHandler<()>> {
ToolRouter::<TestHandler<()>>::new()
.with_route((async_function_tool_attr(), async_function))
.with_route((async_function2_tool_attr(), async_function2))
+ TestHandler::<()>::test_router_1()
+ TestHandler::<()>::test_router_2()
}
#[test]
fn test_disable_route() {
let mut router = build_router();
assert_eq!(router.list_all().len(), 4);
assert!(router.has_route("async_function"));
assert!(router.get("async_function").is_some());
assert!(router.disable_route("async_function"));
assert_eq!(router.list_all().len(), 3);
assert!(!router.has_route("async_function"));
assert!(router.get("async_function").is_none());
assert!(router.is_disabled("async_function"));
// other tools unaffected
assert!(router.has_route("async_function2"));
assert!(router.get("async_function2").is_some());
assert!(!router.is_disabled("async_function2"));
}
#[test]
fn test_enable_route() {
let mut router = build_router();
assert!(router.disable_route("async_function"));
assert!(!router.has_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"));
assert_eq!(router.list_all().len(), 4);
}
#[test]
fn test_with_disabled_builder() {
let router = build_router()
.with_disabled("async_function")
.with_disabled("sync_method");
assert_eq!(router.list_all().len(), 2);
assert!(!router.has_route("async_function"));
assert!(!router.has_route("sync_method"));
assert!(router.has_route("async_function2"));
assert!(router.has_route("async_method"));
}
#[test]
fn test_disabled_tools_survive_merge() {
let mut router_a = ToolRouter::<TestHandler<()>>::new()
.with_route((async_function_tool_attr(), async_function));
assert!(router_a.disable_route("async_function"));
let router_b = ToolRouter::<TestHandler<()>>::new()
.with_route((async_function2_tool_attr(), async_function2));
router_a.merge(router_b);
assert_eq!(router_a.list_all().len(), 1);
assert!(router_a.is_disabled("async_function"));
assert!(router_a.has_route("async_function2"));
}
#[test]
fn test_disable_nonexistent_tool() {
let mut router = build_router();
// 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_preserves_disabled_state() {
let mut router = build_router();
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]
fn test_into_iter_skips_disabled() {
let router = build_router().with_disabled("async_function");
let names: Vec<_> = router
.into_iter()
.map(|r| r.attr.name.to_string())
.collect();
assert_eq!(names.len(), 3);
assert!(!names.contains(&"async_function".to_string()));
}
#[test]
fn test_pre_disable_before_add_route() {
// Disabling a name before adding a route with that name should
// result in the route being disabled once added.
let router = ToolRouter::<TestHandler<()>>::new()
.with_disabled("async_function")
.with_route((async_function_tool_attr(), async_function));
assert_eq!(router.list_all().len(), 0);
assert!(router.is_disabled("async_function"));
assert!(!router.has_route("async_function"));
}
#[test]
fn test_disabled_tool_invisible_across_all_queries() {
let router = build_router().with_disabled("async_function");
// Not listed
let names: Vec<_> = router.list_all().iter().map(|t| t.name.clone()).collect();
assert!(!names.contains(&"async_function".into()));
// Not retrievable
assert!(router.get("async_function").is_none());
// Not routable
assert!(!router.has_route("async_function"));
// 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"));
}
// ── Notifier tests ──────────────────────────────────────────────────────
fn counter_notifier() -> (
impl Fn() + Send + Sync + 'static,
std::sync::Arc<AtomicUsize>,
) {
let counter = std::sync::Arc::new(AtomicUsize::new(0));
let c = counter.clone();
let notifier = move || {
c.fetch_add(1, Ordering::SeqCst);
};
(notifier, counter)
}
#[test]
fn test_notifier_fires_on_disable_and_enable() {
let (notifier, counter) = counter_notifier();
let mut router = build_router();
router.set_notifier(notifier);
assert!(router.disable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 1);
assert!(!router.disable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 1);
assert!(router.enable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 2);
assert!(!router.enable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 2);
}
#[test]
fn test_notifier_skips_nonexistent_tools() {
let (notifier, counter) = counter_notifier();
let mut router = build_router();
router.set_notifier(notifier);
assert!(router.disable_route("does_not_exist"));
assert_eq!(counter.load(Ordering::SeqCst), 0);
assert!(router.enable_route("does_not_exist"));
assert_eq!(counter.load(Ordering::SeqCst), 0);
assert!(router.disable_route("future_tool"));
assert_eq!(counter.load(Ordering::SeqCst), 0);
assert!(router.enable_route("future_tool"));
assert_eq!(counter.load(Ordering::SeqCst), 0);
}
#[test]
fn test_no_notifier_no_panic() {
let mut router = build_router();
assert!(router.disable_route("async_function"));
assert!(router.enable_route("async_function"));
assert!(router.disable_route("async_function"));
assert!(!router.disable_route("async_function"));
}
#[test]
fn test_clone_shares_notifier() {
let (notifier, counter) = counter_notifier();
let mut router = build_router();
router.set_notifier(notifier);
let mut cloned = router.clone();
assert!(cloned.disable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 1);
assert!(router.disable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 2);
cloned.clear_notifier();
assert!(cloned.enable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 2);
assert!(router.enable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 3);
}
#[test]
fn test_pre_init_disable_silent_but_correct() {
let mut router = build_router();
assert!(router.disable_route("async_function"));
assert_eq!(router.list_all().len(), 3);
assert!(!router.has_route("async_function"));
let (notifier, counter) = counter_notifier();
router.set_notifier(notifier);
assert_eq!(counter.load(Ordering::SeqCst), 0);
assert!(router.enable_route("async_function"));
assert_eq!(counter.load(Ordering::SeqCst), 1);
}