feat: implement output schema validation

Author: DaleSeo

crates/rmcp-macros/src/tool.rs

@@ -27,7 +27,14 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
     // First, try direct Json<T>
     if let Some(inner_type) = extract_json_inner_type(ret_type) {
         return syn::parse2::<Expr>(quote! {
-            rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
+            rmcp::handler::server::tool::cached_schema_for_output::<#inner_type>()
+                .unwrap_or_else(|e| {
+                    panic!(
+                        "Invalid output schema for Json<{}>: {}",
+                        std::any::type_name::<#inner_type>(),
+                        e
+                    )
+                })
         })
         .ok();
     }
@@ -57,7 +64,14 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
     let inner_type = extract_json_inner_type(ok_type)?;
 
     syn::parse2::<Expr>(quote! {
-        rmcp::handler::server::tool::cached_schema_for_type::<#inner_type>()
+        rmcp::handler::server::tool::cached_schema_for_output::<#inner_type>()
+            .unwrap_or_else(|e| {
+                panic!(
+                    "Invalid output schema for Result<Json<{}>, E>: {}",
+                    std::any::type_name::<#inner_type>(),
+                    e
+                )
+            })
     })
     .ok()
 }

crates/rmcp/src/handler/server/common.rs

@@ -50,6 +50,50 @@ pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject
     })
 }
 
+/// Generate and validate a JSON schema for outputSchema (must have root type "object").
+pub fn schema_for_output<T: JsonSchema>() -> Result<JsonObject, String> {
+    let schema = schema_for_type::<T>();
+
+    match schema.get("type") {
+        Some(serde_json::Value::String(t)) if t == "object" => Ok(schema),
+        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
+        )),
+    }
+}
+
+/// Call [`schema_for_output`] with a cache.
+pub fn cached_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();
+    };
+    CACHE_FOR_OUTPUT.with(|cache| {
+        if let Some(result) = cache
+            .read()
+            .expect("output schema cache lock poisoned")
+            .get(&TypeId::of::<T>())
+        {
+            result.clone()
+        } else {
+            let result = schema_for_output::<T>().map(Arc::new);
+            cache
+                .write()
+                .expect("output schema cache lock poisoned")
+                .insert(TypeId::of::<T>(), result.clone());
+            result
+        }
+    })
+}
+
 /// 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>;
@@ -143,3 +187,25 @@ pub trait AsRequestContext {
     fn as_request_context(&self) -> &RequestContext<RoleServer>;
     fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer>;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(serde::Serialize, serde::Deserialize, JsonSchema)]
+    struct TestObject {
+        value: i32,
+    }
+
+    #[test]
+    fn test_schema_for_output_rejects_primitive() {
+        let result = schema_for_output::<i32>();
+        assert!(result.is_err(),);
+    }
+
+    #[test]
+    fn test_schema_for_output_accepts_object() {
+        let result = schema_for_output::<TestObject>();
+        assert!(result.is_ok(),);
+    }
+}

crates/rmcp/src/handler/server/tool.rs

@@ -9,7 +9,9 @@ use serde::de::DeserializeOwned;
 
 use super::common::{AsRequestContext, FromContextPart};
 pub use super::{
-    common::{Extension, RequestId, cached_schema_for_type, schema_for_type},
+    common::{
+        Extension, RequestId, cached_schema_for_output, cached_schema_for_type, schema_for_type,
+    },
     router::tool::{ToolRoute, ToolRouter},
 };
 use crate::{

crates/rmcp/src/model/tool.rs

@@ -165,8 +165,14 @@ impl Tool {
     }
 
     /// Set the output schema using a type that implements JsonSchema
+    ///
+    /// # Panics
+    ///
+    /// Panics if the generated schema does not have root type "object" as required by MCP specification.
     pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self {
-        self.output_schema = Some(crate::handler::server::tool::cached_schema_for_type::<T>());
+        let schema = crate::handler::server::tool::cached_schema_for_output::<T>()
+            .unwrap_or_else(|e| panic!("Invalid output schema for tool '{}': {}", self.name, e));
+        self.output_schema = Some(schema);
         self
     }