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>
This commit is contained in:
parent
ba00b15097
commit
9e3de344f4
7 changed files with 262 additions and 53 deletions
|
|
@ -28,13 +28,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
|
|||
if let Some(inner_type) = extract_json_inner_type(ret_type) {
|
||||
return syn::parse2::<Expr>(quote! {
|
||||
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();
|
||||
}
|
||||
|
|
@ -65,13 +58,6 @@ fn extract_schema_from_return_type(ret_type: &syn::Type) -> Option<Expr> {
|
|||
|
||||
syn::parse2::<Expr>(quote! {
|
||||
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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -106,32 +106,40 @@ pub fn schema_for_empty_input() -> Arc<JsonObject> {
|
|||
EMPTY.clone()
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
/// Strip top-level `title` and `description` from a JSON schema for outputSchema.
|
||||
/// Unlike `validate_and_strip`, this performs no validation — output schemas are not
|
||||
/// restricted to `type: "object"` (per SEP-2106).
|
||||
fn strip_output(raw: &Arc<JsonObject>) -> Arc<JsonObject> {
|
||||
let mut object = raw.as_ref().clone();
|
||||
object.remove("title");
|
||||
object.remove("description");
|
||||
Arc::new(object)
|
||||
}
|
||||
|
||||
/// Generate and strip a JSON schema for outputSchema (top-level "title" and
|
||||
/// "description" are removed; output schemas are not restricted to root type "object").
|
||||
pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
|
||||
thread_local! {
|
||||
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Result<Arc<JsonObject>, String>>> = Default::default();
|
||||
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
|
||||
};
|
||||
|
||||
CACHE_FOR_OUTPUT.with(|cache| {
|
||||
// Try to get from cache first
|
||||
if let Some(result) = cache
|
||||
if let Some(schema) = cache
|
||||
.read()
|
||||
.expect("output schema cache lock poisoned")
|
||||
.get(&TypeId::of::<T>())
|
||||
{
|
||||
return result.clone();
|
||||
return schema.clone();
|
||||
}
|
||||
|
||||
// Generate, validate, and strip unnecessary top-level fields
|
||||
let result = validate_and_strip(&schema_for_type::<T>(), "outputSchema");
|
||||
let schema = strip_output(&schema_for_type::<T>());
|
||||
|
||||
// Cache the result (both success and error cases)
|
||||
cache
|
||||
.write()
|
||||
.expect("output schema cache lock poisoned")
|
||||
.insert(TypeId::of::<T>(), result.clone());
|
||||
.insert(TypeId::of::<T>(), schema.clone());
|
||||
|
||||
result
|
||||
schema
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -305,10 +313,69 @@ mod tests {
|
|||
assert!(Arc::ptr_eq(&schema, &cloned));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_accepts_primitive() {
|
||||
let schema = schema_for_output::<i32>();
|
||||
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_strips_description_for_primitive() {
|
||||
let schema = schema_for_output::<i32>();
|
||||
assert!(!schema.contains_key("description"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_accepts_composition() {
|
||||
let schema = schema_for_output::<Option<String>>();
|
||||
let schema_str = serde_json::to_string(&schema).unwrap();
|
||||
assert!(
|
||||
schema_str.contains("anyOf")
|
||||
|| schema_str.contains("oneOf")
|
||||
|| schema_str.contains("null"),
|
||||
"Expected composition schema for Option<String>, got: {schema_str}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_caches_result() {
|
||||
let schema1 = schema_for_output::<i32>();
|
||||
let schema2 = schema_for_output::<i32>();
|
||||
assert!(Arc::ptr_eq(&schema1, &schema2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_input_rejects_array() {
|
||||
let result = schema_for_input::<Vec<i32>>();
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_accepts_unit() {
|
||||
let _schema = schema_for_output::<()>();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_accepts_object() {
|
||||
let schema = schema_for_output::<TestObject>();
|
||||
assert_eq!(schema.get("type"), Some(&serde_json::json!("object")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_strips_top_level_title() {
|
||||
let schema = schema_for_output::<TestObject>();
|
||||
assert!(!schema.contains_key("title"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_schema_for_output_strips_top_level_description() {
|
||||
let schema = schema_for_output::<TestObject>();
|
||||
assert!(!schema.contains_key("description"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::output(schema_for_output::<i32>)]
|
||||
#[case::input(schema_for_input::<i32>)]
|
||||
fn test_schema_for_object_wrappers_reject_primitives(
|
||||
fn test_schema_for_input_rejects_primitives(
|
||||
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
|
||||
) {
|
||||
let result = schema_fn();
|
||||
|
|
@ -316,9 +383,8 @@ mod tests {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::output(schema_for_output::<TestObject>)]
|
||||
#[case::input(schema_for_input::<TestObject>)]
|
||||
fn test_schema_for_object_wrappers_accept_objects(
|
||||
fn test_schema_for_input_accepts_objects(
|
||||
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
|
||||
) {
|
||||
let result = schema_fn();
|
||||
|
|
@ -326,11 +392,9 @@ mod tests {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::output_title(schema_for_output::<TestObject>, "title")]
|
||||
#[case::output_description(schema_for_output::<TestObject>, "description")]
|
||||
#[case::input_title(schema_for_input::<TestObject>, "title")]
|
||||
#[case::input_description(schema_for_input::<TestObject>, "description")]
|
||||
fn test_schema_for_object_wrappers_strip_top_level_metadata(
|
||||
fn test_schema_for_input_strips_top_level_metadata(
|
||||
#[case] schema_fn: fn() -> Result<Arc<JsonObject>, String>,
|
||||
#[case] field: &str,
|
||||
) {
|
||||
|
|
|
|||
|
|
@ -65,13 +65,7 @@ pub trait ToolBase {
|
|||
///
|
||||
/// If the tool does not have any output, you should override this methods to return [`None`].
|
||||
fn output_schema() -> Option<Arc<JsonObject>> {
|
||||
Some(schema_for_output::<Self::Output>().unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"Invalid output schema for ToolBase::Output type `{0}`: {1}",
|
||||
std::any::type_name::<Self::Output>(),
|
||||
e,
|
||||
);
|
||||
}))
|
||||
Some(schema_for_output::<Self::Output>())
|
||||
}
|
||||
|
||||
fn annotations() -> Option<ToolAnnotations> {
|
||||
|
|
|
|||
|
|
@ -315,15 +315,9 @@ 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.
|
||||
#[cfg(feature = "server")]
|
||||
pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self {
|
||||
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.output_schema = Some(crate::handler::server::tool::schema_for_output::<T>());
|
||||
self
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -60,6 +60,28 @@ impl TestServer {
|
|||
pub async fn explicit_schema(&self) -> Result<String, String> {
|
||||
Ok("test".to_string())
|
||||
}
|
||||
|
||||
/// Tool that returns Json<Vec<T>> - array output schema
|
||||
#[tool(name = "with-json-array")]
|
||||
pub async fn with_json_array(&self) -> Result<Json<Vec<TestData>>, String> {
|
||||
Ok(Json(vec![TestData {
|
||||
value: "test".to_string(),
|
||||
}]))
|
||||
}
|
||||
|
||||
/// Tool that returns Result<Json<Vec<T>>, ErrorData> - array output schema
|
||||
#[tool(name = "result-with-json-array")]
|
||||
pub async fn result_with_json_array(&self) -> Result<Json<Vec<TestData>>, rmcp::ErrorData> {
|
||||
Ok(Json(vec![TestData {
|
||||
value: "test".to_string(),
|
||||
}]))
|
||||
}
|
||||
|
||||
/// Tool that returns Json<String> - string output schema
|
||||
#[tool(name = "with-json-string")]
|
||||
pub async fn with_json_string(&self) -> Result<Json<String>, String> {
|
||||
Ok(Json("test".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -113,3 +135,54 @@ async fn test_explicit_schema_override() {
|
|||
"Explicit output_schema attribute should work"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_array_type_generates_schema() {
|
||||
let server = TestServer::new();
|
||||
let tools = server.tool_router.list_all();
|
||||
|
||||
let array_tool = tools.iter().find(|t| t.name == "with-json-array").unwrap();
|
||||
assert!(
|
||||
array_tool.output_schema.is_some(),
|
||||
"Json<Vec<T>> return type should generate output schema"
|
||||
);
|
||||
let schema = array_tool.output_schema.as_ref().unwrap();
|
||||
assert_eq!(
|
||||
schema.get("type").and_then(|v| v.as_str()),
|
||||
Some("array"),
|
||||
"Json<Vec<T>> should produce an array schema"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_result_with_json_array_generates_schema() {
|
||||
let server = TestServer::new();
|
||||
let tools = server.tool_router.list_all();
|
||||
|
||||
let result_array_tool = tools
|
||||
.iter()
|
||||
.find(|t| t.name == "result-with-json-array")
|
||||
.unwrap();
|
||||
assert!(
|
||||
result_array_tool.output_schema.is_some(),
|
||||
"Result<Json<Vec<T>>, ErrorData> return type should generate output schema"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_string_type_generates_schema() {
|
||||
let server = TestServer::new();
|
||||
let tools = server.tool_router.list_all();
|
||||
|
||||
let string_tool = tools.iter().find(|t| t.name == "with-json-string").unwrap();
|
||||
assert!(
|
||||
string_tool.output_schema.is_some(),
|
||||
"Json<String> return type should generate output schema"
|
||||
);
|
||||
let schema = string_tool.output_schema.as_ref().unwrap();
|
||||
assert_eq!(
|
||||
schema.get("type").and_then(|v| v.as_str()),
|
||||
Some("string"),
|
||||
"Json<String> should produce a string schema"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -93,6 +93,27 @@ impl TestServer {
|
|||
Err("User not found".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool that returns a list of calculation results
|
||||
#[tool(
|
||||
name = "calculate-list",
|
||||
description = "Return a list of calculation results"
|
||||
)]
|
||||
pub async fn calculate_list(
|
||||
&self,
|
||||
params: Parameters<CalculationRequest>,
|
||||
) -> Result<Json<Vec<CalculationResult>>, String> {
|
||||
Ok(Json(vec![CalculationResult {
|
||||
sum: params.0.a + params.0.b,
|
||||
product: params.0.a * params.0.b,
|
||||
}]))
|
||||
}
|
||||
|
||||
/// Tool that returns a count
|
||||
#[tool(name = "get-count", description = "Return a count")]
|
||||
pub async fn get_count(&self) -> Result<Json<i32>, String> {
|
||||
Ok(Json(42))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -360,3 +381,38 @@ fn test_call_tool_result_deserialize_without_content() {
|
|||
assert!(result.content.is_empty());
|
||||
assert!(result.structured_content.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_with_array_output_schema() {
|
||||
let server = TestServer::new();
|
||||
let tools = server.tool_router.list_all();
|
||||
|
||||
// Find the calculate-list tool
|
||||
let calculate_list_tool = tools.iter().find(|t| t.name == "calculate-list").unwrap();
|
||||
|
||||
// Verify it has an output schema
|
||||
assert!(calculate_list_tool.output_schema.is_some());
|
||||
|
||||
let schema = calculate_list_tool.output_schema.as_ref().unwrap();
|
||||
|
||||
// Check that the schema contains array type
|
||||
let schema_str = serde_json::to_string(schema).unwrap();
|
||||
assert!(schema_str.contains("array"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_with_primitive_output_schema() {
|
||||
let server = TestServer::new();
|
||||
let tools = server.tool_router.list_all();
|
||||
|
||||
// Find the get-count tool
|
||||
let get_count_tool = tools.iter().find(|t| t.name == "get-count").unwrap();
|
||||
|
||||
// Verify it has an output schema
|
||||
assert!(get_count_tool.output_schema.is_some());
|
||||
|
||||
let schema = get_count_tool.output_schema.as_ref().unwrap();
|
||||
|
||||
// Check that the schema contains integer type
|
||||
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,10 +22,8 @@ fn test_with_output_schema() {
|
|||
|
||||
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"));
|
||||
let schema = tool.output_schema.as_ref().unwrap();
|
||||
assert_eq!(schema.get("type"), Some(&serde_json::json!("object")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -57,7 +55,51 @@ fn test_chained_builder_methods() {
|
|||
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"));
|
||||
let output_schema = tool.output_schema.as_ref().unwrap();
|
||||
assert_eq!(
|
||||
output_schema.get("type"),
|
||||
Some(&serde_json::json!("object"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_output_schema_primitive() {
|
||||
let tool = Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::<i32>();
|
||||
|
||||
assert!(tool.output_schema.is_some());
|
||||
|
||||
let schema = tool.output_schema.as_ref().unwrap();
|
||||
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
|
||||
// title should be stripped from output schema
|
||||
assert!(schema.get("title").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_output_schema_array() {
|
||||
let tool =
|
||||
Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::<Vec<String>>();
|
||||
|
||||
assert!(tool.output_schema.is_some());
|
||||
|
||||
let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap();
|
||||
assert!(schema_str.contains("\"type\":\"array\""));
|
||||
assert!(schema_str.contains("items"));
|
||||
// title should be stripped from output schema
|
||||
assert!(!schema_str.contains("title"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_output_schema_option() {
|
||||
let tool =
|
||||
Tool::new("test", "Test tool", JsonObject::new()).with_output_schema::<Option<String>>();
|
||||
|
||||
assert!(tool.output_schema.is_some());
|
||||
|
||||
let schema_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap();
|
||||
// Option<String> generates a composition schema (anyOf/oneOf/type array with null)
|
||||
assert!(
|
||||
schema_str.contains("anyOf") || schema_str.contains("oneOf") || schema_str.contains("null"),
|
||||
"Expected composition schema for Option<String>, got: {schema_str}"
|
||||
);
|
||||
assert!(!schema_str.contains("title"));
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue