Implements outputSchema validation (#566)

* feat: implement output schema validation

* fix: calculator example comply MCP spec

* refactor: merge cached_schema_for_output into schema_for_output

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::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::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,49 @@ 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 + 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| {
+        // Try to get from cache first
+        if let Some(result) = cache
+            .read()
+            .expect("output schema cache lock poisoned")
+            .get(&TypeId::of::<T>())
+        {
+            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(Arc::new(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
+            )),
+        };
+
+        // Cache the result (both success and error cases)
+        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 +186,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,7 @@ 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_type, schema_for_output, 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::schema_for_output::<T>()
+            .unwrap_or_else(|e| panic!("Invalid output schema for tool '{}': {}", self.name, e));
+        self.output_schema = Some(schema);
         self
     }
 

examples/servers/Cargo.toml

@@ -96,3 +96,7 @@ path = "src/simple_auth_streamhttp.rs"
 [[example]]
 name = "servers_complex_auth_streamhttp"
 path = "src/complex_auth_streamhttp.rs"
+
+[[example]]
+name = "servers_calculator_stdio"
+path = "src/calculator_stdio.rs"

examples/servers/src/calculator_stdio.rs

@@ -0,0 +1,26 @@
+use anyhow::Result;
+use common::calculator::Calculator;
+use rmcp::{ServiceExt, transport::stdio};
+use tracing_subscriber::{self, EnvFilter};
+mod common;
+
+/// npx @modelcontextprotocol/inspector cargo run -p mcp-server-examples --example servers_calculator_stdio
+#[tokio::main]
+async fn main() -> Result<()> {
+    // Initialize the tracing subscriber with file and stdout logging
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env().add_directive(tracing::Level::DEBUG.into()))
+        .with_writer(std::io::stderr)
+        .with_ansi(false)
+        .init();
+
+    tracing::info!("Starting Calculator MCP server");
+
+    // Create an instance of our calculator router
+    let service = Calculator::new().serve(stdio()).await.inspect_err(|e| {
+        tracing::error!("serving error: {:?}", e);
+    })?;
+
+    service.waiting().await?;
+    Ok(())
+}

examples/servers/src/common/calculator.rs

@@ -2,10 +2,7 @@
 
 use rmcp::{
     ServerHandler,
-    handler::server::{
-        router::tool::ToolRouter,
-        wrapper::{Json, Parameters},
-    },
+    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
     model::{ServerCapabilities, ServerInfo},
     schemars, tool, tool_handler, tool_router,
 };
@@ -44,8 +41,8 @@ impl Calculator {
     }
 
     #[tool(description = "Calculate the difference of two numbers")]
-    fn sub(&self, Parameters(SubRequest { a, b }): Parameters<SubRequest>) -> Json<i32> {
-        Json(a - b)
+    fn sub(&self, Parameters(SubRequest { a, b }): Parameters<SubRequest>) -> String {
+        (a - b).to_string()
     }
 }