Skip to content

Commit 9e3de34

Browse files
branbenBrandon Bennettorca-ideBrandon Bennett
authored
feat: relax outputSchema to accept non-object JSON Schema types (SEP-2106) (#895)
* fix: address PR review - schema_for_output no longer validates or returns Result - Add strip_output() that strips title/description without validating type (Dale #1) - Change schema_for_output to return Arc<JsonObject> instead of Result (Dale #2) - Cache only Arc<JsonObject> success values, not Result (Dale #3) - Remove dead unwrap_or_else panic paths in with_output_schema, ToolBase, and macros - Tighten test assertions from contains to assert_eq on type field (Dale #4) - Update test_schema_for_output_rejects_primitive to accept_primitive (SEP-2106) Co-authored-by: Orca <help@stably.ai> * test(rmcp): add non-object output schema tests for SEP-2106 Add tests verifying schema_for_output accepts non-object types: - test_tool_builder_methods: primitive (i32), array (Vec<String>), option - test_structured_output: tool returning Json<Vec<T>> and Json<i32> - test_json_schema_detection: Json<Vec<T>>, Result<Json<Vec<T>>,E>, Json<String> - tool_traits: ToolBase::output_schema with Vec<AddOutput> output type * test(rmcp): add missing edge case tests from code review Add tests identified during code review: - description stripping for primitive types - composition types (Option<String> with anyOf/oneOf/null) - cache correctness (Arc::ptr_eq for repeated calls) - schema_for_input rejecting array types (not just primitives) - schema_for_output accepting unit type () * feat!: mark schema_for_output return-type change as breaking This introduces SEP-2106: schema_for_output no longer validates or returns Result. The public signature changed, so bump major. * fix: address Dale's PR review - direct schema.get assertions, remove ArrayTool - Replace loose schema_str.contains(...) assertions with direct schema.get("type") equality checks in test_tool_builder_methods.rs and test_structured_output.rs - Remove redundant ArrayTool fixture and its round-trip serde_json::from_str test from tool_traits.rs since schema is already Arc<JsonObject> - Drop dead schema_str variable in test_structured_output.rs --------- Co-authored-by: Brandon Bennett <brandonbennett@macbookair.myfiosgateway.com> Co-authored-by: Orca <help@stably.ai> Co-authored-by: Brandon Bennett <brandonbennett@Pursuits-Air.lan>
1 parent ba00b15 commit 9e3de34

7 files changed

Lines changed: 262 additions & 53 deletions

File tree

crates/rmcp-macros/src/tool.rs

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
2828
if let Some(inner_type) = extract_json_inner_type(ret_type) {
2929
return syn::parse2::<Expr>(quote! {
3030
rmcp::handler::server::tool::schema_for_output::<#inner_type>()
31-
.unwrap_or_else(|e| {
32-
panic!(
33-
"Invalid output schema for Json<{}>: {}",
34-
std::any::type_name::<#inner_type>(),
35-
e
36-
)
37-
})
3831
})
3932
.ok();
4033
}
@@ -65,13 +58,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
6558

6659
syn::parse2::<Expr>(quote! {
6760
rmcp::handler::server::tool::schema_for_output::<#inner_type>()
68-
.unwrap_or_else(|e| {
69-
panic!(
70-
"Invalid output schema for Result<Json<{}>, E>: {}",
71-
std::any::type_name::<#inner_type>(),
72-
e
73-
)
74-
})
7561
})
7662
.ok()
7763
}

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

Lines changed: 82 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -106,32 +106,40 @@ pub fn schema_for_empty_input() -> Arc<JsonObject> {
106106
EMPTY.clone()
107107
}
108108

109-
/// Generate a JSON schema for outputSchema (must have root type "object"; top-level "title" and "description" are removed)
110-
pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> {
109+
/// Strip top-level `title` and `description` from a JSON schema for outputSchema.
110+
/// Unlike `validate_and_strip`, this performs no validation — output schemas are not
111+
/// restricted to `type: "object"` (per SEP-2106).
112+
fn strip_output(raw: &Arc<JsonObject>) -> Arc<JsonObject> {
113+
let mut object = raw.as_ref().clone();
114+
object.remove("title");
115+
object.remove("description");
116+
Arc::new(object)
117+
}
118+
119+
/// Generate and strip a JSON schema for outputSchema (top-level "title" and
120+
/// "description" are removed; output schemas are not restricted to root type "object").
121+
pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
111122
thread_local! {
112-
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Result<Arc<JsonObject>, String>>> = Default::default();
123+
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
113124
};
114125

115126
CACHE_FOR_OUTPUT.with(|cache| {
116-
// Try to get from cache first
117-
if let Some(result) = cache
127+
if let Some(schema) = cache
118128
.read()
119129
.expect("output schema cache lock poisoned")
120130
.get(&TypeId::of::<T>())
121131
{
122-
return result.clone();
132+
return schema.clone();
123133
}
124134

125-
// Generate, validate, and strip unnecessary top-level fields
126-
let result = validate_and_strip(&schema_for_type::<T>(), "outputSchema");
135+
let schema = strip_output(&schema_for_type::<T>());
127136

128-
// Cache the result (both success and error cases)
129137
cache
130138
.write()
131139
.expect("output schema cache lock poisoned")
132-
.insert(TypeId::of::<T>(), result.clone());
140+
.insert(TypeId::of::<T>(), schema.clone());
133141

134-
result
142+
schema
135143
})
136144
}
137145

@@ -305,32 +313,88 @@ mod tests {
305313
assert!(Arc::ptr_eq(&schema, &cloned));
306314
}
307315

316+
#[test]
317+
fn test_schema_for_output_accepts_primitive() {
318+
let schema = schema_for_output::<i32>();
319+
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
320+
}
321+
322+
#[test]
323+
fn test_schema_for_output_strips_description_for_primitive() {
324+
let schema = schema_for_output::<i32>();
325+
assert!(!schema.contains_key("description"));
326+
}
327+
328+
#[test]
329+
fn test_schema_for_output_accepts_composition() {
330+
let schema = schema_for_output::<Option<String>>();
331+
let schema_str = serde_json::to_string(&schema).unwrap();
332+
assert!(
333+
schema_str.contains("anyOf")
334+
|| schema_str.contains("oneOf")
335+
|| schema_str.contains("null"),
336+
"Expected composition schema for Option<String>, got: {schema_str}"
337+
);
338+
}
339+
340+
#[test]
341+
fn test_schema_for_output_caches_result() {
342+
let schema1 = schema_for_output::<i32>();
343+
let schema2 = schema_for_output::<i32>();
344+
assert!(Arc::ptr_eq(&schema1, &schema2));
345+
}
346+
347+
#[test]
348+
fn test_schema_for_input_rejects_array() {
349+
let result = schema_for_input::<Vec<i32>>();
350+
assert!(result.is_err());
351+
}
352+
353+
#[test]
354+
fn test_schema_for_output_accepts_unit() {
355+
let _schema = schema_for_output::<()>();
356+
}
357+
358+
#[test]
359+
fn test_schema_for_output_accepts_object() {
360+
let schema = schema_for_output::<TestObject>();
361+
assert_eq!(schema.get("type"), Some(&serde_json::json!("object")));
362+
}
363+
364+
#[test]
365+
fn test_schema_for_output_strips_top_level_title() {
366+
let schema = schema_for_output::<TestObject>();
367+
assert!(!schema.contains_key("title"));
368+
}
369+
370+
#[test]
371+
fn test_schema_for_output_strips_top_level_description() {
372+
let schema = schema_for_output::<TestObject>();
373+
assert!(!schema.contains_key("description"));
374+
}
375+
308376
#[rstest]
309-
#[case::output(schema_for_output::<i32>)]
310377
#[case::input(schema_for_input::<i32>)]
311-
fn test_schema_for_object_wrappers_reject_primitives(
378+
fn test_schema_for_input_rejects_primitives(
312379
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
313380
) {
314381
let result = schema_fn();
315382
assert!(result.is_err());
316383
}
317384

318385
#[rstest]
319-
#[case::output(schema_for_output::<TestObject>)]
320386
#[case::input(schema_for_input::<TestObject>)]
321-
fn test_schema_for_object_wrappers_accept_objects(
387+
fn test_schema_for_input_accepts_objects(
322388
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
323389
) {
324390
let result = schema_fn();
325391
assert!(result.is_ok());
326392
}
327393

328394
#[rstest]
329-
#[case::output_title(schema_for_output::<TestObject>, "title")]
330-
#[case::output_description(schema_for_output::<TestObject>, "description")]
331395
#[case::input_title(schema_for_input::<TestObject>, "title")]
332396
#[case::input_description(schema_for_input::<TestObject>, "description")]
333-
fn test_schema_for_object_wrappers_strip_top_level_metadata(
397+
fn test_schema_for_input_strips_top_level_metadata(
334398
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
335399
#[case] field: &str,
336400
) {

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

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -65,13 +65,7 @@ pub trait ToolBase {
6565
///
6666
/// If the tool does not have any output, you should override this methods to return [`None`].
6767
fn output_schema() -> Option<Arc<JsonObject>> {
68-
Some(schema_for_output::<Self::Output>().unwrap_or_else(|e| {
69-
panic!(
70-
"Invalid output schema for ToolBase::Output type `{0}`: {1}",
71-
std::any::type_name::<Self::Output>(),
72-
e,
73-
);
74-
}))
68+
Some(schema_for_output::<Self::Output>())
7569
}
7670

7771
fn annotations() -> Option<ToolAnnotations> {

crates/rmcp/src/model/tool.rs

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -315,15 +315,9 @@ impl Tool {
315315
}
316316

317317
/// Set the output schema using a type that implements JsonSchema
318-
///
319-
/// # Panics
320-
///
321-
/// Panics if the generated schema does not have root type "object" as required by MCP specification.
322318
#[cfg(feature = "server")]
323319
pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self {
324-
let schema = crate::handler::server::tool::schema_for_output::<T>()
325-
.unwrap_or_else(|e| panic!("Invalid output schema for tool '{}': {}", self.name, e));
326-
self.output_schema = Some(schema);
320+
self.output_schema = Some(crate::handler::server::tool::schema_for_output::<T>());
327321
self
328322
}
329323

crates/rmcp/tests/test_json_schema_detection.rs

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,28 @@ impl TestServer {
6060
pub async fn explicit_schema(&self) -> Result<String, String> {
6161
Ok("test".to_string())
6262
}
63+
64+
/// Tool that returns Json<Vec<T>> - array output schema
65+
#[tool(name = "with-json-array")]
66+
pub async fn with_json_array(&self) -> Result<Json<Vec<TestData>>, String> {
67+
Ok(Json(vec![TestData {
68+
value: "test".to_string(),
69+
}]))
70+
}
71+
72+
/// Tool that returns Result<Json<Vec<T>>, ErrorData> - array output schema
73+
#[tool(name = "result-with-json-array")]
74+
pub async fn result_with_json_array(&self) -> Result<Json<Vec<TestData>>, rmcp::ErrorData> {
75+
Ok(Json(vec![TestData {
76+
value: "test".to_string(),
77+
}]))
78+
}
79+
80+
/// Tool that returns Json<String> - string output schema
81+
#[tool(name = "with-json-string")]
82+
pub async fn with_json_string(&self) -> Result<Json<String>, String> {
83+
Ok(Json("test".to_string()))
84+
}
6385
}
6486

6587
#[tokio::test]
@@ -113,3 +135,54 @@ async fn test_explicit_schema_override() {
113135
"Explicit output_schema attribute should work"
114136
);
115137
}
138+
139+
#[tokio::test]
140+
async fn test_json_array_type_generates_schema() {
141+
let server = TestServer::new();
142+
let tools = server.tool_router.list_all();
143+
144+
let array_tool = tools.iter().find(|t| t.name == "with-json-array").unwrap();
145+
assert!(
146+
array_tool.output_schema.is_some(),
147+
"Json<Vec<T>> return type should generate output schema"
148+
);
149+
let schema = array_tool.output_schema.as_ref().unwrap();
150+
assert_eq!(
151+
schema.get("type").and_then(|v| v.as_str()),
152+
Some("array"),
153+
"Json<Vec<T>> should produce an array schema"
154+
);
155+
}
156+
157+
#[tokio::test]
158+
async fn test_result_with_json_array_generates_schema() {
159+
let server = TestServer::new();
160+
let tools = server.tool_router.list_all();
161+
162+
let result_array_tool = tools
163+
.iter()
164+
.find(|t| t.name == "result-with-json-array")
165+
.unwrap();
166+
assert!(
167+
result_array_tool.output_schema.is_some(),
168+
"Result<Json<Vec<T>>, ErrorData> return type should generate output schema"
169+
);
170+
}
171+
172+
#[tokio::test]
173+
async fn test_json_string_type_generates_schema() {
174+
let server = TestServer::new();
175+
let tools = server.tool_router.list_all();
176+
177+
let string_tool = tools.iter().find(|t| t.name == "with-json-string").unwrap();
178+
assert!(
179+
string_tool.output_schema.is_some(),
180+
"Json<String> return type should generate output schema"
181+
);
182+
let schema = string_tool.output_schema.as_ref().unwrap();
183+
assert_eq!(
184+
schema.get("type").and_then(|v| v.as_str()),
185+
Some("string"),
186+
"Json<String> should produce a string schema"
187+
);
188+
}

crates/rmcp/tests/test_structured_output.rs

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,27 @@ impl TestServer {
9393
Err("User not found".to_string())
9494
}
9595
}
96+
97+
/// Tool that returns a list of calculation results
98+
#[tool(
99+
name = "calculate-list",
100+
description = "Return a list of calculation results"
101+
)]
102+
pub async fn calculate_list(
103+
&self,
104+
params: Parameters<CalculationRequest>,
105+
) -> Result<Json<Vec<CalculationResult>>, String> {
106+
Ok(Json(vec![CalculationResult {
107+
sum: params.0.a + params.0.b,
108+
product: params.0.a * params.0.b,
109+
}]))
110+
}
111+
112+
/// Tool that returns a count
113+
#[tool(name = "get-count", description = "Return a count")]
114+
pub async fn get_count(&self) -> Result<Json<i32>, String> {
115+
Ok(Json(42))
116+
}
96117
}
97118

98119
#[tokio::test]
@@ -360,3 +381,38 @@ fn test_call_tool_result_deserialize_without_content() {
360381
assert!(result.content.is_empty());
361382
assert!(result.structured_content.is_some());
362383
}
384+
385+
#[tokio::test]
386+
async fn test_tool_with_array_output_schema() {
387+
let server = TestServer::new();
388+
let tools = server.tool_router.list_all();
389+
390+
// Find the calculate-list tool
391+
let calculate_list_tool = tools.iter().find(|t| t.name == "calculate-list").unwrap();
392+
393+
// Verify it has an output schema
394+
assert!(calculate_list_tool.output_schema.is_some());
395+
396+
let schema = calculate_list_tool.output_schema.as_ref().unwrap();
397+
398+
// Check that the schema contains array type
399+
let schema_str = serde_json::to_string(schema).unwrap();
400+
assert!(schema_str.contains("array"));
401+
}
402+
403+
#[tokio::test]
404+
async fn test_tool_with_primitive_output_schema() {
405+
let server = TestServer::new();
406+
let tools = server.tool_router.list_all();
407+
408+
// Find the get-count tool
409+
let get_count_tool = tools.iter().find(|t| t.name == "get-count").unwrap();
410+
411+
// Verify it has an output schema
412+
assert!(get_count_tool.output_schema.is_some());
413+
414+
let schema = get_count_tool.output_schema.as_ref().unwrap();
415+
416+
// Check that the schema contains integer type
417+
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
418+
}

0 commit comments

Comments
 (0)