-
Notifications
You must be signed in to change notification settings - Fork 543
feat: relax outputSchema to accept non-object JSON Schema types (SEP-2106) #895
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -50,9 +50,9 @@ pub fn schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> { | |||||
| }) | ||||||
| } | ||||||
|
|
||||||
| /// Validate that the schema root is `type: "object"` (per MCP spec) and strip top-level | ||||||
| /// `title`/`description` (the wrapper type name and doc, which are noise to the LLM). | ||||||
| fn validate_and_strip(raw: &Arc<JsonObject>, purpose: &str) -> Result<Arc<JsonObject>, String> { | ||||||
| /// Validate that the schema root is `type: "object"` (per MCP spec for inputSchema) and | ||||||
| /// strip top-level `title`/`description` (the wrapper type name and doc, which are noise to the LLM). | ||||||
| fn validate_and_strip_input(raw: &Arc<JsonObject>) -> Result<Arc<JsonObject>, String> { | ||||||
| match raw.get("type") { | ||||||
| Some(serde_json::Value::String(t)) if t == "object" => { | ||||||
| let mut object = raw.as_ref().clone(); | ||||||
|
|
@@ -61,17 +61,27 @@ fn validate_and_strip(raw: &Arc<JsonObject>, purpose: &str) -> Result<Arc<JsonOb | |||||
| Ok(Arc::new(object)) | ||||||
| } | ||||||
| Some(serde_json::Value::String(t)) => Err(format!( | ||||||
| "MCP specification requires tool {purpose} to have root type 'object', but found '{t}'." | ||||||
| )), | ||||||
| None => Err(format!( | ||||||
| "Schema is missing 'type' field. MCP specification requires {purpose} to have root type 'object'." | ||||||
| "MCP specification requires tool inputSchema to have root type 'object', but found '{t}'." | ||||||
| )), | ||||||
| None => Err( | ||||||
| "Schema is missing 'type' field. MCP specification requires inputSchema to have root type 'object'.".to_string() | ||||||
| ), | ||||||
| Some(other) => Err(format!( | ||||||
| "Schema 'type' field has unexpected format: {other:?}. Expected \"object\"." | ||||||
| )), | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| /// Strip top-level `title`/`description` from a JSON schema for outputSchema. | ||||||
| /// Unlike inputSchema, outputSchema may have any JSON Schema 2020-12 root type | ||||||
| /// (objects, arrays, primitives, compositions) per SEP-2106. | ||||||
| fn validate_and_strip_output(raw: &Arc<JsonObject>) -> Arc<JsonObject> { | ||||||
| let mut object = raw.as_ref().clone(); | ||||||
| object.remove("title"); | ||||||
| object.remove("description"); | ||||||
| Arc::new(object) | ||||||
| } | ||||||
|
|
||||||
| /// Generate, validate, and strip a JSON schema for inputSchema (must have root type "object"; | ||||||
| /// top-level "title" and "description" are removed). | ||||||
| pub fn schema_for_input<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> { | ||||||
|
|
@@ -86,7 +96,7 @@ pub fn schema_for_input<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObjec | |||||
| { | ||||||
| return result.clone(); | ||||||
| } | ||||||
| let result = validate_and_strip(&schema_for_type::<T>(), "inputSchema"); | ||||||
| let result = validate_and_strip_input(&schema_for_type::<T>()); | ||||||
| cache | ||||||
| .write() | ||||||
| .expect("input schema cache lock poisoned") | ||||||
|
|
@@ -106,7 +116,10 @@ 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) | ||||||
| /// Generate and strip a JSON schema for outputSchema. | ||||||
| /// Unlike inputSchema, outputSchema accepts any JSON Schema 2020-12 root type | ||||||
| /// (objects, arrays, primitives, compositions) per SEP-2106. | ||||||
| /// Top-level "title" and "description" are always removed. | ||||||
| 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(); | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can cache the success value only.
Suggested change
|
||||||
|
|
@@ -122,10 +135,8 @@ pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObje | |||||
| return result.clone(); | ||||||
| } | ||||||
|
|
||||||
| // Generate, validate, and strip unnecessary top-level fields | ||||||
| let result = validate_and_strip(&schema_for_type::<T>(), "outputSchema"); | ||||||
| let result = Ok(validate_and_strip_output(&schema_for_type::<T>())); | ||||||
|
|
||||||
| // Cache the result (both success and error cases) | ||||||
| cache | ||||||
| .write() | ||||||
| .expect("output schema cache lock poisoned") | ||||||
|
|
@@ -306,9 +317,68 @@ mod tests { | |||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn test_schema_for_output_rejects_primitive() { | ||||||
| fn test_schema_for_output_accepts_primitive() { | ||||||
| let result = schema_for_output::<i32>(); | ||||||
| assert!(result.is_err(),); | ||||||
| assert!(result.is_ok()); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn test_schema_for_output_accepts_array() { | ||||||
| let result = schema_for_output::<Vec<i32>>(); | ||||||
| assert!(result.is_ok()); | ||||||
| let schema = result.unwrap(); | ||||||
| assert_eq!(schema.get("type"), Some(&serde_json::json!("array"))); | ||||||
| assert!(schema.contains_key("items")); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn test_schema_for_output_strips_title_for_primitive() { | ||||||
| let schema = schema_for_output::<i32>().unwrap(); | ||||||
| assert!(!schema.contains_key("title")); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn test_schema_for_output_strips_description_for_primitive() { | ||||||
| let schema = schema_for_output::<i32>().unwrap(); | ||||||
| assert!(!schema.contains_key("description")); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn test_schema_for_output_accepts_composition() { | ||||||
| let result = schema_for_output::<Option<String>>(); | ||||||
| assert!(result.is_ok()); | ||||||
| let schema = result.unwrap(); | ||||||
| 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 result1 = schema_for_output::<i32>(); | ||||||
| let result2 = schema_for_output::<i32>(); | ||||||
| assert!(result1.is_ok()); | ||||||
| assert!(result2.is_ok()); | ||||||
| assert!(Arc::ptr_eq( | ||||||
| result1.as_ref().unwrap(), | ||||||
| result2.as_ref().unwrap() | ||||||
| )); | ||||||
| } | ||||||
|
|
||||||
| #[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 result = schema_for_output::<()>(); | ||||||
| assert!(result.is_ok()); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
|
|
||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -318,7 +318,7 @@ impl Tool { | |
| /// | ||
| /// # Panics | ||
| /// | ||
| /// Panics if the generated schema does not have root type "object" as required by MCP specification. | ||
| /// Panics if output schema generation fails. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It won't likely panic because |
||
| #[cfg(feature = "server")] | ||
| pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self { | ||
| let schema = crate::handler::server::tool::schema_for_output::<T>() | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -61,3 +61,45 @@ fn test_chained_builder_methods() { | |||||
| assert!(output_schema_str.contains("greeting")); | ||||||
| assert!(output_schema_str.contains("is_adult")); | ||||||
| } | ||||||
|
|
||||||
| #[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_str = serde_json::to_string(tool.output_schema.as_ref().unwrap()).unwrap(); | ||||||
| assert!(schema_str.contains("\"type\":\"integer\"")); | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. To prevent false positives
Suggested change
|
||||||
| // title should be stripped from output schema | ||||||
| assert!(!schema_str.contains("title")); | ||||||
| } | ||||||
|
|
||||||
| #[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")); | ||||||
| } | ||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This performs no validation.