diff --git a/crates/rmcp/src/handler/server/tool.rs b/crates/rmcp/src/handler/server/tool.rs index a9024066..81171e2c 100644 --- a/crates/rmcp/src/handler/server/tool.rs +++ b/crates/rmcp/src/handler/server/tool.rs @@ -38,6 +38,10 @@ pub struct ToolCallContext<'s, S> { pub service: &'s S, pub name: Cow<'static, str>, pub arguments: Option, + /// Client responses to input requests from the previous MRTR round. + pub input_responses: Option, + /// Opaque state returned by the server during the previous MRTR round. + pub request_state: Option, } impl<'s, S> ToolCallContext<'s, S> { @@ -47,6 +51,8 @@ impl<'s, S> ToolCallContext<'s, S> { meta: _, name, arguments, + input_responses, + request_state, .. }: CallToolRequestParams, request_context: RequestContext, @@ -56,6 +62,8 @@ impl<'s, S> ToolCallContext<'s, S> { service, name, arguments, + input_responses, + request_state, } } pub fn name(&self) -> &str { @@ -98,6 +106,12 @@ impl IntoCallToolResult for InputRequiredResult { } } +impl IntoCallToolResult for CallToolResponse { + fn into_call_tool_result(self) -> Result { + Ok(self) + } +} + impl IntoCallToolResult for crate::ErrorData { fn into_call_tool_result(self) -> Result { Err(self) @@ -193,6 +207,26 @@ impl FromContextPart> for ToolName { } } +/// Extracts the opaque state returned by the server during the previous MRTR round. +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct RequestState(pub Option); + +impl FromContextPart> for RequestState { + fn from_context_part(context: &mut ToolCallContext) -> Result { + Ok(Self(context.request_state.take())) + } +} + +/// Extracts client responses to input requests from the previous MRTR round. +#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] +pub struct InputResponses(pub Option); + +impl FromContextPart> for InputResponses { + fn from_context_part(context: &mut ToolCallContext) -> Result { + Ok(Self(context.input_responses.take())) + } +} + // Special implementation for Parameters that handles tool arguments impl FromContextPart> for Parameters

where diff --git a/crates/rmcp/tests/test_mrtr_behavior.rs b/crates/rmcp/tests/test_mrtr_behavior.rs index 33332c2b..4fb24e18 100644 --- a/crates/rmcp/tests/test_mrtr_behavior.rs +++ b/crates/rmcp/tests/test_mrtr_behavior.rs @@ -13,9 +13,16 @@ use std::sync::{ use rmcp::{ ClientHandler, ServerHandler, + handler::server::{ + tool::{InputResponses as ToolInputResponses, RequestState}, + wrapper::Parameters, + }, model::*, - service::{RequestContext, RoleClient, RoleServer, ServiceError, serve_directly}, + service::{RequestContext, RoleClient, RoleServer, Service, ServiceError, serve_directly}, + tool, tool_handler, tool_router, }; +use schemars::JsonSchema; +use serde::Deserialize; use serde_json::json; /// A `requestState` value with characters that must survive a byte-exact echo: @@ -66,6 +73,54 @@ fn single_elicitation(state: &str) -> InputRequiredResult { InputRequiredResult::new(Some(requests), Some(state.into())) } +#[derive(Clone)] +struct MacroMrtrServer; + +#[derive(Deserialize, JsonSchema)] +struct MacroMrtrArguments { + greeting: String, +} + +#[tool_router] +impl MacroMrtrServer { + #[tool(description = "Greet a user after collecting their name")] + async fn greet( + &self, + Parameters(arguments): Parameters, + RequestState(request_state): RequestState, + ToolInputResponses(input_responses): ToolInputResponses, + ) -> Result { + match request_state.as_deref() { + None => Ok(single_elicitation("macro-state").into()), + Some("macro-state") => { + let name = input_responses + .as_ref() + .and_then(|responses| responses.get("answer")) + .and_then(|response| response["content"]["name"].as_str()) + .ok_or_else(|| ErrorData::invalid_params("missing name response", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text(format!( + "{}, {name}", + arguments.greeting + ))]) + .into()) + } + Some(other) => Err(ErrorData::invalid_params( + format!("unexpected request state {other:?}"), + None, + )), + } + } +} + +#[tool_handler] +impl ServerHandler for MacroMrtrServer { + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build()); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } +} + impl MrtrServer { fn call_tool_impl( &self, @@ -289,12 +344,13 @@ fn server_info(protocol_version: ProtocolVersion) -> ServerInfo { /// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local` /// feature is active) is available, wiring up a connected client/server pair. -async fn with_pair( - server: MrtrServer, +async fn with_pair( + server: S, client_protocol: ProtocolVersion, body: F, ) -> anyhow::Result<()> where + S: Service, F: FnOnce(rmcp::service::RunningService) -> Fut, Fut: std::future::Future>, { @@ -346,6 +402,23 @@ async fn client_auto_fulfills_input_required_tool_call() -> anyhow::Result<()> { .await } +#[tokio::test(flavor = "current_thread")] +async fn tool_macro_receives_mrtr_retry_fields() -> anyhow::Result<()> { + with_pair( + MacroMrtrServer, + ProtocolVersion::V_2026_07_28, + |client| async move { + let arguments = serde_json::from_value(json!({ "greeting": "hello" })).unwrap(); + let result = client + .call_tool(CallToolRequestParams::new("greet").with_arguments(arguments)) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "hello, Ferris"); + Ok(()) + }, + ) + .await +} + #[tokio::test(flavor = "current_thread")] async fn manual_once_returns_input_required_without_retry() -> anyhow::Result<()> { let server = MrtrServer::default();