From 8c40aa81a65585bcbec2e2137a6f3463e0de71b9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:26:12 +0000 Subject: [PATCH 1/2] Initial plan From 28b11e4a71552b45d269aede8073fb01879b446e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:36:22 +0000 Subject: [PATCH 2/2] Add opt-in BM25 tool search index and server wrapper --- crates/rmcp/Cargo.toml | 2 + crates/rmcp/README.md | 49 ++++ crates/rmcp/src/handler/server.rs | 70 +++-- crates/rmcp/src/lib.rs | 3 + crates/rmcp/src/tool_search.rs | 185 +++++++++++++ crates/rmcp/src/tool_search/server.rs | 262 +++++++++++++++++++ crates/rmcp/tests/test_tool_search.rs | 361 ++++++++++++++++++++++++++ 7 files changed, 912 insertions(+), 20 deletions(-) create mode 100644 crates/rmcp/src/tool_search.rs create mode 100644 crates/rmcp/src/tool_search/server.rs create mode 100644 crates/rmcp/tests/test_tool_search.rs diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 38551b433..3e46170ce 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -33,6 +33,7 @@ features = [ "server", "server-side-http", "tower", + "tool-search", "transport-async-rw", "transport-child-process", "transport-io", @@ -138,6 +139,7 @@ libc = "0.2" default = ["base64", "macros", "server"] local = ["rmcp-macros?/local"] client = ["dep:tokio-stream"] +tool-search = [] server = ["transport-async-rw", "schemars", "dep:pastey", "uuid"] macros = ["dep:rmcp-macros", "dep:pastey"] elicitation = ["dep:url"] diff --git a/crates/rmcp/README.md b/crates/rmcp/README.md index 6f57acfd5..e4b33bea0 100644 --- a/crates/rmcp/README.md +++ b/crates/rmcp/README.md @@ -27,6 +27,7 @@ For **getting started**, **usage guides**, and **full MCP feature documentation* | `auth-enterprise-managed` | EMA/XAA refresh-token and ID-JAG exchanges for registered public and confidential clients (includes `auth`) | | | `elicitation` | Elicitation support | | | `skills-fs` | Serve SEP-2640 skills from a directory tree (`SkillCatalog`) | | +| `tool-search` | Dependency-free BM25 tool search and optional server discovery wrapper | | ### Transport features @@ -67,6 +68,54 @@ Any type that implements the [`Transport`](crate::transport::Transport) trait ca 3. A [`Worker`](crate::transport::worker::Worker) implementation 4. A [`Transport`](crate::transport::Transport) implementation directly +## Tool search + +Enable `tool-search` for a standalone, lazily built BM25 Okapi index: + +```rust,ignore +use rmcp::tool_search::ToolSearchIndex; + +let mut index = ToolSearchIndex::new(tools); +let hits = index.search("create a github issue", 5); +// After an upstream tools/list_changed notification, fetch the new catalog: +index.replace_tools(updated_tools); +``` + +Names, titles, descriptions, and input-schema parameter names/descriptions are +searched case-insensitively. Use `with_parameter_indexing(tools, false)` to omit +parameters. Hits include full `Tool` definitions and scores; unmatched queries +return no hits. + +With `server` enabled, wrap your existing `ServerHandler` (including handlers +using `#[tool_handler]` / `ToolRouter`): + +```rust,ignore +use rmcp::tool_search::{ToolSearchConfig, ToolSearchServer}; + +let mut config = ToolSearchConfig::default(); +config.max_results = 5; +config.always_visible = vec!["health_check".into()]; +let server = ToolSearchServer::with_config(my_handler, config)?; +// Serve `server` as usual. +``` + +`tools/list` exposes `search_tools` and `call_tool`, plus `always_visible` tools +which are excluded from search. Search arguments are +`{"query": "create an issue", "max_results": 3}`; the optional limit defaults to +the configured maximum (five by default) and cannot exceed it. Search returns +`{"tools": [...]}` as both structured and text content, including input schemas. +Dispatch arguments are `{"name": "create_issue", "arguments": {...}}`. +Original tools remain directly callable, and dispatch preserves request metadata +and retry state. Search controls discovery, **not authorization**: underlying +handlers must still enforce access checks. + +Use `search_tool_name` / `call_tool_name` to avoid catalog collisions. Invalid +settings are rejected at construction, and catalog collisions fail listing/search. +The wrapper fetches every catalog page on each listing/search and lazily rebuilds +its index only when definitions change, so dynamic or request-specific catalogs +do not require a notification invalidation hook. Other handler methods and +notifications are delegated unchanged. Regex and embedding search are not included. + ## License This project is licensed under the terms specified in the repository's [LICENSE](../../LICENSE) file. diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index c28d66312..14416916a 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -748,7 +748,31 @@ pub trait ServerHandler: Sized + 'static { macro_rules! impl_server_handler_for_wrapper { ($wrapper:ident) => { + impl_server_handler_for_wrapper!($wrapper, { + fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + (**self).call_tool(request, context) + } + + fn list_tools( + &self, + request: Option, + context: RequestContext, + ) -> impl Future> + MaybeSendFuture + '_ { + (**self).list_tools(request, context) + } + + fn get_tool(&self, name: &str) -> Option { + (**self).get_tool(name) + } + }); + }; + ($wrapper:ident, { $($tool_methods:tt)* }) => { impl ServerHandler for $wrapper { + $($tool_methods)* fn ping( &self, context: RequestContext, @@ -869,26 +893,6 @@ macro_rules! impl_server_handler_for_wrapper { (**self).unsubscribe(request, context) } - fn call_tool( - &self, - request: CallToolRequestParams, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - (**self).call_tool(request, context) - } - - fn list_tools( - &self, - request: Option, - context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - (**self).list_tools(request, context) - } - - fn get_tool(&self, name: &str) -> Option { - (**self).get_tool(name) - } - fn on_custom_request( &self, request: CustomRequest, @@ -995,3 +999,29 @@ macro_rules! impl_server_handler_for_wrapper { impl_server_handler_for_wrapper!(Box); impl_server_handler_for_wrapper!(Arc); + +#[cfg(feature = "tool-search")] +use crate::tool_search::ToolSearchServer; + +#[cfg(feature = "tool-search")] +impl_server_handler_for_wrapper!(ToolSearchServer, { + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + self.search_call_tool(request, context).await + } + + async fn list_tools( + &self, + request: Option, + context: RequestContext, + ) -> Result { + self.search_list_tools(request, context).await + } + + fn get_tool(&self, name: &str) -> Option { + self.search_get_tool(name) + } +}); diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index 3be6616ed..3b4135915 100644 --- a/crates/rmcp/src/lib.rs +++ b/crates/rmcp/src/lib.rs @@ -28,6 +28,9 @@ pub use service::{RoleServer, serve_server}; pub mod handler; #[cfg(feature = "server")] pub mod task_manager; +/// Dependency-free BM25 tool discovery, with an optional server wrapper. +#[cfg(feature = "tool-search")] +pub mod tool_search; #[cfg(any(feature = "client", feature = "server"))] pub mod transport; diff --git a/crates/rmcp/src/tool_search.rs b/crates/rmcp/src/tool_search.rs new file mode 100644 index 000000000..620e74d1e --- /dev/null +++ b/crates/rmcp/src/tool_search.rs @@ -0,0 +1,185 @@ +//! Opt-in, in-memory BM25 Okapi search over tool definitions. +//! +//! Matching is case-insensitive, splitting on punctuation (including underscores). +//! Scores use `k1 = 1.5` and `b = 0.75`; only positive matches are returned. +//! Equal scores retain catalog order. This is discovery, not access control. + +use std::{ + collections::{BTreeSet, HashMap}, + sync::OnceLock, +}; + +use crate::model::Tool; + +#[cfg(feature = "server")] +mod server; +#[cfg(feature = "server")] +pub use server::{ToolSearchConfig, ToolSearchServer}; + +/// A matching tool, including its full definition and BM25 score. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct ToolMatch { + /// The complete matching definition, including schemas and metadata. + pub tool: Tool, + /// Positive BM25 Okapi relevance score. + pub score: f64, +} + +/// A lazily built search index usable by servers, clients, or gateways. +#[derive(Debug)] +pub struct ToolSearchIndex { + tools: Vec, + index_parameters: bool, + index: OnceLock, +} + +#[derive(Debug)] +struct Bm25 { + documents: Vec>, + lengths: Vec, + frequencies: HashMap, + average_length: f64, +} + +impl ToolSearchIndex { + /// Index names, titles, descriptions, and parameter names/descriptions. + pub fn new(tools: Vec) -> Self { + Self::with_parameter_indexing(tools, true) + } + + /// Choose whether to include input-schema parameter names/descriptions. + pub fn with_parameter_indexing(tools: Vec, index_parameters: bool) -> Self { + Self { + tools, + index_parameters, + index: OnceLock::new(), + } + } + + /// Replace the catalog and invalidate the index, e.g. after `tools/list_changed`. + pub fn replace_tools(&mut self, tools: Vec) { + self.tools = tools; + self.index.take(); + } + + /// The full catalog backing this index. + pub fn tools(&self) -> &[Tool] { + &self.tools + } + + /// Return up to `max_results` ranked tools. Empty/unmatched queries return no hits. + pub fn search(&self, query: &str, max_results: usize) -> Vec { + let terms: BTreeSet<_> = tokenize(query).collect(); + if terms.is_empty() || max_results == 0 || self.tools.is_empty() { + return Vec::new(); + } + let index = self.index.get_or_init(|| self.build()); + let count = self.tools.len() as f64; + let mut scores: Vec<_> = index + .documents + .iter() + .enumerate() + .filter_map(|(i, document)| { + let mut score = 0.0; + for term in &terms { + if let Some(&frequency) = document.get(term) { + let df = index.frequencies[term] as f64; + let idf = (1.0 + (count - df + 0.5) / (df + 0.5)).ln(); + let tf = frequency as f64; + let length = index.lengths[i] as f64 / index.average_length; + score += idf * tf * 2.5 / (tf + 1.5 * (0.25 + 0.75 * length)); + } + } + (score > 0.0).then_some((i, score)) + }) + .collect(); + scores.sort_by(|(a, sa), (b, sb)| sb.total_cmp(sa).then_with(|| a.cmp(b))); + scores + .into_iter() + .take(max_results) + .map(|(i, score)| ToolMatch { + tool: self.tools[i].clone(), + score, + }) + .collect() + } + + fn build(&self) -> Bm25 { + let mut documents = Vec::with_capacity(self.tools.len()); + let mut lengths = Vec::with_capacity(self.tools.len()); + let mut frequencies = HashMap::new(); + for tool in &self.tools { + let mut text = tool.name.to_string(); + for field in [tool.title.as_deref(), tool.description.as_deref()] + .into_iter() + .flatten() + { + text.push(' '); + text.push_str(field); + } + if self.index_parameters { + // Iterative traversal handles nested schemas without recursive stack growth. + let mut pending: Vec<&serde_json::Value> = tool.input_schema.values().collect(); + if let Some(properties) = tool + .input_schema + .get("properties") + .and_then(|v| v.as_object()) + { + for name in properties.keys() { + text.push(' '); + text.push_str(name); + } + } + while let Some(value) = pending.pop() { + match value { + serde_json::Value::Object(object) => { + if let Some(description) = + object.get("description").and_then(|v| v.as_str()) + { + text.push(' '); + text.push_str(description); + } + if let Some(properties) = + object.get("properties").and_then(|v| v.as_object()) + { + for name in properties.keys() { + text.push(' '); + text.push_str(name); + } + } + pending.extend(object.values()); + } + serde_json::Value::Array(array) => pending.extend(array), + _ => {} + } + } + } + let mut document = HashMap::new(); + let mut length = 0; + for term in tokenize(&text) { + *document.entry(term).or_insert(0) += 1; + length += 1; + } + for term in document.keys() { + *frequencies.entry(term.clone()).or_insert(0) += 1; + } + documents.push(document); + lengths.push(length); + } + let average_length = + (lengths.iter().sum::() as f64 / self.tools.len().max(1) as f64).max(1.0); + Bm25 { + documents, + lengths, + frequencies, + average_length, + } + } +} + +fn tokenize(text: &str) -> impl Iterator + '_ { + text.split(|c: char| !c.is_alphanumeric()) + .filter(|term| !term.is_empty()) + .map(str::to_lowercase) +} diff --git a/crates/rmcp/src/tool_search/server.rs b/crates/rmcp/src/tool_search/server.rs new file mode 100644 index 000000000..0f425994b --- /dev/null +++ b/crates/rmcp/src/tool_search/server.rs @@ -0,0 +1,262 @@ +use std::{collections::HashSet, ops::Deref}; + +use serde::Deserialize; +use serde_json::json; +use tokio::sync::Mutex; + +use super::ToolSearchIndex; +use crate::{ + ErrorData, RoleServer, ServerHandler, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, JsonObject, ListToolsResult, + PaginatedRequestParams, Tool, ToolAnnotations, + }, + service::RequestContext, +}; + +/// Discovery settings. The result limit is also an upper bound on caller-supplied limits. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct ToolSearchConfig { + /// Default and maximum number of search results (must be positive). + pub max_results: usize, + /// Catalog tool names retained in listings and excluded from search. + pub always_visible: Vec, + /// Name of the synthetic search tool. + pub search_tool_name: String, + /// Name of the synthetic dispatch tool. + pub call_tool_name: String, + /// Include input-schema parameter names and descriptions in the index. + pub index_parameters: bool, +} + +impl Default for ToolSearchConfig { + fn default() -> Self { + Self { + max_results: 5, + always_visible: Vec::new(), + search_tool_name: "search_tools".into(), + call_tool_name: "call_tool".into(), + index_parameters: true, + } + } +} + +/// Wrap any `ServerHandler`, including one backed by a `ToolRouter`. +/// +/// The underlying catalog is fetched (including all pages) for every search/list, +/// so dynamic and request-specific catalogs remain current. The BM25 index is +/// built lazily and reused only when the searchable definitions are unchanged. +/// Outgoing `tools/list_changed` notifications need no additional invalidation hook. +/// Original tools remain directly callable; this wrapper does not grant access. +#[derive(Debug)] +pub struct ToolSearchServer { + inner: H, + config: ToolSearchConfig, + index: Mutex, +} + +impl Deref for ToolSearchServer { + type Target = H; + + fn deref(&self) -> &H { + &self.inner + } +} + +impl ToolSearchServer { + /// Use the default synthetic names and a maximum of five results. + pub fn new(inner: H) -> Self { + Self::with_config(inner, ToolSearchConfig::default()) + .expect("default tool search configuration is valid") + } + + /// Validate settings. Catalog name collisions are rejected when listing/searching. + pub fn with_config(inner: H, config: ToolSearchConfig) -> Result { + let valid_name = |name: &str| { + !name.is_empty() + && name.len() <= 128 + && name + .bytes() + .all(|c| c.is_ascii_alphanumeric() || b"_.-".contains(&c)) + }; + if config.max_results == 0 + || !valid_name(&config.search_tool_name) + || !valid_name(&config.call_tool_name) + || config.search_tool_name == config.call_tool_name + || config + .always_visible + .iter() + .any(|name| name == &config.search_tool_name || name == &config.call_tool_name) + { + return Err(ErrorData::invalid_params( + "invalid tool search configuration", + None, + )); + } + let index = Mutex::new(ToolSearchIndex::with_parameter_indexing( + Vec::new(), + config.index_parameters, + )); + Ok(Self { + inner, + config, + index, + }) + } + + async fn catalog( + &self, + context: RequestContext, + ) -> Result { + let mut result = self.inner.list_tools(None, context.clone()).await?; + let mut cursors = HashSet::new(); + while let Some(cursor) = result.next_cursor.take() { + if !cursors.insert(cursor.clone()) { + return Err(ErrorData::internal_error( + "tool catalog repeated a pagination cursor", + None, + )); + } + let page = self + .inner + .list_tools( + Some(PaginatedRequestParams { + cursor: Some(cursor), + ..Default::default() + }), + context.clone(), + ) + .await?; + result.tools.extend(page.tools); + result.next_cursor = page.next_cursor; + } + if result + .tools + .iter() + .any(|tool| self.is_synthetic(&tool.name)) + { + return Err(ErrorData::invalid_params( + "synthetic tool name collides with catalog; configure different names", + None, + )); + } + Ok(result) + } + + fn is_synthetic(&self, name: &str) -> bool { + name == self.config.search_tool_name || name == self.config.call_tool_name + } + + fn synthetic_tools(&self) -> [Tool; 2] { + let mut search = Tool::new( + self.config.search_tool_name.clone(), + "Search for tools using a natural-language query. Returns full tool definitions; call them directly or via the dispatch tool.", + json!({ + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "integer", "minimum": 1, "maximum": self.config.max_results, "default": self.config.max_results} + }, + "required": ["query"], + "additionalProperties": false + }).as_object().unwrap().clone(), + ); + search.annotations = Some(ToolAnnotations::new().read_only(true)); + let call = Tool::new( + self.config.call_tool_name.clone(), + "Call a discovered tool by name, passing its arguments.", + json!({ + "type": "object", + "properties": {"name": {"type": "string"}, "arguments": {"type": "object"}}, + "required": ["name"], + "additionalProperties": false + }) + .as_object() + .unwrap() + .clone(), + ); + [search, call] + } + + pub(crate) fn search_get_tool(&self, name: &str) -> Option { + self.synthetic_tools() + .into_iter() + .find(|tool| tool.name == name) + .or_else(|| self.inner.get_tool(name)) + } + + pub(crate) async fn search_list_tools( + &self, + request: Option, + context: RequestContext, + ) -> Result { + if request.and_then(|request| request.cursor).is_some() { + return Err(ErrorData::invalid_params( + "tool search listing has no pagination cursor", + None, + )); + } + let mut result = self.catalog(context).await?; + result + .tools + .retain(|tool| self.config.always_visible.contains(&tool.name.to_string())); + result.tools.extend(self.synthetic_tools()); + Ok(result) + } + + pub(crate) async fn search_call_tool( + &self, + mut request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name == self.config.search_tool_name { + #[derive(Deserialize)] + #[serde(deny_unknown_fields)] + struct Search { + query: String, + max_results: Option, + } + let args: Search = serde_json::from_value(json!(request.arguments)) + .map_err(|e| ErrorData::invalid_params(e.to_string(), None))?; + let limit = args.max_results.unwrap_or(self.config.max_results); + if limit == 0 || limit > self.config.max_results { + return Err(ErrorData::invalid_params( + "max_results is outside the configured range", + None, + )); + } + let mut tools = self.catalog(context).await?.tools; + tools.retain(|tool| !self.config.always_visible.contains(&tool.name.to_string())); + let mut index = self.index.lock().await; + if index.tools() != tools { + index.replace_tools(tools); + } + let tools: Vec<_> = index + .search(&args.query, limit) + .into_iter() + .map(|hit| hit.tool) + .collect(); + return Ok(CallToolResult::structured(json!({ "tools": tools })).into()); + } + if request.name == self.config.call_tool_name { + #[derive(Deserialize)] + #[serde(deny_unknown_fields)] + struct Call { + name: String, + arguments: Option, + } + let args: Call = serde_json::from_value(json!(request.arguments)) + .map_err(|e| ErrorData::invalid_params(e.to_string(), None))?; + if self.is_synthetic(&args.name) { + return Err(ErrorData::invalid_params( + "cannot dispatch a synthetic tool", + None, + )); + } + request.name = args.name.into(); + request.arguments = args.arguments; + } + self.inner.call_tool(request, context).await + } +} diff --git a/crates/rmcp/tests/test_tool_search.rs b/crates/rmcp/tests/test_tool_search.rs new file mode 100644 index 000000000..b0006f80c --- /dev/null +++ b/crates/rmcp/tests/test_tool_search.rs @@ -0,0 +1,361 @@ +#![cfg(feature = "tool-search")] + +use rmcp::{model::Tool, tool_search::ToolSearchIndex}; +use serde_json::json; + +fn tool(name: &str, description: &str) -> Tool { + Tool::new( + name.to_owned(), + description.to_owned(), + json!({"type": "object", "properties": {}}) + .as_object() + .unwrap() + .clone(), + ) +} + +#[test] +fn ranking_and_full_definitions() { + let tools = vec![ + tool("github_create_issue", "Create a GitHub issue"), + tool("github_list_issues", "List GitHub issues"), + tool("weather", "Get the weather forecast"), + ]; + let index = ToolSearchIndex::new(tools.clone()); + let hits = index.search("create a GITHUB issue", 2); + assert_eq!(hits.len(), 2); + assert_eq!(hits[0].tool, tools[0]); + assert!(hits[0].score > hits[1].score); + assert!( + hits.iter() + .all(|hit| hit.score.is_finite() && hit.score > 0.0) + ); + assert!(index.search("", 5).is_empty()); + assert!(index.search("___", 5).is_empty()); + assert!(index.search("unmatched", 5).is_empty()); + assert!(index.search("github", 0).is_empty()); + assert!(ToolSearchIndex::new(vec![]).search("github", 5).is_empty()); +} + +#[test] +fn bm25_formula_and_stable_ties() { + let index = ToolSearchIndex::new(vec![tool("x", "x"), tool("y", "")]); + let score = index.search("x", 1)[0].score; + let expected = 2.0_f64.ln() * 2.0 * 2.5 / (2.0 + 1.5 * (0.25 + 0.75 * 2.0 / 1.5)); + assert!((score - expected).abs() < 1e-12); + assert_eq!(index.search("x x", 1)[0].score, score); + + let tools = vec![tool("first", "common term"), tool("second", "common term")]; + let index = ToolSearchIndex::new(tools.clone()); + for _ in 0..10 { + assert_eq!(index.search("term common", 2)[0].tool, tools[0]); + } +} + +#[test] +fn title_parameters_unicode_and_replacement() { + let mut definition = tool("opaque", ""); + definition.title = Some("Résumé lookup".into()); + definition.input_schema = std::sync::Arc::new( + json!({ + "type": "object", + "properties": { + "repository_name": { + "type": "object", + "description": "Project identifier", + "properties": {"owner": {"type": "string", "description": "Maintainer"}} + } + } + }) + .as_object() + .unwrap() + .clone(), + ); + let mut index = ToolSearchIndex::new(vec![definition.clone()]); + for query in ["RÉSUMÉ", "repository", "identifier", "owner", "maintainer"] { + assert_eq!(index.search(query, 5)[0].tool, definition); + } + let no_parameters = ToolSearchIndex::with_parameter_indexing(vec![definition], false); + assert!(no_parameters.search("maintainer repository", 5).is_empty()); + index.replace_tools(vec![tool("replacement", "new catalog")]); + assert!(index.search("résumé", 5).is_empty()); + assert_eq!(index.search("catalog", 5)[0].tool.name, "replacement"); +} + +#[cfg(all(feature = "server", feature = "client", not(feature = "local")))] +mod server { + use std::sync::{Arc, RwLock}; + + use rmcp::{ + ErrorData, RoleServer, ServerHandler, ServiceExt, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ListPromptsResult, + ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig, + }, + service::RequestContext, + tool_search::{ToolSearchConfig, ToolSearchServer}, + }; + + use super::*; + + struct Catalog(Arc>>); + + impl ServerHandler for Catalog { + fn get_info(&self) -> ServerConfig { + ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .build(), + ) + } + + async fn list_tools( + &self, + request: Option, + _context: RequestContext, + ) -> Result { + let tools = self.0.read().unwrap(); + let offset: usize = request + .and_then(|r| r.cursor) + .map(|cursor| cursor.parse().unwrap()) + .unwrap_or(0); + Ok(ListToolsResult { + tools: tools.iter().skip(offset).take(1).cloned().collect(), + next_cursor: (offset + 1 < tools.len()).then(|| (offset + 1).to_string()), + ..Default::default() + }) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + if !self + .0 + .read() + .unwrap() + .iter() + .any(|tool| tool.name == request.name) + { + return Err(ErrorData::invalid_params("unknown tool", None)); + } + Ok(CallToolResult::structured(json!({ + "name": request.name, "arguments": request.arguments, "requestState": request.request_state + })).into()) + } + + async fn list_prompts( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + Err(ErrorData::invalid_params("delegated prompts", None)) + } + } + + #[tokio::test] + async fn discovery_dispatch_refresh_and_delegation() -> anyhow::Result<()> { + let definitions = vec![ + tool("visible", "GitHub status"), + tool("create_issue", "Create a GitHub issue"), + tool("list_issues", "List GitHub issues"), + ]; + let catalog = Arc::new(RwLock::new(definitions.clone())); + let mut config = ToolSearchConfig::default(); + config.always_visible = vec!["visible".into()]; + config.search_tool_name = "find".into(); + config.call_tool_name = "dispatch".into(); + let wrapper = ToolSearchServer::with_config(Catalog(catalog.clone()), config)?; + assert_eq!(wrapper.get_tool("find").unwrap().name, "find"); + let (server_io, client_io) = tokio::io::duplex(16384); + let server_task = tokio::spawn(async move { + let server = wrapper.serve(server_io).await.unwrap(); + server.waiting().await.unwrap(); + }); + let client = ().serve(client_io).await?; + let listing = client.list_tools(None).await?; + let names: Vec<_> = listing + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect(); + assert_eq!(names, ["visible", "find", "dispatch"]); + assert!(listing.next_cursor.is_none()); + let search = client + .call_tool( + CallToolRequestParams::new("find") + .with_arguments(json!({"query": "github"}).as_object().unwrap().clone()), + ) + .await?; + let hits: Vec = + serde_json::from_value(search.structured_content.unwrap()["tools"].clone())?; + assert_eq!(hits.len(), 2); + assert!(hits.contains(&definitions[1])); + assert!(hits.contains(&definitions[2])); + for name in ["create_issue", "dispatch"] { + let arguments = if name == "dispatch" { + json!({"name": "create_issue", "arguments": {"title": "test"}}) + } else { + json!({"title": "test"}) + }; + let result = client + .call_tool( + CallToolRequestParams::new(name) + .with_arguments(arguments.as_object().unwrap().clone()) + .with_request_state("opaque-state"), + ) + .await?; + let value = result.structured_content.unwrap(); + assert_eq!(value["name"], "create_issue"); + assert_eq!(value["arguments"], json!({"title": "test"})); + assert_eq!(value["requestState"], "opaque-state"); + } + for name in ["find", "dispatch"] { + assert!( + client + .call_tool( + CallToolRequestParams::new("dispatch") + .with_arguments(json!({"name": name}).as_object().unwrap().clone()), + ) + .await + .is_err() + ); + } + for arguments in [ + json!({"query": "github", "max_results": 0}), + json!({"query": "github", "max_results": 6}), + json!({"query": "github", "max_results": -1}), + json!({"query": 4}), + json!({"query": "github", "extra": true}), + ] { + assert!( + client + .call_tool( + CallToolRequestParams::new("find") + .with_arguments(arguments.as_object().unwrap().clone()), + ) + .await + .is_err() + ); + } + assert!( + client + .list_prompts(None) + .await + .unwrap_err() + .to_string() + .contains("delegated prompts") + ); + *catalog.write().unwrap() = vec![tool("weather", "Forecast")]; + let result = client + .call_tool( + CallToolRequestParams::new("find") + .with_arguments(json!({"query": "github"}).as_object().unwrap().clone()), + ) + .await?; + assert_eq!(result.structured_content.unwrap()["tools"], json!([])); + let result = client + .call_tool( + CallToolRequestParams::new("find") + .with_arguments(json!({"query": "forecast"}).as_object().unwrap().clone()), + ) + .await?; + assert_eq!( + result.structured_content.unwrap()["tools"][0]["name"], + "weather" + ); + *catalog.write().unwrap() = vec![tool("find", "Collision")]; + assert!(client.list_tools(None).await.is_err()); + client.cancel().await?; + server_task.await?; + Ok(()) + } + + #[test] + fn invalid_configuration() { + for (search, call, limit) in [ + ("same", "same", 5), + ("", "call", 5), + ("bad name", "call", 5), + ("search", "call", 0), + ] { + let mut config = ToolSearchConfig::default(); + config.search_tool_name = search.into(); + config.call_tool_name = call.into(); + config.max_results = limit; + assert!( + ToolSearchServer::with_config(Catalog(Arc::new(RwLock::new(vec![]))), config,) + .is_err() + ); + } + } + + #[cfg(feature = "macros")] + #[tokio::test] + async fn tool_router_integration() -> anyhow::Result<()> { + #[derive(Clone)] + struct Routed { + tool_router: rmcp::handler::server::router::tool::ToolRouter, + } + + #[rmcp::tool_router] + impl Routed { + fn new() -> Self { + Self { + tool_router: Self::tool_router(), + } + } + + #[rmcp::tool(description = "Read the weather forecast")] + async fn forecast(&self) -> String { + "sunny".into() + } + } + + #[rmcp::tool_handler(router = self.tool_router)] + impl ServerHandler for Routed { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + } + + let (server_io, client_io) = tokio::io::duplex(16384); + let server_task = tokio::spawn(async move { + ToolSearchServer::new(Routed::new()) + .serve(server_io) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + let client = ().serve(client_io).await?; + let listing = client.list_tools(None).await?; + assert_eq!(listing.tools.len(), 2); + let result = client + .call_tool( + CallToolRequestParams::new("search_tools") + .with_arguments(json!({"query": "weather"}).as_object().unwrap().clone()), + ) + .await?; + assert_eq!( + result.structured_content.unwrap()["tools"][0]["name"], + "forecast" + ); + let result = client + .call_tool( + CallToolRequestParams::new("call_tool") + .with_arguments(json!({"name": "forecast"}).as_object().unwrap().clone()), + ) + .await?; + assert_eq!( + result.content, + vec![rmcp::model::ContentBlock::text("sunny")] + ); + client.cancel().await?; + server_task.await?; + Ok(()) + } +}