From ed2b627bc1c6dccd8fdea002d75403b2bbd9d1ce Mon Sep 17 00:00:00 2001 From: Soheil Norouzi Date: Thu, 16 Apr 2026 13:12:33 -0400 Subject: [PATCH] preserve nested mcp object schemas --- codex-rs/tools/src/json_schema.rs | 75 +++++++++++++++++--- codex-rs/tools/src/json_schema_tests.rs | 93 ++++++++++++++++++++++++- 2 files changed, 158 insertions(+), 10 deletions(-) diff --git a/codex-rs/tools/src/json_schema.rs b/codex-rs/tools/src/json_schema.rs index 22a641491e..29c5c18f22 100644 --- a/codex-rs/tools/src/json_schema.rs +++ b/codex-rs/tools/src/json_schema.rs @@ -147,8 +147,9 @@ impl From for AdditionalProperties { /// Parse the tool `input_schema` or return an error for invalid schema. pub fn parse_tool_input_schema(input_schema: &JsonValue) -> Result { - 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 { // 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, + root_schema: &JsonValue, +) -> Option { + 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) -> Option { + 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, schema_types: &[JsonSchemaPrimitiveType], diff --git a/codex-rs/tools/src/json_schema_tests.rs b/codex-rs/tools/src/json_schema_tests.rs index 3f13df7638..915ba6e81d 100644 --- a/codex-rs/tools/src/json_schema_tests.rs +++ b/codex-rs/tools/src/json_schema_tests.rs @@ -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 )