feat: Add prompt support (#351)
* feat: add prompt support with typed argument handling - Implement #[prompt], #[prompt_router], and #[prompt_handler] macros - Add automatic JSON schema generation from Rust types for arguments - Support flexible async handler signatures with automatic adaptation - Create PromptRouter for efficient prompt dispatch - Include comprehensive tests and example implementation This enables MCP servers to provide reusable prompt templates that LLMs can discover and invoke with strongly-typed parameters, similar to the existing tool system but optimized for prompt use cases. 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com> * refactor: unify parameter handling between tools and prompts - Replace Arguments<T> with Parameters<T> for consistent API - Create shared common module for tool/prompt utilities - Modernize async handling with futures::future::BoxFuture - Move cached_schema_for_type to common module for reuse - Update error types from rmcp::Error to rmcp::ErrorData - Add comprehensive trait implementations for parameter extraction 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com> * refactor: extract Parameters wrapper to shared module and unify trait usage - Move Parameters type from common.rs to dedicated wrapper/parameters.rs module - Unify extractor traits by using FromContextPart for both tools and prompts - Remove duplicate FromToolCallContextPart and FromPromptContextPart traits - Add structured output support to Json wrapper with IntoCallToolResult impl - Improve error messages with more descriptive panic for schema serialization - Update all imports across codebase to use new module path - Clean up trailing whitespace and formatting inconsistencies This consolidates parameter extraction logic and reduces code duplication between tool and prompt handlers while maintaining backward compatibility. * chore: remove committed .egg-info directory and update gitignore * docs: fix documentation formatting
This commit is contained in:
Michael Assaf
GitHub
parent
c0452f93a4
commit
9349f5cb26
11 changed files with 3191 additions and 464 deletions
|
|
@ -4,6 +4,8 @@ use crate::{
|
|||
service::{NotificationContext, RequestContext, RoleServer, Service, ServiceRole},
|
||||
};
|
||||
|
||||
pub mod common;
|
||||
pub mod prompt;
|
||||
mod resource;
|
||||
pub mod router;
|
||||
pub mod tool;
|
||||
|
|
|
|||
146
crates/rmcp/src/handler/server/common.rs
Normal file
146
crates/rmcp/src/handler/server/common.rs
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
//! Common utilities shared between tool and prompt handlers
|
||||
|
||||
use std::{any::TypeId, collections::HashMap, sync::Arc};
|
||||
|
||||
use schemars::JsonSchema;
|
||||
|
||||
use crate::{
|
||||
RoleServer, model::JsonObject, schemars::generate::SchemaSettings, service::RequestContext,
|
||||
};
|
||||
|
||||
/// A shortcut for generating a JSON schema for a type.
|
||||
pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
|
||||
// explicitly to align json schema version to official specifications.
|
||||
// https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.json
|
||||
// TODO: update to 2020-12 waiting for the mcp spec update
|
||||
let mut settings = SchemaSettings::draft07();
|
||||
settings.transforms = vec![Box::new(schemars::transform::AddNullable::default())];
|
||||
let generator = settings.into_generator();
|
||||
let schema = generator.into_root_schema_for::<T>();
|
||||
let object = serde_json::to_value(schema).expect("failed to serialize schema");
|
||||
match object {
|
||||
serde_json::Value::Object(object) => object,
|
||||
_ => panic!(
|
||||
"Schema serialization produced non-object value: expected JSON object but got {:?}",
|
||||
object
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Call [`schema_for_type`] with a cache
|
||||
pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
|
||||
thread_local! {
|
||||
static CACHE_FOR_TYPE: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
|
||||
};
|
||||
CACHE_FOR_TYPE.with(|cache| {
|
||||
if let Some(x) = cache
|
||||
.read()
|
||||
.expect("schema cache lock poisoned")
|
||||
.get(&TypeId::of::<T>())
|
||||
{
|
||||
x.clone()
|
||||
} else {
|
||||
let schema = schema_for_type::<T>();
|
||||
let schema = Arc::new(schema);
|
||||
cache
|
||||
.write()
|
||||
.expect("schema cache lock poisoned")
|
||||
.insert(TypeId::of::<T>(), schema.clone());
|
||||
schema
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Trait for extracting parts from a context, unifying tool and prompt extraction
|
||||
pub trait FromContextPart<C>: Sized {
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData>;
|
||||
}
|
||||
|
||||
/// Common extractors that can be used by both tool and prompt handlers
|
||||
impl<C> FromContextPart<C> for RequestContext<RoleServer>
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.as_request_context().clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> FromContextPart<C> for tokio_util::sync::CancellationToken
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.as_request_context().ct.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> FromContextPart<C> for crate::model::Extensions
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.as_request_context().extensions.clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Extension<T>(pub T);
|
||||
|
||||
impl<C, T> FromContextPart<C> for Extension<T>
|
||||
where
|
||||
C: AsRequestContext,
|
||||
T: Send + Sync + 'static + Clone,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
let extension = context
|
||||
.as_request_context()
|
||||
.extensions
|
||||
.get::<T>()
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
crate::ErrorData::invalid_params(
|
||||
format!("missing extension {}", std::any::type_name::<T>()),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
Ok(Extension(extension))
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> FromContextPart<C> for crate::Peer<RoleServer>
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.as_request_context().peer.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> FromContextPart<C> for crate::model::Meta
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
let request_context = context.as_request_context_mut();
|
||||
let mut meta = crate::model::Meta::default();
|
||||
std::mem::swap(&mut meta, &mut request_context.meta);
|
||||
Ok(meta)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RequestId(pub crate::model::RequestId);
|
||||
|
||||
impl<C> FromContextPart<C> for RequestId
|
||||
where
|
||||
C: AsRequestContext,
|
||||
{
|
||||
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
|
||||
Ok(RequestId(context.as_request_context().id.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for types that can provide access to RequestContext
|
||||
pub trait AsRequestContext {
|
||||
fn as_request_context(&self) -> &RequestContext<RoleServer>;
|
||||
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer>;
|
||||
}
|
||||
366
crates/rmcp/src/handler/server/prompt.rs
Normal file
366
crates/rmcp/src/handler/server/prompt.rs
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
//! Prompt handling infrastructure for MCP servers
|
||||
//!
|
||||
//! This module provides the core types and traits for implementing prompt handlers
|
||||
//! in MCP servers. Prompts allow servers to provide reusable templates for LLM
|
||||
//! interactions with customizable arguments.
|
||||
|
||||
use std::{future::Future, marker::PhantomData};
|
||||
|
||||
use futures::future::{BoxFuture, FutureExt};
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
use super::common::{AsRequestContext, FromContextPart};
|
||||
pub use super::common::{Extension, RequestId};
|
||||
use crate::{
|
||||
RoleServer,
|
||||
handler::server::wrapper::Parameters,
|
||||
model::{GetPromptResult, PromptMessage},
|
||||
service::RequestContext,
|
||||
};
|
||||
|
||||
/// Context for prompt retrieval operations
|
||||
pub struct PromptContext<'a, S> {
|
||||
pub server: &'a S,
|
||||
pub name: String,
|
||||
pub arguments: Option<serde_json::Map<String, serde_json::Value>>,
|
||||
pub context: RequestContext<RoleServer>,
|
||||
}
|
||||
|
||||
impl<'a, S> PromptContext<'a, S> {
|
||||
pub fn new(
|
||||
server: &'a S,
|
||||
name: String,
|
||||
arguments: Option<serde_json::Map<String, serde_json::Value>>,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
server,
|
||||
name,
|
||||
arguments,
|
||||
context,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> AsRequestContext for PromptContext<'_, S> {
|
||||
fn as_request_context(&self) -> &RequestContext<RoleServer> {
|
||||
&self.context
|
||||
}
|
||||
|
||||
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
|
||||
&mut self.context
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for handling prompt retrieval
|
||||
pub trait GetPromptHandler<S, A> {
|
||||
fn handle(
|
||||
self,
|
||||
context: PromptContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>;
|
||||
}
|
||||
|
||||
/// Type alias for dynamic prompt handlers
|
||||
pub type DynGetPromptHandler<S> = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
|
||||
+ Send
|
||||
+ Sync;
|
||||
|
||||
/// Adapter type for async methods that return `Vec<PromptMessage>`
|
||||
pub struct AsyncMethodAdapter<T>(PhantomData<T>);
|
||||
|
||||
/// Adapter type for async methods with parameters that return `Vec<PromptMessage>`
|
||||
pub struct AsyncMethodWithArgsAdapter<T>(PhantomData<T>);
|
||||
|
||||
/// Adapter types for macro-generated implementations
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub struct AsyncPromptAdapter<P, Fut, R>(PhantomData<fn(P) -> fn(Fut) -> R>);
|
||||
pub struct SyncPromptAdapter<P, R>(PhantomData<fn(P) -> R>);
|
||||
pub struct AsyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
|
||||
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>;
|
||||
}
|
||||
|
||||
impl IntoGetPromptResult for GetPromptResult {
|
||||
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoGetPromptResult for Vec<PromptMessage> {
|
||||
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
|
||||
Ok(GetPromptResult {
|
||||
description: None,
|
||||
messages: self,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: IntoGetPromptResult> IntoGetPromptResult for Result<T, crate::ErrorData> {
|
||||
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
|
||||
self.and_then(|v| v.into_get_prompt_result())
|
||||
}
|
||||
}
|
||||
|
||||
// Future wrapper that automatically handles IntoGetPromptResult conversion
|
||||
pin_project_lite::pin_project! {
|
||||
#[project = IntoGetPromptResultFutProj]
|
||||
pub enum IntoGetPromptResultFut<F, R> {
|
||||
Pending {
|
||||
#[pin]
|
||||
fut: F,
|
||||
_marker: PhantomData<R>,
|
||||
},
|
||||
Ready {
|
||||
#[pin]
|
||||
result: futures::future::Ready<Result<GetPromptResult, crate::ErrorData>>,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F, R> Future for IntoGetPromptResultFut<F, R>
|
||||
where
|
||||
F: Future<Output = R>,
|
||||
R: IntoGetPromptResult,
|
||||
{
|
||||
type Output = Result<GetPromptResult, crate::ErrorData>;
|
||||
|
||||
fn poll(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Self::Output> {
|
||||
match self.project() {
|
||||
IntoGetPromptResultFutProj::Pending { fut, _marker } => fut
|
||||
.poll(cx)
|
||||
.map(IntoGetPromptResult::into_get_prompt_result),
|
||||
IntoGetPromptResultFutProj::Ready { result } => result.poll(cx),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Prompt-specific extractor for prompt name
|
||||
pub struct PromptName(pub String);
|
||||
|
||||
impl<S> FromContextPart<PromptContext<'_, S>> for PromptName {
|
||||
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
|
||||
Ok(Self(context.name.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
// Special implementation for Parameters that handles prompt arguments
|
||||
impl<S, P> FromContextPart<PromptContext<'_, S>> for Parameters<P>
|
||||
where
|
||||
P: DeserializeOwned,
|
||||
{
|
||||
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
|
||||
let params = if let Some(args_map) = context.arguments.take() {
|
||||
let args_value = serde_json::Value::Object(args_map);
|
||||
serde_json::from_value::<P>(args_value).map_err(|e| {
|
||||
crate::ErrorData::invalid_params(format!("Failed to parse parameters: {}", e), None)
|
||||
})?
|
||||
} else {
|
||||
// Try to deserialize from empty object for optional fields
|
||||
serde_json::from_value::<P>(serde_json::json!({})).map_err(|e| {
|
||||
crate::ErrorData::invalid_params(
|
||||
format!("Missing required parameters: {}", e),
|
||||
None,
|
||||
)
|
||||
})?
|
||||
};
|
||||
Ok(Parameters(params))
|
||||
}
|
||||
}
|
||||
|
||||
// Macro to generate GetPromptHandler implementations for various parameter combinations
|
||||
macro_rules! impl_prompt_handler_for {
|
||||
($($T: ident)*) => {
|
||||
impl_prompt_handler_for!([] [$($T)*]);
|
||||
};
|
||||
// finished
|
||||
([$($Tn: ident)*] []) => {
|
||||
impl_prompt_handler_for!(@impl $($Tn)*);
|
||||
};
|
||||
([$($Tn: ident)*] [$Tn_1: ident $($Rest: ident)*]) => {
|
||||
impl_prompt_handler_for!(@impl $($Tn)*);
|
||||
impl_prompt_handler_for!([$($Tn)* $Tn_1] [$($Rest)*]);
|
||||
};
|
||||
(@impl $($Tn: ident)*) => {
|
||||
// Implementation for async methods (transformed by #[prompt] macro)
|
||||
impl<$($Tn,)* S, F, R> GetPromptHandler<S, ($($Tn,)*)> for F
|
||||
where
|
||||
$(
|
||||
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
|
||||
)*
|
||||
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R> + Send,
|
||||
R: IntoGetPromptResult + Send + 'static,
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
#[allow(unused_variables, non_snake_case, unused_mut)]
|
||||
fn handle(
|
||||
self,
|
||||
mut context: PromptContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
|
||||
{
|
||||
$(
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
};
|
||||
)*
|
||||
let service = context.server;
|
||||
let fut = self(service, $($Tn,)*);
|
||||
async move {
|
||||
let result = fut.await;
|
||||
result.into_get_prompt_result()
|
||||
}.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Implementation for sync methods
|
||||
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptMethodAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
|
||||
)*
|
||||
F: FnOnce(&S, $($Tn,)*) -> R + Send,
|
||||
R: IntoGetPromptResult + Send,
|
||||
S: Send + Sync,
|
||||
{
|
||||
#[allow(unused_variables, non_snake_case, unused_mut)]
|
||||
fn handle(
|
||||
self,
|
||||
mut context: PromptContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
|
||||
{
|
||||
$(
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
};
|
||||
)*
|
||||
let service = context.server;
|
||||
let result = self(service, $($Tn,)*);
|
||||
std::future::ready(result.into_get_prompt_result()).boxed()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// AsyncPromptAdapter - for standalone functions returning GetPromptResult
|
||||
impl<$($Tn,)* S, F, Fut, R> GetPromptHandler<S, AsyncPromptAdapter<($($Tn,)*), Fut, R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = Result<R, crate::ErrorData>> + Send + 'static,
|
||||
R: IntoGetPromptResult + Send + 'static,
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
#[allow(unused_variables, non_snake_case, unused_mut)]
|
||||
fn handle(
|
||||
self,
|
||||
mut context: PromptContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
|
||||
{
|
||||
// Extract all parameters before moving into the async block
|
||||
$(
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
};
|
||||
)*
|
||||
|
||||
// Since we're dealing with standalone functions that don't take &S,
|
||||
// we can return a 'static future
|
||||
Box::pin(async move {
|
||||
let result = self($($Tn,)*).await?;
|
||||
result.into_get_prompt_result()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// SyncPromptAdapter - for standalone sync functions returning Result
|
||||
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> Result<R, crate::ErrorData> + Send + 'static,
|
||||
R: IntoGetPromptResult + Send + 'static,
|
||||
S: Send + Sync,
|
||||
{
|
||||
#[allow(unused_variables, non_snake_case, unused_mut)]
|
||||
fn handle(
|
||||
self,
|
||||
mut context: PromptContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
|
||||
{
|
||||
$(
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
};
|
||||
)*
|
||||
let result = self($($Tn,)*);
|
||||
std::future::ready(result.and_then(|r| r.into_get_prompt_result())).boxed()
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
}
|
||||
|
||||
// Invoke the macro to generate implementations for up to 16 parameters
|
||||
impl_prompt_handler_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
|
||||
|
||||
/// Extract prompt arguments from a type's JSON schema
|
||||
/// This function analyzes the schema of a type and extracts the properties
|
||||
/// as PromptArgument entries with name, description, and required status
|
||||
pub fn cached_arguments_from_schema<T: schemars::JsonSchema + std::any::Any>()
|
||||
-> Option<Vec<crate::model::PromptArgument>> {
|
||||
let schema = super::common::cached_schema_for_type::<T>();
|
||||
let schema_value = serde_json::Value::Object((*schema).clone());
|
||||
|
||||
let properties = schema_value.get("properties").and_then(|p| p.as_object());
|
||||
|
||||
if let Some(props) = properties {
|
||||
let required = schema_value
|
||||
.get("required")
|
||||
.and_then(|r| r.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str())
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut arguments = Vec::new();
|
||||
for (name, prop_schema) in props {
|
||||
let description = prop_schema
|
||||
.get("description")
|
||||
.and_then(|d| d.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
arguments.push(crate::model::PromptArgument {
|
||||
name: name.clone(),
|
||||
description,
|
||||
required: Some(required.contains(name.as_str())),
|
||||
});
|
||||
}
|
||||
|
||||
if arguments.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(arguments)
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
|
@ -1,18 +1,21 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use prompt::{IntoPromptRoute, PromptRoute};
|
||||
use tool::{IntoToolRoute, ToolRoute};
|
||||
|
||||
use super::ServerHandler;
|
||||
use crate::{
|
||||
RoleServer, Service,
|
||||
model::{ClientRequest, ListToolsResult, ServerResult},
|
||||
model::{ClientRequest, ListPromptsResult, ListToolsResult, ServerResult},
|
||||
service::NotificationContext,
|
||||
};
|
||||
|
||||
pub mod prompt;
|
||||
pub mod tool;
|
||||
|
||||
pub struct Router<S> {
|
||||
pub tool_router: tool::ToolRouter<S>,
|
||||
pub prompt_router: prompt::PromptRouter<S>,
|
||||
pub service: Arc<S>,
|
||||
}
|
||||
|
||||
|
|
@ -23,6 +26,7 @@ where
|
|||
pub fn new(service: S) -> Self {
|
||||
Self {
|
||||
tool_router: tool::ToolRouter::new(),
|
||||
prompt_router: prompt::PromptRouter::new(),
|
||||
service: Arc::new(service),
|
||||
}
|
||||
}
|
||||
|
|
@ -41,6 +45,21 @@ where
|
|||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_prompt<R, A: 'static>(mut self, route: R) -> Self
|
||||
where
|
||||
R: IntoPromptRoute<S, A>,
|
||||
{
|
||||
self.prompt_router.add_route(route.into_prompt_route());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_prompts(mut self, routes: impl IntoIterator<Item = PromptRoute<S>>) -> Self {
|
||||
for route in routes {
|
||||
self.prompt_router.add_route(route);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Service<RoleServer> for Router<S>
|
||||
|
|
@ -86,6 +105,29 @@ where
|
|||
next_cursor: None,
|
||||
}))
|
||||
}
|
||||
ClientRequest::GetPromptRequest(request) => {
|
||||
if self.prompt_router.has_route(request.params.name.as_ref()) {
|
||||
let prompt_context = crate::handler::server::prompt::PromptContext::new(
|
||||
self.service.as_ref(),
|
||||
request.params.name,
|
||||
request.params.arguments,
|
||||
context,
|
||||
);
|
||||
let result = self.prompt_router.get_prompt(prompt_context).await?;
|
||||
Ok(ServerResult::GetPromptResult(result))
|
||||
} else {
|
||||
self.service
|
||||
.handle_request(ClientRequest::GetPromptRequest(request), context)
|
||||
.await
|
||||
}
|
||||
}
|
||||
ClientRequest::ListPromptsRequest(_) => {
|
||||
let prompts = self.prompt_router.list_all();
|
||||
Ok(ServerResult::ListPromptsResult(ListPromptsResult {
|
||||
prompts,
|
||||
next_cursor: None,
|
||||
}))
|
||||
}
|
||||
rest => self.service.handle_request(rest, context).await,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,213 @@
|
|||
use std::{borrow::Cow, sync::Arc};
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::{
|
||||
handler::server::prompt::{DynGetPromptHandler, GetPromptHandler, PromptContext},
|
||||
model::{GetPromptResult, Prompt},
|
||||
};
|
||||
|
||||
pub struct PromptRoute<S> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub get: Arc<DynGetPromptHandler<S>>,
|
||||
pub attr: crate::model::Prompt,
|
||||
}
|
||||
|
||||
impl<S> std::fmt::Debug for PromptRoute<S> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("PromptRoute")
|
||||
.field("name", &self.attr.name)
|
||||
.field("description", &self.attr.description)
|
||||
.field("arguments", &self.attr.arguments)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Clone for PromptRoute<S> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
get: self.get.clone(),
|
||||
attr: self.attr.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Send + Sync + 'static> PromptRoute<S> {
|
||||
pub fn new<H, A: 'static>(attr: impl Into<Prompt>, handler: H) -> Self
|
||||
where
|
||||
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
|
||||
{
|
||||
Self {
|
||||
get: Arc::new(move |context: PromptContext<S>| {
|
||||
let handler = handler.clone();
|
||||
handler.handle(context)
|
||||
}),
|
||||
attr: attr.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_dyn<H>(attr: impl Into<Prompt>, handler: H) -> Self
|
||||
where
|
||||
H: for<'a> Fn(
|
||||
PromptContext<'a, S>,
|
||||
) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
{
|
||||
Self {
|
||||
get: Arc::new(handler),
|
||||
attr: attr.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn name(&self) -> &str {
|
||||
&self.attr.name
|
||||
}
|
||||
}
|
||||
|
||||
pub trait IntoPromptRoute<S, A> {
|
||||
fn into_prompt_route(self) -> PromptRoute<S>;
|
||||
}
|
||||
|
||||
impl<S, H, A, P> IntoPromptRoute<S, A> for (P, H)
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
A: 'static,
|
||||
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
|
||||
P: Into<Prompt>,
|
||||
{
|
||||
fn into_prompt_route(self) -> PromptRoute<S> {
|
||||
PromptRoute::new(self.0.into(), self.1)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> IntoPromptRoute<S, ()> for PromptRoute<S>
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
fn into_prompt_route(self) -> PromptRoute<S> {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Adapter for functions generated by the #\[prompt\] macro
|
||||
pub struct PromptAttrGenerateFunctionAdapter;
|
||||
|
||||
impl<S, F> IntoPromptRoute<S, PromptAttrGenerateFunctionAdapter> for F
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
F: Fn() -> PromptRoute<S>,
|
||||
{
|
||||
fn into_prompt_route(self) -> PromptRoute<S> {
|
||||
(self)()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct PromptRouter<S> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub map: std::collections::HashMap<Cow<'static, str>, PromptRoute<S>>,
|
||||
}
|
||||
|
||||
impl<S> Default for PromptRouter<S> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
map: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Clone for PromptRouter<S> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
map: self.map.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> IntoIterator for PromptRouter<S> {
|
||||
type Item = PromptRoute<S>;
|
||||
type IntoIter = std::collections::hash_map::IntoValues<Cow<'static, str>, PromptRoute<S>>;
|
||||
|
||||
fn into_iter(self) -> Self::IntoIter {
|
||||
self.map.into_values()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> PromptRouter<S>
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
map: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_route<R, A: 'static>(mut self, route: R) -> Self
|
||||
where
|
||||
R: IntoPromptRoute<S, A>,
|
||||
{
|
||||
self.add_route(route.into_prompt_route());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn add_route(&mut self, item: PromptRoute<S>) {
|
||||
self.map.insert(item.attr.name.clone().into(), item);
|
||||
}
|
||||
|
||||
pub fn merge(&mut self, other: PromptRouter<S>) {
|
||||
for item in other.map.into_values() {
|
||||
self.add_route(item);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn remove_route(&mut self, name: &str) {
|
||||
self.map.remove(name);
|
||||
}
|
||||
|
||||
pub fn has_route(&self, name: &str) -> bool {
|
||||
self.map.contains_key(name)
|
||||
}
|
||||
|
||||
pub async fn get_prompt(
|
||||
&self,
|
||||
context: PromptContext<'_, S>,
|
||||
) -> Result<GetPromptResult, crate::ErrorData> {
|
||||
let item = self.map.get(context.name.as_str()).ok_or_else(|| {
|
||||
crate::ErrorData::invalid_params(
|
||||
format!("prompt '{}' not found", context.name),
|
||||
Some(serde_json::json!({
|
||||
"available_prompts": self.list_all().iter().map(|p| &p.name).collect::<Vec<_>>()
|
||||
})),
|
||||
)
|
||||
})?;
|
||||
(item.get)(context).await
|
||||
}
|
||||
|
||||
pub fn list_all(&self) -> Vec<crate::model::Prompt> {
|
||||
self.map.values().map(|item| item.attr.clone()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> std::ops::Add<PromptRouter<S>> for PromptRouter<S>
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
type Output = Self;
|
||||
|
||||
fn add(mut self, other: PromptRouter<S>) -> Self::Output {
|
||||
self.merge(other);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> std::ops::AddAssign<PromptRouter<S>> for PromptRouter<S>
|
||||
where
|
||||
S: Send + Sync + 'static,
|
||||
{
|
||||
fn add_assign(&mut self, other: PromptRouter<S>) {
|
||||
self.merge(other);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,95 +1,23 @@
|
|||
//! Tool handler traits and types for MCP servers.
|
||||
//!
|
||||
//! This module provides the infrastructure for implementing tools that can be called
|
||||
//! by MCP clients. Tools can return either unstructured content (text, images) or
|
||||
//! structured JSON data with schemas.
|
||||
//!
|
||||
//! # Structured Output
|
||||
//!
|
||||
//! Tools can return structured JSON data using the [`Json`] wrapper type.
|
||||
//! When using `Json<T>`, the framework will:
|
||||
//! - Automatically generate a JSON schema for the output type
|
||||
//! - Validate the output against the schema
|
||||
//! - Return the data in the `structured_content` field of [`CallToolResult`]
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use rmcp::{tool, Json};
|
||||
//! use schemars::JsonSchema;
|
||||
//! use serde::{Serialize, Deserialize};
|
||||
//!
|
||||
//! #[derive(Serialize, Deserialize, JsonSchema)]
|
||||
//! struct AnalysisResult {
|
||||
//! score: f64,
|
||||
//! summary: String,
|
||||
//! }
|
||||
//!
|
||||
//! #[tool(name = "analyze")]
|
||||
//! async fn analyze(&self, text: String) -> Result<Json<AnalysisResult>, String> {
|
||||
//! Ok(Json(AnalysisResult {
|
||||
//! score: 0.95,
|
||||
//! summary: "Positive sentiment".to_string(),
|
||||
//! }))
|
||||
//! }
|
||||
//! ```
|
||||
|
||||
use std::{
|
||||
any::TypeId, borrow::Cow, collections::HashMap, future::Ready, marker::PhantomData, sync::Arc,
|
||||
borrow::Cow,
|
||||
future::{Future, Ready},
|
||||
marker::PhantomData,
|
||||
};
|
||||
|
||||
use futures::future::{BoxFuture, FutureExt};
|
||||
use schemars::{JsonSchema, transform::AddNullable};
|
||||
use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
pub use super::router::tool::{ToolRoute, ToolRouter};
|
||||
use super::common::{AsRequestContext, FromContextPart};
|
||||
pub use super::{
|
||||
common::{Extension, RequestId, cached_schema_for_type, schema_for_type},
|
||||
router::tool::{ToolRoute, ToolRouter},
|
||||
};
|
||||
use crate::{
|
||||
RoleServer,
|
||||
handler::server::wrapper::Json,
|
||||
handler::server::wrapper::Parameters,
|
||||
model::{CallToolRequestParam, CallToolResult, IntoContents, JsonObject},
|
||||
schemars::generate::SchemaSettings,
|
||||
service::RequestContext,
|
||||
};
|
||||
/// A shortcut for generating a JSON schema for a type.
|
||||
pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
|
||||
// explicitly to align json schema version to official specifications.
|
||||
// https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.json
|
||||
// TODO: update to 2020-12 waiting for the mcp spec update
|
||||
let mut settings = SchemaSettings::draft07();
|
||||
settings.transforms = vec![Box::new(AddNullable::default())];
|
||||
let generator = settings.into_generator();
|
||||
let schema = generator.into_root_schema_for::<T>();
|
||||
let object = serde_json::to_value(schema).expect("failed to serialize schema");
|
||||
match object {
|
||||
serde_json::Value::Object(object) => object,
|
||||
_ => panic!("unexpected schema value"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Call [`schema_for_type`] with a cache
|
||||
pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
|
||||
thread_local! {
|
||||
static CACHE_FOR_TYPE: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
|
||||
};
|
||||
CACHE_FOR_TYPE.with(|cache| {
|
||||
if let Some(x) = cache
|
||||
.read()
|
||||
.expect("schema cache lock poisoned")
|
||||
.get(&TypeId::of::<T>())
|
||||
{
|
||||
x.clone()
|
||||
} else {
|
||||
let schema = schema_for_type::<T>();
|
||||
let schema = Arc::new(schema);
|
||||
cache
|
||||
.write()
|
||||
.expect("schema cache lock poisoned")
|
||||
.insert(TypeId::of::<T>(), schema.clone());
|
||||
schema
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Deserialize a JSON object into a type
|
||||
pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, crate::ErrorData> {
|
||||
|
|
@ -128,32 +56,18 @@ impl<'s, S> ToolCallContext<'s, S> {
|
|||
}
|
||||
}
|
||||
|
||||
pub trait FromToolCallContextPart<S>: Sized {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData>;
|
||||
impl<S> AsRequestContext for ToolCallContext<'_, S> {
|
||||
fn as_request_context(&self) -> &RequestContext<RoleServer> {
|
||||
&self.request_context
|
||||
}
|
||||
|
||||
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
|
||||
&mut self.request_context
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for converting tool return values into [`CallToolResult`].
|
||||
///
|
||||
/// This trait is automatically implemented for:
|
||||
/// - Types implementing [`IntoContents`] (returns unstructured content)
|
||||
/// - `Result<T, E>` where both `T` and `E` implement [`IntoContents`]
|
||||
/// - [`Json<T>`](crate::handler::server::wrapper::Json) where `T` implements [`Serialize`] (returns structured content)
|
||||
/// - `Result<Json<T>, E>` for structured results with errors
|
||||
///
|
||||
/// The `#[tool]` macro uses this trait to convert tool function return values
|
||||
/// into the appropriate [`CallToolResult`] format.
|
||||
pub trait IntoCallToolResult {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData>;
|
||||
|
||||
/// Returns the output schema for this type, if any.
|
||||
///
|
||||
/// This is used by the macro to automatically generate output schemas
|
||||
/// for tool functions that return structured data.
|
||||
fn output_schema() -> Option<Arc<JsonObject>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: IntoContents> IntoCallToolResult for T {
|
||||
|
|
@ -180,40 +94,6 @@ impl<T: IntoCallToolResult> IntoCallToolResult for Result<T, crate::ErrorData> {
|
|||
}
|
||||
}
|
||||
|
||||
// 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> {
|
||||
let value = serde_json::to_value(self.0).map_err(|e| {
|
||||
crate::ErrorData::internal_error(
|
||||
format!("Failed to serialize structured content: {}", e),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(CallToolResult::structured(value))
|
||||
}
|
||||
|
||||
fn output_schema() -> Option<Arc<JsonObject>> {
|
||||
Some(cached_schema_for_type::<T>())
|
||||
}
|
||||
}
|
||||
|
||||
// Implementation for Result<Json<T>, E>
|
||||
impl<T: Serialize + JsonSchema + 'static, E: IntoContents> IntoCallToolResult
|
||||
for Result<Json<T>, E>
|
||||
{
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
|
||||
match self {
|
||||
Ok(value) => value.into_call_tool_result(),
|
||||
Err(error) => Ok(CallToolResult::error(error.into_contents())),
|
||||
}
|
||||
}
|
||||
|
||||
fn output_schema() -> Option<Arc<JsonObject>> {
|
||||
Json::<T>::output_schema()
|
||||
}
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
#[project = IntoCallToolResultFutProj]
|
||||
pub enum IntoCallToolResultFut<F, R> {
|
||||
|
|
@ -266,47 +146,21 @@ pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFu
|
|||
+ Send
|
||||
+ Sync;
|
||||
|
||||
/// Parameter Extractor
|
||||
///
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct Parameters<P>(pub P);
|
||||
|
||||
impl<P: JsonSchema> JsonSchema for Parameters<P> {
|
||||
fn schema_name() -> Cow<'static, str> {
|
||||
P::schema_name()
|
||||
}
|
||||
|
||||
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
|
||||
P::json_schema(generator)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for CancellationToken {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.request_context.ct.clone())
|
||||
}
|
||||
}
|
||||
|
||||
// Tool-specific extractor for tool name
|
||||
pub struct ToolName(pub Cow<'static, str>);
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for ToolName {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
impl<S> FromContextPart<ToolCallContext<'_, S>> for ToolName {
|
||||
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
|
||||
Ok(Self(context.name.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, P> FromToolCallContextPart<S> for Parameters<P>
|
||||
// Special implementation for Parameters that handles tool arguments
|
||||
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
|
||||
where
|
||||
P: DeserializeOwned,
|
||||
{
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
|
||||
let arguments = context.arguments.take().unwrap_or_default();
|
||||
let value: P =
|
||||
serde_json::from_value(serde_json::Value::Object(arguments)).map_err(|e| {
|
||||
|
|
@ -319,84 +173,14 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for JsonObject {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
// Special implementation for JsonObject that takes tool arguments
|
||||
impl<S> FromContextPart<ToolCallContext<'_, S>> for JsonObject {
|
||||
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
|
||||
let object = context.arguments.take().unwrap_or_default();
|
||||
Ok(object)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for crate::model::Extensions {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
let extensions = context.request_context.extensions.clone();
|
||||
Ok(extensions)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Extension<T>(pub T);
|
||||
|
||||
impl<S, T> FromToolCallContextPart<S> for Extension<T>
|
||||
where
|
||||
T: Send + Sync + 'static + Clone,
|
||||
{
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
let extension = context
|
||||
.request_context
|
||||
.extensions
|
||||
.get::<T>()
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
crate::ErrorData::invalid_params(
|
||||
format!("missing extension {}", std::any::type_name::<T>()),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
Ok(Extension(extension))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for crate::Peer<RoleServer> {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
let peer = context.request_context.peer.clone();
|
||||
Ok(peer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for crate::model::Meta {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
let mut meta = crate::model::Meta::default();
|
||||
std::mem::swap(&mut meta, &mut context.request_context.meta);
|
||||
Ok(meta)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RequestId(pub crate::model::RequestId);
|
||||
impl<S> FromToolCallContextPart<S> for RequestId {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
Ok(RequestId(context.request_context.id.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromToolCallContextPart<S> for RequestContext<RoleServer> {
|
||||
fn from_tool_call_context_part(
|
||||
context: &mut ToolCallContext<S>,
|
||||
) -> Result<Self, crate::ErrorData> {
|
||||
Ok(context.request_context.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'s, S> ToolCallContext<'s, S> {
|
||||
pub fn invoke<H, A>(self, h: H) -> BoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
|
||||
where
|
||||
|
|
@ -428,7 +212,7 @@ macro_rules! impl_for {
|
|||
impl<$($Tn,)* S, F, R> CallToolHandler<S, AsyncMethodAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<S> ,
|
||||
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
|
||||
)*
|
||||
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R>,
|
||||
|
||||
|
|
@ -443,7 +227,7 @@ macro_rules! impl_for {
|
|||
mut context: ToolCallContext<'_, S>,
|
||||
) -> BoxFuture<'_, Result<CallToolResult, crate::ErrorData>>{
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(&mut context);
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
|
|
@ -461,7 +245,7 @@ macro_rules! impl_for {
|
|||
impl<$($Tn,)* S, F, Fut, R> CallToolHandler<S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<S> ,
|
||||
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> Fut + Send + ,
|
||||
Fut: Future<Output = R> + Send + 'static,
|
||||
|
|
@ -474,7 +258,7 @@ macro_rules! impl_for {
|
|||
mut context: ToolCallContext<S>,
|
||||
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>>{
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(&mut context);
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
|
|
@ -491,7 +275,7 @@ macro_rules! impl_for {
|
|||
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncMethodAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<S> + ,
|
||||
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
|
||||
)*
|
||||
F: FnOnce(&S, $($Tn,)*) -> R + Send + ,
|
||||
R: IntoCallToolResult + Send + ,
|
||||
|
|
@ -503,7 +287,7 @@ macro_rules! impl_for {
|
|||
mut context: ToolCallContext<S>,
|
||||
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(&mut context);
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
|
|
@ -516,7 +300,7 @@ macro_rules! impl_for {
|
|||
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<S> + ,
|
||||
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> R + Send + ,
|
||||
R: IntoCallToolResult + Send + ,
|
||||
|
|
@ -528,7 +312,7 @@ macro_rules! impl_for {
|
|||
mut context: ToolCallContext<S>,
|
||||
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(&mut context);
|
||||
let result = $Tn::from_context_part(&mut context);
|
||||
let $Tn = match result {
|
||||
Ok(value) => value,
|
||||
Err(e) => return std::future::ready(Err(e)).boxed(),
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
mod json;
|
||||
mod parameters;
|
||||
pub use json::*;
|
||||
pub use parameters::*;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,12 @@
|
|||
use std::borrow::Cow;
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::{
|
||||
handler::server::tool::IntoCallToolResult,
|
||||
model::{CallToolResult, IntoContents},
|
||||
};
|
||||
|
||||
/// Json wrapper for structured output
|
||||
///
|
||||
|
|
@ -20,3 +26,29 @@ impl<T: JsonSchema> JsonSchema for Json<T> {
|
|||
T::json_schema(generator)
|
||||
}
|
||||
}
|
||||
|
||||
// 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> {
|
||||
let value = serde_json::to_value(self.0).map_err(|e| {
|
||||
crate::ErrorData::internal_error(
|
||||
format!("Failed to serialize structured content: {}", e),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(CallToolResult::structured(value))
|
||||
}
|
||||
}
|
||||
|
||||
// Implementation for Result<Json<T>, E>
|
||||
impl<T: Serialize + JsonSchema + 'static, E: IntoContents> IntoCallToolResult
|
||||
for Result<Json<T>, E>
|
||||
{
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
|
||||
match self {
|
||||
Ok(value) => value.into_call_tool_result(),
|
||||
Err(error) => Ok(CallToolResult::error(error.into_contents())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
55
crates/rmcp/src/handler/server/wrapper/parameters.rs
Normal file
55
crates/rmcp/src/handler/server/wrapper/parameters.rs
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
use schemars::JsonSchema;
|
||||
|
||||
/// Parameter extractor for tools and prompts
|
||||
///
|
||||
/// When used in tool and prompt handlers, this wrapper extracts and deserializes
|
||||
/// parameters from the incoming request. The framework will automatically parse
|
||||
/// the JSON arguments from tool calls or prompt arguments and deserialize them
|
||||
/// into the specified type `P`.
|
||||
///
|
||||
/// The `#[serde(transparent)]` attribute ensures that the wrapper doesn't add
|
||||
/// an extra layer in the JSON structure - it directly delegates serialization
|
||||
/// and deserialization to the inner type `P`.
|
||||
///
|
||||
/// # Usage
|
||||
///
|
||||
/// Use `Parameters<T>` as a parameter in your tool or prompt handler functions:
|
||||
///
|
||||
/// ```rust
|
||||
/// # use rmcp::handler::server::wrapper::Parameters;
|
||||
/// # use schemars::JsonSchema;
|
||||
/// # use serde::{Deserialize, Serialize};
|
||||
/// #[derive(Deserialize, JsonSchema)]
|
||||
/// struct CalculationRequest {
|
||||
/// operation: String,
|
||||
/// a: f64,
|
||||
/// b: f64,
|
||||
/// }
|
||||
///
|
||||
/// // In a tool handler
|
||||
/// async fn calculate(params: Parameters<CalculationRequest>) -> Result<String, String> {
|
||||
/// let request = params.0; // Extract the inner value
|
||||
/// match request.operation.as_str() {
|
||||
/// "add" => Ok((request.a + request.b).to_string()),
|
||||
/// _ => Err("Unknown operation".to_string()),
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// The framework handles the extraction automatically:
|
||||
/// - For tools: Parses the `arguments` field from tool call requests
|
||||
/// - For prompts: Parses the `arguments` field from prompt requests
|
||||
/// - Returns appropriate error responses if deserialization fails
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct Parameters<P>(pub P);
|
||||
|
||||
impl<P: JsonSchema> JsonSchema for Parameters<P> {
|
||||
fn schema_name() -> std::borrow::Cow<'static, str> {
|
||||
P::schema_name()
|
||||
}
|
||||
|
||||
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
|
||||
P::json_schema(generator)
|
||||
}
|
||||
}
|
||||
|
|
@ -53,7 +53,7 @@
|
|||
//! Tools can also return structured JSON data with schemas. Use the [`Json`] wrapper:
|
||||
//!
|
||||
//! ```rust
|
||||
//! # use rmcp::{tool, tool_router, handler::server::tool::{ToolRouter, Parameters}, Json};
|
||||
//! # use rmcp::{tool, tool_router, handler::server::{tool::ToolRouter, wrapper::Parameters}, Json};
|
||||
//! # use schemars::JsonSchema;
|
||||
//! # use serde::{Serialize, Deserialize};
|
||||
//! #
|
||||
|
|
|
|||
|
|
@ -162,12 +162,12 @@ pub trait DynService<R: ServiceRole>: Send + Sync {
|
|||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> BoxFuture<Result<R::Resp, McpError>>;
|
||||
) -> BoxFuture<'_, Result<R::Resp, McpError>>;
|
||||
fn handle_notification(
|
||||
&self,
|
||||
notification: R::PeerNot,
|
||||
context: NotificationContext<R>,
|
||||
) -> BoxFuture<Result<(), McpError>>;
|
||||
) -> BoxFuture<'_, Result<(), McpError>>;
|
||||
fn get_info(&self) -> R::Info;
|
||||
}
|
||||
|
||||
|
|
@ -176,14 +176,14 @@ impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
|
|||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> BoxFuture<Result<R::Resp, McpError>> {
|
||||
) -> BoxFuture<'_, Result<R::Resp, McpError>> {
|
||||
Box::pin(self.handle_request(request, context))
|
||||
}
|
||||
fn handle_notification(
|
||||
&self,
|
||||
notification: R::PeerNot,
|
||||
context: NotificationContext<R>,
|
||||
) -> BoxFuture<Result<(), McpError>> {
|
||||
) -> BoxFuture<'_, Result<(), McpError>> {
|
||||
Box::pin(self.handle_notification(notification, context))
|
||||
}
|
||||
fn get_info(&self) -> R::Info {
|
||||
|
|
|
|||
Loading…
Reference in a new issue