feat: add builder methods to Tool struct for setting schemas

- Add with_output_schema<T>() method to set output schema from type
- Add with_input_schema<T>() method to set input schema from type
- Both methods use cached_schema_for_type internally
- Add comprehensive tests for builder methods

Author: JMLX42

crates/rmcp/src/model/tool.rs

@@ -1,5 +1,6 @@
 use std::{borrow::Cow, sync::Arc};
 
+use schemars::JsonSchema;
 /// Tools represent a routine that a server can execute
 /// Tool calls represent requests from the client to execute one
 use serde::{Deserialize, Serialize};
@@ -151,6 +152,18 @@ impl Tool {
         }
     }
 
+    /// Set the output schema using a type that implements JsonSchema
+    pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self {
+        self.output_schema = Some(crate::handler::server::tool::cached_schema_for_type::<T>());
+        self
+    }
+
+    /// Set the input schema using a type that implements JsonSchema
+    pub fn with_input_schema<T: JsonSchema + 'static>(mut self) -> Self {
+        self.input_schema = crate::handler::server::tool::cached_schema_for_type::<T>();
+        self
+    }
+
     /// Get the schema as json value
     pub fn schema_as_json_value(&self) -> Value {
         Value::Object(self.input_schema.as_ref().clone())

crates/rmcp/tests/test_json_schema_detection.rs

@@ -1,8 +1,6 @@
 //cargo test --test test_json_schema_detection --features "client server macros"
 use rmcp::{
-    Json, ServerHandler,
-    handler::server::router::tool::ToolRouter,
-    tool, tool_handler, tool_router,
+    Json, ServerHandler, handler::server::router::tool::ToolRouter, tool, tool_handler, tool_router,
 };
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -37,7 +35,9 @@ impl TestServer {
     /// Tool that returns Json<T> - should have output schema
     #[tool(name = "with-json")]
     pub async fn with_json(&self) -> Result<Json<TestData>, String> {
-        Ok(Json(TestData { value: "test".to_string() }))
+        Ok(Json(TestData {
+            value: "test".to_string(),
+        }))
     }
 
     /// Tool that returns regular type - should NOT have output schema
@@ -49,7 +49,9 @@ impl TestServer {
     /// Tool that returns Result with inner Json - should have output schema  
     #[tool(name = "result-with-json")]
     pub async fn result_with_json(&self) -> Result<Json<TestData>, rmcp::ErrorData> {
-        Ok(Json(TestData { value: "test".to_string() }))
+        Ok(Json(TestData {
+            value: "test".to_string(),
+        }))
     }
 
     /// Tool with explicit output_schema attribute - should have output schema
@@ -66,7 +68,10 @@ async fn test_json_type_generates_schema() {
 
     // Find the with-json tool
     let json_tool = tools.iter().find(|t| t.name == "with-json").unwrap();
-    assert!(json_tool.output_schema.is_some(), "Json<T> return type should generate output schema");
+    assert!(
+        json_tool.output_schema.is_some(),
+        "Json<T> return type should generate output schema"
+    );
 }
 
 #[tokio::test]
@@ -76,7 +81,10 @@ async fn test_non_json_type_no_schema() {
 
     // Find the without-json tool
     let non_json_tool = tools.iter().find(|t| t.name == "without-json").unwrap();
-    assert!(non_json_tool.output_schema.is_none(), "Regular return type should NOT generate output schema");
+    assert!(
+        non_json_tool.output_schema.is_none(),
+        "Regular return type should NOT generate output schema"
+    );
 }
 
 #[tokio::test]
@@ -86,7 +94,10 @@ async fn test_result_with_json_generates_schema() {
 
     // Find the result-with-json tool
     let result_json_tool = tools.iter().find(|t| t.name == "result-with-json").unwrap();
-    assert!(result_json_tool.output_schema.is_some(), "Result<Json<T>, E> return type should generate output schema");
+    assert!(
+        result_json_tool.output_schema.is_some(),
+        "Result<Json<T>, E> return type should generate output schema"
+    );
 }
 
 #[tokio::test]
@@ -96,5 +107,8 @@ async fn test_explicit_schema_override() {
 
     // Find the explicit-schema tool
     let explicit_tool = tools.iter().find(|t| t.name == "explicit-schema").unwrap();
-    assert!(explicit_tool.output_schema.is_some(), "Explicit output_schema attribute should work");
-}
+    assert!(
+        explicit_tool.output_schema.is_some(),
+        "Explicit output_schema attribute should work"
+    );
+}

crates/rmcp/tests/test_tool_builder_methods.rs

@@ -0,0 +1,62 @@
+//cargo test --test test_tool_builder_methods --features "client server macros"
+use rmcp::model::{JsonObject, Tool};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, JsonSchema)]
+pub struct InputData {
+    pub name: String,
+    pub age: u32,
+}
+
+#[derive(Serialize, Deserialize, JsonSchema)]
+pub struct OutputData {
+    pub greeting: String,
+    pub is_adult: bool,
+}
+
+#[test]
+fn test_with_output_schema() {
+    let tool = Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::<OutputData>();
+
+    assert!(tool.output_schema.is_some());
+
+    // Verify the schema contains expected fields
+    let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap();
+    assert!(schema_str.contains("greeting"));
+    assert!(schema_str.contains("is_adult"));
+}
+
+#[test]
+fn test_with_input_schema() {
+    let tool = Tool::new("test", "Test tool", JsonObject::new()).with_input_schema::<InputData>();
+
+    // Verify the schema contains expected fields
+    let schema_str = serde_json::to_string(&tool.input_schema).unwrap();
+    assert!(schema_str.contains("name"));
+    assert!(schema_str.contains("age"));
+}
+
+#[test]
+fn test_chained_builder_methods() {
+    let tool = Tool::new("test", "Test tool", JsonObject::new())
+        .with_input_schema::<InputData>()
+        .with_output_schema::<OutputData>()
+        .annotate(rmcp::model::ToolAnnotations::new().read_only(true));
+
+    assert!(tool.output_schema.is_some());
+    assert!(tool.annotations.is_some());
+    assert_eq!(
+        tool.annotations.as_ref().unwrap().read_only_hint,
+        Some(true)
+    );
+
+    // Verify both schemas are set correctly
+    let input_schema_str = serde_json::to_string(&tool.input_schema).unwrap();
+    assert!(input_schema_str.contains("name"));
+    assert!(input_schema_str.contains("age"));
+
+    let output_schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap();
+    assert!(output_schema_str.contains("greeting"));
+    assert!(output_schema_str.contains("is_adult"));
+}