From 8dff89c8d145deea1149af91f6414b869c5c2ef6 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 19 May 2026 15:04:48 -0400 Subject: [PATCH] Fix Gemini schema conversion dropping anyOf and nullability (#57160) Follow-up to #49292: - **Bug fix in `push_any_of_constraint`**: when both an `anyOf` and a non-empty `allOf` were already present at the same level, the existing `anyOf` was silently dropped. Now it's always preserved. - **Canonical OpenAPI nullability**: collapse `{nullable: true}`-only entries out of `anyOf` onto the parent so `anyOf: [{type: "string"}, {type: "null"}]` becomes `{type: "string", nullable: true}` (the form Gemini actually expects) instead of `anyOf: [{type: "string"}, {nullable: true}]`. - **Compile-time path consistency**: route `ToJsonSchemaSubsetTransform` through the same `convert_null_in_types_to_nullable` helper so Rust-defined tools using `Option` also keep nullability on Gemini, instead of silently truncating to the non-null type. - Drop an unnecessary clone in `push_any_of_constraint` and simplify `convert_types_to_any_of_defs`. - Add a small `obj()` test helper and a regression test for the `anyOf + allOf + multi-type` case. Release Notes: - Fixed `Option` tool parameters being sent to Gemini without their nullability, and fixed tool schemas with `anyOf` + `allOf` losing constraints during the OpenAPI 3.0 conversion Co-authored-by: Daniel Strobusch <1438302+dastrobu@users.noreply.github.com> --- crates/language_model_core/src/tool_schema.rs | 363 ++++++++++-------- 1 file changed, 209 insertions(+), 154 deletions(-) diff --git a/crates/language_model_core/src/tool_schema.rs b/crates/language_model_core/src/tool_schema.rs index 729d625939..13e6e66524 100644 --- a/crates/language_model_core/src/tool_schema.rs +++ b/crates/language_model_core/src/tool_schema.rs @@ -39,21 +39,31 @@ struct ToJsonSchemaSubsetTransform; impl Transform for ToJsonSchemaSubsetTransform { fn transform(&mut self, schema: &mut Schema) { - // Ensure that the type field is not an array, this happens when we use - // Option, the type will be [T, "null"]. - if let Some(type_field) = schema.get_mut("type") - && let Some(types) = type_field.as_array() - && let Some(first_type) = types.first() - { - *type_field = first_type.clone(); - } + if let Some(obj) = schema.as_object_mut() { + // `Option` produces `type: [T, "null"]`. Convert to OpenAPI 3.0's + // `nullable: true` so nullability isn't silently dropped. + convert_null_in_types_to_nullable(obj); - // oneOf is not supported, use anyOf instead - if let Some(one_of) = schema.remove("oneOf") { - schema.insert("anyOf".to_string(), one_of); + // Any remaining multi-type array (uncommon in Rust-generated schemas) + // is collapsed to its first entry to keep this schema subset-compatible. + if let Some(type_field) = obj.get_mut("type") + && let Some(types) = type_field.as_array() + && let Some(first_type) = types.first().cloned() + { + *type_field = first_type; + } + + // oneOf is not supported, use anyOf instead + if let Some(one_of) = obj.remove("oneOf") { + obj.insert("anyOf".to_string(), one_of); + } } transform_subschemas(self, schema); + + if let Some(obj) = schema.as_object_mut() { + collapse_nullable_only_any_of(obj); + } } } @@ -126,14 +136,16 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { convert_null_in_types_to_nullable(obj); convert_types_to_any_of_defs(obj); - // Ensure that the type field is not an array. This can happen with MCP tool - // schemas that use multiple types (e.g. `["string", "number"]` or `["string", "null"]`). + // After the conversions above, `type` should only ever still be an array + // if a malformed input had a single-element type array (e.g. `["string"]`). + // Collapse it to a single value so downstream consumers see a scalar. if let Some(type_value) = obj.get_mut("type") && let Some(types) = type_value.as_array() && let Some(first_type) = types.first().cloned() { *type_value = first_type; } + if matches!(obj.get("description"), Some(Value::String(_))) && !obj.contains_key("type") && !(obj.contains_key("anyOf") @@ -156,6 +168,11 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { adapt_to_json_schema_subset(value)?; } } + + // Children may have been rewritten from `{"type": "null"}` into + // `{"nullable": true}`. Fold those into the parent so the result matches + // OpenAPI 3.0's convention of `nullable: true` as a sibling of `type`. + collapse_nullable_only_any_of(obj); } else if let Value::Array(arr) = json { for item in arr.iter_mut() { adapt_to_json_schema_subset(item)?; @@ -200,30 +217,30 @@ fn convert_null_in_types_to_nullable(obj: &mut Map) { } fn convert_types_to_any_of_defs(obj: &mut Map) { - if let Some(type_entry) = obj.get_mut("type") { - if let Some(types) = type_entry.as_array_mut() { - if types.len() > 1 { - let remaining_types = std::mem::take(types); - let mut any_of_schemas = Vec::new(); - for t in remaining_types { - any_of_schemas.push(json!({"type": t})); - } - obj.remove("type"); - push_any_of_constraint(obj, Value::Array(any_of_schemas)); - } - } + let is_multi_type = obj + .get("type") + .and_then(|v| v.as_array()) + .is_some_and(|types| types.len() > 1); + if !is_multi_type { + return; } + + let Some(Value::Array(types)) = obj.remove("type") else { + return; + }; + let any_of_schemas = types.into_iter().map(|t| json!({"type": t})).collect(); + push_any_of_constraint(obj, Value::Array(any_of_schemas)); } fn push_any_of_constraint(obj: &mut Map, any_of_schemas: Value) { if let Some(existing_any_of) = obj.remove("anyOf") { - let mut all_of = obj - .remove("allOf") - .and_then(|v| v.as_array().cloned()) - .unwrap_or_default(); - if all_of.is_empty() { - all_of.push(json!({"anyOf": existing_any_of})); - } + let mut all_of = match obj.remove("allOf") { + Some(Value::Array(arr)) => arr, + _ => Vec::new(), + }; + // Always preserve the pre-existing `anyOf` — earlier this push was + // skipped when `allOf` was non-empty, which silently dropped it. + all_of.push(json!({"anyOf": existing_any_of})); all_of.push(json!({"anyOf": any_of_schemas})); obj.insert("allOf".to_string(), Value::Array(all_of)); } else if let Some(all_of) = obj.get_mut("allOf").and_then(|v| v.as_array_mut()) { @@ -233,189 +250,187 @@ fn push_any_of_constraint(obj: &mut Map, any_of_schemas: Value) { } } +/// Folds `{nullable: true}`-only entries out of an `anyOf` array and onto the +/// parent object. This matches OpenAPI 3.0 semantics, where nullability is +/// expressed as a sibling of `type` rather than a separate variant. +fn collapse_nullable_only_any_of(obj: &mut Map) { + let Some(Value::Array(mut any_of)) = obj.remove("anyOf") else { + return; + }; + + let mut found_nullable_only = false; + any_of.retain(|entry| { + let is_nullable_only = entry + .as_object() + .is_some_and(|m| m.len() == 1 && matches!(m.get("nullable"), Some(Value::Bool(true)))); + if is_nullable_only { + found_nullable_only = true; + false + } else { + true + } + }); + + if !found_nullable_only { + obj.insert("anyOf".to_string(), Value::Array(any_of)); + return; + } + + obj.insert("nullable".to_string(), Value::Bool(true)); + + if any_of.is_empty() { + return; + } + + // If a single variant remains and its keys don't collide with the parent's + // existing keys, inline it. `anyOf` with a single entry is equivalent to + // just that entry, and inlining produces the canonical OpenAPI form + // (e.g. `{type: "string", nullable: true}`). + if any_of.len() == 1 + && let Value::Object(entry_obj) = &any_of[0] + && entry_obj.keys().all(|k| !obj.contains_key(k)) + { + let entry = any_of.remove(0); + if let Value::Object(entry_obj) = entry { + for (k, v) in entry_obj { + obj.insert(k, v); + } + } + return; + } + + obj.insert("anyOf".to_string(), Value::Array(any_of)); +} + #[cfg(test)] mod tests { use super::*; use serde_json::json; + fn obj(value: Value) -> Map { + match value { + Value::Object(map) => map, + other => panic!("expected JSON object, got {other}"), + } + } + #[test] fn test_convert_null_in_types_to_nullable() { // ["string", "null"] -> "string", nullable: true - let mut obj = json!({"type": ["string", "null"]}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_null_in_types_to_nullable(&mut obj); - assert_eq!( - obj, - json!({"type": "string", "nullable": true}) - .as_object() - .unwrap() - .to_owned() - ); + let mut o = obj(json!({"type": ["string", "null"]})); + convert_null_in_types_to_nullable(&mut o); + assert_eq!(o, obj(json!({"type": "string", "nullable": true}))); // "null" -> nullable: true - let mut obj = json!({"type": "null"}).as_object_mut().unwrap().to_owned(); - convert_null_in_types_to_nullable(&mut obj); - assert_eq!( - obj, - json!({"nullable": true}).as_object().unwrap().to_owned() - ); + let mut o = obj(json!({"type": "null"})); + convert_null_in_types_to_nullable(&mut o); + assert_eq!(o, obj(json!({"nullable": true}))); // ["string", "number", "null"] -> ["string", "number"], nullable: true (anyOf handled elsewhere) - let mut obj = json!({"type": ["string", "number", "null"]}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_null_in_types_to_nullable(&mut obj); + let mut o = obj(json!({"type": ["string", "number", "null"]})); + convert_null_in_types_to_nullable(&mut o); assert_eq!( - obj, - json!({"type": ["string", "number"], "nullable": true}) - .as_object() - .unwrap() - .to_owned() + o, + obj(json!({"type": ["string", "number"], "nullable": true})) ); // "string" (no change, not nullable) - let mut obj = json!({"type": "string"}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_null_in_types_to_nullable(&mut obj); - assert_eq!( - obj, - json!({"type": "string"}).as_object().unwrap().to_owned() - ); + let mut o = obj(json!({"type": "string"})); + convert_null_in_types_to_nullable(&mut o); + assert_eq!(o, obj(json!({"type": "string"}))); // ["string", "number"] (no change, not nullable) - let mut obj = json!({"type": ["string", "number"]}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_null_in_types_to_nullable(&mut obj); - assert_eq!( - obj, - json!({"type": ["string", "number"]}) - .as_object() - .unwrap() - .to_owned() - ); + let mut o = obj(json!({"type": ["string", "number"]})); + convert_null_in_types_to_nullable(&mut o); + assert_eq!(o, obj(json!({"type": ["string", "number"]}))); // object with other properties, ["boolean", "null"] - let mut obj = json!({ + let mut o = obj(json!({ "description": "A test field", "type": ["boolean", "null"] - }) - .as_object_mut() - .unwrap() - .to_owned(); - convert_null_in_types_to_nullable(&mut obj); + })); + convert_null_in_types_to_nullable(&mut o); assert_eq!( - obj, - json!({ + o, + obj(json!({ "description": "A test field", "type": "boolean", "nullable": true - }) - .as_object() - .unwrap() - .to_owned() + })) ); } #[test] fn test_convert_types_to_any_of_defs() { // ["string", "number"] -> anyOf with string and number - let mut obj = json!({"type": ["string", "number"]}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_types_to_any_of_defs(&mut obj); + let mut o = obj(json!({"type": ["string", "number"]})); + convert_types_to_any_of_defs(&mut o); assert_eq!( - obj, - json!({ + o, + obj(json!({ "anyOf": [ {"type": "string"}, {"type": "number"} ] - }) - .as_object() - .unwrap() - .to_owned() + })) ); // "string" (no change) - let mut obj = json!({"type": "string"}) - .as_object_mut() - .unwrap() - .to_owned(); - convert_types_to_any_of_defs(&mut obj); - assert_eq!( - obj, - json!({"type": "string"}).as_object().unwrap().to_owned() - ); + let mut o = obj(json!({"type": "string"})); + convert_types_to_any_of_defs(&mut o); + assert_eq!(o, obj(json!({"type": "string"}))); + + // single-element array (no change, fallback in caller collapses it) + let mut o = obj(json!({"type": ["string"]})); + convert_types_to_any_of_defs(&mut o); + assert_eq!(o, obj(json!({"type": ["string"]}))); // object with other properties, ["string", "number"] - let mut obj = json!({ + let mut o = obj(json!({ "description": "A test field", "type": ["string", "number"] - }) - .as_object_mut() - .unwrap() - .to_owned(); - convert_types_to_any_of_defs(&mut obj); + })); + convert_types_to_any_of_defs(&mut o); assert_eq!( - obj, - json!({ + o, + obj(json!({ "description": "A test field", "anyOf": [ {"type": "string"}, {"type": "number"} ] - }) - .as_object() - .unwrap() - .to_owned() + })) ); // anyOf already present (no change) - let mut obj = json!({ + let mut o = obj(json!({ "anyOf": [ {"type": "string"}, {"type": "number"} ] - }) - .as_object_mut() - .unwrap() - .to_owned(); - convert_types_to_any_of_defs(&mut obj); + })); + convert_types_to_any_of_defs(&mut o); assert_eq!( - obj, - json!({ + o, + obj(json!({ "anyOf": [ {"type": "string"}, {"type": "number"} ] - }) - .as_object() - .unwrap() - .to_owned() + })) ); // both type array and anyOf present - let mut obj = json!({ + let mut o = obj(json!({ "type": ["string", "number"], "anyOf": [ {"format": "email"} ] - }) - .as_object_mut() - .unwrap() - .to_owned(); - convert_types_to_any_of_defs(&mut obj); + })); + convert_types_to_any_of_defs(&mut o); assert_eq!( - obj, - json!({ + o, + obj(json!({ "allOf": [ { "anyOf": [ @@ -429,10 +444,26 @@ mod tests { ] } ] - }) - .as_object() - .unwrap() - .to_owned() + })) + ); + + // type array + anyOf + pre-existing allOf: pre-existing anyOf must not + // be silently dropped just because allOf is non-empty. + let mut o = obj(json!({ + "type": ["string", "number"], + "anyOf": [{"minLength": 5}], + "allOf": [{"maxLength": 100}] + })); + convert_types_to_any_of_defs(&mut o); + assert_eq!( + o, + obj(json!({ + "allOf": [ + {"maxLength": 100}, + {"anyOf": [{"minLength": 5}]}, + {"anyOf": [{"type": "string"}, {"type": "number"}]} + ] + })) ); } @@ -544,12 +575,38 @@ mod tests { adapt_to_json_schema_subset(&mut json).unwrap(); + // Should collapse to the canonical OpenAPI 3.0 form: `nullable: true` + // as a sibling of `type`, rather than a separate anyOf variant. assert_eq!( json, json!({ + "type": "string", + "nullable": true + }) + ); + } + + #[test] + fn test_transform_null_in_any_of_with_multiple_non_null_variants() { + let mut json = json!({ + "anyOf": [ + { "type": "string" }, + { "type": "number" }, + { "type": "null" } + ] + }); + + adapt_to_json_schema_subset(&mut json).unwrap(); + + // When more than one non-null variant remains, keep the anyOf and just + // hoist `nullable: true` onto the parent. + assert_eq!( + json, + json!({ + "nullable": true, "anyOf": [ { "type": "string" }, - { "nullable": true } + { "type": "number" } ] }) ); @@ -643,10 +700,8 @@ mod tests { "type": "object", "properties": { "nested": { - "anyOf": [ - { "type": "string" }, - { "nullable": true } - ], + "type": "string", + "nullable": true } } })