feat!: add MRTR behavior support (SEP-2322) (#929)

* feat!: add MRTR behavior support

* feat: harden SEP-2322 MRTR support

* ci: diff public API only on features common to base and head

cargo public-api builds both revisions with the same feature set, so a
feature introduced (or removed) by a PR broke the build of the other
revision. Restrict the all-features diff to features present in both the
base and head revisions.
This commit is contained in:
Dale Seo 2026-07-10 12:02:26 -04:00 • committed by GitHub
commit dd30a70f84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 2121 additions and 127 deletions

View file

@ -27,6 +27,9 @@ impl<H: ServerHandler> Service<RoleServer> for H {
) -> Result<<RoleServer as ServiceRole>::Resp, McpError> {
// `context` is moved into the dispatch below, so read the negotiated version first.
let protocol_version = context.protocol_version();
let mrtr_supported = protocol_version
.as_ref()
.is_some_and(|v| v.as_str() >= ProtocolVersion::V_2026_07_28.as_str());
let result = match request {
ClientRequest::InitializeRequest(request) => self
.initialize(request.params, context)
@ -46,7 +49,7 @@ impl<H: ServerHandler> Service<RoleServer> for H {
ClientRequest::GetPromptRequest(request) => self
.get_prompt(request.params, context)
.await
.map(ServerResult::GetPromptResult),
.map(ServerResult::from),
ClientRequest::ListPromptsRequest(request) => self
.list_prompts(request.params, context)
.await
@ -62,7 +65,7 @@ impl<H: ServerHandler> Service<RoleServer> for H {
ClientRequest::ReadResourceRequest(request) => self
.read_resource(request.params, context)
.await
.map(ServerResult::ReadResourceResult),
.map(ServerResult::from),
ClientRequest::SubscribeRequest(request) => self
.subscribe(request.params, context)
.await
@ -105,7 +108,7 @@ impl<H: ServerHandler> Service<RoleServer> for H {
} else {
self.call_tool(request.params, context)
.await
.map(ServerResult::CallToolResult)
.map(ServerResult::from)
}
}
ClientRequest::ListToolsRequest(request) => self
@ -133,6 +136,17 @@ impl<H: ServerHandler> Service<RoleServer> for H {
.await
.map(ServerResult::CancelTaskResult),
};
let result = result.and_then(|result| {
if matches!(result, ServerResult::InputRequiredResult(_)) && !mrtr_supported {
Err(McpError::invalid_request(
"InputRequiredResult requires negotiated protocol version 2026-07-28 or newer",
None,
))
} else {
Ok(result)
}
});
// SEP-2164: peers negotiating 2026-07-28+ get the standard INVALID_PARAMS code for
// resource-not-found; older peers keep RESOURCE_NOT_FOUND. ISO `YYYY-MM-DD` versions
// compare lexically the same as chronologically.
@ -229,7 +243,7 @@ macro_rules! server_handler_methods {
&self,
request: GetPromptRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetPromptResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<GetPromptResponse, McpError>> + MaybeSendFuture + '_ {
std::future::ready(Err(McpError::method_not_found::<GetPromptRequestMethod>()))
}
fn list_prompts(
@ -259,7 +273,7 @@ macro_rules! server_handler_methods {
&self,
request: ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<ReadResourceResponse, McpError>> + MaybeSendFuture + '_ {
std::future::ready(Err(
McpError::method_not_found::<ReadResourceRequestMethod>(),
))
@ -312,7 +326,7 @@ macro_rules! server_handler_methods {
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CallToolResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<CallToolResponse, McpError>> + MaybeSendFuture + '_ {
std::future::ready(Err(McpError::method_not_found::<CallToolRequestMethod>()))
}
fn list_tools(
@ -485,7 +499,7 @@ macro_rules! impl_server_handler_for_wrapper {
&self,
request: GetPromptRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetPromptResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<GetPromptResponse, McpError>> + MaybeSendFuture + '_ {
(**self).get_prompt(request, context)
}
@ -518,7 +532,7 @@ macro_rules! impl_server_handler_for_wrapper {
&self,
request: ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<ReadResourceResponse, McpError>> + MaybeSendFuture + '_ {
(**self).read_resource(request, context)
}
@ -542,7 +556,7 @@ macro_rules! impl_server_handler_for_wrapper {
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CallToolResult, McpError>> + MaybeSendFuture + '_ {
) -> impl Future<Output = Result<CallToolResponse, McpError>> + MaybeSendFuture + '_ {
(**self).call_tool(request, context)
}

View file

@ -15,7 +15,7 @@ pub use super::common::{Extension, RequestId};
use crate::{
RoleServer,
handler::server::wrapper::Parameters,
model::{GetPromptResult, PromptMessage},
model::{GetPromptResponse, GetPromptResult, InputRequiredResult, PromptMessage},
service::{MaybeBoxFuture, MaybeSend, MaybeSendFuture, RequestContext},
};
@ -59,12 +59,12 @@ pub trait GetPromptHandler<S, A> {
fn handle(
self,
context: PromptContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>;
) -> MaybeBoxFuture<'_, Result<GetPromptResponse, crate::ErrorData>>;
}
/// Type alias for dynamic prompt handlers
#[cfg(not(feature = "local"))]
pub type DynGetPromptHandler<S> = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
pub type DynGetPromptHandler<S> = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result<GetPromptResponse, crate::ErrorData>>
+ Send
+ Sync;
@ -73,7 +73,7 @@ pub type DynGetPromptHandler<S> = dyn for<'a> Fn(
PromptContext<'a, S>,
) -> futures::future::LocalBoxFuture<
'a,
Result<GetPromptResult, crate::ErrorData>,
Result<GetPromptResponse, crate::ErrorData>,
>;
/// Adapter type for async methods that return `Vec<PromptMessage>`
@ -91,28 +91,35 @@ pub struct SyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
/// Trait for types that can be converted into GetPromptResult
pub trait IntoGetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData>;
fn into_get_prompt_result(self) -> Result<GetPromptResponse, crate::ErrorData>;
}
impl IntoGetPromptResult for GetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
Ok(self)
fn into_get_prompt_result(self) -> Result<GetPromptResponse, crate::ErrorData> {
Ok(self.into())
}
}
impl IntoGetPromptResult for InputRequiredResult {
fn into_get_prompt_result(self) -> Result<GetPromptResponse, crate::ErrorData> {
Ok(self.into())
}
}
impl IntoGetPromptResult for Vec<PromptMessage> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
fn into_get_prompt_result(self) -> Result<GetPromptResponse, crate::ErrorData> {
Ok(GetPromptResult {
result_type: Default::default(),
description: None,
messages: self,
meta: None,
})
}
.into())
}
}
impl<T: IntoGetPromptResult> IntoGetPromptResult for Result<T, crate::ErrorData> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
fn into_get_prompt_result(self) -> Result<GetPromptResponse, crate::ErrorData> {
self.and_then(|v| v.into_get_prompt_result())
}
}
@ -129,7 +136,7 @@ pin_project_lite::pin_project! {
},
Ready {
#[pin]
result: futures::future::Ready<Result<GetPromptResult, crate::ErrorData>>,
result: futures::future::Ready<Result<GetPromptResponse, crate::ErrorData>>,
}
}
}
@ -139,7 +146,7 @@ where
F: Future<Output = R>,
R: IntoGetPromptResult,
{
type Output = Result<GetPromptResult, crate::ErrorData>;
type Output = Result<GetPromptResponse, crate::ErrorData>;
fn poll(
self: std::pin::Pin<&mut Self>,
@ -216,7 +223,7 @@ macro_rules! impl_prompt_handler_for {
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
) -> MaybeBoxFuture<'_, Result<GetPromptResponse, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
@ -249,7 +256,7 @@ macro_rules! impl_prompt_handler_for {
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
) -> MaybeBoxFuture<'_, Result<GetPromptResponse, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
@ -280,7 +287,7 @@ macro_rules! impl_prompt_handler_for {
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
) -> MaybeBoxFuture<'_, Result<GetPromptResponse, crate::ErrorData>>
{
// Extract all parameters before moving into the async block
$(
@ -315,7 +322,7 @@ macro_rules! impl_prompt_handler_for {
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
) -> MaybeBoxFuture<'_, Result<GetPromptResponse, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);

View file

@ -106,7 +106,7 @@ where
context,
);
let result = self.tool_router.call(tool_call_context).await?;
Ok(ServerResult::CallToolResult(result))
Ok(ServerResult::from(result))
} else {
self.service
.handle_request(ClientRequest::CallToolRequest(request), context)
@ -129,7 +129,7 @@ where
context,
);
let result = self.prompt_router.get_prompt(prompt_context).await?;
Ok(ServerResult::GetPromptResult(result))
Ok(ServerResult::from(result))
} else {
self.service
.handle_request(ClientRequest::GetPromptRequest(request), context)
@ -193,7 +193,7 @@ mod tests {
async fn test_router_deferred_notifier_e2e() {
let mut router = Router::new(DummyHandler).with_tool(tool::ToolRoute::new_dyn(
Tool::new("my_tool", "test", Arc::new(Default::default())),
|_ctx| Box::pin(async { Ok(CallToolResult::default()) }),
|_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }),
));
let id_provider: Arc<dyn RequestIdProvider> =

View file

@ -2,7 +2,7 @@ use std::{borrow::Cow, sync::Arc};
use crate::{
handler::server::prompt::{DynGetPromptHandler, GetPromptHandler, PromptContext},
model::{GetPromptResult, Prompt},
model::{GetPromptResponse, Prompt},
service::{MaybeBoxFuture, MaybeSend},
};
@ -50,7 +50,8 @@ impl<S: MaybeSend + 'static> PromptRoute<S> {
where
H: for<'a> Fn(
PromptContext<'a, S>,
) -> MaybeBoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
)
-> MaybeBoxFuture<'a, Result<GetPromptResponse, crate::ErrorData>>
+ MaybeSend
+ 'static,
{
@ -175,7 +176,7 @@ where
pub async fn get_prompt(
&self,
context: PromptContext<'_, S>,
) -> Result<GetPromptResult, crate::ErrorData> {
) -> Result<GetPromptResponse, crate::ErrorData> {
let item = self.map.get(context.name.as_str()).ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("prompt '{}' not found", context.name),

View file

@ -137,21 +137,19 @@ use crate::{
tool::{CallToolHandler, DynCallToolHandler, ToolCallContext},
tool_name_validation::validate_and_warn_tool_name,
},
model::{CallToolResult, ContentBlock, ErrorCode, Tool, ToolAnnotations},
model::{CallToolResponse, CallToolResult, ContentBlock, ErrorCode, Tool, ToolAnnotations},
service::{MaybeBoxFuture, MaybeSend},
};
const TOOL_ARGUMENT_DESERIALIZATION_ERROR_PREFIX: &str = "failed to deserialize parameters:";
fn into_tool_argument_error(error: crate::ErrorData) -> Result<CallToolResult, crate::ErrorData> {
fn into_tool_argument_error(error: crate::ErrorData) -> Result<CallToolResponse, crate::ErrorData> {
if error.code == ErrorCode::INVALID_PARAMS
&& error
.message
.starts_with(TOOL_ARGUMENT_DESERIALIZATION_ERROR_PREFIX)
{
return Ok(CallToolResult::error(vec![ContentBlock::text(
error.message,
)]));
return Ok(CallToolResult::error(vec![ContentBlock::text(error.message)]).into());
}
Err(error)
@ -200,7 +198,8 @@ impl<S: MaybeSend + 'static> ToolRoute<S> {
where
C: for<'a> Fn(
ToolCallContext<'a, S>,
) -> MaybeBoxFuture<'a, Result<CallToolResult, crate::ErrorData>>
)
-> MaybeBoxFuture<'a, Result<CallToolResponse, crate::ErrorData>>
+ MaybeSend
+ 'static,
{
@ -561,7 +560,7 @@ where
pub async fn call(
&self,
context: ToolCallContext<'_, S>,
) -> Result<CallToolResult, crate::ErrorData> {
) -> Result<crate::model::CallToolResponse, crate::ErrorData> {
let name = context.name();
if self.disabled.contains(name) {
return Err(crate::ErrorData::invalid_params("tool not found", None));
@ -679,6 +678,9 @@ mod tests {
.call(ctx)
.await
.expect("argument validation should be a tool result");
let CallToolResponse::Complete(result) = result else {
panic!("expected complete CallToolResult");
};
assert_eq!(result.is_error, Some(true));
let text = result
@ -696,7 +698,7 @@ mod tests {
let service = DummyService;
let mut router = ToolRouter::new().with_route(ToolRoute::new_dyn(
crate::model::Tool::new("test_tool", "a test tool", Arc::new(Default::default())),
|_ctx| Box::pin(async { Ok(CallToolResult::default()) }),
|_ctx| Box::pin(async { Ok(CallToolResult::default().into()) }),
));
router.disable_route("test_tool");

View file

@ -16,7 +16,10 @@ pub use super::{
use crate::{
RoleServer,
handler::server::wrapper::Parameters,
model::{CallToolRequestParams, CallToolResult, IntoContents, JsonObject},
model::{
CallToolRequestParams, CallToolResponse, CallToolResult, InputRequiredResult, IntoContents,
JsonObject,
},
service::{MaybeBoxFuture, MaybeSend, MaybeSendFuture, RequestContext},
};
@ -77,36 +80,46 @@ impl<S> AsRequestContext for ToolCallContext<'_, S> {
}
pub trait IntoCallToolResult {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData>;
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData>;
}
impl<T: IntoContents> IntoCallToolResult for T {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
Ok(CallToolResult::success(self.into_contents()))
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
Ok(CallToolResult::success(self.into_contents()).into())
}
}
impl IntoCallToolResult for CallToolResult {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
Ok(self)
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
Ok(self.into())
}
}
impl IntoCallToolResult for InputRequiredResult {
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
Ok(self.into())
}
}
impl IntoCallToolResult for crate::ErrorData {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
Err(self)
}
}
impl<T: IntoCallToolResult, E: IntoCallToolResult> IntoCallToolResult for Result<T, E> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
match self {
Ok(value) => value.into_call_tool_result(),
Err(error) => match error.into_call_tool_result() {
Ok(mut result) => {
Ok(CallToolResponse::Complete(mut result)) => {
result.is_error = Some(true);
Ok(result)
Ok(result.into())
}
Ok(CallToolResponse::InputRequired(_)) => Err(crate::ErrorData::internal_error(
"InputRequiredResult cannot be returned from a tool error branch",
None,
)),
Err(e) => Err(e),
},
}
@ -124,7 +137,7 @@ pin_project_lite::pin_project! {
},
Ready {
#[pin]
result: Ready<Result<CallToolResult, crate::ErrorData>>,
result: Ready<Result<CallToolResponse, crate::ErrorData>>,
}
}
}
@ -134,7 +147,7 @@ where
F: Future<Output = R>,
R: IntoCallToolResult,
{
type Output = Result<CallToolResult, crate::ErrorData>;
type Output = Result<CallToolResponse, crate::ErrorData>;
fn poll(
self: std::pin::Pin<&mut Self>,
@ -153,20 +166,21 @@ pub trait CallToolHandler<S, A> {
fn call(
self,
context: ToolCallContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<CallToolResult, crate::ErrorData>>;
) -> MaybeBoxFuture<'_, Result<CallToolResponse, crate::ErrorData>>;
}
#[cfg(not(feature = "local"))]
pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result<CallToolResponse, crate::ErrorData>>
+ Send
+ Sync;
#[cfg(feature = "local")]
pub type DynCallToolHandler<S> =
dyn for<'s> Fn(
ToolCallContext<'s, S>,
)
-> futures::future::LocalBoxFuture<'s, Result<CallToolResult, crate::ErrorData>>;
pub type DynCallToolHandler<S> = dyn for<'s> Fn(
ToolCallContext<'s, S>,
) -> futures::future::LocalBoxFuture<
's,
Result<CallToolResponse, crate::ErrorData>,
>;
// Tool-specific extractor for tool name
#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
@ -205,7 +219,10 @@ impl<S> FromContextPart<ToolCallContext<'_, S>> for JsonObject {
}
impl<'s, S> ToolCallContext<'s, S> {
pub fn invoke<H, A>(self, h: H) -> MaybeBoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
pub fn invoke<H, A>(
self,
h: H,
) -> MaybeBoxFuture<'s, Result<CallToolResponse, crate::ErrorData>>
where
H: CallToolHandler<S, A>,
{
@ -248,7 +265,7 @@ macro_rules! impl_for {
fn call(
self,
mut context: ToolCallContext<'_, S>,
) -> MaybeBoxFuture<'_, Result<CallToolResult, crate::ErrorData>>{
) -> MaybeBoxFuture<'_, Result<CallToolResponse, crate::ErrorData>>{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
@ -279,7 +296,7 @@ macro_rules! impl_for {
fn call(
self,
mut context: ToolCallContext<S>,
) -> MaybeBoxFuture<'static, Result<CallToolResult, crate::ErrorData>>{
) -> MaybeBoxFuture<'static, Result<CallToolResponse, crate::ErrorData>>{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
@ -308,7 +325,7 @@ macro_rules! impl_for {
fn call(
self,
mut context: ToolCallContext<S>,
) -> MaybeBoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
) -> MaybeBoxFuture<'static, Result<CallToolResponse, crate::ErrorData>> {
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
@ -333,7 +350,7 @@ macro_rules! impl_for {
fn call(
self,
mut context: ToolCallContext<S>,
) -> MaybeBoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
) -> MaybeBoxFuture<'static, Result<CallToolResponse, crate::ErrorData>> {
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {

View file

@ -3,7 +3,10 @@ use std::borrow::Cow;
use schemars::JsonSchema;
use serde::Serialize;
use crate::{handler::server::tool::IntoCallToolResult, model::CallToolResult};
use crate::{
handler::server::tool::IntoCallToolResult,
model::{CallToolResponse, CallToolResult},
};
/// Json wrapper for structured output
///
@ -27,7 +30,7 @@ impl<T: JsonSchema> JsonSchema for Json<T> {
// Implementation for Json<T> to create structured content
impl<T: Serialize + JsonSchema + 'static> IntoCallToolResult for Json<T> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
let value = serde_json::to_value(self.0).map_err(|e| {
crate::ErrorData::internal_error(
format!("Failed to serialize structured content: {}", e),
@ -35,6 +38,6 @@ impl<T: Serialize + JsonSchema + 'static> IntoCallToolResult for Json<T> {
)
})?;
Ok(CallToolResult::structured(value))
Ok(CallToolResult::structured(value).into())
}
}

View file

@ -14,6 +14,8 @@ mod extension;
mod meta;
mod mrtr;
mod prompt;
#[cfg(feature = "request-state")]
mod request_state;
mod resource;
mod serde_impl;
mod task;
@ -26,6 +28,8 @@ pub use extension::*;
pub use meta::*;
pub use mrtr::*;
pub use prompt::*;
#[cfg(feature = "request-state")]
pub use request_state::*;
pub use resource::*;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;

View file

@ -10,13 +10,47 @@
//! [`InputRequiredResult`] instead of the normal result. The client fulfills the
//! [`InputRequests`], then retries the original request with [`InputResponses`] and
//! the echoed `requestState`.
//!
//! # Using MRTR
//!
//! **Server:** return an [`InputRequiredResult`] from a tool/prompt/resource
//! handler via the matching outcome enum ([`CallToolResponse`],
//! [`GetPromptResponse`], [`ReadResourceResponse`]). The SDK only lets an
//! `InputRequiredResult` reach a peer that negotiated protocol version
//! `2026-07-28` or newer; older peers get a protocol error instead.
//!
//! **Client:** the high-level `RunningService` helpers — `call_tool`,
//! `get_prompt`, and `read_resource` — automatically fulfil each
//! [`InputRequest`] through the local `ClientHandler` and retry, up to
//! [`DEFAULT_MRTR_MAX_ROUNDS`]. Use the `*_once` variants (e.g.
//! `call_tool_once`) to receive an [`InputRequiredResult`] directly and drive
//! the rounds yourself.
//!
//! # `requestState` is untrusted
//!
//! The client echoes `requestState` back verbatim, so a stateless server that
//! stores meaningful data in it MUST verify integrity before trusting the echoed
//! value. Enable the `request-state` feature and use `RequestStateCodec` to seal
//! and open it, or keep the state server-side and use `requestState` only as an
//! opaque handle.
//!
//! A complete runnable walkthrough lives in the `servers_mrtr` example.
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::{CreateMessageRequest, ElicitRequest, ListRootsRequest, Meta, ResultType};
use super::{
CallToolResult, CreateMessageRequest, ElicitRequest, GetPromptResult, ListRootsRequest, Meta,
ReadResourceResult, ResultType, ServerResult,
};
/// Default maximum number of MRTR rounds a high-level client call will drive.
///
/// This matches the default used by other Tier 1 SDKs and prevents a
/// misbehaving peer from keeping a request alive indefinitely.
pub const DEFAULT_MRTR_MAX_ROUNDS: usize = 10;
/// A server-initiated request that can appear inside [`InputRequests`].
///
@ -53,6 +87,101 @@ pub type InputRequests = BTreeMap<String, InputRequest>;
/// for use as a `BTreeMap` value.
pub type InputResponses = BTreeMap<String, Value>;
/// Result of a `tools/call` request, including the MRTR intermediate result.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum CallToolResponse {
/// The server completed the tool call.
Complete(CallToolResult),
/// The server requires client-side input before the tool call can complete.
InputRequired(InputRequiredResult),
}
impl From<CallToolResult> for CallToolResponse {
fn from(result: CallToolResult) -> Self {
Self::Complete(result)
}
}
impl From<InputRequiredResult> for CallToolResponse {
fn from(result: InputRequiredResult) -> Self {
Self::InputRequired(result)
}
}
impl From<CallToolResponse> for ServerResult {
fn from(response: CallToolResponse) -> Self {
match response {
CallToolResponse::Complete(result) => ServerResult::CallToolResult(result),
CallToolResponse::InputRequired(result) => ServerResult::InputRequiredResult(result),
}
}
}
/// Result of a `prompts/get` request, including the MRTR intermediate result.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum GetPromptResponse {
/// The server completed the prompt request.
Complete(GetPromptResult),
/// The server requires client-side input before the prompt can be returned.
InputRequired(InputRequiredResult),
}
impl From<GetPromptResult> for GetPromptResponse {
fn from(result: GetPromptResult) -> Self {
Self::Complete(result)
}
}
impl From<InputRequiredResult> for GetPromptResponse {
fn from(result: InputRequiredResult) -> Self {
Self::InputRequired(result)
}
}
impl From<GetPromptResponse> for ServerResult {
fn from(response: GetPromptResponse) -> Self {
match response {
GetPromptResponse::Complete(result) => ServerResult::GetPromptResult(result),
GetPromptResponse::InputRequired(result) => ServerResult::InputRequiredResult(result),
}
}
}
/// Result of a `resources/read` request, including the MRTR intermediate result.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ReadResourceResponse {
/// The server completed the resource read.
Complete(ReadResourceResult),
/// The server requires client-side input before the resource can be returned.
InputRequired(InputRequiredResult),
}
impl From<ReadResourceResult> for ReadResourceResponse {
fn from(result: ReadResourceResult) -> Self {
Self::Complete(result)
}
}
impl From<InputRequiredResult> for ReadResourceResponse {
fn from(result: InputRequiredResult) -> Self {
Self::InputRequired(result)
}
}
impl From<ReadResourceResponse> for ServerResult {
fn from(response: ReadResourceResponse) -> Self {
match response {
ReadResourceResponse::Complete(result) => ServerResult::ReadResourceResult(result),
ReadResourceResponse::InputRequired(result) => {
ServerResult::InputRequiredResult(result)
}
}
}
}
/// A result indicating that additional input is needed before the request
/// can be completed.
///

View file

@ -0,0 +1,577 @@
//! Integrity protection for SEP-2322 `requestState`.
//!
//! In the multi round-trip request (MRTR) flow, a server places an opaque
//! `requestState` string in an [`InputRequiredResult`](super::InputRequiredResult)
//! and the client echoes it back verbatim on retry. From the server's point of
//! view the echoed value is **untrusted, attacker-controlled input**: a client
//! can send back anything it likes. Per SEP-2322, a server that lets
//! `requestState` influence authorization, resource access, or business logic
//! MUST protect its integrity and reject values that fail verification.
//!
//! [`RequestStateCodec`] provides an opt-in way to do this. It seals a payload
//! into an opaque string with an HMAC-SHA256 tag and opens it again, rejecting
//! any value that was forged or tampered with.
//!
//! To follow the spec's replay-prevention guidance without hand-rolling the
//! checks, the codec supports two bindings via [`SealOptions`]:
//!
//! * **Associated data** — arbitrary context (e.g. the authenticated principal
//! plus a digest of the originating request) that is mixed into the tag but
//! not stored in the token. [`open_with`](RequestStateCodec::open_with) only
//! succeeds when the caller supplies the same context, so a value cannot be
//! replayed by a different principal or against a different request. This is
//! *fail-closed*: forgetting to pass the context makes verification fail.
//! * **TTL** — a relative expiry stamped into the token; opening a value past
//! its expiry fails with [`RequestStateError::Expired`].
//!
//! Single-use/nonce enforcement (for one-time redemptions) still has to be done
//! server-side, as the spec notes.
//!
//! This helper is only about *integrity*, not *confidentiality*: the sealed
//! payload is signed, not encrypted, so it is base64url-readable by anyone. Do
//! not put secrets in it.
//!
//! Using the codec is entirely optional. A server that keeps its state
//! server-side, or that does not trust `requestState` for anything security
//! sensitive, can keep building the string by hand via
//! [`InputRequiredResult::from_request_state`](super::InputRequiredResult::from_request_state).
//!
//! # Examples
//!
//! ```
//! use rmcp::model::{RequestStateCodec, SealOptions};
//!
//! // Derive the key from a per-process secret; keep it out of client reach.
//! let codec = RequestStateCodec::new(b"a-32-byte-or-longer-secret-key!!!");
//!
//! // Bind the state to the caller and the originating request.
//! let context = b"user:alice|tools/call:weather";
//! let sealed = codec.seal_with(
//! b"step=2",
//! &SealOptions::new().associated_data(context),
//! );
//!
//! // On retry the client echoes `sealed` back untouched; the server re-derives
//! // the same context and opens it.
//! let opened = codec.open_with(&sealed, context).expect("integrity check passes");
//! assert_eq!(opened, b"step=2");
//!
//! // A different principal (different context) is rejected.
//! assert!(codec.open_with(&sealed, b"user:bob|tools/call:weather").is_err());
//! ```
use std::time::Duration;
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use hmac::{Hmac, Mac};
use serde::{Serialize, de::DeserializeOwned};
use sha2::Sha256;
use thiserror::Error;
type HmacSha256 = Hmac<Sha256>;
/// Version tag prefixing every sealed value, so the wire format can evolve.
const VERSION: &str = "rs1";
/// Domain-separation label mixed into the HMAC so a `requestState` tag can never
/// be confused with an HMAC computed for some other purpose using the same key.
const DOMAIN: &[u8] = b"rmcp/mrtr/request-state/v1";
/// Length of the big-endian expiry prefix (unix milliseconds) stored at the
/// front of every sealed body. `0` means "no expiry".
const EXPIRY_LEN: usize = 8;
/// Errors returned when opening a sealed [`RequestStateCodec`] value.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum RequestStateError {
/// The value is not a well-formed sealed request state (wrong prefix or
/// missing sections).
#[error("request state is malformed or uses an unsupported format")]
MalformedFormat,
/// A section of the value was not valid base64url.
#[error("request state is not valid base64url")]
InvalidEncoding,
/// The HMAC tag did not match; the value was forged, tampered with, or
/// opened with the wrong associated data.
#[error("request state failed integrity verification")]
IntegrityCheckFailed,
/// The value carried a TTL that has already elapsed.
#[error("request state has expired")]
Expired,
/// The sealed payload could not be serialized to JSON.
#[error("failed to serialize request state payload: {0}")]
Serialization(#[source] serde_json::Error),
/// The opened payload could not be deserialized from JSON.
#[error("failed to deserialize request state payload: {0}")]
Deserialization(#[source] serde_json::Error),
}
/// Options controlling how a value is sealed by [`RequestStateCodec`].
///
/// Defaults to no associated data and no expiry, which is equivalent to the
/// bare [`seal`](RequestStateCodec::seal) / [`open`](RequestStateCodec::open)
/// methods.
#[derive(Clone, Copy, Debug, Default)]
pub struct SealOptions<'a> {
associated_data: &'a [u8],
ttl: Option<Duration>,
}
impl<'a> SealOptions<'a> {
/// Creates empty options (no associated data, no expiry).
pub fn new() -> Self {
Self::default()
}
/// Binds the sealed value to `associated_data`. The same bytes must be
/// supplied to [`open_with`](RequestStateCodec::open_with); the data is
/// authenticated but not stored in the token.
///
/// Use this to bind the state to the authenticated principal and/or the
/// originating request (e.g. method name plus a digest of its parameters).
pub fn associated_data(mut self, associated_data: &'a [u8]) -> Self {
self.associated_data = associated_data;
self
}
/// Sets a relative time-to-live after which opening the value fails with
/// [`RequestStateError::Expired`].
pub fn ttl(mut self, ttl: Duration) -> Self {
self.ttl = Some(ttl);
self
}
}
/// A keyed codec that seals and opens SEP-2322 `requestState` values with
/// HMAC-SHA256 integrity protection.
///
/// Construct one codec per signing key and reuse it for the lifetime of the
/// key. The same key must be used to [`seal`](Self::seal) and
/// [`open`](Self::open) a value, so it has to survive across the rounds of a
/// single MRTR exchange (e.g. a stable per-process or per-deployment secret).
///
/// The key may be any length; HMAC internally normalizes it. For meaningful
/// security use a high-entropy key of at least 32 bytes.
#[derive(Clone)]
pub struct RequestStateCodec {
key: Box<[u8]>,
}
impl std::fmt::Debug for RequestStateCodec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
// Never leak the signing key through Debug output.
f.debug_struct("RequestStateCodec")
.field("key", &"<redacted>")
.finish()
}
}
impl RequestStateCodec {
/// Creates a codec from a signing key.
pub fn new(key: impl Into<Vec<u8>>) -> Self {
Self {
key: key.into().into_boxed_slice(),
}
}
/// Seals raw bytes into an opaque, integrity-protected string suitable for
/// use as `requestState`.
pub fn seal(&self, payload: &[u8]) -> String {
self.seal_with(payload, &SealOptions::default())
}
/// Seals raw bytes with [`SealOptions`] (associated data and/or TTL).
pub fn seal_with(&self, payload: &[u8], options: &SealOptions<'_>) -> String {
self.seal_at(payload, options, Self::now_ms())
}
/// Seals a serializable value by encoding it as JSON before sealing.
///
/// # Errors
///
/// Returns [`RequestStateError::Serialization`] if `value` cannot be encoded
/// as JSON.
pub fn seal_json<T: Serialize>(&self, value: &T) -> Result<String, RequestStateError> {
self.seal_json_with(value, &SealOptions::default())
}
/// Seals a serializable value with [`SealOptions`].
///
/// # Errors
///
/// Returns [`RequestStateError::Serialization`] if `value` cannot be encoded
/// as JSON.
pub fn seal_json_with<T: Serialize>(
&self,
value: &T,
options: &SealOptions<'_>,
) -> Result<String, RequestStateError> {
let payload = serde_json::to_vec(value).map_err(RequestStateError::Serialization)?;
Ok(self.seal_with(&payload, options))
}
/// Opens a sealed value that was sealed without associated data, verifying
/// its integrity and expiry and returning the original bytes.
///
/// # Errors
///
/// See [`open_with`](Self::open_with).
pub fn open(&self, sealed: &str) -> Result<Vec<u8>, RequestStateError> {
self.open_with(sealed, &[])
}
/// Opens a sealed value, verifying its integrity against `associated_data`
/// and checking its expiry.
///
/// `associated_data` must match the bytes passed to
/// [`SealOptions::associated_data`] when the value was sealed (use `&[]` for
/// values sealed without it).
///
/// # Errors
///
/// - [`RequestStateError::IntegrityCheckFailed`] if the value was not
/// produced by this key or the associated data differs.
/// - [`RequestStateError::Expired`] if the value's TTL has elapsed.
/// - [`RequestStateError::MalformedFormat`] or
/// [`RequestStateError::InvalidEncoding`] if it is not a well-formed sealed
/// value.
pub fn open_with(
&self,
sealed: &str,
associated_data: &[u8],
) -> Result<Vec<u8>, RequestStateError> {
self.open_at(sealed, associated_data, Self::now_ms())
}
/// Opens a sealed value (no associated data) and deserializes its JSON
/// payload.
///
/// # Errors
///
/// See [`open_json_with`](Self::open_json_with).
pub fn open_json<T: DeserializeOwned>(&self, sealed: &str) -> Result<T, RequestStateError> {
self.open_json_with(sealed, &[])
}
/// Opens a sealed value against `associated_data` and deserializes its JSON
/// payload.
///
/// # Errors
///
/// Returns the same integrity, expiry, and format errors as
/// [`open_with`](Self::open_with), plus [`RequestStateError::Deserialization`]
/// if the payload is not valid JSON for `T`.
pub fn open_json_with<T: DeserializeOwned>(
&self,
sealed: &str,
associated_data: &[u8],
) -> Result<T, RequestStateError> {
let payload = self.open_with(sealed, associated_data)?;
serde_json::from_slice(&payload).map_err(RequestStateError::Deserialization)
}
fn seal_at(&self, payload: &[u8], options: &SealOptions<'_>, now_ms: i64) -> String {
let expiry = match options.ttl {
Some(ttl) => now_ms.saturating_add(ttl.as_millis().min(i64::MAX as u128) as i64),
None => 0,
};
// body = big-endian expiry (0 = none) followed by the caller payload.
let mut body = Vec::with_capacity(EXPIRY_LEN + payload.len());
body.extend_from_slice(&expiry.to_be_bytes());
body.extend_from_slice(payload);
let tag = self
.mac_for(options.associated_data, &body)
.finalize()
.into_bytes();
// base64url without padding encodes 3 bytes as 4 chars, rounding up.
let b64_len = |n: usize| n.div_ceil(3) * 4;
let mut out =
String::with_capacity(VERSION.len() + 2 + b64_len(body.len()) + b64_len(tag.len()));
out.push_str(VERSION);
out.push('.');
URL_SAFE_NO_PAD.encode_string(&body, &mut out);
out.push('.');
URL_SAFE_NO_PAD.encode_string(tag.as_slice(), &mut out);
out
}
fn open_at(
&self,
sealed: &str,
associated_data: &[u8],
now_ms: i64,
) -> Result<Vec<u8>, RequestStateError> {
let mut parts = sealed.split('.');
let version = parts.next().ok_or(RequestStateError::MalformedFormat)?;
let body_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?;
let tag_b64 = parts.next().ok_or(RequestStateError::MalformedFormat)?;
if parts.next().is_some() || version != VERSION {
return Err(RequestStateError::MalformedFormat);
}
let body = URL_SAFE_NO_PAD
.decode(body_b64)
.map_err(|_| RequestStateError::InvalidEncoding)?;
let tag = URL_SAFE_NO_PAD
.decode(tag_b64)
.map_err(|_| RequestStateError::InvalidEncoding)?;
// `verify_slice` compares in constant time and rejects wrong-length tags.
self.mac_for(associated_data, &body)
.verify_slice(&tag)
.map_err(|_| RequestStateError::IntegrityCheckFailed)?;
// The body is now authenticated, so its framing can be trusted.
if body.len() < EXPIRY_LEN {
return Err(RequestStateError::MalformedFormat);
}
let expiry = i64::from_be_bytes(body[..EXPIRY_LEN].try_into().expect("checked length"));
if expiry != 0 && now_ms > expiry {
return Err(RequestStateError::Expired);
}
Ok(body[EXPIRY_LEN..].to_vec())
}
/// Builds an HMAC keyed for request-state tags, pre-fed with the
/// domain-separation label, a length-prefixed `associated_data`, and the
/// body. The length prefix keeps the `associated_data`/`body` boundary
/// unambiguous so distinct inputs cannot collide.
fn mac_for(&self, associated_data: &[u8], body: &[u8]) -> HmacSha256 {
let mut mac =
HmacSha256::new_from_slice(&self.key).expect("HMAC accepts keys of any length");
mac.update(DOMAIN);
mac.update(&(associated_data.len() as u64).to_be_bytes());
mac.update(associated_data);
mac.update(body);
mac
}
fn now_ms() -> i64 {
chrono::Utc::now().timestamp_millis()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn seal_open_roundtrips_bytes() {
let codec = RequestStateCodec::new(b"test-key-test-key-test-key-32byte".to_vec());
let sealed = codec.seal(b"hello world");
assert!(sealed.starts_with("rs1."));
assert_eq!(codec.open(&sealed).unwrap(), b"hello world");
}
#[test]
fn seal_open_roundtrips_json() {
#[derive(serde::Serialize, serde::Deserialize, PartialEq, Debug)]
struct State {
tool: String,
round: u32,
}
let codec = RequestStateCodec::new(b"another-strong-signing-key-here!!".to_vec());
let state = State {
tool: "weather".into(),
round: 3,
};
let sealed = codec.seal_json(&state).unwrap();
let opened: State = codec.open_json(&sealed).unwrap();
assert_eq!(opened, state);
}
#[test]
fn empty_payload_roundtrips() {
let codec = RequestStateCodec::new(b"k".to_vec());
let sealed = codec.seal(b"");
assert_eq!(codec.open(&sealed).unwrap(), b"");
}
#[test]
fn tampered_payload_is_rejected() {
let codec = RequestStateCodec::new(b"signing-key-signing-key-signing!!".to_vec());
let sealed = codec.seal(b"amount=100");
// Replace the body section but keep the original tag.
let mut parts: Vec<&str> = sealed.split('.').collect();
let forged_body = URL_SAFE_NO_PAD.encode(b"amount=999");
parts[1] = &forged_body;
let forged = parts.join(".");
assert!(matches!(
codec.open(&forged),
Err(RequestStateError::IntegrityCheckFailed)
));
}
#[test]
fn different_key_is_rejected() {
let signer = RequestStateCodec::new(b"the-real-signing-key-value-here!!".to_vec());
let attacker = RequestStateCodec::new(b"a-totally-different-forged-key!!!".to_vec());
let sealed = signer.seal(b"trusted");
assert!(matches!(
attacker.open(&sealed),
Err(RequestStateError::IntegrityCheckFailed)
));
}
#[test]
fn appended_bytes_are_rejected() {
let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec());
let mut sealed = codec.seal(b"state");
sealed.push('x');
assert!(codec.open(&sealed).is_err());
}
#[test]
fn wrong_version_prefix_is_malformed() {
let codec = RequestStateCodec::new(b"key".to_vec());
let sealed = codec.seal(b"state");
let bumped = sealed.replacen("rs1.", "rs2.", 1);
assert!(matches!(
codec.open(&bumped),
Err(RequestStateError::MalformedFormat)
));
}
#[test]
fn missing_sections_are_malformed() {
let codec = RequestStateCodec::new(b"key".to_vec());
assert!(matches!(
codec.open("rs1"),
Err(RequestStateError::MalformedFormat)
));
assert!(matches!(
codec.open("rs1.onlybody"),
Err(RequestStateError::MalformedFormat)
));
assert!(matches!(
codec.open("rs1.a.b.c"),
Err(RequestStateError::MalformedFormat)
));
}
#[test]
fn non_base64_sections_are_invalid_encoding() {
let codec = RequestStateCodec::new(b"key".to_vec());
assert!(matches!(
codec.open("rs1.!!!!.!!!!"),
Err(RequestStateError::InvalidEncoding)
));
}
#[test]
fn debug_does_not_leak_key() {
let codec = RequestStateCodec::new(b"super-secret-key".to_vec());
let rendered = format!("{codec:?}");
assert!(!rendered.contains("super-secret-key"));
assert!(rendered.contains("redacted"));
}
mod associated_data {
use super::*;
#[test]
fn matching_context_opens() {
let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec());
let ctx = b"user:alice|tools/call:weather";
let sealed = codec.seal_with(b"state", &SealOptions::new().associated_data(ctx));
assert_eq!(codec.open_with(&sealed, ctx).unwrap(), b"state");
}
#[test]
fn different_context_is_rejected() {
let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec());
let sealed =
codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice"));
assert!(matches!(
codec.open_with(&sealed, b"user:bob"),
Err(RequestStateError::IntegrityCheckFailed)
));
}
#[test]
fn missing_context_is_rejected() {
let codec = RequestStateCodec::new(b"key-key-key-key-key-key-key-key!!".to_vec());
let sealed =
codec.seal_with(b"state", &SealOptions::new().associated_data(b"user:alice"));
// Opening without the associated data must fail closed.
assert!(matches!(
codec.open(&sealed),
Err(RequestStateError::IntegrityCheckFailed)
));
}
}
mod ttl {
use super::*;
const KEY: &[u8] = b"ttl-signing-key-ttl-signing-key!!";
#[test]
fn within_ttl_opens() {
let codec = RequestStateCodec::new(KEY.to_vec());
let sealed = codec.seal_at(
b"state",
&SealOptions::new().ttl(Duration::from_secs(60)),
1_000,
);
// 30s later, still valid.
assert_eq!(codec.open_at(&sealed, &[], 31_000).unwrap(), b"state");
}
#[test]
fn past_ttl_is_expired() {
let codec = RequestStateCodec::new(KEY.to_vec());
let sealed = codec.seal_at(
b"state",
&SealOptions::new().ttl(Duration::from_secs(60)),
1_000,
);
// 61s later, expired.
assert!(matches!(
codec.open_at(&sealed, &[], 62_000),
Err(RequestStateError::Expired)
));
}
#[test]
fn no_ttl_never_expires() {
let codec = RequestStateCodec::new(KEY.to_vec());
let sealed = codec.seal_at(b"state", &SealOptions::new(), 1_000);
assert_eq!(codec.open_at(&sealed, &[], i64::MAX).unwrap(), b"state");
}
#[test]
fn ttl_and_associated_data_combine() {
let codec = RequestStateCodec::new(KEY.to_vec());
let ctx = b"user:alice";
let sealed = codec.seal_at(
b"state",
&SealOptions::new()
.associated_data(ctx)
.ttl(Duration::from_secs(60)),
1_000,
);
assert_eq!(codec.open_at(&sealed, ctx, 10_000).unwrap(), b"state");
assert!(matches!(
codec.open_at(&sealed, b"user:bob", 10_000),
Err(RequestStateError::IntegrityCheckFailed)
));
assert!(matches!(
codec.open_at(&sealed, ctx, 99_000),
Err(RequestStateError::Expired)
));
}
}
}

View file

@ -86,6 +86,9 @@ pub enum ServiceError {
Cancelled { reason: Option<String> },
#[error("request timeout after {}", chrono::Duration::from_std(*timeout).unwrap_or_default())]
Timeout { timeout: Duration },
/// The peer kept returning `input_required` beyond the configured round cap.
#[error("input_required did not complete within {max_rounds} MRTR rounds")]
InputRequiredRoundsExceeded { max_rounds: usize },
}
trait TransferObject:

View file

@ -1,24 +1,26 @@
// Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected.
#![expect(deprecated)]
use std::borrow::Cow;
use std::{borrow::Cow, sync::Arc, time::Duration};
use thiserror::Error;
use super::*;
use crate::{
model::{
ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResult,
ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResponse, CallToolResult,
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
ClientNotification, ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams,
CompleteResult, CompletionContext, CompletionInfo, ErrorData, GetPromptRequest,
GetPromptRequestParams, GetPromptResult, InitializeRequest, InitializedNotification,
JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
ListToolsResult, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam,
ReadResourceRequest, ReadResourceRequestParams, ReadResourceResult, Reference, RequestId,
RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification,
ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest,
SubscribeRequestParams, UnsubscribeRequest, UnsubscribeRequestParams,
CompleteResult, CompletionContext, CompletionInfo, DEFAULT_MRTR_MAX_ROUNDS, ErrorData,
GetExtensions, GetMeta, GetPromptRequest, GetPromptRequestParams, GetPromptResponse,
GetPromptResult, InitializeRequest, InitializedNotification, InputRequest,
InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest,
ListPromptsResult, ListResourceTemplatesRequest, ListResourceTemplatesResult,
ListResourcesRequest, ListResourcesResult, ListToolsRequest, ListToolsResult,
NumberOrString, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam,
ReadResourceRequest, ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult,
Reference, RequestId, RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage,
ServerNotification, ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams,
SubscribeRequest, SubscribeRequestParams, UnsubscribeRequest, UnsubscribeRequestParams,
},
transport::DynamicTransportError,
};
@ -361,6 +363,72 @@ macro_rules! method {
}
impl Peer<RoleClient> {
/// Send one `tools/call` request and return either a final result or an MRTR
/// `InputRequiredResult` without driving any follow-up rounds.
pub async fn call_tool_once(
&self,
params: CallToolRequestParams,
) -> Result<CallToolResponse, ServiceError> {
let result = self
.send_request(ClientRequest::CallToolRequest(CallToolRequest {
method: Default::default(),
params,
extensions: Default::default(),
}))
.await?;
match result {
ServerResult::CallToolResult(result) => Ok(CallToolResponse::Complete(result)),
ServerResult::InputRequiredResult(result) => {
Ok(CallToolResponse::InputRequired(result))
}
_ => Err(ServiceError::UnexpectedResponse),
}
}
/// Send one `prompts/get` request and return either a final result or an MRTR
/// `InputRequiredResult` without driving any follow-up rounds.
pub async fn get_prompt_once(
&self,
params: GetPromptRequestParams,
) -> Result<GetPromptResponse, ServiceError> {
let result = self
.send_request(ClientRequest::GetPromptRequest(GetPromptRequest {
method: Default::default(),
params,
extensions: Default::default(),
}))
.await?;
match result {
ServerResult::GetPromptResult(result) => Ok(GetPromptResponse::Complete(result)),
ServerResult::InputRequiredResult(result) => {
Ok(GetPromptResponse::InputRequired(result))
}
_ => Err(ServiceError::UnexpectedResponse),
}
}
/// Send one `resources/read` request and return either a final result or an
/// MRTR `InputRequiredResult` without driving any follow-up rounds.
pub async fn read_resource_once(
&self,
params: ReadResourceRequestParams,
) -> Result<ReadResourceResponse, ServiceError> {
let result = self
.send_request(ClientRequest::ReadResourceRequest(ReadResourceRequest {
method: Default::default(),
params,
extensions: Default::default(),
}))
.await?;
match result {
ServerResult::ReadResourceResult(result) => Ok(ReadResourceResponse::Complete(result)),
ServerResult::InputRequiredResult(result) => {
Ok(ReadResourceResponse::InputRequired(result))
}
_ => Err(ServiceError::UnexpectedResponse),
}
}
method!(peer_req complete CompleteRequest(CompleteRequestParams) => CompleteResult);
method!(
#[deprecated(
@ -558,3 +626,294 @@ impl Peer<RoleClient> {
Ok(completion.values)
}
}
impl<S> RunningService<RoleClient, S>
where
S: Service<RoleClient>,
{
/// Send one `tools/call` request without driving MRTR follow-up rounds.
pub async fn call_tool_once(
&self,
params: CallToolRequestParams,
) -> Result<CallToolResponse, ServiceError> {
self.peer.call_tool_once(params).await
}
/// Send one `prompts/get` request without driving MRTR follow-up rounds.
pub async fn get_prompt_once(
&self,
params: GetPromptRequestParams,
) -> Result<GetPromptResponse, ServiceError> {
self.peer.get_prompt_once(params).await
}
/// Send one `resources/read` request without driving MRTR follow-up rounds.
pub async fn read_resource_once(
&self,
params: ReadResourceRequestParams,
) -> Result<ReadResourceResponse, ServiceError> {
self.peer.read_resource_once(params).await
}
/// High-level `tools/call` helper that automatically fulfils SEP-2322
/// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does
/// not produce a final [`CallToolResult`] within the default MRTR round cap.
/// Other transport, protocol, and local input-handler errors are propagated.
pub async fn call_tool(
&self,
params: CallToolRequestParams,
) -> Result<CallToolResult, ServiceError> {
self.call_tool_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS)
.await
}
/// Same as [`Self::call_tool`], with an explicit MRTR round cap.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds`
/// `input_required` responses have been driven without receiving a final
/// [`CallToolResult`]. Other transport, protocol, and local input-handler
/// errors are propagated.
pub async fn call_tool_with_mrtr_max_rounds(
&self,
mut params: CallToolRequestParams,
max_rounds: usize,
) -> Result<CallToolResult, ServiceError> {
let mut state_only_rounds = 0usize;
for _round in 0..max_rounds {
match self.peer.call_tool_once(params.clone()).await? {
CallToolResponse::Complete(result) => return Ok(result),
CallToolResponse::InputRequired(result) => {
let (input_responses, request_state) = self
.prepare_input_required_retry(result, &mut state_only_rounds)
.await?;
params.input_responses = input_responses;
params.request_state = request_state;
}
}
}
Err(ServiceError::InputRequiredRoundsExceeded { max_rounds })
}
/// High-level `prompts/get` helper that automatically fulfils SEP-2322
/// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does
/// not produce a final [`GetPromptResult`] within the default MRTR round cap.
/// Other transport, protocol, and local input-handler errors are propagated.
pub async fn get_prompt(
&self,
params: GetPromptRequestParams,
) -> Result<GetPromptResult, ServiceError> {
self.get_prompt_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS)
.await
}
/// Same as [`Self::get_prompt`], with an explicit MRTR round cap.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds`
/// `input_required` responses have been driven without receiving a final
/// [`GetPromptResult`]. Other transport, protocol, and local input-handler
/// errors are propagated.
pub async fn get_prompt_with_mrtr_max_rounds(
&self,
mut params: GetPromptRequestParams,
max_rounds: usize,
) -> Result<GetPromptResult, ServiceError> {
let mut state_only_rounds = 0usize;
for _round in 0..max_rounds {
match self.peer.get_prompt_once(params.clone()).await? {
GetPromptResponse::Complete(result) => return Ok(result),
GetPromptResponse::InputRequired(result) => {
let (input_responses, request_state) = self
.prepare_input_required_retry(result, &mut state_only_rounds)
.await?;
params.input_responses = input_responses;
params.request_state = request_state;
}
}
}
Err(ServiceError::InputRequiredRoundsExceeded { max_rounds })
}
/// High-level `resources/read` helper that automatically fulfils SEP-2322
/// `input_required` rounds through the local [`ClientHandler`](crate::ClientHandler) service.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] if the peer does
/// not produce a final [`ReadResourceResult`] within the default MRTR round
/// cap. Other transport, protocol, and local input-handler errors are
/// propagated.
pub async fn read_resource(
&self,
params: ReadResourceRequestParams,
) -> Result<ReadResourceResult, ServiceError> {
self.read_resource_with_mrtr_max_rounds(params, DEFAULT_MRTR_MAX_ROUNDS)
.await
}
/// Same as [`Self::read_resource`], with an explicit MRTR round cap.
///
/// # Errors
///
/// Returns [`ServiceError::InputRequiredRoundsExceeded`] once `max_rounds`
/// `input_required` responses have been driven without receiving a final
/// [`ReadResourceResult`]. Other transport, protocol, and local input-handler
/// errors are propagated.
pub async fn read_resource_with_mrtr_max_rounds(
&self,
mut params: ReadResourceRequestParams,
max_rounds: usize,
) -> Result<ReadResourceResult, ServiceError> {
let mut state_only_rounds = 0usize;
for _round in 0..max_rounds {
match self.peer.read_resource_once(params.clone()).await? {
ReadResourceResponse::Complete(result) => return Ok(result),
ReadResourceResponse::InputRequired(result) => {
let (input_responses, request_state) = self
.prepare_input_required_retry(result, &mut state_only_rounds)
.await?;
params.input_responses = input_responses;
params.request_state = request_state;
}
}
}
Err(ServiceError::InputRequiredRoundsExceeded { max_rounds })
}
async fn prepare_input_required_retry(
&self,
result: InputRequiredResult,
state_only_rounds: &mut usize,
) -> Result<(Option<InputResponses>, Option<String>), ServiceError> {
let had_input_requests = result
.input_requests
.as_ref()
.is_some_and(|requests| !requests.is_empty());
if !had_input_requests && result.request_state.is_none() {
return Err(ServiceError::UnexpectedResponse);
}
let responses = self
.fulfill_input_requests(result.input_requests.unwrap_or_default())
.await?;
if had_input_requests {
*state_only_rounds = 0;
} else {
Self::sleep_state_only_round(*state_only_rounds).await;
*state_only_rounds += 1;
}
Ok((
(!responses.is_empty()).then_some(responses),
result.request_state,
))
}
async fn fulfill_input_requests(
&self,
requests: crate::model::InputRequests,
) -> Result<InputResponses, ServiceError> {
let responses = futures::future::try_join_all(
requests
.into_iter()
.map(|(key, request)| self.fulfill_input_request(key, request)),
)
.await?;
Ok(responses.into_iter().collect())
}
async fn fulfill_input_request(
&self,
key: String,
request: InputRequest,
) -> Result<(String, serde_json::Value), ServiceError> {
let response = match request {
InputRequest::CreateMessage(request) => {
let mut request = ServerRequest::CreateMessageRequest(request);
let context = self.input_request_context(&key, &mut request);
match self
.service
.handle_request(request, context)
.await
.map_err(ServiceError::McpError)?
{
ClientResult::CreateMessageResult(result) => {
serde_json::to_value(result).map_err(Self::serde_to_service_error)?
}
_ => return Err(ServiceError::UnexpectedResponse),
}
}
InputRequest::Elicitation(request) => {
let mut request = ServerRequest::ElicitRequest(request);
let context = self.input_request_context(&key, &mut request);
match self
.service
.handle_request(request, context)
.await
.map_err(ServiceError::McpError)?
{
ClientResult::ElicitResult(result) => {
serde_json::to_value(result).map_err(Self::serde_to_service_error)?
}
_ => return Err(ServiceError::UnexpectedResponse),
}
}
InputRequest::ListRoots(request) => {
let mut request = ServerRequest::ListRootsRequest(request);
let context = self.input_request_context(&key, &mut request);
match self
.service
.handle_request(request, context)
.await
.map_err(ServiceError::McpError)?
{
ClientResult::ListRootsResult(result) => {
serde_json::to_value(result).map_err(Self::serde_to_service_error)?
}
_ => return Err(ServiceError::UnexpectedResponse),
}
}
};
Ok((key, response))
}
fn input_request_context<T>(&self, key: &str, request: &mut T) -> RequestContext<RoleClient>
where
T: GetMeta + GetExtensions,
{
let mut meta = Default::default();
let mut extensions = Default::default();
std::mem::swap(&mut meta, request.get_meta_mut());
std::mem::swap(&mut extensions, request.extensions_mut());
RequestContext {
ct: tokio_util::sync::CancellationToken::new(),
id: NumberOrString::String(Arc::from(key)),
peer: self.peer.clone(),
meta,
extensions,
}
}
async fn sleep_state_only_round(state_only_rounds: usize) {
let millis = (50u64.saturating_mul(1_u64 << state_only_rounds.min(3))).min(250);
tokio::time::sleep(Duration::from_millis(millis)).await;
}
fn serde_to_service_error(error: serde_json::Error) -> ServiceError {
ServiceError::McpError(ErrorData::internal_error(
format!("failed to serialize MRTR input response: {error}"),
None,
))
}
}