preserve nested mcp object schemas

This commit is contained in:
Soheil Norouzi
2026-04-16 13:12:33 -04:00
parent baaf42b2e4
commit ed2b627bc1
2 changed files with 158 additions and 10 deletions

View File

@@ -147,8 +147,9 @@ impl From<JsonSchema> for AdditionalProperties {
/// Parse the tool `input_schema` or return an error for invalid schema.
pub fn parse_tool_input_schema(input_schema: &JsonValue) -> Result<JsonSchema, serde_json::Error> {
let mut input_schema = input_schema.clone();
sanitize_json_schema(&mut input_schema);
let root_schema = input_schema.clone();
let mut input_schema = root_schema.clone();
sanitize_json_schema(&mut input_schema, &root_schema);
let schema: JsonSchema = serde_json::from_value(input_schema)?;
if matches!(
schema.schema_type,
@@ -162,11 +163,13 @@ pub fn parse_tool_input_schema(input_schema: &JsonValue) -> Result<JsonSchema, s
/// Sanitize a JSON Schema (as serde_json::Value) so it can fit our limited
/// schema representation. This function:
/// - Ensures every typed schema object has a `"type"` when required.
/// - Resolves local `$ref` indirections and unwraps single-variant
/// `oneOf`/`anyOf`/`allOf` wrappers before inferring a fallback type.
/// - Preserves explicit `anyOf`.
/// - Collapses `const` into single-value `enum`.
/// - Fills required child fields for object/array schema types, including
/// nullable unions, with permissive defaults when absent.
fn sanitize_json_schema(value: &mut JsonValue) {
fn sanitize_json_schema(value: &mut JsonValue, root_schema: &JsonValue) {
match value {
JsonValue::Bool(_) => {
// JSON Schema boolean form: true/false. Coerce to an accept-all string.
@@ -174,30 +177,42 @@ fn sanitize_json_schema(value: &mut JsonValue) {
}
JsonValue::Array(values) => {
for value in values {
sanitize_json_schema(value);
sanitize_json_schema(value, root_schema);
}
}
JsonValue::Object(map) => {
if let Some(replacement) = resolve_json_schema_reference(map, root_schema) {
*value = replacement;
sanitize_json_schema(value, root_schema);
return;
}
if let Some(replacement) = unwrap_single_variant_combiner(map) {
*value = replacement;
sanitize_json_schema(value, root_schema);
return;
}
if let Some(properties) = map.get_mut("properties")
&& let Some(properties_map) = properties.as_object_mut()
{
for value in properties_map.values_mut() {
sanitize_json_schema(value);
sanitize_json_schema(value, root_schema);
}
}
if let Some(items) = map.get_mut("items") {
sanitize_json_schema(items);
sanitize_json_schema(items, root_schema);
}
if let Some(additional_properties) = map.get_mut("additionalProperties")
&& !matches!(additional_properties, JsonValue::Bool(_))
{
sanitize_json_schema(additional_properties);
sanitize_json_schema(additional_properties, root_schema);
}
if let Some(value) = map.get_mut("prefixItems") {
sanitize_json_schema(value);
sanitize_json_schema(value, root_schema);
}
if let Some(value) = map.get_mut("anyOf") {
sanitize_json_schema(value);
sanitize_json_schema(value, root_schema);
}
if let Some(const_value) = map.remove("const") {
@@ -239,6 +254,48 @@ fn sanitize_json_schema(value: &mut JsonValue) {
}
}
fn resolve_json_schema_reference(
map: &serde_json::Map<String, JsonValue>,
root_schema: &JsonValue,
) -> Option<JsonValue> {
let reference = map.get("$ref")?.as_str()?;
let pointer = reference.strip_prefix('#')?;
let mut replacement = root_schema.pointer(pointer)?.clone();
if let JsonValue::Object(replacement_map) = &mut replacement {
for (key, value) in map {
if key != "$ref" {
replacement_map.insert(key.clone(), value.clone());
}
}
}
Some(replacement)
}
fn unwrap_single_variant_combiner(map: &serde_json::Map<String, JsonValue>) -> Option<JsonValue> {
for combiner in ["oneOf", "anyOf", "allOf"] {
let Some(variants) = map.get(combiner).and_then(JsonValue::as_array) else {
continue;
};
if variants.len() != 1 {
continue;
}
let mut replacement = variants[0].clone();
if let JsonValue::Object(replacement_map) = &mut replacement {
for (key, value) in map {
if key != combiner {
replacement_map
.entry(key.clone())
.or_insert_with(|| value.clone());
}
}
}
return Some(replacement);
}
None
}
fn ensure_default_children_for_schema_types(
map: &mut serde_json::Map<String, JsonValue>,
schema_types: &[JsonSchemaPrimitiveType],

View File

@@ -452,6 +452,97 @@ fn parse_tool_input_schema_fills_default_items_for_nullable_array_union() {
);
}
#[test]
fn parse_tool_input_schema_resolves_local_ref_objects() {
let schema = parse_tool_input_schema(&serde_json::json!({
"type": "object",
"properties": {
"start": { "$ref": "#/$defs/date_time_zone" }
},
"$defs": {
"date_time_zone": {
"type": "object",
"properties": {
"dateTime": { "type": "string" },
"timeZone": { "type": "string" }
},
"required": ["dateTime", "timeZone"]
}
}
}))
.expect("parse schema");
assert_eq!(
schema,
JsonSchema::object(
BTreeMap::from([(
"start".to_string(),
JsonSchema::object(
BTreeMap::from([
(
"dateTime".to_string(),
JsonSchema::string(/*description*/ None),
),
(
"timeZone".to_string(),
JsonSchema::string(/*description*/ None),
),
]),
Some(vec!["dateTime".to_string(), "timeZone".to_string()]),
/*additional_properties*/ None,
),
)]),
/*required*/ None,
/*additional_properties*/ None
)
);
}
#[test]
fn parse_tool_input_schema_unwraps_single_variant_all_of_objects() {
let schema = parse_tool_input_schema(&serde_json::json!({
"type": "object",
"properties": {
"end": {
"allOf": [{
"type": "object",
"properties": {
"dateTime": { "type": "string" },
"timeZone": { "type": "string" }
},
"required": ["dateTime", "timeZone"]
}]
}
}
}))
.expect("parse schema");
assert_eq!(
schema,
JsonSchema::object(
BTreeMap::from([(
"end".to_string(),
JsonSchema::object(
BTreeMap::from([
(
"dateTime".to_string(),
JsonSchema::string(/*description*/ None),
),
(
"timeZone".to_string(),
JsonSchema::string(/*description*/ None),
),
]),
Some(vec!["dateTime".to_string(), "timeZone".to_string()]),
/*additional_properties*/ None,
),
)]),
/*required*/ None,
/*additional_properties*/ None
)
);
}
// Schemas that should be preserved for Responses API compatibility rather than
// being rewritten into a different shape.
@@ -541,7 +632,7 @@ fn parse_tool_input_schema_preserves_nested_nullable_any_of_shape() {
],
/*description*/ None,
),
),]),
)]),
/*required*/ None,
/*additional_properties*/ None
)