FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: expose MRTR state to tool handlers · modelcontextprotocol/rust-sdk@aff4f5c · GitHub

/ rust-sdk Public

Commit aff4f5c

Browse files
DaleSeo
committed
fix: expose MRTR state to tool handlers
1 parent 1f9358e commit aff4f5c

2 files changed

Lines changed: 110 additions & 3 deletions

File tree

‎crates/rmcp/src/handler/server/tool.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ pub struct ToolCallContext<'s, S> {
3838
pub service: &'s S,
3939
pub name: Cow<'static, str>,
4040
pub arguments: Option<JsonObject>,
41+
/// Client responses to input requests from the previous MRTR round.
42+
pub input_responses: Option<crate::model::InputResponses>,
43+
/// Opaque state returned by the server during the previous MRTR round.
44+
pub request_state: Option<String>,
4145
}
4246

4347
impl<'s, S> ToolCallContext<'s, S> {
@@ -47,6 +51,8 @@ impl<'s, S> ToolCallContext<'s, S> {
4751
meta: _,
4852
name,
4953
arguments,
54+
input_responses,
55+
request_state,
5056
..
5157
}: CallToolRequestParams,
5258
request_context: RequestContext<RoleServer>,
@@ -56,6 +62,8 @@ impl<'s, S> ToolCallContext<'s, S> {
5662
service,
5763
name,
5864
arguments,
65+
input_responses,
66+
request_state,
5967
}
6068
}
6169
pub fn name(&self) -> &str {
@@ -98,6 +106,12 @@ impl IntoCallToolResult for InputRequiredResult {
98106
}
99107
}
100108

109+
impl IntoCallToolResult for CallToolResponse {
110+
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
111+
Ok(self)
112+
}
113+
}
114+
101115
impl IntoCallToolResult for crate::ErrorData {
102116
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
103117
Err(self)
@@ -193,6 +207,26 @@ impl<S> FromContextPart<ToolCallContext<'_, S>> for ToolName {
193207
}
194208
}
195209

210+
/// Extracts the opaque state returned by the server during the previous MRTR round.
211+
#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
212+
pub struct RequestState(pub Option<String>);
213+
214+
impl<S> FromContextPart<ToolCallContext<'_, S>> for RequestState {
215+
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
216+
Ok(Self(context.request_state.take()))
217+
}
218+
}
219+
220+
/// Extracts client responses to input requests from the previous MRTR round.
221+
#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
222+
pub struct InputResponses(pub Option<crate::model::InputResponses>);
223+
224+
impl<S> FromContextPart<ToolCallContext<'_, S>> for InputResponses {
225+
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
226+
Ok(Self(context.input_responses.take()))
227+
}
228+
}
229+
196230
// Special implementation for Parameters that handles tool arguments
197231
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
198232
where

‎crates/rmcp/tests/test_mrtr_behavior.rs‎

Lines changed: 76 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,16 @@ use std::sync::{
1313

1414
use rmcp::{
1515
ClientHandler, ServerHandler,
16+
handler::server::{
17+
tool::{InputResponses as ToolInputResponses, RequestState},
18+
wrapper::Parameters,
19+
},
1620
model::*,
17-
service::{RequestContext, RoleClient, RoleServer, ServiceError, serve_directly},
21+
service::{RequestContext, RoleClient, RoleServer, Service, ServiceError, serve_directly},
22+
tool, tool_handler, tool_router,
1823
};
24+
use schemars::JsonSchema;
25+
use serde::Deserialize;
1926
use serde_json::json;
2027

2128
/// A `requestState` value with characters that must survive a byte-exact echo:
@@ -66,6 +73,54 @@ fn single_elicitation(state: &str) -> InputRequiredResult {
6673
InputRequiredResult::new(Some(requests), Some(state.into()))
6774
}
6875

76+
#[derive(Clone)]
77+
struct MacroMrtrServer;
78+
79+
#[derive(Deserialize, JsonSchema)]
80+
struct MacroMrtrArguments {
81+
greeting: String,
82+
}
83+
84+
#[tool_router]
85+
impl MacroMrtrServer {
86+
#[tool(description = "Greet a user after collecting their name")]
87+
async fn greet(
88+
&self,
89+
Parameters(arguments): Parameters<MacroMrtrArguments>,
90+
RequestState(request_state): RequestState,
91+
ToolInputResponses(input_responses): ToolInputResponses,
92+
) -> Result<CallToolResponse, ErrorData> {
93+
match request_state.as_deref() {
94+
None => Ok(single_elicitation("macro-state").into()),
95+
Some("macro-state") => {
96+
let name = input_responses
97+
.as_ref()
98+
.and_then(|responses| responses.get("answer"))
99+
.and_then(|response| response["content"]["name"].as_str())
100+
.ok_or_else(|| ErrorData::invalid_params("missing name response", None))?;
101+
Ok(CallToolResult::success(vec![ContentBlock::text(format!(
102+
"{}, {name}",
103+
arguments.greeting
104+
))])
105+
.into())
106+
}
107+
Some(other) => Err(ErrorData::invalid_params(
108+
format!("unexpected request state {other:?}"),
109+
None,
110+
)),
111+
}
112+
}
113+
}
114+
115+
#[tool_handler]
116+
impl ServerHandler for MacroMrtrServer {
117+
fn get_info(&self) -> ServerInfo {
118+
let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build());
119+
info.protocol_version = ProtocolVersion::V_2026_07_28;
120+
info
121+
}
122+
}
123+
69124
impl MrtrServer {
70125
fn call_tool_impl(
71126
&self,
@@ -289,12 +344,13 @@ fn server_info(protocol_version: ProtocolVersion) -> ServerInfo {
289344

290345
/// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local`
291346
/// feature is active) is available, wiring up a connected client/server pair.
292-
async fn with_pair<F, Fut>(
293-
server: MrtrServer,
347+
async fn with_pair<S, F, Fut>(
348+
server: S,
294349
client_protocol: ProtocolVersion,
295350
body: F,
296351
) -> anyhow::Result<()>
297352
where
353+
S: Service<RoleServer>,
298354
F: FnOnce(rmcp::service::RunningService<RoleClient, MrtrClient>) -> Fut,
299355
Fut: std::future::Future<Output = anyhow::Result<()>>,
300356
{
@@ -346,6 +402,23 @@ async fn client_auto_fulfills_input_required_tool_call() -> anyhow::Result<()> {
346402
.await
347403
}
348404

405+
#[tokio::test(flavor = "current_thread")]
406+
async fn tool_macro_receives_mrtr_retry_fields() -> anyhow::Result<()> {
407+
with_pair(
408+
MacroMrtrServer,
409+
ProtocolVersion::V_2026_07_28,
410+
|client| async move {
411+
let arguments = serde_json::from_value(json!({ "greeting": "hello" })).unwrap();
412+
let result = client
413+
.call_tool(CallToolRequestParams::new("greet").with_arguments(arguments))
414+
.await?;
415+
assert_eq!(result.content[0].as_text().unwrap().text, "hello, Ferris");
416+
Ok(())
417+
},
418+
)
419+
.await
420+
}
421+
349422
#[tokio::test(flavor = "current_thread")]
350423
async fn manual_once_returns_input_required_without_retry() -> anyhow::Result<()> {
351424
let server = MrtrServer::default();

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL