fix: expose MRTR state to tool handlers

Author: DaleSeo

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<JsonObject>,
+    /// Client responses to input requests from the previous MRTR round.
+    pub input_responses: Option<crate::model::InputResponses>,
+    /// Opaque state returned by the server during the previous MRTR round.
+    pub request_state: Option<String>,
 }
 
 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<RoleServer>,
@@ -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<CallToolResponse, crate::ErrorData> {
+        Ok(self)
+    }
+}
+
 impl IntoCallToolResult for crate::ErrorData {
     fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
         Err(self)
@@ -193,6 +207,26 @@ impl<S> FromContextPart<ToolCallContext<'_, S>> 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<String>);
+
+impl<S> FromContextPart<ToolCallContext<'_, S>> for RequestState {
+    fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
+        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<crate::model::InputResponses>);
+
+impl<S> FromContextPart<ToolCallContext<'_, S>> for InputResponses {
+    fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
+        Ok(Self(context.input_responses.take()))
+    }
+}
+
 // Special implementation for Parameters that handles tool arguments
 impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
 where

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<MacroMrtrArguments>,
+        RequestState(request_state): RequestState,
+        ToolInputResponses(input_responses): ToolInputResponses,
+    ) -> Result<CallToolResponse, ErrorData> {
+        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<F, Fut>(
-    server: MrtrServer,
+async fn with_pair<S, F, Fut>(
+    server: S,
     client_protocol: ProtocolVersion,
     body: F,
 ) -> anyhow::Result<()>
 where
+    S: Service<RoleServer>,
     F: FnOnce(rmcp::service::RunningService<RoleClient, MrtrClient>) -> Fut,
     Fut: std::future::Future<Output = anyhow::Result<()>>,
 {
@@ -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();