fix: expose MRTR state to tool handlers (#1104)

* fix: expose MRTR state to tool handlers

* docs: document tool call context fields
This commit is contained in:
Dale Seo 2026-08-05 09:55:09 -04:00 • committed by GitHub
commit c4def55422
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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();