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:
parent
7e14020cb8
commit
c4def55422
2 changed files with 114 additions and 3 deletions
|
|
@ -34,10 +34,18 @@ pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, cr
|
||||||
}
|
}
|
||||||
#[non_exhaustive]
|
#[non_exhaustive]
|
||||||
pub struct ToolCallContext<'s, S> {
|
pub struct ToolCallContext<'s, S> {
|
||||||
|
/// The request-specific context for this tool call.
|
||||||
pub request_context: RequestContext<RoleServer>,
|
pub request_context: RequestContext<RoleServer>,
|
||||||
|
/// The server handling this tool call.
|
||||||
pub service: &'s S,
|
pub service: &'s S,
|
||||||
|
/// The name of the tool being called.
|
||||||
pub name: Cow<'static, str>,
|
pub name: Cow<'static, str>,
|
||||||
|
/// The arguments supplied for the tool call.
|
||||||
pub arguments: Option<JsonObject>,
|
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> {
|
impl<'s, S> ToolCallContext<'s, S> {
|
||||||
|
|
@ -47,6 +55,8 @@ impl<'s, S> ToolCallContext<'s, S> {
|
||||||
meta: _,
|
meta: _,
|
||||||
name,
|
name,
|
||||||
arguments,
|
arguments,
|
||||||
|
input_responses,
|
||||||
|
request_state,
|
||||||
..
|
..
|
||||||
}: CallToolRequestParams,
|
}: CallToolRequestParams,
|
||||||
request_context: RequestContext<RoleServer>,
|
request_context: RequestContext<RoleServer>,
|
||||||
|
|
@ -56,6 +66,8 @@ impl<'s, S> ToolCallContext<'s, S> {
|
||||||
service,
|
service,
|
||||||
name,
|
name,
|
||||||
arguments,
|
arguments,
|
||||||
|
input_responses,
|
||||||
|
request_state,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
pub fn name(&self) -> &str {
|
pub fn name(&self) -> &str {
|
||||||
|
|
@ -98,6 +110,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 {
|
impl IntoCallToolResult for crate::ErrorData {
|
||||||
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
|
fn into_call_tool_result(self) -> Result<CallToolResponse, crate::ErrorData> {
|
||||||
Err(self)
|
Err(self)
|
||||||
|
|
@ -193,6 +211,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
|
// Special implementation for Parameters that handles tool arguments
|
||||||
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
|
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
|
||||||
where
|
where
|
||||||
|
|
|
||||||
|
|
@ -13,9 +13,16 @@ use std::sync::{
|
||||||
|
|
||||||
use rmcp::{
|
use rmcp::{
|
||||||
ClientHandler, ServerHandler,
|
ClientHandler, ServerHandler,
|
||||||
|
handler::server::{
|
||||||
|
tool::{InputResponses as ToolInputResponses, RequestState},
|
||||||
|
wrapper::Parameters,
|
||||||
|
},
|
||||||
model::*,
|
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;
|
use serde_json::json;
|
||||||
|
|
||||||
/// A `requestState` value with characters that must survive a byte-exact echo:
|
/// 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()))
|
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 {
|
impl MrtrServer {
|
||||||
fn call_tool_impl(
|
fn call_tool_impl(
|
||||||
&self,
|
&self,
|
||||||
|
|
@ -289,12 +344,13 @@ fn server_info(protocol_version: ProtocolVersion) -> ServerInfo {
|
||||||
|
|
||||||
/// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local`
|
/// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local`
|
||||||
/// feature is active) is available, wiring up a connected client/server pair.
|
/// feature is active) is available, wiring up a connected client/server pair.
|
||||||
async fn with_pair<F, Fut>(
|
async fn with_pair<S, F, Fut>(
|
||||||
server: MrtrServer,
|
server: S,
|
||||||
client_protocol: ProtocolVersion,
|
client_protocol: ProtocolVersion,
|
||||||
body: F,
|
body: F,
|
||||||
) -> anyhow::Result<()>
|
) -> anyhow::Result<()>
|
||||||
where
|
where
|
||||||
|
S: Service<RoleServer>,
|
||||||
F: FnOnce(rmcp::service::RunningService<RoleClient, MrtrClient>) -> Fut,
|
F: FnOnce(rmcp::service::RunningService<RoleClient, MrtrClient>) -> Fut,
|
||||||
Fut: std::future::Future<Output = anyhow::Result<()>>,
|
Fut: std::future::Future<Output = anyhow::Result<()>>,
|
||||||
{
|
{
|
||||||
|
|
@ -346,6 +402,23 @@ async fn client_auto_fulfills_input_required_tool_call() -> anyhow::Result<()> {
|
||||||
.await
|
.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")]
|
#[tokio::test(flavor = "current_thread")]
|
||||||
async fn manual_once_returns_input_required_without_retry() -> anyhow::Result<()> {
|
async fn manual_once_returns_input_required_without_retry() -> anyhow::Result<()> {
|
||||||
let server = MrtrServer::default();
|
let server = MrtrServer::default();
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue