From 0f3bb9a51f7fd0126582465e75d01fee5ed2fb06 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 6 Aug 2026 11:13:25 +0800 Subject: [PATCH] feat(rmcp): separate model and wire tool names --- crates/rig-agent/src/agent/builder.rs | 96 ++++++++++++- crates/rig-agent/src/tool/rmcp.rs | 196 +++++++++++++++++++++++--- crates/rig-agent/src/tool/server.rs | 27 +++- 3 files changed, 290 insertions(+), 29 deletions(-) diff --git a/crates/rig-agent/src/agent/builder.rs b/crates/rig-agent/src/agent/builder.rs index 9e7626cc5e..10b6e9342e 100644 --- a/crates/rig-agent/src/agent/builder.rs +++ b/crates/rig-agent/src/agent/builder.rs @@ -21,7 +21,7 @@ use crate::{ #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] -use crate::tool::rmcp::McpTool as RmcpTool; +use crate::tool::rmcp::{McpTool as RmcpTool, RmcpToolRegistration}; use super::{Agent, ModelHandle, OutputMode}; @@ -79,12 +79,22 @@ fn build_rmcp_tools( client: rmcp::service::ServerSink, timeout: Option, ) -> Vec<(String, RmcpTool)> { - tools + build_rmcp_tools_from_registrations(tools.into_iter().map(|tool| { + let model_name = tool.name.to_string(); + RmcpToolRegistration::new(model_name, tool, client.clone()).with_timeout(timeout) + })) +} + +#[cfg(all(feature = "rmcp", not(target_family = "wasm")))] +fn build_rmcp_tools_from_registrations( + registrations: impl IntoIterator, +) -> Vec<(String, RmcpTool)> { + registrations .into_iter() - .map(|tool| { - let name = tool.name.to_string(); - let rmcp_tool = RmcpTool::from_mcp_server(tool, client.clone()).with_timeout(timeout); - (name, rmcp_tool) + .map(|registration| { + let model_name = registration.model_name.clone(); + let rmcp_tool = RmcpTool::from_registration(registration); + (model_name, rmcp_tool) }) .collect() } @@ -562,6 +572,33 @@ impl AgentBuilder { self.with_rmcp_toolset(build_rmcp_tools(tools, client, timeout.into())) } + /// Add one MCP registration with a separate model-visible name. + /// + /// `registration.model_name` is used for provider-facing definitions and + /// tool lookup, while `registration.definition.name` is used for MCP + /// `tools/call` requests. + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] + #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] + pub fn rmcp_tool_registration( + self, + registration: RmcpToolRegistration, + ) -> AgentBuilder { + self.rmcp_tool_registrations([registration]) + } + + /// Add MCP registrations with independent model-visible and wire names. + /// + /// The registrations retain their configured timeout and preserve the + /// server-advertised name in each MCP `tools/call` request. + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] + #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] + pub fn rmcp_tool_registrations( + self, + registrations: impl IntoIterator, + ) -> AgentBuilder { + self.with_rmcp_toolset(build_rmcp_tools_from_registrations(registrations)) + } + /// Transition into the `WithBuilderTools` state carrying the given built /// MCP tools. #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] @@ -744,6 +781,30 @@ impl AgentBuilder { self.add_rmcp_tools(build_rmcp_tools(tools, client, timeout.into())) } + /// Add one MCP registration with a separate model-visible name. + /// + /// `registration.model_name` is used for provider-facing definitions and + /// tool lookup, while `registration.definition.name` is used for MCP + /// `tools/call` requests. + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] + #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] + pub fn rmcp_tool_registration(self, registration: RmcpToolRegistration) -> Self { + self.add_rmcp_tools(build_rmcp_tools_from_registrations([registration])) + } + + /// Add MCP registrations with independent model-visible and wire names. + /// + /// The registrations retain their configured timeout and preserve the + /// server-advertised name in each MCP `tools/call` request. + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] + #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] + pub fn rmcp_tool_registrations( + self, + registrations: impl IntoIterator, + ) -> Self { + self.add_rmcp_tools(build_rmcp_tools_from_registrations(registrations)) + } + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] fn add_rmcp_tools(mut self, built: Vec<(String, RmcpTool)>) -> Self { for (_, tool) in built { @@ -1015,6 +1076,29 @@ mod tests { assert!(result.is_error_kind(ToolErrorKind::Timeout)); assert!(result.output().render().contains("timed out")); + // Both builder typestates expose the alias-aware registration path. + let no_tool_config_agent = AgentBuilder::new(MockCompletionModel::text("ok")) + .rmcp_tool_registration( + RmcpToolRegistration::new("model_alias", tool("wire_name"), client.peer().clone()) + .with_timeout(Some(Duration::from_secs(1))), + ) + .build(); + let builder_tools_agent = AgentBuilder::new(MockCompletionModel::text("ok")) + .tool(MockAddTool) + .rmcp_tool_registration( + RmcpToolRegistration::new("model_alias", tool("wire_name"), client.peer().clone()) + .with_timeout(Some(Duration::from_secs(1))), + ) + .build(); + for agent in [no_tool_config_agent, builder_tools_agent] { + let definitions = agent.tool_server_handle.get_tool_defs(None).await.unwrap(); + assert!( + definitions + .iter() + .any(|definition| definition.name == "model_alias") + ); + } + drop(client); server_task.abort(); } diff --git a/crates/rig-agent/src/tool/rmcp.rs b/crates/rig-agent/src/tool/rmcp.rs index c17ad645c4..c8d32b6f7d 100644 --- a/crates/rig-agent/src/tool/rmcp.rs +++ b/crates/rig-agent/src/tool/rmcp.rs @@ -4,6 +4,8 @@ //! `notifications/tools/list_changed` by re-fetching the tool list and updating //! the [`ToolServer`](super::server::ToolServer). Individual MCP tools are //! registered through the agent and tool-server `rmcp_tool` builder methods. +//! [`RmcpToolRegistration`] supports a separate model-visible name when the +//! MCP wire name needs to remain unchanged. //! //! # Example //! @@ -102,9 +104,49 @@ pub const DEFAULT_MCP_REFRESH_TIMEOUT: Duration = Duration::from_secs(30); /// has already exceeded its caller-visible deadline. const MCP_CANCELLATION_GRACE_PERIOD: Duration = Duration::from_secs(1); +/// Public registration data for an MCP tool. +/// +/// `model_name` is used for provider-facing tool definitions and tool lookup, +/// while `definition.name` remains the MCP protocol name sent in +/// `tools/call` requests. Use [`Self::new`] when the names are identical and +/// [`Self::with_timeout`] to override the default per-call timeout. +pub struct RmcpToolRegistration { + /// Name exposed to the model and used by Rig's tool registry. + pub model_name: String, + /// Server-advertised MCP tool definition, including the wire name. + pub definition: rmcp::model::Tool, + /// Connected MCP server peer that receives `tools/call` requests. + pub client: rmcp::service::ServerSink, + /// Per-call timeout. `None` disables the timeout. + pub timeout: Option, +} + +impl RmcpToolRegistration { + /// Create a registration with the default MCP tool timeout. + pub fn new( + model_name: impl Into, + definition: rmcp::model::Tool, + client: rmcp::service::ServerSink, + ) -> Self { + Self { + model_name: model_name.into(), + definition, + client, + timeout: Some(DEFAULT_MCP_TOOL_TIMEOUT), + } + } + + /// Set (or clear) the per-call timeout for this registration. + pub fn with_timeout(mut self, timeout: impl Into>) -> Self { + self.timeout = timeout.into(); + self + } +} + /// Crate-private adapter used by Rig's public MCP registration methods. #[derive(Clone)] pub(crate) struct McpTool { + model_name: String, definition: rmcp::model::Tool, client: rmcp::service::ServerSink, /// Per-call timeout. When `Some`, an MCP `call_tool` that does not complete @@ -125,10 +167,18 @@ impl McpTool { definition: rmcp::model::Tool, client: rmcp::service::ServerSink, ) -> Self { + let model_name = definition.name.to_string(); + Self::from_registration(RmcpToolRegistration::new(model_name, definition, client)) + } + + /// Create an adapter from a public registration, preserving separate + /// model-visible and MCP wire names. + pub(crate) fn from_registration(registration: RmcpToolRegistration) -> Self { Self { - definition, - client, - timeout: Some(DEFAULT_MCP_TOOL_TIMEOUT), + model_name: registration.model_name, + definition: registration.definition, + client: registration.client, + timeout: registration.timeout, } } @@ -292,22 +342,24 @@ impl McpTool { args: String, meta: Option, ) -> WasmBoxedFuture<'_, Result> { - let name = self.definition.name.clone(); + let wire_name = self.definition.name.clone(); + let model_name = self.model_name.clone(); Box::pin(async move { // Validate the JSON arguments before contacting the server: malformed // JSON must surface as an InvalidArgs failure, not a silent no-arg call. let arguments = parse_mcp_arguments(&args).map_err(|error| { ToolExecutionError::invalid_args(format!( - "MCP tool '{name}' received invalid arguments: {error}" + "MCP tool '{model_name}' received invalid arguments: {error}" )) .with_source(error) })?; let mut request = arguments .map(|arguments| { - rmcp::model::CallToolRequestParams::new(name.clone()).with_arguments(arguments) + rmcp::model::CallToolRequestParams::new(wire_name.clone()) + .with_arguments(arguments) }) - .unwrap_or_else(|| rmcp::model::CallToolRequestParams::new(name)); + .unwrap_or_else(|| rmcp::model::CallToolRequestParams::new(wire_name)); request.meta = meta; match call_mcp_tool(&self.client, request, self.timeout).await { @@ -320,14 +372,14 @@ impl McpTool { let timeout = self.timeout.unwrap_or(elapsed_timeout); Err(ToolExecutionError::timeout(format!( "MCP tool '{}' timed out after {timeout:?}", - self.definition.name + model_name )) .with_source(error)) } // A transport/service error before the tool produced a result. Err(error) => Err(ToolExecutionError::provider(format!( "MCP tool '{}' request failed: {error}", - self.definition.name + model_name )) .with_source(error)), } @@ -459,7 +511,7 @@ fn preserve_mcp_result(context: &mut ToolContext, result: &CallToolResult) { impl ErasedTool for McpTool { fn name(&self) -> String { - self.definition.name.to_string() + self.model_name.clone() } fn description(&self) -> String { @@ -498,7 +550,7 @@ impl ErasedTool for McpTool { ToolResult::failed( ToolExecutionError::other(format!( "MCP tool '{}' reported an execution error", - self.definition.name + self.model_name )) .with_model_output(output), ) @@ -832,6 +884,7 @@ mod tests { struct ScenarioServer { scenario: Scenario, seen: Arc>>, + seen_name: Arc>>, cancelled: Arc, } @@ -844,9 +897,10 @@ mod tests { async fn call_tool( &self, - _request: CallToolRequestParams, + request: CallToolRequestParams, context: RequestContext, ) -> Result { + *self.seen_name.write().await = Some(request.name.to_string()); *self.seen.write().await = Some(context.meta.clone()); match self.scenario { Scenario::Success => Ok(CallToolResult::success(vec![ContentBlock::text("ok")])), @@ -894,19 +948,27 @@ mod tests { struct Fixture { handle: ToolServerHandle, seen: Arc>>, + seen_name: Arc>>, cancelled: Arc, - _client: rmcp::service::RunningService, + client: rmcp::service::RunningService, server_task: JoinHandle<()>, } - async fn fixture(scenario: Scenario, timeout: Option) -> Fixture { + async fn fixture_with_names( + scenario: Scenario, + timeout: Option, + model_name: &str, + wire_name: &str, + ) -> Fixture { let seen = Arc::new(RwLock::new(None)); + let seen_name = Arc::new(RwLock::new(None)); let cancelled = Arc::new(Notify::new()); let (client_to_server, server_from_client) = tokio::io::duplex(8192); let (server_to_client, client_from_server) = tokio::io::duplex(8192); let server = ScenarioServer { scenario, seen: seen.clone(), + seen_name: seen_name.clone(), cancelled: cancelled.clone(), }; let server_task = tokio::spawn(async move { @@ -921,22 +983,36 @@ mod tests { .await .expect("client connect"); let definition = Tool::new( - "fixture_tool".to_string(), + wire_name.to_string(), "fixture".to_string(), Arc::new(serde_json::Map::new()), ); - let handle = ToolServer::new() - .rmcp_tool_with_timeout(definition, client.peer().clone(), timeout) - .run(); + let peer = client.peer().clone(); + let handle = if model_name == wire_name { + ToolServer::new() + .rmcp_tool_with_timeout(definition, peer, timeout) + .run() + } else { + ToolServer::new() + .rmcp_tool_registration( + RmcpToolRegistration::new(model_name, definition, peer).with_timeout(timeout), + ) + .run() + }; Fixture { handle, seen, + seen_name, cancelled, - _client: client, + client, server_task, } } + async fn fixture(scenario: Scenario, timeout: Option) -> Fixture { + fixture_with_names(scenario, timeout, "fixture_tool", "fixture_tool").await + } + async fn execute(fixture: &Fixture, args: &str, context: &mut ToolContext) -> ToolResult { tokio::time::timeout( Duration::from_secs(5), @@ -946,6 +1022,88 @@ mod tests { .expect("MCP dispatch exceeded the outer safety timeout") } + #[tokio::test] + async fn registration_uses_model_name_for_lookup_and_wire_name_for_call() { + let fixture = fixture_with_names( + Scenario::Success, + Some(Duration::from_secs(1)), + "model_fixture_tool", + "server_fixture_tool", + ) + .await; + + let definitions = fixture.handle.get_tool_defs(None).await.unwrap(); + assert_eq!(definitions.len(), 1); + assert_eq!(definitions[0].name, "model_fixture_tool"); + + let result = fixture + .handle + .execute("model_fixture_tool", "{}", &mut ToolContext::new()) + .await; + assert!(result.is_success()); + assert_eq!( + *fixture.seen_name.read().await, + Some("server_fixture_tool".to_string()) + ); + + fixture.server_task.abort(); + } + + #[tokio::test] + async fn registrations_allow_two_servers_with_the_same_wire_name() { + let first = fixture(Scenario::Success, Some(Duration::from_secs(1))).await; + let second = fixture(Scenario::Success, Some(Duration::from_secs(1))).await; + let definition = || { + Tool::new( + "fixture_tool".to_string(), + "fixture".to_string(), + Arc::new(serde_json::Map::new()), + ) + }; + + let handle = ToolServer::new() + .rmcp_tool_registration(RmcpToolRegistration::new( + "first_fixture_tool", + definition(), + first.client.peer().clone(), + )) + .rmcp_tool_registration(RmcpToolRegistration::new( + "second_fixture_tool", + definition(), + second.client.peer().clone(), + )) + .run(); + + let definitions = handle.get_tool_defs(None).await.unwrap(); + assert_eq!( + definitions + .iter() + .map(|definition| definition.name.as_str()) + .collect::>(), + vec!["first_fixture_tool", "second_fixture_tool"] + ); + + let first_result = handle + .execute("first_fixture_tool", "{}", &mut ToolContext::new()) + .await; + assert!(first_result.is_success()); + let second_result = handle + .execute("second_fixture_tool", "{}", &mut ToolContext::new()) + .await; + assert!(second_result.is_success()); + assert_eq!( + *first.seen_name.read().await, + Some("fixture_tool".to_string()) + ); + assert_eq!( + *second.seen_name.read().await, + Some("fixture_tool".to_string()) + ); + + first.server_task.abort(); + second.server_task.abort(); + } + #[tokio::test] async fn best_effort_cancellation_drops_stalled_delivery_after_grace_period() { struct DropProbe(Arc); diff --git a/crates/rig-agent/src/tool/server.rs b/crates/rig-agent/src/tool/server.rs index 1a29fe5b7f..7247d5bcca 100644 --- a/crates/rig-agent/src/tool/server.rs +++ b/crates/rig-agent/src/tool/server.rs @@ -194,15 +194,34 @@ impl ToolServer { #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] pub fn rmcp_tool_with_timeout( - mut self, + self, tool: rmcp::model::Tool, client: rmcp::service::ServerSink, timeout: impl Into>, + ) -> Self { + use crate::tool::rmcp::RmcpToolRegistration; + + let model_name = tool.name.to_string(); + self.rmcp_tool_registration( + RmcpToolRegistration::new(model_name, tool, client).with_timeout(timeout), + ) + } + + /// Add an MCP tool with a separate model-visible name. + /// + /// The registration's `model_name` is used for provider-facing tool + /// definitions and tool lookup. The original `definition.name` is kept for + /// MCP `tools/call` requests. + #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))] + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] + pub fn rmcp_tool_registration( + mut self, + registration: crate::tool::rmcp::RmcpToolRegistration, ) -> Self { use crate::tool::rmcp::McpTool; - self.toolset.add_erased(Arc::new( - McpTool::from_mcp_server(tool, client).with_timeout(timeout), - )); + + self.toolset + .add_erased(Arc::new(McpTool::from_registration(registration))); self }