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

* 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
9 changed files with 158 additions and 65 deletions

View file

@ -169,7 +169,7 @@ async fn main() -> anyhow::Result<()> {
}
```
The generated tool `inputSchema` is derived from the fields of `T`. The type name and documentation on `T` are ignored; only field names, field types, and field documentation are used.
The generated tool `inputSchema` and `outputSchema` are derived from the fields of `T`. The type name and documentation on `T` are ignored; only field names, field types, and field documentation are used.
When you need custom server metadata or multiple capabilities (tools + prompts), use explicit `#[tool_handler]`:

View file

@ -232,6 +232,13 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
// if found, use the Parameters schema
syn::parse2::<Expr>(quote! {
rmcp::handler::server::common::schema_for_input::<#params_ty>()
.unwrap_or_else(|e| {
panic!(
"Invalid input schema for `{}`: {}",
std::any::type_name::<#params_ty>(),
e
)
})
})?
} else {
// if not found, use a default empty JSON schema object

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"));
}
}

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 {

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.

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
}

View file

@ -1,6 +1,7 @@
#![cfg(all(feature = "server", feature = "macros", not(feature = "local")))]
use rmcp::{
Json,
handler::server::wrapper::Parameters,
model::{ListToolsResult, NumberOrString, ServerJsonRpcMessage, ServerResult},
};
@ -14,10 +15,17 @@ struct AddRequest {
b: f64,
}
/// Result of adding two numbers.
#[derive(Debug, serde::Serialize, schemars::JsonSchema)]
struct AddResult {
/// The sum of the two numbers.
sum: f64,
}
/// Add two numbers.
#[rmcp::tool]
fn add(Parameters(AddRequest { a, b }): Parameters<AddRequest>) -> String {
(a + b).to_string()
fn add(Parameters(AddRequest { a, b }): Parameters<AddRequest>) -> Json<AddResult> {
Json(AddResult { sum: a + b })
}
#[test]
@ -27,7 +35,7 @@ fn list_tools_result_matches_expected_json() {
let expected: serde_json::Value =
serde_json::from_slice(&expected_json).expect("invalid expected JSON fixture");
assert_eq!(add(Parameters(AddRequest { a: 1.0, b: 2.0 })), "3");
assert_eq!(add(Parameters(AddRequest { a: 1.0, b: 2.0 })).0.sum, 3.0);
let result = ListToolsResult::with_all_items(vec![add_tool_attr()]);
let response = ServerJsonRpcMessage::response(

View file

@ -23,6 +23,20 @@
"a",
"b"
]
},
"outputSchema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"sum": {
"description": "The sum of the two numbers.",
"format": "double",
"type": "number"
}
},
"required": [
"sum"
]
}
}
]

View file

@ -28,6 +28,16 @@ pub struct UserInfo {
pub age: u32,
}
#[derive(Serialize, Deserialize, JsonSchema)]
pub struct GreetingRequest {
pub name: String,
}
#[derive(Serialize, Deserialize, JsonSchema)]
pub struct GetUserRequest {
pub user_id: String,
}
#[tool_handler(router = self.tool_router)]
impl ServerHandler for TestServer {}
@ -64,14 +74,17 @@ impl TestServer {
/// Tool that returns regular string output
#[tool(name = "get-greeting", description = "Get a greeting")]
pub async fn get_greeting(&self, name: Parameters<String>) -> String {
format!("Hello, {}!", name.0)
pub async fn get_greeting(&self, params: Parameters<GreetingRequest>) -> String {
format!("Hello, {}!", params.0.name)
}
/// Tool that returns structured user info
#[tool(name = "get-user", description = "Get user info")]
pub async fn get_user(&self, user_id: Parameters<String>) -> Result<Json<UserInfo>, String> {
if user_id.0 == "123" {
pub async fn get_user(
&self,
params: Parameters<GetUserRequest>,
) -> Result<Json<UserInfo>, String> {
if params.0.user_id == "123" {
Ok(Json(UserInfo {
name: "Alice".to_string(),
age: 30,