Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions crates/rmcp/src/handler/server/tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand All @@ -47,6 +51,8 @@ impl<'s, S> ToolCallContext<'s, S> {
meta: _,
name,
arguments,
input_responses,
request_state,
..
}: CallToolRequestParams,
request_context: RequestContext<RoleServer>,
Expand All @@ -56,6 +62,8 @@ impl<'s, S> ToolCallContext<'s, S> {
service,
name,
arguments,
input_responses,
request_state,
}
}
pub fn name(&self) -> &str {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
79 changes: 76 additions & 3 deletions crates/rmcp/tests/test_mrtr_behavior.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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<()>>,
{
Expand Down Expand Up @@ -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();
Expand Down