modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

fix: strip and validate tool outputSchema and inputSchema (#860)

Browse source 
* fix: remove unnecessary fields from tools' outputSchema

* fix: validate input schema root type per MCP spec
This commit is contained in:
Dale Seo 2026-06-02 12:01:33 -04:00 • committed by GitHub
commit 254f04a764
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 158 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

142
crates/rmcp/src/handler/server/common.rs
View file

@ -1,6 +1,10 @@
//! Common utilities shared between tool and prompt handlers
use std::{any::TypeId, collections::HashMap, sync::Arc};
use std::{
any::TypeId,
collections::HashMap,
sync::{Arc, LazyLock},
};
use schemars::JsonSchema;
@ -30,12 +34,10 @@ pub fn schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
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");
let object = match object {
serde_json::Value::Object(object) => object,
_ => panic!(
"Schema serialization produced non-object value: expected JSON object but got {:?}",
object
),
let serde_json::Value::Object(object) = object else {
panic!(
"Schema serialization produced non-object value: expected JSON object but got {object:?}"
);
};
let schema = Arc::new(object);
cache
@ -48,51 +50,63 @@ pub fn schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
})
}
/// Generate a JSON schema for inputSchema (does not need "title" or "description" fields for the top-level object)
pub fn schema_for_input<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
/// Validate that the schema root is `type: "object"` (per MCP spec) and strip top-level
/// `title`/`description` (the wrapper type name and doc, which are noise to the LLM).
fn validate_and_strip(raw: &Arc<JsonObject>, purpose: &str) -> Result<Arc<JsonObject>, String> {
match raw.get("type") {
Some(serde_json::Value::String(t)) if t == "object" => {
let mut object = raw.as_ref().clone();
object.remove("title");
object.remove("description");
Ok(Arc::new(object))
}
Some(serde_json::Value::String(t)) => Err(format!(
"MCP specification requires tool {purpose} to have root type 'object', but found '{t}'."
)),
None => Err(format!(
"Schema is missing 'type' field. MCP specification requires {purpose} to have root type 'object'."
)),
Some(other) => Err(format!(
"Schema 'type' field has unexpected format: {other:?}. Expected \"object\"."
)),
}
}
/// Generate, validate, and strip a JSON schema for inputSchema (must have root type "object";
/// top-level "title" and "description" are removed).
pub fn schema_for_input<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> {
thread_local! {
static CACHE_FOR_INPUT: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
static CACHE_FOR_INPUT: std::sync::RwLock<HashMap<TypeId, Result<Arc<JsonObject>, String>>> = Default::default();
};
CACHE_FOR_INPUT.with(|cache| {
if let Some(schema) = cache
if let Some(result) = cache
.read()
.expect("input schema cache lock poisoned")
.get(&TypeId::of::<T>())
{
schema.clone()
} else {
let mut schema = schema_for_type::<T>().as_ref().clone();
// Remove unnecessary top-level fields
schema.remove("title");
schema.remove("description");
let schema = Arc::new(schema);
cache
.write()
.expect("input schema cache lock poisoned")
.insert(TypeId::of::<T>(), schema.clone());
schema
return result.clone();
}
let result = validate_and_strip(&schema_for_type::<T>(), "inputSchema");
cache
.write()
.expect("input schema cache lock poisoned")
.insert(TypeId::of::<T>(), result.clone());
result
})
}
// TODO: should be updated according to the new specifications
/// Schema used when input is empty.
pub fn schema_for_empty_input() -> Arc<JsonObject> {
std::sync::Arc::new(
serde_json::json!({
"type": "object",
"properties": {}
})
.as_object()
.unwrap()
.clone(),
)
static EMPTY: LazyLock<Arc<JsonObject>> = LazyLock::new(|| {
let mut object = JsonObject::new();
object.insert("type".into(), serde_json::json!("object"));
object.insert("properties".into(), serde_json::json!({}));
Arc::new(object)
});
EMPTY.clone()
}
/// Generate and validate a JSON schema for outputSchema (must have root type "object").
/// Generate a JSON schema for outputSchema (must have root type "object"; top-level "title" and "description" are removed)
pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> {
thread_local! {
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Result<Arc<JsonObject>, String>>> = Default::default();
@ -108,22 +122,8 @@ pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObje
return result.clone();
}
// Generate and validate schema
let schema = schema_for_type::<T>();
let result = match schema.get("type") {
Some(serde_json::Value::String(t)) if t == "object" => Ok(schema.clone()),
Some(serde_json::Value::String(t)) => Err(format!(
"MCP specification requires tool outputSchema to have root type 'object', but found '{}'.",
t
)),
None => Err(
"Schema is missing 'type' field. MCP specification requires outputSchema to have root type 'object'.".to_string()
),
Some(other) => Err(format!(
"Schema 'type' field has unexpected format: {:?}. Expected \"object\".",
other
)),
};
// Generate, validate, and strip unnecessary top-level fields
let result = validate_and_strip(&schema_for_type::<T>(), "outputSchema");
// Cache the result (both success and error cases)
cache
@ -316,4 +316,40 @@ mod tests {
let result = schema_for_output::<TestObject>();
assert!(result.is_ok(),);
}
#[test]
fn test_schema_for_output_strips_top_level_title() {
let schema = schema_for_output::<TestObject>().unwrap();
assert!(!schema.contains_key("title"));
}
#[test]
fn test_schema_for_output_strips_top_level_description() {
let schema = schema_for_output::<TestObject>().unwrap();
assert!(!schema.contains_key("description"));
}
#[test]
fn test_schema_for_input_rejects_primitive() {
let result = schema_for_input::<i32>();
assert!(result.is_err());
}
#[test]
fn test_schema_for_input_accepts_object() {
let result = schema_for_input::<TestObject>();
assert!(result.is_ok());
}
#[test]
fn test_schema_for_input_strips_top_level_title() {
let schema = schema_for_input::<TestObject>().unwrap();
assert!(!schema.contains_key("title"));
}
#[test]
fn test_schema_for_input_strips_top_level_description() {
let schema = schema_for_input::<TestObject>().unwrap();
assert!(!schema.contains_key("description"));
}
}

11
crates/rmcp/src/handler/server/router/tool.rs
View file

@ -250,7 +250,9 @@ where
attr: Tool::new(
name.into(),
"",
schema_for_input::<crate::model::JsonObject>(),
schema_for_input::<crate::model::JsonObject>().unwrap_or_else(|e| {
panic!("Invalid input schema for JsonObject: {e}");
}),
),
call: self,
_marker: std::marker::PhantomData,
@ -287,7 +289,12 @@ where
self
}
pub fn parameters<T: JsonSchema + 'static>(mut self) -> Self {
self.attr.input_schema = schema_for_input::<T>();
self.attr.input_schema = schema_for_input::<T>().unwrap_or_else(|e| {
panic!(
"Invalid input schema for `{}`: {e}",
std::any::type_name::<T>()
)
});
self
}
pub fn parameters_value(mut self, schema: serde_json::Value) -> Self {

9
crates/rmcp/src/handler/server/router/tool/tool_traits.rs
View file

@ -49,7 +49,14 @@ pub trait ToolBase {
/// If the tool does not have any parameters, you should override this methods to return [`None`],
/// and when invoked, the parameter will get default values.
fn input_schema() -> Option<Arc<JsonObject>> {
Some(schema_for_input::<Parameters<Self::Parameter>>())
Some(
schema_for_input::<Parameters<Self::Parameter>>().unwrap_or_else(|e| {
panic!(
"Invalid input schema for ToolBase::Parameter type `{0}`: {e}",
std::any::type_name::<Self::Parameter>(),
);
}),
)
}
/// Json schema for tool output.

3
crates/rmcp/src/model/tool.rs
View file

@ -330,7 +330,8 @@ impl Tool {
/// Set the input schema using a type that implements JsonSchema
#[cfg(feature = "server")]
pub fn with_input_schema<T: JsonSchema + 'static>(mut self) -> Self {
self.input_schema = crate::handler::server::tool::schema_for_input::<T>();
self.input_schema = crate::handler::server::tool::schema_for_input::<T>()
.unwrap_or_else(|e| panic!("Invalid input schema for tool '{}': {}", self.name, e));
self
}